datafusion-datasource 55.0.0

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Helper struct to manage table schemas with partition columns

use arrow::datatypes::{FieldRef, Fields, SchemaBuilder, SchemaRef};
use std::sync::Arc;

/// The overall schema for potentially partitioned data sources.
///
/// When reading partitioned data (such as Hive-style partitioning), a [`TableSchema`]
/// consists of up to three parts:
/// 1. **File schema**: The schema of the actual data files on disk
/// 2. **Partition columns**: Columns whose values are encoded in the directory structure,
///    but not stored in the files themselves
/// 3. **Virtual columns**: Columns produced by the file reader (e.g. Parquet
///    `row_number`) that are not stored in the files
///
/// The full table schema is composed in that order: file columns, then
/// partition columns, then virtual columns. Consumers that need a different
/// output ordering should use a projection on top of
/// [`TableSchema::table_schema`].
///
/// # Example: Partitioned Table
///
/// Consider a table with the following directory structure:
/// ```text
/// /data/date=2025-10-10/region=us-west/data.parquet
/// /data/date=2025-10-11/region=us-east/data.parquet
/// ```
///
/// In this case:
/// - **File schema**: The schema of `data.parquet` files (e.g., `[user_id, amount]`)
/// - **Partition columns**: `[date, region]` extracted from the directory path
/// - **Table schema**: The full schema combining both (e.g., `[user_id, amount, date, region]`)
///
/// # When to Use
///
/// Use `TableSchema` when:
/// - Reading partitioned data sources (Parquet, CSV, etc. with Hive-style partitioning)
/// - You need to efficiently access different schema representations without reconstructing them
/// - You want to avoid repeatedly concatenating file and partition schemas
///
/// For non-partitioned data or when working with a single schema representation,
/// working directly with Arrow's `Schema` or `SchemaRef` is simpler.
///
/// # Performance
///
/// This struct pre-computes and caches the full table schema, allowing cheap references
/// to any representation without repeated allocations or reconstructions.
#[derive(Debug, Clone)]
pub struct TableSchema {
    /// The schema of the data files themselves, without partition columns.
    ///
    /// For example, if your Parquet files contain `[user_id, amount]`,
    /// this field holds that schema.
    file_schema: SchemaRef,

    /// Columns that are derived from the directory structure (partitioning scheme).
    ///
    /// For Hive-style partitioning like `/date=2025-10-10/region=us-west/`,
    /// this contains the `date` and `region` fields.
    ///
    /// These columns are NOT present in the data files but are appended to each
    /// row during query execution based on the file's location.
    ///
    /// Stored as [`Fields`] (an immutable `Arc<[FieldRef]>`) so that cloning a
    /// `TableSchema` is cheap and the partition columns can be shared zero-copy
    /// with an existing schema.
    table_partition_cols: Fields,

    /// Virtual columns that are generated by the reader rather than read from
    /// the data files or the directory structure.
    ///
    /// For example, a Parquet reader may inject a `row_number` column whose
    /// values are produced per file by the reader. Virtual column fields must
    /// carry an arrow extension type (e.g. `RowNumber`, `RowGroupIndex`) so the
    /// file reader can recognize them.
    ///
    /// Virtual columns are appended at the end of the table schema, after the
    /// file columns and any partition columns (layout: `[file, partition,
    /// virtual]`).
    virtual_columns: Fields,

    /// The complete table schema: file_schema columns, followed by partition
    /// columns, followed by virtual columns.
    ///
    /// This is pre-computed during construction by concatenating the three
    /// parts, so it can be returned as a cheap reference.
    table_schema: SchemaRef,

    /// Schema of file + partition columns, excluding virtual columns.
    ///
    /// Pre-computed during construction so [`Self::schema_without_virtual_columns`]
    /// can return a cheap reference. When there are no virtual columns this
    /// shares the same `Arc` as `table_schema`.
    schema_without_virtual_columns: SchemaRef,
}

