Merge branch 'main' into feat/sort-file-groups-by-statistics

Author: zhuqi-lucas

datafusion/datasource/src/memory.rs

@@ -243,13 +243,22 @@ impl DataSource for MemorySourceConfig {
                     &exprs,
                     self.projection().as_ref().unwrap_or(&all_projections),
                 );
-
-                MemorySourceConfig::try_new(
-                    self.partitions(),
-                    self.original_schema(),
-                    Some(new_projections),
-                )
-                .map(|s| Arc::new(s) as Arc<dyn DataSource>)
+                let projected_schema =
+                    project_schema(&self.schema, Some(&new_projections));
+
+                projected_schema.map(|projected_schema| {
+                    // Clone self to preserve all metadata (fetch, sort_information,
+                    // show_sizes, etc.) then update only the projection-related fields.
+                    let mut new_source = self.clone();
+                    new_source.projection = Some(new_projections);
+                    new_source.projected_schema = projected_schema;
+                    // Project sort information to match the new projection
+                    new_source.sort_information = project_orderings(
+                        &new_source.sort_information,
+                        &new_source.projected_schema,
+                    );
+                    Arc::new(new_source) as Arc<dyn DataSource>
+                })
             })
             .transpose()
     }
@@ -897,6 +906,42 @@ mod tests {
         Ok(())
     }
 
+    /// Test that `try_swapping_with_projection` preserves the `fetch` limit.
+    /// Regression test for <https://github.com/apache/datafusion/issues/21176>
+    #[test]
+    fn try_swapping_with_projection_preserves_fetch() {
+        use datafusion_physical_expr::projection::ProjectionExprs;
+
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int32, false),
+            Field::new("b", DataType::Utf8, false),
+            Field::new("c", DataType::Int64, false),
+        ]));
+        let partitions: Vec<Vec<RecordBatch>> = vec![vec![batch(10)]];
+        let source = MemorySourceConfig::try_new(&partitions, schema.clone(), None)
+            .unwrap()
+            .with_limit(Some(5));
+
+        assert_eq!(source.fetch, Some(5));
+
+        // Create a projection that reorders columns: [c, a] (indices 2, 0)
+        let projection = ProjectionExprs::from_indices(&[2, 0], &schema);
+        let swapped = source
+            .try_swapping_with_projection(&projection)
+            .unwrap()
+            .unwrap();
+        let new_source = swapped
+            .as_any()
+            .downcast_ref::<MemorySourceConfig>()
+            .unwrap();
+
+        assert_eq!(
+            new_source.fetch,
+            Some(5),
+            "fetch limit must be preserved after projection pushdown"
+        );
+    }
+
     #[tokio::test]
     async fn values_empty_case() -> Result<()> {
         let schema = aggr_test_schema();