paimon_datafusion/table/
mod.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Paimon table provider for DataFusion.
19
20use std::fmt::Write as _;
21use std::sync::Arc;
22
23use async_trait::async_trait;
24use datafusion::arrow::datatypes::{
25    DataType as ArrowDataType, Field, Schema, SchemaRef as ArrowSchemaRef,
26};
27use datafusion::catalog::Session;
28use datafusion::datasource::sink::DataSinkExec;
29use datafusion::datasource::{TableProvider, TableType};
30use datafusion::error::{DataFusionError, Result as DFResult};
31use datafusion::logical_expr::dml::InsertOp;
32use datafusion::logical_expr::{Expr, TableProviderFilterPushDown};
33use datafusion::physical_plan::ExecutionPlan;
34use paimon::spec::{
35    BigIntType, CoreOptions, DataField, DataType, ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME,
36};
37use paimon::table::Table;
38
39use crate::physical_plan::PaimonDataSink;
40use crate::BlobReaderRegistry;
41
42use crate::error::to_datafusion_error;
43#[cfg(test)]
44use crate::filter_pushdown::build_pushed_predicate;
45use crate::filter_pushdown::{analyze_filters, classify_filter_pushdown};
46use crate::physical_plan::PaimonTableScan;
47use crate::runtime::await_with_runtime;
48
49const PARQUET_FIELD_ID_META_KEY: &str = "PARQUET:field_id";
50
51pub(crate) fn datafusion_read_fields(table: &Table) -> Vec<DataField> {
52    let mut fields = table.schema().fields().to_vec();
53    if CoreOptions::new(table.schema().options()).data_evolution_enabled() {
54        fields.push(DataField::new(
55            ROW_ID_FIELD_ID,
56            ROW_ID_FIELD_NAME.to_string(),
57            DataType::BigInt(BigIntType::with_nullable(true)),
58        ));
59    }
60    fields
61}
62
63fn datafusion_arrow_schema(
64    fields: &[DataField],
65    schema_force_view_types: bool,
66) -> DFResult<ArrowSchemaRef> {
67    let paimon_schema =
68        paimon::arrow::build_target_arrow_schema(fields).map_err(to_datafusion_error)?;
69    let fields = paimon_schema
70        .fields()
71        .iter()
72        .map(|field| {
73            let mut metadata = field.metadata().clone();
74            metadata.remove(PARQUET_FIELD_ID_META_KEY);
75            let data_type = match field.data_type() {
76                ArrowDataType::Utf8 if schema_force_view_types => ArrowDataType::Utf8View,
77                data_type => data_type.clone(),
78            };
79            Arc::new(
80                field
81                    .as_ref()
82                    .clone()
83                    .with_data_type(data_type)
84                    .with_metadata(metadata),
85            )
86        })
87        .collect::<Vec<_>>();
88
89    Ok(Arc::new(Schema::new_with_metadata(
90        fields,
91        paimon_schema.metadata().clone(),
92    )))
93}
94
95/// Read-only table provider for a Paimon table.
96///
97/// Supports full table scan, column projection, and predicate pushdown for
98/// planning. Partition predicates prune splits eagerly, while supported
99/// non-partition data predicates may also be reused by the Parquet read path
100/// for row-group pruning and partial decode-time filtering.
101///
102/// DataFusion still treats pushed filters as inexact because unsupported
103/// predicates and non-Parquet reads remain residual filters.
104#[derive(Debug, Clone)]
105pub struct PaimonTableProvider {
106    table: Table,
107    schema: ArrowSchemaRef,
108    table_definition: Option<String>,
109}
110
111impl PaimonTableProvider {
112    /// Create a table provider from a Paimon table.
113    ///
114    /// Loads the table schema and converts it to Arrow for DataFusion.
115    pub fn try_new(table: Table) -> DFResult<Self> {
116        let table_definition = build_table_definition(&table)?;
117        Self::try_new_with_table_definition(table, Some(table_definition))
118    }
119
120    fn try_new_with_table_definition(
121        table: Table,
122        table_definition: Option<String>,
123    ) -> DFResult<Self> {
124        let fields = datafusion_read_fields(&table);
125        let schema = datafusion_arrow_schema(&fields, true)?;
126        Ok(Self {
127            table,
128            schema,
129            table_definition,
130        })
131    }
132
133    pub fn try_new_with_blob_reader_registry(
134        table: Table,
135        blob_reader_registry: BlobReaderRegistry,
136    ) -> DFResult<Self> {
137        blob_reader_registry
138            .register_if_absent(table.location().to_string(), table.file_io().clone());
139        Self::try_new(table)
140    }
141
142    pub(crate) fn try_new_with_blob_reader_registry_and_definition(
143        table: Table,
144        blob_reader_registry: BlobReaderRegistry,
145        table_definition: Option<String>,
146    ) -> DFResult<Self> {
147        blob_reader_registry
148            .register_if_absent(table.location().to_string(), table.file_io().clone());
149        Self::try_new_with_table_definition(table, table_definition)
150    }
151
152    pub(crate) fn with_schema_force_view_types(
153        mut self,
154        schema_force_view_types: bool,
155    ) -> DFResult<Self> {
156        if schema_force_view_types {
157            return Ok(self);
158        }
159        let fields = datafusion_read_fields(&self.table);
160        self.schema = datafusion_arrow_schema(&fields, schema_force_view_types)?;
161        Ok(self)
162    }
163
164    pub fn table(&self) -> &Table {
165        &self.table
166    }
167}
168
169/// Build a `CREATE TABLE` DDL string for a Paimon table.
170///
171/// Mirrors the syntax accepted by `SQLContext::handle_create_table`:
172/// `CREATE TABLE <db>.<table> (<col> <type>, ..., PRIMARY KEY (...)) [PARTITIONED BY (...)] [WITH ('k'='v', ...)]`.
173pub(crate) fn build_table_definition(table: &Table) -> DFResult<String> {
174    let identifier = table.identifier();
175    let schema = table.schema();
176    let mut ddl = String::new();
177    let _ = write!(
178        ddl,
179        "CREATE TABLE {}.{} (",
180        quote_identifier(identifier.database()),
181        quote_identifier(identifier.object())
182    );
183
184    for (i, field) in schema.fields().iter().enumerate() {
185        if i > 0 {
186            ddl.push_str(", ");
187        }
188        // `NOT NULL` is a column constraint; render it here at the column
189        // level rather than inside nested type arguments (see `data_type_to_sql`).
190        let ty = data_type_to_sql(field.data_type())?;
191        if field.data_type().is_nullable() {
192            let _ = write!(ddl, "{} {}", quote_identifier(field.name()), ty);
193        } else {
194            let _ = write!(ddl, "{} {} NOT NULL", quote_identifier(field.name()), ty);
195        }
196    }
197
198    let pks = schema.primary_keys();
199    if !pks.is_empty() {
200        ddl.push_str(", PRIMARY KEY (");
201        for (i, pk) in pks.iter().enumerate() {
202            if i > 0 {
203                ddl.push_str(", ");
204            }
205            let _ = write!(ddl, "{}", quote_identifier(pk));
206        }
207        ddl.push(')');
208    }
209    ddl.push(')');
210
211    let partition_keys = schema.partition_keys();
212    if !partition_keys.is_empty() {
213        ddl.push_str(" PARTITIONED BY (");
214        for (i, pk) in partition_keys.iter().enumerate() {
215            if i > 0 {
216                ddl.push_str(", ");
217            }
218            let _ = write!(ddl, "{}", quote_identifier(pk));
219        }
220        ddl.push(')');
221    }
222
223    let mut options: Vec<_> = schema.options().iter().collect();
224    options.sort_by_key(|(left, _)| *left);
225    if !options.is_empty() {
226        ddl.push_str(" WITH (");
227        for (i, (k, v)) in options.iter().enumerate() {
228            if i > 0 {
229                ddl.push_str(", ");
230            }
231            let _ = write!(
232                ddl,
233                "{} = {}",
234                quote_string_literal(k),
235                quote_string_literal(v)
236            );
237        }
238        ddl.push(')');
239    }
240
241    Ok(ddl)
242}
243
244fn quote_identifier(identifier: &str) -> String {
245    format!("\"{}\"", identifier.replace('"', "\"\""))
246}
247
248fn quote_string_literal(text: &str) -> String {
249    format!("'{}'", text.replace('\'', "''"))
250}
251
252/// Render a Paimon [`DataType`] as a SQL type string matching the syntax
253/// accepted by paimon-rust's `CREATE TABLE` parser.
254///
255/// `NOT NULL` is a column constraint, not a type modifier — it is only valid
256/// at the top of a column definition, not nested inside `MAP`, `ARRAY`, or
257/// `STRUCT` arguments. Callers that render a column should append `NOT NULL`
258/// themselves when the field is non-nullable; recursive calls below must not.
259pub(crate) fn data_type_to_sql(data_type: &DataType) -> DFResult<String> {
260    match data_type {
261        DataType::Boolean(_) => Ok("BOOLEAN".to_string()),
262        DataType::TinyInt(_) => Ok("TINYINT".to_string()),
263        DataType::SmallInt(_) => Ok("SMALLINT".to_string()),
264        DataType::Int(_) => Ok("INT".to_string()),
265        DataType::BigInt(_) => Ok("BIGINT".to_string()),
266        DataType::Decimal(t) => Ok(format!("DECIMAL({}, {})", t.precision(), t.scale())),
267        DataType::Double(_) => Ok("DOUBLE".to_string()),
268        DataType::Float(_) => Ok("FLOAT".to_string()),
269        DataType::Binary(t) => Ok(format!("BINARY({})", t.length())),
270        DataType::VarBinary(t) => Ok(format!("VARBINARY({})", t.length())),
271        DataType::Blob(_) => Ok("BLOB".to_string()),
272        DataType::Char(t) => Ok(format!("CHAR({})", t.length())),
273        DataType::VarChar(t) => Ok(format!("VARCHAR({})", t.length())),
274        DataType::Date(_) => Ok("DATE".to_string()),
275        DataType::Time(_) => Err(unsupported_show_create_table_type("TIME")),
276        DataType::Timestamp(t) => Ok(format!("TIMESTAMP({})", t.precision())),
277        DataType::Variant(_) => Ok("VARIANT".to_string()),
278        DataType::LocalZonedTimestamp(t) => {
279            Ok(format!("TIMESTAMP({}) WITH TIME ZONE", t.precision()))
280        }
281        DataType::Array(t) => Ok(format!("ARRAY<{}>", data_type_to_sql(t.element_type())?)),
282        DataType::Map(t) => Ok(format!(
283            "MAP({}, {})",
284            data_type_to_sql(t.key_type())?,
285            data_type_to_sql(t.value_type())?
286        )),
287        DataType::Multiset(_) => Err(unsupported_show_create_table_type("MULTISET")),
288        DataType::Row(t) => {
289            let inner: Vec<String> = t
290                .fields()
291                .iter()
292                .map(|f| {
293                    let ty = data_type_to_sql(f.data_type())?;
294                    if f.name().is_empty() {
295                        Ok(ty)
296                    } else {
297                        Ok(format!("{} {}", quote_identifier(f.name()), ty))
298                    }
299                })
300                .collect::<DFResult<_>>()?;
301            Ok(format!("STRUCT<{}>", inner.join(", ")))
302        }
303        DataType::Vector(_) => Err(unsupported_show_create_table_type("VECTOR")),
304    }
305}
306
307fn unsupported_show_create_table_type(type_name: &str) -> DataFusionError {
308    DataFusionError::NotImplemented(format!(
309        "SHOW CREATE TABLE does not support {type_name} columns because paimon-rust cannot round-trip this type in CREATE TABLE"
310    ))
311}
312
313/// Distribute `items` into `num_buckets` groups using round-robin assignment.
314pub(crate) fn bucket_round_robin<T>(items: Vec<T>, num_buckets: usize) -> Vec<Vec<T>> {
315    let mut buckets: Vec<Vec<T>> = (0..num_buckets).map(|_| Vec::new()).collect();
316    for (index, item) in items.into_iter().enumerate() {
317        buckets[index % num_buckets].push(item);
318    }
319    buckets
320}
321
322/// Build parameters for [`PaimonTableScan`].
323pub(crate) struct PaimonScanBuilder<'a> {
324    pub(crate) table: &'a Table,
325    pub(crate) schema: &'a ArrowSchemaRef,
326    pub(crate) plan: &'a paimon::table::Plan,
327    pub(crate) scan_trace: Option<paimon::table::ScanTrace>,
328    pub(crate) projection: Option<&'a Vec<usize>>,
329    pub(crate) pushed_predicate: Option<paimon::spec::Predicate>,
330    pub(crate) limit: Option<usize>,
331    pub(crate) target_partitions: usize,
332    pub(crate) filter_exact: bool,
333    /// Column-name case sensitivity, carried into the physical scan so execute()
334    /// resolves names the same way planning did.
335    pub(crate) case_sensitive: bool,
336}
337
338impl PaimonScanBuilder<'_> {
339    /// Build a [`PaimonTableScan`] from the configured parameters.
340    pub(crate) fn build(self) -> DFResult<Arc<dyn ExecutionPlan>> {
341        let read_fields = datafusion_read_fields(self.table);
342        self.build_with_read_fields(read_fields)
343    }
344
345    /// Build using an internal read schema which may contain hidden system fields.
346    pub(crate) fn build_with_read_fields(
347        self,
348        read_fields: Vec<DataField>,
349    ) -> DFResult<Arc<dyn ExecutionPlan>> {
350        let (projected_schema, read_type) = if let Some(indices) = self.projection {
351            let fields: Vec<Field> = indices
352                .iter()
353                .map(|&i| self.schema.field(i).clone())
354                .collect();
355            let read_type = indices
356                .iter()
357                .map(|&i| read_fields[i].clone())
358                .collect::<Vec<_>>();
359            (Arc::new(Schema::new(fields)), read_type)
360        } else {
361            (self.schema.clone(), read_fields)
362        };
363
364        let splits = self.plan.splits().to_vec();
365        let planned_partitions: Vec<Arc<[_]>> = if splits.is_empty() {
366            vec![Arc::from(Vec::new())]
367        } else {
368            let num_partitions = splits.len().min(self.target_partitions.max(1));
369            bucket_round_robin(splits, num_partitions)
370                .into_iter()
371                .map(Arc::from)
372                .collect()
373        };
374
375        Ok(Arc::new(PaimonTableScan::new(
376            projected_schema,
377            self.table.clone(),
378            read_type,
379            self.pushed_predicate,
380            planned_partitions,
381            self.limit,
382            self.filter_exact,
383            self.scan_trace,
384            None,
385            self.case_sensitive,
386        )))
387    }
388}
389
390#[async_trait]
391impl TableProvider for PaimonTableProvider {
392    fn schema(&self) -> ArrowSchemaRef {
393        self.schema.clone()
394    }
395
396    fn table_type(&self) -> TableType {
397        TableType::Base
398    }
399
400    fn get_table_definition(&self) -> Option<&str> {
401        self.table_definition.as_deref()
402    }
403
404    async fn scan(
405        &self,
406        state: &dyn Session,
407        projection: Option<&Vec<usize>>,
408        filters: &[Expr],
409        limit: Option<usize>,
410    ) -> DFResult<Arc<dyn ExecutionPlan>> {
411        // Column-name matching is case-sensitive on the DataFusion path.
412        //
413        // DataFusion resolves projection/filter columns against the provider
414        // schema (`schema()`, which exposes the original field casing) during
415        // logical planning, *before* `scan` is called. `enable_ident_normalization`
416        // only lowercases unquoted identifiers at parse time; it does not make
417        // schema resolution case-insensitive. So a genuine case mismatch
418        // (`SELECT name` against a `Name` field) already fails at planning and
419        // never reaches this code — see the negative test in `tests/read_tables.rs`.
420        // Case-insensitive column matching is therefore offered only through the
421        // direct ReadBuilder API (core / C / Python), not via SQL.
422        let case_sensitive = true;
423        // Plan splits eagerly so we know partition count upfront.
424        let filter_analysis =
425            analyze_filters(filters, self.table.schema().fields(), case_sensitive);
426        let mut read_builder = self.table.new_read_builder();
427        read_builder.with_case_sensitive(case_sensitive);
428        if let Some(indices) = projection {
429            let read_fields = datafusion_read_fields(&self.table);
430            let read_type = indices
431                .iter()
432                .map(|&i| read_fields[i].clone())
433                .collect::<Vec<_>>();
434            read_builder.with_read_type(read_type);
435        }
436        if let Some(filter) = filter_analysis.pushed_predicate.clone() {
437            read_builder.with_filter(filter);
438        }
439        let pushed_limit = limit.filter(|_| !filter_analysis.requires_residual);
440        if let Some(limit) = pushed_limit {
441            read_builder.with_limit(limit);
442        }
443        let scan = read_builder.new_scan();
444        // DataFusion's Python FFI may poll `TableProvider::scan()` without an active
445        // Tokio runtime. `scan.plan()` can reach OpenDAL/Tokio filesystem calls while
446        // reading Paimon metadata, so we must provide a runtime here instead of
447        // assuming the caller already entered one.
448        let (plan, scan_trace) = await_with_runtime(scan.plan_with_trace())
449            .await
450            .map_err(to_datafusion_error)?;
451
452        let target = state.config_options().execution.target_partitions;
453        let filter_exact = !filter_analysis.requires_residual
454            && filter_analysis
455                .pushed_predicate
456                .as_ref()
457                .is_none_or(|p| read_builder.is_exact_filter_pushdown(p));
458        PaimonScanBuilder {
459            table: &self.table,
460            schema: &self.schema,
461            plan: &plan,
462            scan_trace: Some(scan_trace),
463            projection,
464            pushed_predicate: filter_analysis.pushed_predicate,
465            limit: pushed_limit,
466            target_partitions: target,
467            filter_exact,
468            case_sensitive,
469        }
470        .build()
471    }
472
473    async fn insert_into(
474        &self,
475        _state: &dyn Session,
476        input: Arc<dyn ExecutionPlan>,
477        insert_op: InsertOp,
478    ) -> DFResult<Arc<dyn ExecutionPlan>> {
479        if self.table.is_branch_reference() {
480            return Err(datafusion::error::DataFusionError::NotImplemented(format!(
481                "Writing to Paimon branch '{}' is not supported",
482                self.table.branch()
483            )));
484        }
485        let overwrite = match insert_op {
486            InsertOp::Append => false,
487            InsertOp::Overwrite => true,
488            other => {
489                return Err(datafusion::error::DataFusionError::NotImplemented(format!(
490                    "{other} is not supported for Paimon tables"
491                )));
492            }
493        };
494        let sink = PaimonDataSink::new(self.table.clone(), self.schema.clone(), overwrite);
495        Ok(Arc::new(DataSinkExec::new(input, Arc::new(sink), None)))
496    }
497
498    fn supports_filters_pushdown(
499        &self,
500        filters: &[&Expr],
501    ) -> DFResult<Vec<TableProviderFilterPushDown>> {
502        let fields = self.table.schema().fields();
503        // SQL reads resolve columns case-sensitively (see `scan`), so classify
504        // pushdown the same way.
505        let case_sensitive = true;
506        let read_builder = self.table.new_read_builder();
507
508        Ok(filters
509            .iter()
510            .map(|filter| {
511                classify_filter_pushdown(filter, fields, case_sensitive, |predicate| {
512                    read_builder.is_exact_filter_pushdown(predicate)
513                })
514            })
515            .collect())
516    }
517}
518
519#[cfg(test)]
520mod tests {
521    use super::*;
522    use std::collections::BTreeSet;
523    use std::sync::Arc;
524
525    use datafusion::datasource::TableProvider;
526    use datafusion::logical_expr::{col, lit, Expr};
527    use datafusion::prelude::{SessionConfig, SessionContext};
528    use paimon::catalog::Identifier;
529    use paimon::spec::{ArrayType, MapType, RowType, VarCharType};
530    use paimon::{Catalog, CatalogOptions, DataSplit, FileSystemCatalog, Options};
531
532    use crate::physical_plan::PaimonTableScan;
533
534    #[test]
535    fn test_bucket_round_robin_distributes_evenly() {
536        let result = bucket_round_robin(vec![0, 1, 2, 3, 4], 3);
537        assert_eq!(result, vec![vec![0, 3], vec![1, 4], vec![2]]);
538    }
539
540    #[test]
541    fn test_bucket_round_robin_fewer_items_than_buckets() {
542        let result = bucket_round_robin(vec![10, 20], 2);
543        assert_eq!(result, vec![vec![10], vec![20]]);
544    }
545
546    #[test]
547    fn test_bucket_round_robin_single_bucket() {
548        let result = bucket_round_robin(vec![1, 2, 3], 1);
549        assert_eq!(result, vec![vec![1, 2, 3]]);
550    }
551
552    fn get_test_warehouse() -> String {
553        std::env::var("PAIMON_TEST_WAREHOUSE")
554            .unwrap_or_else(|_| "/tmp/paimon-warehouse".to_string())
555    }
556
557    fn create_catalog() -> FileSystemCatalog {
558        let warehouse = get_test_warehouse();
559        let mut options = Options::new();
560        options.set(CatalogOptions::WAREHOUSE, warehouse);
561        FileSystemCatalog::new(options).expect("Failed to create catalog")
562    }
563
564    async fn create_provider(table_name: &str) -> PaimonTableProvider {
565        let catalog = create_catalog();
566        let identifier = Identifier::new("default", table_name);
567        let table = catalog
568            .get_table(&identifier)
569            .await
570            .expect("Failed to get table");
571
572        PaimonTableProvider::try_new(table).expect("Failed to create table provider")
573    }
574
575    async fn plan_partitions(
576        provider: &PaimonTableProvider,
577        filters: Vec<Expr>,
578        limit: Option<usize>,
579    ) -> Vec<Arc<[DataSplit]>> {
580        let plan = plan_scan(provider, filters, limit).await;
581        let scan = plan
582            .downcast_ref::<PaimonTableScan>()
583            .expect("Expected PaimonTableScan");
584
585        scan.planned_partitions().to_vec()
586    }
587
588    async fn plan_scan(
589        provider: &PaimonTableProvider,
590        filters: Vec<Expr>,
591        limit: Option<usize>,
592    ) -> Arc<dyn ExecutionPlan> {
593        let config = SessionConfig::new().with_target_partitions(8);
594        let ctx = SessionContext::new_with_config(config);
595        let state = ctx.state();
596        provider
597            .scan(&state, None, &filters, limit)
598            .await
599            .expect("scan() should succeed")
600    }
601
602    fn extract_dt_partition_set(planned_partitions: &[Arc<[DataSplit]>]) -> BTreeSet<String> {
603        planned_partitions
604            .iter()
605            .flat_map(|splits| splits.iter())
606            .map(|split| {
607                split
608                    .partition()
609                    .get_string(0)
610                    .expect("Failed to decode dt")
611                    .to_string()
612            })
613            .collect()
614    }
615
616    fn extract_dt_hr_partition_set(
617        planned_partitions: &[Arc<[DataSplit]>],
618    ) -> BTreeSet<(String, i32)> {
619        planned_partitions
620            .iter()
621            .flat_map(|splits| splits.iter())
622            .map(|split| {
623                let partition = split.partition();
624                (
625                    partition
626                        .get_string(0)
627                        .expect("Failed to decode dt")
628                        .to_string(),
629                    partition.get_int(1).expect("Failed to decode hr"),
630                )
631            })
632            .collect()
633    }
634
635    fn empty_binary_stats_json() -> serde_json::Value {
636        let row = paimon::spec::EMPTY_SERIALIZED_ROW.as_slice().to_vec();
637        serde_json::json!({
638            "_MIN_VALUES": row,
639            "_MAX_VALUES": row,
640            "_NULL_COUNTS": [],
641        })
642    }
643
644    fn data_evolution_file(
645        file_name: &str,
646        file_size: i64,
647        row_count: i64,
648        first_row_id: i64,
649        write_cols: &[&str],
650    ) -> paimon::spec::DataFileMeta {
651        serde_json::from_value(serde_json::json!({
652            "_FILE_NAME": file_name,
653            "_FILE_SIZE": file_size,
654            "_ROW_COUNT": row_count,
655            "_MIN_KEY": [],
656            "_MAX_KEY": [],
657            "_KEY_STATS": empty_binary_stats_json(),
658            "_VALUE_STATS": empty_binary_stats_json(),
659            "_MIN_SEQUENCE_NUMBER": 0,
660            "_MAX_SEQUENCE_NUMBER": 0,
661            "_SCHEMA_ID": 0,
662            "_LEVEL": 1,
663            "_EXTRA_FILES": [],
664            "_CREATION_TIME": null,
665            "_DELETE_ROW_COUNT": null,
666            "_EMBEDDED_FILE_INDEX": null,
667            "_FILE_SOURCE": null,
668            "_VALUE_STATS_COLS": null,
669            "_FIRST_ROW_ID": first_row_id,
670            "_WRITE_COLS": write_cols,
671            "_EXTERNAL_PATH": null,
672        }))
673        .expect("test data file should deserialize")
674    }
675
676    fn manifest_file_meta(
677        file_name: &str,
678        file_size: i64,
679        num_added_files: i64,
680    ) -> paimon::spec::ManifestFileMeta {
681        serde_json::from_value(serde_json::json!({
682            "_VERSION": 2,
683            "_FILE_NAME": file_name,
684            "_FILE_SIZE": file_size,
685            "_NUM_ADDED_FILES": num_added_files,
686            "_NUM_DELETED_FILES": 0,
687            "_PARTITION_STATS": empty_binary_stats_json(),
688            "_SCHEMA_ID": 0,
689        }))
690        .expect("test manifest file meta should deserialize")
691    }
692
693    async fn data_evolution_projection_pruning_provider() -> PaimonTableProvider {
694        use paimon::io::FileIOBuilder;
695        use paimon::spec::{
696            CommitKind, DataType, FileKind, IntType, Manifest, ManifestEntry, ManifestList,
697            Schema as PaimonSchema, Snapshot, TableSchema,
698        };
699        use paimon::table::{SnapshotManager, Table};
700
701        let file_io = FileIOBuilder::new("memory").build().unwrap();
702        let table_path = format!("memory:/df_de_projection_pruning_{}", uuid::Uuid::new_v4());
703        file_io
704            .mkdirs(&format!("{table_path}/snapshot/"))
705            .await
706            .unwrap();
707        file_io
708            .mkdirs(&format!("{table_path}/manifest/"))
709            .await
710            .unwrap();
711
712        let schema = PaimonSchema::builder()
713            .column("id", DataType::Int(IntType::new()))
714            .column("name", DataType::Int(IntType::new()))
715            .option("data-evolution.enabled", "true")
716            .build()
717            .unwrap();
718        let table_schema = TableSchema::new(0, &schema);
719        let table = Table::new(
720            file_io.clone(),
721            Identifier::new("default", "df_de_projection_pruning"),
722            table_path.clone(),
723            table_schema,
724            None,
725        );
726
727        let partition = paimon::spec::EMPTY_SERIALIZED_ROW.as_slice().to_vec();
728        let entries = vec![
729            ManifestEntry::new(
730                FileKind::Add,
731                partition.clone(),
732                0,
733                1,
734                data_evolution_file("id.parquet", 11, 10, 0, &["id"]),
735                2,
736            ),
737            ManifestEntry::new(
738                FileKind::Add,
739                partition,
740                0,
741                1,
742                data_evolution_file("name.parquet", 13, 10, 0, &["name"]),
743                2,
744            ),
745        ];
746
747        let manifest_name = "manifest-de-projection-0";
748        let manifest_path = format!("{table_path}/manifest/{manifest_name}");
749        Manifest::write(&file_io, &manifest_path, &entries)
750            .await
751            .unwrap();
752        let manifest_size = file_io
753            .new_input(&manifest_path)
754            .unwrap()
755            .metadata()
756            .await
757            .unwrap()
758            .size;
759
760        let base_list_name = "base-list-de-projection";
761        let delta_list_name = "delta-list-de-projection";
762        ManifestList::write(
763            &file_io,
764            &format!("{table_path}/manifest/{base_list_name}"),
765            &[manifest_file_meta(
766                manifest_name,
767                manifest_size as i64,
768                entries.len() as i64,
769            )],
770        )
771        .await
772        .unwrap();
773        ManifestList::write(
774            &file_io,
775            &format!("{table_path}/manifest/{delta_list_name}"),
776            &[],
777        )
778        .await
779        .unwrap();
780
781        let snapshot = Snapshot::builder()
782            .version(3)
783            .id(1)
784            .schema_id(0)
785            .base_manifest_list(base_list_name.to_string())
786            .delta_manifest_list(delta_list_name.to_string())
787            .commit_user("test-user".to_string())
788            .commit_identifier(1)
789            .commit_kind(CommitKind::APPEND)
790            .time_millis(1)
791            .total_record_count(Some(10))
792            .delta_record_count(Some(10))
793            .build();
794        let snapshot_manager = SnapshotManager::new(file_io, table_path);
795        assert!(snapshot_manager.commit_snapshot(&snapshot).await.unwrap());
796
797        PaimonTableProvider::try_new(table).expect("provider should be created")
798    }
799
800    #[tokio::test]
801    async fn test_datafusion_schema_hides_paimon_field_ids() {
802        let provider = data_evolution_projection_pruning_provider().await;
803
804        for field in provider.schema().fields() {
805            assert!(
806                !field.metadata().contains_key("PARQUET:field_id"),
807                "storage field id leaked through DataFusion schema for {}",
808                field.name()
809            );
810        }
811    }
812
813    #[test]
814    fn test_datafusion_schema_uses_views_only_for_top_level_strings() {
815        let string_type = || DataType::VarChar(VarCharType::string_type());
816        let schema = datafusion_arrow_schema(
817            &[
818                DataField::new(0, "plain".to_string(), string_type()),
819                DataField::new(
820                    1,
821                    "array".to_string(),
822                    DataType::Array(ArrayType::new(string_type())),
823                ),
824                DataField::new(
825                    2,
826                    "map".to_string(),
827                    DataType::Map(MapType::new(string_type(), string_type())),
828                ),
829                DataField::new(
830                    3,
831                    "row".to_string(),
832                    DataType::Row(RowType::new(vec![DataField::new(
833                        4,
834                        "nested".to_string(),
835                        string_type(),
836                    )])),
837                ),
838            ],
839            true,
840        )
841        .expect("DataFusion schema should be created");
842
843        assert_eq!(schema.field(0).data_type(), &ArrowDataType::Utf8View);
844
845        let ArrowDataType::List(element) = schema.field(1).data_type() else {
846            panic!("array field should map to an Arrow List");
847        };
848        assert_eq!(element.data_type(), &ArrowDataType::Utf8);
849
850        let ArrowDataType::Map(entries, _) = schema.field(2).data_type() else {
851            panic!("map field should map to an Arrow Map");
852        };
853        let ArrowDataType::Struct(map_fields) = entries.data_type() else {
854            panic!("map entries should map to an Arrow Struct");
855        };
856        assert_eq!(map_fields[0].data_type(), &ArrowDataType::Utf8);
857        assert_eq!(map_fields[1].data_type(), &ArrowDataType::Utf8);
858
859        let ArrowDataType::Struct(row_fields) = schema.field(3).data_type() else {
860            panic!("row field should map to an Arrow Struct");
861        };
862        assert_eq!(row_fields[0].data_type(), &ArrowDataType::Utf8);
863    }
864
865    fn planned_file_names(scan: &PaimonTableScan) -> Vec<String> {
866        let mut names = scan
867            .planned_partitions()
868            .iter()
869            .flat_map(|partition| partition.iter())
870            .flat_map(|split| split.data_files().iter())
871            .map(|file| file.file_name.clone())
872            .collect::<Vec<_>>();
873        names.sort();
874        names
875    }
876
877    #[tokio::test]
878    async fn test_scan_partition_filter_plans_matching_partition_set() {
879        let provider = create_provider("partitioned_log_table").await;
880        let planned_partitions =
881            plan_partitions(&provider, vec![col("dt").eq(lit("2024-01-01"))], None).await;
882
883        assert_eq!(
884            extract_dt_partition_set(&planned_partitions),
885            BTreeSet::from(["2024-01-01".to_string()]),
886        );
887    }
888
889    #[tokio::test]
890    async fn test_scan_mixed_and_filter_keeps_partition_pruning() {
891        let provider = create_provider("partitioned_log_table").await;
892        let planned_partitions = plan_partitions(
893            &provider,
894            vec![col("dt").eq(lit("2024-01-01")).and(col("id").gt(lit(1)))],
895            None,
896        )
897        .await;
898
899        assert_eq!(
900            extract_dt_partition_set(&planned_partitions),
901            BTreeSet::from(["2024-01-01".to_string()]),
902        );
903    }
904
905    #[tokio::test]
906    async fn test_scan_multi_partition_filter_plans_exact_partition_set() {
907        let provider = create_provider("multi_partitioned_log_table").await;
908
909        let dt_only_partitions =
910            plan_partitions(&provider, vec![col("dt").eq(lit("2024-01-01"))], None).await;
911        let dt_hr_partitions = plan_partitions(
912            &provider,
913            vec![col("dt").eq(lit("2024-01-01")).and(col("hr").eq(lit(10)))],
914            None,
915        )
916        .await;
917
918        assert_eq!(
919            extract_dt_hr_partition_set(&dt_only_partitions),
920            BTreeSet::from([
921                ("2024-01-01".to_string(), 10),
922                ("2024-01-01".to_string(), 20),
923            ]),
924        );
925        assert_eq!(
926            extract_dt_hr_partition_set(&dt_hr_partitions),
927            BTreeSet::from([("2024-01-01".to_string(), 10)]),
928        );
929    }
930
931    #[tokio::test]
932    async fn test_scan_partially_translated_not_filter_prunes_partitions_but_skips_limit_hint() {
933        let provider = create_provider("multi_partitioned_log_table").await;
934        let filter = col("dt")
935            .eq(lit("2024-01-01"))
936            .and(Expr::Not(Box::new(col("hr").eq(lit(10)))));
937        let full_plan = plan_partitions(&provider, vec![filter.clone()], None).await;
938        let plan = plan_scan(&provider, vec![filter], Some(1)).await;
939        let scan = plan
940            .downcast_ref::<PaimonTableScan>()
941            .expect("Expected PaimonTableScan");
942
943        assert_eq!(scan.limit(), None);
944        assert_eq!(
945            extract_dt_hr_partition_set(scan.planned_partitions()),
946            BTreeSet::from([("2024-01-01".to_string(), 20)]),
947        );
948        assert_eq!(
949            scan.planned_partitions()
950                .iter()
951                .map(|partition| partition.len())
952                .sum::<usize>(),
953            full_plan
954                .iter()
955                .map(|partition| partition.len())
956                .sum::<usize>()
957        );
958    }
959
960    #[tokio::test]
961    async fn test_scan_keeps_pushed_predicate_for_execute() {
962        let provider = create_provider("partitioned_log_table").await;
963        let filter = col("id").gt(lit(1));
964
965        let config = SessionConfig::new().with_target_partitions(8);
966        let ctx = SessionContext::new_with_config(config);
967        let state = ctx.state();
968        let plan = provider
969            .scan(&state, None, std::slice::from_ref(&filter), None)
970            .await
971            .expect("scan() should succeed");
972        let scan = plan
973            .downcast_ref::<PaimonTableScan>()
974            .expect("Expected PaimonTableScan");
975
976        let expected = build_pushed_predicate(&[filter], provider.table().schema().fields())
977            .expect("data filter should translate");
978
979        assert_eq!(scan.pushed_predicate(), Some(&expected));
980    }
981
982    #[tokio::test]
983    async fn test_scan_pushes_not_as_inexact_and_skips_limit_hint() {
984        let provider = data_evolution_projection_pruning_provider().await;
985        let filter = Expr::Not(Box::new(col("id").eq(lit(1))));
986        let plan = plan_scan(&provider, vec![filter.clone()], Some(1)).await;
987        let scan = plan
988            .downcast_ref::<PaimonTableScan>()
989            .expect("Expected PaimonTableScan");
990
991        let expected = build_pushed_predicate(&[filter], provider.table().schema().fields())
992            .expect("NOT filter should translate as inexact pushdown");
993
994        assert_eq!(scan.pushed_predicate(), Some(&expected));
995        assert!(!scan.filter_exact());
996        assert_eq!(scan.limit(), None);
997    }
998
999    #[tokio::test]
1000    async fn test_scan_applies_projection_to_data_evolution_planning() {
1001        let provider = data_evolution_projection_pruning_provider().await;
1002        let config = SessionConfig::new().with_target_partitions(8);
1003        let ctx = SessionContext::new_with_config(config);
1004        let state = ctx.state();
1005
1006        let full_plan = provider
1007            .scan(&state, None, &[], None)
1008            .await
1009            .expect("full scan should succeed");
1010        let full_scan = full_plan
1011            .downcast_ref::<PaimonTableScan>()
1012            .expect("Expected PaimonTableScan");
1013        assert_eq!(
1014            planned_file_names(full_scan),
1015            vec!["id.parquet".to_string(), "name.parquet".to_string()]
1016        );
1017
1018        let projection = vec![1];
1019        let projected_plan = provider
1020            .scan(&state, Some(&projection), &[], None)
1021            .await
1022            .expect("projected scan should succeed");
1023        let projected_scan = projected_plan
1024            .downcast_ref::<PaimonTableScan>()
1025            .expect("Expected PaimonTableScan");
1026
1027        assert_eq!(
1028            planned_file_names(projected_scan),
1029            vec!["name.parquet".to_string()]
1030        );
1031    }
1032
1033    #[tokio::test]
1034    async fn test_scan_applies_limit_hint_only_when_safe() {
1035        let provider = create_provider("partitioned_log_table").await;
1036        let full_plan = plan_partitions(&provider, vec![], None).await;
1037        let plan = plan_scan(&provider, vec![], Some(1)).await;
1038        let scan = plan
1039            .downcast_ref::<PaimonTableScan>()
1040            .expect("Expected PaimonTableScan");
1041
1042        assert_eq!(scan.limit(), Some(1));
1043        assert!(
1044            scan.planned_partitions()
1045                .iter()
1046                .map(|partition| partition.len())
1047                .sum::<usize>()
1048                < full_plan
1049                    .iter()
1050                    .map(|partition| partition.len())
1051                    .sum::<usize>()
1052        );
1053    }
1054
1055    #[tokio::test]
1056    async fn test_scan_keeps_limit_but_skips_limit_pruning_for_data_filters() {
1057        let provider = create_provider("partitioned_log_table").await;
1058        let filter = col("id").gt(lit(1));
1059        let full_plan = plan_partitions(&provider, vec![filter.clone()], None).await;
1060        let plan = plan_scan(&provider, vec![filter], Some(1)).await;
1061        let scan = plan
1062            .downcast_ref::<PaimonTableScan>()
1063            .expect("Expected PaimonTableScan");
1064
1065        assert_eq!(scan.limit(), Some(1));
1066        assert_eq!(
1067            scan.planned_partitions()
1068                .iter()
1069                .map(|partition| partition.len())
1070                .sum::<usize>(),
1071            full_plan
1072                .iter()
1073                .map(|partition| partition.len())
1074                .sum::<usize>()
1075        );
1076    }
1077
1078    #[tokio::test]
1079    async fn test_insert_into_and_read_back() {
1080        use paimon::io::FileIOBuilder;
1081        use paimon::spec::{DataType, IntType, Schema as PaimonSchema, TableSchema};
1082
1083        let file_io = FileIOBuilder::new("memory").build().unwrap();
1084        let table_path = "memory:/test_df_insert_into";
1085        file_io
1086            .mkdirs(&format!("{table_path}/snapshot/"))
1087            .await
1088            .unwrap();
1089        file_io
1090            .mkdirs(&format!("{table_path}/manifest/"))
1091            .await
1092            .unwrap();
1093
1094        let schema = PaimonSchema::builder()
1095            .column("id", DataType::Int(IntType::new()))
1096            .column("value", DataType::Int(IntType::new()))
1097            .build()
1098            .unwrap();
1099        let table_schema = TableSchema::new(0, &schema);
1100        let table = paimon::table::Table::new(
1101            file_io,
1102            Identifier::new("default", "test_insert"),
1103            table_path.to_string(),
1104            table_schema,
1105            None,
1106        );
1107
1108        let provider = PaimonTableProvider::try_new(table).unwrap();
1109        let ctx = SessionContext::new();
1110        ctx.register_table("t", Arc::new(provider)).unwrap();
1111
1112        // INSERT INTO
1113        let result = ctx
1114            .sql("INSERT INTO t VALUES (1, 10), (2, 20), (3, 30)")
1115            .await
1116            .unwrap()
1117            .collect()
1118            .await
1119            .unwrap();
1120
1121        // Verify count output
1122        let count_array = result[0]
1123            .column(0)
1124            .as_any()
1125            .downcast_ref::<datafusion::arrow::array::UInt64Array>()
1126            .unwrap();
1127        assert_eq!(count_array.value(0), 3);
1128
1129        // Read back
1130        let batches = ctx
1131            .sql("SELECT id, value FROM t ORDER BY id")
1132            .await
1133            .unwrap()
1134            .collect()
1135            .await
1136            .unwrap();
1137
1138        let mut rows = Vec::new();
1139        for batch in &batches {
1140            let ids = batch
1141                .column(0)
1142                .as_any()
1143                .downcast_ref::<datafusion::arrow::array::Int32Array>()
1144                .unwrap();
1145            let vals = batch
1146                .column(1)
1147                .as_any()
1148                .downcast_ref::<datafusion::arrow::array::Int32Array>()
1149                .unwrap();
1150            for i in 0..batch.num_rows() {
1151                rows.push((ids.value(i), vals.value(i)));
1152            }
1153        }
1154        assert_eq!(rows, vec![(1, 10), (2, 20), (3, 30)]);
1155    }
1156
1157    #[tokio::test]
1158    async fn test_insert_overwrite() {
1159        use paimon::io::FileIOBuilder;
1160        use paimon::spec::{DataType, IntType, Schema as PaimonSchema, TableSchema, VarCharType};
1161
1162        let file_io = FileIOBuilder::new("memory").build().unwrap();
1163        let table_path = "memory:/test_df_insert_overwrite";
1164        file_io
1165            .mkdirs(&format!("{table_path}/snapshot/"))
1166            .await
1167            .unwrap();
1168        file_io
1169            .mkdirs(&format!("{table_path}/manifest/"))
1170            .await
1171            .unwrap();
1172
1173        let schema = PaimonSchema::builder()
1174            .column("pt", DataType::VarChar(VarCharType::string_type()))
1175            .column("id", DataType::Int(IntType::new()))
1176            .partition_keys(["pt"])
1177            .build()
1178            .unwrap();
1179        let table_schema = TableSchema::new(0, &schema);
1180        let table = paimon::table::Table::new(
1181            file_io,
1182            Identifier::new("default", "test_overwrite"),
1183            table_path.to_string(),
1184            table_schema,
1185            None,
1186        );
1187
1188        let provider = PaimonTableProvider::try_new(table).unwrap();
1189        let ctx = SessionContext::new();
1190        ctx.register_table("t", Arc::new(provider)).unwrap();
1191
1192        // Initial INSERT: partition "a" and "b"
1193        ctx.sql("INSERT INTO t VALUES ('a', 1), ('a', 2), ('b', 3), ('b', 4)")
1194            .await
1195            .unwrap()
1196            .collect()
1197            .await
1198            .unwrap();
1199
1200        // INSERT OVERWRITE with only partition "a" data
1201        // Should overwrite partition "a" but leave partition "b" intact
1202        ctx.sql("INSERT OVERWRITE t VALUES ('a', 10), ('a', 20)")
1203            .await
1204            .unwrap()
1205            .collect()
1206            .await
1207            .unwrap();
1208
1209        // Read back
1210        let batches = ctx
1211            .sql("SELECT pt, id FROM t ORDER BY pt, id")
1212            .await
1213            .unwrap()
1214            .collect()
1215            .await
1216            .unwrap();
1217
1218        let mut rows = Vec::new();
1219        for batch in &batches {
1220            let pts = batch
1221                .column(0)
1222                .as_any()
1223                .downcast_ref::<datafusion::arrow::array::StringViewArray>()
1224                .unwrap();
1225            let ids = batch
1226                .column(1)
1227                .as_any()
1228                .downcast_ref::<datafusion::arrow::array::Int32Array>()
1229                .unwrap();
1230            for i in 0..batch.num_rows() {
1231                rows.push((pts.value(i).to_string(), ids.value(i)));
1232            }
1233        }
1234        // Partition "a" overwritten with new data, partition "b" untouched
1235        assert_eq!(
1236            rows,
1237            vec![
1238                ("a".to_string(), 10),
1239                ("a".to_string(), 20),
1240                ("b".to_string(), 3),
1241                ("b".to_string(), 4),
1242            ]
1243        );
1244    }
1245
1246    #[tokio::test]
1247    async fn test_insert_overwrite_unpartitioned() {
1248        use paimon::io::FileIOBuilder;
1249        use paimon::spec::{DataType, IntType, Schema as PaimonSchema, TableSchema};
1250
1251        let file_io = FileIOBuilder::new("memory").build().unwrap();
1252        let table_path = "memory:/test_df_insert_overwrite_unpart";
1253        file_io
1254            .mkdirs(&format!("{table_path}/snapshot/"))
1255            .await
1256            .unwrap();
1257        file_io
1258            .mkdirs(&format!("{table_path}/manifest/"))
1259            .await
1260            .unwrap();
1261
1262        let schema = PaimonSchema::builder()
1263            .column("id", DataType::Int(IntType::new()))
1264            .column("value", DataType::Int(IntType::new()))
1265            .build()
1266            .unwrap();
1267        let table_schema = TableSchema::new(0, &schema);
1268        let table = paimon::table::Table::new(
1269            file_io,
1270            Identifier::new("default", "test_overwrite_unpart"),
1271            table_path.to_string(),
1272            table_schema,
1273            None,
1274        );
1275
1276        let provider = PaimonTableProvider::try_new(table).unwrap();
1277        let ctx = SessionContext::new();
1278        ctx.register_table("t", Arc::new(provider)).unwrap();
1279
1280        // Initial INSERT
1281        ctx.sql("INSERT INTO t VALUES (1, 10), (2, 20), (3, 30)")
1282            .await
1283            .unwrap()
1284            .collect()
1285            .await
1286            .unwrap();
1287
1288        // INSERT OVERWRITE on unpartitioned table — full table overwrite
1289        ctx.sql("INSERT OVERWRITE t VALUES (4, 40), (5, 50)")
1290            .await
1291            .unwrap()
1292            .collect()
1293            .await
1294            .unwrap();
1295
1296        let batches = ctx
1297            .sql("SELECT id, value FROM t ORDER BY id")
1298            .await
1299            .unwrap()
1300            .collect()
1301            .await
1302            .unwrap();
1303
1304        let mut rows = Vec::new();
1305        for batch in &batches {
1306            let ids = batch
1307                .column(0)
1308                .as_any()
1309                .downcast_ref::<datafusion::arrow::array::Int32Array>()
1310                .unwrap();
1311            let vals = batch
1312                .column(1)
1313                .as_any()
1314                .downcast_ref::<datafusion::arrow::array::Int32Array>()
1315                .unwrap();
1316            for i in 0..batch.num_rows() {
1317                rows.push((ids.value(i), vals.value(i)));
1318            }
1319        }
1320        // Old data fully replaced
1321        assert_eq!(rows, vec![(4, 40), (5, 50)]);
1322    }
1323}
paimon_datafusion/table/mod.rs

paimon_datafusion/table/
mod.rs
