From 2d40c533ad538cd7f19c7e2b2f5563d324e51737 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dapeng=20Sun=28=E5=AD=99=E5=A4=A7=E9=B9=8F=29?= Date: Sat, 1 Aug 2026 00:14:17 +0800 Subject: [PATCH] feat(datafusion): support REST format table partition commands Add Spark-style partition administration to DataFusion for catalog-managed internal Format Tables loaded from REST Catalog. REST partition registrations become the authoritative partition set used by both SQL commands and scans. - SHOW PARTITIONS table [PARTITION (...)] - ALTER TABLE table ADD [IF NOT EXISTS] PARTITION (...) - ALTER TABLE table DROP [IF EXISTS] PARTITION (...) - MSCK REPAIR TABLE table [{ADD|DROP|SYNC} PARTITIONS] DROP PARTITION follows Java Format Table semantics: a statement may carry several specifications, and one that fixes only some of the partition keys expands to every registered partition it matches, whether or not those keys are a leading prefix. One catalog listing serves the whole statement. A catalog-managed scan pushes the leading equality prefix of its filter down as a partitionNamePattern with a 1000-row page size, instead of downloading every registration and pruning locally. The pattern is a hint, so the local match stays as the backstop. Partition boolean values accept every spelling Java accepts (t/true/y/yes/1 and their negatives, case-insensitive), so a partition another engine registered as TRUE is readable here. --- Cargo.lock | 1 + crates/integrations/datafusion/Cargo.toml | 1 + .../datafusion/src/format_partition_repair.rs | 108 +++ crates/integrations/datafusion/src/lib.rs | 1 + .../datafusion/src/sql_context.rs | 878 ++++++++++++++++-- .../tests/rest_format_partition_sql.rs | 408 ++++++++ crates/paimon/src/api/api_request.rs | 120 ++- crates/paimon/src/api/mod.rs | 3 +- crates/paimon/src/api/resource_paths.rs | 14 + crates/paimon/src/api/rest_api.rs | 104 ++- crates/paimon/src/catalog/mod.rs | 32 + .../paimon/src/catalog/rest/rest_catalog.rs | 110 ++- crates/paimon/src/spec/core_options.rs | 63 +- crates/paimon/src/spec/mod.rs | 2 +- crates/paimon/src/spec/partition_utils.rs | 38 +- crates/paimon/src/table/format_partition.rs | 441 +++++++++ crates/paimon/src/table/format_table_scan.rs | 408 ++++---- crates/paimon/src/table/mod.rs | 21 +- crates/paimon/src/table/rest_env.rs | 148 ++- crates/paimon/tests/format_partition_test.rs | 160 ++++ crates/paimon/tests/mock_server.rs | 370 +++++++- crates/paimon/tests/rest_api_test.rs | 145 ++- crates/paimon/tests/rest_catalog_test.rs | 482 ++++++++++ docs/src/sql.md | 81 +- 24 files changed, 3856 insertions(+), 283 deletions(-) create mode 100644 crates/integrations/datafusion/src/format_partition_repair.rs create mode 100644 crates/integrations/datafusion/tests/rest_format_partition_sql.rs create mode 100644 crates/paimon/src/table/format_partition.rs create mode 100644 crates/paimon/tests/format_partition_test.rs diff --git a/Cargo.lock b/Cargo.lock index 69d011fa9..d2dbebede 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4607,6 +4607,7 @@ dependencies = [ "arrow-schema", "arrow-select", "async-trait", + "axum", "bytes", "chrono", "constant_time_eq", diff --git a/crates/integrations/datafusion/Cargo.toml b/crates/integrations/datafusion/Cargo.toml index 17acaea02..f7072fc36 100644 --- a/crates/integrations/datafusion/Cargo.toml +++ b/crates/integrations/datafusion/Cargo.toml @@ -50,6 +50,7 @@ uuid = { version = "1", features = ["v4"] } [dev-dependencies] arrow-array = { workspace = true } arrow-schema = { workspace = true } +axum = { version = "0.7", features = ["macros", "tokio", "http1", "http2"] } bytes = "1.7.1" flate2 = "1" paimon-ftindex-core = "0.1.0" diff --git a/crates/integrations/datafusion/src/format_partition_repair.rs b/crates/integrations/datafusion/src/format_partition_repair.rs new file mode 100644 index 000000000..f0b1738fd --- /dev/null +++ b/crates/integrations/datafusion/src/format_partition_repair.rs @@ -0,0 +1,108 @@ +// 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. + +use std::collections::{BTreeMap, HashMap}; + +use paimon::catalog::{Catalog, Identifier}; +use paimon::spec::CoreOptions; +use paimon::table::{FormatTablePartitionPaths, Table}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum RepairMode { + Add, + Drop, + Sync, +} + +pub(crate) async fn repair( + catalog: &dyn Catalog, + identifier: &Identifier, + table: &Table, + mode: RepairMode, +) -> paimon::Result<()> { + let core_options = CoreOptions::new(table.schema().options()); + let partition_paths = FormatTablePartitionPaths::new( + table.schema().partition_keys().iter().cloned(), + core_options.format_table_partition_only_value_in_path(), + ); + let table_path = table.location(); + + if matches!(mode, RepairMode::Drop | RepairMode::Sync) { + // Discovery treats a missing root as empty. Destructive repair must fail + // instead, or it could unregister every catalog partition. + table.file_io().list_status(table_path).await?; + } + + // Load both views before changing catalog metadata so a listing failure leaves + // metadata unchanged. Discovery preserves raw directory values such as month=01. + let discovered_specs = partition_paths + .discover( + table.file_io(), + table_path, + core_options.partition_default_name(), + ) + .await?; + let registered_partitions = catalog.list_partitions(identifier).await?; + + let discovered_by_name = index_specs_by_name(&partition_paths, discovered_specs)?; + let registered_by_name = index_specs_by_name( + &partition_paths, + registered_partitions + .into_iter() + .map(|partition| partition.spec) + .collect(), + )?; + + let to_register = if matches!(mode, RepairMode::Add | RepairMode::Sync) { + discovered_by_name + .iter() + .filter(|(name, _)| !registered_by_name.contains_key(*name)) + .map(|(_, spec)| spec.clone()) + .collect::>() + } else { + Vec::new() + }; + let to_unregister = if matches!(mode, RepairMode::Drop | RepairMode::Sync) { + registered_by_name + .iter() + .filter(|(name, _)| !discovered_by_name.contains_key(*name)) + .map(|(_, spec)| spec.clone()) + .collect::>() + } else { + Vec::new() + }; + + if !to_register.is_empty() { + catalog + .create_partitions(identifier, to_register, true) + .await?; + } + if !to_unregister.is_empty() { + catalog.drop_partitions(identifier, to_unregister).await?; + } + Ok(()) +} + +fn index_specs_by_name( + partition_paths: &FormatTablePartitionPaths, + specs: Vec>, +) -> paimon::Result>> { + specs + .into_iter() + .map(|spec| Ok((partition_paths.partition_name(&spec)?, spec))) + .collect() +} diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index 05e4037a2..ae8a42cac 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -43,6 +43,7 @@ mod catalog; mod delete; mod error; mod filter_pushdown; +mod format_partition_repair; #[cfg(feature = "fulltext")] mod full_text_search; mod hybrid_search; diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 1803892b2..ada8467e1 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -34,7 +34,10 @@ //! - `ALTER TABLE db.t ALTER COLUMN col TYPE new_type` //! - `ALTER TABLE db.t ALTER COLUMN col SET|DROP NOT NULL` //! - `ALTER TABLE db.t RENAME TO new_name` -//! - `ALTER TABLE db.t DROP PARTITION (col = val, ...)` +//! - `ALTER TABLE db.t ADD [IF NOT EXISTS] PARTITION (...) [PARTITION (...)]` +//! - `ALTER TABLE db.t DROP [IF EXISTS] PARTITION (...)` +//! - `SHOW PARTITIONS db.t [PARTITION (...)]` +//! - `MSCK REPAIR TABLE db.t [{ADD|DROP|SYNC} PARTITIONS]` //! - `CREATE VIEW [IF NOT EXISTS] view [(col, ...)] AS query` //! - `DROP VIEW [IF EXISTS] view` //! - `CREATE FUNCTION name(args) RETURNS type [LANGUAGE SQL] RETURN expression` @@ -52,7 +55,7 @@ use datafusion::arrow::compute::cast; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; -use datafusion::common::TableReference; +use datafusion::common::{ScalarValue, TableReference}; use datafusion::datasource::{MemTable, TableProvider}; use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::execution::SessionStateBuilder; @@ -60,27 +63,32 @@ use datafusion::logical_expr::{Expr as LogicalExpr, LogicalPlan, Volatility}; use datafusion::prelude::{DataFrame, SessionContext}; use datafusion::sql::planner::IdentNormalizer; use datafusion::sql::sqlparser::ast::{ - AlterColumnOperation, AlterTableOperation, BinaryLength, CharacterLength, ColumnDef, - ColumnOption, CreateFunction, CreateFunctionBody, CreateTable, CreateTableOptions, CreateView, - Delete, Expr as SqlExpr, FromTable, FunctionBehavior, FunctionReturnType, Insert, Merge, - ObjectName, ObjectType, RenameTableNameKind, Reset, ResetStatement, Set, ShowCreateObject, - SqlOption, Statement, TableFactor, TableObject, Truncate, Update, Use, Value as SqlValue, + AddDropSync, AlterColumnOperation, AlterTableOperation, BinaryLength, CharacterLength, + ColumnDef, ColumnOption, CreateFunction, CreateFunctionBody, CreateTable, CreateTableOptions, + CreateView, Delete, Expr as SqlExpr, FromTable, FunctionBehavior, FunctionReturnType, Insert, + Merge, Msck, ObjectName, ObjectType, Partition as SqlPartition, RenameTableNameKind, Reset, + ResetStatement, Set, ShowCreateObject, SqlOption, Statement, TableFactor, TableObject, + Truncate, Update, Use, Value as SqlValue, }; use datafusion::sql::sqlparser::dialect::GenericDialect; use datafusion::sql::sqlparser::keywords::Keyword; use datafusion::sql::sqlparser::parser::Parser; use datafusion::sql::sqlparser::tokenizer::{Token, Tokenizer}; use futures::StreamExt; -use paimon::catalog::{parse_object_name, Catalog, Identifier}; +use paimon::catalog::{parse_object_name, Catalog, Identifier, ParsedObjectName}; use paimon::spec::{ ArrayType as PaimonArrayType, BigIntType, BinaryType, BlobType, BooleanType, CharType, - DataField as PaimonDataField, DataType as PaimonDataType, DateType, Datum, DecimalType, - DoubleType, FloatType, IntType, LocalZonedTimestampType, MapType as PaimonMapType, + CoreOptions, DataField as PaimonDataField, DataType as PaimonDataType, DateType, Datum, + DecimalType, DoubleType, FloatType, IntType, LocalZonedTimestampType, MapType as PaimonMapType, RowType as PaimonRowType, SchemaChange, SmallIntType, TimestampType, TinyIntType, VarBinaryType, VarCharType, VariantType, }; +use paimon::table::{ + format_partition_value, parse_format_partition_value, FormatTablePartitionPaths, +}; use crate::error::to_datafusion_error; +use crate::filter_pushdown::scalar_to_datum; use crate::table_loader::load_table_for_read; use crate::{BlobReaderRegistry, DynamicOptions}; @@ -377,7 +385,9 @@ impl SQLContext { // Time-travel queries are not DDL; skip our own parsing and handle directly. return self.handle_time_travel_query(&rewritten_sql).await; } - + if let Some(show_partitions) = parse_show_partitions(&rewritten_sql)? { + return self.handle_show_partitions(&show_partitions).await; + } let statements = parse_sql_statements(&rewritten_sql)?; if statements.len() != 1 { @@ -444,6 +454,15 @@ impl SQLContext { obj_name, } => self.handle_show_create_table(sql, obj_name).await, Statement::AlterTable(alter_table) => { + if alter_table.location.is_some() + && alter_table.operations.iter().any(|operation| { + matches!(operation, AlterTableOperation::AddPartitions { .. }) + }) + { + return Err(DataFusionError::Plan( + "LOCATION is not supported for Format Table partitions".to_string(), + )); + } let (catalog, _catalog_name, _) = self.resolve_catalog_and_table(&alter_table.name)?; self.handle_alter_table( @@ -498,6 +517,7 @@ impl SQLContext { self.ctx.sql(sql).await } Statement::Truncate(truncate) => self.handle_truncate_table(truncate).await, + Statement::Msck(msck) => self.handle_msck(msck).await, Statement::CreateView(create_view) => { if create_view.temporary { // Temporary views are always handled by us (Paimon catalog temp storage) @@ -1165,20 +1185,58 @@ impl SQLContext { operations: &[AlterTableOperation], if_exists: bool, ) -> DFResult { - Self::ensure_main_branch_write_target(name, "ALTER TABLE")?; + let has_partition_operation = operations.iter().any(|operation| { + matches!( + operation, + AlterTableOperation::AddPartitions { .. } + | AlterTableOperation::DropPartitions { .. } + ) + }); + if has_partition_operation { + Self::ensure_partition_command_target(name, "ALTER TABLE")?; + } else { + Self::ensure_main_branch_write_target(name, "ALTER TABLE")?; + } let identifier = self.resolve_table_name(name)?; if operations.len() > 1 - && operations.iter().any(|operation| { - matches!( - operation, - AlterTableOperation::RenameTable { .. } - | AlterTableOperation::DropPartitions { .. } - ) + && operations + .iter() + .any(|operation| matches!(operation, AlterTableOperation::RenameTable { .. })) + { + return Err(DataFusionError::Plan( + "ALTER TABLE RENAME TO must be used alone".to_string(), + )); + } + // A statement may drop several partitions, but only partitions: mixing the drop + // with schema changes would commit two unrelated changes under one statement. + let drop_partition_requests = operations + .iter() + .filter_map(|operation| match operation { + AlterTableOperation::DropPartitions { + partitions, + if_exists: partition_if_exists, + } => Some((partitions.as_slice(), *partition_if_exists)), + _ => None, }) + .collect::>(); + if !drop_partition_requests.is_empty() { + if drop_partition_requests.len() != operations.len() { + return Err(DataFusionError::Plan( + "ALTER TABLE DROP PARTITION must be used alone".to_string(), + )); + } + return self + .handle_drop_partition(catalog, &identifier, &drop_partition_requests, if_exists) + .await; + } + if operations + .iter() + .any(|operation| matches!(operation, AlterTableOperation::AddPartitions { .. })) + && operations.len() != 1 { return Err(DataFusionError::Plan( - "ALTER TABLE RENAME TO and DROP PARTITION must be used alone".to_string(), + "ALTER TABLE ADD PARTITION cannot be combined with other operations".to_string(), )); } @@ -1228,16 +1286,17 @@ impl SQLContext { } } } - AlterTableOperation::DropPartitions { - partitions, - if_exists: partition_if_exists, + AlterTableOperation::AddPartitions { + if_not_exists, + new_partitions, } => { return self - .handle_drop_partitions( + .handle_add_partitions( catalog, &identifier, - partitions, - if_exists || *partition_if_exists, + new_partitions, + *if_not_exists, + if_exists, ) .await; } @@ -1797,34 +1856,50 @@ impl SQLContext { Ok(()) } - async fn handle_drop_partitions( + async fn handle_drop_partition( &self, catalog: &Arc, identifier: &Identifier, - partitions: &[SqlExpr], - if_exists: bool, + requests: &[(&[SqlExpr], bool)], + ignore_if_table_not_exists: bool, ) -> DFResult { - if partitions.is_empty() { + if requests + .iter() + .any(|(expressions, _)| expressions.is_empty()) + { return Err(DataFusionError::Plan( - "DROP PARTITIONS requires at least one partition specification".to_string(), + "DROP PARTITION requires a partition specification".to_string(), )); } let table = match catalog.get_table(identifier).await { - Ok(t) => t, - Err(e) if if_exists && is_table_not_exist(&e) => { + Ok(table) => table, + Err(error) if ignore_if_table_not_exists && is_table_not_exist(&error) => { return ok_result(&self.ctx); } - Err(e) => return Err(to_datafusion_error(e)), + Err(error) => return Err(to_datafusion_error(error)), }; - let partition_values = parse_partition_values( - partitions, - table.schema().fields(), - table.schema().partition_keys(), - )?; + if table.has_catalog_managed_partitions() { + return self + .drop_catalog_managed_partitions(catalog, identifier, &table, requests) + .await; + } + if CoreOptions::new(table.schema().options()).is_format_table() { + ensure_catalog_managed_format_table(&table, "ALTER TABLE DROP PARTITION")?; + } - let wb = table.new_write_builder(); - let commit = wb.try_new_commit().map_err(to_datafusion_error)?; + let mut partition_values = Vec::with_capacity(requests.len()); + for (expressions, _) in requests { + partition_values.extend(parse_partition_values( + expressions, + table.schema().fields(), + table.schema().partition_keys(), + )?); + } + let commit = table + .new_write_builder() + .try_new_commit() + .map_err(to_datafusion_error)?; commit .truncate_partitions(partition_values) .await @@ -1833,6 +1908,262 @@ impl SQLContext { ok_result(&self.ctx) } + /// Unregister catalog-managed partitions and then delete their directories. + /// + /// A specification that fixes only some of the partition keys expands to every + /// registered partition it matches, the way Java Format Tables behave. The keys need + /// not be a leading prefix, so `hh = '10'` alone is a valid specification. One catalog + /// listing serves the whole statement, however many specifications it carries. + async fn drop_catalog_managed_partitions( + &self, + catalog: &Arc, + identifier: &Identifier, + table: &paimon::Table, + requests: &[(&[SqlExpr], bool)], + ) -> DFResult { + ensure_catalog_managed_format_table(table, "ALTER TABLE DROP PARTITION")?; + let partition_key_count = table.schema().partition_keys().len(); + let mut requested = Vec::with_capacity(requests.len()); + for (expressions, ignore_if_not_exists) in requests { + let spec = parse_format_partition_spec(expressions, table, false)?; + let normalized = normalize_catalog_partition_spec(&spec, table, false)?; + requested.push(( + spec.len() == partition_key_count, + spec, + normalized, + *ignore_if_not_exists, + )); + } + + let registered = catalog + .list_partitions(identifier) + .await + .map_err(to_datafusion_error)? + .into_iter() + .map(|partition| { + let normalized = normalize_catalog_partition_spec(&partition.spec, table, true)?; + Ok((partition.spec, normalized)) + }) + .collect::>>()?; + + let core_options = CoreOptions::new(table.schema().options()); + let partition_paths = FormatTablePartitionPaths::new( + table.schema().partition_keys().iter().cloned(), + core_options.format_table_partition_only_value_in_path(), + ); + let table_path = table.location().trim_end_matches('/'); + let mut selected: Vec<(HashMap, String)> = Vec::new(); + let mut selected_paths = HashSet::new(); + for (complete, spec, normalized, ignore_if_not_exists) in &requested { + // A specification that is registered verbatim drops exactly that registration, + // so catalog values that only differ before normalization stay distinguishable. + let verbatim = registered.iter().any(|(registered, _)| registered == spec); + let mut matches = 0usize; + for (registered_spec, registered_normalized) in ®istered { + let selects = if verbatim { + registered_spec == spec + } else { + normalized + .iter() + .all(|(key, value)| registered_normalized.get(key) == Some(value)) + }; + if !selects { + continue; + } + matches += 1; + // Every directory is resolved before the first mutation, so a specification + // that cannot be turned into a path fails the statement as a whole. + let relative_path = partition_paths + .relative_path(registered_spec) + .map_err(to_datafusion_error)?; + let path = format!("{table_path}/{relative_path}"); + if selected_paths.insert(path.clone()) { + selected.push((registered_spec.clone(), path)); + } + } + if *complete && !verbatim && matches > 1 { + return Err(DataFusionError::Plan(format!( + "Partition {spec:?} matches multiple catalog registrations in table {}; use an exact catalog value", + identifier.full_name() + ))); + } + // Only a complete specification names one partition, so only it can be + // reported as missing. A partial one describes a set that is allowed to be + // empty, which is how Java reads it too. + if matches == 0 && *complete && !ignore_if_not_exists { + return Err(DataFusionError::Plan(format!( + "Partition {spec:?} does not exist in table {}", + identifier.full_name() + ))); + } + } + + if selected.is_empty() { + return ok_result(&self.ctx); + } + + catalog + .drop_partitions( + identifier, + selected.iter().map(|(spec, _)| spec.clone()).collect(), + ) + .await + .map_err(to_datafusion_error)?; + for (_, path) in selected { + table + .file_io() + .delete_dir(&path) + .await + .map_err(to_datafusion_error)?; + } + ok_result(&self.ctx) + } + + async fn handle_add_partitions( + &self, + catalog: &Arc, + identifier: &Identifier, + partitions: &[SqlPartition], + ignore_if_exists: bool, + ignore_if_table_not_exists: bool, + ) -> DFResult { + let table = match catalog.get_table(identifier).await { + Ok(table) => table, + Err(error) if ignore_if_table_not_exists && is_table_not_exist(&error) => { + return ok_result(&self.ctx); + } + Err(error) => return Err(to_datafusion_error(error)), + }; + ensure_catalog_managed_format_table(&table, "ALTER TABLE ADD PARTITION")?; + if partitions.is_empty() { + return Err(DataFusionError::Plan( + "ADD PARTITION requires at least one partition specification".to_string(), + )); + } + + let core_options = CoreOptions::new(table.schema().options()); + let partition_paths = FormatTablePartitionPaths::new( + table.schema().partition_keys().iter().cloned(), + core_options.format_table_partition_only_value_in_path(), + ); + let table_path = table.location(); + let mut specs = Vec::with_capacity(partitions.len()); + let mut directories = Vec::with_capacity(partitions.len()); + for partition in partitions { + let expressions = match partition { + SqlPartition::Partitions(expressions) => expressions, + other => { + return Err(DataFusionError::Plan(format!( + "Unsupported ADD PARTITION specification: {other}" + ))) + } + }; + let spec = parse_format_partition_spec(expressions, &table, true)?; + let relative_path = partition_paths + .relative_path(&spec) + .map_err(to_datafusion_error)?; + directories.push(format!( + "{}/{}", + table_path.trim_end_matches('/'), + relative_path + )); + specs.push(spec); + } + + catalog + .create_partitions(identifier, specs, ignore_if_exists) + .await + .map_err(to_datafusion_error)?; + for directory in directories { + table + .file_io() + .mkdirs(&directory) + .await + .map_err(to_datafusion_error)?; + } + ok_result(&self.ctx) + } + + async fn handle_msck(&self, msck: &Msck) -> DFResult { + if !msck.repair { + return Err(DataFusionError::Plan( + "MSCK requires the REPAIR keyword".to_string(), + )); + } + Self::ensure_partition_command_target(&msck.table_name, "MSCK REPAIR TABLE")?; + let (catalog, _catalog_name, identifier) = + self.resolve_catalog_and_table(&msck.table_name)?; + let table = catalog + .get_table(&identifier) + .await + .map_err(to_datafusion_error)?; + ensure_catalog_managed_format_table(&table, "MSCK REPAIR TABLE")?; + let mode = match msck.partition_action { + None | Some(AddDropSync::ADD) => crate::format_partition_repair::RepairMode::Add, + Some(AddDropSync::DROP) => crate::format_partition_repair::RepairMode::Drop, + Some(AddDropSync::SYNC) => crate::format_partition_repair::RepairMode::Sync, + }; + crate::format_partition_repair::repair(catalog.as_ref(), &identifier, &table, mode) + .await + .map_err(to_datafusion_error)?; + ok_result(&self.ctx) + } + + async fn handle_show_partitions( + &self, + show_partitions: &ShowPartitionsStatement, + ) -> DFResult { + Self::ensure_partition_command_target(&show_partitions.table_name, "SHOW PARTITIONS")?; + let (catalog, _catalog_name, identifier) = + self.resolve_catalog_and_table(&show_partitions.table_name)?; + let table = catalog + .get_table(&identifier) + .await + .map_err(to_datafusion_error)?; + ensure_catalog_managed_format_table(&table, "SHOW PARTITIONS")?; + let filter = if show_partitions.partition_filter.is_empty() { + None + } else { + let spec = + parse_format_partition_spec(&show_partitions.partition_filter, &table, false)?; + Some(normalize_catalog_partition_spec(&spec, &table, false)?) + }; + + let partition_paths = FormatTablePartitionPaths::new( + table.schema().partition_keys().iter().cloned(), + CoreOptions::new(table.schema().options()).format_table_partition_only_value_in_path(), + ); + let mut names = Vec::new(); + for partition in catalog + .list_partitions(&identifier) + .await + .map_err(to_datafusion_error)? + { + let normalized = normalize_catalog_partition_spec(&partition.spec, &table, true)?; + if !filter.as_ref().is_none_or(|filter| { + filter + .iter() + .all(|(key, value)| normalized.get(key) == Some(value)) + }) { + continue; + } + let display_spec = display_partition_spec(&normalized, &table)?; + let name = partition_paths + .partition_name(&display_spec) + .map_err(to_datafusion_error)?; + names.push(name); + } + names.sort(); + + let schema = Arc::new(Schema::new(vec![Field::new( + "partition", + ArrowDataType::Utf8, + false, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(StringArray::from(names))])?; + self.ctx.read_batch(batch) + } + /// Returns the name of the current default catalog from DataFusion config. pub(crate) fn current_catalog_name(&self) -> String { self.ctx @@ -2018,13 +2349,7 @@ impl SQLContext { } fn ensure_main_branch_write_target(name: &ObjectName, operation: &str) -> DFResult<()> { - let object = name - .0 - .last() - .and_then(|part| part.as_ident()) - .map(|ident| ident.value.as_str()) - .ok_or_else(|| DataFusionError::Plan(format!("Invalid table reference: {name}")))?; - let parsed = parse_object_name(object).map_err(to_datafusion_error)?; + let parsed = Self::parse_target_object_name(name)?; if let Some(branch) = parsed.branch() { return Err(DataFusionError::NotImplemented(format!( "{operation} on Paimon branch '{branch}' is not supported" @@ -2033,6 +2358,31 @@ impl SQLContext { Ok(()) } + fn ensure_partition_command_target(name: &ObjectName, operation: &str) -> DFResult<()> { + let parsed = Self::parse_target_object_name(name)?; + if let Some(branch) = parsed.branch() { + return Err(DataFusionError::NotImplemented(format!( + "{operation} on Paimon branch '{branch}' is not supported" + ))); + } + if let Some(system_table) = parsed.system_table() { + return Err(DataFusionError::NotImplemented(format!( + "{operation} on Paimon system table '{system_table}' is not supported" + ))); + } + Ok(()) + } + + fn parse_target_object_name(name: &ObjectName) -> DFResult { + let object = name + .0 + .last() + .and_then(|part| part.as_ident()) + .map(|ident| ident.value.as_str()) + .ok_or_else(|| DataFusionError::Plan(format!("Invalid table reference: {name}")))?; + parse_object_name(object).map_err(to_datafusion_error) + } + /// Resolve an ObjectName to just the Identifier (for backward compat in handle_alter_table). fn resolve_table_name(&self, name: &ObjectName) -> DFResult { let (_catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(name)?; @@ -2146,11 +2496,74 @@ fn validate_persistent_create_function(create_function: &CreateFunction) -> DFRe Ok(()) } +#[derive(Debug)] +struct ShowPartitionsStatement { + table_name: ObjectName, + partition_filter: Vec, +} + +fn parse_show_partitions(sql: &str) -> DFResult> { + let dialect = GenericDialect {}; + let tokens = Tokenizer::new(&dialect, sql) + .tokenize_with_location() + .map_err(sql_parse_error)?; + let significant = tokens + .iter() + .filter_map(|token| match &token.token { + Token::Whitespace(_) => None, + Token::Word(word) => Some(word.keyword), + _ => Some(Keyword::NoKeyword), + }) + .take(2) + .collect::>(); + if !matches!(significant.as_slice(), [Keyword::SHOW, Keyword::PARTITIONS]) { + return Ok(None); + } + + let mut parser = Parser::new(&dialect).with_tokens_with_locations(tokens); + parser + .expect_keyword_is(Keyword::SHOW) + .map_err(sql_parse_error)?; + parser + .expect_keyword_is(Keyword::PARTITIONS) + .map_err(sql_parse_error)?; + let table_name = parser.parse_object_name(false).map_err(sql_parse_error)?; + let partition_filter = if parser.parse_keyword(Keyword::PARTITION) { + parser + .expect_token(&Token::LParen) + .map_err(sql_parse_error)?; + let expressions = parser + .parse_comma_separated(Parser::parse_expr) + .map_err(sql_parse_error)?; + parser + .expect_token(&Token::RParen) + .map_err(sql_parse_error)?; + expressions + } else { + Vec::new() + }; + let _ = parser.consume_token(&Token::SemiColon); + if parser.peek_token().token != Token::EOF { + return Err(DataFusionError::Plan(format!( + "SQL parse error: unexpected token {} after SHOW PARTITIONS statement", + parser.peek_token().token + ))); + } + Ok(Some(ShowPartitionsStatement { + table_name, + partition_filter, + })) +} + +fn sql_parse_error(error: impl std::fmt::Display) -> DataFusionError { + DataFusionError::Plan(format!("SQL parse error: {error}")) +} + fn parse_sql_statements(sql: &str) -> DFResult> { let dialect = GenericDialect {}; let mut tokens = Tokenizer::new(&dialect, sql) .tokenize_with_location() - .map_err(|error| DataFusionError::Plan(format!("SQL parse error: {error}")))?; + .map_err(sql_parse_error)?; let significant = tokens .iter() .enumerate() @@ -2196,7 +2609,7 @@ fn parse_sql_statements(sql: &str) -> DFResult> { let mut statements = Parser::new(&dialect) .with_tokens_with_locations(tokens) .parse_statements() - .map_err(|error| DataFusionError::Plan(format!("SQL parse error: {error}")))?; + .map_err(sql_parse_error)?; if create_function_if_not_exists { let Some(Statement::CreateFunction(create_function)) = statements.first_mut() else { return Err(DataFusionError::Plan( @@ -2874,6 +3287,297 @@ fn is_table_not_exist(e: &paimon::Error) -> bool { matches!(e, paimon::Error::TableNotExist { .. }) } +fn ensure_catalog_managed_format_table(table: &paimon::Table, operation: &str) -> DFResult<()> { + if table.schema().partition_keys().is_empty() { + return Err(DataFusionError::Plan(format!( + "{operation} requires a partitioned table, but {} is not partitioned", + table.identifier().full_name() + ))); + } + if !table.has_catalog_managed_partitions() { + return Err(DataFusionError::Plan(format!( + "{operation} is supported only for catalog-managed partitions on an internal Format Table loaded from REST Catalog; table {} does not have this configuration", + table.identifier().full_name() + ))); + } + Ok(()) +} + +fn parse_format_partition_spec( + exprs: &[SqlExpr], + table: &paimon::Table, + require_complete: bool, +) -> DFResult> { + let fields = table + .schema() + .fields() + .iter() + .map(|field| (field.name(), field)) + .collect::>(); + let partition_keys = table.schema().partition_keys(); + let options = CoreOptions::new(table.schema().options()); + let mut spec = HashMap::with_capacity(exprs.len()); + + for expr in exprs { + let (column, value) = partition_assignment(expr)?; + if !partition_keys.contains(&column) { + return Err(DataFusionError::Plan(format!( + "Column '{column}' is not a partition column" + ))); + } + if spec.contains_key(&column) { + return Err(DataFusionError::Plan(format!( + "Duplicate partition column '{column}'" + ))); + } + let field = fields.get(column.as_str()).ok_or_else(|| { + DataFusionError::Plan(format!("Column '{column}' not found in table schema")) + })?; + let value = match parse_format_partition_literal(value, field.data_type())? { + None => options.partition_default_name().to_string(), + Some(datum) => format_partition_value( + &datum, + field.data_type(), + options.partition_default_name(), + options.legacy_partition_name(), + ) + .ok_or_else(|| { + DataFusionError::NotImplemented(format!( + "Partition literals of type {:?} are not supported", + field.data_type() + )) + })?, + }; + spec.insert(column, value); + } + + if require_complete { + let missing = partition_keys + .iter() + .filter(|key| !spec.contains_key(key.as_str())) + .cloned() + .collect::>(); + if !missing.is_empty() { + return Err(DataFusionError::Plan(format!( + "Incomplete partition spec: missing keys [{}]", + missing.join(", ") + ))); + } + } + Ok(spec) +} + +fn parse_format_partition_literal( + expr: &SqlExpr, + data_type: &PaimonDataType, +) -> DFResult> { + if matches!( + expr, + SqlExpr::Value(value) if matches!(&value.value, SqlValue::Null) + ) { + return Ok(None); + } + + let mut value = format_partition_literal_as_string(expr)?; + if !matches!( + data_type, + PaimonDataType::Char(_) | PaimonDataType::VarChar(_) + ) { + value = value.trim().to_string(); + } + let arrow_type = paimon::arrow::paimon_type_to_arrow(data_type).map_err(to_datafusion_error)?; + let scalar = ScalarValue::Utf8(Some(value)) + .cast_to(&arrow_type) + .map_err(|error| { + DataFusionError::Plan(format!( + "Cannot convert partition literal {expr} to {data_type:?}: {error}" + )) + })?; + scalar_to_datum(&scalar, data_type) + .map(Some) + .ok_or_else(|| { + DataFusionError::Plan(format!( + "Partition literals of type {data_type:?} are not supported" + )) + }) +} + +fn format_partition_literal_as_string(expr: &SqlExpr) -> DFResult { + match expr { + SqlExpr::TypedString(typed) + if matches!( + typed.data_type, + datafusion::sql::sqlparser::ast::DataType::Date + ) => + { + typed + .value + .value + .clone() + .into_string() + .ok_or_else(|| DataFusionError::Plan(format!("Invalid partition literal: {expr}"))) + } + SqlExpr::Value(value) => match &value.value { + SqlValue::Number(value, _) => Ok(canonicalize_partition_number(value)), + SqlValue::Boolean(value) => Ok(value.to_string()), + SqlValue::Null => Err(DataFusionError::Internal( + "NULL partition literal reached string conversion".to_string(), + )), + _ => value.value.clone().into_string().ok_or_else(|| { + DataFusionError::Plan(format!("Unsupported partition literal: {expr}")) + }), + }, + SqlExpr::UnaryOp { op, expr } + if matches!( + op, + datafusion::sql::sqlparser::ast::UnaryOperator::Plus + | datafusion::sql::sqlparser::ast::UnaryOperator::Minus + ) => + { + let SqlExpr::Value(value) = expr.as_ref() else { + return Err(DataFusionError::Plan(format!( + "Unary signs are supported only for numeric partition literals: {expr}" + ))); + }; + let SqlValue::Number(value, _) = &value.value else { + return Err(DataFusionError::Plan(format!( + "Unary signs are supported only for numeric partition literals: {expr}" + ))); + }; + let sign = if matches!(op, datafusion::sql::sqlparser::ast::UnaryOperator::Minus) { + "-" + } else { + "" + }; + Ok(format!("{sign}{}", canonicalize_partition_number(value))) + } + _ => Err(DataFusionError::Plan(format!( + "Unsupported partition value expression: {expr}" + ))), + } +} + +fn canonicalize_partition_number(value: &str) -> String { + value + .parse::() + .map(|value| value.to_string()) + .or_else(|_| value.parse::().map(|value| value.to_string())) + .unwrap_or_else(|_| value.to_string()) +} + +fn partition_assignment(expr: &SqlExpr) -> DFResult<(String, &SqlExpr)> { + let SqlExpr::BinaryOp { + left, + op: datafusion::sql::sqlparser::ast::BinaryOperator::Eq, + right, + } = expr + else { + return Err(DataFusionError::Plan(format!( + "Expected 'column = value' in partition spec, got: {expr}" + ))); + }; + let SqlExpr::Identifier(identifier) = left.as_ref() else { + return Err(DataFusionError::Plan(format!( + "Expected column name in partition spec, got: {left}" + ))); + }; + Ok(( + IdentNormalizer::default().normalize(identifier.clone()), + right.as_ref(), + )) +} + +fn normalize_catalog_partition_spec( + spec: &HashMap, + table: &paimon::Table, + require_complete: bool, +) -> DFResult>> { + let partition_keys = table.schema().partition_keys(); + if spec.keys().any(|key| !partition_keys.contains(key)) { + return Err(DataFusionError::Plan(format!( + "Invalid partition spec {spec:?} for table {}", + table.identifier().full_name() + ))); + } + if require_complete + && (spec.len() != partition_keys.len() + || partition_keys.iter().any(|key| !spec.contains_key(key))) + { + return Err(DataFusionError::Plan(format!( + "Invalid partition spec {spec:?} for table {}: expected keys {partition_keys:?}", + table.identifier().full_name() + ))); + } + + let fields = table + .schema() + .partition_fields() + .into_iter() + .map(|field| (field.name().to_string(), field)) + .collect::>(); + let options = CoreOptions::new(table.schema().options()); + let default_partition_name = options.partition_default_name(); + let mut normalized = HashMap::with_capacity(spec.len()); + for key in partition_keys.iter().filter(|key| spec.contains_key(*key)) { + let raw = &spec[key]; + let value = if raw == default_partition_name { + None + } else { + let field = fields.get(key).ok_or_else(|| { + DataFusionError::Plan(format!( + "Partition column '{key}' is missing from table schema" + )) + })?; + Some( + parse_format_partition_value(raw, field.data_type()).ok_or_else(|| { + DataFusionError::Plan(format!( + "Invalid catalog partition value {raw:?} for column '{key}' with type {:?}", + field.data_type() + )) + })?, + ) + }; + normalized.insert(key.clone(), value); + } + Ok(normalized) +} + +fn display_partition_spec( + spec: &HashMap>, + table: &paimon::Table, +) -> DFResult> { + let fields = table + .schema() + .partition_fields() + .into_iter() + .map(|field| (field.name().to_string(), field)) + .collect::>(); + let options = CoreOptions::new(table.schema().options()); + let default_partition_name = options.partition_default_name(); + spec.iter() + .map(|(key, value)| { + let value = match value { + None => "null".to_string(), + Some(datum) => { + let field = fields.get(key).ok_or_else(|| { + DataFusionError::Plan(format!( + "Partition column '{key}' is missing from table schema" + )) + })?; + format_partition_value(datum, field.data_type(), default_partition_name, false) + .ok_or_else(|| { + DataFusionError::NotImplemented(format!( + "SHOW PARTITIONS does not support type {:?}", + field.data_type() + )) + })? + } + }; + Ok((key.clone(), value)) + }) + .collect() +} + /// Parse partition expressions (`col = val, ...`) into partition value maps /// suitable for `TableCommit::truncate_partitions`. /// @@ -2889,30 +3593,9 @@ fn parse_partition_values( let mut partition = HashMap::new(); for expr in exprs { - let (col_name, val_expr) = match expr { - SqlExpr::BinaryOp { - left, - op: datafusion::sql::sqlparser::ast::BinaryOperator::Eq, - right, - } => { - let col = match left.as_ref() { - SqlExpr::Identifier(ident) => ident.value.clone(), - other => { - return Err(DataFusionError::Plan(format!( - "Expected column name in partition spec, got: {other}" - ))) - } - }; - (col, right.as_ref()) - } - other => { - return Err(DataFusionError::Plan(format!( - "Expected 'column = value' in partition spec, got: {other}" - ))) - } - }; + let (col_name, val_expr) = partition_assignment(expr)?; - if !partition_keys.iter().any(|k| k == &col_name) { + if !partition_keys.contains(&col_name) { return Err(DataFusionError::Plan(format!( "Column '{col_name}' is not a partition column" ))); @@ -3006,6 +3689,21 @@ fn parse_static_partitions( /// Convert a SQL literal expression to a Paimon Datum. fn sql_expr_to_datum(expr: &SqlExpr, data_type: &PaimonDataType) -> DFResult { + if let SqlExpr::TypedString(typed) = expr { + if matches!( + typed.data_type, + datafusion::sql::sqlparser::ast::DataType::Date + ) && matches!(data_type, PaimonDataType::Date(_)) + { + if let SqlValue::SingleQuotedString(value) = &typed.value.value { + return parse_date_datum(value); + } + } + return Err(DataFusionError::Plan(format!( + "Cannot convert {expr} to {data_type:?}" + ))); + } + let (value, negate) = match expr { SqlExpr::Value(v) => (&v.value, false), SqlExpr::UnaryOp { @@ -3029,14 +3727,13 @@ fn sql_expr_to_datum(expr: &SqlExpr, data_type: &PaimonDataType) -> DFResult parse_number_datum(n, data_type, negate), - (SqlValue::SingleQuotedString(s), PaimonDataType::VarChar(_)) if !negate => { + (SqlValue::SingleQuotedString(s), PaimonDataType::Char(_) | PaimonDataType::VarChar(_)) + if !negate => + { Ok(Datum::String(s.clone())) } (SqlValue::SingleQuotedString(s), PaimonDataType::Date(_)) if !negate => { - let date = chrono::NaiveDate::parse_from_str(s, "%Y-%m-%d") - .map_err(|e| DataFusionError::Plan(format!("Invalid DATE '{s}': {e}")))?; - let epoch = chrono::NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); - Ok(Datum::Date((date - epoch).num_days() as i32)) + parse_date_datum(s) } (SqlValue::Boolean(b), PaimonDataType::Boolean(_)) if !negate => Ok(Datum::Bool(*b)), _ if negate => Err(DataFusionError::Plan(format!( @@ -3048,6 +3745,19 @@ fn sql_expr_to_datum(expr: &SqlExpr, data_type: &PaimonDataType) -> DFResult DFResult { + let date = chrono::NaiveDate::parse_from_str(value, "%Y-%m-%d") + .map_err(|e| DataFusionError::Plan(format!("Invalid DATE '{value}': {e}")))?; + let unix_epoch = + chrono::NaiveDate::from_ymd_opt(1970, 1, 1).expect("1970-01-01 is a valid date"); + let epoch_days = i32::try_from((date - unix_epoch).num_days()).map_err(|e| { + DataFusionError::Plan(format!( + "DATE '{value}' is outside the supported partition range: {e}" + )) + })?; + Ok(Datum::Date(epoch_days)) +} + fn parse_number_datum(n: &str, data_type: &PaimonDataType, negate: bool) -> DFResult { let s: String = if negate { format!("-{n}") @@ -7212,6 +7922,20 @@ mod tests { .unwrap(); } + #[tokio::test] + async fn test_drop_if_exists_partition_does_not_ignore_missing_table() { + let (_tmp, sql_context) = setup_fs_sql_context().await; + + let err = sql_context + .sql("ALTER TABLE paimon.test_db.nonexistent DROP IF EXISTS PARTITION (pt = 'a')") + .await + .unwrap_err(); + assert!( + err.to_string().contains("does not exist"), + "Expected table-not-exist error, got: {err}" + ); + } + #[tokio::test] async fn test_drop_partition_incomplete_spec() { let (_tmp, sql_context) = setup_fs_sql_context().await; diff --git a/crates/integrations/datafusion/tests/rest_format_partition_sql.rs b/crates/integrations/datafusion/tests/rest_format_partition_sql.rs new file mode 100644 index 000000000..f879aa83c --- /dev/null +++ b/crates/integrations/datafusion/tests/rest_format_partition_sql.rs @@ -0,0 +1,408 @@ +// 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. + +#[path = "../../../paimon/tests/mock_server.rs"] +mod mock_server; + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_array::StringArray; +use paimon::api::ConfigResponse; +use paimon::catalog::RESTCatalog; +use paimon::spec::{BigIntType, BooleanType, DataType, DateType, IntType, Schema, VarCharType}; +use paimon::{CatalogOptions, Options}; +use paimon_datafusion::SQLContext; +use tempfile::TempDir; + +use mock_server::{start_mock_server, RESTServer}; + +const DATABASE: &str = "default"; +const TABLE: &str = "events"; +const WAREHOUSE: &str = "test_warehouse"; + +async fn setup_rest_table(temp_dir: &TempDir, schema: Schema) -> (RESTServer, SQLContext) { + let server = start_mock_server( + WAREHOUSE.to_string(), + temp_dir.path().to_string_lossy().into_owned(), + ConfigResponse::new(HashMap::from([( + CatalogOptions::PREFIX.to_string(), + "mock-test".to_string(), + )])), + vec![DATABASE.to_string()], + ) + .await; + server.add_table_with_schema( + DATABASE, + TABLE, + schema, + &format!("file://{}", temp_dir.path().display()), + ); + server.set_table_external(DATABASE, TABLE, false); + + let mut options = Options::new(); + options.set(CatalogOptions::URI, server.url().unwrap()); + options.set(CatalogOptions::WAREHOUSE, WAREHOUSE); + options.set(CatalogOptions::TOKEN_PROVIDER, "bear"); + options.set(CatalogOptions::TOKEN, "test-token"); + let catalog = Arc::new(RESTCatalog::new(options, true).await.unwrap()); + let mut context = SQLContext::new(); + context.register_catalog("paimon", catalog).await.unwrap(); + (server, context) +} + +fn format_table_schema(partition_columns: &[(&str, DataType)]) -> Schema { + let partition_keys = partition_columns + .iter() + .map(|(name, _)| (*name).to_string()) + .collect::>(); + partition_columns + .iter() + .fold(Schema::builder(), |builder, (name, data_type)| { + builder.column(*name, data_type.clone()) + }) + .column("id", DataType::BigInt(BigIntType::new())) + .partition_keys(partition_keys) + .option("type", "format-table") + .option("file.format", "parquet") + .option("metastore.partitioned-table", "true") + .build() + .unwrap() +} + +fn assert_partition_directories(temp_dir: &TempDir, expected: &[(&str, bool)]) { + for (path, exists) in expected { + assert_eq!( + temp_dir.path().join(path).exists(), + *exists, + "partition directory {path}" + ); + } +} + +#[cfg(not(windows))] +#[tokio::test] +async fn test_partition_commands_update_rest_metadata_and_directories() { + let temp_dir = tempfile::tempdir().unwrap(); + std::fs::create_dir_all(temp_dir.path().join("dt=2026-07-21")).unwrap(); + let schema = format_table_schema(&[("dt", DataType::VarChar(VarCharType::new(255).unwrap()))]); + let (_server, context) = setup_rest_table(&temp_dir, schema).await; + + context + .sql("ALTER TABLE paimon.default.events ADD PARTITION (dt = '2026-07-22')") + .await + .unwrap(); + assert_partitions(&context, &["dt=2026-07-22"]).await; + assert_partition_directories(&temp_dir, &[("dt=2026-07-22", true)]); + + context + .sql("MSCK REPAIR TABLE paimon.default.events ADD PARTITIONS") + .await + .unwrap(); + assert_partitions(&context, &["dt=2026-07-21", "dt=2026-07-22"]).await; + + std::fs::remove_dir_all(temp_dir.path().join("dt=2026-07-22")).unwrap(); + context + .sql("MSCK REPAIR TABLE paimon.default.events SYNC PARTITIONS") + .await + .unwrap(); + assert_partitions(&context, &["dt=2026-07-21"]).await; + + context + .sql("ALTER TABLE paimon.default.events ADD PARTITION (dt = '2026-07-22')") + .await + .unwrap(); + context + .sql("ALTER TABLE paimon.default.events DROP PARTITION (dt = '2026-07-21')") + .await + .unwrap(); + context + .sql("ALTER TABLE paimon.default.events DROP PARTITION (dt = '2026-07-22')") + .await + .unwrap(); + assert_partitions(&context, &[]).await; + assert_partition_directories( + &temp_dir, + &[("dt=2026-07-21", false), ("dt=2026-07-22", false)], + ); +} + +#[cfg(not(windows))] +#[tokio::test] +async fn test_partition_literals_and_default_path() { + let temp_dir = tempfile::tempdir().unwrap(); + let schema = format_table_schema(&[ + ("dt", DataType::Date(DateType::new())), + ("month", DataType::Int(IntType::new())), + ("active", DataType::Boolean(BooleanType::new())), + ("label", DataType::VarChar(VarCharType::new(255).unwrap())), + ]); + let (_server, context) = setup_rest_table(&temp_dir, schema).await; + + context + .sql( + "ALTER TABLE paimon.default.events ADD PARTITION (\ + dt = DATE '2026-07-22', month = '01', active = 'TRUE', label = 20260722)", + ) + .await + .unwrap(); + assert_partitions( + &context, + &["dt=2026-07-22/month=1/active=true/label=20260722"], + ) + .await; + // Non-legacy DATE partition paths use Unix epoch days. + assert_partition_directories( + &temp_dir, + &[("dt=20656/month=1/active=true/label=20260722", true)], + ); + + context + .sql( + "ALTER TABLE paimon.default.events DROP PARTITION (\ + dt = DATE '2026-07-22', month = '01', active = 'TRUE', label = 20260722)", + ) + .await + .unwrap(); + context + .sql( + "ALTER TABLE paimon.default.events ADD PARTITION (\ + dt = NULL, month = NULL, active = NULL, label = NULL)", + ) + .await + .unwrap(); + assert_partitions(&context, &["dt=null/month=null/active=null/label=null"]).await; + assert_partition_directories( + &temp_dir, + &[( + "dt=__DEFAULT_PARTITION__/month=__DEFAULT_PARTITION__/\ + active=__DEFAULT_PARTITION__/label=__DEFAULT_PARTITION__", + true, + )], + ); +} + +#[cfg(not(windows))] +#[tokio::test] +async fn test_drop_partition_takes_several_specs_and_expands_partial_ones() { + let temp_dir = tempfile::tempdir().unwrap(); + let schema = format_table_schema(&[ + ("dt", DataType::VarChar(VarCharType::new(255).unwrap())), + ("hh", DataType::VarChar(VarCharType::new(255).unwrap())), + ]); + let (_server, context) = setup_rest_table(&temp_dir, schema).await; + add_partitions(&context, &[("20260722", "10"), ("20260722", "11")]).await; + + // A partial spec expands to every registered partition it matches. + context + .sql("ALTER TABLE paimon.default.events DROP PARTITION (dt = '20260722')") + .await + .unwrap(); + assert_partitions(&context, &[]).await; + assert_partition_directories( + &temp_dir, + &[("dt=20260722/hh=10", false), ("dt=20260722/hh=11", false)], + ); + + // The fixed keys need not be a leading prefix. + add_partitions( + &context, + &[("20260722", "10"), ("20260722", "11"), ("20260723", "10")], + ) + .await; + context + .sql("ALTER TABLE paimon.default.events DROP PARTITION (hh = '10')") + .await + .unwrap(); + assert_partitions(&context, &["dt=20260722/hh=11"]).await; + assert_partition_directories( + &temp_dir, + &[ + ("dt=20260722/hh=10", false), + ("dt=20260723/hh=10", false), + ("dt=20260722/hh=11", true), + ], + ); + + // One statement may carry several specifications. + add_partitions(&context, &[("20260723", "10"), ("20260724", "12")]).await; + context + .sql( + "ALTER TABLE paimon.default.events \ + DROP PARTITION (dt = '20260722', hh = '11'), DROP PARTITION (dt = '20260723')", + ) + .await + .unwrap(); + assert_partitions(&context, &["dt=20260724/hh=12"]).await; +} + +#[cfg(not(windows))] +#[tokio::test] +async fn test_drop_partition_reports_a_specification_that_matches_nothing() { + let temp_dir = tempfile::tempdir().unwrap(); + let schema = format_table_schema(&[ + ("dt", DataType::VarChar(VarCharType::new(255).unwrap())), + ("hh", DataType::VarChar(VarCharType::new(255).unwrap())), + ]); + let (_server, context) = setup_rest_table(&temp_dir, schema).await; + add_partitions(&context, &[("20260722", "10")]).await; + + // A complete specification names one partition, so a missing one is an error. + let error = context + .sql("ALTER TABLE paimon.default.events DROP PARTITION (dt = '20260723', hh = '10')") + .await + .unwrap_err() + .to_string(); + assert!(error.contains("does not exist"), "{error}"); + assert_partitions(&context, &["dt=20260722/hh=10"]).await; + + context + .sql( + "ALTER TABLE paimon.default.events \ + DROP IF EXISTS PARTITION (dt = '20260723', hh = '10')", + ) + .await + .unwrap(); + assert_partitions(&context, &["dt=20260722/hh=10"]).await; + + // A partial specification describes a set that is allowed to come out empty, so it is + // a no-op rather than an error, and it does not take the statement down with it. + context + .sql( + "ALTER TABLE paimon.default.events \ + DROP PARTITION (dt = '20260723'), DROP PARTITION (dt = '20260722')", + ) + .await + .unwrap(); + assert_partitions(&context, &[]).await; + + // One failing specification must leave the whole statement unapplied. + add_partitions(&context, &[("20260722", "10")]).await; + let error = context + .sql( + "ALTER TABLE paimon.default.events \ + DROP PARTITION (dt = '20260722', hh = '10'), \ + DROP PARTITION (dt = '20260723', hh = '10')", + ) + .await + .unwrap_err() + .to_string(); + assert!(error.contains("does not exist"), "{error}"); + assert_partitions(&context, &["dt=20260722/hh=10"]).await; + + // Dropping partitions cannot be combined with a schema change. + let error = context + .sql("ALTER TABLE paimon.default.events DROP PARTITION (dt = '20260722'), ADD COLUMN c INT") + .await + .unwrap_err() + .to_string(); + assert!(error.contains("must be used alone"), "{error}"); + assert_partitions(&context, &["dt=20260722/hh=10"]).await; +} + +#[cfg(not(windows))] +// Planning a SELECT resolves the table on a blocking catalog-access thread, so the mock +// server needs a runtime thread of its own to answer while that one waits. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_catalog_managed_scan_pushes_a_partition_name_pattern() { + let temp_dir = tempfile::tempdir().unwrap(); + let schema = format_table_schema(&[ + ("dt", DataType::VarChar(VarCharType::new(255).unwrap())), + ("hh", DataType::VarChar(VarCharType::new(255).unwrap())), + ]); + let (server, context) = setup_rest_table(&temp_dir, schema).await; + add_partitions( + &context, + &[("20260722", "10"), ("20260722", "11"), ("20260723", "10")], + ) + .await; + + for (predicate, expected) in [ + ("dt = '20260722' AND hh = '10'", Some("dt=20260722/hh=10")), + ("dt = '20260722'", Some("dt=20260722/%")), + ( + "dt = '20260722' AND hh IN ('10', '11')", + Some("dt=20260722/%"), + ), + // Only a leading run of equalities becomes a prefix pattern; anything else has to + // list every partition and prune locally. + ("hh = '10'", None), + ("dt > '20260722'", None), + ] { + let seen = server + .table_partition_list_name_patterns(DATABASE, TABLE) + .len(); + context + .sql(&format!( + "SELECT * FROM paimon.default.events WHERE {predicate}" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + let patterns = server.table_partition_list_name_patterns(DATABASE, TABLE); + let pushed = &patterns[seen..]; + assert!(!pushed.is_empty(), "{predicate} listed no partitions"); + assert!( + pushed.iter().all(|pattern| pattern.as_deref() == expected), + "{predicate} pushed {pushed:?}, expected {expected:?}" + ); + } +} + +async fn add_partitions(context: &SQLContext, partitions: &[(&str, &str)]) { + for (dt, hh) in partitions { + context + .sql(&format!( + "ALTER TABLE paimon.default.events \ + ADD IF NOT EXISTS PARTITION (dt = '{dt}', hh = '{hh}')" + )) + .await + .unwrap(); + } +} + +async fn show_partitions(context: &SQLContext) -> Vec { + let batches = context + .sql("SHOW PARTITIONS paimon.default.events") + .await + .unwrap() + .collect() + .await + .unwrap(); + let mut partitions = Vec::new(); + for batch in batches { + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + partitions.extend((0..batch.num_rows()).map(|index| values.value(index).to_string())); + } + partitions +} + +async fn assert_partitions(context: &SQLContext, expected: &[&str]) { + assert_eq!( + show_partitions(context).await, + expected + .iter() + .map(|partition| (*partition).to_string()) + .collect::>() + ); +} diff --git a/crates/paimon/src/api/api_request.rs b/crates/paimon/src/api/api_request.rs index 84ee7c5c0..ea292b036 100644 --- a/crates/paimon/src/api/api_request.rs +++ b/crates/paimon/src/api/api_request.rs @@ -19,7 +19,7 @@ //! //! This module contains all request structures used in REST API calls. -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Deserializer, Serialize}; use std::collections::HashMap; use crate::{ @@ -167,6 +167,65 @@ impl AlterTableRequest { } } +/// Request to create table partitions. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CreatePartitionsRequest { + /// Partition specs to register. + pub partition_specs: Vec>, + /// Whether already registered partitions should be ignored. + #[serde( + default = "default_true", + deserialize_with = "deserialize_null_to_true" + )] + pub ignore_if_exists: bool, +} + +impl CreatePartitionsRequest { + /// Create a request to register partitions. + pub fn new(partition_specs: Vec>, ignore_if_exists: bool) -> Self { + Self { + partition_specs, + ignore_if_exists, + } + } +} + +/// Request to drop (unregister) table partitions. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct DropPartitionsRequest { + /// Partition specs to unregister. + pub partition_specs: Vec>, + /// Whether missing partitions should be ignored. + #[serde( + default = "default_true", + deserialize_with = "deserialize_null_to_true" + )] + pub ignore_if_not_exists: bool, +} + +impl DropPartitionsRequest { + /// Create a request to unregister partitions. + pub fn new(partition_specs: Vec>, ignore_if_not_exists: bool) -> Self { + Self { + partition_specs, + ignore_if_not_exists, + } + } +} + +fn default_true() -> bool { + true +} + +fn deserialize_null_to_true<'de, D>(deserializer: D) -> Result +where + D: Deserializer<'de>, +{ + Ok(Option::::deserialize(deserializer)?.unwrap_or(true)) +} + /// Request for auth table query: the projected columns of the query. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct AuthTableQueryRequest { @@ -220,6 +279,65 @@ mod tests { assert_eq!(serde_json::to_string(&req).unwrap(), "{}"); } + #[test] + fn test_create_partitions_request_serialization() { + let req = CreatePartitionsRequest::new( + vec![HashMap::from([ + ("dt".to_string(), "2026-07-22".to_string()), + ("hour".to_string(), "10".to_string()), + ])], + false, + ); + + assert_eq!( + serde_json::to_value(req).unwrap(), + serde_json::json!({ + "partitionSpecs": [{ + "dt": "2026-07-22", + "hour": "10" + }], + "ignoreIfExists": false + }) + ); + } + + #[test] + fn test_drop_partitions_request_serialization() { + let req = DropPartitionsRequest::new( + vec![HashMap::from([( + "dt".to_string(), + "2026-07-22".to_string(), + )])], + false, + ); + + assert_eq!( + serde_json::to_value(req).unwrap(), + serde_json::json!({ + "partitionSpecs": [{"dt": "2026-07-22"}], + "ignoreIfNotExists": false + }) + ); + } + + #[test] + fn test_partition_request_flags_default_to_true() { + for json in [ + serde_json::json!({"partitionSpecs": []}), + serde_json::json!({"partitionSpecs": [], "ignoreIfExists": null}), + ] { + let request: CreatePartitionsRequest = serde_json::from_value(json).unwrap(); + assert!(request.ignore_if_exists); + } + for json in [ + serde_json::json!({"partitionSpecs": []}), + serde_json::json!({"partitionSpecs": [], "ignoreIfNotExists": null}), + ] { + let request: DropPartitionsRequest = serde_json::from_value(json).unwrap(); + assert!(request.ignore_if_not_exists); + } + } + #[test] fn test_rename_table_request_serialization() { let source = Identifier::new("db1".to_string(), "table1".to_string()); diff --git a/crates/paimon/src/api/mod.rs b/crates/paimon/src/api/mod.rs index 809b90211..5773a85a7 100644 --- a/crates/paimon/src/api/mod.rs +++ b/crates/paimon/src/api/mod.rs @@ -32,7 +32,8 @@ mod api_response; // Re-export request types pub use api_request::{ AlterDatabaseRequest, AlterTableRequest, AuthTableQueryRequest, CreateDatabaseRequest, - CreateFunctionRequest, CreateTableRequest, CreateViewRequest, RenameTableRequest, + CreateFunctionRequest, CreatePartitionsRequest, CreateTableRequest, CreateViewRequest, + DropPartitionsRequest, RenameTableRequest, }; // Re-export response types diff --git a/crates/paimon/src/api/resource_paths.rs b/crates/paimon/src/api/resource_paths.rs index 38c39c9a6..569d1b92a 100644 --- a/crates/paimon/src/api/resource_paths.rs +++ b/crates/paimon/src/api/resource_paths.rs @@ -216,6 +216,11 @@ impl ResourcePaths { Self::PARTITIONS ) } + + /// Get the endpoint path for dropping table partitions. + pub fn drop_partitions(&self, database_name: &str, table_name: &str) -> String { + format!("{}/drop", self.partitions(database_name, table_name)) + } } #[cfg(test)] @@ -296,4 +301,13 @@ mod tests { "/v1/catalog/databases/analytics/functions/rectangle+area" ); } + + #[test] + fn test_drop_partitions_path_encodes_names() { + let paths = ResourcePaths::new("catalog"); + assert_eq!( + paths.drop_partitions("analytics db", "user events"), + "/v1/catalog/databases/analytics+db/tables/user+events/partitions/drop" + ); + } } diff --git a/crates/paimon/src/api/rest_api.rs b/crates/paimon/src/api/rest_api.rs index cdde2b23d..c3117ea7e 100644 --- a/crates/paimon/src/api/rest_api.rs +++ b/crates/paimon/src/api/rest_api.rs @@ -20,7 +20,7 @@ //! This module provides a REST API client for interacting with //! Paimon rest catalog services, supporting database operations. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use crate::api::rest_client::HttpClient; use crate::catalog::{Function, Identifier, ViewSchema}; @@ -30,7 +30,8 @@ use crate::Result; use super::api_request::{ AlterDatabaseRequest, AlterTableRequest, AuthTableQueryRequest, CreateDatabaseRequest, - CreateFunctionRequest, CreateTableRequest, CreateViewRequest, RenameTableRequest, + CreateFunctionRequest, CreatePartitionsRequest, CreateTableRequest, CreateViewRequest, + DropPartitionsRequest, RenameTableRequest, }; use super::api_response::{ AuthTableQueryResponse, ConfigResponse, GetDatabaseResponse, GetFunctionResponse, @@ -91,6 +92,9 @@ impl RESTApi { pub const TABLE_NAME_PATTERN: &'static str = "tableNamePattern"; pub const VIEW_NAME_PATTERN: &'static str = "viewNamePattern"; pub const FUNCTION_NAME_PATTERN: &'static str = "functionNamePattern"; + pub const PARTITION_NAME_PATTERN: &'static str = "partitionNamePattern"; + /// Bounds one request: catalog services cap the partitions a single call may carry. + const PARTITION_REQUEST_SIZE: u32 = 1000; pub const TABLE_TYPE: &'static str = "tableType"; /// Create a new RESTApi from options. @@ -556,36 +560,121 @@ impl RESTApi { // ==================== Partition Operations ==================== + /// Create table partitions in a single REST request. + pub async fn create_partitions( + &self, + identifier: &Identifier, + partition_specs: Vec>, + ignore_if_exists: bool, + ) -> Result<()> { + let database = identifier.database(); + let table = identifier.object(); + validate_non_empty_multi(&[(database, "database name"), (table, "table name")])?; + let path = self.resource_paths.partitions(database, table); + let request = CreatePartitionsRequest::new(partition_specs, ignore_if_exists); + let _resp: serde_json::Value = self.client.post(&path, &request).await?; + Ok(()) + } + + /// Unregister table partitions in a single REST request. + /// + /// The REST service removes metadata only; it does not delete partition + /// directories or data files. + pub async fn drop_partitions( + &self, + identifier: &Identifier, + partition_specs: Vec>, + ignore_if_not_exists: bool, + ) -> Result<()> { + let database = identifier.database(); + let table = identifier.object(); + validate_non_empty_multi(&[(database, "database name"), (table, "table name")])?; + let path = self.resource_paths.drop_partitions(database, table); + let request = DropPartitionsRequest::new(partition_specs, ignore_if_not_exists); + let _resp: serde_json::Value = self.client.post(&path, &request).await?; + Ok(()) + } + /// List all partitions of a table, paging internally. pub async fn list_partitions(&self, identifier: &Identifier) -> Result> { + self.drain_partitions(identifier, None, None).await + } + + /// List partitions, asking the catalog to return only those whose partition name + /// matches `partition_name_pattern`. + /// + /// The pattern is a pushdown hint: a catalog may apply it partially or not at all, so + /// the result is a superset of the matching partitions and never misses one. Callers + /// keep applying their own filter to what comes back. + pub async fn list_partitions_by_name_pattern( + &self, + identifier: &Identifier, + partition_name_pattern: Option<&str>, + ) -> Result> { + self.drain_partitions( + identifier, + Some(Self::PARTITION_REQUEST_SIZE), + partition_name_pattern, + ) + .await + } + + async fn drain_partitions( + &self, + identifier: &Identifier, + max_results: Option, + partition_name_pattern: Option<&str>, + ) -> Result> { let database = identifier.database(); let table = identifier.object(); validate_non_empty_multi(&[(database, "database name"), (table, "table name")])?; let mut results = Vec::new(); let mut page_token: Option = None; + let mut seen_page_tokens = HashSet::new(); loop { let paged = self - .list_partitions_paged(identifier, None, page_token.as_deref()) + .list_partitions_paged( + identifier, + max_results, + page_token.as_deref(), + partition_name_pattern, + ) .await?; - let is_empty = paged.elements.is_empty(); results.extend(paged.elements); - page_token = paged.next_page_token; - if page_token.is_none() || is_empty { + + let Some(next_page_token) = paged.next_page_token.filter(|token| !token.is_empty()) + else { break; + }; + if !seen_page_tokens.insert(next_page_token.clone()) { + return Err(crate::Error::UnexpectedError { + message: format!( + "REST catalog returned partition page token '{next_page_token}' more than \ + once for table {}", + identifier.full_name() + ), + source: None, + }); } + page_token = Some(next_page_token); } Ok(results) } /// List partitions with pagination. + /// + /// `partition_name_pattern` is a SQL LIKE pattern over partition names, where `%` is + /// the only wildcard. Like [`Self::list_partitions_by_name_pattern`], it is a hint the + /// catalog may ignore. pub async fn list_partitions_paged( &self, identifier: &Identifier, max_results: Option, page_token: Option<&str>, + partition_name_pattern: Option<&str>, ) -> Result> { let database = identifier.database(); let table = identifier.object(); @@ -599,6 +688,9 @@ impl RESTApi { if let Some(token) = page_token { params.push((Self::PAGE_TOKEN, token.to_string())); } + if let Some(pattern) = partition_name_pattern.filter(|pattern| !pattern.is_empty()) { + params.push((Self::PARTITION_NAME_PATTERN, pattern.to_string())); + } let response: ListPartitionsResponse = if params.is_empty() { self.client.get(&path, None::<&[(&str, &str)]>).await? diff --git a/crates/paimon/src/catalog/mod.rs b/crates/paimon/src/catalog/mod.rs index 38a414b97..bc9487b6b 100644 --- a/crates/paimon/src/catalog/mod.rs +++ b/crates/paimon/src/catalog/mod.rs @@ -445,6 +445,38 @@ pub trait Catalog: Send + Sync { }) } + // ======================= partition methods =============================== + + /// Register table partition specs in the catalog. + /// + /// When `ignore_if_exists` is false, implementations must reject the + /// entire request if any supplied spec already exists. When true, existing + /// specs are ignored so callers can safely retry the request. + async fn create_partitions( + &self, + _identifier: &Identifier, + _partition_specs: Vec>, + _ignore_if_exists: bool, + ) -> Result<()> { + Err(Error::Unsupported { + message: "Catalog does not support creating partitions".to_string(), + }) + } + + /// Unregister table partition metadata from the catalog. + /// + /// This does not delete partition directories or data files. Missing specs + /// are ignored so callers can safely retry the request. + async fn drop_partitions( + &self, + _identifier: &Identifier, + _partition_specs: Vec>, + ) -> Result<()> { + Err(Error::Unsupported { + message: "Catalog does not support dropping partitions".to_string(), + }) + } + /// List partitions for a table. /// /// Default impl scans the table's manifest entries via diff --git a/crates/paimon/src/catalog/rest/rest_catalog.rs b/crates/paimon/src/catalog/rest/rest_catalog.rs index 7a19ba12d..049a25930 100644 --- a/crates/paimon/src/catalog/rest/rest_catalog.rs +++ b/crates/paimon/src/catalog/rest/rest_catalog.rs @@ -38,6 +38,8 @@ use crate::spec::{Partition, Schema, SchemaChange}; use crate::table::{RESTEnv, Table}; use crate::Result; +const PARTITION_BATCH_SIZE: usize = 1000; + /// REST catalog implementation. /// /// This catalog communicates with a Paimon REST catalog server @@ -373,16 +375,64 @@ impl Catalog for RESTCatalog { )) } + async fn create_partitions( + &self, + identifier: &Identifier, + partition_specs: Vec>, + ignore_if_exists: bool, + ) -> Result<()> { + if partition_specs.is_empty() { + return Ok(()); + } + if !ignore_if_exists { + return self + .api + .create_partitions(identifier, partition_specs, false) + .await + .map_err(|error| map_rest_error_for_create_partitions(error, identifier)); + } + + for batch in partition_specs.chunks(PARTITION_BATCH_SIZE) { + self.api + .create_partitions(identifier, batch.to_vec(), true) + .await + .map_err(|error| map_rest_error_for_create_partitions(error, identifier))?; + } + Ok(()) + } + + async fn drop_partitions( + &self, + identifier: &Identifier, + partition_specs: Vec>, + ) -> Result<()> { + if partition_specs.is_empty() { + return Ok(()); + } + for batch in partition_specs.chunks(PARTITION_BATCH_SIZE) { + self.api + .drop_partitions(identifier, batch.to_vec(), true) + .await + .map_err(|error| map_rest_error_for_partition_request(error, identifier))?; + } + Ok(()) + } + async fn list_partitions(&self, identifier: &Identifier) -> Result> { match self.api.list_partitions(identifier).await { - Ok(parts) => Ok(parts), - Err(Error::RestApi { - source: RestError::NotImplemented { .. }, - }) => { + Ok(partitions) => Ok(partitions), + Err( + error @ Error::RestApi { + source: RestError::NotImplemented { .. }, + }, + ) => { let table = self.get_table(identifier).await?; + if table.has_catalog_managed_partitions() { + return Err(error); + } list_partitions_from_file_system(&table).await } - Err(e) => Err(map_rest_error_for_table(e, identifier)), + Err(error) => Err(map_rest_error_for_table(error, identifier)), } } @@ -394,21 +444,27 @@ impl Catalog for RESTCatalog { ) -> Result> { match self .api - .list_partitions_paged(identifier, max_results, page_token) + .list_partitions_paged(identifier, max_results, page_token, None) .await { Ok(page) => Ok(page), - Err(Error::RestApi { - source: RestError::NotImplemented { .. }, - }) => { + Err( + error @ Error::RestApi { + source: RestError::NotImplemented { .. }, + }, + ) => { let table = self.get_table(identifier).await?; - let parts = list_partitions_from_file_system(&table).await?; - Ok(PagedList::new(parts, None)) + if table.has_catalog_managed_partitions() { + return Err(error); + } + let partitions = list_partitions_from_file_system(&table).await?; + Ok(PagedList::new(partitions, None)) } - Err(e) => Err(map_rest_error_for_table(e, identifier)), + Err(error) => Err(map_rest_error_for_table(error, identifier)), } } } + // ============================================================================ // Error mapping helpers // ============================================================================ @@ -455,6 +511,36 @@ fn map_rest_error_for_table(err: Error, identifier: &Identifier) -> Error { } } +fn map_rest_error_for_create_partitions(err: Error, identifier: &Identifier) -> Error { + match err { + Error::RestApi { + source: RestError::AlreadyExists { message, .. }, + } => Error::DataInvalid { + message: format!( + "One or more partitions already exist for table {}: {message}", + identifier.full_name() + ), + source: None, + }, + other => map_rest_error_for_partition_request(other, identifier), + } +} + +fn map_rest_error_for_partition_request(err: Error, identifier: &Identifier) -> Error { + match err { + Error::RestApi { + source: RestError::BadRequest { message }, + } => Error::DataInvalid { + message: format!( + "Invalid partition request for table {}: {message}", + identifier.full_name() + ), + source: None, + }, + other => map_rest_error_for_table(other, identifier), + } +} + /// Map a REST API error from creating a persistent view. fn map_rest_error_for_create_view(err: Error, identifier: &Identifier) -> Error { match err { diff --git a/crates/paimon/src/spec/core_options.rs b/crates/paimon/src/spec/core_options.rs index 507471c88..cf29e44bd 100644 --- a/crates/paimon/src/spec/core_options.rs +++ b/crates/paimon/src/spec/core_options.rs @@ -33,7 +33,9 @@ const SOURCE_SPLIT_TARGET_SIZE_OPTION: &str = "source.split.target-size"; const SOURCE_SPLIT_OPEN_FILE_COST_OPTION: &str = "source.split.open-file-cost"; const PARTITION_DEFAULT_NAME_OPTION: &str = "partition.default-name"; const PARTITION_LEGACY_NAME_OPTION: &str = "partition.legacy-name"; -const FORMAT_TABLE_PARTITION_PATH_ONLY_VALUE_OPTION: &str = +pub(crate) const METASTORE_PARTITIONED_TABLE_OPTION: &str = "metastore.partitioned-table"; +pub(crate) const FORMAT_TABLE_IMPLEMENTATION_OPTION: &str = "format-table.implementation"; +pub(crate) const FORMAT_TABLE_PARTITION_PATH_ONLY_VALUE_OPTION: &str = "format-table.partition-path-only-value"; pub(crate) const BUCKET_KEY_OPTION: &str = "bucket-key"; const BUCKET_FUNCTION_TYPE_OPTION: &str = "bucket-function.type"; @@ -127,6 +129,12 @@ pub const BLOB_VIEW_RESOLVE_ENABLED_OPTION: &str = "blob-view.resolve.enabled"; const PK_VECTOR_INDEX_COLUMNS_OPTION: &str = "pk-vector.index.columns"; const PK_FULL_TEXT_INDEX_COLUMNS_OPTION: &str = "pk-full-text.index.columns"; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum FormatTableImplementation { + Paimon, + Engine, +} + /// Merge engine for primary-key tables. /// /// Reference: Java `CoreOptions.MergeEngine`. @@ -550,6 +558,43 @@ impl<'a> CoreOptions<'a> { .unwrap_or(false) } + pub(crate) fn try_format_table_partition_only_value_in_path(&self) -> crate::Result { + self.try_boolean_option(FORMAT_TABLE_PARTITION_PATH_ONLY_VALUE_OPTION) + } + + pub(crate) fn partitioned_table_in_metastore(&self) -> crate::Result { + self.try_boolean_option(METASTORE_PARTITIONED_TABLE_OPTION) + } + + pub(crate) fn format_table_implementation(&self) -> crate::Result { + match self + .options + .get(FORMAT_TABLE_IMPLEMENTATION_OPTION) + .map(String::as_str) + .unwrap_or("paimon") + { + value if value.eq_ignore_ascii_case("paimon") => Ok(FormatTableImplementation::Paimon), + value if value.eq_ignore_ascii_case("engine") => Ok(FormatTableImplementation::Engine), + value => Err(crate::Error::ConfigInvalid { + message: format!( + "Invalid value '{value}' for {FORMAT_TABLE_IMPLEMENTATION_OPTION}; \ + expected 'paimon' or 'engine'" + ), + }), + } + } + + fn try_boolean_option(&self, key: &str) -> crate::Result { + match self.options.get(key).map(String::as_str) { + None => Ok(false), + Some(value) if value.eq_ignore_ascii_case("true") => Ok(true), + Some(value) if value.eq_ignore_ascii_case("false") => Ok(false), + Some(value) => Err(crate::Error::ConfigInvalid { + message: format!("Invalid value '{value}' for {key}; expected 'true' or 'false'"), + }), + } + } + pub fn global_index_enabled(&self) -> bool { self.options .get(GLOBAL_INDEX_ENABLED_OPTION) @@ -1636,6 +1681,22 @@ mod tests { assert!(core.format_table_partition_only_value_in_path()); } + #[test] + fn test_partitioned_table_in_metastore_option() { + let empty = HashMap::new(); + assert!(!CoreOptions::new(&empty) + .partitioned_table_in_metastore() + .unwrap()); + + let options = HashMap::from([( + METASTORE_PARTITIONED_TABLE_OPTION.to_string(), + "TrUe".to_string(), + )]); + assert!(CoreOptions::new(&options) + .partitioned_table_in_metastore() + .unwrap()); + } + #[test] fn test_try_time_travel_selector_rejects_conflicting_selectors() { let options = HashMap::from([ diff --git a/crates/paimon/src/spec/mod.rs b/crates/paimon/src/spec/mod.rs index f1d3147f2..e79a158bf 100644 --- a/crates/paimon/src/spec/mod.rs +++ b/crates/paimon/src/spec/mod.rs @@ -94,7 +94,7 @@ pub use types::*; mod partition; pub use partition::Partition; mod partition_utils; -pub(crate) use partition_utils::PartitionComputer; +pub(crate) use partition_utils::{escape_path_name, unescape_path_name, PartitionComputer}; mod predicate; pub(crate) use predicate::datum_cmp; pub(crate) use predicate::eval_row; diff --git a/crates/paimon/src/spec/partition_utils.rs b/crates/paimon/src/spec/partition_utils.rs index 3db1884c0..56ab1d3dd 100644 --- a/crates/paimon/src/spec/partition_utils.rs +++ b/crates/paimon/src/spec/partition_utils.rs @@ -464,7 +464,7 @@ fn format_timestamp_non_legacy(dt: NaiveDateTime, precision: u32) -> String { /// Escape a path component following Java `PartitionPathUtils.escapePathName`. /// /// Characters that need escaping are encoded as `%XX` (uppercase hex). -fn escape_path_name(path: &str) -> String { +pub(crate) fn escape_path_name(path: &str) -> String { if !path.chars().any(needs_escaping) { return path.to_string(); } @@ -488,6 +488,34 @@ fn escape_path_name(path: &str) -> String { sb } +/// Unescape a path component following Java `PartitionPathUtils.unescapePathName`. +pub(crate) fn unescape_path_name(path: &str) -> Option { + let bytes = path.as_bytes(); + let mut result = Vec::with_capacity(bytes.len()); + let mut index = 0; + while index < bytes.len() { + if bytes[index] == b'%' { + let high = hex_value(*bytes.get(index + 1)?)?; + let low = hex_value(*bytes.get(index + 2)?)?; + result.push((high << 4) | low); + index += 3; + } else { + result.push(bytes[index]); + index += 1; + } + } + String::from_utf8(result).ok() +} + +fn hex_value(byte: u8) -> Option { + match byte { + b'0'..=b'9' => Some(byte - b'0'), + b'a'..=b'f' => Some(byte - b'a' + 10), + b'A'..=b'F' => Some(byte - b'A' + 10), + _ => None, + } +} + /// Check if a character needs escaping in partition path names. /// /// Matches Java `PartitionPathUtils.CHAR_TO_ESCAPE`: @@ -657,6 +685,14 @@ mod tests { assert_eq!(escape_path_name("a\x7Fb"), "a%7Fb"); } + #[test] + fn test_unescape_path_name() { + assert_eq!(unescape_path_name("a%2Fb"), Some("a/b".to_string())); + assert_eq!(unescape_path_name("%E4%B8%AD"), Some("中".to_string())); + assert_eq!(unescape_path_name("a%ZZb"), None); + assert_eq!(unescape_path_name("a%b"), None); + } + // ======================== PartitionComputer tests ======================== #[test] diff --git a/crates/paimon/src/table/format_partition.rs b/crates/paimon/src/table/format_partition.rs new file mode 100644 index 000000000..a70476c51 --- /dev/null +++ b/crates/paimon/src/table/format_partition.rs @@ -0,0 +1,441 @@ +// 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. + +//! Format Table partition path helpers shared by readers and SQL administration. + +use std::collections::HashMap; + +use chrono::NaiveDate; + +use crate::io::FileIO; +use crate::spec::{escape_path_name, unescape_path_name, DataType, Datum}; + +const UNIX_EPOCH_DAYS_FROM_CE: i32 = 719_163; + +/// Generates canonical names and physical paths for Format Table partitions. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FormatTablePartitionPaths { + partition_keys: Vec, + only_value_in_path: bool, +} + +impl FormatTablePartitionPaths { + /// Create a helper for the declared partition-key order and physical layout. + pub fn new(partition_keys: I, only_value_in_path: bool) -> Self + where + I: IntoIterator, + S: Into, + { + Self { + partition_keys: partition_keys.into_iter().map(Into::into).collect(), + only_value_in_path, + } + } + + /// Return the canonical logical partition name (`key=value/...`). + pub fn partition_name(&self, spec: &HashMap) -> crate::Result { + let values = self.ordered_values(spec)?; + Ok(self + .partition_keys + .iter() + .zip(values) + .map(|(key, value)| format!("{}={}", escape_path_name(key), escape_path_name(value))) + .collect::>() + .join("/")) + } + + /// Build the partition-name pattern that selects the partitions whose leading values + /// are `leading_values`, to push down to a partition-managing catalog. + /// + /// Pattern contract, shared by every engine talking to the catalog: partition names are + /// the escaped `key=value` form joined by `/`, `%` is the only wildcard and has no + /// escape sequence, and `_` stays literal. Values covering every partition key give the + /// one exact name; a shorter prefix is suffixed with `%`. + /// + /// `None` means pushdown has to be skipped and the caller must list every partition: + /// there are no leading values, one of them is blank, or escaping produced a literal + /// `%` that the contract cannot express. + /// + /// Mirrors Java `PartitionPathUtils.buildPartitionNamePrefixPattern`. + pub fn name_prefix_pattern(&self, leading_values: &[String]) -> Option { + if leading_values.is_empty() || leading_values.len() > self.partition_keys.len() { + return None; + } + let mut segments = Vec::with_capacity(leading_values.len()); + for (key, value) in self.partition_keys.iter().zip(leading_values) { + if value.trim().is_empty() { + return None; + } + segments.push(format!( + "{}={}", + escape_path_name(key), + escape_path_name(value) + )); + } + let pattern = segments.join("/"); + if pattern.contains('%') { + return None; + } + if leading_values.len() == self.partition_keys.len() { + Some(pattern) + } else { + Some(format!("{pattern}/%")) + } + } + + /// Return the physical partition path relative to the table location. + pub fn relative_path(&self, spec: &HashMap) -> crate::Result { + if !self.only_value_in_path { + return self.partition_name(spec); + } + Ok(self + .ordered_values(spec)? + .into_iter() + .map(escape_path_name) + .collect::>() + .join("/")) + } + + /// Discover complete raw partition specs from the table directory. + /// + /// Hidden directories whose names begin with `.` or `_`, segments that do + /// not match the configured layout, and paths shallower than the declared + /// partition depth are ignored. In value-only layouts, the configured + /// default-partition directory is the only hidden-name exception. + /// A matching segment that is malformed or not canonically escaped returns + /// an error rather than being skipped, because its catalog spec would not + /// resolve back to the same physical path. + /// Results are sorted and deduplicated by canonical `key=value/...` name. + pub async fn discover( + &self, + file_io: &FileIO, + table_path: &str, + default_partition_name: &str, + ) -> crate::Result>> { + let default_partition_path_name = self + .only_value_in_path + .then(|| escape_path_name(default_partition_name)); + let mut frontier = vec![(table_path.trim_end_matches('/').to_string(), HashMap::new())]; + for key in &self.partition_keys { + let mut next = Vec::new(); + for (path, spec) in frontier { + let statuses = match file_io.list_status(&path).await { + Ok(statuses) => statuses, + Err(error) if is_storage_not_found(&error) => continue, + Err(error) => return Err(error), + }; + for status in statuses { + if !status.is_dir { + continue; + } + let Some(segment) = last_path_segment(&status.path) else { + continue; + }; + let is_value_only_default = + default_partition_path_name.as_deref() == Some(segment); + if (segment.starts_with('.') || segment.starts_with('_')) + && !is_value_only_default + { + continue; + } + let Some(value) = self.partition_value_from_segment(key, segment)? else { + continue; + }; + let mut child_spec = spec.clone(); + child_spec.insert(key.clone(), value); + next.push((status.path.trim_end_matches('/').to_string(), child_spec)); + } + } + frontier = next; + } + + let mut partitions = frontier + .into_iter() + .map(|(_, spec)| Ok((self.partition_name(&spec)?, spec))) + .collect::>>()?; + partitions.sort_by(|left, right| left.0.cmp(&right.0)); + partitions.dedup_by(|left, right| left.0 == right.0); + Ok(partitions.into_iter().map(|(_, spec)| spec).collect()) + } + + fn ordered_values<'a>(&self, spec: &'a HashMap) -> crate::Result> { + if spec.len() != self.partition_keys.len() { + return Err(crate::Error::DataInvalid { + message: self.invalid_partition_keys_message(spec), + source: None, + }); + } + + let mut values = Vec::with_capacity(self.partition_keys.len()); + for key in &self.partition_keys { + let Some(value) = spec.get(key).map(String::as_str) else { + return Err(crate::Error::DataInvalid { + message: self.invalid_partition_keys_message(spec), + source: None, + }); + }; + if value.is_empty() || (self.only_value_in_path && matches!(value, "." | "..")) { + return Err(crate::Error::DataInvalid { + message: format!( + "Partition value {value:?} cannot be used as a partition path component" + ), + source: None, + }); + } + values.push(value); + } + Ok(values) + } + + fn invalid_partition_keys_message(&self, spec: &HashMap) -> String { + let mut actual_keys = spec.keys().collect::>(); + actual_keys.sort(); + format!( + "Partition spec must contain exactly keys {:?}, but contains {actual_keys:?}", + self.partition_keys + ) + } + + fn partition_value_from_segment( + &self, + key: &str, + segment: &str, + ) -> crate::Result> { + let value = if self.only_value_in_path { + decode_canonical_path_name(segment)? + } else { + let Some((segment_key, value)) = segment.split_once('=') else { + return Ok(None); + }; + let decoded_key = decode_canonical_path_name(segment_key)?; + if decoded_key != key { + return Ok(None); + } + decode_canonical_path_name(value)? + }; + Ok(Some(value)) + } +} + +fn decode_canonical_path_name(value: &str) -> crate::Result { + let decoded = unescape_path_name(value).ok_or_else(|| crate::Error::DataInvalid { + message: format!("Invalid escaped partition path segment {value:?}"), + source: None, + })?; + let canonical = escape_path_name(&decoded); + if canonical != value { + return Err(crate::Error::DataInvalid { + message: format!( + "Partition path segment {value:?} cannot round-trip through catalog metadata; \ + its canonical escaped form is {canonical:?}" + ), + source: None, + }); + } + Ok(decoded) +} + +/// Parse a raw Format Table partition value from a path or catalog registration. +pub fn parse_format_partition_value(value: &str, data_type: &DataType) -> Option { + match data_type { + DataType::Boolean(_) => parse_partition_bool(value).map(Datum::Bool), + DataType::TinyInt(_) => value.parse::().ok().map(Datum::TinyInt), + DataType::SmallInt(_) => value.parse::().ok().map(Datum::SmallInt), + DataType::Int(_) => value.parse::().ok().map(Datum::Int), + DataType::BigInt(_) => value.parse::().ok().map(Datum::Long), + DataType::Char(_) | DataType::VarChar(_) => Some(Datum::String(value.to_string())), + DataType::Date(_) => parse_partition_date(value).map(Datum::Date), + DataType::Time(_) => value.parse::().ok().map(Datum::Time), + _ => None, + } +} + +/// Format a typed value for Format Table partition metadata and paths. +pub fn format_partition_value( + datum: &Datum, + data_type: &DataType, + default_partition_name: &str, + legacy_partition_name: bool, +) -> Option { + match (datum, data_type) { + (Datum::Bool(value), DataType::Boolean(_)) => Some(value.to_string()), + (Datum::TinyInt(value), DataType::TinyInt(_)) => Some(value.to_string()), + (Datum::SmallInt(value), DataType::SmallInt(_)) => Some(value.to_string()), + (Datum::Int(value), DataType::Int(_)) => Some(value.to_string()), + (Datum::Long(value), DataType::BigInt(_)) => Some(value.to_string()), + (Datum::String(value), DataType::Char(_) | DataType::VarChar(_)) => { + if value.trim().is_empty() { + Some(default_partition_name.to_string()) + } else { + Some(value.clone()) + } + } + (Datum::Date(value), DataType::Date(_)) => { + if legacy_partition_name { + Some(value.to_string()) + } else { + format_partition_date(*value) + } + } + (Datum::Time(value), DataType::Time(_)) => Some(value.to_string()), + _ => None, + } +} + +/// Accept the boolean spellings Java accepts, so a partition another engine registered as +/// `TRUE`, `t`, `yes` or `1` reads back here instead of failing as invalid metadata. +/// +/// Mirrors Java `BinaryStringUtils.toBoolean`. +fn parse_partition_bool(value: &str) -> Option { + const TRUE_VALUES: [&str; 5] = ["t", "true", "y", "yes", "1"]; + const FALSE_VALUES: [&str; 5] = ["f", "false", "n", "no", "0"]; + if TRUE_VALUES + .iter() + .any(|candidate| value.eq_ignore_ascii_case(candidate)) + { + return Some(true); + } + if FALSE_VALUES + .iter() + .any(|candidate| value.eq_ignore_ascii_case(candidate)) + { + return Some(false); + } + None +} + +fn parse_partition_date(value: &str) -> Option { + if let Ok(epoch_days) = value.parse::() { + return Some(epoch_days); + } + let date = NaiveDate::parse_from_str(value, "%Y-%m-%d").ok()?; + let epoch = NaiveDate::from_num_days_from_ce_opt(UNIX_EPOCH_DAYS_FROM_CE)?; + date.signed_duration_since(epoch).num_days().try_into().ok() +} + +pub(crate) fn format_partition_date(epoch_days: i32) -> Option { + NaiveDate::from_num_days_from_ce_opt(epoch_days.checked_add(UNIX_EPOCH_DAYS_FROM_CE)?) + .map(|date| date.format("%Y-%m-%d").to_string()) +} + +pub(crate) fn is_storage_not_found(error: &crate::Error) -> bool { + matches!( + error, + crate::Error::IoUnexpected { source, .. } + if source.kind() == opendal::ErrorKind::NotFound + ) +} + +fn last_path_segment(path: &str) -> Option<&str> { + path.trim_end_matches('/').rsplit('/').next() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::{BooleanType, DateType}; + + #[test] + fn test_parse_format_partition_value() { + assert_eq!( + parse_format_partition_value("true", &DataType::Boolean(BooleanType::new())), + Some(Datum::Bool(true)) + ); + assert_eq!( + parse_format_partition_value("2026-07-22", &DataType::Date(DateType::new())), + Some(Datum::Date(20_656)) + ); + assert_eq!( + parse_format_partition_value("20656", &DataType::Date(DateType::new())), + Some(Datum::Date(20_656)) + ); + } + + #[test] + fn test_parse_format_partition_bool_accepts_every_java_spelling() { + let boolean = DataType::Boolean(BooleanType::new()); + for value in ["t", "T", "true", "TRUE", "True", "y", "YES", "1"] { + assert_eq!( + parse_format_partition_value(value, &boolean), + Some(Datum::Bool(true)), + "{value} should read as true" + ); + } + for value in ["f", "F", "false", "FALSE", "False", "n", "NO", "0"] { + assert_eq!( + parse_format_partition_value(value, &boolean), + Some(Datum::Bool(false)), + "{value} should read as false" + ); + } + for value in ["", "2", "tru", "yes please", "null"] { + assert_eq!( + parse_format_partition_value(value, &boolean), + None, + "{value} is not a boolean" + ); + } + } + + #[test] + fn test_name_prefix_pattern() { + let paths = FormatTablePartitionPaths::new(["dt".to_string(), "hh".to_string()], false); + + // A complete prefix names one partition, a shorter one takes the wildcard. + assert_eq!( + paths.name_prefix_pattern(&["20260722".to_string(), "10".to_string()]), + Some("dt=20260722/hh=10".to_string()) + ); + assert_eq!( + paths.name_prefix_pattern(&["20260722".to_string()]), + Some("dt=20260722/%".to_string()) + ); + + // The physical layout never changes the pattern: catalog names are always key=value. + let value_only = FormatTablePartitionPaths::new(["dt".to_string(), "hh".to_string()], true); + assert_eq!( + value_only.name_prefix_pattern(&["20260722".to_string()]), + Some("dt=20260722/%".to_string()) + ); + + // Nothing to push down. + assert_eq!(paths.name_prefix_pattern(&[]), None); + assert_eq!(paths.name_prefix_pattern(&[" ".to_string()]), None); + assert_eq!( + paths.name_prefix_pattern(&["a".to_string(), "b".to_string(), "c".to_string()]), + None + ); + + // Escaping a value would inject the pattern's only wildcard, so pushdown is skipped + // rather than silently widened. + assert_eq!(paths.name_prefix_pattern(&["2026/07".to_string()]), None); + } + + #[test] + fn test_storage_not_found_matches_only_not_found() { + for (kind, expected) in [ + (opendal::ErrorKind::NotFound, true), + (opendal::ErrorKind::PermissionDenied, false), + ] { + let error = crate::Error::IoUnexpected { + message: "list partition directory".to_string(), + source: Box::new(opendal::Error::new(kind, "test")), + }; + assert_eq!(is_storage_not_found(&error), expected); + } + } +} diff --git a/crates/paimon/src/table/format_table_scan.rs b/crates/paimon/src/table/format_table_scan.rs index df2a48e02..1bfb6ee88 100644 --- a/crates/paimon/src/table/format_table_scan.rs +++ b/crates/paimon/src/table/format_table_scan.rs @@ -17,15 +17,21 @@ //! Scan implementation for Java-compatible `type=format-table` metadata. -use super::{Plan, ScanTrace, Table}; +use std::collections::{HashMap, HashSet}; + +use super::format_partition::{ + format_partition_value, is_storage_not_found, parse_format_partition_value, + FormatTablePartitionPaths, +}; +use super::rest_env::LoadedFormatTablePartitionOptions; +use super::{Plan, RESTEnv, ScanTrace, Table}; use crate::spec::stats::BinaryTableStats; use crate::spec::{ - extract_datum, BinaryRow, BinaryRowBuilder, CoreOptions, DataField, DataFileMeta, DataType, - Datum, PartitionComputer, Predicate, PredicateOperator, + escape_path_name, extract_datum, unescape_path_name, BinaryRow, BinaryRowBuilder, CoreOptions, + DataField, DataFileMeta, Datum, PartitionComputer, Predicate, PredicateOperator, }; use crate::table::partition_filter::PartitionFilter; use crate::table::source::DataSplitBuilder; -use chrono::NaiveDate; #[derive(Debug, Clone)] pub(crate) struct FormatTableScan<'a> { @@ -65,17 +71,25 @@ impl<'a> FormatTableScan<'a> { async fn plan_inner(&self, trace: Option<&mut ScanTrace>) -> crate::Result { let core_options = CoreOptions::new(self.table.schema().options()); - let format_extension = supported_format_table_extension(core_options.file_format())?; + let managed_options = self + .table + .rest_env() + .and_then(RESTEnv::catalog_managed_partition_options); + let file_format = managed_options + .map(|options| options.file_format.as_str()) + .unwrap_or_else(|| core_options.file_format()); + let format_extension = supported_format_table_extension(file_format)?; let schema_id = self.table.schema().id(); - let table_path = core_options - .path() + let table_path = managed_options + .map(|options| options.table_path.as_str()) + .or_else(|| core_options.path()) .unwrap_or_else(|| self.table.location()) .trim_end_matches('/') .to_string(); let partition_fields = self.table.schema().partition_fields(); let mut splits = Vec::new(); - for scan_root in self.scan_roots(&core_options, &table_path)? { + for scan_root in self.scan_roots(&core_options, &table_path).await? { let statuses = self .list_status_recursive_if_exists(&scan_root.path) .await?; @@ -111,7 +125,7 @@ impl<'a> FormatTableScan<'a> { Ok(Plan::new(splits)) } - fn scan_roots( + async fn scan_roots( &self, core_options: &CoreOptions<'_>, table_path: &str, @@ -124,6 +138,19 @@ impl<'a> FormatTableScan<'a> { partition: BinaryRow::new(0), }]); } + if let Some(rest_env) = self.table.rest_env() { + if let Some(managed_options) = rest_env.catalog_managed_partition_options() { + return self + .catalog_managed_scan_roots( + rest_env, + table_path, + partition_keys, + &partition_fields, + managed_options, + ) + .await; + } + } let Some(PartitionFilter::PartitionSet { partitions, .. }) = &self.partition_filter else { if let Some(PartitionFilter::Predicate(predicate)) = &self.partition_filter { @@ -178,19 +205,81 @@ impl<'a> FormatTableScan<'a> { Ok(roots) } + async fn catalog_managed_scan_roots( + &self, + rest_env: &RESTEnv, + table_path: &str, + partition_keys: &[String], + partition_fields: &[DataField], + managed_options: &LoadedFormatTablePartitionOptions, + ) -> crate::Result> { + let only_value_in_path = managed_options.only_value_in_path; + let partition_paths = + FormatTablePartitionPaths::new(partition_keys.iter().cloned(), only_value_in_path); + let core_options = CoreOptions::new(self.table.schema().options()); + let default_partition_name = core_options.partition_default_name(); + // Ask the catalog only for the partitions the filter can reach. Downloading every + // registration of a table with many partitions is what dominates planning time, + // and the local match below still decides what is actually scanned. + let pattern = match &self.partition_filter { + Some(filter) => { + let leading_values = leading_equality_partition_values( + filter, + partition_fields, + default_partition_name, + core_options.legacy_partition_name(), + )?; + partition_paths.name_prefix_pattern(&leading_values) + } + None => None, + }; + let partitions = rest_env + .api() + .list_partitions_by_name_pattern(rest_env.identifier(), pattern.as_deref()) + .await?; + let mut seen_paths = HashSet::with_capacity(partitions.len()); + let mut roots = Vec::with_capacity(partitions.len()); + for partition in partitions { + let partition_path = partition_paths + .relative_path(&partition.spec) + .map_err(|error| self.invalid_catalog_partition_metadata(error))?; + if !seen_paths.insert(partition_path.clone()) { + continue; + } + let path = join_path(table_path, &partition_path); + let partition = partition_row_from_catalog_spec( + &partition.spec, + partition_fields, + partition_keys, + default_partition_name, + ) + .map_err(|error| self.invalid_catalog_partition_metadata(error))?; + if self.partition_matches(&partition)? { + roots.push(ScanRoot { path, partition }); + } + } + roots.sort_by(|left, right| left.path.cmp(&right.path)); + Ok(roots) + } + + fn invalid_catalog_partition_metadata(&self, source: crate::Error) -> crate::Error { + crate::Error::DataInvalid { + message: format!( + "Catalog returned invalid partition metadata for Format Table {}", + self.table.identifier().full_name() + ), + source: Some(Box::new(source)), + } + } + async fn list_status_recursive_if_exists( &self, path: &str, ) -> crate::Result> { match self.table.file_io().list_status_recursive(path).await { Ok(statuses) => Ok(statuses), - Err(err) => { - if !self.table.file_io().exists(path).await.unwrap_or(true) { - Ok(Vec::new()) - } else { - Err(err) - } - } + Err(err) if is_storage_not_found(&err) => Ok(Vec::new()), + Err(err) => Err(err), } } @@ -319,8 +408,96 @@ fn leading_equality_partition_path( legacy_partition_name: bool, only_value_in_path: bool, ) -> Option { + let values = leading_equality_values_from_predicate( + predicate, + partition_fields, + default_partition_name, + legacy_partition_name, + ); + if values.is_empty() { + return None; + } + let segments = partition_keys + .iter() + .zip(&values) + .map(|(key, value)| { + if only_value_in_path { + escape_path_name(value) + } else { + format!("{}={}", escape_path_name(key), escape_path_name(value)) + } + }) + .collect::>(); + Some(join_path(table_path, &segments.join("/"))) +} + +/// The leading run of partition values this filter pins to a single value, in +/// partition-key order, formatted the way the catalog and the partition path spell them. +/// +/// Only a leading run is useful: a partition path prefix and the partition-name pattern a +/// catalog prunes on can express nothing else. A null value has no such spelling, so the +/// run stops there. +fn leading_equality_partition_values( + filter: &PartitionFilter, + partition_fields: &[DataField], + default_partition_name: &str, + legacy_partition_name: bool, +) -> crate::Result> { + match filter { + PartitionFilter::Predicate(predicate) => Ok(leading_equality_values_from_predicate( + predicate, + partition_fields, + default_partition_name, + legacy_partition_name, + )), + // An enumerated partition set still pins a prefix whenever its rows agree on one, + // which is what `dt = 'a' AND hh IN ('10', '11')` collapses to. + PartitionFilter::PartitionSet { partitions, .. } => { + let mut common: Option> = None; + for serialized in partitions { + let row = BinaryRow::from_serialized_bytes(serialized)?; + let mut values = Vec::with_capacity(partition_fields.len()); + for (index, field) in partition_fields.iter().enumerate() { + let Some(datum) = extract_datum(&row, index, field.data_type())? else { + break; + }; + let Some(value) = format_partition_value( + &datum, + field.data_type(), + default_partition_name, + legacy_partition_name, + ) else { + break; + }; + values.push(value); + } + common = Some(match common { + None => values, + Some(common) => common + .into_iter() + .zip(values) + .take_while(|(left, right)| left == right) + .map(|(left, _)| left) + .collect(), + }); + if common.as_ref().is_some_and(Vec::is_empty) { + break; + } + } + Ok(common.unwrap_or_default()) + } + } +} + +/// Mirrors Java `FormatTableScan.extractLeadingEqualityPartitionSpecWhenOnlyAnd`. +fn leading_equality_values_from_predicate( + predicate: &Predicate, + partition_fields: &[DataField], + default_partition_name: &str, + legacy_partition_name: bool, +) -> Vec { + let mut pinned: Vec> = vec![None; partition_fields.len()]; let predicates = predicate.clone().split_and(); - let mut values: Vec> = vec![None; partition_keys.len()]; for predicate in &predicates { let Predicate::Leaf { index, @@ -331,38 +508,27 @@ fn leading_equality_partition_path( else { continue; }; - if *index < values.len() { - values[*index] = literals.first(); + if *index < pinned.len() { + pinned[*index] = literals.first(); } } - let mut segments = Vec::new(); - for (idx, key) in partition_keys.iter().enumerate() { - let Some(datum) = values[idx] else { + let mut values = Vec::new(); + for (field, datum) in partition_fields.iter().zip(pinned) { + let Some(datum) = datum else { break; }; - let value = partition_value_from_datum( + let Some(value) = format_partition_value( datum, - partition_fields[idx].data_type(), + field.data_type(), default_partition_name, legacy_partition_name, - )?; - if only_value_in_path { - segments.push(escape_path_name(&value)); - } else { - segments.push(format!( - "{}={}", - escape_path_name(key), - escape_path_name(&value) - )); - } - } - - if segments.is_empty() { - None - } else { - Some(join_path(table_path, &segments.join("/"))) + ) else { + break; + }; + values.push(value); } + values } fn partition_path_from_row( @@ -376,7 +542,7 @@ fn partition_path_from_row( for (idx, field) in partition_fields.iter().enumerate() { let value = match extract_datum(row, idx, field.data_type())? { None => default_partition_name.to_string(), - Some(datum) => partition_value_from_datum( + Some(datum) => format_partition_value( &datum, field.data_type(), default_partition_name, @@ -402,78 +568,6 @@ fn partition_path_from_row( Ok(segments.join("/")) } -fn partition_value_from_datum( - datum: &Datum, - data_type: &DataType, - default_partition_name: &str, - legacy_partition_name: bool, -) -> Option { - match (datum, data_type) { - (Datum::Bool(value), DataType::Boolean(_)) => Some(value.to_string()), - (Datum::TinyInt(value), DataType::TinyInt(_)) => Some(value.to_string()), - (Datum::SmallInt(value), DataType::SmallInt(_)) => Some(value.to_string()), - (Datum::Int(value), DataType::Int(_)) => Some(value.to_string()), - (Datum::Long(value), DataType::BigInt(_)) => Some(value.to_string()), - (Datum::String(value), DataType::Char(_) | DataType::VarChar(_)) => { - if value.trim().is_empty() { - Some(default_partition_name.to_string()) - } else { - Some(value.clone()) - } - } - (Datum::Date(value), DataType::Date(_)) => { - if legacy_partition_name { - Some(value.to_string()) - } else { - Some(format_partition_date(*value)) - } - } - (Datum::Time(value), DataType::Time(_)) => Some(value.to_string()), - _ => None, - } -} - -fn format_partition_date(epoch_days: i32) -> String { - let date = NaiveDate::from_num_days_from_ce_opt(epoch_days + 719_163) - .unwrap_or(NaiveDate::from_ymd_opt(1970, 1, 1).unwrap()); - date.format("%Y-%m-%d").to_string() -} - -fn escape_path_name(path: &str) -> String { - let mut result = String::with_capacity(path.len()); - for byte in path.bytes() { - if should_escape(byte) { - result.push('%'); - result.push_str(&format!("{byte:02X}")); - } else { - result.push(byte as char); - } - } - result -} - -fn should_escape(byte: u8) -> bool { - byte <= 0x1F - || byte >= 0x7F - || matches!( - byte, - b'"' | b'#' - | b'%' - | b'\'' - | b'*' - | b'/' - | b':' - | b'=' - | b'?' - | b'\\' - | b'\x7F' - | b'{' - | b'[' - | b']' - | b'^' - ) -} - fn partition_row_from_path( table_path: &str, file_parent: &str, @@ -526,7 +620,8 @@ fn partition_row_from_path( builder.set_null_at(idx); continue; } - let Some(datum) = parse_partition_datum(value, partition_fields[idx].data_type()) else { + let Some(datum) = parse_format_partition_value(value, partition_fields[idx].data_type()) + else { return Ok(None); }; builder.write_datum(idx, &datum, partition_fields[idx].data_type()); @@ -534,6 +629,36 @@ fn partition_row_from_path( Ok(Some(builder.build())) } +fn partition_row_from_catalog_spec( + spec: &HashMap, + partition_fields: &[DataField], + partition_keys: &[String], + default_partition_name: &str, +) -> crate::Result { + let mut builder = BinaryRowBuilder::new(partition_fields.len() as i32); + for (index, (key, field)) in partition_keys.iter().zip(partition_fields).enumerate() { + let value = spec.get(key).ok_or_else(|| crate::Error::DataInvalid { + message: format!("Catalog partition is missing column '{key}'"), + source: None, + })?; + if value == default_partition_name { + builder.set_null_at(index); + continue; + } + let datum = parse_format_partition_value(value, field.data_type()).ok_or_else(|| { + crate::Error::DataInvalid { + message: format!( + "Invalid catalog partition value {value:?} for column '{key}' with type {:?}", + field.data_type() + ), + source: None, + } + })?; + builder.write_datum(index, &datum, field.data_type()); + } + Ok(builder.build()) +} + fn partition_segment_value(segment: &str, key: &str) -> Option { let (segment_key, segment_value) = segment.split_once('=')?; if unescape_path_name(segment_key)? == key { @@ -543,61 +668,6 @@ fn partition_segment_value(segment: &str, key: &str) -> Option { } } -fn unescape_path_name(value: &str) -> Option { - let bytes = value.as_bytes(); - let mut out = Vec::with_capacity(bytes.len()); - let mut i = 0; - while i < bytes.len() { - if bytes[i] == b'%' { - if i + 2 >= bytes.len() { - return None; - } - let hi = hex_value(bytes[i + 1])?; - let lo = hex_value(bytes[i + 2])?; - out.push((hi << 4) | lo); - i += 3; - } else { - out.push(bytes[i]); - i += 1; - } - } - String::from_utf8(out).ok() -} - -fn hex_value(byte: u8) -> Option { - match byte { - b'0'..=b'9' => Some(byte - b'0'), - b'a'..=b'f' => Some(byte - b'a' + 10), - b'A'..=b'F' => Some(byte - b'A' + 10), - _ => None, - } -} - -fn parse_partition_datum(value: &str, data_type: &DataType) -> Option { - match data_type { - DataType::Boolean(_) => value.parse::().ok().map(Datum::Bool), - DataType::TinyInt(_) => value.parse::().ok().map(Datum::TinyInt), - DataType::SmallInt(_) => value.parse::().ok().map(Datum::SmallInt), - DataType::Int(_) => value.parse::().ok().map(Datum::Int), - DataType::BigInt(_) => value.parse::().ok().map(Datum::Long), - DataType::Char(_) | DataType::VarChar(_) => Some(Datum::String(value.to_string())), - DataType::Date(_) => parse_partition_date(value).map(Datum::Date), - DataType::Time(_) => value.parse::().ok().map(Datum::Time), - _ => None, - } -} - -fn parse_partition_date(value: &str) -> Option { - if let Ok(epoch_days) = value.parse::() { - return Some(epoch_days); - } - let date = NaiveDate::parse_from_str(value, "%Y-%m-%d").ok()?; - date.signed_duration_since(NaiveDate::from_ymd_opt(1970, 1, 1).unwrap()) - .num_days() - .try_into() - .ok() -} - fn supported_format_table_extension(format: &str) -> crate::Result<&'static str> { match format.to_ascii_lowercase().as_str() { "parquet" => Ok(".parquet"), diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 4078e3a23..dbfef0fc7 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -38,6 +38,7 @@ pub mod data_evolution_writer; mod data_file_reader; mod data_file_writer; mod dedicated_format_file_writer; +mod format_partition; mod format_read_builder; mod format_table_read; mod format_table_scan; @@ -105,6 +106,9 @@ pub use btree_global_index_build_builder::BTreeGlobalIndexBuildBuilder; pub use commit_message::CommitMessage; pub use cow_writer::{CopyOnWriteMergeWriter, FileInfo}; pub use data_evolution_writer::{DataEvolutionDeleteWriter, DataEvolutionWriter}; +pub use format_partition::{ + format_partition_value, parse_format_partition_value, FormatTablePartitionPaths, +}; #[cfg(feature = "fulltext")] pub use full_text_search_builder::FullTextSearchBuilder; use futures::stream::BoxStream; @@ -303,7 +307,16 @@ impl Table { } pub(crate) fn is_format_table(&self) -> bool { - CoreOptions::new(self.schema.options()).is_format_table() + self.has_catalog_managed_partitions() + || CoreOptions::new(self.schema.options()).is_format_table() + } + + /// Whether this table uses catalog-managed Format Table partitions. + pub fn has_catalog_managed_partitions(&self) -> bool { + self.rest_env + .as_ref() + .and_then(RESTEnv::catalog_managed_partition_options) + .is_some() } /// Create a read builder for scan/read. @@ -365,6 +378,9 @@ impl Table { /// `FileStoreTable.copyWithoutTimeTravel`. Use /// [`Table::copy_with_time_travel`] when the options may select a /// historical snapshot whose schema should be used for reading. + /// + /// Catalog-managed Format Table scans keep the partition source, table + /// path, file format, and path layout loaded from REST metadata. pub fn copy_with_options(&self, extra: HashMap) -> Self { // Changing the time-travel selector invalidates the resolved snapshot // (a time-travelled schema then has no matching snapshot anymore, and @@ -408,6 +424,9 @@ impl Table { /// schema (the `if let Ok` below swallows them); an invalid selector /// still fails later at scan planning. pub async fn copy_with_time_travel(&self, extra: HashMap) -> Result { + if let Some(rest_env) = &self.rest_env { + rest_env.validate_dynamic_format_table_partition_options(&extra)?; + } let mut table = self.copy_with_options(extra); // Reject unimplemented scan options on the merged view before any IO, so // both table-level and per-read options are covered. diff --git a/crates/paimon/src/table/rest_env.rs b/crates/paimon/src/table/rest_env.rs index 4f2a9df76..df6712e4a 100644 --- a/crates/paimon/src/table/rest_env.rs +++ b/crates/paimon/src/table/rest_env.rs @@ -17,6 +17,9 @@ //! REST environment for REST-backed table operations. +use std::collections::HashMap; +use std::sync::Arc; + use crate::api::rest_api::RESTApi; use crate::api::rest_error::RestError; use crate::catalog::{Identifier, RESTTokenFileIO}; @@ -24,14 +27,23 @@ use crate::common::Options; use crate::error::Error; use crate::io::cache::LocalCache; use crate::io::FileIO; -use crate::spec::{CoreOptions, TableSchema, PATH_OPTION}; +use crate::spec::{ + CoreOptions, FormatTableImplementation, TableSchema, FORMAT_TABLE_IMPLEMENTATION_OPTION, + FORMAT_TABLE_PARTITION_PATH_ONLY_VALUE_OPTION, METASTORE_PARTITIONED_TABLE_OPTION, PATH_OPTION, +}; use crate::table::snapshot_commit::{RESTSnapshotCommit, SnapshotCommit}; use crate::table::Table; use crate::Result; -use std::sync::Arc; -/// REST environment that holds the REST API client, identifier, and uuid -/// needed to create a `RESTSnapshotCommit`. +#[derive(Clone, Debug)] +pub(crate) struct LoadedFormatTablePartitionOptions { + pub(crate) table_path: String, + pub(crate) file_format: String, + pub(crate) only_value_in_path: bool, + partitioned_table_in_metastore: bool, +} + +/// REST-backed table context used by snapshot commits and managed Format Table scans. #[derive(Clone)] pub struct RESTEnv { identifier: Identifier, @@ -40,6 +52,7 @@ pub struct RESTEnv { options: Options, data_token_enabled: bool, local_cache: Option>, + loaded_format_table_partition_options: Option, } impl std::fmt::Debug for RESTEnv { @@ -68,6 +81,7 @@ impl RESTEnv { options, data_token_enabled, local_cache, + loaded_format_table_partition_options: None, } } @@ -76,6 +90,53 @@ impl RESTEnv { self.local_cache.is_some() } + pub(crate) fn catalog_managed_partition_options( + &self, + ) -> Option<&LoadedFormatTablePartitionOptions> { + self.loaded_format_table_partition_options + .as_ref() + .filter(|options| options.partitioned_table_in_metastore) + } + + pub(crate) fn validate_dynamic_format_table_partition_options( + &self, + extra: &HashMap, + ) -> Result<()> { + let Some(loaded_options) = &self.loaded_format_table_partition_options else { + return Ok(()); + }; + let options = CoreOptions::new(extra); + if extra.contains_key(METASTORE_PARTITIONED_TABLE_OPTION) { + ensure_partition_option_unchanged( + &self.identifier, + METASTORE_PARTITIONED_TABLE_OPTION, + loaded_options.partitioned_table_in_metastore, + options.partitioned_table_in_metastore()?, + )?; + } + if loaded_options.partitioned_table_in_metastore + && extra.contains_key(FORMAT_TABLE_PARTITION_PATH_ONLY_VALUE_OPTION) + { + ensure_partition_option_unchanged( + &self.identifier, + FORMAT_TABLE_PARTITION_PATH_ONLY_VALUE_OPTION, + loaded_options.only_value_in_path, + options.try_format_table_partition_only_value_in_path()?, + )?; + } + if loaded_options.partitioned_table_in_metastore + && extra.contains_key(FORMAT_TABLE_IMPLEMENTATION_OPTION) + { + ensure_partition_option_unchanged( + &self.identifier, + FORMAT_TABLE_IMPLEMENTATION_OPTION, + FormatTableImplementation::Paimon, + options.format_table_implementation()?, + )?; + } + Ok(()) + } + /// Get the REST API client. pub fn api(&self) -> &Arc { &self.api @@ -143,7 +204,6 @@ impl RESTEnv { ), source: None, })?; - let uuid = response.id.ok_or_else(|| Error::DataInvalid { message: format!( "Table {} response missing id (uuid)", @@ -170,7 +230,7 @@ impl RESTEnv { builder.build()? }; - let rest_env = RESTEnv::new( + let mut rest_env = RESTEnv::new( identifier.clone(), uuid, api, @@ -178,6 +238,8 @@ impl RESTEnv { data_token_enabled, local_cache, ); + rest_env.loaded_format_table_partition_options = + load_format_table_partition_options(identifier, &table_schema, is_external)?; Ok(Table::new( file_io, @@ -198,6 +260,80 @@ impl RESTEnv { } } +fn load_format_table_partition_options( + identifier: &Identifier, + table_schema: &TableSchema, + is_external: bool, +) -> Result> { + let core_options = CoreOptions::new(table_schema.options()); + if !core_options.is_format_table() { + return Ok(None); + } + + let partitioned_table_in_metastore = core_options.partitioned_table_in_metastore()?; + let only_value_in_path = if partitioned_table_in_metastore { + if core_options.format_table_implementation()? == FormatTableImplementation::Engine { + return Err(Error::DataInvalid { + message: format!( + "Format Table {} cannot set metastore.partitioned-table=true when \ + format-table.implementation=engine", + identifier.full_name() + ), + source: None, + }); + } + if is_external { + return Err(Error::DataInvalid { + message: format!( + "Catalog-managed partitions require an internal Format Table, but {} is external", + identifier.full_name() + ), + source: None, + }); + } + core_options.try_format_table_partition_only_value_in_path()? + } else { + core_options.format_table_partition_only_value_in_path() + }; + + Ok(Some(LoadedFormatTablePartitionOptions { + table_path: core_options + .path() + .ok_or_else(|| Error::DataInvalid { + message: format!( + "REST Format Table {} is missing option '{PATH_OPTION}'", + identifier.full_name() + ), + source: None, + })? + .to_string(), + file_format: core_options.file_format().to_string(), + partitioned_table_in_metastore, + only_value_in_path, + })) +} + +fn ensure_partition_option_unchanged( + identifier: &Identifier, + key: &str, + loaded_value: T, + dynamic_value: T, +) -> Result<()> +where + T: PartialEq, +{ + if loaded_value == dynamic_value { + return Ok(()); + } + Err(Error::DataInvalid { + message: format!( + "Dynamic option '{key}' must match the value returned by the REST catalog for Format Table {}", + identifier.full_name() + ), + source: None, + }) +} + fn map_rest_error_for_table(err: Error, identifier: &Identifier) -> Error { match err { Error::RestApi { diff --git a/crates/paimon/tests/format_partition_test.rs b/crates/paimon/tests/format_partition_test.rs new file mode 100644 index 000000000..a0da749cd --- /dev/null +++ b/crates/paimon/tests/format_partition_test.rs @@ -0,0 +1,160 @@ +// 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. + +use std::collections::HashMap; +#[cfg(not(windows))] +use std::path::Path; + +use paimon::io::FileIO; +use paimon::table::FormatTablePartitionPaths; + +#[cfg(not(windows))] +async fn discover_partitions( + root: &Path, + partition_keys: &[&str], + only_value_in_path: bool, + default_partition_name: &str, +) -> paimon::Result>> { + let table_path = format!("file://{}", root.display()); + let file_io = FileIO::from_path(&table_path)?.build()?; + FormatTablePartitionPaths::new(partition_keys.iter().copied(), only_value_in_path) + .discover(&file_io, &table_path, default_partition_name) + .await +} + +#[test] +fn format_partition_paths_escape_values_and_honor_layout() { + let spec = HashMap::from([ + ("dt".to_string(), "2026/07=22".to_string()), + ("hour".to_string(), "10".to_string()), + ]); + + let keyed = FormatTablePartitionPaths::new(["dt", "hour"], false); + assert_eq!( + keyed.partition_name(&spec).unwrap(), + "dt=2026%2F07%3D22/hour=10" + ); + assert_eq!( + keyed.relative_path(&spec).unwrap(), + "dt=2026%2F07%3D22/hour=10" + ); + + let value_only = FormatTablePartitionPaths::new(["dt", "hour"], true); + assert_eq!( + value_only.partition_name(&spec).unwrap(), + "dt=2026%2F07%3D22/hour=10" + ); + assert_eq!( + value_only.relative_path(&spec).unwrap(), + "2026%2F07%3D22/10" + ); +} + +#[cfg(not(windows))] +#[tokio::test] +async fn discover_format_partitions_returns_complete_sorted_specs() { + let tmp = tempfile::tempdir().unwrap(); + for relative in [ + "dt=2026-07-22/hour=10", + "dt=2026-07-21/hour=09", + "dt=2026-07-20", + ".staging/dt=2026-07-19/hour=08", + ] { + std::fs::create_dir_all(tmp.path().join(relative)).unwrap(); + } + + let partitions = + discover_partitions(tmp.path(), &["dt", "hour"], false, "__DEFAULT_PARTITION__") + .await + .unwrap(); + + assert_eq!( + partitions, + vec![ + HashMap::from([ + ("dt".to_string(), "2026-07-21".to_string()), + ("hour".to_string(), "09".to_string()), + ]), + HashMap::from([ + ("dt".to_string(), "2026-07-22".to_string()), + ("hour".to_string(), "10".to_string()), + ]), + ] + ); +} + +#[cfg(not(windows))] +#[tokio::test] +async fn discover_format_partitions_treats_missing_root_as_empty() { + let tmp = tempfile::tempdir().unwrap(); + + let partitions = discover_partitions( + &tmp.path().join("missing"), + &["dt"], + false, + "__DEFAULT_PARTITION__", + ) + .await + .unwrap(); + + assert!(partitions.is_empty()); +} + +#[cfg(not(windows))] +#[tokio::test] +async fn discover_value_only_partitions_rejects_parent_traversal_value() { + let tmp = tempfile::tempdir().unwrap(); + std::fs::create_dir_all(tmp.path().join("%2E%2E")).unwrap(); + + let error = discover_partitions(tmp.path(), &["dt"], true, "__DEFAULT_PARTITION__") + .await + .unwrap_err(); + + assert!( + error.to_string().contains(".."), + "expected traversal value in error, got: {error}" + ); +} + +#[cfg(not(windows))] +#[tokio::test] +async fn discover_value_only_partitions_keeps_hidden_default_directory() { + for (default_name, directory, ignored_directory) in [ + ("__DEFAULT_PARTITION__", "__DEFAULT_PARTITION__", None), + (".NULL", ".NULL", Some(".staging")), + ("_NULL/%NA", "_NULL%2F%25NA", Some("_temporary")), + ] { + let tmp = tempfile::tempdir().unwrap(); + std::fs::create_dir_all(tmp.path().join(directory)).unwrap(); + if let Some(ignored_directory) = ignored_directory { + std::fs::create_dir_all(tmp.path().join(ignored_directory)).unwrap(); + } + + let partitions = discover_partitions(tmp.path(), &["dt"], true, default_name) + .await + .unwrap(); + + assert_eq!( + partitions, + vec![HashMap::from([( + "dt".to_string(), + default_name.to_string() + )])], + "{default_name}" + ); + } +} diff --git a/crates/paimon/tests/mock_server.rs b/crates/paimon/tests/mock_server.rs index 6c8a0048a..6d9b58939 100644 --- a/crates/paimon/tests/mock_server.rs +++ b/crates/paimon/tests/mock_server.rs @@ -35,12 +35,16 @@ use tokio::task::JoinHandle; use paimon::api::{ AlterDatabaseRequest, AlterTableRequest, AuditRESTResponse, ConfigResponse, - CreateFunctionRequest, CreateViewRequest, ErrorResponse, GetDatabaseResponse, - GetFunctionResponse, GetTableResponse, GetViewResponse, ListDatabasesResponse, - ListFunctionsResponse, ListTablesResponse, ListViewsResponse, RenameTableRequest, - ResourcePaths, + CreateFunctionRequest, CreatePartitionsRequest, CreateViewRequest, DropPartitionsRequest, + ErrorResponse, GetDatabaseResponse, GetFunctionResponse, GetTableResponse, GetViewResponse, + ListDatabasesResponse, ListFunctionsResponse, ListPartitionsResponse, ListTablesResponse, + ListViewsResponse, RenameTableRequest, ResourcePaths, }; use paimon::catalog::{Function, Identifier}; +use paimon::spec::Partition; + +type PartitionPageResponse = (Vec, Option); +type PartitionSpecPageResponse = (Vec>, Option); #[derive(Clone, Debug, Default)] struct MockState { @@ -48,17 +52,60 @@ struct MockState { tables: HashMap, views: HashMap, functions: HashMap, + partitions: HashMap>, + partition_page_responses: HashMap>, + partition_list_call_counts: HashMap, + partition_list_name_patterns: HashMap>>, view_function_endpoints_unsupported: bool, drop_view_error_status: Option, list_page_size: Option, no_permission_databases: HashSet, no_permission_tables: HashSet, + create_partitions_calls: Vec<(String, String, CreatePartitionsRequest)>, + drop_partitions_calls: Vec<(String, String, DropPartitionsRequest)>, + create_partitions_error_status: Option, + list_partitions_error_status: Option, /// ECS metadata role name (for token loader testing) ecs_role_name: Option, /// ECS metadata token (for token loader testing) ecs_token: Option, } +/// Match a partition spec against a partition-name pattern the way a catalog would. +/// +/// Only the pattern shape clients send is understood: `key=value` segments joined by `/`, +/// optionally ending in a `/%` wildcard. Values are compared unescaped, so a test that +/// needs escaping should assert the pattern the client sent instead of the rows returned. +fn partition_spec_matches_name_pattern(spec: &HashMap, pattern: &str) -> bool { + let prefix = pattern.strip_suffix("/%"); + let segments = prefix.unwrap_or(pattern).split('/').collect::>(); + if prefix.is_none() && segments.len() != spec.len() { + return false; + } + segments.iter().all(|segment| { + segment + .split_once('=') + .is_some_and(|(key, value)| spec.get(key).map(String::as_str) == Some(value)) + }) +} + +fn partition_from_spec(spec: HashMap) -> Partition { + Partition { + spec, + record_count: 0, + file_size_in_bytes: 0, + file_count: 0, + last_file_creation_time: 0, + total_buckets: 0, + done: false, + created_at: None, + created_by: None, + updated_at: None, + updated_by: None, + options: None, + } +} + fn paginate_names( names: Vec, params: &HashMap, @@ -282,6 +329,11 @@ impl RESTServer { // Also remove all tables in this database let prefix = format!("{name}."); s.tables.retain(|key, _| !key.starts_with(&prefix)); + s.partitions.retain(|key, _| !key.starts_with(&prefix)); + s.partition_page_responses + .retain(|key, _| !key.starts_with(&prefix)); + s.partition_list_call_counts + .retain(|key, _| !key.starts_with(&prefix)); s.no_permission_tables .retain(|key| !key.starts_with(&prefix)); (StatusCode::OK, Json(serde_json::json!(""))).into_response() @@ -724,6 +776,9 @@ impl RESTServer { } if s.tables.remove(&key).is_some() { + s.partitions.remove(&key); + s.partition_page_responses.remove(&key); + s.partition_list_call_counts.remove(&key); s.no_permission_tables.remove(&key); (StatusCode::OK, Json(serde_json::json!(""))).into_response() } else { @@ -771,6 +826,200 @@ impl RESTServer { } } + /// Handle POST /databases/:db/tables/:table/partitions - create partitions. + pub async fn create_partitions( + Path((db, table)): Path<(String, String)>, + Extension(state): Extension>, + Json(request): Json, + ) -> impl IntoResponse { + let mut inner = state.inner.lock().unwrap(); + inner + .create_partitions_calls + .push((db.clone(), table.clone(), request.clone())); + if let Some(status) = inner.create_partitions_error_status { + let message = if status == StatusCode::CONFLICT { + "Some partitions already exist" + } else { + "Invalid partition request" + }; + let error = ErrorResponse::new( + Some("partition".to_string()), + Some(table.clone()), + Some(message.to_string()), + Some(status.as_u16() as i32), + ); + return (status, Json(error)).into_response(); + } + + let key = format!("{db}.{table}"); + if !inner.tables.contains_key(&key) { + let error = ErrorResponse::new( + Some("table".to_string()), + Some(table), + Some("Not Found".to_string()), + Some(404), + ); + return (StatusCode::NOT_FOUND, Json(error)).into_response(); + } + + let registered_partitions = inner.partitions.entry(key).or_default(); + let has_conflict = request + .partition_specs + .iter() + .enumerate() + .any(|(index, spec)| { + registered_partitions + .iter() + .any(|partition| partition.spec == *spec) + || request.partition_specs[..index].contains(spec) + }); + if has_conflict && !request.ignore_if_exists { + let error = ErrorResponse::new( + Some("partition".to_string()), + Some(table), + Some("Some partitions already exist".to_string()), + Some(StatusCode::CONFLICT.as_u16() as i32), + ); + return (StatusCode::CONFLICT, Json(error)).into_response(); + } + + for spec in request.partition_specs { + if !registered_partitions + .iter() + .any(|partition| partition.spec == spec) + { + registered_partitions.push(partition_from_spec(spec)); + } + } + let response = json!({"success": true}); + (StatusCode::OK, Json(response)).into_response() + } + + /// Handle GET /databases/:db/tables/:table/partitions - list partitions. + pub async fn list_partitions( + Path((db, table)): Path<(String, String)>, + Query(params): Query>, + Extension(state): Extension>, + ) -> impl IntoResponse { + let mut inner = state.inner.lock().unwrap(); + let key = format!("{db}.{table}"); + let name_pattern = params.get("partitionNamePattern").cloned(); + inner + .partition_list_name_patterns + .entry(key.clone()) + .or_default() + .push(name_pattern.clone()); + if !inner.tables.contains_key(&key) { + let error = ErrorResponse::new( + Some("table".to_string()), + Some(table), + Some("Not Found".to_string()), + Some(404), + ); + return (StatusCode::NOT_FOUND, Json(error)).into_response(); + } + if let Some(status) = inner.list_partitions_error_status { + let error = ErrorResponse::new( + Some("partition".to_string()), + Some(table), + Some("Partition listing is not implemented".to_string()), + Some(status.as_u16() as i32), + ); + return (status, Json(error)).into_response(); + } + if inner.partition_page_responses.contains_key(&key) { + let request_index = { + let calls = inner + .partition_list_call_counts + .entry(key.clone()) + .or_default(); + let request_index = *calls; + *calls += 1; + request_index + }; + let responses = &inner.partition_page_responses[&key]; + let expected_token = request_index + .checked_sub(1) + .and_then(|index| responses.get(index)) + .and_then(|(_, token)| token.clone()); + let actual_token = params.get("pageToken").cloned(); + if actual_token != expected_token || request_index >= responses.len() { + let error = ErrorResponse::new( + Some("partition".to_string()), + Some(table), + Some(format!( + "Invalid page token: expected {expected_token:?}, got {actual_token:?}" + )), + Some(StatusCode::BAD_REQUEST.as_u16() as i32), + ); + return (StatusCode::BAD_REQUEST, Json(error)).into_response(); + } + let (partitions, next_page_token) = responses[request_index].clone(); + return ( + StatusCode::OK, + Json(ListPartitionsResponse::new( + Some(partitions), + next_page_token, + )), + ) + .into_response(); + } + let mut partitions = inner.partitions.get(&key).cloned().unwrap_or_default(); + if let Some(pattern) = name_pattern { + partitions + .retain(|partition| partition_spec_matches_name_pattern(&partition.spec, &pattern)); + } + ( + StatusCode::OK, + Json(ListPartitionsResponse::new(Some(partitions), None)), + ) + .into_response() + } + + /// Handle POST /databases/:db/tables/:table/partitions/drop - drop partitions. + pub async fn drop_partitions( + Path((db, table)): Path<(String, String)>, + Extension(state): Extension>, + Json(request): Json, + ) -> impl IntoResponse { + let mut inner = state.inner.lock().unwrap(); + inner + .drop_partitions_calls + .push((db.clone(), table.clone(), request.clone())); + + let key = format!("{db}.{table}"); + if !inner.tables.contains_key(&key) { + let error = ErrorResponse::new( + Some("table".to_string()), + Some(table), + Some("Not Found".to_string()), + Some(404), + ); + return (StatusCode::NOT_FOUND, Json(error)).into_response(); + } + + let registered_partitions = inner.partitions.entry(key).or_default(); + let has_missing = request.partition_specs.iter().any(|spec| { + !registered_partitions + .iter() + .any(|partition| partition.spec == *spec) + }); + if has_missing && !request.ignore_if_not_exists { + let error = ErrorResponse::new( + Some("partition".to_string()), + Some(table), + Some("Some partitions do not exist".to_string()), + Some(StatusCode::NOT_FOUND.as_u16() as i32), + ); + return (StatusCode::NOT_FOUND, Json(error)).into_response(); + } + + registered_partitions + .retain(|partition| !request.partition_specs.contains(&partition.spec)); + let response = json!({"success": true}); + (StatusCode::OK, Json(response)).into_response() + } + /// Handle POST /rename-table - rename a table. pub async fn rename_table( Extension(state): Extension>, @@ -827,7 +1076,6 @@ impl RESTServer { if s.no_permission_tables.remove(&source_key) { s.no_permission_tables.insert(dest_key.clone()); } - (StatusCode::OK, Json(serde_json::json!(""))).into_response() } else { let err = ErrorResponse::new( @@ -930,6 +1178,16 @@ impl RESTServer { self.inner.lock().unwrap().drop_view_error_status = status; } + /// Make the create-partitions endpoint return the given status. + pub fn set_create_partitions_error_status(&self, status: Option) { + self.inner.lock().unwrap().create_partitions_error_status = status; + } + + /// Make the list-partitions endpoint return the given status. + pub fn set_list_partitions_error_status(&self, status: Option) { + self.inner.lock().unwrap().list_partitions_error_status = status; + } + /// Add a table with schema and path to the server state. /// /// This is needed for `RESTCatalog::get_table` which requires @@ -972,6 +1230,100 @@ impl RESTServer { let mut s = self.inner.lock().unwrap(); s.no_permission_tables.insert(format!("{database}.{table}")); } + + /// Set whether a stored table is external. + pub fn set_table_external(&self, database: &str, table: &str, is_external: bool) { + let key = format!("{database}.{table}"); + let mut state = self.inner.lock().unwrap(); + state + .tables + .get_mut(&key) + .unwrap_or_else(|| panic!("table {key} does not exist")) + .is_external = Some(is_external); + } + + /// Set the catalog-registered partitions for a table. + pub fn set_table_partitions( + &self, + database: &str, + table: &str, + partition_specs: Vec>, + ) { + let partitions = partition_specs + .into_iter() + .map(partition_from_spec) + .collect(); + let key = format!("{database}.{table}"); + let mut inner = self.inner.lock().unwrap(); + assert!( + inner.tables.contains_key(&key), + "table {key} does not exist" + ); + inner.partitions.insert(key.clone(), partitions); + inner.partition_page_responses.remove(&key); + inner.partition_list_call_counts.remove(&key); + } + + /// Set explicit list-partitions pages and response tokens in request order. + pub fn set_table_partition_page_responses( + &self, + database: &str, + table: &str, + responses: Vec, + ) { + let responses = responses + .into_iter() + .map(|(specs, token)| (specs.into_iter().map(partition_from_spec).collect(), token)) + .collect(); + let key = format!("{database}.{table}"); + let mut inner = self.inner.lock().unwrap(); + assert!( + inner.tables.contains_key(&key), + "table {key} does not exist" + ); + inner + .partition_page_responses + .insert(key.clone(), responses); + inner.partition_list_call_counts.remove(&key); + } + + /// Return the `partitionNamePattern` of every partition-list request the table + /// received, in order, with `None` where the client sent no pattern. + pub fn table_partition_list_name_patterns( + &self, + database: &str, + table: &str, + ) -> Vec> { + self.inner + .lock() + .unwrap() + .partition_list_name_patterns + .get(&format!("{database}.{table}")) + .cloned() + .unwrap_or_default() + } + + /// Return how many paged partition-list requests the table received. + pub fn table_partition_list_call_count(&self, database: &str, table: &str) -> usize { + self.inner + .lock() + .unwrap() + .partition_list_call_counts + .get(&format!("{database}.{table}")) + .copied() + .unwrap_or_default() + } + + /// Return all create-partitions calls received by the server. + pub fn create_partitions_calls(&self) -> Vec<(String, String, CreatePartitionsRequest)> { + self.inner.lock().unwrap().create_partitions_calls.clone() + } + + /// Return all drop-partitions calls received by the server. + pub fn drop_partitions_calls(&self) -> Vec<(String, String, DropPartitionsRequest)> { + self.inner.lock().unwrap().drop_partitions_calls.clone() + } + /// Get the server URL. pub fn url(&self) -> Option { self.addr.map(|a| format!("http://{a}")) @@ -1089,6 +1441,14 @@ pub async fn start_mock_server( .post(RESTServer::alter_table) .delete(RESTServer::drop_table), ) + .route( + &format!("{prefix}/databases/:db/tables/:table/partitions"), + get(RESTServer::list_partitions).post(RESTServer::create_partitions), + ) + .route( + &format!("{prefix}/databases/:db/tables/:table/partitions/drop"), + post(RESTServer::drop_partitions), + ) .route( &format!("{prefix}/databases/:db/views"), get(RESTServer::list_views).post(RESTServer::create_view), diff --git a/crates/paimon/tests/rest_api_test.rs b/crates/paimon/tests/rest_api_test.rs index 99b952ce8..fc8791194 100644 --- a/crates/paimon/tests/rest_api_test.rs +++ b/crates/paimon/tests/rest_api_test.rs @@ -24,7 +24,7 @@ use std::collections::HashMap; use paimon::api::auth::{DLFECSTokenLoader, DLFToken, DLFTokenLoader}; use paimon::api::rest_api::RESTApi; -use paimon::api::ConfigResponse; +use paimon::api::{ConfigResponse, CreatePartitionsRequest, DropPartitionsRequest}; use paimon::catalog::{Function, FunctionDefinition, Identifier, ViewSchema}; use paimon::common::Options; use paimon::spec::DataField; @@ -71,6 +71,13 @@ async fn setup_test_server(initial_dbs: Vec<&str>) -> TestContext { TestContext { server, api, url } } +async fn setup_partition_api() -> (TestContext, Identifier) { + let ctx = setup_test_server(vec!["default"]).await; + let identifier = Identifier::new("default", "managed_table"); + ctx.server.add_table("default", "managed_table"); + (ctx, identifier) +} + // ==================== Database Tests ==================== #[tokio::test] async fn test_list_databases() { @@ -594,6 +601,142 @@ async fn test_drop_table_no_permission() { assert!(result.is_err(), "dropping no-permission table should fail"); } +// ==================== Partition Tests ==================== + +#[tokio::test] +async fn test_partition_mutations_post_expected_requests() { + let (ctx, identifier) = setup_partition_api().await; + let partition_specs = vec![HashMap::from([ + ("dt".to_string(), "2026-07-22".to_string()), + ("hour".to_string(), "10".to_string()), + ])]; + + ctx.api + .create_partitions(&identifier, partition_specs.clone(), false) + .await + .unwrap(); + + assert_eq!( + ctx.server.create_partitions_calls(), + vec![( + "default".to_string(), + "managed_table".to_string(), + CreatePartitionsRequest::new(partition_specs.clone(), false), + )] + ); + + ctx.api + .drop_partitions(&identifier, partition_specs.clone(), false) + .await + .unwrap(); + + assert_eq!( + ctx.server.drop_partitions_calls(), + vec![( + "default".to_string(), + "managed_table".to_string(), + DropPartitionsRequest::new(partition_specs, false), + )] + ); +} + +#[tokio::test] +async fn test_list_partitions_follows_non_empty_token_after_empty_page() { + let (ctx, identifier) = setup_partition_api().await; + let expected = HashMap::from([("dt".to_string(), "2026-07-22".to_string())]); + ctx.server.set_table_partition_page_responses( + "default", + "managed_table", + vec![ + (Vec::new(), Some("1".to_string())), + (vec![expected.clone()], None), + ], + ); + + let partitions = ctx.api.list_partitions(&identifier).await.unwrap(); + + assert_eq!(partitions.len(), 1); + assert_eq!(partitions[0].spec, expected); +} + +#[tokio::test] +async fn test_list_partitions_stops_on_empty_next_page_token() { + let (ctx, identifier) = setup_partition_api().await; + let expected = HashMap::from([("dt".to_string(), "2026-07-22".to_string())]); + let unexpected = HashMap::from([("dt".to_string(), "2026-07-23".to_string())]); + ctx.server.set_table_partition_page_responses( + "default", + "managed_table", + vec![ + (vec![expected.clone()], Some(String::new())), + (vec![unexpected], None), + ], + ); + + let partitions = ctx.api.list_partitions(&identifier).await.unwrap(); + + assert_eq!( + partitions + .into_iter() + .map(|partition| partition.spec) + .collect::>(), + vec![expected] + ); + assert_eq!( + ctx.server + .table_partition_list_call_count("default", "managed_table"), + 1 + ); +} + +#[tokio::test] +async fn test_list_partitions_rejects_repeated_page_token() { + let (ctx, identifier) = setup_partition_api().await; + ctx.server.set_table_partition_page_responses( + "default", + "managed_table", + vec![ + ( + vec![HashMap::from([( + "dt".to_string(), + "2026-07-20".to_string(), + )])], + Some("a".to_string()), + ), + ( + vec![HashMap::from([( + "dt".to_string(), + "2026-07-21".to_string(), + )])], + Some("b".to_string()), + ), + ( + vec![HashMap::from([( + "dt".to_string(), + "2026-07-22".to_string(), + )])], + Some("a".to_string()), + ), + ], + ); + + let error = ctx.api.list_partitions(&identifier).await.unwrap_err(); + + let paimon::Error::UnexpectedError { message, .. } = error else { + panic!("expected a repeated-page-token error, got: {error}"); + }; + assert_eq!( + message, + "REST catalog returned partition page token 'a' more than once for table \ + default.managed_table" + ); + assert_eq!( + ctx.server + .table_partition_list_call_count("default", "managed_table"), + 3 + ); +} + // ==================== Rename Table Tests ==================== #[tokio::test] diff --git a/crates/paimon/tests/rest_catalog_test.rs b/crates/paimon/tests/rest_catalog_test.rs index 2000f9e40..c46434741 100644 --- a/crates/paimon/tests/rest_catalog_test.rs +++ b/crates/paimon/tests/rest_catalog_test.rs @@ -84,6 +84,30 @@ fn test_schema() -> Schema { .expect("Failed to build schema") } +fn format_table_schema(options: &[(&str, &str)]) -> Schema { + let mut builder = Schema::builder() + .column("dt", DataType::VarChar(VarCharType::new(255).unwrap())) + .column("id", DataType::BigInt(BigIntType::new())) + .partition_keys(["dt"]) + .option("type", "format-table") + .option("file.format", "parquet") + .option("metastore.partitioned-table", "true"); + for (key, value) in options { + builder = builder.option(*key, *value); + } + builder.build().unwrap() +} + +fn add_internal_table_with_schema( + server: &RESTServer, + table: &str, + schema: Schema, + location: &str, +) { + server.add_table_with_schema("default", table, schema, location); + server.set_table_external("default", table, false); +} + fn blob_schema(options: &[(&str, &str)]) -> Schema { let mut builder = Schema::builder() .column("id", DataType::Int(IntType::new())) @@ -167,6 +191,162 @@ fn collect_blob_rows(batches: &[RecordBatch]) -> Vec<(i32, String, Option>(); + + ctx.catalog + .create_partitions(&identifier, Vec::new(), false) + .await + .unwrap(); + ctx.catalog + .create_partitions(&identifier, partition_specs.clone(), true) + .await + .unwrap(); + + let calls = ctx.server.create_partitions_calls(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].2.partition_specs, partition_specs[..1000]); + assert_eq!(calls[1].2.partition_specs, partition_specs[1000..]); + assert!(calls.iter().all(|(_, _, request)| request.ignore_if_exists)); +} + +#[tokio::test] +async fn test_rest_catalog_keeps_non_idempotent_create_in_one_request() { + let ctx = setup_catalog(vec!["default"]).await; + let identifier = Identifier::new("default", "managed_table"); + ctx.server.add_table("default", "managed_table"); + let partition_specs = (0..2500) + .map(|value| HashMap::from([("dt".to_string(), value.to_string())])) + .collect::>(); + + ctx.catalog + .create_partitions(&identifier, partition_specs.clone(), false) + .await + .unwrap(); + + let calls = ctx.server.create_partitions_calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].2.partition_specs, partition_specs); + assert!(!calls[0].2.ignore_if_exists); +} + +#[tokio::test] +async fn test_rest_catalog_maps_create_partition_conflict() { + let ctx = setup_catalog(vec!["default"]).await; + let identifier = Identifier::new("default", "managed_table"); + ctx.server + .set_create_partitions_error_status(Some(StatusCode::CONFLICT)); + + let error = ctx + .catalog + .create_partitions( + &identifier, + vec![HashMap::from([( + "dt".to_string(), + "2026-07-22".to_string(), + )])], + false, + ) + .await + .unwrap_err(); + + assert!(matches!( + error, + paimon::Error::DataInvalid { message, .. } + if message.contains("default.managed_table") + && message.contains("already exist") + )); +} + +#[tokio::test] +async fn test_rest_catalog_maps_invalid_partition_request() { + let ctx = setup_catalog(vec!["default"]).await; + let identifier = Identifier::new("default", "managed_table"); + ctx.server + .set_create_partitions_error_status(Some(StatusCode::BAD_REQUEST)); + + let error = ctx + .catalog + .create_partitions( + &identifier, + vec![HashMap::from([( + "unknown".to_string(), + "value".to_string(), + )])], + false, + ) + .await + .unwrap_err(); + + assert!(matches!( + error, + paimon::Error::DataInvalid { message, .. } + if message.contains("default.managed_table") + )); +} + +#[tokio::test] +async fn test_rest_catalog_skips_empty_and_batches_drop_partitions() { + let ctx = setup_catalog(vec!["default"]).await; + let identifier = Identifier::new("default", "managed_table"); + ctx.server.add_table("default", "managed_table"); + let partition_specs = (0..1001) + .map(|value| HashMap::from([("dt".to_string(), value.to_string())])) + .collect::>(); + ctx.server + .set_table_partitions("default", "managed_table", partition_specs.clone()); + + ctx.catalog + .drop_partitions(&identifier, Vec::new()) + .await + .unwrap(); + ctx.catalog + .drop_partitions(&identifier, partition_specs.clone()) + .await + .unwrap(); + + let calls = ctx.server.drop_partitions_calls(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].2.partition_specs, partition_specs[..1000]); + assert_eq!(calls[1].2.partition_specs, partition_specs[1000..]); + assert!(calls + .iter() + .all(|(_, _, request)| request.ignore_if_not_exists)); +} + +#[tokio::test] +async fn test_rest_catalog_partition_operations_map_missing_table() { + let ctx = setup_catalog(vec!["default"]).await; + let identifier = Identifier::new("default", "missing"); + let partition_specs = vec![HashMap::from([( + "dt".to_string(), + "2026-07-22".to_string(), + )])]; + + let create_error = ctx + .catalog + .create_partitions(&identifier, partition_specs.clone(), false) + .await + .unwrap_err(); + let drop_error = ctx + .catalog + .drop_partitions(&identifier, partition_specs.clone()) + .await + .unwrap_err(); + for error in [create_error, drop_error] { + assert!(matches!( + error, + paimon::Error::TableNotExist { full_name } if full_name == "default.missing" + )); + } +} + // ==================== Database Tests ==================== #[tokio::test] @@ -579,6 +759,308 @@ async fn test_rest_catalog_reads_format_table() { ); } +#[tokio::test] +async fn test_rest_catalog_validates_dynamic_managed_partition_options() { + let ctx = setup_catalog(vec!["default"]).await; + let schema = format_table_schema(&[]); + let identifier = Identifier::new("default", "managed_format"); + add_internal_table_with_schema( + &ctx.server, + "managed_format", + schema, + "memory:/managed_format", + ); + + let table = ctx.catalog.get_table(&identifier).await.unwrap(); + + assert!(table.has_catalog_managed_partitions()); + for (key, value) in [ + ("metastore.partitioned-table", "false"), + ("format-table.partition-path-only-value", "true"), + ("format-table.implementation", "engine"), + ] { + let error = table + .copy_with_time_travel(HashMap::from([(key.to_string(), value.to_string())])) + .await + .unwrap_err(); + assert!( + matches!(error, paimon::Error::DataInvalid { .. }) && error.to_string().contains(key), + "expected {key} validation error, got: {error}" + ); + } + + let copied = table + .copy_with_time_travel(HashMap::from([ + ( + "metastore.partitioned-table".to_string(), + "TRUE".to_string(), + ), + ( + "format-table.partition-path-only-value".to_string(), + "false".to_string(), + ), + ( + "format-table.implementation".to_string(), + "PAIMON".to_string(), + ), + ("type".to_string(), "table".to_string()), + ])) + .await + .unwrap(); + assert!(copied.has_catalog_managed_partitions()); + assert!(copied + .new_read_builder() + .new_scan() + .plan() + .await + .unwrap() + .splits() + .is_empty()); +} + +#[tokio::test] +async fn test_rest_catalog_rejects_enabling_managed_partitions_in_dynamic_options() { + let ctx = setup_catalog(vec!["default"]).await; + let schema = format_table_schema(&[("metastore.partitioned-table", "false")]); + let identifier = Identifier::new("default", "unmanaged_format"); + add_internal_table_with_schema( + &ctx.server, + "unmanaged_format", + schema, + "memory:/unmanaged_format", + ); + + let table = ctx.catalog.get_table(&identifier).await.unwrap(); + assert!(!table.has_catalog_managed_partitions()); + + let error = table + .copy_with_time_travel(HashMap::from([( + "metastore.partitioned-table".to_string(), + "true".to_string(), + )])) + .await + .unwrap_err(); + assert!( + matches!(error, paimon::Error::DataInvalid { .. }) + && error.to_string().contains("metastore.partitioned-table"), + "expected partition source validation error, got: {error}" + ); +} + +#[tokio::test] +async fn test_rest_catalog_rejects_invalid_catalog_managed_partition_options() { + let ctx = setup_catalog(vec!["default"]).await; + for (table, option, bad_value) in [ + ( + "invalid_partition_source", + "metastore.partitioned-table", + "tru", + ), + ( + "invalid_format_implementation", + "format-table.implementation", + "unknown", + ), + ( + "invalid_partition_layout", + "format-table.partition-path-only-value", + "tru", + ), + ] { + let schema = format_table_schema(&[(option, bad_value)]); + add_internal_table_with_schema(&ctx.server, table, schema, &format!("memory:/{table}")); + + let error = ctx + .catalog + .get_table(&Identifier::new("default", table)) + .await + .unwrap_err(); + + assert!( + matches!(error, paimon::Error::ConfigInvalid { .. }) + && error.to_string().contains(option) + && error.to_string().contains(bad_value), + "expected {option}={bad_value} validation error, got: {error}" + ); + } +} + +#[tokio::test] +async fn test_rest_catalog_rejects_external_catalog_managed_format_table() { + let ctx = setup_catalog(vec!["default"]).await; + let schema = format_table_schema(&[]); + let identifier = Identifier::new("default", "external_managed_format"); + ctx.server.add_table_with_schema( + "default", + "external_managed_format", + schema, + "memory:/external_managed_format", + ); + + let error = ctx.catalog.get_table(&identifier).await.unwrap_err(); + + assert!(matches!( + error, + paimon::Error::DataInvalid { message, .. } + if message.contains("default.external_managed_format") + && message.contains("internal") + )); +} + +#[tokio::test] +async fn test_rest_catalog_rejects_engine_managed_format_table() { + let ctx = setup_catalog(vec!["default"]).await; + let schema = format_table_schema(&[("format-table.implementation", "engine")]); + let identifier = Identifier::new("default", "engine_managed_format"); + add_internal_table_with_schema( + &ctx.server, + "engine_managed_format", + schema, + "memory:/engine_managed_format", + ); + + let error = ctx.catalog.get_table(&identifier).await.unwrap_err(); + + assert!(matches!( + error, + paimon::Error::DataInvalid { message, .. } + if message.contains("metastore.partitioned-table") + && message.contains("format-table.implementation=engine") + )); +} + +#[cfg(not(windows))] +#[tokio::test] +async fn test_managed_format_scan_uses_registered_partition_paths_once() { + let tmp = tempfile::tempdir().unwrap(); + for dt in ["2026-07-21", "2026-07-22"] { + let partition_dir = tmp.path().join(format!("dt={dt}")); + std::fs::create_dir_all(&partition_dir).unwrap(); + std::fs::write(partition_dir.join("part-0.parquet"), b"data").unwrap(); + } + let table_path = format!("file://{}", tmp.path().display()); + + let ctx = setup_catalog(vec!["default"]).await; + let schema = format_table_schema(&[]); + let identifier = Identifier::new("default", "managed_visibility"); + add_internal_table_with_schema(&ctx.server, "managed_visibility", schema, &table_path); + ctx.server.set_table_partitions( + "default", + "managed_visibility", + vec![ + HashMap::from([("dt".to_string(), "2026-07-22".to_string())]), + HashMap::from([("dt".to_string(), "2026-07-22".to_string())]), + ], + ); + + let table = ctx.catalog.get_table(&identifier).await.unwrap(); + let plan = table.new_read_builder().new_scan().plan().await.unwrap(); + + assert_eq!(plan.splits().len(), 1); + assert!(plan.splits()[0].bucket_path().ends_with("/dt=2026-07-22")); +} + +#[cfg(not(windows))] +#[tokio::test] +async fn test_managed_format_scan_treats_registered_missing_directory_as_empty() { + let tmp = tempfile::tempdir().unwrap(); + let table_path = format!("file://{}", tmp.path().display()); + + let ctx = setup_catalog(vec!["default"]).await; + let schema = format_table_schema(&[]); + let identifier = Identifier::new("default", "managed_missing_directory"); + add_internal_table_with_schema( + &ctx.server, + "managed_missing_directory", + schema, + &table_path, + ); + ctx.server.set_table_partitions( + "default", + "managed_missing_directory", + vec![HashMap::from([( + "dt".to_string(), + "2026-07-22".to_string(), + )])], + ); + + let table = ctx.catalog.get_table(&identifier).await.unwrap(); + let plan = table.new_read_builder().new_scan().plan().await.unwrap(); + + assert!(plan.splits().is_empty()); +} + +#[cfg(not(windows))] +#[tokio::test] +async fn test_managed_format_scan_propagates_partition_listing_error() { + let tmp = tempfile::tempdir().unwrap(); + let table_path = format!("file://{}", tmp.path().display()); + + let ctx = setup_catalog(vec!["default"]).await; + let schema = format_table_schema(&[]); + let identifier = Identifier::new("default", "managed_scan_error"); + add_internal_table_with_schema(&ctx.server, "managed_scan_error", schema, &table_path); + ctx.server + .set_list_partitions_error_status(Some(StatusCode::NOT_IMPLEMENTED)); + + let table = ctx.catalog.get_table(&identifier).await.unwrap(); + let error = table + .new_read_builder() + .new_scan() + .plan() + .await + .unwrap_err(); + + assert!(matches!( + error, + paimon::Error::RestApi { + source: paimon::api::RestError::NotImplemented { .. } + } + )); +} + +#[cfg(not(windows))] +#[tokio::test] +async fn test_managed_format_scan_reports_table_for_malformed_partition_metadata() { + let tmp = tempfile::tempdir().unwrap(); + let table_path = format!("file://{}", tmp.path().display()); + + let ctx = setup_catalog(vec!["default"]).await; + let schema = format_table_schema(&[]); + let identifier = Identifier::new("default", "managed_corrupt_metadata"); + add_internal_table_with_schema(&ctx.server, "managed_corrupt_metadata", schema, &table_path); + ctx.server.set_table_partitions( + "default", + "managed_corrupt_metadata", + vec![HashMap::from([ + ("dt".to_string(), "2026-07-22".to_string()), + ("unexpected".to_string(), "value".to_string()), + ])], + ); + + let table = ctx.catalog.get_table(&identifier).await.unwrap(); + let error = table + .new_read_builder() + .new_scan() + .plan() + .await + .unwrap_err(); + + match error { + paimon::Error::DataInvalid { + message, + source: Some(source), + } => { + assert!(message.contains("invalid partition metadata")); + assert!(message.contains("default.managed_corrupt_metadata")); + let cause = source.to_string(); + assert!(cause.contains("unexpected"), "unexpected cause: {cause}"); + assert!(cause.contains("dt"), "unexpected cause: {cause}"); + } + other => panic!("expected invalid partition metadata, got: {other}"), + } +} + #[cfg(not(windows))] #[tokio::test] async fn test_rest_catalog_prunes_format_table_partition_filter() { diff --git a/docs/src/sql.md b/docs/src/sql.md index b0da06bf8..6fa5b355d 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -40,7 +40,7 @@ Mosaic support is always available and currently read-only. SQL queries can read SQL support has two layers: - DataFusion provides the parser, query planner, optimizer, execution engine, expressions, scalar functions, aggregate functions, and window functions. SQL statements that `SQLContext` does not intercept are delegated to DataFusion. This includes the DataFusion SQL surface for `SELECT` queries, CTEs (including recursive CTEs), subqueries, joins including `LATERAL` joins, SQL lambda functions, grouping, `HAVING`, window clauses, `QUALIFY`, set operations, `ORDER BY`, `LIMIT`/`OFFSET`, `EXPLAIN`, information-schema commands such as `SHOW TABLES`, `DESCRIBE`, `COPY`, and ordinary `INSERT`. -- Paimon-specific table management and row-level writes are implemented by `SQLContext`. This includes Paimon `CREATE TABLE`, `ALTER TABLE`, `DROP TABLE`, `CREATE TEMPORARY TABLE`, `CREATE TEMPORARY VIEW`, REST Catalog persistent `CREATE VIEW`, `DROP VIEW`, and `CREATE FUNCTION`, `DROP TEMPORARY TABLE` / `VIEW`, `INSERT OVERWRITE ... PARTITION`, `UPDATE`, `DELETE`, `MERGE INTO`, `TRUNCATE TABLE`, `ALTER TABLE ... DROP PARTITION`, `CALL sys.*`, Paimon time travel, and `SET` / `RESET 'paimon.*'`. +- Paimon-specific table management and row-level writes are implemented by `SQLContext`. This includes Paimon `CREATE TABLE`, `ALTER TABLE`, `DROP TABLE`, `CREATE TEMPORARY TABLE`, `CREATE TEMPORARY VIEW`, REST Catalog persistent `CREATE VIEW`, `DROP VIEW`, and `CREATE FUNCTION`, `DROP TEMPORARY TABLE` / `VIEW`, `INSERT OVERWRITE ... PARTITION`, `UPDATE`, `DELETE`, `MERGE INTO`, `TRUNCATE TABLE`, `ALTER TABLE ... ADD PARTITION`, `ALTER TABLE ... DROP PARTITION`, `SHOW PARTITIONS`, `MSCK REPAIR TABLE`, `CALL sys.*`, Paimon time travel, and `SET` / `RESET 'paimon.*'`. Not every DataFusion DDL/DML statement maps to a Paimon table operation. For Paimon catalogs, `CREATE EXTERNAL TABLE`, `LOCATION`, `CREATE MATERIALIZED VIEW`, and persistent `CREATE TABLE AS SELECT` are rejected or not implemented. Persistent `CREATE FUNCTION` is supported only for the REST Catalog SQL scalar form documented below. DataFusion `COPY` can export query results to files; it does not create or commit Paimon table files. @@ -875,6 +875,85 @@ Multiple partition key-value pairs can be specified: ALTER TABLE paimon.my_db.events DROP PARTITION (dt = '2024-01-01', region = 'us'); ``` +## Format Table Partitions + +A `type=format-table` table loaded from a REST Catalog can have its partitions managed +by the catalog instead of discovered from the directory layout. The catalog registration +is then the authoritative partition set: it decides what a scan reads, and a directory +nobody registered is not part of the table. + +A table opts in with `'metastore.partitioned-table' = 'true'`. The statements below apply +only to such a table — a partitioned internal Format Table, loaded from a REST Catalog, +with a non-`engine` `format-table.implementation`. + +### SHOW PARTITIONS + +```sql +SHOW PARTITIONS paimon.my_db.events; +SHOW PARTITIONS paimon.my_db.events PARTITION (dt = '2024-01-01'); +``` + +Partition names are returned in the escaped `key=value/...` form, sorted. The optional +`PARTITION` clause keeps only the partitions matching the given values; it may fix any +subset of the partition keys. + +### ADD PARTITION + +```sql +ALTER TABLE paimon.my_db.events ADD PARTITION (dt = '2024-01-01', region = 'us'); +ALTER TABLE paimon.my_db.events ADD IF NOT EXISTS PARTITION (dt = '2024-01-01') + PARTITION (dt = '2024-01-02'); +``` + +Every specification must fix all partition keys. The partitions are registered with the +catalog first and their directories are created afterwards, so a failure to create a +directory leaves the registration in place and re-running the statement with +`IF NOT EXISTS` completes it. Without `IF NOT EXISTS`, an already registered partition is +an error. A custom `LOCATION` is not supported: the directory always follows the table's +partition layout. + +### DROP PARTITION + +```sql +ALTER TABLE paimon.my_db.events DROP PARTITION (dt = '2024-01-01', region = 'us'); +ALTER TABLE paimon.my_db.events DROP IF EXISTS PARTITION (dt = '2024-01-01'); +``` + +A specification that fixes only some of the partition keys drops every registered +partition it matches, and the fixed keys need not be a leading prefix — on a +`(dt, region)` table, `DROP PARTITION (region = 'us')` drops the `us` partition of every +date. A specification that fixes all keys names one partition, so a missing one is an +error unless `IF EXISTS` is given; a partial one describes a set that is allowed to come +out empty. + +One statement may carry several specifications, each `DROP PARTITION` in its own clause: + +```sql +ALTER TABLE paimon.my_db.events DROP PARTITION (dt = '2024-01-01'), + DROP PARTITION (dt = '2024-01-02'); +``` + +The catalog registration is removed first, then the directory is deleted by the client — +the catalog never deletes data. If the deletion fails the partition is already invisible +and the directory may survive; repair the file system and remove it there rather than +re-registering it. + +### MSCK REPAIR TABLE + +Reconcile the catalog registrations with the directories that actually exist: + +```sql +MSCK REPAIR TABLE paimon.my_db.events; -- same as ADD PARTITIONS +MSCK REPAIR TABLE paimon.my_db.events ADD PARTITIONS; -- register discovered directories +MSCK REPAIR TABLE paimon.my_db.events DROP PARTITIONS; -- unregister vanished directories +MSCK REPAIR TABLE paimon.my_db.events SYNC PARTITIONS; -- both +``` + +Repair is metadata-only: it never deletes data. Directory discovery and the catalog +listing both complete before any change is made, so a listing failure cannot turn a +truncated view of the table into a `DROP` diff. There is no dry-run and no scope +argument — repair always covers the whole table. + ## Procedures Use `CALL` to invoke built-in procedures. All procedures are under the `sys` namespace.