Refactor TableProvider::scan into TableProvider::scan_with_args

adriangb · adriangb · commit 96ea7e87b07f · 2025-09-04T12:44:49.000-05:00
diff --git a/datafusion/catalog/src/table.rs b/datafusion/catalog/src/table.rs
@@ -25,7 +25,7 @@ use arrow::datatypes::SchemaRef;
 use async_trait::async_trait;
 use datafusion_common::Result;
 use datafusion_common::{not_impl_err, Constraints, Statistics};
-use datafusion_expr::Expr;
+use datafusion_expr::{Expr, SortExpr};
 
 use datafusion_expr::dml::InsertOp;
 use datafusion_expr::{
@@ -171,6 +171,41 @@ pub trait TableProvider: Debug + Sync + Send {
         limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>>;
 
+    /// Create an [`ExecutionPlan`] for scanning the table using structured arguments.
+    ///
+    /// This method uses [`ScanArgs`] to pass scan parameters in a structured way
+    /// and returns a [`ScanResult`] containing the execution plan. This approach
+    /// allows for extensible parameter passing and result handling.
+    ///
+    /// Table providers can override this method to take advantage of additional
+    /// parameters like `preferred_ordering` that may not be available through
+    /// other scan methods.
+    ///
+    /// # Arguments
+    /// * `state` - The session state containing configuration and context
+    /// * `args` - Structured scan arguments including projection, filters, limit, and ordering preferences
+    ///
+    /// # Returns
+    /// A [`ScanResult`] containing the [`ExecutionPlan`] for scanning the table
+    ///
+    /// See [`Self::scan`] for detailed documentation about projection, filters, and limits.
+    async fn scan_with_args(
+        &self,
+        state: &dyn Session,
+        args: ScanArgs,
+    ) -> Result<ScanResult> {
+        let ScanArgs {
+            filters,
+            projection,
+            limit,
+        } = args;
+        let filters = filters.unwrap_or_default();
+        let plan = self
+            .scan(state, projection.as_ref(), &filters, limit)
+            .await?;
+        Ok(ScanResult::new(plan))
+    }
+
     /// Specify if DataFusion should provide filter expressions to the
     /// TableProvider to apply *during* the scan.
     ///
@@ -299,6 +334,119 @@ pub trait TableProvider: Debug + Sync + Send {
     }
 }
 
+/// Arguments for scanning a table with [`TableProvider::scan_with_args`].
+///
+/// `ScanArgs` provides a structured way to pass scan parameters to table providers,
+/// replacing the multiple individual parameters used by [`TableProvider::scan`].
+/// This struct uses the builder pattern for convenient construction.
+///
+/// # Examples
+///
+/// ```
+/// # use datafusion_catalog::ScanArgs;
+/// # use datafusion_expr::Expr;
+/// let args = ScanArgs::default()
+///     .with_projection(Some(vec![0, 2, 4]))
+///     .with_limit(Some(1000));
+/// ```
+#[derive(Debug, Clone, Default)]
+pub struct ScanArgs {
+    filters: Option<Vec<Expr>>,
+    projection: Option<Vec<usize>>,
+    limit: Option<usize>,
+}
+
+impl ScanArgs {
+    /// Set the column projection for the scan.
+    ///
+    /// The projection is a list of column indices from [`TableProvider::schema`]
+    /// that should be included in the scan results. If `None`, all columns are included.
+    ///
+    /// # Arguments
+    /// * `projection` - Optional list of column indices to project
+    pub fn with_projection(mut self, projection: Option<Vec<usize>>) -> Self {
+        self.projection = projection;
+        self
+    }
+
+    /// Get the column projection for the scan.
+    ///
+    /// Returns a cloned copy of the projection column indices, or `None` if
+    /// no projection was specified (meaning all columns should be included).
+    pub fn projection(&self) -> Option<Vec<usize>> {
+        self.projection.clone()
+    }
+
+    /// Set the filter expressions for the scan.
+    ///
+    /// Filters are boolean expressions that should be evaluated during the scan
+    /// to reduce the number of rows returned. All expressions are combined with AND logic.
+    /// Whether filters are actually pushed down depends on [`TableProvider::supports_filters_pushdown`].
+    ///
+    /// # Arguments
+    /// * `filters` - Optional list of filter expressions
+    pub fn with_filters(mut self, filters: Option<Vec<Expr>>) -> Self {
+        self.filters = filters;
+        self
+    }
+
+    /// Get the filter expressions for the scan.
+    ///
+    /// Returns a reference to the filter expressions, or `None` if no filters were specified.
+    pub fn filters(&self) -> Option<&[Expr]> {
+        self.filters.as_deref()
+    }
+
+    /// Set the maximum number of rows to return from the scan.
+    ///
+    /// If specified, the scan should return at most this many rows. This is typically
+    /// used to optimize queries with `LIMIT` clauses.
+    ///
+    /// # Arguments
+    /// * `limit` - Optional maximum number of rows to return
+    pub fn with_limit(mut self, limit: Option<usize>) -> Self {
+        self.limit = limit;
+        self
+    }
+
+    /// Get the maximum number of rows to return from the scan.
+    ///
+    /// Returns the row limit, or `None` if no limit was specified.
+    pub fn limit(&self) -> Option<usize> {
+        self.limit
+    }
+}
+
+/// Result of a table scan operation from [`TableProvider::scan_with_args`].
+///
+/// `ScanResult` encapsulates the [`ExecutionPlan`] produced by a table scan,
+/// providing a typed return value instead of returning the plan directly.
+/// This allows for future extensibility of scan results without breaking
+/// the API.
+#[derive(Debug, Clone)]
+pub struct ScanResult {
+    /// The ExecutionPlan to run.
+    plan: Arc<dyn ExecutionPlan>,
+}
+
+impl ScanResult {
+    /// Create a new `ScanResult` with the given execution plan.
+    ///
+    /// # Arguments
+    /// * `plan` - The execution plan that will perform the table scan
+    pub fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
+        Self { plan }
+    }
+
+    /// Get the execution plan for this scan result.
+    ///
+    /// Returns a cloned reference to the [`ExecutionPlan`] that will perform
+    /// the actual table scanning and data retrieval.
+    pub fn plan(&self) -> Arc<dyn ExecutionPlan> {
+        Arc::clone(&self.plan)
+    }
+}
+
 /// A factory which creates [`TableProvider`]s at runtime given a URL.
 ///
 /// For example, this can be used to create a table "on the fly"
diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs
@@ -29,7 +29,7 @@ use crate::{
 use arrow::datatypes::{DataType, Field, SchemaBuilder, SchemaRef};
 use arrow_schema::Schema;
 use async_trait::async_trait;
-use datafusion_catalog::{Session, TableProvider};
+use datafusion_catalog::{ScanArgs, ScanResult, Session, TableProvider};
 use datafusion_common::{
     config_datafusion_err, config_err, internal_err, plan_err, project_schema,
     stats::Precision, Constraints, DataFusionError, Result, SchemaExt,
@@ -1166,6 +1166,22 @@ impl TableProvider for ListingTable {
         filters: &[Expr],
         limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
+        let options = ScanArgs::default()
+            .with_projection(projection.cloned())
+            .with_filters(Some(filters.to_vec()))
+            .with_limit(limit);
+        Ok(self.scan_with_args(state, options).await?.plan())
+    }
+
+    async fn scan_with_args(
+        &self,
+        state: &dyn Session,
+        args: ScanArgs,
+    ) -> Result<ScanResult> {
+        let projection = args.projection();
+        let filters = args.filters().map(|f| f.to_vec()).unwrap_or_default();
+        let limit = args.limit();
+
         // extract types of partition columns
         let table_partition_cols = self
             .options
@@ -1178,6 +1194,7 @@ impl TableProvider for ListingTable {
             .iter()
             .map(|field| field.name().as_str())
             .collect::<Vec<_>>();
+
         // If the filters can be resolved using only partition cols, there is no need to
         // pushdown it to TableScan, otherwise, `unhandled` pruning predicates will be generated
         let (partition_filters, filters): (Vec<_>, Vec<_>) =
@@ -1195,8 +1212,8 @@ impl TableProvider for ListingTable {
 
         // if no files need to be read, return an `EmptyExec`
         if partitioned_file_lists.is_empty() {
-            let projected_schema = project_schema(&self.schema(), projection)?;
-            return Ok(Arc::new(EmptyExec::new(projected_schema)));
+            let projected_schema = project_schema(&self.schema(), projection.as_ref())?;
+            return Ok(ScanResult::new(Arc::new(EmptyExec::new(projected_schema))));
         }
 
         let output_ordering = self.try_create_output_ordering()?;
@@ -1230,13 +1247,16 @@ impl TableProvider for ListingTable {
         let Some(object_store_url) =
             self.table_paths.first().map(ListingTableUrl::object_store)
         else {
-            return Ok(Arc::new(EmptyExec::new(Arc::new(Schema::empty()))));
+            return Ok(ScanResult::new(Arc::new(EmptyExec::new(Arc::new(
+                Schema::empty(),
+            )))));
         };
 
         let file_source = self.create_file_source_with_schema_adapter()?;
 
         // create the execution plan
-        self.options
+        let plan = self
+            .options
             .format
             .create_physical_plan(
                 state,
@@ -1248,14 +1268,16 @@ impl TableProvider for ListingTable {
                 .with_file_groups(partitioned_file_lists)
                 .with_constraints(self.constraints.clone())
                 .with_statistics(statistics)
-                .with_projection(projection.cloned())
+                .with_projection(projection)
                 .with_limit(limit)
                 .with_output_ordering(output_ordering)
                 .with_table_partition_cols(table_partition_cols)
                 .with_expr_adapter(self.expr_adapter_factory.clone())
                 .build(),
             )
-            .await
+            .await?;
+
+        Ok(ScanResult::new(plan))
     }
 
     fn supports_filters_pushdown(
diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs
@@ -60,6 +60,7 @@ use crate::schema_equivalence::schema_satisfied_by;
 use arrow::array::{builder::StringBuilder, RecordBatch};
 use arrow::compute::SortOptions;
 use arrow::datatypes::{Schema, SchemaRef};
+use datafusion_catalog::ScanArgs;
 use datafusion_common::display::ToStringifiedPlan;
 use datafusion_common::tree_node::{
     Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeVisitor,
@@ -459,9 +460,12 @@ impl DefaultPhysicalPlanner {
                 // doesn't know (nor should care) how the relation was
                 // referred to in the query
                 let filters = unnormalize_cols(filters.iter().cloned());
-                source
-                    .scan(session_state, projection.as_ref(), &filters, *fetch)
-                    .await?
+                let opts = ScanArgs::default()
+                    .with_projection(projection.clone())
+                    .with_filters(Some(filters))
+                    .with_limit(*fetch);
+                let res = source.scan_with_args(session_state, opts).await?;
+                res.plan()
             }
             LogicalPlan::Values(Values { values, schema }) => {
                 let exec_schema = schema.as_ref().to_owned().into();
