paimon-datafusion 0.3.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.

//! Paimon table provider for DataFusion.

use std::fmt::Write as _;
use std::sync::Arc;

use async_trait::async_trait;
use datafusion::arrow::datatypes::{
    DataType as ArrowDataType, Field, Schema, SchemaRef as ArrowSchemaRef,
};
use datafusion::catalog::Session;
use datafusion::datasource::sink::DataSinkExec;
use datafusion::datasource::{TableProvider, TableType};
use datafusion::error::{DataFusionError, Result as DFResult};
use datafusion::logical_expr::dml::InsertOp;
use datafusion::logical_expr::{Expr, TableProviderFilterPushDown};
use datafusion::physical_plan::ExecutionPlan;
use paimon::spec::{
    BigIntType, CoreOptions, DataField, DataType, ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME,
};
use paimon::table::Table;

use crate::physical_plan::PaimonDataSink;
use crate::BlobReaderRegistry;

use crate::error::to_datafusion_error;
#[cfg(test)]
use crate::filter_pushdown::build_pushed_predicate;
use crate::filter_pushdown::{analyze_filters, classify_filter_pushdown};
use crate::physical_plan::PaimonTableScan;
use crate::runtime::await_with_runtime;

const PARQUET_FIELD_ID_META_KEY: &str = "PARQUET:field_id";

pub(crate) fn datafusion_read_fields(table: &Table) -> Vec<DataField> {
    let mut fields = table.schema().fields().to_vec();
    if CoreOptions::new(table.schema().options()).data_evolution_enabled() {
        fields.push(DataField::new(
            ROW_ID_FIELD_ID,
            ROW_ID_FIELD_NAME.to_string(),
            DataType::BigInt(BigIntType::with_nullable(true)),
        ));
    }
    fields
}

fn datafusion_arrow_schema(
    fields: &[DataField],
    schema_force_view_types: bool,
) -> DFResult<ArrowSchemaRef> {
    let paimon_schema =
        paimon::arrow::build_target_arrow_schema(fields).map_err(to_datafusion_error)?;
    let fields = paimon_schema
        .fields()
        .iter()
        .map(|field| {
            let mut metadata = field.metadata().clone();
            metadata.remove(PARQUET_FIELD_ID_META_KEY);
            let data_type = match field.data_type() {
                ArrowDataType::Utf8 if schema_force_view_types => ArrowDataType::Utf8View,
                data_type => data_type.clone(),
            };
            Arc::new(
                field
                    .as_ref()
                    .clone()
                    .with_data_type(data_type)
                    .with_metadata(metadata),
            )
        })
        .collect::<Vec<_>>();

    Ok(Arc::new(Schema::new_with_metadata(
        fields,
        paimon_schema.metadata().clone(),
    )))
}

/// Read-only table provider for a Paimon table.
///
/// Supports full table scan, column projection, and predicate pushdown for
/// planning. Partition predicates prune splits eagerly, while supported
/// non-partition data predicates may also be reused by the Parquet read path
/// for row-group pruning and partial decode-time filtering.
///
/// DataFusion still treats pushed filters as inexact because unsupported
/// predicates and non-Parquet reads remain residual filters.
#[derive(Debug, Clone)]
pub struct PaimonTableProvider {
    table: Table,
    schema: ArrowSchemaRef,
    table_definition: Option<String>,
}

impl PaimonTableProvider {
    /// Create a table provider from a Paimon table.
    ///
    /// Loads the table schema and converts it to Arrow for DataFusion.
    pub fn try_new(table: Table) -> DFResult<Self> {
        let table_definition = build_table_definition(&table)?;
        Self::try_new_with_table_definition(table, Some(table_definition))
    }

    fn try_new_with_table_definition(
        table: Table,
        table_definition: Option<String>,
    ) -> DFResult<Self> {
        let fields = datafusion_read_fields(&table);
        let schema = datafusion_arrow_schema(&fields, true)?;
        Ok(Self {
            table,
            schema,
            table_definition,
        })
    }

    pub fn try_new_with_blob_reader_registry(
        table: Table,
        blob_reader_registry: BlobReaderRegistry,
    ) -> DFResult<Self> {
        blob_reader_registry
            .register_if_absent(table.location().to_string(), table.file_io().clone());
        Self::try_new(table)
    }

    pub(crate) fn try_new_with_blob_reader_registry_and_definition(
        table: Table,
        blob_reader_registry: BlobReaderRegistry,
        table_definition: Option<String>,
    ) -> DFResult<Self> {
        blob_reader_registry
            .register_if_absent(table.location().to_string(), table.file_io().clone());
        Self::try_new_with_table_definition(table, table_definition)
    }

    pub(crate) fn with_schema_force_view_types(
        mut self,
        schema_force_view_types: bool,
    ) -> DFResult<Self> {
        if schema_force_view_types {
            return Ok(self);
        }
        let fields = datafusion_read_fields(&self.table);
        self.schema = datafusion_arrow_schema(&fields, schema_force_view_types)?;
        Ok(self)
    }

    pub fn table(&self) -> &Table {
        &self.table
    }
}

