feat: add TableSchemaBuilder and store partition columns as Fields (apache#22496)

## Which issue does this PR close?

- No separate issue. Follows up on apache#22372 (panic fix in
`TableSchema::with_table_partition_cols`) and the API discussion it
spawned, and is informed by apache#22026 (which adds a third column group,
virtual columns, to `TableSchema`).

## Rationale for this change

`TableSchema` has one required input (the file schema) and a growing set
of *optional* column groups: partition columns today, virtual columns in
apache#22026. The current API expresses this awkwardly:

- `new(file_schema, partition_cols)` privileges partition columns with a
positional slot while virtual columns only get a builder method — an
asymmetry that grows with every new column kind.
- `TableSchema` eagerly recomputes and caches the concatenated table
schema on *every* incremental setter call, so
`from_file_schema(s).with_table_partition_cols(p)` rebuilds it twice
(three times once virtual columns are added). This is exactly why
`new()`'s docs told callers to avoid the builder-style chain.
- The setter mutated an inner `Arc<Vec<FieldRef>>` in place, which is
what caused the shared-`Arc` panic fixed in apache#22372.

A dedicated builder addresses all three, and mirrors the existing
`FileScanConfigBuilder` (the type that *owns* a `TableSchema`).

## What changes are included in this PR?

- **`TableSchemaBuilder`**: `new(file_schema)` →
`.with_table_partition_cols(impl Into<Fields>)` → `.build()`. The
concatenated table schema is computed exactly **once**, in `build()`.
The setter takes `impl Into<Fields>`, so an existing schema's `Fields`
is accepted zero-copy.
- **Partition columns are now stored as `arrow::datatypes::Fields`** (an
immutable `Arc<[FieldRef]>`) instead of `Arc<Vec<FieldRef>>`: one fewer
indirection, shareable zero-copy, and — being immutable — the
shared-`Arc` mutation panic is structurally impossible.
- **`TableSchema::table_partition_cols()` and the delegating
`FileScanConfig::table_partition_cols()` now return `&Fields`.**
`Fields` derefs to `&[FieldRef]`, so iteration/indexing/`len`/`is_empty`
are unchanged; only the arrow `FileFormat` path needed `.to_vec()`.
- **`TableSchema::with_table_partition_cols` is deprecated** in favor of
the builder. It now **replaces** rather than appends. (Note: `main`
currently *appends* here — the replace change in apache#22372 was not captured
by that PR's squash merge — so this also restores the intended replace
semantics.)
- `new` / `from_file_schema` are kept as conveniences that route through
the builder.
- Documented in the 54.0.0 upgrade guide.

This intentionally leaves virtual columns out; apache#22026 should extend the
builder with `with_virtual_columns` once it lands.

## Are these changes tested?

Yes. New unit tests cover building with partition columns,
replace-on-repeat, zero-copy `Fields` input, and the deprecated setter's
behavior; existing `TableSchema` / `FileScanConfig` tests and doctests
pass. `cargo clippy --all-targets -- -D warnings` is clean across the
datasource/proto/arrow/parquet/catalog-listing crates.

## Are there any user-facing changes?

Yes — please apply the `api change` label:

- `TableSchema::table_partition_cols()` /
`FileScanConfig::table_partition_cols()` return `&Fields` instead of
`&Vec<FieldRef>` (source-compatible for most uses via `Deref`).
- `TableSchema::with_table_partition_cols` is deprecated (use the
builder) and now replaces rather than appends.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Author: 3 people

datafusion/catalog-listing/src/table.rs

@@ -32,7 +32,7 @@ use datafusion_datasource::file_sink_config::{FileOutputMode, FileSinkConfig};
 #[expect(deprecated)]
 use datafusion_datasource::schema_adapter::SchemaAdapterFactory;
 use datafusion_datasource::{
-    ListingTableUrl, PartitionedFile, TableSchema, compute_all_files_statistics,
+    ListingTableUrl, PartitionedFile, TableSchemaBuilder, compute_all_files_statistics,
 };
 use datafusion_execution::cache::TableScopedPath;
 use datafusion_execution::cache::cache_manager::FileStatisticsCache;
