feat: add support for pyspark equivalent explode_outer, posexplode/posexplode_outer (via explode_with_index)

Author: bcallender

protos/logical_plan/v1/plans.proto

@@ -42,6 +42,7 @@ message LogicalPlan {
     Union union = 24;
     Limit limit = 25;
     Explode explode = 26;
+    ExplodeWithIndex explode_with_index = 34;
     DropDuplicates drop_duplicates = 27;
     Sort sort = 28;
     Unnest unnest = 29;
@@ -143,6 +144,15 @@ message Limit {
 message Explode {
   LogicalPlan input = 1;
   LogicalExpr expr = 2;
+  bool keep_null_and_empty = 3;
+}
+
+message ExplodeWithIndex {
+  LogicalPlan input = 1;
+  LogicalExpr expr = 2;
+  string index_name = 3;
+  optional string value_name = 4;
+  bool keep_null_and_empty = 5;
 }
 
 message DropDuplicates {

src/fenic/_backends/local/physical_plan/__init__.py

@@ -37,6 +37,9 @@
     DropDuplicatesExec as DropDuplicatesExec,
 )
 from fenic._backends.local.physical_plan.transform import ExplodeExec as ExplodeExec
+from fenic._backends.local.physical_plan.transform import (
+    ExplodeWithIndexExec as ExplodeWithIndexExec,
+)
 from fenic._backends.local.physical_plan.transform import FilterExec as FilterExec
 from fenic._backends.local.physical_plan.transform import LimitExec as LimitExec
 from fenic._backends.local.physical_plan.transform import (
@@ -65,6 +68,7 @@
     "InMemorySourceExec",
     "DropDuplicatesExec",
     "ExplodeExec",
+    "ExplodeWithIndexExec",
     "FilterExec",
     "LimitExec",
     "ProjectionExec",

src/fenic/_backends/local/physical_plan/transform.py

@@ -174,21 +174,25 @@ def __init__(
         child: PhysicalPlan,
         physical_expr: pl.Expr,
         col_name: str,
+        keep_null_and_empty: bool,
         cache_info: Optional[CacheInfo],
         session_state: LocalSessionState,
     ):
         super().__init__([child], cache_info=cache_info, session_state=session_state)
         self.physical_expr = physical_expr
         self.col_name = col_name
+        self.keep_null_and_empty = keep_null_and_empty
 
     def execute_node(self, child_dfs: List[pl.DataFrame]) -> pl.DataFrame:
         if len(child_dfs) != 1:
             raise ValueError("Unreachable: ExplodeExec expects 1 child")
         child_df = child_dfs[0]
         child_df = child_df.with_columns(self.physical_expr)
         exploded_df = child_df.explode(self.col_name)
-        # Optionally filter out rows where the exploded column is null.
-        return exploded_df.filter(pl.col(self.col_name).is_not_null())
+        # Filter out nulls unless keep_null_and_empty is True
+        if not self.keep_null_and_empty:
+            exploded_df = exploded_df.filter(pl.col(self.col_name).is_not_null())
+        return exploded_df
 
     def with_children(self, children: List[PhysicalPlan]) -> PhysicalPlan:
         if len(children) != 1:
@@ -197,6 +201,7 @@ def with_children(self, children: List[PhysicalPlan]) -> PhysicalPlan:
             child=children[0],
             physical_expr=self.physical_expr,
             col_name=self.col_name,
+            keep_null_and_empty=self.keep_null_and_empty,
             cache_info=self.cache_info,
             session_state=self.session_state,
         )
@@ -207,6 +212,109 @@ def build_node_lineage(
     ) -> Tuple[OperatorLineage, pl.DataFrame]:
         child_operator, child_df = self.children[0].build_node_lineage(leaf_nodes)
         exploded_df = child_df.explode(self.col_name)
+        # Filter out nulls unless keep_null_and_empty is True
+        if not self.keep_null_and_empty:
+            exploded_df = exploded_df.filter(pl.col(self.col_name).is_not_null())
+        exploded_df = exploded_df.with_columns(
+            pl.col("_uuid").alias("_backwards_uuid"),
+        )
+        exploded_df = _with_lineage_uuid(exploded_df)
+        backwards_df = exploded_df.select(["_uuid", "_backwards_uuid"])
+
+        materialize_df = exploded_df.drop("_backwards_uuid")
+
+        operator = self._build_unary_operator_lineage(
+            materialize_df=materialize_df,
+            child=(child_operator, backwards_df),
+        )
+
+        return operator, materialize_df
+
+
+class ExplodeWithIndexExec(PhysicalPlan):
+    def __init__(
+        self,
+        child: PhysicalPlan,
+        physical_expr: pl.Expr,
+        col_name: str,
+        index_name: str,
+        value_name: str,
+        keep_null_and_empty: bool,
+        cache_info: Optional[CacheInfo],
+        session_state: LocalSessionState,
+    ):
+        super().__init__([child], cache_info=cache_info, session_state=session_state)
+        self.physical_expr = physical_expr
+        self.col_name = col_name
+        self.index_name = index_name
+        self.value_name = value_name
+        self.keep_null_and_empty = keep_null_and_empty
+
+    def execute_node(self, child_dfs: List[pl.DataFrame]) -> pl.DataFrame:
+        if len(child_dfs) != 1:
+            raise ValueError("Unreachable: ExplodeWithIndexExec expects 1 child")
+        child_df = child_dfs[0]
+
+        # Add the array column if it's an expression
+        child_df = child_df.with_columns(self.physical_expr)
+
+        # Add a temporary row index to track original rows
+        child_df = child_df.with_row_index("__explode_row_id")
+
+        # Explode the array column
+        exploded_df = child_df.explode(self.col_name)
+
+        if self.keep_null_and_empty:
+            # For outer explode, we need to handle null/empty arrays specially
+            # Add the position column, but set it to null for null/empty arrays
+            exploded_df = exploded_df.with_columns(
+                pl.when(pl.col(self.col_name).is_not_null())
+                .then(pl.int_range(pl.len(), dtype=pl.Int64).over("__explode_row_id", mapping_strategy="group_to_rows"))
+                .otherwise(None)
+                .alias(self.index_name)
+            )
+        else:
+            # Filter out nulls in the exploded column for regular explode
+            exploded_df = exploded_df.filter(pl.col(self.col_name).is_not_null())
+            # Add the position column (0-based index within each original row)
+            exploded_df = exploded_df.with_columns(
+                pl.int_range(pl.len(), dtype=pl.Int64)
+                .over("__explode_row_id", mapping_strategy="group_to_rows")
+                .alias(self.index_name)
+            )
+
+        # Drop the temporary row index
+        exploded_df = exploded_df.drop("__explode_row_id")
+
+        # Rename the exploded column if needed
+        if self.value_name != self.col_name:
+            exploded_df = exploded_df.rename({self.col_name: self.value_name})
+
+        return exploded_df
+
+    def with_children(self, children: List[PhysicalPlan]) -> PhysicalPlan:
+        if len(children) != 1:
+            raise InternalError("Unreachable: ExplodeWithIndexExec expects 1 child")
+        return ExplodeWithIndexExec(
+            child=children[0],
+            physical_expr=self.physical_expr,
+            col_name=self.col_name,
+            index_name=self.index_name,
+            value_name=self.value_name,
+            keep_null_and_empty=self.keep_null_and_empty,
+            cache_info=self.cache_info,
+            session_state=self.session_state,
+        )
+
+    def build_node_lineage(
+        self,
+        leaf_nodes: List[OperatorLineage],
+    ) -> Tuple[OperatorLineage, pl.DataFrame]:
+        child_operator, child_df = self.children[0].build_node_lineage(leaf_nodes)
+        exploded_df = child_df.explode(self.col_name)
+        # Filter out nulls unless keep_null_and_empty is True
+        if not self.keep_null_and_empty:
+            exploded_df = exploded_df.filter(pl.col(self.col_name).is_not_null())
         exploded_df = exploded_df.with_columns(
             pl.col("_uuid").alias("_backwards_uuid"),
         )

src/fenic/_backends/local/transpiler/plan_converter.py

@@ -10,6 +10,7 @@
     DuckDBTableSinkExec,
     DuckDBTableSourceExec,
     ExplodeExec,
+    ExplodeWithIndexExec,
     FileSinkExec,
     FileSourceExec,
     FilterExec,
@@ -41,6 +42,7 @@
     DocSource,
     DropDuplicates,
     Explode,
+    ExplodeWithIndex,
     FileSink,
     FileSource,
     Filter,
@@ -335,6 +337,30 @@ def _convert_to_physical_plan(self, logical: LogicalPlan, cache_keys: set[str])
                 child_physical,
                 physical_expr,
                 target_field.name,
+                logical.keep_null_and_empty,
+                cache_info=logical.cache_info,
+                session_state=self.session_state,
+            )
+
+        elif isinstance(logical, ExplodeWithIndex):
+            child_logical = logical.children()[0]
+            physical_expr = self.expr_converter.convert(
+                logical._expr
+            )
+            child_physical = self._convert_to_physical_plan(
+                child_logical,
+                cache_keys,
+            )
+            target_field = logical._expr.to_column_field(child_logical, self.session_state)
+            # Determine the actual value name
+            actual_value_name = logical.value_name if logical.value_name is not None else target_field.name
+            return ExplodeWithIndexExec(
+                child_physical,
+                physical_expr,
+                target_field.name,
+                logical.index_name,
+                actual_value_name,
+                logical.keep_null_and_empty,
                 cache_info=logical.cache_info,
                 session_state=self.session_state,
             )

src/fenic/_gen/protos/logical_plan/v1/plans_pb2.py

@@ -31,7 +31,7 @@ class CacheInfo(_message.Message):
     def __init__(self, cache_key: _Optional[str] = ...) -> None: ...
 
 class LogicalPlan(_message.Message):
-    __slots__ = ("schema", "cache_info", "in_memory_source", "file_source", "table_source", "doc_source", "projection", "filter", "join", "aggregate", "union", "limit", "explode", "drop_duplicates", "sort", "unnest", "sql", "semantic_cluster", "semantic_join", "semantic_similarity_join", "file_sink", "table_sink")
+    __slots__ = ("schema", "cache_info", "in_memory_source", "file_source", "table_source", "doc_source", "projection", "filter", "join", "aggregate", "union", "limit", "explode", "explode_with_index", "drop_duplicates", "sort", "unnest", "sql", "semantic_cluster", "semantic_join", "semantic_similarity_join", "file_sink", "table_sink")
     SCHEMA_FIELD_NUMBER: _ClassVar[int]
     CACHE_INFO_FIELD_NUMBER: _ClassVar[int]
     IN_MEMORY_SOURCE_FIELD_NUMBER: _ClassVar[int]
@@ -45,6 +45,7 @@ class LogicalPlan(_message.Message):
     UNION_FIELD_NUMBER: _ClassVar[int]
     LIMIT_FIELD_NUMBER: _ClassVar[int]
     EXPLODE_FIELD_NUMBER: _ClassVar[int]
+    EXPLODE_WITH_INDEX_FIELD_NUMBER: _ClassVar[int]
     DROP_DUPLICATES_FIELD_NUMBER: _ClassVar[int]
     SORT_FIELD_NUMBER: _ClassVar[int]
     UNNEST_FIELD_NUMBER: _ClassVar[int]
@@ -67,6 +68,7 @@ class LogicalPlan(_message.Message):
     union: Union
     limit: Limit
     explode: Explode
+    explode_with_index: ExplodeWithIndex
     drop_duplicates: DropDuplicates
     sort: Sort
     unnest: Unnest
@@ -76,7 +78,7 @@ class LogicalPlan(_message.Message):
     semantic_similarity_join: SemanticSimilarityJoin
     file_sink: FileSink
     table_sink: TableSink
-    def __init__(self, schema: _Optional[_Union[FenicSchema, _Mapping]] = ..., cache_info: _Optional[_Union[CacheInfo, _Mapping]] = ..., in_memory_source: _Optional[_Union[InMemorySource, _Mapping]] = ..., file_source: _Optional[_Union[FileSource, _Mapping]] = ..., table_source: _Optional[_Union[TableSource, _Mapping]] = ..., doc_source: _Optional[_Union[DocSource, _Mapping]] = ..., projection: _Optional[_Union[Projection, _Mapping]] = ..., filter: _Optional[_Union[Filter, _Mapping]] = ..., join: _Optional[_Union[Join, _Mapping]] = ..., aggregate: _Optional[_Union[Aggregate, _Mapping]] = ..., union: _Optional[_Union[Union, _Mapping]] = ..., limit: _Optional[_Union[Limit, _Mapping]] = ..., explode: _Optional[_Union[Explode, _Mapping]] = ..., drop_duplicates: _Optional[_Union[DropDuplicates, _Mapping]] = ..., sort: _Optional[_Union[Sort, _Mapping]] = ..., unnest: _Optional[_Union[Unnest, _Mapping]] = ..., sql: _Optional[_Union[SQL, _Mapping]] = ..., semantic_cluster: _Optional[_Union[SemanticCluster, _Mapping]] = ..., semantic_join: _Optional[_Union[SemanticJoin, _Mapping]] = ..., semantic_similarity_join: _Optional[_Union[SemanticSimilarityJoin, _Mapping]] = ..., file_sink: _Optional[_Union[FileSink, _Mapping]] = ..., table_sink: _Optional[_Union[TableSink, _Mapping]] = ...) -> None: ...
+    def __init__(self, schema: _Optional[_Union[FenicSchema, _Mapping]] = ..., cache_info: _Optional[_Union[CacheInfo, _Mapping]] = ..., in_memory_source: _Optional[_Union[InMemorySource, _Mapping]] = ..., file_source: _Optional[_Union[FileSource, _Mapping]] = ..., table_source: _Optional[_Union[TableSource, _Mapping]] = ..., doc_source: _Optional[_Union[DocSource, _Mapping]] = ..., projection: _Optional[_Union[Projection, _Mapping]] = ..., filter: _Optional[_Union[Filter, _Mapping]] = ..., join: _Optional[_Union[Join, _Mapping]] = ..., aggregate: _Optional[_Union[Aggregate, _Mapping]] = ..., union: _Optional[_Union[Union, _Mapping]] = ..., limit: _Optional[_Union[Limit, _Mapping]] = ..., explode: _Optional[_Union[Explode, _Mapping]] = ..., explode_with_index: _Optional[_Union[ExplodeWithIndex, _Mapping]] = ..., drop_duplicates: _Optional[_Union[DropDuplicates, _Mapping]] = ..., sort: _Optional[_Union[Sort, _Mapping]] = ..., unnest: _Optional[_Union[Unnest, _Mapping]] = ..., sql: _Optional[_Union[SQL, _Mapping]] = ..., semantic_cluster: _Optional[_Union[SemanticCluster, _Mapping]] = ..., semantic_join: _Optional[_Union[SemanticJoin, _Mapping]] = ..., semantic_similarity_join: _Optional[_Union[SemanticSimilarityJoin, _Mapping]] = ..., file_sink: _Optional[_Union[FileSink, _Mapping]] = ..., table_sink: _Optional[_Union[TableSink, _Mapping]] = ...) -> None: ...
 
 class InMemorySource(_message.Message):
     __slots__ = ("source",)
@@ -217,12 +219,28 @@ class Limit(_message.Message):
     def __init__(self, input: _Optional[_Union[LogicalPlan, _Mapping]] = ..., n: _Optional[int] = ...) -> None: ...
 
 class Explode(_message.Message):
-    __slots__ = ("input", "expr")
+    __slots__ = ("input", "expr", "keep_null_and_empty")
     INPUT_FIELD_NUMBER: _ClassVar[int]
     EXPR_FIELD_NUMBER: _ClassVar[int]
+    KEEP_NULL_AND_EMPTY_FIELD_NUMBER: _ClassVar[int]
     input: LogicalPlan
     expr: _expressions_pb2.LogicalExpr
-    def __init__(self, input: _Optional[_Union[LogicalPlan, _Mapping]] = ..., expr: _Optional[_Union[_expressions_pb2.LogicalExpr, _Mapping]] = ...) -> None: ...
+    keep_null_and_empty: bool
+    def __init__(self, input: _Optional[_Union[LogicalPlan, _Mapping]] = ..., expr: _Optional[_Union[_expressions_pb2.LogicalExpr, _Mapping]] = ..., keep_null_and_empty: bool = ...) -> None: ...
+
+class ExplodeWithIndex(_message.Message):
+    __slots__ = ("input", "expr", "index_name", "value_name", "keep_null_and_empty")
+    INPUT_FIELD_NUMBER: _ClassVar[int]
+    EXPR_FIELD_NUMBER: _ClassVar[int]
+    INDEX_NAME_FIELD_NUMBER: _ClassVar[int]
+    VALUE_NAME_FIELD_NUMBER: _ClassVar[int]
+    KEEP_NULL_AND_EMPTY_FIELD_NUMBER: _ClassVar[int]
+    input: LogicalPlan
+    expr: _expressions_pb2.LogicalExpr
+    index_name: str
+    value_name: str
+    keep_null_and_empty: bool
+    def __init__(self, input: _Optional[_Union[LogicalPlan, _Mapping]] = ..., expr: _Optional[_Union[_expressions_pb2.LogicalExpr, _Mapping]] = ..., index_name: _Optional[str] = ..., value_name: _Optional[str] = ..., keep_null_and_empty: bool = ...) -> None: ...
 
 class DropDuplicates(_message.Message):
     __slots__ = ("input", "subset")