Bug-fix: MemoryExec sort expressions do NOT refer to the projected schema

datafusion/core/src/datasource/memory.rs

@@ -37,14 +37,18 @@ use crate::physical_planner::create_physical_sort_exprs;
 
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
+use datafusion_catalog::Session;
 use datafusion_common::{not_impl_err, plan_err, Constraints, DFSchema, SchemaExt};
 use datafusion_execution::TaskContext;
 use datafusion_expr::dml::InsertOp;
+use datafusion_expr::SortExpr;
+use datafusion_physical_expr::equivalence::ProjectionMapping;
+use datafusion_physical_expr::expressions::Column;
+use datafusion_physical_expr::EquivalenceProperties;
 use datafusion_physical_plan::metrics::MetricsSet;
 
 use async_trait::async_trait;
-use datafusion_catalog::Session;
-use datafusion_expr::SortExpr;
+use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use futures::StreamExt;
 use log::debug;
 use parking_lot::Mutex;
@@ -231,7 +235,7 @@ impl TableProvider for MemTable {
         if !sort_order.is_empty() {
             let df_schema = DFSchema::try_from(self.schema.as_ref().clone())?;
 
-            let file_sort_order = sort_order
+            let mut file_sort_order = sort_order
                 .iter()
                 .map(|sort_exprs| {
                     create_physical_sort_exprs(
@@ -241,6 +245,31 @@ impl TableProvider for MemTable {
                     )
                 })
                 .collect::<Result<Vec<_>>>()?;
+
+            // If there is a projection on the source, we also need to project orderings
+            if let Some(projection) = projection {
+                let base_eqp = EquivalenceProperties::new_with_orderings(
+                    self.schema(),
+                    &file_sort_order,
+                );
+                let proj_exprs = projection
+                    .iter()
+                    .map(|idx| {
+                        let name = self.schema.field(*idx).name();
+                        (
+                            Arc::new(Column::new(name, *idx)) as Arc<dyn PhysicalExpr>,
+                            name.to_string(),
+                        )
+                    })
+                    .collect::<Vec<_>>();
+                let projection_mapping =
+                    ProjectionMapping::try_new(&proj_exprs, &self.schema)?;
+                file_sort_order = base_eqp
+                    .project(&projection_mapping, exec.schema())
+                    .oeq_class
+                    .orderings;
+            }
+
             exec = exec.with_sort_information(file_sort_order);
         }
 

datafusion/physical-plan/src/memory.rs

@@ -35,6 +35,7 @@ use datafusion_execution::memory_pool::MemoryReservation;
 use datafusion_execution::TaskContext;
 use datafusion_physical_expr::{EquivalenceProperties, LexOrdering};
 
+use datafusion_physical_expr::utils::collect_columns;
 use futures::Stream;
 
 /// Execution plan for reading in-memory batches of data
@@ -207,6 +208,20 @@ impl MemoryExec {
     /// [`EquivalenceProperties`], we can keep track of these equivalences
     /// and treat `a ASC` and `b DESC` as the same ordering requirement.
     pub fn with_sort_information(mut self, sort_information: Vec<LexOrdering>) -> Self {
+        // All sort expressions must refer to the projected schema
+        debug_assert!({
+            let fields = self.projected_schema.fields();
+            sort_information
+                .iter()
+                .flatten()
+                .flat_map(|expr| collect_columns(&expr.expr))
+                .all(|col| {
+                    fields
+                        .get(col.index())
+                        .map(|field| field.name() == col.name())
+                        .unwrap_or(false)
+                })
+        });
         self.sort_information = sort_information;
 
         // We need to update equivalence properties when updating sort information.