Upgrade DataFusion fork to 55.0.0 and preserve fork patches

datafusion/core/src/physical_planner.rs

@@ -57,7 +57,6 @@ use crate::physical_plan::{
 use crate::schema_equivalence::schema_satisfied_by;
 
 use arrow::array::{RecordBatch, builder::StringBuilder};
-use arrow::compute::SortOptions;
 use arrow::datatypes::Schema;
 use arrow_schema::Field;
 use datafusion_catalog::ScanArgs;
@@ -1708,15 +1707,27 @@ impl DefaultPhysicalPlanner {
                 // implement null-aware anti-join semantics and would return wrong
                 // results when the right side contains a null join key.
                 {
-                    // Use SortMergeJoin if hash join is not preferred
-                    let join_on_len = join_on.len();
+                    let sort_options = join_on
+                        .iter()
+                        .map(|(left_col, _)| {
+                            physical_left
+                                .output_ordering()
+                                .and_then(|ordering| {
+                                    ordering
+                                        .iter()
+                                        .find(|sort_expr| sort_expr.expr.eq(left_col))
+                                        .map(|sort_expr| sort_expr.options)
+                                })
+                                .unwrap_or_default()
+                        })
+                        .collect();
                     Arc::new(SortMergeJoinExec::try_new(
                         physical_left,
                         physical_right,
                         join_on,
                         join_filter,
                         *join_type,
-                        vec![SortOptions::default(); join_on_len],
+                        sort_options,
                         *null_equality,
                     )?)
                 } else if session_state.config().target_partitions() > 1
@@ -3316,8 +3327,8 @@ mod tests {
     use datafusion_expr::{
         Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, HigherOrderUDF,
         LogicalPlanBuilder, Partitioning as LogicalPartitioning, RangePartitioning,
-        ScalarUDF, Signature, TableSource, UserDefinedLogicalNodeCore, Volatility,
-        WindowFunctionDefinition, WindowUDF, col, lit, scalar_subquery,
+        ScalarUDF, Signature, SortExpr, TableSource, UserDefinedLogicalNodeCore,
+        Volatility, WindowFunctionDefinition, WindowUDF, col, lit, scalar_subquery,
     };
     use datafusion_functions_aggregate::count::{count_all, count_udaf};
     use datafusion_functions_aggregate::expr_fn::sum;
@@ -3820,6 +3831,134 @@ mod tests {
         ctx.sql(query).await?.collect().await
     }
 
+    fn smj_test_context(
+        left_ordering: Option<Vec<SortExpr>>,
+        right_ordering: Option<Vec<SortExpr>>,
+    ) -> Result<SessionContext> {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int32, false),
+            Field::new("b", DataType::Int32, false),
+        ]));
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![
+                Arc::new(Int32Array::from(vec![3, 2, 1])),
+                Arc::new(Int32Array::from(vec![10, 20, 30])),
+            ],
+        )?;
+
+        let table = |batch, ordering| -> Result<MemTable> {
+            let table = MemTable::try_new(Arc::clone(&schema), vec![vec![batch]])?;
+            Ok(match ordering {
+                Some(ordering) => table.with_sort_order(vec![ordering]),
+                None => table,
+            })
+        };
+        let left = table(batch.clone(), left_ordering)?;
+        let right = table(batch, right_ordering)?;
+
+        let config = SessionConfig::new()
+            .with_target_partitions(4)
+            .set_bool("datafusion.optimizer.prefer_hash_join", false)
+            .set_bool("datafusion.optimizer.skip_failed_rules", false);
+        let state = SessionStateBuilder::new()
+            .with_config(config)
+            .with_runtime_env(Arc::new(RuntimeEnv::default()))
+            .with_default_features()
+            .build();
+        let ctx = SessionContext::new_with_state(state);
+        ctx.register_table("left_table", Arc::new(left))?;
+        ctx.register_table("right_table", Arc::new(right))?;
+        Ok(ctx)
+    }
+
+    async fn smj_plan(ctx: &SessionContext, sql: &str) -> Result<String> {
+        let state = ctx.state();
+        let logical_plan = state.create_logical_plan(sql).await?;
+        let logical_plan = state.optimize(&logical_plan)?;
+        let plan = DefaultPhysicalPlanner::default()
+            .create_physical_plan(&logical_plan, &state)
+            .await?;
+        Ok(format!("{}", displayable(plan.as_ref()).indent(false)))
+    }
+
+    #[tokio::test]
+    async fn smj_uses_left_descending_sort_options() -> Result<()> {
+        let descending = vec![col("a").sort(false, false)];
+        let ctx = smj_test_context(Some(descending), None)?;
+        let plan = smj_plan(
+            &ctx,
+            "SELECT * FROM left_table JOIN right_table ON left_table.a = right_table.a",
+        )
+        .await?;
+
+        insta::assert_snapshot!(plan, @r"
+        SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)]
+          DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST
+          SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false]
+            DataSourceExec: partitions=1, partition_sizes=[1]
+        ");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn smj_defaults_to_ascending_when_left_is_unsorted() -> Result<()> {
+        let ctx = smj_test_context(None, None)?;
+        let plan = smj_plan(
+            &ctx,
+            "SELECT * FROM left_table JOIN right_table ON left_table.a = right_table.a",
+        )
+        .await?;
+
+        insta::assert_snapshot!(plan, @r"
+        SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)]
+          SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
+            DataSourceExec: partitions=1, partition_sizes=[1]
+          SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
+            DataSourceExec: partitions=1, partition_sizes=[1]
+        ");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn smj_defaults_only_unmatched_join_columns_to_ascending() -> Result<()> {
+        let descending = vec![col("a").sort(false, false)];
+        let ctx = smj_test_context(Some(descending), None)?;
+        let plan = smj_plan(
+            &ctx,
+            "SELECT * FROM left_table JOIN right_table \
+             ON left_table.a = right_table.a AND left_table.b = right_table.b",
+        )
+        .await?;
+
+        insta::assert_snapshot!(plan, @r"
+        SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0), (b@1, b@1)]
+          SortExec: expr=[a@0 DESC NULLS LAST, b@1 ASC], preserve_partitioning=[false]
+            DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST
+          SortExec: expr=[a@0 DESC NULLS LAST, b@1 ASC], preserve_partitioning=[false]
+            DataSourceExec: partitions=1, partition_sizes=[1]
+        ");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn smj_reuses_matching_descending_order_on_both_sides() -> Result<()> {
+        let descending = vec![col("a").sort(false, false)];
+        let ctx = smj_test_context(Some(descending.clone()), Some(descending))?;
+        let plan = smj_plan(
+            &ctx,
+            "SELECT * FROM left_table JOIN right_table ON left_table.a = right_table.a",
+        )
+        .await?;
+
+        insta::assert_snapshot!(plan, @r"
+        SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)]
+          DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST
+          DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST
+        ");
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_all_operators() -> Result<()> {
         let logical_plan = test_csv_scan()

datafusion/datasource-parquet/src/opener/mod.rs

@@ -1749,7 +1749,11 @@ mod test {
         CachedParquetFileReaderFactory, DefaultParquetFileReaderFactory,
         ParquetFileReaderFactory, ParquetRowSelection, RowGroupAccess,
     };
-    use arrow::array::{RecordBatch, record_batch};
+    use arrow::array::{
+        Array, ArrayRef, Date32Array, Int32Array, ListArray, RecordBatch, StringArray,
+        record_batch,
+    };
+    use arrow::buffer::OffsetBuffer;
     use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
     use bytes::{BufMut, BytesMut};
     use datafusion_common::{
@@ -1774,8 +1778,8 @@ mod test {
     };
     use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
     use datafusion_pruning::MAX_IN_LIST_SIZE;
-    use futures::StreamExt;
     use futures::stream::BoxStream;
+    use futures::{StreamExt, TryStreamExt};
     use object_store::{ObjectStore, ObjectStoreExt, memory::InMemory, path::Path};
     use parquet::arrow::{ArrowSchemaConverter, ArrowWriter};
     use parquet::file::metadata::{ColumnChunkMetaData, FileMetaData, ParquetMetaData};
@@ -2340,13 +2344,12 @@ mod test {
     async fn collect_int32_values(
         mut stream: BoxStream<'static, Result<RecordBatch>>,
     ) -> Vec<i32> {
-        use arrow::array::Array;
         let mut values = vec![];
         while let Some(Ok(batch)) = stream.next().await {
             let array = batch
                 .column(0)
                 .as_any()
-                .downcast_ref::<arrow::array::Int32Array>()
+                .downcast_ref::<Int32Array>()
                 .unwrap();
             for i in 0..array.len() {
                 if !array.is_null(i) {
@@ -2365,6 +2368,133 @@ mod test {
         write_parquet_batches(store, filename, vec![batch], None).await
     }
 
+    async fn read_batches_with_schema(
+        store: Arc<dyn ObjectStore>,
+        filename: &str,
+        data_size: usize,
+        schema: SchemaRef,
+    ) -> Result<Vec<RecordBatch>> {
+        let projection_indices = (0..schema.fields().len()).collect::<Vec<_>>();
+        let file =
+            PartitionedFile::new(filename.to_string(), u64::try_from(data_size).unwrap());
+        let morselizer = ParquetMorselizerBuilder::new()
+            .with_store(store)
+            .with_schema(schema)
+            .with_projection_indices(&projection_indices)
+            .build();
+        open_file(&morselizer, file).await?.try_collect().await
+    }
+
+    #[tokio::test]
+    async fn schema_adaptation_casts_utf8_to_date32() -> Result<()> {
+        let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+        let file_schema =
+            Arc::new(Schema::new(vec![Field::new("date", DataType::Utf8, true)]));
+        let batch = RecordBatch::try_new(
+            file_schema,
+            vec![Arc::new(StringArray::from(vec![
+                "2026-01-01",
+                "2026-02-01",
+            ]))],
+        )?;
+        let data_size =
+            write_parquet(Arc::clone(&store), "schema_adapt_dates.parquet", batch).await;
+
+        let logical_schema = Arc::new(Schema::new(vec![Field::new(
+            "date",
+            DataType::Date32,
+            true,
+        )]));
+        let batches = read_batches_with_schema(
+            store,
+            "schema_adapt_dates.parquet",
+            data_size,
+            logical_schema,
+        )
+        .await?;
+
+        assert_eq!(batches.len(), 1);
+        assert_eq!(batches[0].num_rows(), 2);
+        let dates = batches[0]
+            .column(0)
+            .as_any()
+            .downcast_ref::<Date32Array>()
+            .unwrap();
+        assert_eq!(dates.values(), &[20_454, 20_485]);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn schema_adaptation_normalizes_list_fields() -> Result<()> {
+        let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+        let inner_field = Arc::new(Field::new("conditions", DataType::Int32, false));
+        let file_schema = Arc::new(Schema::new(vec![Field::new(
+            "conditions",
+            DataType::List(Arc::clone(&inner_field)),
+            true,
+        )]));
+        let values = Int32Array::from(vec![1, 2, 3, 4]);
+        let offsets = OffsetBuffer::from_lengths([2, 2]);
+        let list = ListArray::new(inner_field, offsets, Arc::new(values), None);
+        let batch = RecordBatch::try_new(file_schema, vec![Arc::new(list) as ArrayRef])?;
+        let data_size =
+            write_parquet(Arc::clone(&store), "schema_adapt_list.parquet", batch).await;
+
+        let logical_inner = Arc::new(Field::new("element", DataType::Int32, true));
+        let logical_schema = Arc::new(Schema::new(vec![Field::new(
+            "conditions",
+            DataType::List(logical_inner),
+            true,
+        )]));
+        let batches = read_batches_with_schema(
+            store,
+            "schema_adapt_list.parquet",
+            data_size,
+            Arc::clone(&logical_schema),
+        )
+        .await?;
+
+        assert_eq!(batches.len(), 1);
+        assert_eq!(batches[0].num_rows(), 2);
+        assert_eq!(batches[0].schema(), logical_schema);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn schema_adaptation_preserves_logical_nullability() -> Result<()> {
+        let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+        let file_schema =
+            Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
+        let batch = RecordBatch::try_new(
+            file_schema,
+            vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
+        )?;
+        let data_size =
+            write_parquet(Arc::clone(&store), "schema_adapt_nullable.parquet", batch)
+                .await;
+
+        let logical_schema =
+            Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, true)]));
+        let batches = read_batches_with_schema(
+            store,
+            "schema_adapt_nullable.parquet",
+            data_size,
+            logical_schema,
+        )
+        .await?;
+
+        assert_eq!(batches.len(), 1);
+        assert_eq!(batches[0].num_rows(), 3);
+        assert!(batches[0].schema().field(0).is_nullable());
+        let ids = batches[0]
+            .column(0)
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .unwrap();
+        assert_eq!(ids.values(), &[1, 2, 3]);
+        Ok(())
+    }
+
     /// Write multiple batches to a parquet file with optional writer properties
     async fn write_parquet_batches(
         store: Arc<dyn ObjectStore>,
@@ -3742,7 +3872,7 @@ mod test {
             let part = batch
                 .column(1)
                 .as_any()
-                .downcast_ref::<arrow::array::Int32Array>()
+                .downcast_ref::<Int32Array>()
                 .unwrap();
             assert!(part.iter().all(|v| v == Some(5)));