use TableSchema

adriangb · adriangb · commit 5f605a91f361 · 2025-10-23T13:42:15.000-05:00
diff --git a/datafusion-examples/examples/csv_json_opener.rs b/datafusion-examples/examples/csv_json_opener.rs
@@ -31,7 +31,9 @@ use datafusion::{
     test_util::aggr_test_schema,
 };
 
-use datafusion::datasource::physical_plan::FileScanConfigBuilder;
+use datafusion::datasource::{
+    physical_plan::FileScanConfigBuilder, table_schema::TableSchema,
+};
 use futures::StreamExt;
 use object_store::{local::LocalFileSystem, memory::InMemory, ObjectStore};
 
@@ -67,7 +69,7 @@ async fn csv_opener() -> Result<()> {
 
     let config = CsvSource::new(true, b',', b'"')
         .with_comment(Some(b'#'))
-        .with_schema(schema)
+        .with_schema(TableSchema::new(schema, vec![]))
         .with_batch_size(8192)
         .with_projection(&scan_config);
 
diff --git a/datafusion/core/src/datasource/mod.rs b/datafusion/core/src/datasource/mod.rs
@@ -45,6 +45,7 @@ pub use datafusion_catalog::view;
 pub use datafusion_datasource::schema_adapter;
 pub use datafusion_datasource::sink;
 pub use datafusion_datasource::source;
+pub use datafusion_datasource::table_schema;
 pub use datafusion_execution::object_store;
 pub use datafusion_physical_expr::create_ordering;
 
diff --git a/datafusion/core/src/datasource/physical_plan/parquet.rs b/datafusion/core/src/datasource/physical_plan/parquet.rs
@@ -54,7 +54,7 @@ mod tests {
     use datafusion_datasource::source::DataSourceExec;
 
     use datafusion_datasource::file::FileSource;
-    use datafusion_datasource::{FileRange, PartitionedFile};
+    use datafusion_datasource::{FileRange, PartitionedFile, TableSchema};
     use datafusion_datasource_parquet::source::ParquetSource;
     use datafusion_datasource_parquet::{
         DefaultParquetFileReaderFactory, ParquetFileReaderFactory, ParquetFormat,
@@ -186,7 +186,7 @@ mod tests {
                 source = source.with_bloom_filter_on_read(false);
             }
 
-            source.with_schema(Arc::clone(&table_schema))
+            source.with_schema(TableSchema::new(Arc::clone(&table_schema), vec![]))
         }
 
         fn build_parquet_exec(
diff --git a/datafusion/core/src/test_util/parquet.rs b/datafusion/core/src/test_util/parquet.rs
@@ -186,7 +186,10 @@ impl TestParquetFile {
                 ParquetSource::new(parquet_options)
                     .with_predicate(Arc::clone(&physical_filter_expr)),
             )
-            .with_schema(Arc::clone(&self.schema));
+            .with_schema(datafusion_datasource::TableSchema::new(
+                Arc::clone(&self.schema),
+                vec![],
+            ));
             let config = scan_config_builder.with_source(source).build();
             let parquet_exec = DataSourceExec::from_data_source(config);
 
diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown/util.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown/util.rs
@@ -24,6 +24,7 @@ use datafusion_datasource::{
     file_scan_config::FileScanConfigBuilder, file_stream::FileOpenFuture,
     file_stream::FileOpener, schema_adapter::DefaultSchemaAdapterFactory,
     schema_adapter::SchemaAdapterFactory, source::DataSourceExec, PartitionedFile,
+    TableSchema,
 };
 use datafusion_physical_expr_common::physical_expr::fmt_sql;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
@@ -156,9 +157,9 @@ impl FileSource for TestSource {
         })
     }
 
-    fn with_schema(&self, schema: SchemaRef) -> Arc<dyn FileSource> {
+    fn with_schema(&self, schema: TableSchema) -> Arc<dyn FileSource> {
         Arc::new(TestSource {
-            schema: Some(schema),
+            schema: Some(schema.file_schema().clone()),
             ..self.clone()
         })
     }
diff --git a/datafusion/datasource-arrow/src/source.rs b/datafusion/datasource-arrow/src/source.rs
@@ -20,9 +20,9 @@ use std::sync::Arc;
 
 use datafusion_datasource::as_file_source;
 use datafusion_datasource::schema_adapter::SchemaAdapterFactory;
+use datafusion_datasource::TableSchema;
 
 use arrow::buffer::Buffer;
-use arrow::datatypes::SchemaRef;
 use arrow_ipc::reader::FileDecoder;
 use datafusion_common::error::Result;
 use datafusion_common::{exec_datafusion_err, Statistics};
@@ -73,7 +73,7 @@ impl FileSource for ArrowSource {
         Arc::new(Self { ..self.clone() })
     }
 
-    fn with_schema(&self, _schema: SchemaRef) -> Arc<dyn FileSource> {
+    fn with_schema(&self, _schema: TableSchema) -> Arc<dyn FileSource> {
         Arc::new(Self { ..self.clone() })
     }
     fn with_statistics(&self, statistics: Statistics) -> Arc<dyn FileSource> {
diff --git a/datafusion/datasource-avro/src/source.rs b/datafusion/datasource-avro/src/source.rs
@@ -29,6 +29,7 @@ use datafusion_datasource::file::FileSource;
 use datafusion_datasource::file_scan_config::FileScanConfig;
 use datafusion_datasource::file_stream::FileOpener;
 use datafusion_datasource::schema_adapter::SchemaAdapterFactory;
+use datafusion_datasource::TableSchema;
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
 
@@ -84,9 +85,9 @@ impl FileSource for AvroSource {
         Arc::new(conf)
     }
 
-    fn with_schema(&self, schema: SchemaRef) -> Arc<dyn FileSource> {
+    fn with_schema(&self, schema: TableSchema) -> Arc<dyn FileSource> {
         let mut conf = self.clone();
-        conf.schema = Some(schema);
+        conf.schema = Some(schema.file_schema().clone());
         Arc::new(conf)
     }
     fn with_statistics(&self, statistics: Statistics) -> Arc<dyn FileSource> {
diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs
@@ -29,7 +29,7 @@ use datafusion_datasource::file_compression_type::FileCompressionType;
 use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};
 use datafusion_datasource::{
     as_file_source, calculate_range, FileRange, ListingTableUrl, PartitionedFile,
-    RangeCalculation,
+    RangeCalculation, TableSchema,
 };
 
 use arrow::csv;
@@ -258,9 +258,9 @@ impl FileSource for CsvSource {
         Arc::new(conf)
     }
 
-    fn with_schema(&self, schema: SchemaRef) -> Arc<dyn FileSource> {
+    fn with_schema(&self, schema: TableSchema) -> Arc<dyn FileSource> {
         let mut conf = self.clone();
-        conf.file_schema = Some(schema);
+        conf.file_schema = Some(schema.file_schema().clone());
         Arc::new(conf)
     }
 
diff --git a/datafusion/datasource-json/src/source.rs b/datafusion/datasource-json/src/source.rs
@@ -32,6 +32,7 @@ use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};
 use datafusion_datasource::schema_adapter::SchemaAdapterFactory;
 use datafusion_datasource::{
     as_file_source, calculate_range, ListingTableUrl, PartitionedFile, RangeCalculation,
+    TableSchema,
 };
 use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 
@@ -122,7 +123,7 @@ impl FileSource for JsonSource {
         Arc::new(conf)
     }
 
-    fn with_schema(&self, _schema: SchemaRef) -> Arc<dyn FileSource> {
+    fn with_schema(&self, _schema: TableSchema) -> Arc<dyn FileSource> {
         Arc::new(Self { ..self.clone() })
     }
     fn with_statistics(&self, statistics: Statistics) -> Arc<dyn FileSource> {
diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs
@@ -463,8 +463,7 @@ impl FileFormat for ParquetFormat {
             metadata_size_hint = Some(metadata);
         }
 
-        let mut source = ParquetSource::new(self.options.clone())
-            .with_table_partition_cols(conf.table_partition_cols().clone());
+        let mut source = ParquetSource::new(self.options.clone());
 
         // Use the CachedParquetFileReaderFactory
         let metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache();
diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs
@@ -26,8 +26,6 @@ use crate::opener::ParquetOpener;
 use crate::row_filter::can_expr_be_pushed_down_with_schemas;
 use crate::DefaultParquetFileReaderFactory;
 use crate::ParquetFileReaderFactory;
-use arrow::datatypes::Field;
-use arrow::datatypes::Schema;
 use datafusion_common::config::ConfigOptions;
 #[cfg(feature = "parquet_encryption")]
 use datafusion_common::config::EncryptionFactoryOptions;
@@ -37,11 +35,12 @@ use datafusion_datasource::schema_adapter::{
     DefaultSchemaAdapterFactory, SchemaAdapterFactory,
 };
 
-use arrow::datatypes::{SchemaRef, TimeUnit};
+use arrow::datatypes::TimeUnit;
 use datafusion_common::config::TableParquetOptions;
 use datafusion_common::{DataFusionError, Statistics};
 use datafusion_datasource::file::FileSource;
 use datafusion_datasource::file_scan_config::FileScanConfig;
+use datafusion_datasource::TableSchema;
 use datafusion_physical_expr::conjunction;
 use datafusion_physical_expr_adapter::DefaultPhysicalExprAdapterFactory;
 use datafusion_physical_expr_common::physical_expr::fmt_sql;
@@ -277,7 +276,7 @@ pub struct ParquetSource {
     /// The schema of the file.
     /// In particular, this is the schema of the table without partition columns,
     /// *not* the physical schema of the file.
-    pub(crate) file_schema: Option<SchemaRef>,
+    pub(crate) table_schema: Option<TableSchema>,
     /// Optional predicate for row filtering during parquet scan
     pub(crate) predicate: Option<Arc<dyn PhysicalExpr>>,
     /// Optional user defined parquet file reader factory
@@ -289,7 +288,6 @@ pub struct ParquetSource {
     /// Optional hint for the size of the parquet metadata
     pub(crate) metadata_size_hint: Option<usize>,
     pub(crate) projected_statistics: Option<Statistics>,
-    pub(crate) table_partition_cols: Vec<Arc<Field>>,
     #[cfg(feature = "parquet_encryption")]
     pub(crate) encryption_factory: Option<Arc<dyn EncryptionFactory>>,
 }
@@ -323,11 +321,6 @@ impl ParquetSource {
         conf
     }
 
-    pub fn with_table_partition_cols(mut self, partition_cols: Vec<Arc<Field>>) -> Self {
-        self.table_partition_cols = partition_cols;
-        self
-    }
-
     /// Set the encryption factory to use to generate file decryption properties
     #[cfg(feature = "parquet_encryption")]
     pub fn with_encryption_factory(
@@ -609,9 +602,9 @@ impl FileSource for ParquetSource {
         Arc::new(conf)
     }
 
-    fn with_schema(&self, schema: SchemaRef) -> Arc<dyn FileSource> {
+    fn with_schema(&self, schema: TableSchema) -> Arc<dyn FileSource> {
         Arc::new(Self {
-            file_schema: Some(schema),
+            table_schema: Some(schema),
             ..self.clone()
         })
     }
@@ -669,9 +662,10 @@ impl FileSource for ParquetSource {
                 // the actual predicates are built in reference to the physical schema of
                 // each file, which we do not have at this point and hence cannot use.
                 // Instead we use the logical schema of the file (the table schema without partition columns).
-                if let (Some(file_schema), Some(predicate)) =
-                    (&self.file_schema, &self.predicate)
-                {
+                if let (Some(file_schema), Some(predicate)) = (
+                    &self.table_schema.as_ref().map(|ts| ts.file_schema()),
+                    &self.predicate,
+                ) {
                     let predicate_creation_errors = Count::new();
                     if let (Some(pruning_predicate), _) = build_pruning_predicates(
                         Some(predicate),
@@ -708,7 +702,12 @@ impl FileSource for ParquetSource {
         filters: Vec<Arc<dyn PhysicalExpr>>,
         config: &ConfigOptions,
     ) -> datafusion_common::Result<FilterPushdownPropagation<Arc<dyn FileSource>>> {
-        let Some(file_schema) = self.file_schema.clone() else {
+        let Some(table_schema) = self
+            .table_schema
+            .as_ref()
+            .map(|ts| ts.table_schema())
+            .cloned()
+        else {
             return Ok(FilterPushdownPropagation::with_parent_pushdown_result(
                 vec![PushedDown::No; filters.len()],
             ));
@@ -728,19 +727,7 @@ impl FileSource for ParquetSource {
         let filters: Vec<PushedDownPredicate> = filters
             .into_iter()
             .map(|filter| {
-                let schema_with_partition_cols = Arc::new(Schema::new(
-                    file_schema
-                        .fields()
-                        .iter()
-                        .cloned()
-                        .chain(self.table_partition_cols.iter().cloned())
-                        .collect_vec(),
-                ));
-
-                if can_expr_be_pushed_down_with_schemas(
-                    &filter,
-                    &schema_with_partition_cols,
-                ) {
+                if can_expr_be_pushed_down_with_schemas(&filter, &table_schema) {
                     PushedDownPredicate::supported(filter)
                 } else {
                     PushedDownPredicate::unsupported(filter)
diff --git a/datafusion/datasource/src/file.rs b/datafusion/datasource/src/file.rs
@@ -26,7 +26,7 @@ use crate::file_groups::FileGroupPartitioner;
 use crate::file_scan_config::FileScanConfig;
 use crate::file_stream::FileOpener;
 use crate::schema_adapter::SchemaAdapterFactory;
-use arrow::datatypes::SchemaRef;
+use crate::TableSchema;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::{not_impl_err, Result, Statistics};
 use datafusion_physical_expr::{LexOrdering, PhysicalExpr};
@@ -64,7 +64,7 @@ pub trait FileSource: Send + Sync {
     /// Initialize new type with batch size configuration
     fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource>;
     /// Initialize new instance with a new schema
-    fn with_schema(&self, schema: SchemaRef) -> Arc<dyn FileSource>;
+    fn with_schema(&self, schema: TableSchema) -> Arc<dyn FileSource>;
     /// Initialize new instance with projection information
     fn with_projection(&self, config: &FileScanConfig) -> Arc<dyn FileSource>;
     /// Initialize new instance with projected statistics
diff --git a/datafusion/datasource/src/file_scan_config.rs b/datafusion/datasource/src/file_scan_config.rs
@@ -87,6 +87,7 @@ use log::{debug, warn};
 /// # use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
 /// # use datafusion_datasource::file_stream::FileOpener;
 /// # use datafusion_datasource::source::DataSourceExec;
+/// # use datafusion_datasource::table_schema::TableSchema;
 /// # use datafusion_execution::object_store::ObjectStoreUrl;
 /// # use datafusion_physical_plan::ExecutionPlan;
 /// # use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
@@ -107,7 +108,7 @@ use log::{debug, warn};
 /// #  fn create_file_opener(&self, _: Arc<dyn ObjectStore>, _: &FileScanConfig, _: usize) -> Arc<dyn FileOpener> { unimplemented!() }
 /// #  fn as_any(&self) -> &dyn Any { self  }
 /// #  fn with_batch_size(&self, _: usize) -> Arc<dyn FileSource> { unimplemented!() }
-/// #  fn with_schema(&self, _: SchemaRef) -> Arc<dyn FileSource> { Arc::new(self.clone()) as Arc<dyn FileSource> }
+/// #  fn with_schema(&self, _: TableSchema) -> Arc<dyn FileSource> { Arc::new(self.clone()) as Arc<dyn FileSource> }
 /// #  fn with_projection(&self, _: &FileScanConfig) -> Arc<dyn FileSource> { unimplemented!() }
 /// #  fn with_statistics(&self, statistics: Statistics) -> Arc<dyn FileSource> { Arc::new(Self {projected_statistics: Some(statistics), schema_adapter_factory: self.schema_adapter_factory.clone()} ) }
 /// #  fn metrics(&self) -> &ExecutionPlanMetricsSet { unimplemented!() }
@@ -446,14 +447,17 @@ impl FileScanConfigBuilder {
         let statistics =
             statistics.unwrap_or_else(|| Statistics::new_unknown(&file_schema));
 
+        // Create TableSchema from file_schema and table_partition_cols
+        let table_schema =
+            TableSchema::new(Arc::clone(&file_schema), table_partition_cols.clone());
+
         let file_source = file_source
             .with_statistics(statistics.clone())
-            .with_schema(Arc::clone(&file_schema));
+            .with_schema(table_schema.clone());
         let file_compression_type =
             file_compression_type.unwrap_or(FileCompressionType::UNCOMPRESSED);
         let new_lines_in_values = new_lines_in_values.unwrap_or(false);
 
-        // Create TableSchema from file_schema and table_partition_cols
         let table_schema = TableSchema::new(file_schema, table_partition_cols);
 
         FileScanConfig {
diff --git a/datafusion/datasource/src/test_util.rs b/datafusion/datasource/src/test_util.rs
@@ -22,7 +22,8 @@ use crate::{
 
 use std::sync::Arc;
 
-use arrow::datatypes::{Schema, SchemaRef};
+use crate::TableSchema;
+use arrow::datatypes::Schema;
 use datafusion_common::{Result, Statistics};
 use datafusion_physical_expr::{expressions::Column, PhysicalExpr};
 use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
@@ -66,7 +67,7 @@ impl FileSource for MockSource {
         Arc::new(Self { ..self.clone() })
     }
 
-    fn with_schema(&self, _schema: SchemaRef) -> Arc<dyn FileSource> {
+    fn with_schema(&self, _schema: TableSchema) -> Arc<dyn FileSource> {
         Arc::new(Self { ..self.clone() })
     }
 

Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,7 @@ mod tests {`
`54`	`54`	`use datafusion_datasource::source::DataSourceExec;`
`55`	`55`
`56`	`56`	`use datafusion_datasource::file::FileSource;`
`57`		`- use datafusion_datasource::{FileRange, PartitionedFile};`
	`57`	`+ use datafusion_datasource::{FileRange, PartitionedFile, TableSchema};`
`58`	`58`	`use datafusion_datasource_parquet::source::ParquetSource;`
`59`	`59`	`use datafusion_datasource_parquet::{`
`60`	`60`	`DefaultParquetFileReaderFactory, ParquetFileReaderFactory, ParquetFormat,`
`@@ -186,7 +186,7 @@ mod tests {`
`186`	`186`	`source = source.with_bloom_filter_on_read(false);`
`187`	`187`	`}`
`188`	`188`
`189`		`- source.with_schema(Arc::clone(&table_schema))`
	`189`	`+ source.with_schema(TableSchema::new(Arc::clone(&table_schema), vec![]))`
`190`	`190`	`}`
`191`	`191`
`192`	`192`	`fn build_parquet_exec(`
Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,7 @@ use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};`
`32`	`32`	`use datafusion_datasource::schema_adapter::SchemaAdapterFactory;`
`33`	`33`	`use datafusion_datasource::{`
`34`	`34`	`as_file_source, calculate_range, ListingTableUrl, PartitionedFile, RangeCalculation,`
	`35`	`+ TableSchema,`
`35`	`36`	`};`
`36`	`37`	`use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};`
`37`	`38`
`@@ -122,7 +123,7 @@ impl FileSource for JsonSource {`
`122`	`123`	`Arc::new(conf)`
`123`	`124`	`}`
`124`	`125`
`125`		`- fn with_schema(&self, _schema: SchemaRef) -> Arc<dyn FileSource> {`
	`126`	`+ fn with_schema(&self, _schema: TableSchema) -> Arc<dyn FileSource> {`
`126`	`127`	`Arc::new(Self { ..self.clone() })`
`127`	`128`	`}`
`128`	`129`	`fn with_statistics(&self, statistics: Statistics) -> Arc<dyn FileSource> {`
Original file line number	Diff line number	Diff line change
`@@ -463,8 +463,7 @@ impl FileFormat for ParquetFormat {`
`463`	`463`	`metadata_size_hint = Some(metadata);`
`464`	`464`	`}`
`465`	`465`
`466`		`- let mut source = ParquetSource::new(self.options.clone())`
`467`		`- .with_table_partition_cols(conf.table_partition_cols().clone());`
	`466`	`+ let mut source = ParquetSource::new(self.options.clone());`
`468`	`467`
`469`	`468`	`// Use the CachedParquetFileReaderFactory`
`470`	`469`	`let metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache();`