diff --git a/Cargo.lock b/Cargo.lock index ffcb3c7a..3260dab1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4606,6 +4606,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 17acaea0..f7072fc3 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 00000000..f0b1738f --- /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 05e4037a..ae8a42ca 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 3190487b..0102e5f5 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -29,7 +29,10 @@ //! - `ALTER TABLE db.t DROP COLUMN col` //! - `ALTER TABLE db.t RENAME COLUMN old TO new` //! - `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` @@ -47,7 +50,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; @@ -55,27 +58,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::{ - 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, Value as SqlValue, + AddDropSync, 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, + 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}; @@ -350,8 +358,8 @@ impl SQLContext { &self.dynamic_options } - /// Execute a SQL statement. ALTER TABLE is handled by Paimon directly; - /// everything else is delegated to DataFusion. + /// Execute SQL, intercepting Paimon-specific DDL/DML and delegating the + /// remaining statements to DataFusion. pub async fn sql(&self, sql: &str) -> DFResult { let is_create_table = looks_like_create_table(sql); let (rewritten_sql, partition_keys) = if is_create_table { @@ -363,7 +371,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 { @@ -388,6 +398,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( @@ -442,6 +461,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) @@ -1027,9 +1047,54 @@ 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 let [AlterTableOperation::DropPartitions { + partitions, + if_exists: partition_if_exists, + }] = operations + { + return self + .handle_drop_partition( + catalog, + &identifier, + partitions, + *partition_if_exists, + if_exists, + ) + .await; + } + if operations + .iter() + .any(|operation| matches!(operation, AlterTableOperation::DropPartitions { .. })) + { + return Err(DataFusionError::Plan( + "ALTER TABLE DROP PARTITION supports one partition specification per statement" + .to_string(), + )); + } + if operations + .iter() + .any(|operation| matches!(operation, AlterTableOperation::AddPartitions { .. })) + && operations.len() != 1 + { + return Err(DataFusionError::Plan( + "ALTER TABLE ADD PARTITION cannot be combined with other operations".to_string(), + )); + } + let mut changes = Vec::new(); let mut rename_to: Option = None; @@ -1073,16 +1138,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; } @@ -1642,34 +1708,109 @@ impl SQLContext { Ok(()) } - async fn handle_drop_partitions( + async fn handle_drop_partition( &self, catalog: &Arc, identifier: &Identifier, - partitions: &[SqlExpr], - if_exists: bool, + expressions: &[SqlExpr], + ignore_if_not_exists: bool, + ignore_if_table_not_exists: bool, ) -> DFResult { - if partitions.is_empty() { + if 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)), }; + if table.has_catalog_managed_partitions() { + ensure_catalog_managed_format_table(&table, "ALTER TABLE DROP PARTITION")?; + let requested = parse_format_partition_spec(expressions, &table, true)?; + let requested_normalized = normalize_catalog_partition_spec(&requested, &table, true)?; + 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 has_exact_match = registered.iter().any(|(spec, _)| spec == &requested); + 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::new(); + for (spec, normalized) in registered { + let matches = if has_exact_match { + spec == requested + } else { + normalized == requested_normalized + }; + if matches { + let relative_path = partition_paths + .relative_path(&spec) + .map_err(to_datafusion_error)?; + selected.push((spec, format!("{table_path}/{relative_path}"))); + } + } + if !has_exact_match && selected.len() > 1 { + return Err(DataFusionError::Plan(format!( + "Partition {requested:?} matches multiple catalog registrations in table {}; use an exact catalog value", + identifier.full_name() + ))); + } + + if selected.is_empty() { + if ignore_if_not_exists { + return ok_result(&self.ctx); + } + return Err(DataFusionError::Plan(format!( + "Partition {requested:?} does not exist in table {}", + identifier.full_name() + ))); + } + + 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)?; + } + return ok_result(&self.ctx); + } + if CoreOptions::new(table.schema().options()).is_format_table() { + ensure_catalog_managed_format_table(&table, "ALTER TABLE DROP PARTITION")?; + } + let partition_values = parse_partition_values( - partitions, + expressions, table.schema().fields(), table.schema().partition_keys(), )?; - - let wb = table.new_write_builder(); - let commit = wb.try_new_commit().map_err(to_datafusion_error)?; + let commit = table + .new_write_builder() + .try_new_commit() + .map_err(to_datafusion_error)?; commit .truncate_partitions(partition_values) .await @@ -1678,6 +1819,151 @@ impl SQLContext { 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 @@ -1830,13 +2116,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" @@ -1845,6 +2125,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)?; @@ -1958,11 +2263,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() @@ -2008,7 +2376,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( @@ -2651,6 +3019,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`. /// @@ -2666,30 +3325,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" ))); @@ -2783,6 +3421,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 { @@ -2806,14 +3459,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!( @@ -2825,6 +3477,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}") @@ -6872,6 +7537,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 00000000..cb0cdff9 --- /dev/null +++ b/crates/integrations/datafusion/tests/rest_format_partition_sql.rs @@ -0,0 +1,227 @@ +// 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, + )], + ); +} + +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 84ee7c5c..ea292b03 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 809b9021..5773a85a 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 38c39c9a..569d1b92 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 cdde2b23..48a38105 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, @@ -556,6 +557,41 @@ 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> { let database = identifier.database(); @@ -564,17 +600,29 @@ impl RESTApi { 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()) .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) diff --git a/crates/paimon/src/catalog/mod.rs b/crates/paimon/src/catalog/mod.rs index 38a414b9..bc9487b6 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 2057d3f5..6db8e13f 100644 --- a/crates/paimon/src/catalog/rest/rest_catalog.rs +++ b/crates/paimon/src/catalog/rest/rest_catalog.rs @@ -37,6 +37,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 @@ -363,16 +365,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)), } } @@ -388,17 +438,23 @@ impl Catalog for RESTCatalog { .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 // ============================================================================ @@ -445,6 +501,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 fce04169..4f95a102 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"; @@ -125,6 +127,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`. @@ -537,6 +545,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) @@ -1556,6 +1601,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 2902ee72..e03d2e7d 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 3db1884c..56ab1d3d 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 00000000..7fdd1cc7 --- /dev/null +++ b/crates/paimon/src/table/format_partition.rs @@ -0,0 +1,324 @@ +// 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("/")) + } + + /// 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(_) => 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, + } +} + +/// 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, + } +} + +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)) + ); + assert_eq!( + parse_format_partition_value("yes", &DataType::Boolean(BooleanType::new())), + 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 28907664..61856663 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,66 @@ 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 partitions = rest_env + .api() + .list_partitions(rest_env.identifier()) + .await?; + 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(); + 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), } } @@ -341,7 +415,7 @@ fn leading_equality_partition_path( let Some(datum) = values[idx] else { break; }; - let value = partition_value_from_datum( + let value = format_partition_value( datum, partition_fields[idx].data_type(), default_partition_name, @@ -376,7 +450,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 +476,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 +528,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 +537,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 +576,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 0b76d283..8b7298c1 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; @@ -97,6 +98,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; @@ -255,7 +259,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. @@ -317,6 +330,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 @@ -360,6 +376,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 b7771001..b97c953a 100644 --- a/crates/paimon/src/table/rest_env.rs +++ b/crates/paimon/src/table/rest_env.rs @@ -17,20 +17,32 @@ //! 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}; use crate::common::Options; use crate::error::Error; 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, @@ -38,6 +50,7 @@ pub struct RESTEnv { api: Arc, options: Options, data_token_enabled: bool, + loaded_format_table_partition_options: Option, } impl std::fmt::Debug for RESTEnv { @@ -64,7 +77,55 @@ impl RESTEnv { api, options, data_token_enabled, + loaded_format_table_partition_options: None, + } + } + + 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. @@ -132,7 +193,6 @@ impl RESTEnv { ), source: None, })?; - let uuid = response.id.ok_or_else(|| Error::DataInvalid { message: format!( "Table {} response missing id (uuid)", @@ -151,7 +211,9 @@ impl RESTEnv { builder.build()? }; - let rest_env = RESTEnv::new(identifier.clone(), uuid, api, options, data_token_enabled); + let mut rest_env = RESTEnv::new(identifier.clone(), uuid, api, options, data_token_enabled); + rest_env.loaded_format_table_partition_options = + load_format_table_partition_options(identifier, &table_schema, is_external)?; Ok(Table::new( file_io, @@ -172,6 +234,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 00000000..a0da749c --- /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 6c8a0048..b5ffb724 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,41 @@ struct MockState { tables: HashMap, views: HashMap, functions: HashMap, + partitions: HashMap>, + partition_page_responses: HashMap>, + partition_list_call_counts: 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, } +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 +310,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 +757,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 +807,190 @@ 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}"); + 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 partitions = inner.partitions.get(&key).cloned().unwrap_or_default(); + ( + 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 +1047,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 +1149,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 +1201,84 @@ 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 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 +1396,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 99b952ce..fc879119 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 2000f9e4..c4643474 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() {