refactor: scope extension-relation derivers to the ExtensionRegistry

src/substrait/builders/extended_expression.py

@@ -447,7 +447,8 @@ def resolve(
         ]
 
         expression_schemas = [
-            infer_extended_expression_schema(b) for b in bound_expressions
+            infer_extended_expression_schema(b, registry=registry)
+            for b in bound_expressions
         ]
 
         signature = [typ for es in expression_schemas for typ in es.types]
@@ -543,7 +544,8 @@ def resolve(
         ]
 
         expression_schemas = [
-            infer_extended_expression_schema(b) for b in bound_expressions
+            infer_extended_expression_schema(b, registry=registry)
+            for b in bound_expressions
         ]
 
         signature = [typ for es in expression_schemas for typ in es.types]
@@ -640,7 +642,8 @@ def resolve(
         ]
 
         expression_schemas = [
-            infer_extended_expression_schema(b) for b in bound_expressions
+            infer_extended_expression_schema(b, registry=registry)
+            for b in bound_expressions
         ]
 
         signature = [typ for es in expression_schemas for typ in es.types]

src/substrait/builders/plan.py

@@ -203,7 +203,7 @@ def project(
 
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         _plan = plan if isinstance(plan, stp.Plan) else plan(registry)
-        ns = infer_plan_schema(_plan)
+        ns = infer_plan_schema(_plan, registry=registry)
         bound_expressions: Iterable[stee.ExtendedExpression] = [
             resolve_expression(e, ns, registry) for e in expressions
         ]
@@ -252,7 +252,7 @@ def select(
 
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         _plan = plan if isinstance(plan, stp.Plan) else plan(registry)
-        ns = infer_plan_schema(_plan)
+        ns = infer_plan_schema(_plan, registry=registry)
         bound_expressions: Iterable[stee.ExtendedExpression] = [
             resolve_expression(e, ns, registry) for e in expressions
         ]
@@ -294,7 +294,7 @@ def filter(
 ) -> UnboundPlan:
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry)
-        ns = infer_plan_schema(bound_plan)
+        ns = infer_plan_schema(bound_plan, registry=registry)
         bound_expression: stee.ExtendedExpression = resolve_expression(
             expression, ns, registry
         )
@@ -330,7 +330,7 @@ def sort(
 ) -> UnboundPlan:
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry)
-        ns = infer_plan_schema(bound_plan)
+        ns = infer_plan_schema(bound_plan, registry=registry)
 
         bound_expressions = [
             (e, stalg.SortField.SORT_DIRECTION_ASC_NULLS_LAST)
@@ -397,7 +397,7 @@ def fetch(
 ) -> UnboundPlan:
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry)
-        ns = infer_plan_schema(bound_plan)
+        ns = infer_plan_schema(bound_plan, registry=registry)
 
         bound_offset = resolve_expression(offset, ns, registry) if offset else None
         # count=None means "all remaining rows" (FetchRel leaves count_expr unset).