@@ -321,14 +321,15 @@ impl ListingTable {
 
     /// Creates a file source for this table
     fn create_file_source(&self) -> Arc<dyn FileSource> {
-        let table_schema = TableSchema::new(
-            Arc::clone(&self.file_schema),
-            self.options
-                .table_partition_cols
-                .iter()
-                .map(|(col, field)| Arc::new(Field::new(col, field.clone(), false)))
-                .collect(),
-        );
+        let table_schema = TableSchemaBuilder::from(&self.file_schema)
+            .with_table_partition_cols(
+                self.options
+                    .table_partition_cols
+                    .iter()
+                    .map(|(col, field)| Arc::new(Field::new(col, field.clone(), false)))
+                    .collect::<Vec<_>>(),
+            )
+            .build();
 
         self.options.format.file_source(table_schema)
     }

datafusion/core/src/datasource/file_format/mod.rs

@@ -67,7 +67,7 @@ pub(crate) mod test_util {
                 .await?
         };
 
-        let table_schema = TableSchema::new(file_schema.clone(), vec![]);
+        let table_schema = TableSchema::from(&file_schema);
 
         let statistics = format
             .infer_stats(state, &store, file_schema.clone(), &meta)

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

@@ -34,7 +34,7 @@ mod tests {
     use datafusion_common::{Result, ScalarValue, test_util};
     use datafusion_datasource::file_format::FileFormat;
     use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
-    use datafusion_datasource::{PartitionedFile, TableSchema};
+    use datafusion_datasource::{PartitionedFile, TableSchemaBuilder};
     use datafusion_datasource_avro::AvroFormat;
     use datafusion_datasource_avro::source::AvroSource;
     use datafusion_execution::object_store::ObjectStoreUrl;
@@ -223,10 +223,13 @@ mod tests {
         partitioned_file.partition_values = vec![ScalarValue::from("2021-10-26")];
 
         let projection = Some(vec![0, 1, file_schema.fields().len(), 2]);
-        let table_schema = TableSchema::new(
-            file_schema.clone(),
-            vec![Arc::new(Field::new("date", DataType::Utf8, false))],
-        );
+        let table_schema = TableSchemaBuilder::from(file_schema)
+            .with_table_partition_cols(vec![Arc::new(Field::new(
+                "date",
+                DataType::Utf8,
+                false,
+            ))])
+            .build();
         let source = Arc::new(AvroSource::new(table_schema.clone()));
         let conf = FileScanConfigBuilder::new(object_store_url, source)
             // select specific columns of the files as well as the partitioning

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

@@ -122,7 +122,7 @@ mod tests {
             quote: b'"',
             ..Default::default()
         };
-        let table_schema = TableSchema::from_file_schema(Arc::clone(&file_schema));
+        let table_schema = TableSchema::from(&file_schema);
         let source =
             Arc::new(CsvSource::new(table_schema.clone()).with_csv_options(options));
         let config =
@@ -194,7 +194,7 @@ mod tests {
             quote: b'"',
             ..Default::default()
         };
-        let table_schema = TableSchema::from_file_schema(Arc::clone(&file_schema));
+        let table_schema = TableSchema::from(&file_schema);
         let source =
             Arc::new(CsvSource::new(table_schema.clone()).with_csv_options(options));
         let config =
@@ -265,7 +265,7 @@ mod tests {
             quote: b'"',
             ..Default::default()
         };
-        let table_schema = TableSchema::from_file_schema(Arc::clone(&file_schema));
+        let table_schema = TableSchema::from(&file_schema);
         let source =
             Arc::new(CsvSource::new(table_schema.clone()).with_csv_options(options));
         let config =
@@ -335,7 +335,7 @@ mod tests {
             quote: b'"',
             ..Default::default()
         };
-        let table_schema = TableSchema::from_file_schema(Arc::clone(&file_schema));
+        let table_schema = TableSchema::from(&file_schema);
         let source =
             Arc::new(CsvSource::new(table_schema.clone()).with_csv_options(options));
         let config =
@@ -371,7 +371,7 @@ mod tests {
         file_compression_type: FileCompressionType,
     ) -> Result<()> {
         use datafusion_common::ScalarValue;
-        use datafusion_datasource::TableSchema;
+        use datafusion_datasource::TableSchemaBuilder;
 
         let session_ctx = SessionContext::new();
         let task_ctx = session_ctx.task_ctx();
@@ -400,10 +400,13 @@ mod tests {
             quote: b'"',
             ..Default::default()
         };
-        let table_schema = TableSchema::new(
-            Arc::clone(&file_schema),
-            vec![Arc::new(Field::new("date", DataType::Utf8, false))],
-        );
+        let table_schema = TableSchemaBuilder::from(&file_schema)
+            .with_table_partition_cols(vec![Arc::new(Field::new(
+                "date",
+                DataType::Utf8,
+                false,
+            ))])
+            .build();
         let source =
             Arc::new(CsvSource::new(table_schema.clone()).with_csv_options(options));
         let config =
@@ -508,7 +511,7 @@ mod tests {
             quote: b'"',
             ..Default::default()
         };
-        let table_schema = TableSchema::from_file_schema(Arc::clone(&file_schema));
+        let table_schema = TableSchema::from(&file_schema);
         let source =
             Arc::new(CsvSource::new(table_schema.clone()).with_csv_options(options));
         let config =

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::{PartitionedFile, TableSchema};
+    use datafusion_datasource::{PartitionedFile, TableSchemaBuilder};
     use datafusion_datasource_parquet::source::ParquetSource;
     use datafusion_datasource_parquet::{
         DefaultParquetFileReaderFactory, ParquetFileReaderFactory, ParquetFormat,
@@ -1642,9 +1642,8 @@ mod tests {
             ),
         ]);
 
-        let table_schema = TableSchema::new(
-            Arc::clone(&schema),
-            vec![
+        let table_schema = TableSchemaBuilder::from(&schema)
+            .with_table_partition_cols(vec![
                 Arc::new(Field::new("year", DataType::Utf8, false)),
                 Arc::new(Field::new("month", DataType::UInt8, false)),
                 Arc::new(Field::new(
@@ -1655,8 +1654,8 @@ mod tests {
                     ),
                     false,
                 )),
-            ],
-        );
+            ])
+            .build();
         let source = Arc::new(ParquetSource::new(table_schema.clone()));
         let config = FileScanConfigBuilder::new(object_store_url, source)
             .with_file(partitioned_file)

datafusion/core/src/test/mod.rs

@@ -103,7 +103,7 @@ pub fn scan_partitioned_csv(
         quote: b'"',
         ..Default::default()
     };
-    let table_schema = TableSchema::from_file_schema(schema);
+    let table_schema = TableSchema::from(schema);
     let source = Arc::new(CsvSource::new(table_schema.clone()).with_csv_options(options));
     let config =
         FileScanConfigBuilder::from(partitioned_csv_config(file_groups, source)?)

datafusion/core/tests/physical_optimizer/projection_pushdown.rs

@@ -25,7 +25,7 @@ use datafusion::datasource::physical_plan::CsvSource;
 use datafusion::datasource::source::DataSourceExec;
 use datafusion_common::config::{ConfigOptions, CsvOptions};
 use datafusion_common::{JoinSide, JoinType, NullEquality, Result, ScalarValue};
-use datafusion_datasource::TableSchema;
+use datafusion_datasource::TableSchemaBuilder;
 use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
 use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion_execution::{SendableRecordBatchStream, TaskContext};
@@ -1574,10 +1574,13 @@ fn partitioned_data_source() -> Arc<DataSourceExec> {
         quote: b'"',
         ..Default::default()
     };
-    let table_schema = TableSchema::new(
-        Arc::clone(&file_schema),
-        vec![Arc::new(Field::new("partition_col", DataType::Utf8, true))],
-    );
+    let table_schema = TableSchemaBuilder::from(&file_schema)
+        .with_table_partition_cols(vec![Arc::new(Field::new(
+            "partition_col",
+            DataType::Utf8,
+            true,
+        ))])
+        .build();
     let config = FileScanConfigBuilder::new(
         ObjectStoreUrl::parse("test:///").unwrap(),
         Arc::new(CsvSource::new(table_schema).with_csv_options(options)),

datafusion/core/tests/physical_optimizer/pushdown_utils.rs

@@ -111,7 +111,7 @@ pub struct TestSource {
 
 impl TestSource {
     pub fn new(schema: SchemaRef, support: bool, batches: Vec<RecordBatch>) -> Self {
-        let table_schema = datafusion_datasource::TableSchema::new(schema, vec![]);
+        let table_schema = datafusion_datasource::TableSchema::from(schema);
         Self {
             support,
             metrics: ExecutionPlanMetricsSet::new(),

datafusion/datasource-arrow/src/file_format.rs

@@ -37,14 +37,14 @@ use datafusion_common::{
     internal_datafusion_err, not_impl_err,
 };
 use datafusion_common_runtime::{JoinSet, SpawnedTask};
-use datafusion_datasource::TableSchema;
 use datafusion_datasource::display::FileGroupDisplay;
 use datafusion_datasource::file::FileSource;
 use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
 use datafusion_datasource::sink::{DataSink, DataSinkExec};
 use datafusion_datasource::write::{
     ObjectWriterBuilder, SharedBuffer, get_writer_schema,
 };
+use datafusion_datasource::{TableSchema, TableSchemaBuilder};
 use datafusion_execution::{SendableRecordBatchStream, TaskContext};
 use datafusion_expr::dml::InsertOp;
 use datafusion_physical_expr_common::sort_expr::LexRequirement;
@@ -197,10 +197,9 @@ impl FileFormat for ArrowFormat {
             .object_meta
             .location;
 
-        let table_schema = TableSchema::new(
-            Arc::clone(conf.file_schema()),
-            conf.table_partition_cols().clone(),
-        );
+        let table_schema = TableSchemaBuilder::from(conf.file_schema())
+            .with_table_partition_cols(conf.table_partition_cols().clone())
+            .build();
 
         let mut source: Arc<dyn FileSource> =
             match is_object_in_arrow_ipc_file_format(object_store, object_location).await

datafusion/datasource-parquet/src/opener/mod.rs

@@ -1415,7 +1415,7 @@ mod test {
         stats::Precision,
     };
     use datafusion_datasource::morsel::{Morsel, Morselizer};
-    use datafusion_datasource::{PartitionedFile, TableSchema};
+    use datafusion_datasource::{PartitionedFile, TableSchema, TableSchemaBuilder};
     use datafusion_expr::{col, lit};
     use datafusion_physical_expr::{
         PhysicalExpr,
@@ -1495,7 +1495,7 @@ mod test {
 
         /// Create a simple table schema from a file schema (for files without partition columns).
         fn with_schema(mut self, file_schema: SchemaRef) -> Self {
-            self.table_schema = Some(TableSchema::from_file_schema(file_schema));
+            self.table_schema = Some(TableSchema::from(file_schema));
             self
         }
 
@@ -1882,10 +1882,13 @@ mod test {
             Field::new("a", DataType::Int32, false),
         ]));
 
-        let table_schema_for_opener = TableSchema::new(
-            file_schema.clone(),
-            vec![Arc::new(Field::new("part", DataType::Int32, false))],
-        );
+        let table_schema_for_opener = TableSchemaBuilder::from(&file_schema)
+            .with_table_partition_cols(vec![Arc::new(Field::new(
+                "part",
+                DataType::Int32,
+                false,
+            ))])
+            .build();
         let make_opener = |predicate| {
             ParquetMorselizerBuilder::new()
                 .with_store(Arc::clone(&store))
@@ -1951,10 +1954,13 @@ mod test {
             Field::new("a", DataType::Int32, false),
             Field::new("b", DataType::Float32, true),
         ]));
-        let table_schema_for_opener = TableSchema::new(
-            file_schema.clone(),
-            vec![Arc::new(Field::new("part", DataType::Int32, false))],
-        );
+        let table_schema_for_opener = TableSchemaBuilder::from(&file_schema)
+            .with_table_partition_cols(vec![Arc::new(Field::new(
+                "part",
+                DataType::Int32,
+                false,
+            ))])
+            .build();
         let make_opener = |predicate| {
             ParquetMorselizerBuilder::new()
                 .with_store(Arc::clone(&store))
@@ -2023,10 +2029,13 @@ mod test {
             Field::new("a", DataType::Int32, false),
         ]));
 
-        let table_schema_for_opener = TableSchema::new(
-            file_schema.clone(),
-            vec![Arc::new(Field::new("part", DataType::Int32, false))],
-        );
+        let table_schema_for_opener = TableSchemaBuilder::from(&file_schema)
+            .with_table_partition_cols(vec![Arc::new(Field::new(
+                "part",
+                DataType::Int32,
+                false,
+            ))])
+            .build();
         let make_opener = |predicate| {
             ParquetMorselizerBuilder::new()
                 .with_store(Arc::clone(&store))
@@ -2104,10 +2113,13 @@ mod test {
             Field::new("part", DataType::Int32, false),
         ]));
 
-        let table_schema_for_opener = TableSchema::new(
-            file_schema.clone(),
-            vec![Arc::new(Field::new("part", DataType::Int32, false))],
-        );
+        let table_schema_for_opener = TableSchemaBuilder::from(&file_schema)
+            .with_table_partition_cols(vec![Arc::new(Field::new(
+                "part",
+                DataType::Int32,
+                false,
+            ))])
+            .build();
         let make_opener = |predicate| {
             ParquetMorselizerBuilder::new()
                 .with_store(Arc::clone(&store))