Add TableSchema helper to encapsulate file schema + partition fields (#18178)

Hoping this helps with #14993

---------

Co-authored-by: Andrew Lamb <[email protected]>

Author: adriangb

datafusion/core/src/datasource/physical_plan/csv.rs

@@ -121,7 +121,7 @@ mod tests {
         .with_projection(Some(vec![0, 2, 4]))
         .build();
 
-        assert_eq!(13, config.file_schema.fields().len());
+        assert_eq!(13, config.file_schema().fields().len());
         let csv = DataSourceExec::from_data_source(config);
 
         assert_eq!(3, csv.schema().fields().len());
@@ -185,7 +185,7 @@ mod tests {
         .with_file_compression_type(file_compression_type.to_owned())
         .with_projection(Some(vec![4, 0, 2]))
         .build();
-        assert_eq!(13, config.file_schema.fields().len());
+        assert_eq!(13, config.file_schema().fields().len());
         let csv = DataSourceExec::from_data_source(config);
         assert_eq!(3, csv.schema().fields().len());
 
@@ -250,7 +250,7 @@ mod tests {
         .with_file_compression_type(file_compression_type.to_owned())
         .with_limit(Some(5))
         .build();
-        assert_eq!(13, config.file_schema.fields().len());
+        assert_eq!(13, config.file_schema().fields().len());
         let csv = DataSourceExec::from_data_source(config);
         assert_eq!(13, csv.schema().fields().len());
 
@@ -313,7 +313,7 @@ mod tests {
         .with_file_compression_type(file_compression_type.to_owned())
         .with_limit(Some(5))
         .build();
-        assert_eq!(14, config.file_schema.fields().len());
+        assert_eq!(14, config.file_schema().fields().len());
         let csv = DataSourceExec::from_data_source(config);
         assert_eq!(14, csv.schema().fields().len());
 
@@ -349,38 +349,37 @@ mod tests {
         let filename = "aggregate_test_100.csv";
         let tmp_dir = TempDir::new()?;
 
-        let file_groups = partitioned_file_groups(
+        let mut file_groups = partitioned_file_groups(
             path.as_str(),
             filename,
             1,
             Arc::new(CsvFormat::default()),
             file_compression_type.to_owned(),
             tmp_dir.path(),
         )?;
+        // Add partition columns / values
+        file_groups[0][0].partition_values = vec![ScalarValue::from("2021-10-26")];
+
+        let num_file_schema_fields = file_schema.fields().len();
 
         let source = Arc::new(CsvSource::new(true, b',', b'"'));
-        let mut config = FileScanConfigBuilder::from(partitioned_csv_config(
+        let config = FileScanConfigBuilder::from(partitioned_csv_config(
             file_schema,
             file_groups,
             source,
         ))
         .with_newlines_in_values(false)
         .with_file_compression_type(file_compression_type.to_owned())
-        .build();
-
-        // Add partition columns
-        config.table_partition_cols =
-            vec![Arc::new(Field::new("date", DataType::Utf8, false))];
-        config.file_groups[0][0].partition_values = vec![ScalarValue::from("2021-10-26")];
-
+        .with_table_partition_cols(vec![Field::new("date", DataType::Utf8, false)])
         // We should be able to project on the partition column
         // Which is supposed to be after the file fields
-        config.projection = Some(vec![0, config.file_schema.fields().len()]);
+        .with_projection(Some(vec![0, num_file_schema_fields]))
+        .build();
 
         // we don't have `/date=xx/` in the path but that is ok because
         // partitions are resolved during scan anyway
 
-        assert_eq!(13, config.file_schema.fields().len());
+        assert_eq!(13, config.file_schema().fields().len());
         let csv = DataSourceExec::from_data_source(config);
         assert_eq!(2, csv.schema().fields().len());
 

datafusion/datasource-parquet/src/source.rs

@@ -497,7 +497,7 @@ impl FileSource for ParquetSource {
     ) -> Arc<dyn FileOpener> {
         let projection = base_config
             .file_column_projection_indices()
-            .unwrap_or_else(|| (0..base_config.file_schema.fields().len()).collect());
+            .unwrap_or_else(|| (0..base_config.file_schema().fields().len()).collect());
 
         let (expr_adapter_factory, schema_adapter_factory) = match (
             base_config.expr_adapter_factory.as_ref(),
@@ -566,8 +566,8 @@ impl FileSource for ParquetSource {
                 .expect("Batch size must set before creating ParquetOpener"),
             limit: base_config.limit,
             predicate: self.predicate.clone(),
-            logical_file_schema: Arc::clone(&base_config.file_schema),
-            partition_fields: base_config.table_partition_cols.clone(),
+            logical_file_schema: Arc::clone(base_config.file_schema()),
+            partition_fields: base_config.table_partition_cols().clone(),
             metadata_size_hint: self.metadata_size_hint,
             metrics: self.metrics().clone(),
             parquet_file_reader_factory,

datafusion/datasource/src/file_scan_config.rs

@@ -24,7 +24,7 @@ use crate::schema_adapter::SchemaAdapterFactory;
 use crate::{
     display::FileGroupsDisplay, file::FileSource,
     file_compression_type::FileCompressionType, file_stream::FileStream,
-    source::DataSource, statistics::MinMaxStatistics, PartitionedFile,
+    source::DataSource, statistics::MinMaxStatistics, PartitionedFile, TableSchema,
 };
 use arrow::datatypes::FieldRef;
 use arrow::{
@@ -153,15 +153,11 @@ pub struct FileScanConfig {
     /// [`RuntimeEnv::register_object_store`]: datafusion_execution::runtime_env::RuntimeEnv::register_object_store
     /// [`RuntimeEnv::object_store`]: datafusion_execution::runtime_env::RuntimeEnv::object_store
     pub object_store_url: ObjectStoreUrl,
-    /// Schema before `projection` is applied. It contains the all columns that may
-    /// appear in the files. It does not include table partition columns
-    /// that may be added.
-    /// Note that this is **not** the schema of the physical files.
-    /// This is the schema that the physical file schema will be
-    /// mapped onto, and the schema that the [`DataSourceExec`] will return.
+    /// Schema information including the file schema, table partition columns,
+    /// and the combined table schema.
     ///
     /// [`DataSourceExec`]: crate::source::DataSourceExec
-    pub file_schema: SchemaRef,
+    pub table_schema: TableSchema,
     /// List of files to be processed, grouped into partitions
     ///
     /// Each file must have a schema of `file_schema` or a subset. If
@@ -180,8 +176,6 @@ pub struct FileScanConfig {
     /// The maximum number of records to read from this plan. If `None`,
     /// all records after filtering are returned.
     pub limit: Option<usize>,
-    /// The partitioning columns
-    pub table_partition_cols: Vec<FieldRef>,
     /// All equivalent lexicographical orderings that describe the schema.
     pub output_ordering: Vec<LexOrdering>,
     /// File compression type
@@ -459,13 +453,15 @@ impl FileScanConfigBuilder {
             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 {
             object_store_url,
-            file_schema,
+            table_schema,
             file_source,
             limit,
             projection,
-            table_partition_cols,
             constraints,
             file_groups,
             output_ordering,
@@ -481,7 +477,7 @@ impl From<FileScanConfig> for FileScanConfigBuilder {
     fn from(config: FileScanConfig) -> Self {
         Self {
             object_store_url: config.object_store_url,
-            file_schema: config.file_schema,
+            file_schema: Arc::clone(config.table_schema.file_schema()),
             file_source: Arc::<dyn FileSource>::clone(&config.file_source),
             file_groups: config.file_groups,
             statistics: config.file_source.statistics().ok(),
@@ -490,7 +486,7 @@ impl From<FileScanConfig> for FileScanConfigBuilder {
             new_lines_in_values: Some(config.new_lines_in_values),
             limit: config.limit,
             projection: config.projection,
-            table_partition_cols: config.table_partition_cols,
+            table_partition_cols: config.table_schema.table_partition_cols().clone(),
             constraints: Some(config.constraints),
             batch_size: config.batch_size,
             expr_adapter_factory: config.expr_adapter_factory,
@@ -635,7 +631,7 @@ impl DataSource for FileScanConfig {
                 .expr
                 .as_any()
                 .downcast_ref::<Column>()
-                .map(|expr| expr.index() >= self.file_schema.fields().len())
+                .map(|expr| expr.index() >= self.file_schema().fields().len())
                 .unwrap_or(false)
         });
 
@@ -650,7 +646,7 @@ impl DataSource for FileScanConfig {
                 &file_scan
                     .projection
                     .clone()
-                    .unwrap_or_else(|| (0..self.file_schema.fields().len()).collect()),
+                    .unwrap_or_else(|| (0..self.file_schema().fields().len()).collect()),
             );
 
             Arc::new(
@@ -691,11 +687,21 @@ impl DataSource for FileScanConfig {
 }
 
 impl FileScanConfig {
+    /// Get the file schema (schema of the files without partition columns)
+    pub fn file_schema(&self) -> &SchemaRef {
+        self.table_schema.file_schema()
+    }
+
+    /// Get the table partition columns
+    pub fn table_partition_cols(&self) -> &Vec<FieldRef> {
+        self.table_schema.table_partition_cols()
+    }
+
     fn projection_indices(&self) -> Vec<usize> {
         match &self.projection {
             Some(proj) => proj.clone(),
-            None => (0..self.file_schema.fields().len()
-                + self.table_partition_cols.len())
+            None => (0..self.file_schema().fields().len()
+                + self.table_partition_cols().len())
                 .collect(),
         }
     }
@@ -707,7 +713,7 @@ impl FileScanConfig {
             .projection_indices()
             .into_iter()
             .map(|idx| {
-                if idx < self.file_schema.fields().len() {
+                if idx < self.file_schema().fields().len() {
                     statistics.column_statistics[idx].clone()
                 } else {
                     // TODO provide accurate stat for partition column (#1186)
@@ -729,20 +735,20 @@ impl FileScanConfig {
             .projection_indices()
             .into_iter()
             .map(|idx| {
-                if idx < self.file_schema.fields().len() {
-                    self.file_schema.field(idx).clone()
+                if idx < self.file_schema().fields().len() {
+                    self.file_schema().field(idx).clone()
                 } else {
-                    let partition_idx = idx - self.file_schema.fields().len();
+                    let partition_idx = idx - self.file_schema().fields().len();
                     Arc::unwrap_or_clone(Arc::clone(
-                        &self.table_partition_cols[partition_idx],
+                        &self.table_partition_cols()[partition_idx],
                     ))
                 }
             })
             .collect();
 
         Arc::new(Schema::new_with_metadata(
             table_fields,
-            self.file_schema.metadata().clone(),
+            self.file_schema().metadata().clone(),
         ))
     }
 
@@ -790,9 +796,9 @@ impl FileScanConfig {
 
     /// Project the schema, constraints, and the statistics on the given column indices
     pub fn project(&self) -> (SchemaRef, Constraints, Statistics, Vec<LexOrdering>) {
-        if self.projection.is_none() && self.table_partition_cols.is_empty() {
+        if self.projection.is_none() && self.table_partition_cols().is_empty() {
             return (
-                Arc::clone(&self.file_schema),
+                Arc::clone(self.file_schema()),
                 self.constraints.clone(),
                 self.file_source.statistics().unwrap().clone(),
                 self.output_ordering.clone(),
@@ -811,8 +817,8 @@ impl FileScanConfig {
     pub fn projected_file_column_names(&self) -> Option<Vec<String>> {
         self.projection.as_ref().map(|p| {
             p.iter()
-                .filter(|col_idx| **col_idx < self.file_schema.fields().len())
-                .map(|col_idx| self.file_schema.field(*col_idx).name())
+                .filter(|col_idx| **col_idx < self.file_schema().fields().len())
+                .map(|col_idx| self.file_schema().field(*col_idx).name())
                 .cloned()
                 .collect()
         })
@@ -823,17 +829,17 @@ impl FileScanConfig {
         let fields = self.file_column_projection_indices().map(|indices| {
             indices
                 .iter()
-                .map(|col_idx| self.file_schema.field(*col_idx))
+                .map(|col_idx| self.file_schema().field(*col_idx))
                 .cloned()
                 .collect::<Vec<_>>()
         });
 
         fields.map_or_else(
-            || Arc::clone(&self.file_schema),
+            || Arc::clone(self.file_schema()),
             |f| {
                 Arc::new(Schema::new_with_metadata(
                     f,
-                    self.file_schema.metadata.clone(),
+                    self.file_schema().metadata.clone(),
                 ))
             },
         )
@@ -842,7 +848,7 @@ impl FileScanConfig {
     pub fn file_column_projection_indices(&self) -> Option<Vec<usize>> {
         self.projection.as_ref().map(|p| {
             p.iter()
-                .filter(|col_idx| **col_idx < self.file_schema.fields().len())
+                .filter(|col_idx| **col_idx < self.file_schema().fields().len())
                 .copied()
                 .collect()
         })
@@ -2182,11 +2188,11 @@ mod tests {
 
         // Verify the built config has all the expected values
         assert_eq!(config.object_store_url, object_store_url);
-        assert_eq!(config.file_schema, file_schema);
+        assert_eq!(*config.file_schema(), file_schema);
         assert_eq!(config.limit, Some(1000));
         assert_eq!(config.projection, Some(vec![0, 1]));
-        assert_eq!(config.table_partition_cols.len(), 1);
-        assert_eq!(config.table_partition_cols[0].name(), "date");
+        assert_eq!(config.table_partition_cols().len(), 1);
+        assert_eq!(config.table_partition_cols()[0].name(), "date");
         assert_eq!(config.file_groups.len(), 1);
         assert_eq!(config.file_groups[0].len(), 1);
         assert_eq!(
@@ -2265,10 +2271,10 @@ mod tests {
 
         // Verify default values
         assert_eq!(config.object_store_url, object_store_url);
-        assert_eq!(config.file_schema, file_schema);
+        assert_eq!(*config.file_schema(), file_schema);
         assert_eq!(config.limit, None);
         assert_eq!(config.projection, None);
-        assert!(config.table_partition_cols.is_empty());
+        assert!(config.table_partition_cols().is_empty());
         assert!(config.file_groups.is_empty());
         assert_eq!(
             config.file_compression_type,
@@ -2339,10 +2345,10 @@ mod tests {
         // Verify properties match
         let partition_cols = partition_cols.into_iter().map(Arc::new).collect::<Vec<_>>();
         assert_eq!(new_config.object_store_url, object_store_url);
-        assert_eq!(new_config.file_schema, schema);
+        assert_eq!(*new_config.file_schema(), schema);
         assert_eq!(new_config.projection, Some(vec![0, 2]));
         assert_eq!(new_config.limit, Some(10));
-        assert_eq!(new_config.table_partition_cols, partition_cols);
+        assert_eq!(*new_config.table_partition_cols(), partition_cols);
         assert_eq!(new_config.file_groups.len(), 1);
         assert_eq!(new_config.file_groups[0].len(), 1);
         assert_eq!(

datafusion/datasource/src/file_stream.rs

@@ -80,7 +80,7 @@ impl FileStream {
         let pc_projector = PartitionColumnProjector::new(
             Arc::clone(&projected_schema),
             &config
-                .table_partition_cols
+                .table_partition_cols()
                 .iter()
                 .map(|x| x.name().clone())
                 .collect::<Vec<_>>(),

datafusion/datasource/src/mod.rs

@@ -41,6 +41,7 @@ pub mod schema_adapter;
 pub mod sink;
 pub mod source;
 mod statistics;
+pub mod table_schema;
 
 #[cfg(test)]
 pub mod test_util;
@@ -57,6 +58,7 @@ use datafusion_common::{ScalarValue, Statistics};
 use futures::{Stream, StreamExt};
 use object_store::{path::Path, ObjectMeta};
 use object_store::{GetOptions, GetRange, ObjectStore};
+pub use table_schema::TableSchema;
 // Remove when add_row_stats is remove
 #[allow(deprecated)]
 pub use statistics::add_row_stats;