@@ -445,8 +445,8 @@ def join(
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_left = left if isinstance(left, stp.Plan) else left(registry)
         bound_right = right if isinstance(right, stp.Plan) else right(registry)
-        left_ns = infer_plan_schema(bound_left)
-        right_ns = infer_plan_schema(bound_right)
+        left_ns = infer_plan_schema(bound_left, registry=registry)
+        right_ns = infer_plan_schema(bound_right, registry=registry)
 
         ns = stt.NamedStruct(
             struct=stt.Type.Struct(
@@ -500,8 +500,8 @@ def cross(
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_left = left if isinstance(left, stp.Plan) else left(registry)
         bound_right = right if isinstance(right, stp.Plan) else right(registry)
-        left_ns = infer_plan_schema(bound_left)
-        right_ns = infer_plan_schema(bound_right)
+        left_ns = infer_plan_schema(bound_left, registry=registry)
+        right_ns = infer_plan_schema(bound_right, registry=registry)
 
         ns = stt.NamedStruct(
             struct=stt.Type.Struct(
@@ -548,7 +548,7 @@ def aggregate(
 
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_input = input if isinstance(input, stp.Plan) else input(registry)
-        ns = infer_plan_schema(bound_input)
+        ns = infer_plan_schema(bound_input, registry=registry)
 
         bound_grouping_expressions = [
             resolve_expression(e, ns, registry) for e in grouping_expressions
@@ -618,7 +618,7 @@ def write_named_table(
 ) -> UnboundPlan:
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_input = input if isinstance(input, stp.Plan) else input(registry)
-        ns = infer_plan_schema(bound_input)
+        ns = infer_plan_schema(bound_input, registry=registry)
         _table_names = [table_names] if isinstance(table_names, str) else table_names
         _create_mode = create_mode or stalg.WriteRel.CREATE_MODE_ERROR_IF_EXISTS
         _op = op if op is not None else stalg.WriteRel.WRITE_OP_CTAS
@@ -673,7 +673,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             view_rel = view_plan.relations[-1].root.input
             merge_sources.append(view_plan)
             if schema is None:
-                schema = infer_plan_schema(view_plan)
+                schema = infer_plan_schema(view_plan, registry=registry)
 
         ddl_rel = stalg.Rel(
             ddl=stalg.DdlRel(
@@ -764,7 +764,7 @@ def consistent_partition_window(
 ) -> UnboundPlan:
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry)
-        ns = infer_plan_schema(bound_plan)
+        ns = infer_plan_schema(bound_plan, registry=registry)
 
         bound_partitions = [
             resolve_expression(e, ns, registry) for e in partition_expressions
@@ -859,7 +859,7 @@ def expand(
 
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_input = plan if isinstance(plan, stp.Plan) else plan(registry)
-        ns = infer_plan_schema(bound_input)
+        ns = infer_plan_schema(bound_input, registry=registry)
 
         expand_fields = []
         merge_sources = [bound_input]
@@ -910,8 +910,8 @@ def nested_loop_join(
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_left = left if isinstance(left, stp.Plan) else left(registry)
         bound_right = right if isinstance(right, stp.Plan) else right(registry)
-        left_ns = infer_plan_schema(bound_left)
-        right_ns = infer_plan_schema(bound_right)
+        left_ns = infer_plan_schema(bound_left, registry=registry)
+        right_ns = infer_plan_schema(bound_right, registry=registry)
 
         ns = stt.NamedStruct(
             struct=stt.Type.Struct(
@@ -987,8 +987,8 @@ def builder(
         def resolve(registry: ExtensionRegistry) -> stp.Plan:
             bound_left = left if isinstance(left, stp.Plan) else left(registry)
             bound_right = right if isinstance(right, stp.Plan) else right(registry)
-            left_ns = infer_plan_schema(bound_left)
-            right_ns = infer_plan_schema(bound_right)
+            left_ns = infer_plan_schema(bound_left, registry=registry)
+            right_ns = infer_plan_schema(bound_right, registry=registry)
             keys = _comparison_join_keys(
                 list(left_keys), list(right_keys), left_ns, right_ns, registry
             )
@@ -1061,7 +1061,7 @@ def extension_single(plan: PlanOrUnbound, detail) -> UnboundPlan:
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry)
         if hasattr(detail, "derive_schema"):
-            input_struct = infer_plan_schema(bound_plan).struct
+            input_struct = infer_plan_schema(bound_plan, registry=registry).struct
             names = list(detail.derive_schema(input_struct).names)
         else:
             names = list(bound_plan.relations[-1].root.names)
@@ -1084,7 +1084,9 @@ def extension_multi(inputs: Iterable[PlanOrUnbound], detail) -> UnboundPlan:
 
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_inputs = [i if isinstance(i, stp.Plan) else i(registry) for i in inputs]
-        input_structs = [infer_plan_schema(b).struct for b in bound_inputs]
+        input_structs = [
+            infer_plan_schema(b, registry=registry).struct for b in bound_inputs
+        ]
         names = list(detail.derive_schema(input_structs).names)
         rel = stalg.Rel(
             extension_multi=stalg.ExtensionMultiRel(
@@ -1159,7 +1161,7 @@ def top_n(
 
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry)
-        ns = infer_plan_schema(bound_plan)
+        ns = infer_plan_schema(bound_plan, registry=registry)
         bound_sorts = [
             (resolve_expression(e, ns, registry), direction) for e, direction in sorts
         ]

src/substrait/dataframe/expr.py

@@ -242,7 +242,9 @@ def _resolve_over_urns(
     identically and their error text cannot drift apart.
     """
     signature = [
-        typ for b in bound for typ in infer_extended_expression_schema(b).types
+        typ
+        for b in bound
+        for typ in infer_extended_expression_schema(b, registry=registry).types
     ]
     for urn in urns:
         if registry.lookup_function(urn, name, signature):
@@ -293,9 +295,13 @@ def bind(operand):
 
         peer = None
         if left_is_expr:
-            peer = infer_extended_expression_schema(left_val).types[0]
+            peer = infer_extended_expression_schema(left_val, registry=registry).types[
+                0
+            ]
         elif right_is_expr:
-            peer = infer_extended_expression_schema(right_val).types[0]
+            peer = infer_extended_expression_schema(right_val, registry=registry).types[
+                0
+            ]
 
         def as_bound(value, is_expr):
             if is_expr:
@@ -635,7 +641,9 @@ def _higher_order(self, function: str, callback) -> "Expr":
         def resolve(base_schema, registry):
             bound_list = list_unbound(base_schema, registry)
             element_type = (
-                infer_extended_expression_schema(bound_list).types[0].list.type
+                infer_extended_expression_schema(bound_list, registry=registry)
+                .types[0]
+                .list.type
             )
             param_struct = stp.Type.Struct(
                 types=[element_type], nullability=stp.Type.NULLABILITY_REQUIRED

src/substrait/dataframe/frame.py

@@ -215,7 +215,7 @@ def rename(self, mapping: dict) -> "DataFrame":
 
         def resolve(registry: ExtensionRegistry):
             bound = inner(registry)
-            names = list(infer_plan_schema(bound).names)
+            names = list(infer_plan_schema(bound, registry=registry).names)
             unknown = set(mapping) - set(names)
             if unknown:
                 raise ValueError(f"rename got unknown columns: {sorted(unknown)}")
@@ -234,7 +234,7 @@ def drop(self, *columns: str) -> "DataFrame":
 
         def resolve(registry: ExtensionRegistry):
             bound = inner(registry)
-            names = list(infer_plan_schema(bound).names)
+            names = list(infer_plan_schema(bound, registry=registry).names)
             unknown = drop_set - set(names)
             if unknown:
                 raise ValueError(f"drop got unknown columns: {sorted(unknown)}")

src/substrait/extension_registry/registry.py

@@ -53,13 +53,15 @@ def register_extension_relation(self, detail_cls) -> None:
         Enables schema inference for ``ExtensionLeaf/Single/MultiRel`` built with
         an instance of ``detail_cls``: inference reconstructs the detail from the
         plan's ``Any`` and calls its ``derive_schema``. See
-        :mod:`substrait.dataframe.extension_relations`. Registration is process-global
-        (type_urls are globally unique), so inference works on any plan.
+        :mod:`substrait.dataframe.extension_relations`. Registration is scoped to
+        this ``ExtensionRegistry`` instance, so inference must be given this same
+        registry (as it is when a plan is built or re-inferred through it).
         """
-        from substrait.type_inference import register_extension_relation
-
         self._extension_relations[detail_cls.type_url] = detail_cls
-        register_extension_relation(detail_cls)
+
+    def lookup_extension_relation(self, type_url: str):
+        """The extension-relation detail class registered for ``type_url``, or None."""
+        return self._extension_relations.get(type_url)
 
     def register_extension_dict(self, definitions: dict) -> None:
         """Register extensions from a dictionary (parsed YAML).