#17411 Relax constraint that file sort order must only reference individual columns

Author: pepijnve

datafusion/catalog/src/stream.rs

@@ -28,18 +28,21 @@ use std::sync::Arc;
 use crate::{Session, TableProvider, TableProviderFactory};
 use arrow::array::{RecordBatch, RecordBatchReader, RecordBatchWriter};
 use arrow::datatypes::SchemaRef;
-use datafusion_common::{config_err, plan_err, Constraints, DataFusionError, Result};
+use datafusion_common::{
+    config_err, plan_err, Constraints, DFSchema, DataFusionError, Result,
+};
 use datafusion_common_runtime::SpawnedTask;
 use datafusion_datasource::sink::{DataSink, DataSinkExec};
 use datafusion_execution::{SendableRecordBatchStream, TaskContext};
 use datafusion_expr::dml::InsertOp;
 use datafusion_expr::{CreateExternalTable, Expr, SortExpr, TableType};
-use datafusion_physical_expr::create_ordering;
+use datafusion_physical_expr::create_lex_orderings;
 use datafusion_physical_plan::stream::RecordBatchReceiverStreamBuilder;
 use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec};
 use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan};
 
 use async_trait::async_trait;
+use datafusion_physical_expr::equivalence::project_orderings;
 use futures::StreamExt;
 
 /// A [`TableProviderFactory`] for [`StreamTable`]
