Skip to main content

datafusion_datasource/file_scan_config/
mod.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! [`FileScanConfig`] to configure scanning of possibly partitioned
19//! file sources.
20
21pub(crate) mod sort_pushdown;
22
23/// Shared `FileScanConfig` <-> proto conversion, gated on the `proto` feature.
24/// Attaches inherent `try_to_proto` / `try_from_proto` /
25/// `parse_table_schema_from_proto` helpers to [`FileScanConfig`] used by every
26/// file source's `try_to_proto` hook.
27#[cfg(feature = "proto")]
28mod proto;
29
30use crate::file_groups::FileGroup;
31use crate::{
32    PartitionedFile, display::FileGroupsDisplay, file::FileSource,
33    file_compression_type::FileCompressionType, file_stream::FileStreamBuilder,
34    file_stream::work_source::SharedWorkSource, source::DataSource,
35    statistics::MinMaxStatistics,
36};
37use arrow::datatypes::Fields;
38use arrow::datatypes::{DataType, Schema, SchemaRef};
39use datafusion_common::config::ConfigOptions;
40use datafusion_common::tree_node::TreeNodeRecursion;
41use datafusion_common::{
42    Constraints, Result, ScalarValue, Statistics, internal_datafusion_err, internal_err,
43};
44use datafusion_execution::{
45    SendableRecordBatchStream, TaskContext, object_store::ObjectStoreUrl,
46};
47use datafusion_expr::Operator;
48
49use crate::source::OpenArgs;
50use datafusion_common::stats::Precision;
51use datafusion_physical_expr::expressions::{BinaryExpr, Column};
52use datafusion_physical_expr::projection::{ProjectionExprs, ProjectionMapping};
53use datafusion_physical_expr::utils::reassign_expr_columns;
54use datafusion_physical_expr::{EquivalenceProperties, Partitioning, split_conjunction};
55use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory;
56use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, is_volatile};
57use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
58use datafusion_physical_plan::SortOrderPushdownResult;
59use datafusion_physical_plan::coop::cooperative;
60use datafusion_physical_plan::execution_plan::SchedulingType;
61use datafusion_physical_plan::{
62    DisplayAs, DisplayFormatType,
63    display::{ProjectSchemaDisplay, display_orderings},
64    filter_pushdown::FilterPushdownPropagation,
65    metrics::ExecutionPlanMetricsSet,
66};
67use log::{debug, warn};
68use std::any::Any;
69use std::{fmt::Debug, fmt::Formatter, fmt::Result as FmtResult, sync::Arc};
70
71/// [`FileScanConfig`] represents scanning data from a group of files
72///
73/// `FileScanConfig` is used to create a [`DataSourceExec`], the physical plan
74/// for scanning files with a particular file format.
75///
76/// The [`FileSource`] (e.g. `ParquetSource`, `CsvSource`, etc.) is responsible
77/// for creating the actual execution plan to read the files based on a
78/// `FileScanConfig`. Fields in a `FileScanConfig` such as Statistics represent
79/// information about the files **before** any projection or filtering is
80/// applied in the file source.
81///
82/// Use [`FileScanConfigBuilder`] to construct a `FileScanConfig`.
83///
84/// Use [`DataSourceExec::from_data_source`] to create a [`DataSourceExec`] from
85/// a `FileScanConfig`.
86///
87/// # Example
88/// ```
89/// # use std::sync::Arc;
90/// # use arrow::datatypes::{Field, Fields, DataType, Schema, SchemaRef};
91/// # use object_store::ObjectStore;
92/// # use datafusion_common::Result;
93/// # use datafusion_common::tree_node::TreeNodeRecursion;
94/// # use datafusion_datasource::file::FileSource;
95/// # use datafusion_physical_plan::PhysicalExpr;
96/// # use datafusion_datasource::file_groups::FileGroup;
97/// # use datafusion_datasource::PartitionedFile;
98/// # use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
99/// # use datafusion_datasource::file_stream::FileOpener;
100/// # use datafusion_datasource::source::DataSourceExec;
101/// # use datafusion_datasource::table_schema::TableSchema;
102/// # use datafusion_execution::object_store::ObjectStoreUrl;
103/// # use datafusion_physical_expr::projection::ProjectionExprs;
104/// # use datafusion_physical_plan::ExecutionPlan;
105/// # use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
106/// # let file_schema = Arc::new(Schema::new(vec![
107/// #  Field::new("c1", DataType::Int32, false),
108/// #  Field::new("c2", DataType::Int32, false),
109/// #  Field::new("c3", DataType::Int32, false),
110/// #  Field::new("c4", DataType::Int32, false),
111/// # ]));
112/// # // Note: crate mock ParquetSource, as ParquetSource is not in the datasource crate
113/// #[derive(Clone)]
114/// # struct ParquetSource {
115/// #    table_schema: TableSchema,
116/// # };
117/// # impl FileSource for ParquetSource {
118/// #  fn create_file_opener(&self, _: Arc<dyn ObjectStore>, _: &FileScanConfig, _: usize) -> Result<Arc<dyn FileOpener>> { unimplemented!() }
119/// #  fn table_schema(&self) -> &TableSchema { &self.table_schema }
120/// #  fn with_batch_size(&self, _: usize) -> Arc<dyn FileSource> { unimplemented!() }
121/// #  fn metrics(&self) -> &ExecutionPlanMetricsSet { unimplemented!() }
122/// #  fn file_type(&self) -> &str { "parquet" }
123/// #  // Note that this implementation drops the projection on the floor, it is not complete!
124/// #  fn try_pushdown_projection(&self, projection: &ProjectionExprs) -> Result<Option<Arc<dyn FileSource>>> { Ok(Some(Arc::new(self.clone()) as Arc<dyn FileSource>)) }
125/// #  fn apply_expressions(&self, _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>) -> Result<TreeNodeRecursion> { Ok(TreeNodeRecursion::Continue) }
126/// #  }
127/// # impl ParquetSource {
128/// #  fn new(table_schema: impl Into<TableSchema>) -> Self { Self {table_schema: table_schema.into()} }
129/// # }
130/// // create FileScan config for reading parquet files from file://
131/// let object_store_url = ObjectStoreUrl::local_filesystem();
132/// let file_source = Arc::new(ParquetSource::new(file_schema.clone()));
133/// let config = FileScanConfigBuilder::new(object_store_url, file_source)
134///   .with_limit(Some(1000))            // read only the first 1000 records
135///   .with_projection_indices(Some(vec![2, 3])) // project columns 2 and 3
136///   .expect("Failed to push down projection")
137///    // Read /tmp/file1.parquet with known size of 1234 bytes in a single group
138///   .with_file(PartitionedFile::new("file1.parquet", 1234))
139///   // Read /tmp/file2.parquet 56 bytes and /tmp/file3.parquet 78 bytes
140///   // in a  single row group
141///   .with_file_group(FileGroup::new(vec![
142///    PartitionedFile::new("file2.parquet", 56),
143///    PartitionedFile::new("file3.parquet", 78),
144///   ])).build();
145/// // create an execution plan from the config
146/// let plan: Arc<dyn ExecutionPlan> = DataSourceExec::from_data_source(config);
147/// ```
148///
149/// [`DataSourceExec`]: crate::source::DataSourceExec
150/// [`DataSourceExec::from_data_source`]: crate::source::DataSourceExec::from_data_source
151#[derive(Clone)]
152pub struct FileScanConfig {
153    /// Object store URL, used to get an [`ObjectStore`] instance from
154    /// [`RuntimeEnv::object_store`]
155    ///
156    /// This `ObjectStoreUrl` should be the prefix of the absolute url for files
157    /// as `file://` or `s3://my_bucket`. It should not include the path to the
158    /// file itself. The relevant URL prefix must be registered via
159    /// [`RuntimeEnv::register_object_store`]
160    ///
161    /// [`ObjectStore`]: object_store::ObjectStore
162    /// [`RuntimeEnv::register_object_store`]: datafusion_execution::runtime_env::RuntimeEnv::register_object_store
163    /// [`RuntimeEnv::object_store`]: datafusion_execution::runtime_env::RuntimeEnv::object_store
164    pub object_store_url: ObjectStoreUrl,
165    /// List of files to be processed, grouped into partitions
166    ///
167    /// Each file must have a schema of `file_schema` or a subset. If
168    /// a particular file has a subset, the missing columns are
169    /// padded with NULLs.
170    ///
171    /// DataFusion may attempt to read each partition of files
172    /// concurrently, however files *within* a partition will be read
173    /// sequentially, one after the next.
174    ///
175    /// Note that when `datafusion.execution.enable_file_stream_work_stealing`
176    /// is enabled (the default), files may be reassigned to a different
177    /// partition at runtime unless `preserve_order` or
178    /// `partitioned_by_file_group` is set, so a file is not guaranteed to be
179    /// read by the partition it is grouped under here.
180    pub file_groups: Vec<FileGroup>,
181    /// Table constraints
182    pub constraints: Constraints,
183    /// The maximum number of records to read from this plan. If `None`,
184    /// all records after filtering are returned.
185    pub limit: Option<usize>,
186    /// Whether the scan's limit is order sensitive
187    /// When `true`, files must be read in the exact order specified to produce
188    /// correct results (e.g., for `ORDER BY ... LIMIT` queries). When `false`,
189    /// DataFusion may reorder file processing for optimization without affecting correctness.
190    pub preserve_order: bool,
191    /// All equivalent lexicographical output orderings of this file scan, in terms of
192    /// [`FileSource::table_schema`]. See [`FileScanConfigBuilder::with_output_ordering`] for more
193    /// details.
194    ///
195    /// [`Self::eq_properties`] uses this information along with projection
196    /// and filtering information to compute the effective
197    /// [`EquivalenceProperties`]
198    pub output_ordering: Vec<LexOrdering>,
199    /// File compression type
200    pub file_compression_type: FileCompressionType,
201    /// File source such as `ParquetSource`, `CsvSource`, `JsonSource`, etc.
202    pub file_source: Arc<dyn FileSource>,
203    /// Batch size while creating new batches
204    /// Defaults to [`datafusion_common::config::ExecutionOptions`] batch_size.
205    pub batch_size: Option<usize>,
206    /// Expression adapter used to adapt filters and projections that are pushed down into the scan
207    /// from the logical schema to the physical schema of the file.
208    pub expr_adapter_factory: Option<Arc<dyn PhysicalExprAdapterFactory>>,
209    /// Statistics for the entire table (file schema + partition columns).
210    /// See [`FileScanConfigBuilder::with_statistics`] for more details.
211    ///
212    /// The effective statistics are computed on-demand via
213    /// [`ProjectionExprs::project_statistics`].
214    ///
215    /// Note that this field is pub(crate) because accessing it directly from outside
216    /// would be incorrect if there are filters being applied, thus this should be accessed
217    /// via [`FileScanConfig::statistics`].
218    pub(crate) statistics: Statistics,
219    /// Declared physical output partitioning for this scan.
220    ///
221    /// Expressions are against the full table schema, before scan projection or
222    /// filtering. `ListingTable` validates partition count before building the
223    /// scan, and direct builders with mismatched counts fall back to
224    /// `UnknownPartitioning`.
225    pub output_partitioning: Option<Partitioning>,
226}
227
228/// A builder for [`FileScanConfig`]'s.
229///
230/// Example:
231///
232/// ```rust
233/// # use std::sync::Arc;
234/// # use arrow::datatypes::{DataType, Field, Schema};
235/// # use datafusion_datasource::file_scan_config::{FileScanConfigBuilder, FileScanConfig};
236/// # use datafusion_datasource::file_compression_type::FileCompressionType;
237/// # use datafusion_datasource::file_groups::FileGroup;
238/// # use datafusion_datasource::PartitionedFile;
239/// # use datafusion_datasource::table_schema::TableSchema;
240/// # use datafusion_execution::object_store::ObjectStoreUrl;
241/// # use datafusion_common::Statistics;
242/// # use datafusion_datasource::file::FileSource;
243///
244/// # fn main() {
245/// # fn with_source(file_source: Arc<dyn FileSource>) {
246///     // Create a schema for our Parquet files
247///     let file_schema = Arc::new(Schema::new(vec![
248///         Field::new("id", DataType::Int32, false),
249///         Field::new("value", DataType::Utf8, false),
250///     ]));
251///
252///     // Create partition columns
253///     let partition_cols = vec![
254///         Arc::new(Field::new("date", DataType::Utf8, false)),
255///     ];
256///
257///     // Create table schema with file schema and partition columns
258///     let table_schema = TableSchema::builder(file_schema)
259///         .with_table_partition_cols(partition_cols)
260///         .build();
261///
262///     // Create a builder for scanning Parquet files from a local filesystem
263///     let config = FileScanConfigBuilder::new(
264///         ObjectStoreUrl::local_filesystem(),
265///         file_source,
266///     )
267///     // Set a limit of 1000 rows
268///     .with_limit(Some(1000))
269///     // Project only the first column
270///     .with_projection_indices(Some(vec![0]))
271///     .expect("Failed to push down projection")
272///     // Add a file group with two files
273///     .with_file_group(FileGroup::new(vec![
274///         PartitionedFile::new("data/date=2024-01-01/file1.parquet", 1024),
275///         PartitionedFile::new("data/date=2024-01-01/file2.parquet", 2048),
276///     ]))
277///     // Set compression type
278///     .with_file_compression_type(FileCompressionType::UNCOMPRESSED)
279///     // Build the final config
280///     .build();
281/// # }
282/// # }
283/// ```
284#[derive(Clone)]
285pub struct FileScanConfigBuilder {
286    object_store_url: ObjectStoreUrl,
287    file_source: Arc<dyn FileSource>,
288    limit: Option<usize>,
289    preserve_order: bool,
290    constraints: Option<Constraints>,
291    file_groups: Vec<FileGroup>,
292    statistics: Option<Statistics>,
293    output_ordering: Vec<LexOrdering>,
294    output_partitioning: Option<Partitioning>,
295    file_compression_type: Option<FileCompressionType>,
296    batch_size: Option<usize>,
297    expr_adapter_factory: Option<Arc<dyn PhysicalExprAdapterFactory>>,
298}
299
300impl FileScanConfigBuilder {
301    /// Create a new [`FileScanConfigBuilder`] with default settings for scanning files.
302    ///
303    /// # Parameters:
304    /// * `object_store_url`: See [`FileScanConfig::object_store_url`]
305    /// * `file_source`: See [`FileScanConfig::file_source`]. The file source must have
306    ///   a schema set via its constructor.
307    pub fn new(
308        object_store_url: ObjectStoreUrl,
309        file_source: Arc<dyn FileSource>,
310    ) -> Self {
311        Self {
312            object_store_url,
313            file_source,
314            file_groups: vec![],
315            statistics: None,
316            output_ordering: vec![],
317            output_partitioning: None,
318            file_compression_type: None,
319            limit: None,
320            preserve_order: false,
321            constraints: None,
322            batch_size: None,
323            expr_adapter_factory: None,
324        }
325    }
326
327    /// Set the maximum number of records to read from this plan.
328    ///
329    /// If `None`, all records after filtering are returned.
330    pub fn with_limit(mut self, limit: Option<usize>) -> Self {
331        self.limit = limit;
332        self
333    }
334
335    /// Set whether the limit should be order-sensitive.
336    ///
337    /// When `true`, files must be read in the exact order specified to produce
338    /// correct results (e.g., for `ORDER BY ... LIMIT` queries). When `false`,
339    /// DataFusion may reorder file processing for optimization without
340    /// affecting correctness.
341    pub fn with_preserve_order(mut self, order_sensitive: bool) -> Self {
342        self.preserve_order = order_sensitive;
343        self
344    }
345
346    /// Set the file source for scanning files.
347    ///
348    /// This method allows you to change the file source implementation (e.g.
349    /// ParquetSource, CsvSource, etc.) after the builder has been created.
350    pub fn with_source(mut self, file_source: Arc<dyn FileSource>) -> Self {
351        self.file_source = file_source;
352        self
353    }
354
355    /// Return the table schema
356    pub fn table_schema(&self) -> &SchemaRef {
357        self.file_source.table_schema().table_schema()
358    }
359
360    /// Set the columns on which to project the data. Indexes that are higher than the
361    /// number of columns of `file_schema` refer to `table_partition_cols`.
362    ///
363    /// # Deprecated
364    /// Use [`Self::with_projection_indices`] instead. This method will be removed in a future release.
365    #[deprecated(since = "51.0.0", note = "Use with_projection_indices instead")]
366    pub fn with_projection(self, indices: Option<Vec<usize>>) -> Self {
367        match self.clone().with_projection_indices(indices) {
368            Ok(builder) => builder,
369            Err(e) => {
370                warn!(
371                    "Failed to push down projection in FileScanConfigBuilder::with_projection: {e}"
372                );
373                self
374            }
375        }
376    }
377
378    /// Set the columns on which to project the data using column indices.
379    ///
380    /// This method attempts to push down the projection to the underlying file
381    /// source if supported. If the file source does not support projection
382    /// pushdown, an error is returned.
383    ///
384    /// Indexes that are higher than the number of columns of `file_schema`
385    /// refer to `table_partition_cols`.
386    pub fn with_projection_indices(
387        mut self,
388        indices: Option<Vec<usize>>,
389    ) -> Result<Self> {
390        let projection_exprs = indices.map(|indices| {
391            ProjectionExprs::from_indices(
392                &indices,
393                self.file_source.table_schema().table_schema(),
394            )
395        });
396        let Some(projection_exprs) = projection_exprs else {
397            return Ok(self);
398        };
399        let new_source = self
400            .file_source
401            .try_pushdown_projection(&projection_exprs)
402            .map_err(|e| {
403                internal_datafusion_err!(
404                    "Failed to push down projection in FileScanConfigBuilder::build: {e}"
405                )
406            })?;
407        if let Some(new_source) = new_source {
408            self.file_source = new_source;
409        } else {
410            internal_err!(
411                "FileSource {} does not support projection pushdown",
412                self.file_source.file_type()
413            )?;
414        }
415        Ok(self)
416    }
417
418    /// Set the table constraints
419    pub fn with_constraints(mut self, constraints: Constraints) -> Self {
420        self.constraints = Some(constraints);
421        self
422    }
423
424    /// Set the statistics of the files, including partition
425    /// columns. Defaults to [`Statistics::new_unknown`].
426    ///
427    /// These statistics are for the entire table (file schema + partition
428    /// columns) before any projection or filtering is applied. Projections are
429    /// applied when statistics are retrieved, and if a filter is present,
430    /// [`FileScanConfig::statistics`] will mark the statistics as inexact
431    /// (counts are not adjusted).
432    ///
433    /// Projections and filters may be applied by the file source, either by
434    /// [`Self::with_projection_indices`] or a preexisting
435    /// [`FileSource::projection`] or [`FileSource::filter`].
436    pub fn with_statistics(mut self, statistics: Statistics) -> Self {
437        self.statistics = Some(statistics);
438        self
439    }
440
441    /// Set the list of files to be processed, grouped into partitions.
442    ///
443    /// Each file must have a schema of `file_schema` or a subset. If
444    /// a particular file has a subset, the missing columns are
445    /// padded with NULLs.
446    ///
447    /// DataFusion may attempt to read each partition of files
448    /// concurrently, however files *within* a partition will be read
449    /// sequentially, one after the next.
450    pub fn with_file_groups(mut self, file_groups: Vec<FileGroup>) -> Self {
451        self.file_groups = file_groups;
452        self
453    }
454
455    /// Add a new file group
456    ///
457    /// See [`Self::with_file_groups`] for more information
458    pub fn with_file_group(mut self, file_group: FileGroup) -> Self {
459        self.file_groups.push(file_group);
460        self
461    }
462
463    /// Add a file as a single group
464    ///
465    /// See [`Self::with_file_groups`] for more information.
466    pub fn with_file(self, partitioned_file: PartitionedFile) -> Self {
467        self.with_file_group(FileGroup::new(vec![partitioned_file]))
468    }
469
470    /// Set the output ordering of the files
471    ///
472    /// The expressions are in terms of the entire table schema (file schema +
473    /// partition columns), before any projection or filtering from the file
474    /// scan is applied.
475    ///
476    /// This is used for optimization purposes, e.g. to determine if a file scan
477    /// can satisfy an `ORDER BY` without an additional sort.
478    pub fn with_output_ordering(mut self, output_ordering: Vec<LexOrdering>) -> Self {
479        self.output_ordering = output_ordering;
480        self
481    }
482
483    /// Set declared physical output partitioning for this scan.
484    pub fn with_output_partitioning(
485        mut self,
486        output_partitioning: Option<Partitioning>,
487    ) -> Self {
488        self.output_partitioning = output_partitioning;
489        self
490    }
491
492    /// Set the file compression type
493    pub fn with_file_compression_type(
494        mut self,
495        file_compression_type: FileCompressionType,
496    ) -> Self {
497        self.file_compression_type = Some(file_compression_type);
498        self
499    }
500
501    /// Set the batch_size property
502    pub fn with_batch_size(mut self, batch_size: Option<usize>) -> Self {
503        self.batch_size = batch_size;
504        self
505    }
506
507    /// Register an expression adapter used to adapt filters and projections that are pushed down into the scan
508    /// from the logical schema to the physical schema of the file.
509    /// This can include things like:
510    /// - Column ordering changes
511    /// - Handling of missing columns
512    /// - Rewriting expression to use pre-computed values or file format specific optimizations
513    pub fn with_expr_adapter(
514        mut self,
515        expr_adapter: Option<Arc<dyn PhysicalExprAdapterFactory>>,
516    ) -> Self {
517        self.expr_adapter_factory = expr_adapter;
518        self
519    }
520
521    /// Build the final [`FileScanConfig`] with all the configured settings.
522    ///
523    /// This method takes ownership of the builder and returns the constructed `FileScanConfig`.
524    /// Any unset optional fields will use their default values.
525    ///
526    /// # Errors
527    /// Returns an error if projection pushdown fails or if schema operations fail.
528    pub fn build(self) -> FileScanConfig {
529        let Self {
530            object_store_url,
531            file_source,
532            limit,
533            preserve_order,
534            constraints,
535            file_groups,
536            statistics,
537            output_ordering,
538            output_partitioning,
539            file_compression_type,
540            batch_size,
541            expr_adapter_factory: expr_adapter,
542        } = self;
543
544        let constraints = constraints.unwrap_or_default();
545        let statistics = statistics.unwrap_or_else(|| {
546            Statistics::new_unknown(file_source.table_schema().table_schema())
547        });
548        let file_compression_type =
549            file_compression_type.unwrap_or(FileCompressionType::UNCOMPRESSED);
550
551        // If there is an output ordering, we should preserve it.
552        let preserve_order = preserve_order || !output_ordering.is_empty();
553
554        FileScanConfig {
555            object_store_url,
556            file_source,
557            limit,
558            preserve_order,
559            constraints,
560            file_groups,
561            output_ordering,
562            file_compression_type,
563            batch_size,
564            expr_adapter_factory: expr_adapter,
565            statistics,
566            output_partitioning,
567        }
568    }
569}
570
571impl From<FileScanConfig> for FileScanConfigBuilder {
572    fn from(config: FileScanConfig) -> Self {
573        Self {
574            object_store_url: config.object_store_url,
575            file_source: Arc::<dyn FileSource>::clone(&config.file_source),
576            file_groups: config.file_groups,
577            statistics: Some(config.statistics),
578            output_ordering: config.output_ordering,
579            output_partitioning: config.output_partitioning,
580            file_compression_type: Some(config.file_compression_type),
581            limit: config.limit,
582            preserve_order: config.preserve_order,
583            constraints: Some(config.constraints),
584            batch_size: config.batch_size,
585            expr_adapter_factory: config.expr_adapter_factory,
586        }
587    }
588}
589
590/// Builds output partitioning over `partition_cols` (resolved to their indices in
591/// `schema`) with `partition_count` partitions. Returns `None` when there are no
592/// partition columns. Callers use this to declare the output partitioning of a scan
593/// whose file groups are organized by partition column values.
594pub fn output_partitioning_from_partition_fields(
595    schema: &Schema,
596    partition_cols: &Fields,
597    partition_count: usize,
598) -> Option<Partitioning> {
599    if partition_cols.is_empty() {
600        return None;
601    }
602
603    let mut exprs: Vec<Arc<dyn PhysicalExpr>> = Vec::with_capacity(partition_cols.len());
604    for partition_col in partition_cols {
605        let name = partition_col.name();
606        let idx = schema
607            .fields()
608            .iter()
609            .position(|field| field.name() == name)?;
610        exprs.push(Arc::new(Column::new(name, idx)));
611    }
612
613    Some(Partitioning::Hash(exprs, partition_count))
614}
615
616fn project_output_partitioning(
617    partitioning: &Partitioning,
618    mapping: &ProjectionMapping,
619    input_schema: &SchemaRef,
620    partition_count: usize,
621) -> Partitioning {
622    let input_eq_properties = EquivalenceProperties::new(Arc::clone(input_schema));
623    match partitioning {
624        Partitioning::Hash(exprs, _) => {
625            let projected_exprs = input_eq_properties
626                .project_expressions(exprs, mapping)
627                .collect::<Option<Vec<_>>>();
628            projected_exprs
629                .map(|exprs| Partitioning::Hash(exprs, partition_count))
630                .unwrap_or_else(|| Partitioning::UnknownPartitioning(partition_count))
631        }
632        Partitioning::Range(_)
633        | Partitioning::RoundRobinBatch(_)
634        | Partitioning::UnknownPartitioning(_) => {
635            partitioning.project(mapping, &input_eq_properties)
636        }
637    }
638}
639
640/// Returns `true` if merging `outer` into `inner` would duplicate a volatile or
641/// non-trivial expression that CSE deduplicated; the caller should then decline
642/// the merge.
643///
644/// Merging substitutes each `inner` expression into every `outer` reference to
645/// it. Since the logical optimizer extracts a repeated expression into a single
646/// `inner` entry referenced by column, re-inlining it at more than one
647/// reference site undoes that deduplication. An `inner` expression referenced
648/// more than once is therefore blocked when it is either:
649///
650/// - **volatile** (e.g. `random()`) — evaluating it independently at each site
651///   makes references that should share one "locked-in" value diverge (the
652///   correctness guard the physical `ProjectionPushdown` and `FilterPushdown`
653///   rules also apply via
654///   `datafusion_physical_expr_common::physical_expr::is_volatile`); or
655/// - **not cheap to recompute** — its placement is not push-to-leaves
656///   (`KeepInPlace`: arithmetic, casts, most scalar functions). Leaf-pushable
657///   expressions (columns, `get_field`, `input_file_name`) still merge. This
658///   matches `try_collapse_projection_chain`.
659///
660/// References are counted with multiplicity, so `r + r` counts as two; an
661/// expression referenced exactly once has nothing to duplicate.
662fn would_duplicate_costly_exprs(
663    inner: &ProjectionExprs,
664    outer: &ProjectionExprs,
665) -> bool {
666    use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
667
668    let inner_exprs = inner.as_ref();
669
670    let mut ref_counts = vec![0usize; inner_exprs.len()];
671    for proj_expr in outer.as_ref() {
672        proj_expr
673            .expr
674            .apply(|e| {
675                if let Some(col) = e.as_ref().downcast_ref::<Column>()
676                    && let Some(count) = ref_counts.get_mut(col.index())
677                {
678                    *count += 1;
679                }
680                Ok(TreeNodeRecursion::Continue)
681            })
682            .expect("infallible closure should not fail");
683    }
684
685    ref_counts.iter().enumerate().any(|(idx, &count)| {
686        let expr = &inner_exprs[idx].expr;
687        count > 1 && (is_volatile(expr) || !expr.placement().should_push_to_leaves())
688    })
689}
690
691impl DataSource for FileScanConfig {
692    fn open(
693        &self,
694        partition: usize,
695        context: Arc<TaskContext>,
696    ) -> Result<SendableRecordBatchStream> {
697        self.open_with_args(OpenArgs::new(partition, context))
698    }
699
700    fn open_with_args(&self, args: OpenArgs) -> Result<SendableRecordBatchStream> {
701        let OpenArgs {
702            partition,
703            context,
704            sibling_state,
705        } = args;
706        let object_store = context.runtime_env().object_store(&self.object_store_url)?;
707        let batch_size = self
708            .batch_size
709            .unwrap_or_else(|| context.session_config().batch_size());
710
711        let source = self.file_source.with_batch_size(batch_size);
712
713        let morselizer = source.create_morselizer(object_store, self, partition)?;
714
715        // Extract the shared work source from the sibling state if it exists.
716        // This allows multiple sibling streams to steal work from a single
717        // shared queue of unopened files.
718        let shared_work_source = sibling_state
719            .as_ref()
720            .and_then(|state| state.downcast_ref::<SharedWorkSource>())
721            .cloned();
722
723        let stream = FileStreamBuilder::new(self)
724            .with_partition(partition)
725            .with_shared_work_source(shared_work_source)
726            .with_morselizer(morselizer)
727            .with_metrics(source.metrics())
728            .build()?;
729        Ok(Box::pin(cooperative(stream)))
730    }
731
732    fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> FmtResult {
733        match t {
734            DisplayFormatType::Default | DisplayFormatType::Verbose => {
735                let schema = self.projected_schema().map_err(|_| std::fmt::Error {})?;
736                let orderings =
737                    sort_pushdown::get_projected_output_ordering(self, &schema);
738
739                write!(f, "file_groups=")?;
740                FileGroupsDisplay(&self.file_groups).fmt_as(t, f)?;
741
742                if !schema.fields().is_empty() {
743                    if let Some(projection) = self.file_source.projection() {
744                        // This matches what ProjectionExec does.
745                        // TODO: can we put this into ProjectionExprs so that it's shared code?
746                        let expr: Vec<String> = projection
747                            .as_ref()
748                            .iter()
749                            .map(|proj_expr| {
750                                if let Some(column) =
751                                    proj_expr.expr.downcast_ref::<Column>()
752                                {
753                                    if column.name() == proj_expr.alias {
754                                        column.name().to_string()
755                                    } else {
756                                        format!(
757                                            "{} as {}",
758                                            proj_expr.expr, proj_expr.alias
759                                        )
760                                    }
761                                } else {
762                                    format!("{} as {}", proj_expr.expr, proj_expr.alias)
763                                }
764                            })
765                            .collect();
766                        write!(f, ", projection=[{}]", expr.join(", "))?;
767                    } else {
768                        write!(f, ", projection={}", ProjectSchemaDisplay(&schema))?;
769                    }
770                }
771
772                if let Some(limit) = self.limit {
773                    write!(f, ", limit={limit}")?;
774                }
775
776                display_orderings(f, &orderings)?;
777
778                if self.output_partitioning.is_some() {
779                    write!(f, ", output_partitioning={}", self.output_partitioning())?;
780                }
781
782                if !self.constraints.is_empty() {
783                    write!(f, ", {}", self.constraints)?;
784                }
785
786                self.fmt_file_source(t, f)
787            }
788            DisplayFormatType::TreeRender => {
789                writeln!(f, "format={}", self.file_source.file_type())?;
790                self.file_source.fmt_extra(t, f)?;
791                let num_files = self.file_groups.iter().map(|fg| fg.len()).sum::<usize>();
792                writeln!(f, "files={num_files}")?;
793                Ok(())
794            }
795        }
796    }
797
798    /// If supported by the underlying [`FileSource`], redistribute files across partitions according to their size.
799    fn repartitioned(
800        &self,
801        target_partitions: usize,
802        repartition_file_min_size: usize,
803        output_ordering: Option<LexOrdering>,
804    ) -> Result<Option<Arc<dyn DataSource>>> {
805        // When file groups define output partitioning, repartitioning files
806        // would invalidate the partition-to-file-group mapping.
807        if self.output_partitioning.is_some() {
808            return Ok(None);
809        }
810
811        let source = self.file_source.repartitioned(
812            target_partitions,
813            repartition_file_min_size,
814            output_ordering,
815            self,
816        )?;
817
818        Ok(source.map(|s| Arc::new(s) as _))
819    }
820
821    /// Returns the output partitioning for this file scan.
822    ///
823    /// When `output_partitioning` is set, this returns the declared partitioning
824    /// after applying scan projection, allowing the optimizer to skip hash
825    /// repartitioning for aggregates and joins on the partitioning columns.
826    ///
827    /// If projection or partition count validation fails, this returns
828    /// `UnknownPartitioning`.
829    ///
830    /// Tradeoffs
831    /// - Benefit: Eliminates `RepartitionExec` and `SortExec` for queries whose
832    ///   required distribution is satisfied by the scan's output partitioning.
833    /// - Cost: Files are grouped by partition values rather than split by byte
834    ///   ranges, which may reduce I/O parallelism when partition sizes are uneven.
835    ///   For simple aggregations without `ORDER BY`, this cost may outweigh the benefit.
836    ///
837    /// Follow-up Work
838    /// - Idea: Could allow byte-range splitting within partition-aware groups,
839    ///   preserving I/O parallelism while maintaining partition semantics.
840    fn output_partitioning(&self) -> Partitioning {
841        let Some(output_partitioning) = self.output_partitioning.clone() else {
842            return Partitioning::UnknownPartitioning(self.file_groups.len());
843        };
844        if output_partitioning.partition_count() != self.file_groups.len() {
845            warn!(
846                "Declared output partitioning has {} partitions, but file scan has {} file groups. Falling back to UnknownPartitioning.",
847                output_partitioning.partition_count(),
848                self.file_groups.len()
849            );
850            return Partitioning::UnknownPartitioning(self.file_groups.len());
851        }
852
853        if let Some(projection) = self.file_source.projection() {
854            let schema = self.file_source.table_schema().table_schema();
855            return match projection.projection_mapping(schema) {
856                Ok(mapping) => project_output_partitioning(
857                    &output_partitioning,
858                    &mapping,
859                    schema,
860                    self.file_groups.len(),
861                ),
862                Err(e) => {
863                    debug!(
864                        "Could not project output partitioning, falling back to UnknownPartitioning: {e}"
865                    );
866                    Partitioning::UnknownPartitioning(self.file_groups.len())
867                }
868            };
869        }
870
871        output_partitioning
872    }
873
874    /// Computes the effective equivalence properties of this file scan, taking
875    /// into account the file schema, any projections or filters applied by the
876    /// file source, and the output ordering.
877    fn eq_properties(&self) -> EquivalenceProperties {
878        let schema = self.file_source.table_schema().table_schema();
879        let mut eq_properties = EquivalenceProperties::new_with_orderings(
880            Arc::clone(schema),
881            self.validated_output_ordering(),
882        )
883        .with_constraints(self.constraints.clone());
884
885        if let Some(filter) = self.file_source.filter() {
886            // We need to remap column indexes to match the projected schema since that's what the equivalence properties deal with.
887            // Note that this will *ignore* any non-projected columns: these don't factor into ordering / equivalence.
888            match Self::add_filter_equivalence_info(&filter, &mut eq_properties, schema) {
889                Ok(()) => {}
890                Err(e) => {
891                    warn!("Failed to add filter equivalence info: {e}");
892                    #[cfg(debug_assertions)]
893                    panic!("Failed to add filter equivalence info: {e}");
894                }
895            }
896        }
897
898        if let Some(projection) = self.file_source.projection() {
899            match (
900                projection.project_schema(schema),
901                projection.projection_mapping(schema),
902            ) {
903                (Ok(output_schema), Ok(mapping)) => {
904                    eq_properties =
905                        eq_properties.project(&mapping, Arc::new(output_schema));
906                }
907                (Err(e), _) | (_, Err(e)) => {
908                    warn!("Failed to project equivalence properties: {e}");
909                    #[cfg(debug_assertions)]
910                    panic!("Failed to project equivalence properties: {e}");
911                }
912            }
913        }
914
915        eq_properties
916    }
917
918    fn scheduling_type(&self) -> SchedulingType {
919        SchedulingType::Cooperative
920    }
921
922    fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
923        if let Some(partition) = partition {
924            // Get statistics for a specific partition
925            // Note: FileGroup statistics include partition columns (computed from partition_values)
926            if let Some(file_group) = self.file_groups.get(partition)
927                && let Some(stat) = file_group.file_statistics(None)
928            {
929                // Project the statistics based on the projection
930                let output_schema = self.projected_schema()?;
931                return if let Some(projection) = self.file_source.projection() {
932                    Ok(Arc::new(
933                        projection.project_statistics(stat.clone(), &output_schema)?,
934                    ))
935                } else {
936                    Ok(Arc::new(stat.clone()))
937                };
938            }
939            // If no statistics available for this partition, return unknown
940            Ok(Arc::new(Statistics::new_unknown(
941                self.projected_schema()?.as_ref(),
942            )))
943        } else {
944            // Return aggregate statistics across all partitions
945            let statistics = self.statistics();
946            let projection = self.file_source.projection();
947            let output_schema = self.projected_schema()?;
948            if let Some(projection) = &projection {
949                Ok(Arc::new(
950                    projection.project_statistics(statistics.clone(), &output_schema)?,
951                ))
952            } else {
953                Ok(Arc::new(statistics))
954            }
955        }
956    }
957
958    fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn DataSource>> {
959        let source = FileScanConfigBuilder::from(self.clone())
960            .with_limit(limit)
961            .build();
962        Some(Arc::new(source))
963    }
964
965    fn fetch(&self) -> Option<usize> {
966        self.limit
967    }
968
969    fn metrics(&self) -> ExecutionPlanMetricsSet {
970        self.file_source.metrics().clone()
971    }
972
973    fn try_swapping_with_projection(
974        &self,
975        projection: &ProjectionExprs,
976    ) -> Result<Option<Arc<dyn DataSource>>> {
977        // Don't merge a projection into the scan if it would inline a volatile
978        // or expensive expression referenced more than once. For a volatile
979        // expression (e.g. `random()` aliased in a subquery) this would turn a
980        // single "locked-in" value into multiple independent evaluations (see
981        // #23220); for an expensive scalar function it would undo CSE and
982        // re-evaluate the expression at every reference site.
983        if let Some(inner) = self.file_source.projection()
984            && would_duplicate_costly_exprs(inner, projection)
985        {
986            return Ok(None);
987        }
988        match self.file_source.try_pushdown_projection(projection)? {
989            Some(new_source) => {
990                let mut new_file_scan_config = self.clone();
991                new_file_scan_config.file_source = new_source;
992                Ok(Some(Arc::new(new_file_scan_config) as Arc<dyn DataSource>))
993            }
994            None => Ok(None),
995        }
996    }
997
998    fn try_pushdown_filters(
999        &self,
1000        filters: Vec<Arc<dyn PhysicalExpr>>,
1001        config: &ConfigOptions,
1002    ) -> Result<FilterPushdownPropagation<Arc<dyn DataSource>>> {
1003        // Remap filter Column indices to match the table schema (file + partition columns).
1004        // This is necessary because filters refer to the output schema of this `DataSource`
1005        // (e.g., after projection pushdown has been applied) and need to be remapped to the table schema
1006        // before being passed to the file source
1007        //
1008        // For example, consider a filter `c1_c2 > 5` being pushed down. If the
1009        // `DataSource` has a projection `c1 + c2 as c1_c2`, the filter must be rewritten
1010        // to refer to the table schema `c1 + c2 > 5`
1011        let table_schema = self.file_source.table_schema().table_schema();
1012        let filters_to_remap = if let Some(projection) = self.file_source.projection() {
1013            filters
1014                .into_iter()
1015                .map(|filter| projection.unproject_expr(&filter))
1016                .collect::<Result<Vec<_>>>()?
1017        } else {
1018            filters
1019        };
1020        // Now remap column indices to match the table schema.
1021        let remapped_filters = filters_to_remap
1022            .into_iter()
1023            .map(|filter| reassign_expr_columns(filter, table_schema))
1024            .collect::<Result<Vec<_>>>()?;
1025
1026        let result = self
1027            .file_source
1028            .try_pushdown_filters(remapped_filters, config)?;
1029        match result.updated_node {
1030            Some(new_file_source) => {
1031                let mut new_file_scan_config = self.clone();
1032                new_file_scan_config.file_source = new_file_source;
1033                Ok(FilterPushdownPropagation {
1034                    filters: result.filters,
1035                    updated_node: Some(Arc::new(new_file_scan_config) as _),
1036                })
1037            }
1038            None => {
1039                // If the file source does not support filter pushdown, return the original config
1040                Ok(FilterPushdownPropagation {
1041                    filters: result.filters,
1042                    updated_node: None,
1043                })
1044            }
1045        }
1046    }
1047
1048    /// Push sort requirements into file-based data sources.
1049    ///
1050    /// # Sort Pushdown Architecture
1051    ///
1052    /// When a partition (file group) contains multiple files in wrong order,
1053    /// `validated_output_ordering()` strips the ordering and `EnforceSorting`
1054    /// inserts a `SortExec`. This optimizer fixes the file order by sorting
1055    /// files within each group by min/max statistics, enabling sort elimination.
1056    ///
1057    /// This applies to both single-partition and multi-partition plans — any
1058    /// file group with multiple files in wrong order benefits.
1059    ///
1060    /// ```text
1061    /// PushdownSort optimizer finds SortExec
1062    ///   │
1063    ///   ▼
1064    /// FileScanConfig::try_pushdown_sort()
1065    ///   │
1066    ///   ├─► FileSource returns Exact
1067    ///   │     (natural ordering satisfies request)
1068    ///   │     → rebuild_with_source: sort files by stats, verify non-overlapping
1069    ///   │     → SortExec removed, fetch (LIMIT) pushed to DataSourceExec
1070    ///   │
1071    ///   ├─► FileSource returns Inexact
1072    ///   │     (e.g. column_in_file_schema: opener will reorder RGs at runtime)
1073    ///   │     → rebuild_with_source: sort files by stats; if the post-sort
1074    ///   │       file groups are non-overlapping AND the request now validates
1075    ///   │       AND no NULLs sit in the sort columns of non-last files,
1076    ///   │       upgrade back to Exact (SortExec removed). Otherwise stays
1077    ///   │       Inexact and SortExec is kept while the scan is still
1078    ///   │       optimised via `sort_order_for_reorder` / `reverse_row_groups`.
1079    ///   │
1080    ///   └─► FileSource returns Unsupported
1081    ///         (e.g. expression sort key or partition column)
1082    ///         → try_sort_file_groups_by_statistics():
1083    ///           1. Sort files within each group by min/max statistics
1084    ///           2. Re-check: non-overlapping + ordering valid + no NULLs?
1085    ///              YES → Exact → SortExec removed
1086    ///              NO  → Inexact (files reordered, Sort stays)
1087    /// ```
1088    fn try_pushdown_sort(
1089        &self,
1090        order: &[PhysicalSortExpr],
1091    ) -> Result<SortOrderPushdownResult<Arc<dyn DataSource>>> {
1092        let pushdown_result = self
1093            .file_source
1094            .try_pushdown_sort(order, &self.eq_properties())?;
1095
1096        match pushdown_result {
1097            SortOrderPushdownResult::Exact { inner } => {
1098                let config = self.rebuild_with_source(inner, true, order)?;
1099                // rebuild_with_source keeps output_ordering only when all groups
1100                // are non-overlapping. If output_ordering was cleared, files
1101                // overlap despite within-file ordering → downgrade to Inexact.
1102                if config.output_ordering.is_empty() {
1103                    Ok(SortOrderPushdownResult::Inexact {
1104                        inner: Arc::new(config),
1105                    })
1106                } else {
1107                    Ok(SortOrderPushdownResult::Exact {
1108                        inner: Arc::new(config),
1109                    })
1110                }
1111            }
1112            SortOrderPushdownResult::Inexact { inner } => {
1113                let mut config = self.rebuild_with_source(inner, false, order)?;
1114                // `rebuild_with_source` reorders files by stats; if the
1115                // post-sort files are non-overlapping AND the request now
1116                // validates against the new file groups, `output_ordering`
1117                // is preserved and we can upgrade back to Exact. This
1118                // restores the sort-elimination behaviour that lived in
1119                // the `Unsupported` → `try_sort_file_groups_by_statistics`
1120                // path before #21956 routed `column_in_file_schema` cases
1121                // here.
1122                if config.output_ordering.is_empty() {
1123                    return Ok(SortOrderPushdownResult::Inexact {
1124                        inner: Arc::new(config),
1125                    });
1126                }
1127                // Upgrading to Exact: the post-sort file groups are
1128                // non-overlapping and each file's declared ordering
1129                // re-validates, so reading the files in their natural
1130                // (declared-sorted) order already yields the requested
1131                // ordering — exactly like the `Unsupported` → Exact path,
1132                // which reads files in natural order too.
1133                //
1134                // Drop the runtime row-group reorder hints the Inexact
1135                // source carried (`sort_order_for_reorder` /
1136                // `reverse_row_groups`) by restoring the original,
1137                // hint-free source. With the `SortExec` removed those
1138                // hints are not just redundant but unsafe: for a DESC
1139                // request the opener sorts row groups ASC-by-min and then
1140                // reverses them, which reorders two row groups within a
1141                // single file that share the same `min` incorrectly
1142                // (e.g. a file `[10,8,8,8]` whose row groups are
1143                // `[10,8]` and `[8,8]` would stream as `8,8,10,8`).
1144                // The `SortExec` used to mask this; once it is gone the
1145                // reordered stream is the final, wrong answer.
1146                config.file_source = Arc::clone(&self.file_source);
1147                Ok(SortOrderPushdownResult::Exact {
1148                    inner: Arc::new(config),
1149                })
1150            }
1151            SortOrderPushdownResult::Unsupported => {
1152                self.try_sort_file_groups_by_statistics(order)
1153            }
1154        }
1155    }
1156
1157    fn with_preserve_order(&self, preserve_order: bool) -> Option<Arc<dyn DataSource>> {
1158        if self.preserve_order == preserve_order {
1159            return Some(Arc::new(self.clone()));
1160        }
1161
1162        let new_config = FileScanConfig {
1163            preserve_order,
1164            ..self.clone()
1165        };
1166        Some(Arc::new(new_config))
1167    }
1168
1169    fn apply_expressions(
1170        &self,
1171        f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
1172    ) -> Result<TreeNodeRecursion> {
1173        // Delegate to the file source
1174        self.file_source.apply_expressions(f)
1175    }
1176
1177    /// Create any shared state that should be passed between sibling streams
1178    /// during one execution.
1179    ///
1180    /// This returns `None` when sibling streams must not share work, such as
1181    /// when file order must be preserved, the file groups define the output
1182    /// partitioning needed for the rest of the plan, or work stealing is
1183    /// disabled via
1184    /// `datafusion.execution.enable_file_stream_work_stealing`.
1185    fn create_sibling_state(
1186        &self,
1187        config: &ConfigOptions,
1188    ) -> Option<Arc<dyn Any + Send + Sync>> {
1189        if self.preserve_order
1190            || self.output_partitioning.is_some()
1191            || !config.execution.enable_file_stream_work_stealing
1192        {
1193            return None;
1194        }
1195
1196        Some(Arc::new(SharedWorkSource::from_config(self)) as Arc<dyn Any + Send + Sync>)
1197    }
1198
1199    /// Serialize this file scan by delegating to the concrete
1200    /// [`FileSource`]'s
1201    /// [`try_to_proto`](crate::file::FileSource::try_to_proto) hook, passing
1202    /// `self` as the shared spine it needs to emit the base config.
1203    #[cfg(feature = "proto")]
1204    fn try_to_proto(
1205        &self,
1206        ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
1207    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
1208        self.file_source().try_to_proto(self, ctx)
1209    }
1210}
1211
1212impl FileScanConfig {
1213    /// Returns only the output orderings that are validated against actual
1214    /// file group statistics.
1215    ///
1216    /// For example, individual files may be ordered by `col1 ASC`,
1217    /// but if we have files with these min/max statistics in a single partition / file group:
1218    ///
1219    /// - file1: min(col1) = 10, max(col1) = 20
1220    /// - file2: min(col1) = 5, max(col1) = 15
1221    ///
1222    /// Because reading file1 followed by file2 would produce out-of-order output (there is overlap
1223    /// in the ranges), we cannot retain `col1 ASC` as a valid output ordering.
1224    ///
1225    /// Similarly this would not be a valid order (non-overlapping ranges but not ordered):
1226    ///
1227    /// - file1: min(col1) = 20, max(col1) = 30
1228    /// - file2: min(col1) = 10, max(col1) = 15
1229    ///
1230    /// On the other hand if we had:
1231    ///
1232    /// - file1: min(col1) = 5, max(col1) = 15
1233    /// - file2: min(col1) = 16, max(col1) = 25
1234    ///
1235    /// Then we know that reading file1 followed by file2 will produce ordered output,
1236    /// so `col1 ASC` would be retained.
1237    ///
1238    /// Note that we are checking for ordering *within* *each* file group / partition,
1239    /// files in different partitions are read independently and do not affect each other's ordering.
1240    /// Merging of the multiple partition streams into a single ordered stream is handled
1241    /// upstream e.g. by `SortPreservingMergeExec`.
1242    fn validated_output_ordering(&self) -> Vec<LexOrdering> {
1243        let schema = self.file_source.table_schema().table_schema();
1244        sort_pushdown::validate_orderings(
1245            &self.output_ordering,
1246            schema,
1247            &self.file_groups,
1248            None,
1249        )
1250    }
1251
1252    /// Get the file schema (schema of the files without partition columns)
1253    pub fn file_schema(&self) -> &SchemaRef {
1254        self.file_source.table_schema().file_schema()
1255    }
1256
1257    /// Get the table partition columns
1258    pub fn table_partition_cols(&self) -> &Fields {
1259        self.file_source.table_schema().table_partition_cols()
1260    }
1261
1262    /// Returns the unprojected table statistics, marking them as inexact if filters are present.
1263    ///
1264    /// When filters are pushed down (including pruning predicates and bloom filters),
1265    /// we can't guarantee the statistics are exact because we don't know how many
1266    /// rows will be filtered out.
1267    pub fn statistics(&self) -> Statistics {
1268        let filter_may_change_row_count = self.file_source.filter().is_some()
1269            && self.statistics.num_rows != Precision::Exact(0);
1270        if filter_may_change_row_count {
1271            self.statistics.clone().to_inexact()
1272        } else {
1273            self.statistics.clone()
1274        }
1275    }
1276
1277    pub fn projected_schema(&self) -> Result<Arc<Schema>> {
1278        let schema = self.file_source.table_schema().table_schema();
1279        match self.file_source.projection() {
1280            Some(proj) => Ok(Arc::new(proj.project_schema(schema)?)),
1281            None => Ok(Arc::clone(schema)),
1282        }
1283    }
1284
1285    fn add_filter_equivalence_info(
1286        filter: &Arc<dyn PhysicalExpr>,
1287        eq_properties: &mut EquivalenceProperties,
1288        schema: &Schema,
1289    ) -> Result<()> {
1290        // Gather valid equality pairs from the filter expression
1291        let equal_pairs = split_conjunction(filter).into_iter().filter_map(|expr| {
1292            // Ignore any binary expressions that reference non-existent columns in the current schema
1293            // (e.g. due to unnecessary projections being removed)
1294            reassign_expr_columns(Arc::clone(expr), schema)
1295                .ok()
1296                .and_then(|expr| match expr.downcast_ref::<BinaryExpr>() {
1297                    Some(expr) if expr.op() == &Operator::Eq => {
1298                        Some((Arc::clone(expr.left()), Arc::clone(expr.right())))
1299                    }
1300                    _ => None,
1301                })
1302        });
1303
1304        for (lhs, rhs) in equal_pairs {
1305            eq_properties.add_equal_conditions(lhs, rhs)?
1306        }
1307
1308        Ok(())
1309    }
1310
1311    /// Returns whether newlines in values are supported.
1312    ///
1313    /// This method always returns `false`. The actual newlines_in_values setting
1314    /// has been moved to [`CsvSource`] and should be accessed via
1315    /// [`CsvSource::csv_options()`] instead.
1316    ///
1317    /// [`CsvSource`]: https://docs.rs/datafusion/latest/datafusion/datasource/physical_plan/struct.CsvSource.html
1318    /// [`CsvSource::csv_options()`]: https://docs.rs/datafusion/latest/datafusion/datasource/physical_plan/struct.CsvSource.html#method.csv_options
1319    #[deprecated(
1320        since = "52.0.0",
1321        note = "newlines_in_values has moved to CsvSource. Access it via CsvSource::csv_options().newlines_in_values instead. It will be removed in 58.0.0 or 6 months after 52.0.0 is released, whichever comes first."
1322    )]
1323    pub fn newlines_in_values(&self) -> bool {
1324        false
1325    }
1326
1327    #[deprecated(
1328        since = "52.0.0",
1329        note = "This method is no longer used, use eq_properties instead. It will be removed in 58.0.0 or 6 months after 52.0.0 is released, whichever comes first."
1330    )]
1331    pub fn projected_constraints(&self) -> Constraints {
1332        let props = self.eq_properties();
1333        props.constraints().clone()
1334    }
1335
1336    #[deprecated(
1337        since = "52.0.0",
1338        note = "This method is no longer used, use eq_properties instead. It will be removed in 58.0.0 or 6 months after 52.0.0 is released, whichever comes first."
1339    )]
1340    pub fn file_column_projection_indices(&self) -> Option<Vec<usize>> {
1341        #[expect(deprecated)]
1342        self.file_source.projection().as_ref().map(|p| {
1343            p.ordered_column_indices()
1344                .into_iter()
1345                .filter(|&i| i < self.file_schema().fields().len())
1346                .collect::<Vec<_>>()
1347        })
1348    }
1349
1350    /// Splits file groups into new groups based on statistics to enable efficient parallel processing.
1351    ///
1352    /// The method distributes files across a target number of partitions while ensuring
1353    /// files within each partition maintain sort order based on their min/max statistics.
1354    ///
1355    /// The algorithm works by:
1356    /// 1. Takes files sorted by minimum values
1357    /// 2. For each file:
1358    ///   - Finds eligible groups (empty or where file's min > group's last max)
1359    ///   - Selects the smallest eligible group
1360    ///   - Creates a new group if needed
1361    ///
1362    /// # Parameters
1363    /// * `table_schema`: Schema containing information about the columns
1364    /// * `file_groups`: The original file groups to split
1365    /// * `sort_order`: The lexicographical ordering to maintain within each group
1366    /// * `target_partitions`: The desired number of output partitions
1367    ///
1368    /// # Returns
1369    /// A new set of file groups, where files within each group are non-overlapping with respect to
1370    /// their min/max statistics and maintain the specified sort order.
1371    pub fn split_groups_by_statistics_with_target_partitions(
1372        table_schema: &SchemaRef,
1373        file_groups: &[FileGroup],
1374        sort_order: &LexOrdering,
1375        target_partitions: usize,
1376    ) -> Result<Vec<FileGroup>> {
1377        if target_partitions == 0 {
1378            return Err(internal_datafusion_err!(
1379                "target_partitions must be greater than 0"
1380            ));
1381        }
1382
1383        let flattened_files = file_groups
1384            .iter()
1385            .flat_map(FileGroup::iter)
1386            .collect::<Vec<_>>();
1387
1388        if flattened_files.is_empty() {
1389            return Ok(vec![]);
1390        }
1391
1392        let statistics = MinMaxStatistics::new_from_files(
1393            sort_order,
1394            table_schema,
1395            None,
1396            flattened_files.iter().copied(),
1397        )?;
1398
1399        let indices_sorted_by_min = statistics.min_values_sorted();
1400
1401        // Initialize with target_partitions empty groups
1402        let mut file_groups_indices: Vec<Vec<usize>> = vec![vec![]; target_partitions];
1403
1404        for (idx, min) in indices_sorted_by_min {
1405            if let Some((_, group)) = file_groups_indices
1406                .iter_mut()
1407                .enumerate()
1408                .filter(|(_, group)| {
1409                    group.is_empty()
1410                        || min
1411                            > statistics
1412                                .max(*group.last().expect("groups should not be empty"))
1413                })
1414                .min_by_key(|(_, group)| group.len())
1415            {
1416                group.push(idx);
1417            } else {
1418                // Create a new group if no existing group fits
1419                file_groups_indices.push(vec![idx]);
1420            }
1421        }
1422
1423        // Remove any empty groups
1424        file_groups_indices.retain(|group| !group.is_empty());
1425
1426        // Assemble indices back into groups of PartitionedFiles
1427        Ok(file_groups_indices
1428            .into_iter()
1429            .map(|file_group_indices| {
1430                FileGroup::new(
1431                    file_group_indices
1432                        .into_iter()
1433                        .map(|idx| flattened_files[idx].clone())
1434                        .collect(),
1435                )
1436            })
1437            .collect())
1438    }
1439
1440    /// Attempts to do a bin-packing on files into file groups, such that any two files
1441    /// in a file group are ordered and non-overlapping with respect to their statistics.
1442    /// It will produce the smallest number of file groups possible.
1443    pub fn split_groups_by_statistics(
1444        table_schema: &SchemaRef,
1445        file_groups: &[FileGroup],
1446        sort_order: &LexOrdering,
1447    ) -> Result<Vec<FileGroup>> {
1448        let flattened_files = file_groups
1449            .iter()
1450            .flat_map(FileGroup::iter)
1451            .collect::<Vec<_>>();
1452        // First Fit:
1453        // * Choose the first file group that a file can be placed into.
1454        // * If it fits into no existing file groups, create a new one.
1455        //
1456        // By sorting files by min values and then applying first-fit bin packing,
1457        // we can produce the smallest number of file groups such that
1458        // files within a group are in order and non-overlapping.
1459        //
1460        // Source: Applied Combinatorics (Keller and Trotter), Chapter 6.8
1461        // https://www.appliedcombinatorics.org/book/s_posets_dilworth-intord.html
1462
1463        if flattened_files.is_empty() {
1464            return Ok(vec![]);
1465        }
1466
1467        let statistics = MinMaxStatistics::new_from_files(
1468            sort_order,
1469            table_schema,
1470            None,
1471            flattened_files.iter().copied(),
1472        )
1473        .map_err(|e| {
1474            e.context("construct min/max statistics for split_groups_by_statistics")
1475        })?;
1476
1477        let indices_sorted_by_min = statistics.min_values_sorted();
1478        let mut file_groups_indices: Vec<Vec<usize>> = vec![];
1479
1480        for (idx, min) in indices_sorted_by_min {
1481            let file_group_to_insert = file_groups_indices.iter_mut().find(|group| {
1482                // If our file is non-overlapping and comes _after_ the last file,
1483                // it fits in this file group.
1484                min > statistics.max(
1485                    *group
1486                        .last()
1487                        .expect("groups should be nonempty at construction"),
1488                )
1489            });
1490            match file_group_to_insert {
1491                Some(group) => group.push(idx),
1492                None => file_groups_indices.push(vec![idx]),
1493            }
1494        }
1495
1496        // Assemble indices back into groups of PartitionedFiles
1497        Ok(file_groups_indices
1498            .into_iter()
1499            .map(|file_group_indices| {
1500                file_group_indices
1501                    .into_iter()
1502                    .map(|idx| flattened_files[idx].clone())
1503                    .collect()
1504            })
1505            .collect())
1506    }
1507
1508    /// Write the data_type based on file_source
1509    fn fmt_file_source(&self, t: DisplayFormatType, f: &mut Formatter) -> FmtResult {
1510        write!(f, ", file_type={}", self.file_source.file_type())?;
1511        self.file_source.fmt_extra(t, f)
1512    }
1513
1514    /// Returns the file_source
1515    pub fn file_source(&self) -> &Arc<dyn FileSource> {
1516        &self.file_source
1517    }
1518
1519    // Sort pushdown methods (rebuild_with_source, try_sort_file_groups_by_statistics,
1520    // sort_files_within_groups_by_statistics, any_file_has_nulls_in_sort_columns)
1521    // are in crate::sort_pushdown module.
1522}
1523
1524impl Debug for FileScanConfig {
1525    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
1526        write!(f, "FileScanConfig {{")?;
1527        write!(f, "object_store_url={:?}, ", self.object_store_url)?;
1528
1529        write!(f, "statistics={:?}, ", self.statistics())?;
1530
1531        DisplayAs::fmt_as(self, DisplayFormatType::Verbose, f)?;
1532        write!(f, "}}")
1533    }
1534}
1535
1536impl DisplayAs for FileScanConfig {
1537    fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> FmtResult {
1538        let schema = self.projected_schema().map_err(|_| std::fmt::Error {})?;
1539        let orderings = sort_pushdown::get_projected_output_ordering(self, &schema);
1540
1541        write!(f, "file_groups=")?;
1542        FileGroupsDisplay(&self.file_groups).fmt_as(t, f)?;
1543
1544        if !schema.fields().is_empty() {
1545            write!(f, ", projection={}", ProjectSchemaDisplay(&schema))?;
1546        }
1547
1548        if let Some(limit) = self.limit {
1549            write!(f, ", limit={limit}")?;
1550        }
1551
1552        display_orderings(f, &orderings)?;
1553
1554        if !self.constraints.is_empty() {
1555            write!(f, ", {}", self.constraints)?;
1556        }
1557
1558        Ok(())
1559    }
1560}
1561
1562/// Convert type to a type suitable for use as a `ListingTable`
1563/// partition column. Returns `Dictionary(UInt16, val_type)`, which is
1564/// a reasonable trade off between a reasonable number of partition
1565/// values and space efficiency.
1566///
1567/// This use this to specify types for partition columns. However
1568/// you MAY also choose not to dictionary-encode the data or to use a
1569/// different dictionary type.
1570///
1571/// Use [`wrap_partition_value_in_dict`] to wrap a [`ScalarValue`] in the same say.
1572pub fn wrap_partition_type_in_dict(val_type: DataType) -> DataType {
1573    DataType::Dictionary(Box::new(DataType::UInt16), Box::new(val_type))
1574}
1575
1576/// Convert a [`ScalarValue`] of partition columns to a type, as
1577/// described in the documentation of [`wrap_partition_type_in_dict`],
1578/// which can wrap the types.
1579pub fn wrap_partition_value_in_dict(val: ScalarValue) -> ScalarValue {
1580    ScalarValue::Dictionary(Box::new(DataType::UInt16), Box::new(val))
1581}
1582
1583#[cfg(test)]
1584mod tests {
1585    use std::collections::HashMap;
1586
1587    use super::*;
1588    use crate::source::DataSourceExec;
1589    use crate::test_util::col;
1590    use crate::{TableSchema, TableSchemaBuilder};
1591    use crate::{
1592        generate_test_files, test_util::MockSource, tests::aggr_test_schema,
1593        verify_sort_integrity,
1594    };
1595
1596    use arrow::array::{Int32Array, RecordBatch};
1597    use arrow::datatypes::Field;
1598    use datafusion_common::ColumnStatistics;
1599    use datafusion_common::stats::Precision;
1600    use datafusion_common::tree_node::TreeNodeRecursion;
1601    use datafusion_common::{Result, assert_batches_eq, internal_err};
1602    use datafusion_execution::TaskContext;
1603    use datafusion_expr::SortExpr;
1604    use datafusion_physical_expr::PhysicalExpr;
1605
1606    #[cfg(feature = "proto")]
1607    use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};
1608    use datafusion_physical_expr::create_physical_sort_expr;
1609    use datafusion_physical_expr::expressions::Literal;
1610    use datafusion_physical_expr::projection::ProjectionExpr;
1611    use datafusion_physical_expr::projection::ProjectionExprs;
1612    use datafusion_physical_plan::ExecutionPlan;
1613    use datafusion_physical_plan::execution_plan::collect;
1614    #[cfg(feature = "proto")]
1615    use datafusion_physical_plan::proto::{ExecutionPlanEncode, ExecutionPlanEncodeCtx};
1616    #[cfg(feature = "proto")]
1617    use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode};
1618    use futures::FutureExt as _;
1619    use futures::StreamExt as _;
1620    use futures::stream;
1621    use object_store::ObjectStore;
1622    use std::fmt::Debug;
1623
1624    #[derive(Clone)]
1625    struct InexactSortPushdownSource {
1626        metrics: ExecutionPlanMetricsSet,
1627        table_schema: TableSchema,
1628    }
1629
1630    impl InexactSortPushdownSource {
1631        fn new(table_schema: TableSchema) -> Self {
1632            Self {
1633                metrics: ExecutionPlanMetricsSet::new(),
1634                table_schema,
1635            }
1636        }
1637    }
1638
1639    impl FileSource for InexactSortPushdownSource {
1640        fn create_file_opener(
1641            &self,
1642            _object_store: Arc<dyn ObjectStore>,
1643            _base_config: &FileScanConfig,
1644            _partition: usize,
1645        ) -> Result<Arc<dyn crate::file_stream::FileOpener>> {
1646            unimplemented!()
1647        }
1648
1649        fn table_schema(&self) -> &TableSchema {
1650            &self.table_schema
1651        }
1652
1653        fn with_batch_size(&self, _batch_size: usize) -> Arc<dyn FileSource> {
1654            Arc::new(self.clone())
1655        }
1656
1657        fn metrics(&self) -> &ExecutionPlanMetricsSet {
1658            &self.metrics
1659        }
1660
1661        fn file_type(&self) -> &str {
1662            "mock"
1663        }
1664
1665        fn try_pushdown_sort(
1666            &self,
1667            _order: &[PhysicalSortExpr],
1668            _eq_properties: &EquivalenceProperties,
1669        ) -> Result<SortOrderPushdownResult<Arc<dyn FileSource>>> {
1670            Ok(SortOrderPushdownResult::Inexact {
1671                inner: Arc::new(self.clone()) as Arc<dyn FileSource>,
1672            })
1673        }
1674
1675        fn apply_expressions(
1676            &self,
1677            _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
1678        ) -> Result<TreeNodeRecursion> {
1679            Ok(TreeNodeRecursion::Continue)
1680        }
1681    }
1682
1683    #[cfg(feature = "proto")]
1684    #[derive(Clone)]
1685    struct ProtoHookSource {
1686        metrics: ExecutionPlanMetricsSet,
1687        table_schema: TableSchema,
1688    }
1689
1690    #[cfg(feature = "proto")]
1691    impl ProtoHookSource {
1692        fn new(table_schema: TableSchema) -> Self {
1693            Self {
1694                metrics: ExecutionPlanMetricsSet::new(),
1695                table_schema,
1696            }
1697        }
1698    }
1699
1700    #[cfg(feature = "proto")]
1701    impl FileSource for ProtoHookSource {
1702        fn create_file_opener(
1703            &self,
1704            _object_store: Arc<dyn ObjectStore>,
1705            _base_config: &FileScanConfig,
1706            _partition: usize,
1707        ) -> Result<Arc<dyn crate::file_stream::FileOpener>> {
1708            internal_err!("not needed for proto delegation test")
1709        }
1710
1711        fn table_schema(&self) -> &TableSchema {
1712            &self.table_schema
1713        }
1714
1715        fn with_batch_size(&self, _batch_size: usize) -> Arc<dyn FileSource> {
1716            Arc::new(self.clone())
1717        }
1718
1719        fn metrics(&self) -> &ExecutionPlanMetricsSet {
1720            &self.metrics
1721        }
1722
1723        fn file_type(&self) -> &str {
1724            "proto-hook-test"
1725        }
1726
1727        fn apply_expressions(
1728            &self,
1729            _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
1730        ) -> Result<TreeNodeRecursion> {
1731            Ok(TreeNodeRecursion::Continue)
1732        }
1733
1734        fn try_to_proto(
1735            &self,
1736            _base: &FileScanConfig,
1737            _ctx: &ExecutionPlanEncodeCtx<'_>,
1738        ) -> Result<Option<PhysicalPlanNode>> {
1739            Ok(Some(PhysicalPlanNode::default()))
1740        }
1741    }
1742
1743    #[cfg(feature = "proto")]
1744    struct UnusedPlanEncoder;
1745
1746    #[cfg(feature = "proto")]
1747    impl ExecutionPlanEncode for UnusedPlanEncoder {
1748        fn encode_plan(
1749            &self,
1750            _plan: &Arc<dyn ExecutionPlan>,
1751        ) -> Result<PhysicalPlanNode> {
1752            internal_err!("not needed for proto delegation test")
1753        }
1754
1755        fn encode_expr(&self, _expr: &Arc<dyn PhysicalExpr>) -> Result<PhysicalExprNode> {
1756            internal_err!("not needed for proto delegation test")
1757        }
1758
1759        fn encode_udf(&self, _udf: &ScalarUDF) -> Result<Option<Vec<u8>>> {
1760            internal_err!("not needed for proto delegation test")
1761        }
1762
1763        fn encode_udaf(&self, _udaf: &AggregateUDF) -> Result<Option<Vec<u8>>> {
1764            internal_err!("not needed for proto delegation test")
1765        }
1766
1767        fn encode_udwf(&self, _udwf: &WindowUDF) -> Result<Option<Vec<u8>>> {
1768            internal_err!("not needed for proto delegation test")
1769        }
1770    }
1771
1772    #[cfg(feature = "proto")]
1773    #[test]
1774    fn data_source_exec_delegates_proto_to_file_source() -> Result<()> {
1775        let schema = Arc::new(Schema::new(vec![Field::new(
1776            "value",
1777            DataType::Int32,
1778            false,
1779        )]));
1780        let source = Arc::new(ProtoHookSource::new(TableSchema::from(&schema)));
1781        let config =
1782            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source)
1783                .build();
1784        let exec = DataSourceExec::from_data_source(config);
1785        let encoder = UnusedPlanEncoder;
1786        let ctx = ExecutionPlanEncodeCtx::new(&encoder);
1787
1788        assert_eq!(exec.try_to_proto(&ctx)?, Some(PhysicalPlanNode::default()));
1789        Ok(())
1790    }
1791
1792    #[test]
1793    fn physical_plan_config_no_projection_tab_cols_as_field() {
1794        let file_schema = aggr_test_schema();
1795
1796        // make a table_partition_col as a field
1797        let table_partition_col =
1798            Field::new("date", wrap_partition_type_in_dict(DataType::Utf8), true)
1799                .with_metadata(HashMap::from_iter(vec![(
1800                    "key_whatever".to_owned(),
1801                    "value_whatever".to_owned(),
1802                )]));
1803
1804        let conf = config_for_projection(
1805            Arc::clone(&file_schema),
1806            None,
1807            Statistics::new_unknown(&file_schema),
1808            vec![table_partition_col.clone()],
1809        );
1810
1811        // verify the proj_schema includes the last column and exactly the same the field it is defined
1812        let proj_schema = conf.projected_schema().unwrap();
1813        assert_eq!(proj_schema.fields().len(), file_schema.fields().len() + 1);
1814        assert_eq!(
1815            *proj_schema.field(file_schema.fields().len()),
1816            table_partition_col,
1817            "partition columns are the last columns and ust have all values defined in created field"
1818        );
1819    }
1820
1821    #[test]
1822    fn test_split_groups_by_statistics() -> Result<()> {
1823        use chrono::TimeZone;
1824        use datafusion_common::DFSchema;
1825        use datafusion_expr::execution_props::ExecutionProps;
1826        use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
1827        use object_store::{ObjectMeta, path::Path};
1828
1829        struct File {
1830            name: &'static str,
1831            date: &'static str,
1832            statistics: Vec<Option<(Option<f64>, Option<f64>)>>,
1833        }
1834        impl File {
1835            fn new(
1836                name: &'static str,
1837                date: &'static str,
1838                statistics: Vec<Option<(f64, f64)>>,
1839            ) -> Self {
1840                Self::new_nullable(
1841                    name,
1842                    date,
1843                    statistics
1844                        .into_iter()
1845                        .map(|opt| opt.map(|(min, max)| (Some(min), Some(max))))
1846                        .collect(),
1847                )
1848            }
1849
1850            fn new_nullable(
1851                name: &'static str,
1852                date: &'static str,
1853                statistics: Vec<Option<(Option<f64>, Option<f64>)>>,
1854            ) -> Self {
1855                Self {
1856                    name,
1857                    date,
1858                    statistics,
1859                }
1860            }
1861        }
1862
1863        struct TestCase {
1864            name: &'static str,
1865            file_schema: Schema,
1866            files: Vec<File>,
1867            sort: Vec<SortExpr>,
1868            expected_result: Result<Vec<Vec<&'static str>>, &'static str>,
1869        }
1870
1871        use datafusion_expr::col;
1872        let cases = vec![
1873            TestCase {
1874                name: "test sort",
1875                file_schema: Schema::new(vec![Field::new(
1876                    "value".to_string(),
1877                    DataType::Float64,
1878                    false,
1879                )]),
1880                files: vec![
1881                    File::new("0", "2023-01-01", vec![Some((0.00, 0.49))]),
1882                    File::new("1", "2023-01-01", vec![Some((0.50, 1.00))]),
1883                    File::new("2", "2023-01-02", vec![Some((0.00, 1.00))]),
1884                ],
1885                sort: vec![col("value").sort(true, false)],
1886                expected_result: Ok(vec![vec!["0", "1"], vec!["2"]]),
1887            },
1888            // same input but file '2' is in the middle
1889            // test that we still order correctly
1890            TestCase {
1891                name: "test sort with files ordered differently",
1892                file_schema: Schema::new(vec![Field::new(
1893                    "value".to_string(),
1894                    DataType::Float64,
1895                    false,
1896                )]),
1897                files: vec![
1898                    File::new("0", "2023-01-01", vec![Some((0.00, 0.49))]),
1899                    File::new("2", "2023-01-02", vec![Some((0.00, 1.00))]),
1900                    File::new("1", "2023-01-01", vec![Some((0.50, 1.00))]),
1901                ],
1902                sort: vec![col("value").sort(true, false)],
1903                expected_result: Ok(vec![vec!["0", "1"], vec!["2"]]),
1904            },
1905            TestCase {
1906                name: "reverse sort",
1907                file_schema: Schema::new(vec![Field::new(
1908                    "value".to_string(),
1909                    DataType::Float64,
1910                    false,
1911                )]),
1912                files: vec![
1913                    File::new("0", "2023-01-01", vec![Some((0.00, 0.49))]),
1914                    File::new("1", "2023-01-01", vec![Some((0.50, 1.00))]),
1915                    File::new("2", "2023-01-02", vec![Some((0.00, 1.00))]),
1916                ],
1917                sort: vec![col("value").sort(false, true)],
1918                expected_result: Ok(vec![vec!["1", "0"], vec!["2"]]),
1919            },
1920            TestCase {
1921                name: "nullable sort columns, nulls last",
1922                file_schema: Schema::new(vec![Field::new(
1923                    "value".to_string(),
1924                    DataType::Float64,
1925                    true,
1926                )]),
1927                files: vec![
1928                    File::new_nullable(
1929                        "0",
1930                        "2023-01-01",
1931                        vec![Some((Some(0.00), Some(0.49)))],
1932                    ),
1933                    File::new_nullable("1", "2023-01-01", vec![Some((Some(0.50), None))]),
1934                    File::new_nullable("2", "2023-01-02", vec![Some((Some(0.00), None))]),
1935                ],
1936                sort: vec![col("value").sort(true, false)],
1937                expected_result: Ok(vec![vec!["0", "1"], vec!["2"]]),
1938            },
1939            TestCase {
1940                name: "nullable sort columns, nulls first",
1941                file_schema: Schema::new(vec![Field::new(
1942                    "value".to_string(),
1943                    DataType::Float64,
1944                    true,
1945                )]),
1946                files: vec![
1947                    File::new_nullable("0", "2023-01-01", vec![Some((None, Some(0.49)))]),
1948                    File::new_nullable(
1949                        "1",
1950                        "2023-01-01",
1951                        vec![Some((Some(0.50), Some(1.00)))],
1952                    ),
1953                    File::new_nullable("2", "2023-01-02", vec![Some((None, Some(1.00)))]),
1954                ],
1955                sort: vec![col("value").sort(true, true)],
1956                expected_result: Ok(vec![vec!["0", "1"], vec!["2"]]),
1957            },
1958            TestCase {
1959                name: "all three non-overlapping",
1960                file_schema: Schema::new(vec![Field::new(
1961                    "value".to_string(),
1962                    DataType::Float64,
1963                    false,
1964                )]),
1965                files: vec![
1966                    File::new("0", "2023-01-01", vec![Some((0.00, 0.49))]),
1967                    File::new("1", "2023-01-01", vec![Some((0.50, 0.99))]),
1968                    File::new("2", "2023-01-02", vec![Some((1.00, 1.49))]),
1969                ],
1970                sort: vec![col("value").sort(true, false)],
1971                expected_result: Ok(vec![vec!["0", "1", "2"]]),
1972            },
1973            TestCase {
1974                name: "all three overlapping",
1975                file_schema: Schema::new(vec![Field::new(
1976                    "value".to_string(),
1977                    DataType::Float64,
1978                    false,
1979                )]),
1980                files: vec![
1981                    File::new("0", "2023-01-01", vec![Some((0.00, 0.49))]),
1982                    File::new("1", "2023-01-01", vec![Some((0.00, 0.49))]),
1983                    File::new("2", "2023-01-02", vec![Some((0.00, 0.49))]),
1984                ],
1985                sort: vec![col("value").sort(true, false)],
1986                expected_result: Ok(vec![vec!["0"], vec!["1"], vec!["2"]]),
1987            },
1988            TestCase {
1989                name: "empty input",
1990                file_schema: Schema::new(vec![Field::new(
1991                    "value".to_string(),
1992                    DataType::Float64,
1993                    false,
1994                )]),
1995                files: vec![],
1996                sort: vec![col("value").sort(true, false)],
1997                expected_result: Ok(vec![]),
1998            },
1999            TestCase {
2000                name: "one file missing statistics",
2001                file_schema: Schema::new(vec![Field::new(
2002                    "value".to_string(),
2003                    DataType::Float64,
2004                    false,
2005                )]),
2006                files: vec![
2007                    File::new("0", "2023-01-01", vec![Some((0.00, 0.49))]),
2008                    File::new("1", "2023-01-01", vec![Some((0.00, 0.49))]),
2009                    File::new("2", "2023-01-02", vec![None]),
2010                ],
2011                sort: vec![col("value").sort(true, false)],
2012                expected_result: Err(
2013                    "construct min/max statistics for split_groups_by_statistics\ncaused by\ncollect min/max values\ncaused by\nget min/max for column: 'value'\ncaused by\nError during planning: statistics not found",
2014                ),
2015            },
2016        ];
2017
2018        for case in cases {
2019            let table_schema = Arc::new(Schema::new(
2020                case.file_schema
2021                    .fields()
2022                    .clone()
2023                    .into_iter()
2024                    .cloned()
2025                    .chain(Some(Arc::new(Field::new(
2026                        "date".to_string(),
2027                        DataType::Utf8,
2028                        false,
2029                    ))))
2030                    .collect::<Vec<_>>(),
2031            ));
2032            let Some(sort_order) = LexOrdering::new(
2033                case.sort
2034                    .into_iter()
2035                    .map(|expr| {
2036                        create_physical_sort_expr(
2037                            &expr,
2038                            &DFSchema::try_from(Arc::clone(&table_schema))?,
2039                            &ExecutionProps::default(),
2040                            &PhysicalPlanningContext::default(),
2041                        )
2042                    })
2043                    .collect::<Result<Vec<_>>>()?,
2044            ) else {
2045                return internal_err!("This test should always use an ordering");
2046            };
2047
2048            let partitioned_files = FileGroup::new(
2049                case.files.into_iter().map(From::from).collect::<Vec<_>>(),
2050            );
2051            let result = FileScanConfig::split_groups_by_statistics(
2052                &table_schema,
2053                std::slice::from_ref(&partitioned_files),
2054                &sort_order,
2055            );
2056            let results_by_name = result
2057                .as_ref()
2058                .map(|file_groups| {
2059                    file_groups
2060                        .iter()
2061                        .map(|file_group| {
2062                            file_group
2063                                .iter()
2064                                .map(|file| {
2065                                    partitioned_files
2066                                        .iter()
2067                                        .find_map(|f| {
2068                                            if f.object_meta == file.object_meta {
2069                                                Some(
2070                                                    f.object_meta
2071                                                        .location
2072                                                        .as_ref()
2073                                                        .rsplit('/')
2074                                                        .next()
2075                                                        .unwrap()
2076                                                        .trim_end_matches(".parquet"),
2077                                                )
2078                                            } else {
2079                                                None
2080                                            }
2081                                        })
2082                                        .unwrap()
2083                                })
2084                                .collect::<Vec<_>>()
2085                        })
2086                        .collect::<Vec<_>>()
2087                })
2088                .map_err(|e| e.strip_backtrace().leak() as &'static str);
2089
2090            assert_eq!(results_by_name, case.expected_result, "{}", case.name);
2091        }
2092
2093        return Ok(());
2094
2095        impl From<File> for PartitionedFile {
2096            fn from(file: File) -> Self {
2097                let object_meta = ObjectMeta {
2098                    location: Path::from(format!(
2099                        "data/date={}/{}.parquet",
2100                        file.date, file.name
2101                    )),
2102                    last_modified: chrono::Utc.timestamp_nanos(0),
2103                    size: 0,
2104                    e_tag: None,
2105                    version: None,
2106                };
2107                let statistics = Arc::new(Statistics {
2108                    num_rows: Precision::Absent,
2109                    total_byte_size: Precision::Absent,
2110                    column_statistics: file
2111                        .statistics
2112                        .into_iter()
2113                        .map(|stats| {
2114                            stats
2115                                .map(|(min, max)| ColumnStatistics {
2116                                    min_value: Precision::Exact(ScalarValue::Float64(
2117                                        min,
2118                                    )),
2119                                    max_value: Precision::Exact(ScalarValue::Float64(
2120                                        max,
2121                                    )),
2122                                    ..Default::default()
2123                                })
2124                                .unwrap_or_default()
2125                        })
2126                        .collect::<Vec<_>>(),
2127                });
2128                PartitionedFile::new_from_meta(object_meta)
2129                    .with_partition_values(vec![ScalarValue::from(file.date)])
2130                    .with_statistics(statistics)
2131            }
2132        }
2133    }
2134
2135    // sets default for configs that play no role in projections
2136    fn config_for_projection(
2137        file_schema: SchemaRef,
2138        projection: Option<Vec<usize>>,
2139        statistics: Statistics,
2140        table_partition_cols: Vec<Field>,
2141    ) -> FileScanConfig {
2142        let table_schema = TableSchema::builder(file_schema)
2143            .with_table_partition_cols(
2144                table_partition_cols
2145                    .into_iter()
2146                    .map(Arc::new)
2147                    .collect::<Fields>(),
2148            )
2149            .build();
2150        FileScanConfigBuilder::new(
2151            ObjectStoreUrl::parse("test:///").unwrap(),
2152            Arc::new(MockSource::new(table_schema.clone())),
2153        )
2154        .with_projection_indices(projection)
2155        .unwrap()
2156        .with_statistics(statistics)
2157        .build()
2158    }
2159
2160    #[test]
2161    fn test_file_scan_config_builder() {
2162        let file_schema = aggr_test_schema();
2163        let object_store_url = ObjectStoreUrl::parse("test:///").unwrap();
2164
2165        let table_schema = TableSchemaBuilder::from(&file_schema)
2166            .with_table_partition_cols(vec![Arc::new(Field::new(
2167                "date",
2168                wrap_partition_type_in_dict(DataType::Utf8),
2169                false,
2170            ))])
2171            .build();
2172
2173        let file_source: Arc<dyn FileSource> =
2174            Arc::new(MockSource::new(table_schema.clone()));
2175
2176        // Create a builder with required parameters
2177        let builder = FileScanConfigBuilder::new(
2178            object_store_url.clone(),
2179            Arc::clone(&file_source),
2180        );
2181
2182        // Build with various configurations
2183        let config = builder
2184            .with_limit(Some(1000))
2185            .with_projection_indices(Some(vec![0, 1]))
2186            .unwrap()
2187            .with_statistics(Statistics::new_unknown(&file_schema))
2188            .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new(
2189                "test.parquet".to_string(),
2190                1024,
2191            )])])
2192            .with_output_ordering(vec![
2193                [PhysicalSortExpr::new_default(Arc::new(Column::new(
2194                    "date", 0,
2195                )))]
2196                .into(),
2197            ])
2198            .with_file_compression_type(FileCompressionType::UNCOMPRESSED)
2199            .build();
2200
2201        // Verify the built config has all the expected values
2202        assert_eq!(config.object_store_url, object_store_url);
2203        assert_eq!(*config.file_schema(), file_schema);
2204        assert_eq!(config.limit, Some(1000));
2205        assert_eq!(
2206            config
2207                .file_source
2208                .projection()
2209                .as_ref()
2210                .map(|p| p.column_indices()),
2211            Some(vec![0, 1])
2212        );
2213        assert_eq!(config.table_partition_cols().len(), 1);
2214        assert_eq!(config.table_partition_cols()[0].name(), "date");
2215        assert_eq!(config.file_groups.len(), 1);
2216        assert_eq!(config.file_groups[0].len(), 1);
2217        assert_eq!(
2218            config.file_groups[0][0].object_meta.location.as_ref(),
2219            "test.parquet"
2220        );
2221        assert_eq!(
2222            config.file_compression_type,
2223            FileCompressionType::UNCOMPRESSED
2224        );
2225        assert_eq!(config.output_ordering.len(), 1);
2226    }
2227
2228    #[test]
2229    fn equivalence_properties_after_schema_change() {
2230        let file_schema = aggr_test_schema();
2231        let object_store_url = ObjectStoreUrl::parse("test:///").unwrap();
2232
2233        let table_schema = TableSchema::from(&file_schema);
2234
2235        // Create a file source with a filter
2236        let file_source: Arc<dyn FileSource> = Arc::new(
2237            MockSource::new(table_schema.clone()).with_filter(Arc::new(BinaryExpr::new(
2238                col("c2", &file_schema).unwrap(),
2239                Operator::Eq,
2240                Arc::new(Literal::new(ScalarValue::Int32(Some(10)))),
2241            ))),
2242        );
2243
2244        let config = FileScanConfigBuilder::new(
2245            object_store_url.clone(),
2246            Arc::clone(&file_source),
2247        )
2248        .with_projection_indices(Some(vec![0, 1, 2]))
2249        .unwrap()
2250        .build();
2251
2252        // Simulate projection being updated. Since the filter has already been pushed down,
2253        // the new projection won't include the filtered column.
2254        let exprs = ProjectionExprs::new(vec![ProjectionExpr::new(
2255            col("c1", &file_schema).unwrap(),
2256            "c1",
2257        )]);
2258        let data_source = config
2259            .try_swapping_with_projection(&exprs)
2260            .unwrap()
2261            .unwrap();
2262
2263        // Gather the equivalence properties from the new data source. There should
2264        // be no equivalence class for column c2 since it was removed by the projection.
2265        let eq_properties = data_source.eq_properties();
2266        let eq_group = eq_properties.eq_group();
2267
2268        for class in eq_group.iter() {
2269            for expr in class.iter() {
2270                if let Some(col) = expr.downcast_ref::<Column>() {
2271                    assert_ne!(
2272                        col.name(),
2273                        "c2",
2274                        "c2 should not be present in any equivalence class"
2275                    );
2276                }
2277            }
2278        }
2279    }
2280
2281    #[test]
2282    fn test_file_scan_config_builder_defaults() {
2283        let file_schema = aggr_test_schema();
2284        let object_store_url = ObjectStoreUrl::parse("test:///").unwrap();
2285
2286        let table_schema = TableSchema::from(&file_schema);
2287
2288        let file_source: Arc<dyn FileSource> =
2289            Arc::new(MockSource::new(table_schema.clone()));
2290
2291        // Create a builder with only required parameters and build without any additional configurations
2292        let config = FileScanConfigBuilder::new(
2293            object_store_url.clone(),
2294            Arc::clone(&file_source),
2295        )
2296        .build();
2297
2298        // Verify default values
2299        assert_eq!(config.object_store_url, object_store_url);
2300        assert_eq!(*config.file_schema(), file_schema);
2301        assert_eq!(config.limit, None);
2302        // When no projection is specified, the file source should have an unprojected projection
2303        // (i.e., all columns)
2304        let expected_projection: Vec<usize> = (0..file_schema.fields().len()).collect();
2305        assert_eq!(
2306            config
2307                .file_source
2308                .projection()
2309                .as_ref()
2310                .map(|p| p.column_indices()),
2311            Some(expected_projection)
2312        );
2313        assert!(config.table_partition_cols().is_empty());
2314        assert!(config.file_groups.is_empty());
2315        assert_eq!(
2316            config.file_compression_type,
2317            FileCompressionType::UNCOMPRESSED
2318        );
2319        assert!(config.output_ordering.is_empty());
2320        assert!(config.constraints.is_empty());
2321
2322        // Verify statistics are set to unknown
2323        assert_eq!(config.statistics().num_rows, Precision::Absent);
2324        assert_eq!(config.statistics().total_byte_size, Precision::Absent);
2325        assert_eq!(
2326            config.statistics().column_statistics.len(),
2327            file_schema.fields().len()
2328        );
2329        for stat in config.statistics().column_statistics {
2330            assert_eq!(stat.distinct_count, Precision::Absent);
2331            assert_eq!(stat.min_value, Precision::Absent);
2332            assert_eq!(stat.max_value, Precision::Absent);
2333            assert_eq!(stat.null_count, Precision::Absent);
2334        }
2335    }
2336
2337    #[test]
2338    fn test_file_scan_config_builder_new_from() {
2339        let schema = aggr_test_schema();
2340        let object_store_url = ObjectStoreUrl::parse("test:///").unwrap();
2341        let partition_cols = vec![Field::new(
2342            "date",
2343            wrap_partition_type_in_dict(DataType::Utf8),
2344            false,
2345        )];
2346        let file = PartitionedFile::new("test_file.parquet", 100);
2347
2348        let table_schema = TableSchemaBuilder::from(&schema)
2349            .with_table_partition_cols(
2350                partition_cols
2351                    .iter()
2352                    .map(|f| Arc::new(f.clone()))
2353                    .collect::<Fields>(),
2354            )
2355            .build();
2356
2357        let file_source: Arc<dyn FileSource> =
2358            Arc::new(MockSource::new(table_schema.clone()));
2359
2360        // Create a config with non-default values
2361        let original_config = FileScanConfigBuilder::new(
2362            object_store_url.clone(),
2363            Arc::clone(&file_source),
2364        )
2365        .with_projection_indices(Some(vec![0, 2]))
2366        .unwrap()
2367        .with_limit(Some(10))
2368        .with_file(file.clone())
2369        .with_constraints(Constraints::default())
2370        .build();
2371
2372        // Create a new builder from the config
2373        let new_builder = FileScanConfigBuilder::from(original_config);
2374
2375        // Build a new config from this builder
2376        let new_config = new_builder.build();
2377
2378        // Verify properties match
2379        let partition_cols = partition_cols.into_iter().map(Arc::new).collect::<Vec<_>>();
2380        assert_eq!(new_config.object_store_url, object_store_url);
2381        assert_eq!(*new_config.file_schema(), schema);
2382        assert_eq!(
2383            new_config
2384                .file_source
2385                .projection()
2386                .as_ref()
2387                .map(|p| p.column_indices()),
2388            Some(vec![0, 2])
2389        );
2390        assert_eq!(new_config.limit, Some(10));
2391        assert_eq!(
2392            *new_config.table_partition_cols(),
2393            Fields::from(partition_cols)
2394        );
2395        assert_eq!(new_config.file_groups.len(), 1);
2396        assert_eq!(new_config.file_groups[0].len(), 1);
2397        assert_eq!(
2398            new_config.file_groups[0][0].object_meta.location.as_ref(),
2399            "test_file.parquet"
2400        );
2401        assert_eq!(new_config.constraints, Constraints::default());
2402    }
2403
2404    #[test]
2405    fn test_split_groups_by_statistics_with_target_partitions() -> Result<()> {
2406        use datafusion_common::DFSchema;
2407        use datafusion_expr::{
2408            col, execution_props::ExecutionProps,
2409            physical_planning_context::PhysicalPlanningContext,
2410        };
2411
2412        let schema = Arc::new(Schema::new(vec![Field::new(
2413            "value",
2414            DataType::Float64,
2415            false,
2416        )]));
2417
2418        // Setup sort expression
2419        let exec_props = ExecutionProps::new();
2420        let df_schema = DFSchema::try_from_qualified_schema("test", schema.as_ref())?;
2421        let sort_expr = [col("value").sort(true, false)];
2422        let sort_ordering = sort_expr
2423            .map(|expr| {
2424                create_physical_sort_expr(
2425                    &expr,
2426                    &df_schema,
2427                    &exec_props,
2428                    &PhysicalPlanningContext::default(),
2429                )
2430                .unwrap()
2431            })
2432            .into();
2433
2434        // Test case parameters
2435        struct TestCase {
2436            name: String,
2437            file_count: usize,
2438            overlap_factor: f64,
2439            target_partitions: usize,
2440            expected_partition_count: usize,
2441        }
2442
2443        let test_cases = vec![
2444            // Basic cases
2445            TestCase {
2446                name: "no_overlap_10_files_4_partitions".to_string(),
2447                file_count: 10,
2448                overlap_factor: 0.0,
2449                target_partitions: 4,
2450                expected_partition_count: 4,
2451            },
2452            TestCase {
2453                name: "medium_overlap_20_files_5_partitions".to_string(),
2454                file_count: 20,
2455                overlap_factor: 0.5,
2456                target_partitions: 5,
2457                expected_partition_count: 5,
2458            },
2459            TestCase {
2460                name: "high_overlap_30_files_3_partitions".to_string(),
2461                file_count: 30,
2462                overlap_factor: 0.8,
2463                target_partitions: 3,
2464                expected_partition_count: 7,
2465            },
2466            // Edge cases
2467            TestCase {
2468                name: "fewer_files_than_partitions".to_string(),
2469                file_count: 3,
2470                overlap_factor: 0.0,
2471                target_partitions: 10,
2472                expected_partition_count: 3, // Should only create as many partitions as files
2473            },
2474            TestCase {
2475                name: "single_file".to_string(),
2476                file_count: 1,
2477                overlap_factor: 0.0,
2478                target_partitions: 5,
2479                expected_partition_count: 1, // Should create only one partition
2480            },
2481            TestCase {
2482                name: "empty_files".to_string(),
2483                file_count: 0,
2484                overlap_factor: 0.0,
2485                target_partitions: 3,
2486                expected_partition_count: 0, // Empty result for empty input
2487            },
2488        ];
2489
2490        for case in test_cases {
2491            println!("Running test case: {}", case.name);
2492
2493            // Generate files using bench utility function
2494            let file_groups = generate_test_files(case.file_count, case.overlap_factor);
2495
2496            // Call the function under test
2497            let result =
2498                FileScanConfig::split_groups_by_statistics_with_target_partitions(
2499                    &schema,
2500                    &file_groups,
2501                    &sort_ordering,
2502                    case.target_partitions,
2503                )?;
2504
2505            // Verify results
2506            println!(
2507                "Created {} partitions (target was {})",
2508                result.len(),
2509                case.target_partitions
2510            );
2511
2512            // Check partition count
2513            assert_eq!(
2514                result.len(),
2515                case.expected_partition_count,
2516                "Case '{}': Unexpected partition count",
2517                case.name
2518            );
2519
2520            // Verify sort integrity
2521            assert!(
2522                verify_sort_integrity(&result),
2523                "Case '{}': Files within partitions are not properly ordered",
2524                case.name
2525            );
2526
2527            // Distribution check for partitions
2528            if case.file_count > 1 && case.expected_partition_count > 1 {
2529                let group_sizes: Vec<usize> = result.iter().map(FileGroup::len).collect();
2530                let max_size = *group_sizes.iter().max().unwrap();
2531                let min_size = *group_sizes.iter().min().unwrap();
2532
2533                // Check partition balancing - difference shouldn't be extreme
2534                let avg_files_per_partition =
2535                    case.file_count as f64 / case.expected_partition_count as f64;
2536                assert!(
2537                    (max_size as f64) < 2.0 * avg_files_per_partition,
2538                    "Case '{}': Unbalanced distribution. Max partition size {} exceeds twice the average {}",
2539                    case.name,
2540                    max_size,
2541                    avg_files_per_partition
2542                );
2543
2544                println!("Distribution - min files: {min_size}, max files: {max_size}");
2545            }
2546        }
2547
2548        // Test error case: zero target partitions
2549        let empty_groups: Vec<FileGroup> = vec![];
2550        let err = FileScanConfig::split_groups_by_statistics_with_target_partitions(
2551            &schema,
2552            &empty_groups,
2553            &sort_ordering,
2554            0,
2555        )
2556        .unwrap_err();
2557
2558        assert!(
2559            err.to_string()
2560                .contains("target_partitions must be greater than 0"),
2561            "Expected error for zero target partitions"
2562        );
2563
2564        Ok(())
2565    }
2566
2567    #[test]
2568    fn test_partition_statistics_projection() {
2569        // This test verifies that partition_statistics applies projection correctly.
2570        // The old implementation had a bug where it returned file group statistics
2571        // without applying the projection, returning all column statistics instead
2572        // of just the projected ones.
2573
2574        use crate::source::DataSourceExec;
2575        use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext};
2576
2577        let schema = Arc::new(Schema::new(vec![
2578            Field::new("col0", DataType::Int32, false),
2579            Field::new("col1", DataType::Int32, false),
2580            Field::new("col2", DataType::Int32, false),
2581            Field::new("col3", DataType::Int32, false),
2582        ]));
2583
2584        // Create statistics for all 4 columns
2585        let file_group_stats = Statistics {
2586            num_rows: Precision::Exact(100),
2587            total_byte_size: Precision::Exact(1024),
2588            column_statistics: vec![
2589                ColumnStatistics {
2590                    null_count: Precision::Exact(0),
2591                    ..ColumnStatistics::new_unknown()
2592                },
2593                ColumnStatistics {
2594                    null_count: Precision::Exact(5),
2595                    ..ColumnStatistics::new_unknown()
2596                },
2597                ColumnStatistics {
2598                    null_count: Precision::Exact(10),
2599                    ..ColumnStatistics::new_unknown()
2600                },
2601                ColumnStatistics {
2602                    null_count: Precision::Exact(15),
2603                    ..ColumnStatistics::new_unknown()
2604                },
2605            ],
2606        };
2607
2608        // Create a file group with statistics
2609        let file_group = FileGroup::new(vec![PartitionedFile::new("test.parquet", 1024)])
2610            .with_statistics(Arc::new(file_group_stats));
2611
2612        let table_schema = TableSchema::from(&schema);
2613
2614        // Create a FileScanConfig with projection: only keep columns 0 and 2
2615        let config = FileScanConfigBuilder::new(
2616            ObjectStoreUrl::parse("test:///").unwrap(),
2617            Arc::new(MockSource::new(table_schema.clone())),
2618        )
2619        .with_projection_indices(Some(vec![0, 2]))
2620        .unwrap() // Only project columns 0 and 2
2621        .with_file_groups(vec![file_group])
2622        .build();
2623
2624        // Create a DataSourceExec from the config
2625        let exec = DataSourceExec::from_data_source(config);
2626
2627        // Get statistics for partition 0
2628        let partition_stats = StatisticsContext::new()
2629            .compute(
2630                exec.as_ref(),
2631                &StatisticsArgs::new().with_partition(Some(0)),
2632            )
2633            .unwrap();
2634
2635        // Verify that only 2 columns are in the statistics (the projected ones)
2636        assert_eq!(
2637            partition_stats.column_statistics.len(),
2638            2,
2639            "Expected 2 column statistics (projected), but got {}",
2640            partition_stats.column_statistics.len()
2641        );
2642
2643        // Verify the column statistics are for columns 0 and 2
2644        assert_eq!(
2645            partition_stats.column_statistics[0].null_count,
2646            Precision::Exact(0),
2647            "First projected column should be col0 with 0 nulls"
2648        );
2649        assert_eq!(
2650            partition_stats.column_statistics[1].null_count,
2651            Precision::Exact(10),
2652            "Second projected column should be col2 with 10 nulls"
2653        );
2654
2655        // Verify row count and byte size
2656        assert_eq!(partition_stats.num_rows, Precision::Exact(100));
2657        assert_eq!(partition_stats.total_byte_size, Precision::Exact(800));
2658    }
2659
2660    #[test]
2661    fn test_statistics_with_filter() {
2662        assert_num_rows_with_filter(Precision::Absent, Precision::Absent);
2663        assert_num_rows_with_filter(Precision::Exact(100), Precision::Inexact(100));
2664        assert_num_rows_with_filter(Precision::Inexact(100), Precision::Inexact(100));
2665        assert_num_rows_with_filter(Precision::Exact(0), Precision::Exact(0));
2666
2667        /// Creates a [`FileScanConfig`] with a filter and calls [`FileScanConfig::statistics`].
2668        /// Then the function checks the output num_rows stats, given the input num_rows stats.
2669        fn assert_num_rows_with_filter(
2670            input_num_rows: Precision<usize>,
2671            expected_num_rows: Precision<usize>,
2672        ) {
2673            let schema = Arc::new(Schema::new(vec![Field::new(
2674                "col0",
2675                DataType::Int32,
2676                false,
2677            )]));
2678
2679            let stats =
2680                Statistics::new_unknown(schema.as_ref()).with_num_rows(input_num_rows);
2681            let file_group =
2682                FileGroup::new(vec![PartitionedFile::new("test.parquet", 1024)]);
2683
2684            let table_schema = TableSchema::from(&schema);
2685            let config = FileScanConfigBuilder::new(
2686                ObjectStoreUrl::parse("test:///").unwrap(),
2687                Arc::new(MockSource::new(table_schema.clone()).with_filter(Arc::new(
2688                    Literal::new(ScalarValue::Boolean(Some(true))),
2689                ))),
2690            )
2691            .with_file_groups(vec![file_group])
2692            .with_statistics(stats)
2693            .build();
2694
2695            assert_eq!(config.statistics().num_rows, expected_num_rows,);
2696        }
2697    }
2698
2699    /// Regression test for reusing a `DataSourceExec` after its execution-local
2700    /// shared work queue has been drained.
2701    ///
2702    /// This test uses a single file group with two files so the scan creates a
2703    /// shared unopened-file queue. Executing after `reset_state` must recreate
2704    /// the shared queue and return the same rows again.
2705    #[tokio::test]
2706    async fn reset_state_recreates_shared_work_source() -> Result<()> {
2707        let schema = Arc::new(Schema::new(vec![Field::new(
2708            "value",
2709            DataType::Int32,
2710            false,
2711        )]));
2712        let file_source = Arc::new(
2713            MockSource::new(Arc::clone(&schema))
2714                .with_file_opener(Arc::new(ResetStateTestFileOpener { schema })),
2715        );
2716
2717        let config =
2718            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
2719                .with_file_group(FileGroup::new(vec![
2720                    PartitionedFile::new("file1.parquet", 100),
2721                    PartitionedFile::new("file2.parquet", 100),
2722                ]))
2723                .build();
2724
2725        let exec: Arc<dyn ExecutionPlan> = DataSourceExec::from_data_source(config);
2726        let task_ctx = Arc::new(TaskContext::default());
2727
2728        // Running the same scan after resetting the state, should
2729        // produce the same answer.
2730        let first_run = collect(Arc::clone(&exec), Arc::clone(&task_ctx)).await?;
2731        let reset_exec = exec.reset_state()?;
2732        let second_run = collect(reset_exec, task_ctx).await?;
2733
2734        let expected = [
2735            "+-------+",
2736            "| value |",
2737            "+-------+",
2738            "| 1     |",
2739            "| 2     |",
2740            "+-------+",
2741        ];
2742        assert_batches_eq!(expected, &first_run);
2743        assert_batches_eq!(expected, &second_run);
2744
2745        Ok(())
2746    }
2747
2748    /// Test-only `FileOpener` that turns file names like `file1.parquet` into a
2749    /// single-batch stream containing that numeric value
2750    #[derive(Debug)]
2751    struct ResetStateTestFileOpener {
2752        schema: SchemaRef,
2753    }
2754
2755    impl crate::file_stream::FileOpener for ResetStateTestFileOpener {
2756        fn open(
2757            &self,
2758            file: PartitionedFile,
2759        ) -> Result<crate::file_stream::FileOpenFuture> {
2760            let value = file
2761                .object_meta
2762                .location
2763                .as_ref()
2764                .trim_start_matches("file")
2765                .trim_end_matches(".parquet")
2766                .parse::<i32>()
2767                .expect("invalid test file name");
2768            let schema = Arc::clone(&self.schema);
2769            Ok(async move {
2770                let batch = RecordBatch::try_new(
2771                    schema,
2772                    vec![Arc::new(Int32Array::from(vec![value]))],
2773                )
2774                .expect("test batch should be valid");
2775                Ok(stream::iter(vec![Ok(batch)]).boxed())
2776            }
2777            .boxed())
2778        }
2779    }
2780
2781    #[test]
2782    fn test_output_partitioning_not_partitioned_by_file_group() {
2783        let file_schema = aggr_test_schema();
2784        let partition_col =
2785            Field::new("date", wrap_partition_type_in_dict(DataType::Utf8), false);
2786
2787        let config = config_for_projection(
2788            Arc::clone(&file_schema),
2789            None,
2790            Statistics::new_unknown(&file_schema),
2791            vec![partition_col],
2792        );
2793
2794        // output_partitioning defaults to None
2795        let partitioning = config.output_partitioning();
2796        assert!(matches!(partitioning, Partitioning::UnknownPartitioning(_)));
2797    }
2798
2799    #[test]
2800    fn test_declared_output_partitioning_projects_with_scan() {
2801        let file_schema = aggr_test_schema();
2802        let output_partitioning =
2803            Partitioning::Hash(vec![Arc::new(Column::new("c2", 1))], 4);
2804
2805        let mut config = config_for_projection(
2806            Arc::clone(&file_schema),
2807            Some(vec![1, 2]),
2808            Statistics::new_unknown(&file_schema),
2809            vec![],
2810        );
2811        config.file_groups = vec![
2812            FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]),
2813            FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]),
2814            FileGroup::new(vec![PartitionedFile::new("f3.parquet".to_string(), 1024)]),
2815            FileGroup::new(vec![PartitionedFile::new("f4.parquet".to_string(), 1024)]),
2816        ];
2817        config.output_partitioning = Some(output_partitioning);
2818
2819        match config.output_partitioning() {
2820            Partitioning::Hash(exprs, num_partitions) => {
2821                assert_eq!(num_partitions, 4);
2822                assert_eq!(exprs.len(), 1);
2823                let column = exprs[0].downcast_ref::<Column>().unwrap();
2824                assert_eq!(column.name(), "c2");
2825                assert_eq!(column.index(), 0);
2826            }
2827            _ => panic!("Expected Hash partitioning"),
2828        }
2829
2830        let mut config = config_for_projection(
2831            Arc::clone(&file_schema),
2832            Some(vec![2]),
2833            Statistics::new_unknown(&file_schema),
2834            vec![],
2835        );
2836        config.file_groups = vec![
2837            FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]),
2838            FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]),
2839            FileGroup::new(vec![PartitionedFile::new("f3.parquet".to_string(), 1024)]),
2840            FileGroup::new(vec![PartitionedFile::new("f4.parquet".to_string(), 1024)]),
2841        ];
2842        config.output_partitioning =
2843            Some(Partitioning::Hash(vec![Arc::new(Column::new("c2", 1))], 4));
2844
2845        assert!(matches!(
2846            config.output_partitioning(),
2847            Partitioning::UnknownPartitioning(4)
2848        ));
2849    }
2850
2851    #[test]
2852    fn test_output_partitioning_no_partition_columns() {
2853        let file_schema = aggr_test_schema();
2854        let config = config_for_projection(
2855            Arc::clone(&file_schema),
2856            None,
2857            Statistics::new_unknown(&file_schema),
2858            vec![], // No partition columns
2859        );
2860
2861        let partitioning = config.output_partitioning();
2862        assert!(matches!(partitioning, Partitioning::UnknownPartitioning(_)));
2863    }
2864
2865    #[test]
2866    fn test_output_partitioning_with_partition_columns() {
2867        let file_schema = aggr_test_schema();
2868
2869        // Test single partition column
2870        let single_partition_col = vec![Field::new(
2871            "date",
2872            wrap_partition_type_in_dict(DataType::Utf8),
2873            false,
2874        )];
2875
2876        let mut config = config_for_projection(
2877            Arc::clone(&file_schema),
2878            None,
2879            Statistics::new_unknown(&file_schema),
2880            single_partition_col,
2881        );
2882        config.file_groups = vec![
2883            FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]),
2884            FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]),
2885            FileGroup::new(vec![PartitionedFile::new("f3.parquet".to_string(), 1024)]),
2886        ];
2887        config.output_partitioning = output_partitioning_from_partition_fields(
2888            config.file_source.table_schema().table_schema(),
2889            config.table_partition_cols(),
2890            config.file_groups.len(),
2891        );
2892
2893        let partitioning = config.output_partitioning();
2894        match partitioning {
2895            Partitioning::Hash(exprs, num_partitions) => {
2896                assert_eq!(num_partitions, 3);
2897                assert_eq!(exprs.len(), 1);
2898                assert_eq!(exprs[0].downcast_ref::<Column>().unwrap().name(), "date");
2899            }
2900            _ => panic!("Expected Hash partitioning"),
2901        }
2902
2903        // Test multiple partition columns
2904        let multiple_partition_cols = vec![
2905            Field::new("year", wrap_partition_type_in_dict(DataType::Utf8), false),
2906            Field::new("month", wrap_partition_type_in_dict(DataType::Utf8), false),
2907        ];
2908
2909        config = config_for_projection(
2910            Arc::clone(&file_schema),
2911            None,
2912            Statistics::new_unknown(&file_schema),
2913            multiple_partition_cols,
2914        );
2915        config.file_groups = vec![
2916            FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]),
2917            FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]),
2918        ];
2919        config.output_partitioning = output_partitioning_from_partition_fields(
2920            config.file_source.table_schema().table_schema(),
2921            config.table_partition_cols(),
2922            config.file_groups.len(),
2923        );
2924
2925        let partitioning = config.output_partitioning();
2926        match partitioning {
2927            Partitioning::Hash(exprs, num_partitions) => {
2928                assert_eq!(num_partitions, 2);
2929                assert_eq!(exprs.len(), 2);
2930                let col_names: Vec<_> = exprs
2931                    .iter()
2932                    .map(|e| e.downcast_ref::<Column>().unwrap().name())
2933                    .collect();
2934                assert_eq!(col_names, vec!["year", "month"]);
2935            }
2936            _ => panic!("Expected Hash partitioning"),
2937        }
2938    }
2939
2940    #[test]
2941    fn try_pushdown_sort_reverses_file_groups_only_when_requested_is_reverse()
2942    -> Result<()> {
2943        let file_schema =
2944            Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
2945
2946        let table_schema = TableSchema::from(&file_schema);
2947        let file_source = Arc::new(InexactSortPushdownSource::new(table_schema));
2948
2949        let file_groups = vec![FileGroup::new(vec![
2950            PartitionedFile::new("file1", 1),
2951            PartitionedFile::new("file2", 1),
2952        ])];
2953
2954        let sort_expr_asc = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
2955        let config =
2956            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
2957                .with_file_groups(file_groups)
2958                .with_output_ordering(vec![
2959                    LexOrdering::new(vec![sort_expr_asc.clone()]).unwrap(),
2960                ])
2961                .build();
2962
2963        let requested_asc = vec![sort_expr_asc.clone()];
2964        let result = config.try_pushdown_sort(&requested_asc)?;
2965        let SortOrderPushdownResult::Inexact { inner } = result else {
2966            panic!("Expected Inexact result");
2967        };
2968        let pushed_config = inner
2969            .downcast_ref::<FileScanConfig>()
2970            .expect("Expected FileScanConfig");
2971        let pushed_files = pushed_config.file_groups[0].files();
2972        assert_eq!(pushed_files[0].object_meta.location.as_ref(), "file1");
2973        assert_eq!(pushed_files[1].object_meta.location.as_ref(), "file2");
2974
2975        let requested_desc = vec![sort_expr_asc.reverse()];
2976        let result = config.try_pushdown_sort(&requested_desc)?;
2977        let SortOrderPushdownResult::Inexact { inner } = result else {
2978            panic!("Expected Inexact result");
2979        };
2980        let pushed_config = inner
2981            .downcast_ref::<FileScanConfig>()
2982            .expect("Expected FileScanConfig");
2983        let pushed_files = pushed_config.file_groups[0].files();
2984        assert_eq!(pushed_files[0].object_meta.location.as_ref(), "file2");
2985        assert_eq!(pushed_files[1].object_meta.location.as_ref(), "file1");
2986
2987        Ok(())
2988    }
2989
2990    fn make_file_with_stats(name: &str, min: f64, max: f64) -> PartitionedFile {
2991        PartitionedFile::new(name.to_string(), 1024).with_statistics(Arc::new(
2992            Statistics {
2993                num_rows: Precision::Exact(100),
2994                total_byte_size: Precision::Exact(1024),
2995                column_statistics: vec![ColumnStatistics {
2996                    null_count: Precision::Exact(0),
2997                    min_value: Precision::Exact(ScalarValue::Float64(Some(min))),
2998                    max_value: Precision::Exact(ScalarValue::Float64(Some(max))),
2999                    ..Default::default()
3000                }],
3001            },
3002        ))
3003    }
3004
3005    #[derive(Clone)]
3006    struct ExactSortPushdownSource {
3007        metrics: ExecutionPlanMetricsSet,
3008        table_schema: TableSchema,
3009    }
3010
3011    impl ExactSortPushdownSource {
3012        fn new(table_schema: TableSchema) -> Self {
3013            Self {
3014                metrics: ExecutionPlanMetricsSet::new(),
3015                table_schema,
3016            }
3017        }
3018    }
3019
3020    impl FileSource for ExactSortPushdownSource {
3021        fn create_file_opener(
3022            &self,
3023            _object_store: Arc<dyn ObjectStore>,
3024            _base_config: &FileScanConfig,
3025            _partition: usize,
3026        ) -> Result<Arc<dyn crate::file_stream::FileOpener>> {
3027            unimplemented!()
3028        }
3029
3030        fn table_schema(&self) -> &TableSchema {
3031            &self.table_schema
3032        }
3033
3034        fn with_batch_size(&self, _batch_size: usize) -> Arc<dyn FileSource> {
3035            Arc::new(self.clone())
3036        }
3037
3038        fn metrics(&self) -> &ExecutionPlanMetricsSet {
3039            &self.metrics
3040        }
3041
3042        fn file_type(&self) -> &str {
3043            "mock_exact"
3044        }
3045
3046        fn try_pushdown_sort(
3047            &self,
3048            _order: &[PhysicalSortExpr],
3049            _eq_properties: &EquivalenceProperties,
3050        ) -> Result<SortOrderPushdownResult<Arc<dyn FileSource>>> {
3051            Ok(SortOrderPushdownResult::Exact {
3052                inner: Arc::new(self.clone()) as Arc<dyn FileSource>,
3053            })
3054        }
3055
3056        fn apply_expressions(
3057            &self,
3058            _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
3059        ) -> Result<TreeNodeRecursion> {
3060            Ok(TreeNodeRecursion::Continue)
3061        }
3062    }
3063
3064    #[test]
3065    fn sort_pushdown_unsupported_source_files_get_sorted() -> Result<()> {
3066        let file_schema =
3067            Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3068        let table_schema = TableSchema::from(&file_schema);
3069        let file_source = Arc::new(MockSource::new(table_schema));
3070
3071        let file_groups = vec![FileGroup::new(vec![
3072            make_file_with_stats("file3", 20.0, 30.0),
3073            make_file_with_stats("file1", 0.0, 9.0),
3074            make_file_with_stats("file2", 10.0, 19.0),
3075        ])];
3076
3077        let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3078        let config =
3079            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3080                .with_file_groups(file_groups)
3081                .build();
3082
3083        let result = config.try_pushdown_sort(&[sort_expr])?;
3084        let SortOrderPushdownResult::Inexact { inner } = result else {
3085            panic!("Expected Inexact result, got {result:?}");
3086        };
3087        let pushed_config = inner
3088            .downcast_ref::<FileScanConfig>()
3089            .expect("Expected FileScanConfig");
3090        let files = pushed_config.file_groups[0].files();
3091        assert_eq!(files[0].object_meta.location.as_ref(), "file1");
3092        assert_eq!(files[1].object_meta.location.as_ref(), "file2");
3093        assert_eq!(files[2].object_meta.location.as_ref(), "file3");
3094        assert!(pushed_config.output_ordering.is_empty());
3095        Ok(())
3096    }
3097
3098    #[test]
3099    fn sort_pushdown_unsupported_source_already_sorted() -> Result<()> {
3100        let file_schema =
3101            Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3102        let table_schema = TableSchema::from(&file_schema);
3103        let file_source = Arc::new(MockSource::new(table_schema));
3104
3105        let file_groups = vec![FileGroup::new(vec![
3106            make_file_with_stats("file1", 0.0, 9.0),
3107            make_file_with_stats("file2", 10.0, 19.0),
3108            make_file_with_stats("file3", 20.0, 30.0),
3109        ])];
3110
3111        let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3112        let config =
3113            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3114                .with_file_groups(file_groups)
3115                .build();
3116
3117        let result = config.try_pushdown_sort(&[sort_expr])?;
3118        assert!(matches!(result, SortOrderPushdownResult::Unsupported));
3119        Ok(())
3120    }
3121
3122    #[test]
3123    fn sort_pushdown_unsupported_source_descending_sort() -> Result<()> {
3124        let file_schema =
3125            Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3126        let table_schema = TableSchema::from(&file_schema);
3127        let file_source = Arc::new(MockSource::new(table_schema));
3128
3129        let file_groups = vec![FileGroup::new(vec![
3130            make_file_with_stats("file1", 0.0, 9.0),
3131            make_file_with_stats("file3", 20.0, 30.0),
3132            make_file_with_stats("file2", 10.0, 19.0),
3133        ])];
3134
3135        let sort_expr = PhysicalSortExpr::new(
3136            Arc::new(Column::new("a", 0)),
3137            arrow::compute::SortOptions {
3138                descending: true,
3139                nulls_first: true,
3140            },
3141        );
3142        let config =
3143            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3144                .with_file_groups(file_groups)
3145                .build();
3146
3147        let result = config.try_pushdown_sort(&[sort_expr])?;
3148        let SortOrderPushdownResult::Inexact { inner } = result else {
3149            panic!("Expected Inexact result");
3150        };
3151        let pushed_config = inner
3152            .downcast_ref::<FileScanConfig>()
3153            .expect("Expected FileScanConfig");
3154        let files = pushed_config.file_groups[0].files();
3155        assert_eq!(files[0].object_meta.location.as_ref(), "file3");
3156        assert_eq!(files[1].object_meta.location.as_ref(), "file2");
3157        assert_eq!(files[2].object_meta.location.as_ref(), "file1");
3158        Ok(())
3159    }
3160
3161    #[test]
3162    fn sort_pushdown_exact_source_non_overlapping_returns_exact() -> Result<()> {
3163        let file_schema =
3164            Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3165        let table_schema = TableSchema::from(&file_schema);
3166        let file_source = Arc::new(ExactSortPushdownSource::new(table_schema));
3167
3168        let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3169
3170        let file_groups = vec![FileGroup::new(vec![
3171            make_file_with_stats("file1", 0.0, 9.0),
3172            make_file_with_stats("file2", 10.0, 19.0),
3173            make_file_with_stats("file3", 20.0, 30.0),
3174        ])];
3175
3176        let config =
3177            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3178                .with_file_groups(file_groups)
3179                .with_output_ordering(vec![
3180                    LexOrdering::new(vec![sort_expr.clone()]).unwrap(),
3181                ])
3182                .build();
3183
3184        let result = config.try_pushdown_sort(&[sort_expr])?;
3185        let SortOrderPushdownResult::Exact { inner } = result else {
3186            panic!("Expected Exact result, got {result:?}");
3187        };
3188        let pushed_config = inner
3189            .downcast_ref::<FileScanConfig>()
3190            .expect("Expected FileScanConfig");
3191        assert!(!pushed_config.output_ordering.is_empty());
3192        Ok(())
3193    }
3194
3195    #[test]
3196    fn sort_pushdown_exact_source_overlapping_downgraded_to_inexact() -> Result<()> {
3197        let file_schema =
3198            Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3199        let table_schema = TableSchema::from(&file_schema);
3200        let file_source = Arc::new(ExactSortPushdownSource::new(table_schema));
3201
3202        let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3203
3204        let file_groups = vec![FileGroup::new(vec![
3205            make_file_with_stats("file1", 0.0, 15.0),
3206            make_file_with_stats("file2", 10.0, 25.0),
3207            make_file_with_stats("file3", 20.0, 30.0),
3208        ])];
3209
3210        let config =
3211            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3212                .with_file_groups(file_groups)
3213                .with_output_ordering(vec![
3214                    LexOrdering::new(vec![sort_expr.clone()]).unwrap(),
3215                ])
3216                .build();
3217
3218        let result = config.try_pushdown_sort(&[sort_expr])?;
3219        let SortOrderPushdownResult::Inexact { inner } = result else {
3220            panic!("Expected Inexact (downgraded), got {result:?}");
3221        };
3222        let pushed_config = inner
3223            .downcast_ref::<FileScanConfig>()
3224            .expect("Expected FileScanConfig");
3225        assert!(pushed_config.output_ordering.is_empty());
3226        Ok(())
3227    }
3228
3229    #[test]
3230    fn sort_pushdown_exact_source_out_of_order_returns_exact() -> Result<()> {
3231        let file_schema =
3232            Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3233        let table_schema = TableSchema::from(&file_schema);
3234        let file_source = Arc::new(ExactSortPushdownSource::new(table_schema));
3235
3236        let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3237
3238        let file_groups = vec![FileGroup::new(vec![
3239            make_file_with_stats("file3", 20.0, 30.0),
3240            make_file_with_stats("file1", 0.0, 9.0),
3241            make_file_with_stats("file2", 10.0, 19.0),
3242        ])];
3243
3244        let config =
3245            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3246                .with_file_groups(file_groups)
3247                .with_output_ordering(vec![
3248                    LexOrdering::new(vec![sort_expr.clone()]).unwrap(),
3249                ])
3250                .build();
3251
3252        let result = config.try_pushdown_sort(&[sort_expr])?;
3253        let SortOrderPushdownResult::Exact { inner } = result else {
3254            panic!("Expected Exact result, got {result:?}");
3255        };
3256        let pushed_config = inner
3257            .downcast_ref::<FileScanConfig>()
3258            .expect("Expected FileScanConfig");
3259        let files = pushed_config.file_groups[0].files();
3260        assert_eq!(files[0].object_meta.location.as_ref(), "file1");
3261        assert_eq!(files[1].object_meta.location.as_ref(), "file2");
3262        assert_eq!(files[2].object_meta.location.as_ref(), "file3");
3263        assert!(!pushed_config.output_ordering.is_empty());
3264        Ok(())
3265    }
3266
3267    #[test]
3268    fn sort_pushdown_unsupported_source_single_file_groups() -> Result<()> {
3269        let file_schema =
3270            Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3271        let table_schema = TableSchema::from(&file_schema);
3272        let file_source = Arc::new(MockSource::new(table_schema));
3273
3274        let file_groups = vec![
3275            FileGroup::new(vec![make_file_with_stats("file1", 0.0, 9.0)]),
3276            FileGroup::new(vec![make_file_with_stats("file2", 10.0, 19.0)]),
3277        ];
3278
3279        let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3280        let config =
3281            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3282                .with_file_groups(file_groups)
3283                .build();
3284
3285        let result = config.try_pushdown_sort(&[sort_expr])?;
3286        assert!(
3287            matches!(result, SortOrderPushdownResult::Unsupported),
3288            "Expected Unsupported for single-file groups"
3289        );
3290        Ok(())
3291    }
3292
3293    #[test]
3294    fn sort_pushdown_unsupported_source_multiple_groups() -> Result<()> {
3295        let file_schema =
3296            Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3297        let table_schema = TableSchema::from(&file_schema);
3298        let file_source = Arc::new(MockSource::new(table_schema));
3299
3300        let file_groups = vec![
3301            FileGroup::new(vec![
3302                make_file_with_stats("file_b", 10.0, 19.0),
3303                make_file_with_stats("file_a", 0.0, 9.0),
3304            ]),
3305            FileGroup::new(vec![
3306                make_file_with_stats("file_d", 30.0, 39.0),
3307                make_file_with_stats("file_c", 20.0, 29.0),
3308            ]),
3309        ];
3310
3311        let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3312        let config =
3313            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3314                .with_file_groups(file_groups)
3315                .build();
3316
3317        let result = config.try_pushdown_sort(&[sort_expr])?;
3318        let SortOrderPushdownResult::Inexact { inner } = result else {
3319            panic!("Expected Inexact result");
3320        };
3321        let pushed_config = inner
3322            .downcast_ref::<FileScanConfig>()
3323            .expect("Expected FileScanConfig");
3324        let files0 = pushed_config.file_groups[0].files();
3325        assert_eq!(files0[0].object_meta.location.as_ref(), "file_a");
3326        assert_eq!(files0[1].object_meta.location.as_ref(), "file_b");
3327        let files1 = pushed_config.file_groups[1].files();
3328        assert_eq!(files1[0].object_meta.location.as_ref(), "file_c");
3329        assert_eq!(files1[1].object_meta.location.as_ref(), "file_d");
3330        Ok(())
3331    }
3332
3333    #[test]
3334    fn sort_pushdown_unsupported_source_partial_statistics() -> Result<()> {
3335        let file_schema =
3336            Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3337        let table_schema = TableSchema::from(&file_schema);
3338        let file_source = Arc::new(MockSource::new(table_schema));
3339
3340        let file_groups = vec![
3341            FileGroup::new(vec![
3342                make_file_with_stats("file_b", 10.0, 19.0),
3343                make_file_with_stats("file_a", 0.0, 9.0),
3344            ]),
3345            FileGroup::new(vec![
3346                PartitionedFile::new("file_d".to_string(), 1024),
3347                PartitionedFile::new("file_c".to_string(), 1024),
3348            ]),
3349        ];
3350
3351        let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3352        let config =
3353            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3354                .with_file_groups(file_groups)
3355                .build();
3356
3357        let result = config.try_pushdown_sort(&[sort_expr])?;
3358        let SortOrderPushdownResult::Inexact { inner } = result else {
3359            panic!("Expected Inexact result");
3360        };
3361        let pushed_config = inner
3362            .downcast_ref::<FileScanConfig>()
3363            .expect("Expected FileScanConfig");
3364        let files0 = pushed_config.file_groups[0].files();
3365        assert_eq!(files0[0].object_meta.location.as_ref(), "file_a");
3366        assert_eq!(files0[1].object_meta.location.as_ref(), "file_b");
3367        let files1 = pushed_config.file_groups[1].files();
3368        assert_eq!(files1[0].object_meta.location.as_ref(), "file_d");
3369        assert_eq!(files1[1].object_meta.location.as_ref(), "file_c");
3370        Ok(())
3371    }
3372
3373    #[test]
3374    fn sort_pushdown_inexact_source_with_statistics_sorting() -> Result<()> {
3375        let file_schema =
3376            Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3377        let table_schema = TableSchema::from(&file_schema);
3378        let file_source = Arc::new(InexactSortPushdownSource::new(table_schema));
3379
3380        let file_groups = vec![FileGroup::new(vec![
3381            make_file_with_stats("file2", 10.0, 19.0),
3382            make_file_with_stats("file1", 0.0, 9.0),
3383        ])];
3384
3385        let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3386        let config =
3387            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3388                .with_file_groups(file_groups)
3389                .build();
3390
3391        let result = config.try_pushdown_sort(&[sort_expr])?;
3392        let SortOrderPushdownResult::Inexact { inner } = result else {
3393            panic!("Expected Inexact result");
3394        };
3395        let pushed_config = inner
3396            .downcast_ref::<FileScanConfig>()
3397            .expect("Expected FileScanConfig");
3398        let files = pushed_config.file_groups[0].files();
3399        assert_eq!(files[0].object_meta.location.as_ref(), "file1");
3400        assert_eq!(files[1].object_meta.location.as_ref(), "file2");
3401        assert!(pushed_config.output_ordering.is_empty());
3402        Ok(())
3403    }
3404
3405    #[test]
3406    fn sort_pushdown_exact_multi_group_preserves_parallelism() -> Result<()> {
3407        // ExactSortPushdownSource + 4 non-overlapping files in 2 interleaved groups.
3408        // Groups should NOT be redistributed — interleaved groups allow SPM to
3409        // pull from both partitions concurrently, keeping parallel I/O active.
3410        // Redistributing consecutively would make SPM read one partition at a
3411        // time (all values in group 0 < group 1), degrading to single-threaded I/O.
3412        let file_schema =
3413            Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3414        let table_schema = TableSchema::from(&file_schema);
3415        let file_source = Arc::new(ExactSortPushdownSource::new(table_schema));
3416
3417        let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3418
3419        // 2 groups with interleaved ranges (simulating bin-packing result):
3420        // Group 0: [file_01(0-9), file_03(20-29)]
3421        // Group 1: [file_02(10-19), file_04(30-39)]
3422        let file_groups = vec![
3423            FileGroup::new(vec![
3424                make_file_with_stats("file_01", 0.0, 9.0),
3425                make_file_with_stats("file_03", 20.0, 29.0),
3426            ]),
3427            FileGroup::new(vec![
3428                make_file_with_stats("file_02", 10.0, 19.0),
3429                make_file_with_stats("file_04", 30.0, 39.0),
3430            ]),
3431        ];
3432
3433        let config =
3434            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3435                .with_file_groups(file_groups)
3436                .with_output_ordering(vec![
3437                    LexOrdering::new(vec![sort_expr.clone()]).unwrap(),
3438                ])
3439                .build();
3440
3441        let result = config.try_pushdown_sort(&[sort_expr])?;
3442        let SortOrderPushdownResult::Exact { inner } = result else {
3443            panic!("Expected Exact result, got {result:?}");
3444        };
3445        let pushed_config = inner
3446            .downcast_ref::<FileScanConfig>()
3447            .expect("Expected FileScanConfig");
3448
3449        // 2 groups preserved (parallelism maintained)
3450        assert_eq!(pushed_config.file_groups.len(), 2);
3451
3452        // Files within each group are sorted by stats, but groups are NOT
3453        // redistributed — interleaved assignment from bin-packing is kept
3454        let files0 = pushed_config.file_groups[0].files();
3455        assert_eq!(files0[0].object_meta.location.as_ref(), "file_01");
3456        assert_eq!(files0[1].object_meta.location.as_ref(), "file_03");
3457        let files1 = pushed_config.file_groups[1].files();
3458        assert_eq!(files1[0].object_meta.location.as_ref(), "file_02");
3459        assert_eq!(files1[1].object_meta.location.as_ref(), "file_04");
3460
3461        // output_ordering preserved (Exact, each group internally non-overlapping)
3462        assert!(!pushed_config.output_ordering.is_empty());
3463        Ok(())
3464    }
3465
3466    #[test]
3467    fn sort_pushdown_reverse_preserves_file_order_with_stats() -> Result<()> {
3468        // Reverse scan should reverse file order but NOT apply statistics-based
3469        // sorting (which would undo the reversal). The result is Inexact.
3470        let file_schema =
3471            Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3472        let table_schema = TableSchema::from(&file_schema);
3473        let file_source = Arc::new(InexactSortPushdownSource::new(table_schema));
3474
3475        let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3476
3477        // Files with stats, in ASC order. Output ordering is [a ASC].
3478        let file_groups = vec![FileGroup::new(vec![
3479            make_file_with_stats("file1", 0.0, 9.0),
3480            make_file_with_stats("file2", 10.0, 19.0),
3481            make_file_with_stats("file3", 20.0, 30.0),
3482        ])];
3483
3484        let config =
3485            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3486                .with_file_groups(file_groups)
3487                .with_output_ordering(vec![
3488                    LexOrdering::new(vec![sort_expr.clone()]).unwrap(),
3489                ])
3490                .build();
3491
3492        // Request DESC → reverse path
3493        let result = config.try_pushdown_sort(&[sort_expr.reverse()])?;
3494        let SortOrderPushdownResult::Inexact { inner } = result else {
3495            panic!("Expected Inexact for reverse scan, got {result:?}");
3496        };
3497        let pushed_config = inner
3498            .downcast_ref::<FileScanConfig>()
3499            .expect("Expected FileScanConfig");
3500
3501        // Files should be reversed (not re-sorted by stats)
3502        let files = pushed_config.file_groups[0].files();
3503        assert_eq!(files[0].object_meta.location.as_ref(), "file3");
3504        assert_eq!(files[1].object_meta.location.as_ref(), "file2");
3505        assert_eq!(files[2].object_meta.location.as_ref(), "file1");
3506
3507        // output_ordering cleared (Inexact)
3508        assert!(pushed_config.output_ordering.is_empty());
3509        Ok(())
3510    }
3511
3512    /// Helper: create a PartitionedFile with stats including null count
3513    fn make_file_with_null_stats(
3514        name: &str,
3515        min: f64,
3516        max: f64,
3517        null_count: usize,
3518    ) -> PartitionedFile {
3519        PartitionedFile::new(name.to_string(), 1024).with_statistics(Arc::new(
3520            Statistics {
3521                num_rows: Precision::Exact(100),
3522                total_byte_size: Precision::Exact(1024),
3523                column_statistics: vec![ColumnStatistics {
3524                    null_count: Precision::Exact(null_count),
3525                    min_value: Precision::Exact(ScalarValue::Float64(Some(min))),
3526                    max_value: Precision::Exact(ScalarValue::Float64(Some(max))),
3527                    ..Default::default()
3528                }],
3529            },
3530        ))
3531    }
3532
3533    #[test]
3534    fn sort_pushdown_unsupported_with_nulls_does_not_upgrade_to_exact() -> Result<()> {
3535        // Files are non-overlapping but one has NULLs.
3536        // Should NOT upgrade to Exact — NULLs would appear in wrong position.
3537        let file_schema =
3538            Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, true)]));
3539        let table_schema = TableSchema::from(&file_schema);
3540        let file_source = Arc::new(MockSource::new(table_schema));
3541
3542        let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3543
3544        // Files in wrong order (high min first) to trigger reordering
3545        let file_groups = vec![FileGroup::new(vec![
3546            make_file_with_null_stats("b_no_nulls", 10.0, 19.0, 0),
3547            make_file_with_null_stats("a_with_nulls", 0.0, 9.0, 5), // has NULLs
3548        ])];
3549
3550        let config =
3551            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3552                .with_file_groups(file_groups)
3553                .with_output_ordering(vec![
3554                    LexOrdering::new(vec![sort_expr.clone()]).unwrap(),
3555                ])
3556                .build();
3557
3558        let result = config.try_pushdown_sort(&[sort_expr])?;
3559        // Should be Inexact (not Exact) because of NULLs
3560        assert!(
3561            matches!(result, SortOrderPushdownResult::Inexact { .. }),
3562            "Expected Inexact due to NULLs, got {result:?}"
3563        );
3564        Ok(())
3565    }
3566
3567    #[test]
3568    fn sort_pushdown_unsupported_no_nulls_upgrades_to_exact() -> Result<()> {
3569        // Files are non-overlapping, no NULLs → should upgrade to Exact
3570        let file_schema =
3571            Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, true)]));
3572        let table_schema = TableSchema::from(&file_schema);
3573        let file_source = Arc::new(MockSource::new(table_schema));
3574
3575        let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3576
3577        let file_groups = vec![FileGroup::new(vec![
3578            make_file_with_null_stats("b_high", 10.0, 19.0, 0),
3579            make_file_with_null_stats("a_low", 0.0, 9.0, 0),
3580        ])];
3581
3582        let config =
3583            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3584                .with_file_groups(file_groups)
3585                .with_output_ordering(vec![
3586                    LexOrdering::new(vec![sort_expr.clone()]).unwrap(),
3587                ])
3588                .build();
3589
3590        let result = config.try_pushdown_sort(&[sort_expr])?;
3591        assert!(
3592            matches!(result, SortOrderPushdownResult::Exact { .. }),
3593            "Expected Exact (no NULLs), got {result:?}"
3594        );
3595        Ok(())
3596    }
3597
3598    /// Helper: build a `ProjectionExprs` from `(expr, alias)` pairs.
3599    fn make_projection(pairs: Vec<(Arc<dyn PhysicalExpr>, &str)>) -> ProjectionExprs {
3600        ProjectionExprs::new(
3601            pairs
3602                .into_iter()
3603                .map(|(expr, alias)| ProjectionExpr::new(expr, alias)),
3604        )
3605    }
3606
3607    /// Helper: create a volatile (non-deterministic) function expression,
3608    /// e.g. `random()`.
3609    fn make_volatile_expr() -> Arc<dyn PhysicalExpr> {
3610        use datafusion_common::config::ConfigOptions;
3611        use datafusion_expr::ScalarUDF;
3612        use datafusion_functions::math::random::RandomFunc;
3613        use datafusion_physical_expr::ScalarFunctionExpr;
3614
3615        Arc::new(ScalarFunctionExpr::new(
3616            "random",
3617            Arc::new(ScalarUDF::from(RandomFunc::new())),
3618            vec![],
3619            Arc::new(Field::new("random", DataType::Float64, false)),
3620            Arc::new(ConfigOptions::default()),
3621        ))
3622    }
3623
3624    /// Helper: create a deterministic but expensive scalar-function
3625    /// expression, e.g. `abs(<arg>)`.
3626    fn make_udf_expr(args: Vec<Arc<dyn PhysicalExpr>>) -> Arc<dyn PhysicalExpr> {
3627        use datafusion_common::config::ConfigOptions;
3628        use datafusion_expr::ScalarUDF;
3629        use datafusion_functions::math::abs::AbsFunc;
3630        use datafusion_physical_expr::ScalarFunctionExpr;
3631
3632        Arc::new(ScalarFunctionExpr::new(
3633            "abs",
3634            Arc::new(ScalarUDF::from(AbsFunc::new())),
3635            args,
3636            Arc::new(Field::new("abs", DataType::Int32, false)),
3637            Arc::new(ConfigOptions::default()),
3638        ))
3639    }
3640
3641    /// Helper: create a cheap, leaf-pushable scalar function — struct field
3642    /// access `get_field(s, 'x')`, whose placement is `MoveTowardsLeafNodes`
3643    /// when the base is a column and the key is a literal.
3644    fn make_leaf_pushable_expr() -> Arc<dyn PhysicalExpr> {
3645        use datafusion_common::config::ConfigOptions;
3646        use datafusion_expr::ScalarUDF;
3647        use datafusion_functions::core::getfield::GetFieldFunc;
3648        use datafusion_physical_expr::ScalarFunctionExpr;
3649        use datafusion_physical_expr::expressions::Literal;
3650
3651        Arc::new(ScalarFunctionExpr::new(
3652            "get_field",
3653            Arc::new(ScalarUDF::from(GetFieldFunc::new())),
3654            vec![
3655                Arc::new(Column::new("s", 0)),
3656                Arc::new(Literal::new(ScalarValue::Utf8(Some("x".to_string())))),
3657            ],
3658            Arc::new(Field::new("x", DataType::Int32, true)),
3659            Arc::new(ConfigOptions::default()),
3660        ))
3661    }
3662
3663    /// Column-only inner projections always merge safely, even when
3664    /// the outer projection references them multiple times.
3665    #[test]
3666    fn test_would_duplicate_allows_column_only_inner() {
3667        let col_a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
3668        let col_b: Arc<dyn PhysicalExpr> = Arc::new(Column::new("b", 1));
3669
3670        let inner =
3671            make_projection(vec![(Arc::clone(&col_a), "a"), (Arc::clone(&col_b), "b")]);
3672
3673        // Outer references col 0 twice
3674        let outer = make_projection(vec![
3675            (Arc::new(Column::new("a", 0)), "x"),
3676            (Arc::new(Column::new("a", 0)), "y"),
3677        ]);
3678
3679        assert!(!would_duplicate_costly_exprs(&inner, &outer));
3680    }
3681
3682    /// A non-trivial computed expression (arithmetic, `KeepInPlace`) referenced
3683    /// multiple times blocks the merge — recomputing it per site is wasteful.
3684    #[test]
3685    fn test_would_duplicate_blocks_computed_multi_ref() {
3686        let col_a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
3687        let col_b: Arc<dyn PhysicalExpr> = Arc::new(Column::new("b", 1));
3688        // Inner: [a + b, b]  (index 0 is a non-trivial computed expression)
3689        let inner = make_projection(vec![
3690            (
3691                Arc::new(BinaryExpr::new(
3692                    Arc::clone(&col_a),
3693                    Operator::Plus,
3694                    Arc::clone(&col_b),
3695                )),
3696                "sum",
3697            ),
3698            (Arc::clone(&col_b), "b"),
3699        ]);
3700
3701        // Outer references index 0 twice
3702        let outer = make_projection(vec![
3703            (Arc::new(Column::new("sum", 0)), "x"),
3704            (Arc::new(Column::new("sum", 0)), "y"),
3705        ]);
3706
3707        assert!(would_duplicate_costly_exprs(&inner, &outer));
3708    }
3709
3710    /// A volatile expression the outer projection does not reference is
3711    /// safe to merge (it is projected away, not duplicated).
3712    #[test]
3713    fn test_would_duplicate_allows_unreferenced_volatile() {
3714        let col_a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
3715        // Inner: [random(), a]
3716        let inner =
3717            make_projection(vec![(make_volatile_expr(), "r"), (Arc::clone(&col_a), "a")]);
3718
3719        // Outer references only index 1 (the column), not the volatile expr
3720        let outer = make_projection(vec![(Arc::new(Column::new("a", 1)), "a")]);
3721
3722        assert!(!would_duplicate_costly_exprs(&inner, &outer));
3723    }
3724
3725    /// A volatile expression referenced multiple times must block merge:
3726    /// this is the #23220 regression (`random()` aliased then referenced as
3727    /// `x` and `y`).
3728    #[test]
3729    fn test_would_duplicate_blocks_multi_ref_volatile() {
3730        // Inner: [random()]
3731        let inner = make_projection(vec![(make_volatile_expr(), "r")]);
3732
3733        // Outer references index 0 twice
3734        let outer = make_projection(vec![
3735            (Arc::new(Column::new("r", 0)), "x"),
3736            (Arc::new(Column::new("r", 0)), "y"),
3737        ]);
3738
3739        assert!(would_duplicate_costly_exprs(&inner, &outer));
3740    }
3741
3742    /// A volatile expression referenced exactly once has nothing to duplicate,
3743    /// so the merge is allowed.
3744    #[test]
3745    fn test_would_duplicate_allows_single_ref_volatile() {
3746        let col_a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
3747        // Inner: [random(), a]
3748        let inner =
3749            make_projection(vec![(make_volatile_expr(), "r"), (Arc::clone(&col_a), "a")]);
3750
3751        // Outer references the volatile expression exactly once
3752        let outer = make_projection(vec![
3753            (Arc::new(Column::new("r", 0)), "x"),
3754            (Arc::new(Column::new("a", 1)), "a"),
3755        ]);
3756
3757        assert!(!would_duplicate_costly_exprs(&inner, &outer));
3758    }
3759
3760    /// References are counted with multiplicity, so a single outer expression
3761    /// that duplicates the value (e.g. `r + r`) still blocks the merge.
3762    #[test]
3763    fn test_would_duplicate_blocks_single_expr_self_ref_volatile() {
3764        // Inner: [random()]
3765        let inner = make_projection(vec![(make_volatile_expr(), "r")]);
3766
3767        // Outer: [r + r] — one expression referencing `random()` twice
3768        let outer = make_projection(vec![(
3769            Arc::new(BinaryExpr::new(
3770                Arc::new(Column::new("r", 0)),
3771                Operator::Plus,
3772                Arc::new(Column::new("r", 0)),
3773            )),
3774            "x",
3775        )]);
3776
3777        assert!(would_duplicate_costly_exprs(&inner, &outer));
3778    }
3779
3780    /// A volatile expression buried inside a larger expression (e.g.
3781    /// `random() + 1`) is still detected and blocks merge.
3782    #[test]
3783    fn test_would_duplicate_blocks_volatile_nested_in_arithmetic() {
3784        // Inner: [random() + 1]
3785        let inner = make_projection(vec![(
3786            Arc::new(BinaryExpr::new(
3787                make_volatile_expr(),
3788                Operator::Plus,
3789                Arc::new(Literal::new(ScalarValue::Float64(Some(1.0)))),
3790            )),
3791            "expr",
3792        )]);
3793
3794        // Outer references index 0 twice
3795        let outer = make_projection(vec![
3796            (Arc::new(Column::new("expr", 0)), "x"),
3797            (Arc::new(Column::new("expr", 0)), "y"),
3798        ]);
3799
3800        assert!(would_duplicate_costly_exprs(&inner, &outer));
3801    }
3802
3803    /// Empty projections should not block merging.
3804    #[test]
3805    fn test_would_duplicate_empty_projections() {
3806        let inner = make_projection(vec![]);
3807        let outer = make_projection(vec![]);
3808        assert!(!would_duplicate_costly_exprs(&inner, &outer));
3809    }
3810
3811    /// An expensive (scalar-function) expression referenced more than once
3812    /// must block the merge to preserve CSE.
3813    #[test]
3814    fn test_would_duplicate_blocks_multi_ref_expensive() {
3815        let col_a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
3816        // Inner: [abs(a)]
3817        let inner = make_projection(vec![(make_udf_expr(vec![col_a]), "abs_a")]);
3818
3819        // Outer references index 0 twice
3820        let outer = make_projection(vec![
3821            (Arc::new(Column::new("abs_a", 0)), "x"),
3822            (Arc::new(Column::new("abs_a", 0)), "y"),
3823        ]);
3824
3825        assert!(would_duplicate_costly_exprs(&inner, &outer));
3826    }
3827
3828    /// An expensive expression referenced only once has nothing to duplicate,
3829    /// so the merge is allowed.
3830    #[test]
3831    fn test_would_duplicate_allows_single_ref_expensive() {
3832        let col_a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
3833        // Inner: [abs(a), a]
3834        let inner = make_projection(vec![
3835            (make_udf_expr(vec![Arc::clone(&col_a)]), "abs_a"),
3836            (Arc::clone(&col_a), "a"),
3837        ]);
3838
3839        // Outer references each inner column once
3840        let outer = make_projection(vec![
3841            (Arc::new(Column::new("abs_a", 0)), "out"),
3842            (Arc::new(Column::new("a", 1)), "a"),
3843        ]);
3844
3845        assert!(!would_duplicate_costly_exprs(&inner, &outer));
3846    }
3847
3848    /// A cheap, leaf-pushable scalar function (placement
3849    /// `MoveTowardsLeafNodes`, e.g. `get_field` / `input_file_name`) still
3850    /// merges even when referenced multiple times — it is meant to be pushed
3851    /// into the scan, so blocking would defeat that optimization.
3852    #[test]
3853    fn test_would_duplicate_allows_leaf_pushable_scalar_function() {
3854        // Inner: [input_file_name()]
3855        let inner = make_projection(vec![(make_leaf_pushable_expr(), "f")]);
3856
3857        // Outer references index 0 twice
3858        let outer = make_projection(vec![
3859            (Arc::new(Column::new("f", 0)), "x"),
3860            (Arc::new(Column::new("f", 0)), "y"),
3861        ]);
3862
3863        assert!(!would_duplicate_costly_exprs(&inner, &outer));
3864    }
3865}