/// Build a `CREATE TABLE` DDL string for a Paimon table.
///
/// Mirrors the syntax accepted by `SQLContext::handle_create_table`:
/// `CREATE TABLE <db>.<table> (<col> <type>, ..., PRIMARY KEY (...)) [PARTITIONED BY (...)] [WITH ('k'='v', ...)]`.
pub(crate) fn build_table_definition(table: &Table) -> DFResult<String> {
    let identifier = table.identifier();
    let schema = table.schema();
    let mut ddl = String::new();
    let _ = write!(
        ddl,
        "CREATE TABLE {}.{} (",
        quote_identifier(identifier.database()),
        quote_identifier(identifier.object())
    );

    for (i, field) in schema.fields().iter().enumerate() {
        if i > 0 {
            ddl.push_str(", ");
        }
        // `NOT NULL` is a column constraint; render it here at the column
        // level rather than inside nested type arguments (see `data_type_to_sql`).
        let ty = data_type_to_sql(field.data_type())?;
        if field.data_type().is_nullable() {
            let _ = write!(ddl, "{} {}", quote_identifier(field.name()), ty);
        } else {
            let _ = write!(ddl, "{} {} NOT NULL", quote_identifier(field.name()), ty);
        }
    }

    let pks = schema.primary_keys();
    if !pks.is_empty() {
        ddl.push_str(", PRIMARY KEY (");
        for (i, pk) in pks.iter().enumerate() {
            if i > 0 {
                ddl.push_str(", ");
            }
            let _ = write!(ddl, "{}", quote_identifier(pk));
        }
        ddl.push(')');
    }
    ddl.push(')');

    let partition_keys = schema.partition_keys();
    if !partition_keys.is_empty() {
        ddl.push_str(" PARTITIONED BY (");
        for (i, pk) in partition_keys.iter().enumerate() {
            if i > 0 {
                ddl.push_str(", ");
            }
            let _ = write!(ddl, "{}", quote_identifier(pk));
        }
        ddl.push(')');
    }

    let mut options: Vec<_> = schema.options().iter().collect();
    options.sort_by_key(|(left, _)| *left);
    if !options.is_empty() {
        ddl.push_str(" WITH (");
        for (i, (k, v)) in options.iter().enumerate() {
            if i > 0 {
                ddl.push_str(", ");
            }
            let _ = write!(
                ddl,
                "{} = {}",
                quote_string_literal(k),
                quote_string_literal(v)
            );
        }
        ddl.push(')');
    }

    Ok(ddl)
}

fn quote_identifier(identifier: &str) -> String {
    format!("\"{}\"", identifier.replace('"', "\"\""))
}

fn quote_string_literal(text: &str) -> String {
    format!("'{}'", text.replace('\'', "''"))
}

/// Render a Paimon [`DataType`] as a SQL type string matching the syntax
/// accepted by paimon-rust's `CREATE TABLE` parser.
///
/// `NOT NULL` is a column constraint, not a type modifier — it is only valid
/// at the top of a column definition, not nested inside `MAP`, `ARRAY`, or
/// `STRUCT` arguments. Callers that render a column should append `NOT NULL`
/// themselves when the field is non-nullable; recursive calls below must not.
pub(crate) fn data_type_to_sql(data_type: &DataType) -> DFResult<String> {
    match data_type {
        DataType::Boolean(_) => Ok("BOOLEAN".to_string()),
        DataType::TinyInt(_) => Ok("TINYINT".to_string()),
        DataType::SmallInt(_) => Ok("SMALLINT".to_string()),
        DataType::Int(_) => Ok("INT".to_string()),
        DataType::BigInt(_) => Ok("BIGINT".to_string()),
        DataType::Decimal(t) => Ok(format!("DECIMAL({}, {})", t.precision(), t.scale())),
        DataType::Double(_) => Ok("DOUBLE".to_string()),
        DataType::Float(_) => Ok("FLOAT".to_string()),
        DataType::Binary(t) => Ok(format!("BINARY({})", t.length())),
        DataType::VarBinary(t) => Ok(format!("VARBINARY({})", t.length())),
        DataType::Blob(_) => Ok("BLOB".to_string()),
        DataType::Char(t) => Ok(format!("CHAR({})", t.length())),
        DataType::VarChar(t) => Ok(format!("VARCHAR({})", t.length())),
        DataType::Date(_) => Ok("DATE".to_string()),
        DataType::Time(_) => Err(unsupported_show_create_table_type("TIME")),
        DataType::Timestamp(t) => Ok(format!("TIMESTAMP({})", t.precision())),
        DataType::Variant(_) => Ok("VARIANT".to_string()),
        DataType::LocalZonedTimestamp(t) => {
            Ok(format!("TIMESTAMP({}) WITH TIME ZONE", t.precision()))
        }
        DataType::Array(t) => Ok(format!("ARRAY<{}>", data_type_to_sql(t.element_type())?)),
        DataType::Map(t) => Ok(format!(
            "MAP({}, {})",
            data_type_to_sql(t.key_type())?,
            data_type_to_sql(t.value_type())?
        )),
        DataType::Multiset(_) => Err(unsupported_show_create_table_type("MULTISET")),
        DataType::Row(t) => {
            let inner: Vec<String> = t
                .fields()
                .iter()
                .map(|f| {
                    let ty = data_type_to_sql(f.data_type())?;
                    if f.name().is_empty() {
                        Ok(ty)
                    } else {
                        Ok(format!("{} {}", quote_identifier(f.name()), ty))
                    }
                })
                .collect::<DFResult<_>>()?;
            Ok(format!("STRUCT<{}>", inner.join(", ")))
        }
        DataType::Vector(_) => Err(unsupported_show_create_table_type("VECTOR")),
    }
}

fn unsupported_show_create_table_type(type_name: &str) -> DataFusionError {
    DataFusionError::NotImplemented(format!(
        "SHOW CREATE TABLE does not support {type_name} columns because paimon-rust cannot round-trip this type in CREATE TABLE"
    ))
}