@@ -321,24 +324,26 @@ impl TableProvider for StreamTable {
 
     async fn scan(
         &self,
-        _state: &dyn Session,
+        state: &dyn Session,
         projection: Option<&Vec<usize>>,
         _filters: &[Expr],
         limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        let projected_schema = match projection {
-            Some(p) => {
-                let projected = self.0.source.schema().project(p)?;
-                create_ordering(&projected, &self.0.order)?
-            }
-            None => create_ordering(self.0.source.schema(), &self.0.order)?,
+        let schema = self.0.source.schema();
+        let df_schema = DFSchema::try_from(Arc::clone(schema))?;
+        let mut sort_information =
+            create_lex_orderings(&self.0.order, &df_schema, state.execution_props())?;
+
+        // If there is a projection on the source, we also need to project orderings
+        if let Some(projection) = projection {
+            sort_information = project_orderings(sort_information, schema, projection)?;
         };
 
         Ok(Arc::new(StreamingTableExec::try_new(
             Arc::clone(self.0.source.schema()),
             vec![Arc::new(StreamRead(Arc::clone(&self.0))) as _],
             projection,
-            projected_schema,
+            sort_information,
             true,
             limit,
         )?))
@@ -351,7 +356,9 @@ impl TableProvider for StreamTable {
         _insert_op: InsertOp,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let schema = self.0.source.schema();
-        let orders = create_ordering(schema, &self.0.order)?;
+        let df_schema = DFSchema::try_from(Arc::clone(schema))?;
+        let orders =
+            create_lex_orderings(&self.0.order, &df_schema, _state.execution_props())?;
         // It is sufficient to pass only one of the equivalent orderings:
         let ordering = orders.into_iter().next().map(Into::into);
 

datafusion/core/src/datasource/listing/table.rs

@@ -23,7 +23,7 @@ use super::{
 };
 use crate::{
     datasource::file_format::{file_compression_type::FileCompressionType, FileFormat},
-    datasource::{create_ordering, physical_plan::FileSinkConfig},
+    datasource::physical_plan::FileSinkConfig,
     execution::context::SessionState,
 };
 use arrow::datatypes::{DataType, Field, SchemaBuilder, SchemaRef};
@@ -32,7 +32,7 @@ use async_trait::async_trait;
 use datafusion_catalog::{Session, TableProvider};
 use datafusion_common::{
     config_datafusion_err, config_err, internal_err, plan_err, project_schema,
-    stats::Precision, Constraints, DataFusionError, Result, SchemaExt,
+    stats::Precision, Constraints, DFSchema, DataFusionError, Result, SchemaExt,
 };
 use datafusion_datasource::{
     compute_all_files_statistics,
@@ -45,16 +45,19 @@ use datafusion_execution::{
     cache::{cache_manager::FileStatisticsCache, cache_unit::DefaultFileStatisticsCache},
     config::SessionConfig,
 };
+use datafusion_expr::execution_props::ExecutionProps;
 use datafusion_expr::{
     dml::InsertOp, Expr, SortExpr, TableProviderFilterPushDown, TableType,
 };
+use datafusion_physical_expr::create_lex_orderings;
 use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory;
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use datafusion_physical_plan::{empty::EmptyExec, ExecutionPlan, Statistics};
 use futures::{future, stream, Stream, StreamExt, TryStreamExt};
 use itertools::Itertools;
 use object_store::ObjectStore;
 use std::{any::Any, collections::HashMap, str::FromStr, sync::Arc};
+
 /// Indicates the source of the schema for a [`ListingTable`]
 // PartialEq required for assert_eq! in tests
 #[derive(Debug, Clone, Copy, PartialEq, Default)]
@@ -1126,8 +1129,12 @@ impl ListingTable {
     }
 
     /// If file_sort_order is specified, creates the appropriate physical expressions
-    fn try_create_output_ordering(&self) -> Result<Vec<LexOrdering>> {
-        create_ordering(&self.table_schema, &self.options.file_sort_order)
+    fn try_create_output_ordering(
+        &self,
+        execution_props: &ExecutionProps,
+    ) -> Result<Vec<LexOrdering>> {
+        let df_schema = DFSchema::try_from(Arc::clone(&self.table_schema))?;
+        create_lex_orderings(&self.options.file_sort_order, &df_schema, execution_props)
     }
 }
 
@@ -1199,7 +1206,7 @@ impl TableProvider for ListingTable {
             return Ok(Arc::new(EmptyExec::new(projected_schema)));
         }
 
-        let output_ordering = self.try_create_output_ordering()?;
+        let output_ordering = self.try_create_output_ordering(state.execution_props())?;
         match state
             .config_options()
             .execution
@@ -1334,7 +1341,7 @@ impl TableProvider for ListingTable {
             file_extension: self.options().format.get_ext(),
         };
 
-        let orderings = self.try_create_output_ordering()?;
+        let orderings = self.try_create_output_ordering(state.execution_props())?;
         // It is sufficient to pass only one of the equivalent orderings:
         let order_requirements = orderings.into_iter().next().map(Into::into);
 
@@ -1562,6 +1569,7 @@ mod tests {
         SchemaAdapter, SchemaAdapterFactory, SchemaMapper,
     };
     use datafusion_expr::{BinaryExpr, LogicalPlanBuilder, Operator};
+    use datafusion_physical_expr::expressions::binary;
     use datafusion_physical_expr::PhysicalSortExpr;
     use datafusion_physical_plan::{collect, ExecutionPlanProperties};
     use rstest::rstest;
@@ -1694,29 +1702,44 @@ mod tests {
 
         use crate::datasource::file_format::parquet::ParquetFormat;
         use datafusion_physical_plan::expressions::col as physical_col;
+        use datafusion_physical_plan::expressions::lit as physical_lit;
         use std::ops::Add;
 
         // (file_sort_order, expected_result)
         let cases = vec![
-            (vec![], Ok(Vec::<LexOrdering>::new())),
+            (
+                vec![],
+                Ok::<Vec<LexOrdering>, DataFusionError>(Vec::<LexOrdering>::new()),
+            ),
             // sort expr, but non column
             (
-                vec![vec![
-                    col("int_col").add(lit(1)).sort(true, true),
-                ]],
-                Err("Expected single column reference in sort_order[0][0], got int_col + Int32(1)"),
+                vec![vec![col("int_col").add(lit(1)).sort(true, true)]],
+                Ok(vec![[PhysicalSortExpr {
+                    expr: binary(
+                        physical_col("int_col", &schema).unwrap(),
+                        Operator::Plus,
+                        physical_lit(1),
+                        &schema,
+                    )
+                    .unwrap(),
+                    options: SortOptions {
+                        descending: false,
+                        nulls_first: true,
+                    },
+                }]
+                .into()]),
             ),
             // ok with one column
             (
                 vec![vec![col("string_col").sort(true, false)]],
                 Ok(vec![[PhysicalSortExpr {
-                            expr: physical_col("string_col", &schema).unwrap(),
-                            options: SortOptions {
-                                descending: false,
-                                nulls_first: false,
-                            },
-                        }].into(),
-                ])
+                    expr: physical_col("string_col", &schema).unwrap(),
+                    options: SortOptions {
+                        descending: false,
+                        nulls_first: false,
+                    },
+                }]
+                .into()]),
             ),
             // ok with two columns, different options
             (
@@ -1725,14 +1748,18 @@ mod tests {
                     col("int_col").sort(false, true),
                 ]],
                 Ok(vec![[
-                            PhysicalSortExpr::new_default(physical_col("string_col", &schema).unwrap())
-                                        .asc()
-                                        .nulls_last(),
-                            PhysicalSortExpr::new_default(physical_col("int_col", &schema).unwrap())
-                                        .desc()
-                                        .nulls_first()
-                        ].into(),
-                ])
+                    PhysicalSortExpr::new_default(
+                        physical_col("string_col", &schema).unwrap(),
+                    )
+                    .asc()
+                    .nulls_last(),
+                    PhysicalSortExpr::new_default(
+                        physical_col("int_col", &schema).unwrap(),
+                    )
+                    .desc()
+                    .nulls_first(),
+                ]
+                .into()]),
             ),
         ];
 
