chore: bump datafusion to 54.1.0

Cargo.toml

@@ -150,7 +150,7 @@ prost-types = "0.14"
 # The `axum` version must match the one used in `tonic` (replace `RELEASE` with the release we are using):
 #   https://github.com/hyperium/tonic/blob/vRELEASE/tonic/Cargo.toml
 axum = "0.8.9"
-datafusion = { version = "54.0.0", default-features = false, features = [
+datafusion = { version = "54.1.0", default-features = false, features = [
     "nested_expressions",
     "crypto_expressions",
     "datetime_expressions",
@@ -164,18 +164,18 @@ datafusion = { version = "54.0.0", default-features = false, features = [
     "avro",
     "recursive_protection",
 ] }
-datafusion-common = { version = "54.0.0", features = ["object_store"] }
-datafusion-datasource = { version = "54.0.0" }
-datafusion-datasource-avro = { version = "54.0.0" }
-datafusion-datasource-json = { version = "54.0.0" }
-datafusion-expr = { version = "54.0.0", default-features = false }
-datafusion-expr-common = { version = "54.0.0" }
-datafusion-proto = { version = "54.0.0" }
-datafusion-functions = { version = "54.0.0" }
-datafusion-functions-nested = { version = "54.0.0", default-features = false }
-datafusion-physical-expr = { version = "54.0.0" }
-datafusion-session = { version = "54.0.0" }
-datafusion-spark = { version = "54.0.0", features = ["core"] }
+datafusion-common = { version = "54.1.0", features = ["object_store"] }
+datafusion-datasource = { version = "54.1.0" }
+datafusion-datasource-avro = { version = "54.1.0" }
+datafusion-datasource-json = { version = "54.1.0" }
+datafusion-expr = { version = "54.1.0", default-features = false }
+datafusion-expr-common = { version = "54.1.0" }
+datafusion-proto = { version = "54.1.0" }
+datafusion-functions = { version = "54.1.0" }
+datafusion-functions-nested = { version = "54.1.0", default-features = false }
+datafusion-physical-expr = { version = "54.1.0" }
+datafusion-session = { version = "54.1.0" }
+datafusion-spark = { version = "54.1.0", features = ["core"] }
 pyo3 = { version = "0.29.0", features = ["serde"] }
 jiter = { version = "0.15.0", default-features = false }
 arrow = { version = "58.3.0", features = ["chrono-tz"] }

crates/sail-execution/proto/sail/plan/physical.proto

@@ -657,6 +657,7 @@ message RecursiveQueryExecNode {
   bytes static_term = 2;
   bytes recursive_term = 3;
   bool is_distinct = 4;
+  bytes output_schema = 5;
 }
 
 message SortMergeJoinExecNode {

crates/sail-execution/src/proto/codec.rs

@@ -616,11 +616,14 @@ impl PhysicalExtensionCodec for RemoteExecutionCodec {
                 static_term,
                 recursive_term,
                 is_distinct,
+                output_schema,
             }) => {
                 let static_term = try_decode_physical_plan(ctx, self, &static_term)?;
                 let recursive_term = try_decode_physical_plan(ctx, self, &recursive_term)?;
+                let output_schema = Arc::new(try_decode_schema(&output_schema)?);
                 Ok(Arc::new(RecursiveQueryExec::try_new(
                     name,
+                    output_schema,
                     static_term,
                     recursive_term,
                     is_distinct,
@@ -1530,11 +1533,13 @@ impl PhysicalExtensionCodec for RemoteExecutionCodec {
             let recursive_term =
                 try_encode_physical_plan(self, recursive_query.recursive_term().clone())?;
             let is_distinct = recursive_query.is_distinct();
+            let output_schema = try_encode_schema(recursive_query.schema().as_ref())?;
             NodeKind::RecursiveQuery(r#gen::RecursiveQueryExecNode {
                 name,
                 static_term,
                 recursive_term,
                 is_distinct,
+                output_schema,
             })
         } else if let Some(sort_merge_join) = node.downcast_ref::<SortMergeJoinExec>() {
             let left = try_encode_physical_plan(self, sort_merge_join.left().clone())?;
@@ -4421,6 +4426,46 @@ mod tests {
         codec.try_decode_udwf(&name, &buf)
     }
 
+    #[test]
+    fn test_round_trip_recursive_query_preserves_output_schema() -> Result<()> {
+        use datafusion::arrow::datatypes::{DataType, Field};
+        use datafusion::physical_plan::empty::EmptyExec;
+
+        let static_schema = Arc::new(Schema::new(vec![Field::new(
+            "value",
+            DataType::Int32,
+            false,
+        )]));
+        let output_schema = Arc::new(Schema::new(vec![Field::new(
+            "value",
+            DataType::Int32,
+            true,
+        )]));
+        let static_term: Arc<dyn ExecutionPlan> =
+            Arc::new(EmptyExec::new(Arc::clone(&static_schema)));
+        let recursive_term: Arc<dyn ExecutionPlan> =
+            Arc::new(EmptyExec::new(Arc::clone(&output_schema)));
+        let plan = Arc::new(RecursiveQueryExec::try_new(
+            "numbers".to_string(),
+            Arc::clone(&output_schema),
+            static_term,
+            recursive_term,
+            false,
+        )?);
+
+        let codec = RemoteExecutionCodec;
+        let bytes = try_encode_physical_plan(&codec, plan)?;
+        let decoded = try_decode_physical_plan(&TaskContext::default(), &codec, &bytes)?;
+        let recursive_query = decoded
+            .downcast_ref::<RecursiveQueryExec>()
+            .ok_or_else(|| plan_datafusion_err!("decoded plan is not a RecursiveQueryExec"))?;
+
+        assert_eq!(recursive_query.schema(), output_schema);
+        assert_eq!(recursive_query.static_term().schema(), output_schema);
+        assert_eq!(recursive_query.recursive_term().schema(), output_schema);
+        Ok(())
+    }
+
     #[test]
     fn test_struct_field_matching_proto_round_trip_and_default() -> Result<()> {
         for mode in [

crates/sail-function/src/scalar/array/spark_array_compact.rs

@@ -84,6 +84,7 @@ fn array_compact_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
 
 fn array_compact_generic<O: OffsetSizeTrait>(list: &GenericListArray<O>) -> Result<ArrayRef> {
     let values = list.values();
+    let values_nulls = values.logical_nulls();
     let offsets = list.offsets();
 
     // Build a boolean filter mask for the flattened values array.
@@ -106,7 +107,7 @@ fn array_compact_generic<O: OffsetSizeTrait>(list: &GenericListArray<O>) -> Resu
         } else {
             // Non-null row: keep only non-null values.
             for j in start..end {
-                let keep_value = values.is_valid(j);
+                let keep_value = values_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(j));
                 keep.push(keep_value);
                 if keep_value {
                     total_kept += 1;

python/pysail/tests/spark/__snapshots__/test_tpcds.plan.yaml

@@ -342,36 +342,35 @@ data:
                   AggregateExec: mode=FinalPartitioned, gby=[#8@0 as #8], aggr=[count(Int64(1))]
                     RepartitionExec: partitioning=Hash([#8@0], 4), input_partitions=4
                       AggregateExec: mode=Partial, gby=[#8@0 as #8], aggr=[count(Int64(1))]
-                        FilterExec: CAST(#87@1 AS Decimal128(11, 7)) > Some(12),2,1 * #155@2, projection=[#8@0]
-                          HashJoinExec: mode=CollectLeft, join_type=Left, on=[(#94@2, #145@1)], projection=[#8@0, #87@1, #155@3]
-                            CoalescePartitionsExec
-                              HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(#82@0, #33@1)], projection=[#8@3, #87@1, #94@2]
-                                ProjectionExec: expr=[_0@0 as #82, _5@1 as #87, _12@2 as #94]
-                                  DataSourceExec: partitions=1, partition_sizes=[1]
-                                HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(#31@1, CAST(d.#54 AS Float64)@1)], projection=[#8@0, #33@2]
-                                  CoalescePartitionsExec
-                                    HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(CAST(c.#13 AS Float64)@2, #34@2)], projection=[#8@0, #31@3, #33@4]
-                                      CoalescePartitionsExec
-                                        ProjectionExec: expr=[#8@0 as #8, #13@1 as #13, CAST(#13@1 AS Float64) as CAST(c.#13 AS Float64)]
-                                          RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
-                                            HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(#17@1, #0@0)], projection=[#8@3, #13@0]
-                                              ProjectionExec: expr=[_0@0 as #13, _4@1 as #17]
-                                                DataSourceExec: partitions=1, partition_sizes=[1]
-                                              ProjectionExec: expr=[_0@0 as #0, _8@1 as #8]
-                                                DataSourceExec: partitions=1, partition_sizes=[1]
-                                      RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
-                                        ProjectionExec: expr=[_0@0 as #31, _2@1 as #33, _3@2 as #34]
-                                          DataSourceExec: partitions=1, partition_sizes=[1]
-                                  ProjectionExec: expr=[_0@0 as #54, CAST(_0@0 AS Float64) as CAST(d.#54 AS Float64)]
-                                    FilterExec: _3@1 = scalar_subquery(<pending>), projection=[_0@0]
-                                      RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+                        HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(#94@2, #145@1)], filter=CAST(#87@0 AS Decimal128(11, 7)) > Some(12),2,1 * #155@1, projection=[#8@0]
+                          CoalescePartitionsExec
+                            HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(#82@0, #33@1)], projection=[#8@3, #87@1, #94@2]
+                              ProjectionExec: expr=[_0@0 as #82, _5@1 as #87, _12@2 as #94]
+                                DataSourceExec: partitions=1, partition_sizes=[1]
+                              HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(#31@1, CAST(d.#54 AS Float64)@1)], projection=[#8@0, #33@2]
+                                CoalescePartitionsExec
+                                  HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(CAST(c.#13 AS Float64)@2, #34@2)], projection=[#8@0, #31@3, #33@4]
+                                    CoalescePartitionsExec
+                                      ProjectionExec: expr=[#8@0 as #8, #13@1 as #13, CAST(#13@1 AS Float64) as CAST(c.#13 AS Float64)]
+                                        RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+                                          HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(#17@1, #0@0)], projection=[#8@3, #13@0]
+                                            ProjectionExec: expr=[_0@0 as #13, _4@1 as #17]
+                                              DataSourceExec: partitions=1, partition_sizes=[1]
+                                            ProjectionExec: expr=[_0@0 as #0, _8@1 as #8]
+                                              DataSourceExec: partitions=1, partition_sizes=[1]
+                                    RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+                                      ProjectionExec: expr=[_0@0 as #31, _2@1 as #33, _3@2 as #34]
                                         DataSourceExec: partitions=1, partition_sizes=[1]
-                            ProjectionExec: expr=[avg(j.#138)@1 as #155, #145@0 as #145]
-                              AggregateExec: mode=FinalPartitioned, gby=[#145@0 as #145], aggr=[avg(j.#138)]
-                                RepartitionExec: partitioning=Hash([#145@0], 4), input_partitions=1
-                                  AggregateExec: mode=Partial, gby=[#145@1 as #145], aggr=[avg(j.#138)]
-                                    ProjectionExec: expr=[_5@0 as #138, _12@1 as #145]
+                                ProjectionExec: expr=[_0@0 as #54, CAST(_0@0 AS Float64) as CAST(d.#54 AS Float64)]
+                                  FilterExec: _3@1 = scalar_subquery(<pending>), projection=[_0@0]
+                                    RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
                                       DataSourceExec: partitions=1, partition_sizes=[1]
+                          ProjectionExec: expr=[avg(j.#138)@1 as #155, #145@0 as #145]
+                            AggregateExec: mode=FinalPartitioned, gby=[#145@0 as #145], aggr=[avg(j.#138)]
+                              RepartitionExec: partitioning=Hash([#145@0], 4), input_partitions=1
+                                AggregateExec: mode=Partial, gby=[#145@1 as #145], aggr=[avg(j.#138)]
+                                  ProjectionExec: expr=[_5@0 as #138, _12@1 as #145]
+                                    DataSourceExec: partitions=1, partition_sizes=[1]
         AggregateExec: mode=FinalPartitioned, gby=[#132@0 as #132], aggr=[]
           RepartitionExec: partitioning=Hash([#132@0], 4), input_partitions=4
             AggregateExec: mode=Partial, gby=[#132@0 as #132], aggr=[]

python/pysail/tests/spark/function/features/array_compact.feature

@@ -38,6 +38,15 @@ Feature: array_compact() removes null values from an array
         | result |
         | []     |
 
+    Scenario: array_compact with untyped null values returns empty array
+      When query
+        """
+        SELECT array_compact(array(NULL, NULL)) AS result
+        """
+      Then query result
+        | result |
+        | []     |
+
     Scenario: array_compact with null at beginning
       When query
         """

python/pysail/tests/spark/function/features/width_bucket.feature

@@ -0,0 +1,15 @@
+Feature: width_bucket() assigns values to buckets
+
+  Scenario: width_bucket returns a Spark BIGINT result
+    When query
+      """
+      SELECT width_bucket(5.0, 0.0, 10.0, 5) AS result
+      """
+    Then query result
+      | result |
+      | 3      |
+    And query schema
+      """
+      root
+       |-- result: long (nullable = true)
+      """