impl TableSchema {
    /// Start building a [`TableSchema`] from its (required) file schema.
    ///
    /// Partition columns are optional and added with
    /// [`TableSchemaBuilder::with_table_partition_cols`]; the full table schema
    /// is computed once by [`TableSchemaBuilder::build`]. This is the preferred
    /// way to construct a `TableSchema`.
    ///
    /// # Example
    ///
    /// ```
    /// # use std::sync::Arc;
    /// # use arrow::datatypes::{Schema, Field, DataType};
    /// # use datafusion_datasource::TableSchema;
    /// let file_schema = Arc::new(Schema::new(vec![
    ///     Field::new("user_id", DataType::Int64, false),
    ///     Field::new("amount", DataType::Float64, false),
    /// ]));
    ///
    /// let table_schema = TableSchema::builder(file_schema)
    ///     .with_table_partition_cols(vec![
    ///         Arc::new(Field::new("date", DataType::Utf8, false)),
    ///         Arc::new(Field::new("region", DataType::Utf8, false)),
    ///     ])
    ///     .build();
    ///
    /// // Table schema will have 4 columns: user_id, amount, date, region
    /// assert_eq!(table_schema.table_schema().fields().len(), 4);
    /// ```
    pub fn builder(file_schema: SchemaRef) -> TableSchemaBuilder {
        TableSchemaBuilder::new(file_schema)
    }

    /// Create a new TableSchema from a file schema and partition columns.
    ///
    /// This is a convenience for
    /// `TableSchema::builder(file_schema).with_table_partition_cols(cols).build()`.
    #[deprecated(
        since = "55.0.0",
        note = "use TableSchema::builder(file_schema).with_table_partition_cols(cols).build() (or TableSchema::from(file_schema) for no partition columns)"
    )]
    pub fn new(file_schema: SchemaRef, table_partition_cols: Vec<FieldRef>) -> Self {
        TableSchemaBuilder::new(file_schema)
            .with_table_partition_cols(table_partition_cols)
            .build()
    }

    /// Create a new TableSchema with no partition columns.
    #[deprecated(
        since = "55.0.0",
        note = "use TableSchema::from(file_schema) / file_schema.into()"
    )]
    pub fn from_file_schema(file_schema: SchemaRef) -> Self {
        TableSchemaBuilder::new(file_schema).build()
    }

    /// Return a new `TableSchema` with `partition_cols` as its partition columns,
    /// replacing any existing ones. Existing virtual columns are preserved.
    #[deprecated(
        since = "55.0.0",
        note = "use TableSchema::builder(file_schema).with_table_partition_cols(cols).build()"
    )]
    pub fn with_table_partition_cols(self, partition_cols: Vec<FieldRef>) -> Self {
        TableSchemaBuilder::new(self.file_schema)
            .with_table_partition_cols(partition_cols)
            .with_virtual_columns(self.virtual_columns)
            .build()
    }

    /// Get the file schema (without partition columns).
    ///
    /// This is the schema of the actual data files on disk.
    pub fn file_schema(&self) -> &SchemaRef {
        &self.file_schema
    }

    /// Get the table partition columns.
    ///
    /// These are the columns derived from the directory structure that
    /// will be appended to each row during query execution.
    pub fn table_partition_cols(&self) -> &Fields {
        &self.table_partition_cols
    }

    /// Get the virtual columns.
    ///
    /// Virtual columns are produced by the file reader (e.g. Parquet
    /// `row_number`) and are not stored in the data files or derived from
    /// partition paths.
    pub fn virtual_columns(&self) -> &Fields {
        &self.virtual_columns
    }

    /// Get the full table schema (file schema + partition columns + virtual columns).
    ///
    /// This is the complete schema that will be seen by queries. Fields appear
    /// in the order: file columns, partition columns, virtual columns.
    pub fn table_schema(&self) -> &SchemaRef {
        &self.table_schema
    }

    /// Schema of columns that can be referenced by predicates pushed into the
    /// file reader: file columns plus partition columns, excluding virtual
    /// columns.
    ///
    /// Virtual columns are produced by the reader itself (e.g. Parquet
    /// `row_number`) and cannot be referenced inside the reader's row filter,
    /// so predicates that reference them must stay above the scan. Callers
    /// deciding which filters to push down should check against this schema
    /// rather than [`Self::table_schema`].
    ///
    /// When there are no virtual columns this returns the same schema as
    /// [`Self::table_schema`].
    pub fn schema_without_virtual_columns(&self) -> &SchemaRef {
        &self.schema_without_virtual_columns
    }
}

impl From<SchemaRef> for TableSchema {
    fn from(schema: SchemaRef) -> Self {
        TableSchemaBuilder::new(schema).build()
    }
}

