remove downcast

Author: xudong963

datafusion/datasource-parquet/src/opener.rs

@@ -69,13 +69,13 @@ use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader, RowGroupMe
 /// Implements [`FileOpener`] for a parquet file
 pub(super) struct ParquetOpener {
     /// Execution partition index
-    pub partition_index: usize,
+    pub(crate) partition_index: usize,
     /// Projection to apply on top of the table schema (i.e. can reference partition columns).
     pub projection: ProjectionExprs,
     /// Target number of rows in each output RecordBatch
     pub batch_size: usize,
     /// Optional limit on the number of rows to read
-    pub limit: Option<usize>,
+    pub(crate) limit: Option<usize>,
     /// If should keep the output rows in order
     pub preserve_order: bool,
     /// Optional predicate to apply during the scan

datafusion/datasource/src/file_scan_config.rs

@@ -484,6 +484,9 @@ impl FileScanConfigBuilder {
         let file_compression_type =
             file_compression_type.unwrap_or(FileCompressionType::UNCOMPRESSED);
 
+        // If there is an output ordering, we should preserve it.
+        let preserve_order = preserve_order || !output_ordering.is_empty();
+
         FileScanConfig {
             object_store_url,
             file_source,
@@ -869,6 +872,18 @@ impl DataSource for FileScanConfig {
             }
         }
     }
+
+    fn with_preserve_order(&self, preserve_order: bool) -> Option<Arc<dyn DataSource>> {
+        if self.preserve_order == preserve_order {
+            return Some(Arc::new(self.clone()));
+        }
+
+        let new_config = FileScanConfig {
+            preserve_order,
+            ..self.clone()
+        };
+        Some(Arc::new(new_config))
+    }
 }
 
 impl FileScanConfig {

datafusion/datasource/src/source.rs

@@ -210,6 +210,11 @@ pub trait DataSource: Send + Sync + Debug {
     ) -> Result<SortOrderPushdownResult<Arc<dyn DataSource>>> {
         Ok(SortOrderPushdownResult::Unsupported)
     }
+
+    /// Returns a variant of this `DataSource` that is aware of order-sensitivity.
+    fn with_preserve_order(&self, _preserve_order: bool) -> Option<Arc<dyn DataSource>> {
+        None
+    }
 }
 
 /// [`ExecutionPlan`] that reads one or more files
@@ -393,6 +398,18 @@ impl ExecutionPlan for DataSourceExec {
                 Ok(Arc::new(new_exec) as Arc<dyn ExecutionPlan>)
             })
     }
+
+    fn with_preserve_order(
+        &self,
+        preserve_order: bool,
+    ) -> Option<Arc<dyn ExecutionPlan>> {
+        self.data_source
+            .with_preserve_order(preserve_order)
+            .map(|new_data_source| {
+                Arc::new(self.clone().with_data_source(new_data_source))
+                    as Arc<dyn ExecutionPlan>
+            })
+    }
 }
 
 impl DataSourceExec {

datafusion/physical-optimizer/src/limit_pushdown.rs

@@ -27,8 +27,6 @@ use datafusion_common::config::ConfigOptions;
 use datafusion_common::error::Result;
 use datafusion_common::tree_node::{Transformed, TreeNodeRecursion};
 use datafusion_common::utils::combine_limit;
-use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
-use datafusion_datasource::source::DataSourceExec;
 use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
 use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
@@ -52,7 +50,7 @@ pub struct GlobalRequirements {
     fetch: Option<usize>,
     skip: usize,
     satisfied: bool,
-    order_sensitive: bool,
+    preserve_order: bool,
 }
 
 impl LimitPushdown {
@@ -72,7 +70,7 @@ impl PhysicalOptimizerRule for LimitPushdown {
             fetch: None,
             skip: 0,
             satisfied: false,
-            order_sensitive: false,
+            preserve_order: false,
         };
         pushdown_limits(plan, global_state)
     }
@@ -116,7 +114,7 @@ impl LimitExec {
         }
     }
 