@@ -1745,7 +1772,8 @@ mod tests {
 
             let table =
                 ListingTable::try_new(config.clone()).expect("Creating the table");
-            let ordering_result = table.try_create_output_ordering();
+            let ordering_result =
+                table.try_create_output_ordering(state.execution_props());
 
             match (expected_result, ordering_result) {
                 (Ok(expected), Ok(result)) => {

datafusion/datasource/src/file_scan_config.rs

@@ -18,11 +18,6 @@
 //! [`FileScanConfig`] to configure scanning of possibly partitioned
 //! file sources.
 
-use std::{
-    any::Any, borrow::Cow, collections::HashMap, fmt::Debug, fmt::Formatter,
-    fmt::Result as FmtResult, marker::PhantomData, sync::Arc,
-};
-
 use crate::file_groups::FileGroup;
 #[allow(unused_imports)]
 use crate::schema_adapter::SchemaAdapterFactory;
@@ -52,7 +47,7 @@ use datafusion_physical_expr::{expressions::Column, utils::reassign_predicate_co
 use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
 use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory;
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
-use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
+use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use datafusion_physical_plan::projection::ProjectionExpr;
 use datafusion_physical_plan::{
     display::{display_orderings, ProjectSchemaDisplay},
@@ -63,7 +58,13 @@ use datafusion_physical_plan::{
 use datafusion_physical_plan::{
     filter::collect_columns_from_predicate, filter_pushdown::FilterPushdownPropagation,
 };
+use std::ops::Range;
+use std::{
+    any::Any, borrow::Cow, collections::HashMap, fmt::Debug, fmt::Formatter,
+    fmt::Result as FmtResult, marker::PhantomData, sync::Arc,
+};
 
+use datafusion_physical_expr::equivalence::project_orderings;
 use datafusion_physical_plan::coop::cooperative;
 use datafusion_physical_plan::execution_plan::SchedulingType;
 use log::{debug, warn};
@@ -703,12 +704,14 @@ impl FileScanConfig {
     fn projection_indices(&self) -> Vec<usize> {
         match &self.projection {
             Some(proj) => proj.clone(),
-            None => (0..self.file_schema.fields().len()
-                + self.table_partition_cols.len())
-                .collect(),
+            None => self.column_indices().collect(),
         }
     }
 
+    fn column_indices(&self) -> Range<usize> {
+        0..self.file_schema.fields().len() + self.table_partition_cols.len()
+    }
+
     pub fn projected_stats(&self) -> Statistics {
         let statistics = self.file_source.statistics().unwrap();
 
@@ -734,9 +737,17 @@ impl FileScanConfig {
     }
 
     pub fn projected_schema(&self) -> Arc<Schema> {
+        let full_schema = self.full_schema();
+        if let Some(projection) = &self.projection {
+            Arc::new(full_schema.project(projection).unwrap())
+        } else {
+            full_schema
+        }
+    }
+
+    fn full_schema(&self) -> Arc<Schema> {
         let table_fields: Vec<_> = self
-            .projection_indices()
-            .into_iter()
+            .column_indices()
             .map(|idx| {
                 if idx < self.file_schema.fields().len() {
                     self.file_schema.field(idx).clone()
@@ -1365,30 +1376,16 @@ fn get_projected_output_ordering(
     base_config: &FileScanConfig,
     projected_schema: &SchemaRef,
 ) -> Vec<LexOrdering> {
-    let mut all_orderings = vec![];
-    for output_ordering in &base_config.output_ordering {
-        let mut new_ordering = vec![];
-        for PhysicalSortExpr { expr, options } in output_ordering.iter() {
-            if let Some(col) = expr.as_any().downcast_ref::<Column>() {
-                let name = col.name();
-                if let Some((idx, _)) = projected_schema.column_with_name(name) {
-                    // Compute the new sort expression (with correct index) after projection:
-                    new_ordering.push(PhysicalSortExpr::new(
-                        Arc::new(Column::new(name, idx)),
-                        *options,
-                    ));
-                    continue;
-                }
-            }
-            // Cannot find expression in the projected_schema, stop iterating
-            // since rest of the orderings are violated
-            break;
-        }
-
-        let Some(new_ordering) = LexOrdering::new(new_ordering) else {
-            continue;
-        };
+    let Ok(projected_orderings) = project_orderings(
+        base_config.output_ordering.clone(),
+        &base_config.full_schema(),
+        &base_config.projection_indices(),
+    ) else {
+        return vec![];
+    };
 
+    let mut all_orderings = vec![];
+    for output_ordering in projected_orderings {
         // Check if any file groups are not sorted
         if base_config.file_groups.iter().any(|group| {
             if group.len() <= 1 {
@@ -1397,7 +1394,7 @@ fn get_projected_output_ordering(
             }
 
             let statistics = match MinMaxStatistics::new_from_files(
-                &new_ordering,
+                &output_ordering,
                 projected_schema,
                 base_config.projection.as_deref(),
                 group.iter(),
@@ -1420,7 +1417,7 @@ fn get_projected_output_ordering(
             continue;
         }
 
-        all_orderings.push(new_ordering);
+        all_orderings.push(output_ordering);
     }
     all_orderings
 }
@@ -1459,6 +1456,7 @@ mod tests {
     use datafusion_common::{assert_batches_eq, internal_err};
     use datafusion_expr::SortExpr;
     use datafusion_physical_expr::create_physical_sort_expr;
+    use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
 
     /// Returns the column names on the schema
     pub fn columns(schema: &Schema) -> Vec<String> {