impl From<&SchemaRef> for TableSchema {
    fn from(schema: &SchemaRef) -> Self {
        TableSchemaBuilder::new(Arc::clone(schema)).build()
    }
}

/// Builder for [`TableSchema`].
///
/// The file schema is the only required input; partition columns and virtual
/// columns are optional. Unlike calling [`TableSchema`]'s setters repeatedly,
/// the builder computes the concatenated table schema exactly once, in
/// [`TableSchemaBuilder::build`].
///
/// ```
/// # use std::sync::Arc;
/// # use arrow::datatypes::{Schema, Field, DataType};
/// # use datafusion_datasource::TableSchemaBuilder;
/// # let file_schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
/// let table_schema = TableSchemaBuilder::new(file_schema)
///     .with_table_partition_cols(vec![Arc::new(Field::new("date", DataType::Utf8, false))])
///     .build();
/// assert_eq!(table_schema.table_partition_cols().len(), 1);
/// ```
#[derive(Debug, Clone)]
pub struct TableSchemaBuilder {
    file_schema: SchemaRef,
    table_partition_cols: Fields,
    virtual_columns: Fields,
}

impl TableSchemaBuilder {
    /// Create a builder for a `TableSchema` over the given file schema, with no
    /// partition or virtual columns yet.
    pub fn new(file_schema: SchemaRef) -> Self {
        Self {
            file_schema,
            table_partition_cols: Fields::empty(),
            virtual_columns: Fields::empty(),
        }
    }

    /// Set the partition columns, replacing any previously set.
    ///
    /// Accepts anything convertible into [`Fields`] (e.g. `Vec<FieldRef>` or an
    /// existing schema's `Fields`, which is shared zero-copy).
    pub fn with_table_partition_cols(
        mut self,
        table_partition_cols: impl Into<Fields>,
    ) -> Self {
        self.table_partition_cols = table_partition_cols.into();
        self
    }

    /// Set the virtual columns, replacing any previously set.
    ///
    /// Virtual columns are produced by the file reader (e.g. Parquet
    /// `row_number`) and appended at the end of the table schema. Each field
    /// must carry an arrow virtual extension type so the reader can recognize
    /// it.
    ///
    /// Accepts anything convertible into [`Fields`] (e.g. `Vec<FieldRef>`).
    pub fn with_virtual_columns(mut self, virtual_columns: impl Into<Fields>) -> Self {
        self.virtual_columns = virtual_columns.into();
        self
    }

    /// Build the [`TableSchema`], computing the full
    /// `file + partition + virtual` schema once.
    pub fn build(self) -> TableSchema {
        debug_assert!(
            self.virtual_columns.iter().enumerate().all(|(i, v)| {
                let name = v.name();
                !self.file_schema.fields().iter().any(|f| f.name() == name)
                    && !self.table_partition_cols.iter().any(|p| p.name() == name)
                    && !self.virtual_columns[..i].iter().any(|w| w.name() == name)
            }),
            "virtual column name collides with an existing file, partition, or virtual column"
        );

        let mut builder = SchemaBuilder::from(self.file_schema.as_ref());
        builder.extend(self.table_partition_cols.iter().cloned());
        let (table_schema, schema_without_virtual_columns) =
            if self.virtual_columns.is_empty() {
                let schema = Arc::new(builder.finish());
                (Arc::clone(&schema), schema)
            } else {
                let without_virtual = Arc::new(builder.finish());
                let mut builder = SchemaBuilder::from(without_virtual.as_ref());
                builder.extend(self.virtual_columns.iter().cloned());
                (Arc::new(builder.finish()), without_virtual)
            };
        TableSchema {
            file_schema: self.file_schema,
            table_partition_cols: self.table_partition_cols,
            virtual_columns: self.virtual_columns,
            table_schema,
            schema_without_virtual_columns,
        }
    }
}

impl From<SchemaRef> for TableSchemaBuilder {
    fn from(schema: SchemaRef) -> Self {
        TableSchemaBuilder::new(schema)
    }
}

impl From<&SchemaRef> for TableSchemaBuilder {
    fn from(schema: &SchemaRef) -> Self {
        TableSchemaBuilder::new(Arc::clone(schema))
    }
}

#[cfg(test)]
mod tests {
    use super::{TableSchema, TableSchemaBuilder};
    use arrow::datatypes::{DataType, Field, Schema};
    use std::sync::Arc;