/// Distribute `items` into `num_buckets` groups using round-robin assignment.
pub(crate) fn bucket_round_robin<T>(items: Vec<T>, num_buckets: usize) -> Vec<Vec<T>> {
    let mut buckets: Vec<Vec<T>> = (0..num_buckets).map(|_| Vec::new()).collect();
    for (index, item) in items.into_iter().enumerate() {
        buckets[index % num_buckets].push(item);
    }
    buckets
}

/// Build parameters for [`PaimonTableScan`].
pub(crate) struct PaimonScanBuilder<'a> {
    pub(crate) table: &'a Table,
    pub(crate) schema: &'a ArrowSchemaRef,
    pub(crate) plan: &'a paimon::table::Plan,
    pub(crate) scan_trace: Option<paimon::table::ScanTrace>,
    pub(crate) projection: Option<&'a Vec<usize>>,
    pub(crate) pushed_predicate: Option<paimon::spec::Predicate>,
    pub(crate) limit: Option<usize>,
    pub(crate) target_partitions: usize,
    pub(crate) filter_exact: bool,
    /// Column-name case sensitivity, carried into the physical scan so execute()
    /// resolves names the same way planning did.
    pub(crate) case_sensitive: bool,
}

impl PaimonScanBuilder<'_> {
    /// Build a [`PaimonTableScan`] from the configured parameters.
    pub(crate) fn build(self) -> DFResult<Arc<dyn ExecutionPlan>> {
        let read_fields = datafusion_read_fields(self.table);
        self.build_with_read_fields(read_fields)
    }

    /// Build using an internal read schema which may contain hidden system fields.
    pub(crate) fn build_with_read_fields(
        self,
        read_fields: Vec<DataField>,
    ) -> DFResult<Arc<dyn ExecutionPlan>> {
        let (projected_schema, read_type) = if let Some(indices) = self.projection {
            let fields: Vec<Field> = indices
                .iter()
                .map(|&i| self.schema.field(i).clone())
                .collect();
            let read_type = indices
                .iter()
                .map(|&i| read_fields[i].clone())
                .collect::<Vec<_>>();
            (Arc::new(Schema::new(fields)), read_type)
        } else {
            (self.schema.clone(), read_fields)
        };

        let splits = self.plan.splits().to_vec();
        let planned_partitions: Vec<Arc<[_]>> = if splits.is_empty() {
            vec![Arc::from(Vec::new())]
        } else {
            let num_partitions = splits.len().min(self.target_partitions.max(1));
            bucket_round_robin(splits, num_partitions)
                .into_iter()
                .map(Arc::from)
                .collect()
        };

        Ok(Arc::new(PaimonTableScan::new(
            projected_schema,
            self.table.clone(),
            read_type,
            self.pushed_predicate,
            planned_partitions,
            self.limit,
            self.filter_exact,
            self.scan_trace,
            None,
            self.case_sensitive,
        )))
    }
}

#[async_trait]
impl TableProvider for PaimonTableProvider {
    fn schema(&self) -> ArrowSchemaRef {
        self.schema.clone()
    }

    fn table_type(&self) -> TableType {
        TableType::Base
    }

    fn get_table_definition(&self) -> Option<&str> {
        self.table_definition.as_deref()
    }

    async fn scan(
        &self,
        state: &dyn Session,
        projection: Option<&Vec<usize>>,
        filters: &[Expr],
        limit: Option<usize>,
    ) -> DFResult<Arc<dyn ExecutionPlan>> {
        // Column-name matching is case-sensitive on the DataFusion path.
        //
        // DataFusion resolves projection/filter columns against the provider
        // schema (`schema()`, which exposes the original field casing) during
        // logical planning, *before* `scan` is called. `enable_ident_normalization`
        // only lowercases unquoted identifiers at parse time; it does not make
        // schema resolution case-insensitive. So a genuine case mismatch
        // (`SELECT name` against a `Name` field) already fails at planning and
        // never reaches this code — see the negative test in `tests/read_tables.rs`.
        // Case-insensitive column matching is therefore offered only through the
        // direct ReadBuilder API (core / C / Python), not via SQL.
        let case_sensitive = true;
        // Plan splits eagerly so we know partition count upfront.
        let filter_analysis =
            analyze_filters(filters, self.table.schema().fields(), case_sensitive);
        let mut read_builder = self.table.new_read_builder();
        read_builder.with_case_sensitive(case_sensitive);
        if let Some(indices) = projection {
            let read_fields = datafusion_read_fields(&self.table);
            let read_type = indices
                .iter()
                .map(|&i| read_fields[i].clone())
                .collect::<Vec<_>>();
            read_builder.with_read_type(read_type);
        }
        if let Some(filter) = filter_analysis.pushed_predicate.clone() {
            read_builder.with_filter(filter);
        }
        let pushed_limit = limit.filter(|_| !filter_analysis.requires_residual);
        if let Some(limit) = pushed_limit {
            read_builder.with_limit(limit);
        }
        let scan = read_builder.new_scan();
        // DataFusion's Python FFI may poll `TableProvider::scan()` without an active
        // Tokio runtime. `scan.plan()` can reach OpenDAL/Tokio filesystem calls while
        // reading Paimon metadata, so we must provide a runtime here instead of
        // assuming the caller already entered one.
        let (plan, scan_trace) = await_with_runtime(scan.plan_with_trace())
            .await
            .map_err(to_datafusion_error)?;

        let target = state.config_options().execution.target_partitions;
        let filter_exact = !filter_analysis.requires_residual
            && filter_analysis
                .pushed_predicate
                .as_ref()
                .is_none_or(|p| read_builder.is_exact_filter_pushdown(p));
        PaimonScanBuilder {
            table: &self.table,
            schema: &self.schema,
            plan: &plan,
            scan_trace: Some(scan_trace),
            projection,
            pushed_predicate: filter_analysis.pushed_predicate,
            limit: pushed_limit,
            target_partitions: target,
            filter_exact,
            case_sensitive,
        }
        .build()
    }