-    fn order_sensitive(&self) -> bool {
+    fn preserve_order(&self) -> bool {
         match self {
             Self::Global(global) => global.required_ordering().is_some(),
             Self::Local(local) => local.required_ordering().is_some(),
@@ -156,7 +154,7 @@ pub fn pushdown_limit_helper(
         );
         global_state.skip = skip;
         global_state.fetch = fetch;
-        global_state.order_sensitive = limit_exec.order_sensitive();
+        global_state.preserve_order = limit_exec.preserve_order();
 
         // Now the global state has the most recent information, we can remove
         // the `LimitExec` plan. We will decide later if we should add it again
@@ -253,21 +251,19 @@ pub fn pushdown_limit_helper(
         let maybe_fetchable = pushdown_plan.with_fetch(skip_and_fetch);
         if global_state.satisfied {
             if let Some(plan_with_fetch) = maybe_fetchable {
-                let plan_with_preserve_order = ensure_preserve_order_if_needed(
-                    plan_with_fetch,
-                    global_state.order_sensitive,
-                );
+                let plan_with_preserve_order = plan_with_fetch
+                    .with_preserve_order(global_state.preserve_order)
+                    .unwrap_or(plan_with_fetch);
                 Ok((Transformed::yes(plan_with_preserve_order), global_state))
             } else {
                 Ok((Transformed::no(pushdown_plan), global_state))
             }
         } else {
             global_state.satisfied = true;
             pushdown_plan = if let Some(plan_with_fetch) = maybe_fetchable {
-                let plan_with_preserve_order = ensure_preserve_order_if_needed(
-                    plan_with_fetch,
-                    global_state.order_sensitive,
-                );
+                let plan_with_preserve_order = plan_with_fetch
+                    .with_preserve_order(global_state.preserve_order)
+                    .unwrap_or(plan_with_fetch);
 
                 if global_skip > 0 {
                     add_global_limit(
@@ -362,37 +358,4 @@ fn add_global_limit(
     Arc::new(GlobalLimitExec::new(pushdown_plan, skip, fetch)) as _
 }
 
-/// Helper function to handle DataSourceExec preserve_order setting
-fn ensure_preserve_order_if_needed(
-    plan: Arc<dyn ExecutionPlan>,
-    order_sensitive: bool,
-) -> Arc<dyn ExecutionPlan> {
-    if !order_sensitive {
-        return plan;
-    }
-
-    let Some(data_source_exec) = plan.as_any().downcast_ref::<DataSourceExec>() else {
-        return plan;
-    };
-
-    let Some(file_scan_config) = data_source_exec
-        .data_source()
-        .as_any()
-        .downcast_ref::<FileScanConfig>()
-    else {
-        return plan;
-    };
-
-    if file_scan_config.preserve_order {
-        return plan;
-    }
-
-    let new_config = FileScanConfigBuilder::from(file_scan_config.clone())
-        .with_preserve_order(true)
-        .build();
-
-    let new_data_source_exec = DataSourceExec::new(Arc::new(new_config));
-    Arc::new(new_data_source_exec) as Arc<dyn ExecutionPlan>
-}
-
 // See tests in datafusion/core/tests/physical_optimizer

datafusion/physical-plan/src/coalesce_partitions.rs

@@ -278,6 +278,19 @@ impl ExecutionPlan for CoalescePartitionsExec {
         }))
     }
 
+    fn with_preserve_order(
+        &self,
+        preserve_order: bool,
+    ) -> Option<Arc<dyn ExecutionPlan>> {
+        self.input
+            .with_preserve_order(preserve_order)
+            .and_then(|new_input| {
+                Arc::new(self.clone())
+                    .with_new_children(vec![new_input])
+                    .ok()
+            })
+    }
+
     fn gather_filters_for_pushdown(
         &self,
         _phase: FilterPushdownPhase,

datafusion/physical-plan/src/execution_plan.rs

@@ -708,6 +708,19 @@ pub trait ExecutionPlan: Debug + DisplayAs + Send + Sync {
     ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> {
         Ok(SortOrderPushdownResult::Unsupported)
     }
+
+    /// Returns a variant of this `ExecutionPlan` that is aware of order-sensitivity.
+    ///
+    /// This is used to signal to data sources that the output ordering must be
+    /// preserved, even if it might be more efficient to ignore it (e.g. by
+    /// skipping some row groups in Parquet).
+    ///
+    fn with_preserve_order(
+        &self,
+        _preserve_order: bool,
+    ) -> Option<Arc<dyn ExecutionPlan>> {
+        None
+    }
 }
 
 /// [`ExecutionPlan`] Invariant Level

datafusion/physical-plan/src/filter.rs

@@ -615,6 +615,19 @@ impl ExecutionPlan for FilterExec {
             fetch,
         }))
     }
+
+    fn with_preserve_order(
+        &self,
+        preserve_order: bool,
+    ) -> Option<Arc<dyn ExecutionPlan>> {
+        self.input
+            .with_preserve_order(preserve_order)
+            .and_then(|new_input| {
+                Arc::new(self.clone())
+                    .with_new_children(vec![new_input])
+                    .ok()
+            })
+    }
 }
 
 impl EmbeddedProjection for FilterExec {

datafusion/physical-plan/src/projection.rs

@@ -427,6 +427,19 @@ impl ExecutionPlan for ProjectionExec {
             }
         }
     }
+
+    fn with_preserve_order(
+        &self,
+        preserve_order: bool,
+    ) -> Option<Arc<dyn ExecutionPlan>> {
+        self.input
+            .with_preserve_order(preserve_order)
+            .and_then(|new_input| {
+                Arc::new(self.clone())
+                    .with_new_children(vec![new_input])
+                    .ok()
+            })
+    }
 }
 
 impl ProjectionStream {

datafusion/physical-plan/src/sorts/sort_preserving_merge.rs

@@ -245,6 +245,19 @@ impl ExecutionPlan for SortPreservingMergeExec {
         }))
     }
 
+    fn with_preserve_order(
+        &self,
+        preserve_order: bool,
+    ) -> Option<Arc<dyn ExecutionPlan>> {
+        self.input
+            .with_preserve_order(preserve_order)
+            .and_then(|new_input| {
+                Arc::new(self.clone())
+                    .with_new_children(vec![new_input])
+                    .ok()
+            })
+    }
+
     fn required_input_distribution(&self) -> Vec<Distribution> {
         vec![Distribution::UnspecifiedDistribution]
     }