feat(builders): add ConsistentPartitionWindowRel builder

Add builder, type inference, display, and comprehensive tests for
ConsistentPartitionWindowRel, enabling window function plans with
shared partitioning and ordering.

Author: bvolpato-dd

src/substrait/builders/plan.py

@@ -498,3 +498,97 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         )
 
     return resolve
+
+
+def consistent_partition_window(
+    plan: PlanOrUnbound,
+    window_functions: Iterable[ExtendedExpressionOrUnbound],
+    partition_expressions: Iterable[ExtendedExpressionOrUnbound] = (),
+    sorts: Iterable[
+        Union[
+            ExtendedExpressionOrUnbound,
+            tuple[ExtendedExpressionOrUnbound, stalg.SortField.SortDirection.ValueType],
+        ]
+    ] = (),
+    extension: Optional[AdvancedExtension] = None,
+) -> UnboundPlan:
+    def resolve(registry: ExtensionRegistry) -> stp.Plan:
+        bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry)
+        ns = infer_plan_schema(bound_plan)
+
+        bound_partitions = [
+            resolve_expression(e, ns, registry) for e in partition_expressions
+        ]
+
+        bound_sorts = [
+            (e, stalg.SortField.SORT_DIRECTION_ASC_NULLS_LAST)
+            if not isinstance(e, tuple)
+            else e
+            for e in sorts
+        ]
+        bound_sorts = [
+            (resolve_expression(e[0], ns, registry), e[1]) for e in bound_sorts
+        ]
+
+        bound_window_fns = [
+            resolve_expression(e, ns, registry) for e in window_functions
+        ]
+
+        window_rel_functions = []
+        for wf_ee in bound_window_fns:
+            wf_expr = wf_ee.referred_expr[0].expression.window_function
+            window_rel_functions.append(
+                stalg.ConsistentPartitionWindowRel.WindowRelFunction(
+                    function_reference=wf_expr.function_reference,
+                    arguments=list(wf_expr.arguments),
+                    options=list(wf_expr.options),
+                    output_type=wf_expr.output_type,
+                    phase=wf_expr.phase,
+                    invocation=wf_expr.invocation,
+                    lower_bound=wf_expr.lower_bound
+                    if wf_expr.HasField("lower_bound")
+                    else None,
+                    upper_bound=wf_expr.upper_bound
+                    if wf_expr.HasField("upper_bound")
+                    else None,
+                    bounds_type=wf_expr.bounds_type,
+                )
+            )
+
+        names = list(bound_plan.relations[-1].root.names) + [
+            wf_ee.referred_expr[0].output_names[0]
+            if wf_ee.referred_expr[0].output_names
+            else f"window_{i}"
+            for i, wf_ee in enumerate(bound_window_fns)
+        ]
+
+        rel = stalg.Rel(
+            window=stalg.ConsistentPartitionWindowRel(
+                input=bound_plan.relations[-1].root.input,
+                window_functions=window_rel_functions,
+                partition_expressions=[
+                    e.referred_expr[0].expression for e in bound_partitions
+                ],
+                sorts=[
+                    stalg.SortField(
+                        expr=e[0].referred_expr[0].expression,
+                        direction=e[1],
+                    )
+                    for e in bound_sorts
+                ],
+                advanced_extension=extension,
+            )
+        )
+
+        return stp.Plan(
+            version=default_version,
+            relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
+            **_merge_extensions(
+                bound_plan,
+                *bound_partitions,
+                *[e[0] for e in bound_sorts],
+                *bound_window_fns,
+            ),
+        )
+
+    return resolve

src/substrait/type_inference.py

@@ -342,6 +342,14 @@ def infer_rel_schema(rel: stalg.Rel) -> stt.Type.Struct:
             raise Exception(f"Unhandled join_type {rel.join.type}")
 
         (common, struct) = (rel.join.common, raw_schema)
+    elif rel_type == "window":
+        parent_schema = infer_rel_schema(rel.window.input)
+        window_output_types = [wf.output_type for wf in rel.window.window_functions]
+        raw_schema = stt.Type.Struct(
+            types=list(parent_schema.types) + window_output_types,
+            nullability=parent_schema.nullability,
+        )
+        (common, struct) = (rel.window.common, raw_schema)
     else:
         raise Exception(f"Unhandled rel_type {rel_type}")
 

src/substrait/utils/display.py

@@ -171,6 +171,8 @@ def _stream_rel(self, rel: stalg.Rel, stream, depth: int):
             self._stream_extension_single_rel(rel.extension_single, stream, depth)
         elif rel.HasField("extension_multi"):
             self._stream_extension_multi_rel(rel.extension_multi, stream, depth)
+        elif rel.HasField("window"):
+            self._stream_window_rel(rel.window, stream, depth)
         else:
             stream.write(f"{indent}<unknown_relation>\n")
 
@@ -401,6 +403,43 @@ def _stream_extension_multi_rel(
                     f"{self._get_indent_with_arrow(depth + 2)}<unpackable_detail>\n"
                 )
 
+    def _stream_window_rel(
+        self, window: stalg.ConsistentPartitionWindowRel, stream, depth: int
+    ):
+        """Print a consistent partition window relation concisely"""
+        indent = " " * (depth * self.indent_size)
+
+        stream.write(
+            f"{indent}{self._color('window', Colors.MAGENTA)}: "
+            f"{self._color(str(len(window.window_functions)), Colors.YELLOW)} functions\n"
+        )
+        stream.write(
+            f"{self._get_indent_with_arrow(depth + 1)}{self._color('input:', Colors.BLUE)}\n"
+        )
+        self._stream_rel(window.input, stream, depth + 1)
+
+        if window.partition_expressions:
+            stream.write(
+                f"{self._get_indent_with_arrow(depth + 1)}"
+                f"{self._color('partitions:', Colors.BLUE)} "
+                f"{self._color(str(len(window.partition_expressions)), Colors.YELLOW)}\n"
+            )
+
+        if window.sorts:
+            stream.write(
+                f"{self._get_indent_with_arrow(depth + 1)}"
+                f"{self._color('sorts:', Colors.BLUE)} "
+                f"{self._color(str(len(window.sorts)), Colors.YELLOW)}\n"
+            )
+
+        for i, wf in enumerate(window.window_functions):
+            stream.write(
+                f"{self._get_indent_with_arrow(depth + 1)}"
+                f"{self._color('window_fn', Colors.BLUE)}"
+                f"[{self._color(str(i), Colors.CYAN)}]: "
+                f"func_ref={wf.function_reference}\n"
+            )
+
     def _stream_expression(self, expression: stalg.Expression, stream, depth: int):
         """Print an expression concisely"""
         indent = " " * (depth * self.indent_size)