    #[test]
    fn test_table_schema_creation() {
        let file_schema = Arc::new(Schema::new(vec![
            Field::new("user_id", DataType::Int64, false),
            Field::new("amount", DataType::Float64, false),
        ]));

        let partition_cols = vec![
            Arc::new(Field::new("date", DataType::Utf8, false)),
            Arc::new(Field::new("region", DataType::Utf8, false)),
        ];

        let table_schema = TableSchema::builder(file_schema.clone())
            .with_table_partition_cols(partition_cols.clone())
            .build();

        // Verify file schema
        assert_eq!(table_schema.file_schema().as_ref(), file_schema.as_ref());

        // Verify partition columns
        assert_eq!(table_schema.table_partition_cols().len(), 2);
        assert_eq!(table_schema.table_partition_cols()[0], partition_cols[0]);
        assert_eq!(table_schema.table_partition_cols()[1], partition_cols[1]);

        // Verify full table schema
        let expected_fields = vec![
            Field::new("user_id", DataType::Int64, false),
            Field::new("amount", DataType::Float64, false),
            Field::new("date", DataType::Utf8, false),
            Field::new("region", DataType::Utf8, false),
        ];
        let expected_schema = Schema::new(expected_fields);
        assert_eq!(table_schema.table_schema().as_ref(), &expected_schema);
    }