    async fn insert_into(
        &self,
        _state: &dyn Session,
        input: Arc<dyn ExecutionPlan>,
        insert_op: InsertOp,
    ) -> DFResult<Arc<dyn ExecutionPlan>> {
        if self.table.is_branch_reference() {
            return Err(datafusion::error::DataFusionError::NotImplemented(format!(
                "Writing to Paimon branch '{}' is not supported",
                self.table.branch()
            )));
        }
        let overwrite = match insert_op {
            InsertOp::Append => false,
            InsertOp::Overwrite => true,
            other => {
                return Err(datafusion::error::DataFusionError::NotImplemented(format!(
                    "{other} is not supported for Paimon tables"
                )));
            }
        };
        let sink = PaimonDataSink::new(self.table.clone(), self.schema.clone(), overwrite);
        Ok(Arc::new(DataSinkExec::new(input, Arc::new(sink), None)))
    }

    fn supports_filters_pushdown(
        &self,
        filters: &[&Expr],
    ) -> DFResult<Vec<TableProviderFilterPushDown>> {
        let fields = self.table.schema().fields();
        // SQL reads resolve columns case-sensitively (see `scan`), so classify
        // pushdown the same way.
        let case_sensitive = true;
        let read_builder = self.table.new_read_builder();

        Ok(filters
            .iter()
            .map(|filter| {
                classify_filter_pushdown(filter, fields, case_sensitive, |predicate| {
                    read_builder.is_exact_filter_pushdown(predicate)
                })
            })
            .collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::BTreeSet;
    use std::sync::Arc;

    use datafusion::datasource::TableProvider;
    use datafusion::logical_expr::{col, lit, Expr};
    use datafusion::prelude::{SessionConfig, SessionContext};
    use paimon::catalog::Identifier;
    use paimon::spec::{ArrayType, MapType, RowType, VarCharType};
    use paimon::{Catalog, CatalogOptions, DataSplit, FileSystemCatalog, Options};

    use crate::physical_plan::PaimonTableScan;

    #[test]
    fn test_bucket_round_robin_distributes_evenly() {
        let result = bucket_round_robin(vec![0, 1, 2, 3, 4], 3);
        assert_eq!(result, vec![vec![0, 3], vec![1, 4], vec![2]]);
    }

    #[test]
    fn test_bucket_round_robin_fewer_items_than_buckets() {
        let result = bucket_round_robin(vec![10, 20], 2);
        assert_eq!(result, vec![vec![10], vec![20]]);
    }

    #[test]
    fn test_bucket_round_robin_single_bucket() {
        let result = bucket_round_robin(vec![1, 2, 3], 1);
        assert_eq!(result, vec![vec![1, 2, 3]]);
    }

    fn get_test_warehouse() -> String {
        std::env::var("PAIMON_TEST_WAREHOUSE")
            .unwrap_or_else(|_| "/tmp/paimon-warehouse".to_string())
    }

    fn create_catalog() -> FileSystemCatalog {
        let warehouse = get_test_warehouse();
        let mut options = Options::new();
        options.set(CatalogOptions::WAREHOUSE, warehouse);
        FileSystemCatalog::new(options).expect("Failed to create catalog")
    }

    async fn create_provider(table_name: &str) -> PaimonTableProvider {
        let catalog = create_catalog();
        let identifier = Identifier::new("default", table_name);
        let table = catalog
            .get_table(&identifier)
            .await
            .expect("Failed to get table");

        PaimonTableProvider::try_new(table).expect("Failed to create table provider")
    }

    async fn plan_partitions(
        provider: &PaimonTableProvider,
        filters: Vec<Expr>,
        limit: Option<usize>,
    ) -> Vec<Arc<[DataSplit]>> {
        let plan = plan_scan(provider, filters, limit).await;
        let scan = plan
            .downcast_ref::<PaimonTableScan>()
            .expect("Expected PaimonTableScan");

        scan.planned_partitions().to_vec()
    }

    async fn plan_scan(
        provider: &PaimonTableProvider,
        filters: Vec<Expr>,
        limit: Option<usize>,
    ) -> Arc<dyn ExecutionPlan> {
        let config = SessionConfig::new().with_target_partitions(8);
        let ctx = SessionContext::new_with_config(config);
        let state = ctx.state();
        provider
            .scan(&state, None, &filters, limit)
            .await
            .expect("scan() should succeed")
    }

    fn extract_dt_partition_set(planned_partitions: &[Arc<[DataSplit]>]) -> BTreeSet<String> {
        planned_partitions
            .iter()
            .flat_map(|splits| splits.iter())
            .map(|split| {
                split
                    .partition()
                    .get_string(0)
                    .expect("Failed to decode dt")
                    .to_string()
            })
            .collect()
    }

    fn extract_dt_hr_partition_set(
        planned_partitions: &[Arc<[DataSplit]>],
    ) -> BTreeSet<(String, i32)> {
        planned_partitions
            .iter()
            .flat_map(|splits| splits.iter())
            .map(|split| {
                let partition = split.partition();
                (
                    partition
                        .get_string(0)
                        .expect("Failed to decode dt")
                        .to_string(),
                    partition.get_int(1).expect("Failed to decode hr"),
                )
            })
            .collect()
    }

    fn empty_binary_stats_json() -> serde_json::Value {
        let row = paimon::spec::EMPTY_SERIALIZED_ROW.as_slice().to_vec();
        serde_json::json!({
            "_MIN_VALUES": row,
            "_MAX_VALUES": row,
            "_NULL_COUNTS": [],
        })
    }

    fn data_evolution_file(
        file_name: &str,
        file_size: i64,
        row_count: i64,
        first_row_id: i64,
        write_cols: &[&str],
    ) -> paimon::spec::DataFileMeta {
        serde_json::from_value(serde_json::json!({
            "_FILE_NAME": file_name,
            "_FILE_SIZE": file_size,
            "_ROW_COUNT": row_count,
            "_MIN_KEY": [],
            "_MAX_KEY": [],
            "_KEY_STATS": empty_binary_stats_json(),
            "_VALUE_STATS": empty_binary_stats_json(),
            "_MIN_SEQUENCE_NUMBER": 0,
            "_MAX_SEQUENCE_NUMBER": 0,
            "_SCHEMA_ID": 0,
            "_LEVEL": 1,
            "_EXTRA_FILES": [],
            "_CREATION_TIME": null,
            "_DELETE_ROW_COUNT": null,
            "_EMBEDDED_FILE_INDEX": null,
            "_FILE_SOURCE": null,
            "_VALUE_STATS_COLS": null,
            "_FIRST_ROW_ID": first_row_id,
            "_WRITE_COLS": write_cols,
            "_EXTERNAL_PATH": null,
        }))
        .expect("test data file should deserialize")
    }

    fn manifest_file_meta(
        file_name: &str,
        file_size: i64,
        num_added_files: i64,
    ) -> paimon::spec::ManifestFileMeta {
        serde_json::from_value(serde_json::json!({
            "_VERSION": 2,
            "_FILE_NAME": file_name,
            "_FILE_SIZE": file_size,
            "_NUM_ADDED_FILES": num_added_files,
            "_NUM_DELETED_FILES": 0,
            "_PARTITION_STATS": empty_binary_stats_json(),
            "_SCHEMA_ID": 0,
        }))
        .expect("test manifest file meta should deserialize")
    }

    async fn data_evolution_projection_pruning_provider() -> PaimonTableProvider {
        use paimon::io::FileIOBuilder;
        use paimon::spec::{
            CommitKind, DataType, FileKind, IntType, Manifest, ManifestEntry, ManifestList,
            Schema as PaimonSchema, Snapshot, TableSchema,
        };
        use paimon::table::{SnapshotManager, Table};

        let file_io = FileIOBuilder::new("memory").build().unwrap();
        let table_path = format!("memory:/df_de_projection_pruning_{}", uuid::Uuid::new_v4());
        file_io
            .mkdirs(&format!("{table_path}/snapshot/"))
            .await
            .unwrap();
        file_io
            .mkdirs(&format!("{table_path}/manifest/"))
            .await
            .unwrap();

        let schema = PaimonSchema::builder()
            .column("id", DataType::Int(IntType::new()))
            .column("name", DataType::Int(IntType::new()))
            .option("data-evolution.enabled", "true")
            .build()
            .unwrap();
        let table_schema = TableSchema::new(0, &schema);
        let table = Table::new(
            file_io.clone(),
            Identifier::new("default", "df_de_projection_pruning"),
            table_path.clone(),
            table_schema,
            None,
        );

        let partition = paimon::spec::EMPTY_SERIALIZED_ROW.as_slice().to_vec();
        let entries = vec![
            ManifestEntry::new(
                FileKind::Add,
                partition.clone(),
                0,
                1,
                data_evolution_file("id.parquet", 11, 10, 0, &["id"]),
                2,
            ),
            ManifestEntry::new(
                FileKind::Add,
                partition,
                0,
                1,
                data_evolution_file("name.parquet", 13, 10, 0, &["name"]),
                2,
            ),
        ];

        let manifest_name = "manifest-de-projection-0";
        let manifest_path = format!("{table_path}/manifest/{manifest_name}");
        Manifest::write(&file_io, &manifest_path, &entries)
            .await
            .unwrap();
        let manifest_size = file_io
            .new_input(&manifest_path)
            .unwrap()
            .metadata()
            .await
            .unwrap()
            .size;

        let base_list_name = "base-list-de-projection";
        let delta_list_name = "delta-list-de-projection";
        ManifestList::write(
            &file_io,
            &format!("{table_path}/manifest/{base_list_name}"),
            &[manifest_file_meta(
                manifest_name,
                manifest_size as i64,
                entries.len() as i64,
            )],
        )
        .await
        .unwrap();
        ManifestList::write(
            &file_io,
            &format!("{table_path}/manifest/{delta_list_name}"),
            &[],
        )
        .await
        .unwrap();

        let snapshot = Snapshot::builder()
            .version(3)
            .id(1)
            .schema_id(0)
            .base_manifest_list(base_list_name.to_string())
            .delta_manifest_list(delta_list_name.to_string())
            .commit_user("test-user".to_string())
            .commit_identifier(1)
            .commit_kind(CommitKind::APPEND)
            .time_millis(1)
            .total_record_count(Some(10))
            .delta_record_count(Some(10))
            .build();
        let snapshot_manager = SnapshotManager::new(file_io, table_path);
        assert!(snapshot_manager.commit_snapshot(&snapshot).await.unwrap());

        PaimonTableProvider::try_new(table).expect("provider should be created")
    }

    #[tokio::test]
    async fn test_datafusion_schema_hides_paimon_field_ids() {
        let provider = data_evolution_projection_pruning_provider().await;

        for field in provider.schema().fields() {
            assert!(
                !field.metadata().contains_key("PARQUET:field_id"),
                "storage field id leaked through DataFusion schema for {}",
                field.name()
            );
        }
    }

    #[test]
    fn test_datafusion_schema_uses_views_only_for_top_level_strings() {
        let string_type = || DataType::VarChar(VarCharType::string_type());
        let schema = datafusion_arrow_schema(
            &[
                DataField::new(0, "plain".to_string(), string_type()),
                DataField::new(
                    1,
                    "array".to_string(),
                    DataType::Array(ArrayType::new(string_type())),
                ),
                DataField::new(
                    2,
                    "map".to_string(),
                    DataType::Map(MapType::new(string_type(), string_type())),
                ),
                DataField::new(
                    3,
                    "row".to_string(),
                    DataType::Row(RowType::new(vec![DataField::new(
                        4,
                        "nested".to_string(),
                        string_type(),
                    )])),
                ),
            ],
            true,
        )
        .expect("DataFusion schema should be created");

        assert_eq!(schema.field(0).data_type(), &ArrowDataType::Utf8View);

        let ArrowDataType::List(element) = schema.field(1).data_type() else {
            panic!("array field should map to an Arrow List");
        };
        assert_eq!(element.data_type(), &ArrowDataType::Utf8);

        let ArrowDataType::Map(entries, _) = schema.field(2).data_type() else {
            panic!("map field should map to an Arrow Map");
        };
        let ArrowDataType::Struct(map_fields) = entries.data_type() else {
            panic!("map entries should map to an Arrow Struct");
        };
        assert_eq!(map_fields[0].data_type(), &ArrowDataType::Utf8);
        assert_eq!(map_fields[1].data_type(), &ArrowDataType::Utf8);

        let ArrowDataType::Struct(row_fields) = schema.field(3).data_type() else {
            panic!("row field should map to an Arrow Struct");
        };
        assert_eq!(row_fields[0].data_type(), &ArrowDataType::Utf8);
    }

    fn planned_file_names(scan: &PaimonTableScan) -> Vec<String> {
        let mut names = scan
            .planned_partitions()
            .iter()
            .flat_map(|partition| partition.iter())
            .flat_map(|split| split.data_files().iter())
            .map(|file| file.file_name.clone())
            .collect::<Vec<_>>();
        names.sort();
        names
    }

    #[tokio::test]
    async fn test_scan_partition_filter_plans_matching_partition_set() {
        let provider = create_provider("partitioned_log_table").await;
        let planned_partitions =
            plan_partitions(&provider, vec![col("dt").eq(lit("2024-01-01"))], None).await;

        assert_eq!(
            extract_dt_partition_set(&planned_partitions),
            BTreeSet::from(["2024-01-01".to_string()]),
        );
    }

    #[tokio::test]
    async fn test_scan_mixed_and_filter_keeps_partition_pruning() {
        let provider = create_provider("partitioned_log_table").await;
        let planned_partitions = plan_partitions(
            &provider,
            vec![col("dt").eq(lit("2024-01-01")).and(col("id").gt(lit(1)))],
            None,
        )
        .await;

        assert_eq!(
            extract_dt_partition_set(&planned_partitions),
            BTreeSet::from(["2024-01-01".to_string()]),
        );
    }

    #[tokio::test]
    async fn test_scan_multi_partition_filter_plans_exact_partition_set() {
        let provider = create_provider("multi_partitioned_log_table").await;

        let dt_only_partitions =
            plan_partitions(&provider, vec![col("dt").eq(lit("2024-01-01"))], None).await;
        let dt_hr_partitions = plan_partitions(
            &provider,
            vec![col("dt").eq(lit("2024-01-01")).and(col("hr").eq(lit(10)))],
            None,
        )
        .await;

        assert_eq!(
            extract_dt_hr_partition_set(&dt_only_partitions),
            BTreeSet::from([
                ("2024-01-01".to_string(), 10),
                ("2024-01-01".to_string(), 20),
            ]),
        );
        assert_eq!(
            extract_dt_hr_partition_set(&dt_hr_partitions),
            BTreeSet::from([("2024-01-01".to_string(), 10)]),
        );
    }

    #[tokio::test]
    async fn test_scan_partially_translated_not_filter_prunes_partitions_but_skips_limit_hint() {
        let provider = create_provider("multi_partitioned_log_table").await;
        let filter = col("dt")
            .eq(lit("2024-01-01"))
            .and(Expr::Not(Box::new(col("hr").eq(lit(10)))));
        let full_plan = plan_partitions(&provider, vec![filter.clone()], None).await;
        let plan = plan_scan(&provider, vec![filter], Some(1)).await;
        let scan = plan
            .downcast_ref::<PaimonTableScan>()
            .expect("Expected PaimonTableScan");

        assert_eq!(scan.limit(), None);
        assert_eq!(
            extract_dt_hr_partition_set(scan.planned_partitions()),
            BTreeSet::from([("2024-01-01".to_string(), 20)]),
        );
        assert_eq!(
            scan.planned_partitions()
                .iter()
                .map(|partition| partition.len())
                .sum::<usize>(),
            full_plan
                .iter()
                .map(|partition| partition.len())
                .sum::<usize>()
        );
    }

    #[tokio::test]
    async fn test_scan_keeps_pushed_predicate_for_execute() {
        let provider = create_provider("partitioned_log_table").await;
        let filter = col("id").gt(lit(1));

        let config = SessionConfig::new().with_target_partitions(8);
        let ctx = SessionContext::new_with_config(config);
        let state = ctx.state();
        let plan = provider
            .scan(&state, None, std::slice::from_ref(&filter), None)
            .await
            .expect("scan() should succeed");
        let scan = plan
            .downcast_ref::<PaimonTableScan>()
            .expect("Expected PaimonTableScan");

        let expected = build_pushed_predicate(&[filter], provider.table().schema().fields())
            .expect("data filter should translate");

        assert_eq!(scan.pushed_predicate(), Some(&expected));
    }

    #[tokio::test]
    async fn test_scan_pushes_not_as_inexact_and_skips_limit_hint() {
        let provider = data_evolution_projection_pruning_provider().await;
        let filter = Expr::Not(Box::new(col("id").eq(lit(1))));
        let plan = plan_scan(&provider, vec![filter.clone()], Some(1)).await;
        let scan = plan
            .downcast_ref::<PaimonTableScan>()
            .expect("Expected PaimonTableScan");

        let expected = build_pushed_predicate(&[filter], provider.table().schema().fields())
            .expect("NOT filter should translate as inexact pushdown");

        assert_eq!(scan.pushed_predicate(), Some(&expected));
        assert!(!scan.filter_exact());
        assert_eq!(scan.limit(), None);
    }

    #[tokio::test]
    async fn test_scan_applies_projection_to_data_evolution_planning() {
        let provider = data_evolution_projection_pruning_provider().await;
        let config = SessionConfig::new().with_target_partitions(8);
        let ctx = SessionContext::new_with_config(config);
        let state = ctx.state();

        let full_plan = provider
            .scan(&state, None, &[], None)
            .await
            .expect("full scan should succeed");
        let full_scan = full_plan
            .downcast_ref::<PaimonTableScan>()
            .expect("Expected PaimonTableScan");
        assert_eq!(
            planned_file_names(full_scan),
            vec!["id.parquet".to_string(), "name.parquet".to_string()]
        );

        let projection = vec![1];
        let projected_plan = provider
            .scan(&state, Some(&projection), &[], None)
            .await
            .expect("projected scan should succeed");
        let projected_scan = projected_plan
            .downcast_ref::<PaimonTableScan>()
            .expect("Expected PaimonTableScan");

        assert_eq!(
            planned_file_names(projected_scan),
            vec!["name.parquet".to_string()]
        );
    }

    #[tokio::test]
    async fn test_scan_applies_limit_hint_only_when_safe() {
        let provider = create_provider("partitioned_log_table").await;
        let full_plan = plan_partitions(&provider, vec![], None).await;
        let plan = plan_scan(&provider, vec![], Some(1)).await;
        let scan = plan
            .downcast_ref::<PaimonTableScan>()
            .expect("Expected PaimonTableScan");

        assert_eq!(scan.limit(), Some(1));
        assert!(
            scan.planned_partitions()
                .iter()
                .map(|partition| partition.len())
                .sum::<usize>()
                < full_plan
                    .iter()
                    .map(|partition| partition.len())
                    .sum::<usize>()
        );
    }

    #[tokio::test]
    async fn test_scan_keeps_limit_but_skips_limit_pruning_for_data_filters() {
        let provider = create_provider("partitioned_log_table").await;
        let filter = col("id").gt(lit(1));
        let full_plan = plan_partitions(&provider, vec![filter.clone()], None).await;
        let plan = plan_scan(&provider, vec![filter], Some(1)).await;
        let scan = plan
            .downcast_ref::<PaimonTableScan>()
            .expect("Expected PaimonTableScan");

        assert_eq!(scan.limit(), Some(1));
        assert_eq!(
            scan.planned_partitions()
                .iter()
                .map(|partition| partition.len())
                .sum::<usize>(),
            full_plan
                .iter()
                .map(|partition| partition.len())
                .sum::<usize>()
        );
    }

    #[tokio::test]
    async fn test_insert_into_and_read_back() {
        use paimon::io::FileIOBuilder;
        use paimon::spec::{DataType, IntType, Schema as PaimonSchema, TableSchema};

        let file_io = FileIOBuilder::new("memory").build().unwrap();
        let table_path = "memory:/test_df_insert_into";
        file_io
            .mkdirs(&format!("{table_path}/snapshot/"))
            .await
            .unwrap();
        file_io
            .mkdirs(&format!("{table_path}/manifest/"))
            .await
            .unwrap();

        let schema = PaimonSchema::builder()
            .column("id", DataType::Int(IntType::new()))
            .column("value", DataType::Int(IntType::new()))
            .build()
            .unwrap();
        let table_schema = TableSchema::new(0, &schema);
        let table = paimon::table::Table::new(
            file_io,
            Identifier::new("default", "test_insert"),
            table_path.to_string(),
            table_schema,
            None,
        );

        let provider = PaimonTableProvider::try_new(table).unwrap();
        let ctx = SessionContext::new();
        ctx.register_table("t", Arc::new(provider)).unwrap();

        // INSERT INTO
        let result = ctx
            .sql("INSERT INTO t VALUES (1, 10), (2, 20), (3, 30)")
            .await
            .unwrap()
            .collect()
            .await
            .unwrap();

        // Verify count output
        let count_array = result[0]
            .column(0)
            .as_any()
            .downcast_ref::<datafusion::arrow::array::UInt64Array>()
            .unwrap();
        assert_eq!(count_array.value(0), 3);

        // Read back
        let batches = ctx
            .sql("SELECT id, value FROM t ORDER BY id")
            .await
            .unwrap()
            .collect()
            .await
            .unwrap();

        let mut rows = Vec::new();
        for batch in &batches {
            let ids = batch
                .column(0)
                .as_any()
                .downcast_ref::<datafusion::arrow::array::Int32Array>()
                .unwrap();
            let vals = batch
                .column(1)
                .as_any()
                .downcast_ref::<datafusion::arrow::array::Int32Array>()
                .unwrap();
            for i in 0..batch.num_rows() {
                rows.push((ids.value(i), vals.value(i)));
            }
        }
        assert_eq!(rows, vec![(1, 10), (2, 20), (3, 30)]);
    }

    #[tokio::test]
    async fn test_insert_overwrite() {
        use paimon::io::FileIOBuilder;
        use paimon::spec::{DataType, IntType, Schema as PaimonSchema, TableSchema, VarCharType};

        let file_io = FileIOBuilder::new("memory").build().unwrap();
        let table_path = "memory:/test_df_insert_overwrite";
        file_io
            .mkdirs(&format!("{table_path}/snapshot/"))
            .await
            .unwrap();
        file_io
            .mkdirs(&format!("{table_path}/manifest/"))
            .await
            .unwrap();

        let schema = PaimonSchema::builder()
            .column("pt", DataType::VarChar(VarCharType::string_type()))
            .column("id", DataType::Int(IntType::new()))
            .partition_keys(["pt"])
            .build()
            .unwrap();
        let table_schema = TableSchema::new(0, &schema);
        let table = paimon::table::Table::new(
            file_io,
            Identifier::new("default", "test_overwrite"),
            table_path.to_string(),
            table_schema,
            None,
        );

        let provider = PaimonTableProvider::try_new(table).unwrap();
        let ctx = SessionContext::new();
        ctx.register_table("t", Arc::new(provider)).unwrap();

        // Initial INSERT: partition "a" and "b"
        ctx.sql("INSERT INTO t VALUES ('a', 1), ('a', 2), ('b', 3), ('b', 4)")
            .await
            .unwrap()
            .collect()
            .await
            .unwrap();

        // INSERT OVERWRITE with only partition "a" data
        // Should overwrite partition "a" but leave partition "b" intact
        ctx.sql("INSERT OVERWRITE t VALUES ('a', 10), ('a', 20)")
            .await
            .unwrap()
            .collect()
            .await
            .unwrap();

        // Read back
        let batches = ctx
            .sql("SELECT pt, id FROM t ORDER BY pt, id")
            .await
            .unwrap()
            .collect()
            .await
            .unwrap();

        let mut rows = Vec::new();
        for batch in &batches {
            let pts = batch
                .column(0)
                .as_any()
                .downcast_ref::<datafusion::arrow::array::StringViewArray>()
                .unwrap();
            let ids = batch
                .column(1)
                .as_any()
                .downcast_ref::<datafusion::arrow::array::Int32Array>()
                .unwrap();
            for i in 0..batch.num_rows() {
                rows.push((pts.value(i).to_string(), ids.value(i)));
            }
        }
        // Partition "a" overwritten with new data, partition "b" untouched
        assert_eq!(
            rows,
            vec![
                ("a".to_string(), 10),
                ("a".to_string(), 20),
                ("b".to_string(), 3),
                ("b".to_string(), 4),
            ]
        );
    }

    #[tokio::test]
    async fn test_insert_overwrite_unpartitioned() {
        use paimon::io::FileIOBuilder;
        use paimon::spec::{DataType, IntType, Schema as PaimonSchema, TableSchema};

        let file_io = FileIOBuilder::new("memory").build().unwrap();
        let table_path = "memory:/test_df_insert_overwrite_unpart";
        file_io
            .mkdirs(&format!("{table_path}/snapshot/"))
            .await
            .unwrap();
        file_io
            .mkdirs(&format!("{table_path}/manifest/"))
            .await
            .unwrap();

        let schema = PaimonSchema::builder()
            .column("id", DataType::Int(IntType::new()))
            .column("value", DataType::Int(IntType::new()))
            .build()
            .unwrap();
        let table_schema = TableSchema::new(0, &schema);
        let table = paimon::table::Table::new(
            file_io,
            Identifier::new("default", "test_overwrite_unpart"),
            table_path.to_string(),
            table_schema,
            None,
        );

        let provider = PaimonTableProvider::try_new(table).unwrap();
        let ctx = SessionContext::new();
        ctx.register_table("t", Arc::new(provider)).unwrap();

        // Initial INSERT
        ctx.sql("INSERT INTO t VALUES (1, 10), (2, 20), (3, 30)")
            .await
            .unwrap()
            .collect()
            .await
            .unwrap();

        // INSERT OVERWRITE on unpartitioned table — full table overwrite
        ctx.sql("INSERT OVERWRITE t VALUES (4, 40), (5, 50)")
            .await
            .unwrap()
            .collect()
            .await
            .unwrap();

        let batches = ctx
            .sql("SELECT id, value FROM t ORDER BY id")
            .await
            .unwrap()
            .collect()
            .await
            .unwrap();

        let mut rows = Vec::new();
        for batch in &batches {
            let ids = batch
                .column(0)
                .as_any()
                .downcast_ref::<datafusion::arrow::array::Int32Array>()
                .unwrap();
            let vals = batch
                .column(1)
                .as_any()
                .downcast_ref::<datafusion::arrow::array::Int32Array>()
                .unwrap();
            for i in 0..batch.num_rows() {
                rows.push((ids.value(i), vals.value(i)));
            }
        }
        // Old data fully replaced
        assert_eq!(rows, vec![(4, 40), (5, 50)]);
    }
}