    #[test]
    fn test_builder_with_partition_cols() {
        let file_schema =
            Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));

        let table_schema = TableSchemaBuilder::new(Arc::clone(&file_schema))
            .with_table_partition_cols(vec![
                Arc::new(Field::new("country", DataType::Utf8, false)),
                Arc::new(Field::new("year", DataType::Int32, false)),
            ])
            .build();

        // File schema is preserved and the partition columns are appended.
        assert_eq!(table_schema.file_schema().as_ref(), file_schema.as_ref());
        assert_eq!(table_schema.table_partition_cols().len(), 2);
        assert_eq!(table_schema.table_partition_cols()[0].name(), "country");
        assert_eq!(table_schema.table_partition_cols()[1].name(), "year");

        let expected_schema = Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("country", DataType::Utf8, false),
            Field::new("year", DataType::Int32, false),
        ]);
        assert_eq!(table_schema.table_schema().as_ref(), &expected_schema);
    }

    #[test]
    fn test_builder_with_table_partition_cols_replaces() {
        // Calling the setter more than once replaces rather than appends.
        let file_schema =
            Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));

        let table_schema = TableSchemaBuilder::new(file_schema)
            .with_table_partition_cols(vec![Arc::new(Field::new(
                "country",
                DataType::Utf8,
                false,
            ))])
            .with_table_partition_cols(vec![Arc::new(Field::new(
                "city",
                DataType::Utf8,
                false,
            ))])
            .build();

        assert_eq!(table_schema.table_partition_cols().len(), 1);
        assert_eq!(table_schema.table_partition_cols()[0].name(), "city");
    }

    #[test]
    fn test_builder_accepts_fields_zero_copy() {
        // `with_table_partition_cols` accepts an existing schema's `Fields`
        // directly (shared via `Arc`, no `Vec` round-trip).
        let file_schema =
            Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
        let partition_schema =
            Schema::new(vec![Field::new("date", DataType::Utf8, false)]);

        let table_schema = TableSchemaBuilder::new(file_schema)
            .with_table_partition_cols(partition_schema.fields().clone())
            .build();

        assert_eq!(table_schema.table_partition_cols().len(), 1);
        assert_eq!(table_schema.table_partition_cols()[0].name(), "date");
    }

    #[test]
    #[expect(deprecated)]
    fn test_deprecated_with_table_partition_cols_replaces() {
        // The deprecated setter still works and replaces the partition columns.
        // It is safe on a shared clone because partition columns are immutable.
        let file_schema =
            Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
        let original = TableSchema::builder(file_schema)
            .with_table_partition_cols(vec![Arc::new(Field::new(
                "country",
                DataType::Utf8,
                false,
            ))])
            .build();

        let replaced =
            original
                .clone()
                .with_table_partition_cols(vec![Arc::new(Field::new(
                    "city",
                    DataType::Utf8,
                    false,
                ))]);

        assert_eq!(replaced.table_partition_cols().len(), 1);
        assert_eq!(replaced.table_partition_cols()[0].name(), "city");

        // The original is untouched.
        assert_eq!(original.table_partition_cols().len(), 1);
        assert_eq!(original.table_partition_cols()[0].name(), "country");
    }

    #[test]
    fn test_builder_with_virtual_columns_layout() {
        let file_schema = Arc::new(Schema::new(vec![
            Field::new("user_id", DataType::Int64, false),
            Field::new("amount", DataType::Float64, false),
        ]));

        let virtual_cols =
            vec![Arc::new(Field::new("row_number", DataType::Int64, true))];

        let partition_cols = vec![Arc::new(Field::new("date", DataType::Utf8, false))];

        // Apply virtual columns and partition columns in either order on the
        // builder; the resulting table schema should always be
        // [file, partition, virtual].
        let built_virtual_first = TableSchemaBuilder::new(Arc::clone(&file_schema))
            .with_virtual_columns(virtual_cols.clone())
            .with_table_partition_cols(partition_cols.clone())
            .build();

        let built_partition_first = TableSchemaBuilder::new(Arc::clone(&file_schema))
            .with_table_partition_cols(partition_cols.clone())
            .with_virtual_columns(virtual_cols.clone())
            .build();

        let expected = Schema::new(vec![
            Field::new("user_id", DataType::Int64, false),
            Field::new("amount", DataType::Float64, false),
            Field::new("date", DataType::Utf8, false),
            Field::new("row_number", DataType::Int64, true),
        ]);

        for ts in [built_virtual_first, built_partition_first] {
            assert_eq!(ts.table_schema().as_ref(), &expected);
            assert_eq!(ts.virtual_columns().len(), 1);
            assert_eq!(ts.virtual_columns()[0].name(), "row_number");
            assert_eq!(ts.table_partition_cols().len(), 1);
            assert_eq!(ts.file_schema().fields().len(), 2);
        }
    }

    #[test]
    #[should_panic(expected = "virtual column name collides")]
    #[cfg(debug_assertions)]
    fn test_virtual_column_collides_with_file_schema_panics_in_debug() {
        let file_schema = Arc::new(Schema::new(vec![Field::new(
            "row_number",
            DataType::Int64,
            false,
        )]));
        let _ = TableSchemaBuilder::new(file_schema)
            .with_virtual_columns(vec![Arc::new(Field::new(
                "row_number",
                DataType::Int64,
                true,
            ))])
            .build();
    }

    #[test]
    #[should_panic(expected = "virtual column name collides")]
    #[cfg(debug_assertions)]
    fn test_virtual_column_collides_with_partition_panics_in_debug() {
        let file_schema = Arc::new(Schema::new(vec![Field::new(
            "user_id",
            DataType::Int64,
            false,
        )]));
        let partition_cols =
            vec![Arc::new(Field::new("row_number", DataType::Utf8, false))];
        let _ = TableSchemaBuilder::new(file_schema)
            .with_table_partition_cols(partition_cols)
            .with_virtual_columns(vec![Arc::new(Field::new(
                "row_number",
                DataType::Int64,
                true,
            ))])
            .build();
    }

    #[test]
    #[should_panic(expected = "virtual column name collides")]
    #[cfg(debug_assertions)]
    fn test_duplicate_virtual_columns_panic_in_debug() {
        let file_schema = Arc::new(Schema::new(vec![Field::new(
            "user_id",
            DataType::Int64,
            false,
        )]));
        let _ = TableSchemaBuilder::new(file_schema)
            .with_virtual_columns(vec![
                Arc::new(Field::new("vc", DataType::Int64, true)),
                Arc::new(Field::new("vc", DataType::Int64, true)),
            ])
            .build();
    }

    #[test]
    #[should_panic(expected = "virtual column name collides")]
    #[cfg(debug_assertions)]
    fn test_partition_column_added_after_colliding_virtual_panics_in_debug() {
        // Builder order is irrelevant: collision check runs in build().
        let file_schema = Arc::new(Schema::new(vec![Field::new(
            "user_id",
            DataType::Int64,
            false,
        )]));
        let _ = TableSchemaBuilder::new(file_schema)
            .with_virtual_columns(vec![Arc::new(Field::new(
                "row_number",
                DataType::Int64,
                true,
            ))])
            .with_table_partition_cols(vec![Arc::new(Field::new(
                "row_number",
                DataType::Utf8,
                false,
            ))])
            .build();
    }
}