diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index 340cf07c2..06f34c3c4 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -693,11 +693,17 @@ impl SchemaProvider for PaimonSchemaProvider { let object = system_tables::parse_object_name_for_datafusion(name)?; if let Some(system_name) = object.system_table().map(str::to_string) { + let dynamic_options = self + .dynamic_options + .read() + .unwrap_or_else(|e| e.into_inner()) + .clone(); return await_with_runtime(system_tables::load( Arc::clone(&self.catalog), self.database.clone(), object, system_name, + dynamic_options, )) .await; } diff --git a/crates/integrations/datafusion/src/physical_plan/audit_log.rs b/crates/integrations/datafusion/src/physical_plan/audit_log.rs new file mode 100644 index 000000000..f98098ec8 --- /dev/null +++ b/crates/integrations/datafusion/src/physical_plan/audit_log.rs @@ -0,0 +1,240 @@ +// 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. + +//! Audit execution policy layered over the shared scan mechanics. + +use std::sync::Arc; + +use datafusion::common::{stats::Precision, Statistics}; +use datafusion::config::ConfigOptions; +use datafusion::error::Result as DFResult; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr::utils::collect_columns; +use datafusion::physical_plan::filter_pushdown::{ + ChildPushdownResult, FilterPushdownPhase, FilterPushdownPropagation, +}; +use datafusion::physical_plan::{DisplayAs, ExecutionPlan, PlanProperties}; +use paimon::table::AuditLogRead; + +use super::PaimonTableScan; + +/// Retains retract rows and keeps logical audit columns out of physical pushdown. +#[derive(Debug, Clone)] +pub(crate) struct PaimonAuditLogScan { + inner: PaimonTableScan, +} + +impl PaimonAuditLogScan { + pub(crate) fn new(inner: PaimonTableScan) -> Self { + Self { inner } + } +} + +impl ExecutionPlan for PaimonAuditLogScan { + fn name(&self) -> &str { + "PaimonAuditLogScan" + } + + fn properties(&self) -> &Arc { + self.inner.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DFResult> { + Ok(self) + } + + fn handle_child_pushdown_result( + &self, + _phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> DFResult>> { + let result = self + .inner + .pushdown_filters(child_pushdown_result, |filter| { + // Audit system-table names are case sensitive. Synthetic columns + // have no counterpart in the underlying data files. + collect_columns(filter).iter().all(|column| { + self.inner + .table() + .schema() + .fields() + .iter() + .any(|field| field.name() == column.name()) + }) + })?; + Ok(FilterPushdownPropagation { + filters: result.filters, + updated_node: result + .updated_node + .map(|scan| Arc::new(Self::new(scan)) as Arc), + }) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> DFResult { + self.inner.execute_with(partition, |read, splits| { + AuditLogRead::new(read)?.to_arrow(splits) + }) + } + + fn partition_statistics(&self, partition: Option) -> DFResult> { + let mut statistics = self.inner.partition_statistics(partition)?; + Arc::make_mut(&mut statistics).num_rows = Precision::Absent; + Ok(statistics) + } +} + +impl DisplayAs for PaimonAuditLogScan { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + self.inner.fmt_scan(self.name(), f) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::table::{datafusion_arrow_schema, PaimonScanBuilder}; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{lit, BinaryExpr, Column}; + use datafusion::physical_expr::PhysicalExpr; + use datafusion::physical_plan::filter_pushdown::{ChildFilterPushdownResult, PushedDown}; + use paimon::catalog::Identifier; + use paimon::table::Table; + use paimon::DataSplitBuilder; + + fn first_row_audit_scan() -> PaimonAuditLogScan { + let file_io = paimon::io::FileIOBuilder::new("memory").build().unwrap(); + let schema = paimon::spec::Schema::builder() + .column( + "id", + paimon::spec::DataType::Int(paimon::spec::IntType::new()), + ) + .primary_key(["id"]) + .option("bucket", "1") + .option("merge-engine", "first-row") + .build() + .unwrap(); + let table = Table::new( + file_io, + Identifier::new("default", "first_row_audit"), + "memory:/first-row-audit".to_string(), + paimon::spec::TableSchema::new(0, &schema), + None, + ); + let split = |snapshot| { + DataSplitBuilder::new() + .with_snapshot(snapshot) + .with_partition(paimon::spec::BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path("memory:/first-row-audit/bucket-0".to_string()) + .with_total_buckets(1) + .with_data_files(vec![]) + .build() + .unwrap() + }; + let read_fields = paimon::table::AuditLogTable::new(table.clone()) + .fields() + .unwrap(); + let arrow_schema = datafusion_arrow_schema(&read_fields, true).unwrap(); + let plan = PaimonScanBuilder { + table: &table, + schema: &arrow_schema, + plan: paimon::table::Plan::new(vec![split(1), split(2)]), + scan_trace: None, + projection: None, + pushed_predicate: None, + limit: None, + target_partitions: 8, + filter_exact: false, + case_sensitive: true, + } + .build_scan(read_fields) + .unwrap(); + PaimonAuditLogScan::new(plan) + } + + #[test] + fn test_first_row_audit_distributes_independent_splits() { + let scan = first_row_audit_scan(); + + assert_eq!(scan.inner.planned_partitions().len(), 2); + assert!(scan + .inner + .planned_partitions() + .iter() + .all(|splits| splits.len() == 1)); + } + + #[test] + fn test_audit_policy_survives_filter_pushdown() { + let scan = first_row_audit_scan(); + let filters: Vec> = vec![ + Arc::new(BinaryExpr::new( + Arc::new(Column::new("id", 1)), + Operator::Gt, + lit(1_i32), + )), + Arc::new(BinaryExpr::new( + Arc::new(Column::new("rowkind", 0)), + Operator::Eq, + lit("-D"), + )), + ]; + let result = scan + .handle_child_pushdown_result( + FilterPushdownPhase::Post, + ChildPushdownResult { + parent_filters: filters + .into_iter() + .map(|filter| ChildFilterPushdownResult { + filter, + child_results: Vec::new(), + }) + .collect(), + self_filters: Vec::new(), + }, + &ConfigOptions::default(), + ) + .unwrap(); + + assert!(matches!( + result.filters.as_slice(), + [PushedDown::Yes, PushedDown::No] + )); + let updated = result.updated_node.unwrap(); + assert!(updated.downcast_ref::().is_some()); + assert_eq!( + updated.partition_statistics(None).unwrap().num_rows, + Precision::Absent + ); + } +} diff --git a/crates/integrations/datafusion/src/physical_plan/mod.rs b/crates/integrations/datafusion/src/physical_plan/mod.rs index 2d1905035..e0e3c8878 100644 --- a/crates/integrations/datafusion/src/physical_plan/mod.rs +++ b/crates/integrations/datafusion/src/physical_plan/mod.rs @@ -15,10 +15,12 @@ // specific language governing permissions and limitations // under the License. +mod audit_log; pub(crate) mod scan; mod search_score; pub(crate) mod sink; +pub(crate) use audit_log::PaimonAuditLogScan; pub use scan::PaimonTableScan; pub(crate) use search_score::{SearchScoreExec, SearchScoreOutputColumn}; pub use sink::PaimonDataSink; diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs b/crates/integrations/datafusion/src/physical_plan/scan.rs index fe0a38590..52659126e 100644 --- a/crates/integrations/datafusion/src/physical_plan/scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/scan.rs @@ -51,7 +51,7 @@ use datafusion::physical_plan::{DisplayAs, ExecutionPlan, Partitioning, PlanProp use futures::{FutureExt, StreamExt, TryStreamExt}; use paimon::arrow::ParquetReadBudget; use paimon::spec::{DataField, Datum, MergeEngine, Predicate, PredicateBuilder, PredicateOperator}; -use paimon::table::{ScanTrace, Table}; +use paimon::table::{ArrowRecordBatchStream, ScanTrace, Table, TableRead}; use paimon::DataSplit; use crate::error::to_datafusion_error; @@ -969,34 +969,12 @@ impl PaimonTableScan { .map(|(accumulator, field)| accumulator.finish(field.data_type(), exact_null_counts)) .collect() } -} - -impl ExecutionPlan for PaimonTableScan { - fn name(&self) -> &str { - "PaimonTableScan" - } - - fn properties(&self) -> &Arc { - &self.plan_properties - } - - fn children(&self) -> Vec<&Arc> { - vec![] - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> DFResult> { - Ok(self) - } - fn handle_child_pushdown_result( + pub(crate) fn pushdown_filters( &self, - _phase: FilterPushdownPhase, child_pushdown_result: ChildPushdownResult, - _config: &ConfigOptions, - ) -> DFResult>> { + supported: impl Fn(&Arc) -> bool, + ) -> DFResult> { let filters = child_pushdown_result .parent_filters .into_iter() @@ -1007,13 +985,14 @@ impl ExecutionPlan for PaimonTableScan { Vec::new(), )); } - let schema = self.schema(); let mut accepted = Vec::new(); let parent_filter_handled = filters .into_iter() .map(|filter| { - if can_expr_be_pushed_down_with_schemas(&filter, schema.as_ref()) { + if supported(&filter) + && can_expr_be_pushed_down_with_schemas(&filter, schema.as_ref()) + { accepted.push(filter); // This scan evaluates accepted expressions exactly, so the // parent FilterExec can be removed. @@ -1048,14 +1027,16 @@ impl ExecutionPlan for PaimonTableScan { } Ok( FilterPushdownPropagation::with_parent_pushdown_result(parent_filter_handled) - .with_updated_node(Arc::new(scan)), + .with_updated_node(scan), ) } - fn execute( + pub(crate) fn execute_with( &self, partition: usize, - _context: Arc, + read_splits: impl FnOnce(TableRead<'_>, &[DataSplit]) -> paimon::Result + + Send + + 'static, ) -> DFResult { let splits = Arc::clone(self.planned_partitions.get(partition).ok_or_else(|| { datafusion::error::DataFusionError::Internal(format!( @@ -1098,12 +1079,11 @@ impl ExecutionPlan for PaimonTableScan { Arc::clone(&schema), ))); } - let stream = read.to_arrow(&splits).map_err(to_datafusion_error)?; + let stream = read_splits(read, &splits).map_err(to_datafusion_error)?; let batch_schema = Arc::clone(&schema); let stream = stream.map(move |result| { - let mut batch = result - .map_err(to_datafusion_error) - .and_then(|batch| to_datafusion_batch(batch, &batch_schema))?; + let batch = result.map_err(to_datafusion_error)?; + let mut batch = to_datafusion_batch(batch, &batch_schema)?; // The decoder hook is an optimization and may be unavailable // for a file/path. Retain every original live expression as // the exact fallback; evaluating it on decoder survivors is @@ -1136,6 +1116,95 @@ impl ExecutionPlan for PaimonTableScan { ))) } + pub(crate) fn fmt_scan(&self, name: &str, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "{}: table={}", name, self.table.identifier())?; + + let total_splits: usize = self.planned_partitions.iter().map(|p| p.len()).sum(); + let total_files: usize = self + .planned_partitions + .iter() + .flat_map(|p| p.iter()) + .map(|s| s.data_files().len()) + .sum(); + write!( + f, + ", partitions={}, splits={total_splits}, files={total_files}", + self.planned_partitions.len() + )?; + + let columns = self + .read_type + .iter() + .map(|field| field.name()) + .collect::>(); + write!(f, ", projection=[{}]", columns.join(", "))?; + if let Some(ref predicate) = self.pushed_predicate { + write!(f, ", predicate={predicate}")?; + } + if let Some(limit) = self.limit { + write!(f, ", limit={limit}")?; + } + if let Some(ref trace) = self.scan_trace { + write!(f, ", trace={trace}")?; + } + if let Some(ref pushed_variants) = self.pushed_variants { + write!(f, ", PushedVariants=[{pushed_variants}]")?; + } + if !self.runtime_filters.is_empty() { + let filters = self + .runtime_filters + .iter() + .map(ToString::to_string) + .collect::>(); + write!(f, ", runtime_filters=[{}]", filters.join(" AND "))?; + } + Ok(()) + } +} + +impl ExecutionPlan for PaimonTableScan { + fn name(&self) -> &str { + "PaimonTableScan" + } + + fn properties(&self) -> &Arc { + &self.plan_properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DFResult> { + Ok(self) + } + + fn handle_child_pushdown_result( + &self, + _phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> DFResult>> { + let result = self.pushdown_filters(child_pushdown_result, |_| true)?; + Ok(FilterPushdownPropagation { + filters: result.filters, + updated_node: result + .updated_node + .map(|scan| Arc::new(scan) as Arc), + }) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> DFResult { + self.execute_with(partition, |read, splits| read.to_arrow(splits)) + } + fn partition_statistics(&self, partition: Option) -> DFResult> { let partitions: &[Arc<[DataSplit]>] = match partition { Some(idx) => std::slice::from_ref(&self.planned_partitions[idx]), @@ -1183,48 +1252,7 @@ impl DisplayAs for PaimonTableScan { _t: datafusion::physical_plan::DisplayFormatType, f: &mut std::fmt::Formatter, ) -> std::fmt::Result { - write!(f, "PaimonTableScan: table={}", self.table.identifier())?; - - let total_splits: usize = self.planned_partitions.iter().map(|p| p.len()).sum(); - let total_files: usize = self - .planned_partitions - .iter() - .flat_map(|p| p.iter()) - .map(|s| s.data_files().len()) - .sum(); - write!( - f, - ", partitions={}, splits={total_splits}, files={total_files}", - self.planned_partitions.len() - )?; - - let columns = self - .read_type - .iter() - .map(|field| field.name()) - .collect::>(); - write!(f, ", projection=[{}]", columns.join(", "))?; - if let Some(ref predicate) = self.pushed_predicate { - write!(f, ", predicate={predicate}")?; - } - if let Some(limit) = self.limit { - write!(f, ", limit={limit}")?; - } - if let Some(ref trace) = self.scan_trace { - write!(f, ", trace={trace}")?; - } - if let Some(ref pushed_variants) = self.pushed_variants { - write!(f, ", PushedVariants=[{pushed_variants}]")?; - } - if !self.runtime_filters.is_empty() { - let filters = self - .runtime_filters - .iter() - .map(ToString::to_string) - .collect::>(); - write!(f, ", runtime_filters=[{}]", filters.join(" AND "))?; - } - Ok(()) + self.fmt_scan(self.name(), f) } } diff --git a/crates/integrations/datafusion/src/system_tables/audit_log.rs b/crates/integrations/datafusion/src/system_tables/audit_log.rs new file mode 100644 index 000000000..b464afd4a --- /dev/null +++ b/crates/integrations/datafusion/src/system_tables/audit_log.rs @@ -0,0 +1,132 @@ +// 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. + +//! Mirrors Java [AuditLogTable](https://github.com/apache/paimon/blob/master/paimon-core/src/main/java/org/apache/paimon/table/system/AuditLogTable.java). + +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::catalog::Session; +use datafusion::datasource::{TableProvider, TableType}; +use datafusion::error::Result as DFResult; +use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; +use datafusion::physical_plan::ExecutionPlan; +use paimon::spec::DataField; +use paimon::table::{AuditLogTable as PaimonAuditLogTable, Table}; + +use crate::error::to_datafusion_error; +use crate::filter_pushdown::{analyze_filters, classify_filter_pushdown}; +use crate::physical_plan::PaimonAuditLogScan; +use crate::runtime::await_with_runtime; +use crate::table::{datafusion_arrow_schema, PaimonScanBuilder}; + +pub(super) fn build(table: Table) -> DFResult> { + let fields = PaimonAuditLogTable::new(table.clone()) + .fields() + .map_err(to_datafusion_error)?; + let schema = datafusion_arrow_schema(&fields, true)?; + Ok(Arc::new(AuditLogTable { + table, + fields, + schema, + })) +} + +#[derive(Debug)] +struct AuditLogTable { + table: Table, + fields: Vec, + schema: SchemaRef, +} + +#[async_trait] +impl TableProvider for AuditLogTable { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn table_type(&self) -> TableType { + TableType::View + } + + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> DFResult> { + let filter_analysis = analyze_filters(filters, self.table.schema().fields(), true); + let pushed_limit = limit.filter(|_| !filter_analysis.requires_residual); + let mut read_builder = self.table.new_read_builder(); + if let Some(indices) = projection { + read_builder.with_read_type( + indices + .iter() + .map(|&index| self.fields[index].clone()) + .filter(|field| { + !matches!( + field.id(), + paimon::spec::ROW_KIND_FIELD_ID + | paimon::spec::SEQUENCE_NUMBER_FIELD_ID + ) + }) + .collect(), + ); + } + if let Some(predicate) = filter_analysis.pushed_predicate.clone() { + read_builder.with_filter(predicate); + } + if let Some(limit) = pushed_limit { + read_builder.with_limit(limit); + } + let (plan, trace) = await_with_runtime(read_builder.new_audit_scan().plan_with_trace()) + .await + .map_err(to_datafusion_error)?; + + let scan = PaimonScanBuilder { + table: &self.table, + schema: &self.schema, + plan, + scan_trace: Some(trace), + projection, + pushed_predicate: filter_analysis.pushed_predicate, + limit: pushed_limit, + target_partitions: state.config_options().execution.target_partitions, + filter_exact: false, + case_sensitive: true, + } + .build_scan(self.fields.clone())?; + Ok(Arc::new(PaimonAuditLogScan::new(scan))) + } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> DFResult> { + let read_builder = self.table.new_read_builder(); + Ok(filters + .iter() + .map(|filter| { + classify_filter_pushdown(filter, self.table.schema().fields(), true, |predicate| { + read_builder.is_exact_filter_pushdown(predicate) + }) + }) + .collect()) + } +} diff --git a/crates/integrations/datafusion/src/system_tables/mod.rs b/crates/integrations/datafusion/src/system_tables/mod.rs index 3b3a0b098..5c23565d6 100644 --- a/crates/integrations/datafusion/src/system_tables/mod.rs +++ b/crates/integrations/datafusion/src/system_tables/mod.rs @@ -20,6 +20,7 @@ //! Mirrors Java [SystemTableLoader](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/table/system/SystemTableLoader.java): //! `TABLES` maps each system-table name to its builder function. +use std::collections::HashMap; use std::sync::Arc; use datafusion::datasource::TableProvider; @@ -29,6 +30,7 @@ use paimon::table::Table; use crate::error::to_datafusion_error; +mod audit_log; mod branches; mod consumers; mod files; @@ -49,6 +51,7 @@ type Builder = fn(Table) -> DFResult>; // in `load` because it needs the catalog handle (for metastore-tracked audit // metadata via `Catalog::list_partitions`). const TABLES: &[(&str, Builder)] = &[ + ("audit_log", audit_log::build), ("branches", branches::build), ("consumers", consumers::build), ("files", files::build), @@ -63,6 +66,7 @@ const TABLES: &[(&str, Builder)] = &[ ]; const SYSTEM_TABLE_NAMES: &[&str] = &[ + "audit_log", "branches", "consumers", "files", @@ -77,6 +81,21 @@ const SYSTEM_TABLE_NAMES: &[&str] = &[ "tags", ]; +// Reject system tables whose contents can expose protected table data or +// persisted credentials until Rust can apply row filters and column masks. +const QUERY_AUTH_UNSUPPORTED_TABLES: &[&str] = &[ + "audit_log", + "files", + "file_key_ranges", + "binlog", + "statistics", + "options", + "schemas", + "partitions", + "manifests", + "table_indexes", +]; + /// Parse a Paimon object name into table, branch, and optional system table. /// /// Mirrors Java [Identifier.splitObjectName](https://github.com/apache/paimon/blob/release-1.3/paimon-api/src/main/java/org/apache/paimon/catalog/Identifier.java). @@ -104,6 +123,21 @@ fn wrap_to_system_table(name: &str, base_table: Table) -> Option, + name: &str, +) -> DFResult<()> { + if QUERY_AUTH_UNSUPPORTED_TABLES + .iter() + .any(|candidate| name.eq_ignore_ascii_case(candidate)) + { + paimon::spec::CoreOptions::new(options) + .ensure_read_authorized() + .map_err(to_datafusion_error)?; + } + Ok(()) +} + pub(crate) fn provider_for_table( catalog: Arc, identifier: Identifier, @@ -113,10 +147,8 @@ pub(crate) fn provider_for_table( if !is_registered(system_name) { return Ok(None); } - // Fail closed: system tables expose file metadata the client can't authorize. - paimon::spec::CoreOptions::new(table.schema().options()) - .ensure_read_authorized() - .map_err(to_datafusion_error)?; + crate::table_loader::ensure_paimon_served(&table, &identifier)?; + ensure_system_table_read_supported(table.schema().options(), system_name)?; if system_name.eq_ignore_ascii_case("partitions") { return partitions::build(catalog, identifier, table).map(Some); } @@ -136,13 +168,25 @@ pub(crate) async fn load( database: String, object: ParsedObjectName, system_name: String, + dynamic_options: HashMap, ) -> DFResult>> { if !is_registered(&system_name) { return Ok(None); } + if system_name.eq_ignore_ascii_case("audit_log") + && paimon::spec::CoreOptions::new(&dynamic_options).table_read_sequence_number_enabled() + { + return Err(DataFusionError::Plan( + "table-read.sequence-number.enabled is not supported by dynamic options for $audit_log" + .to_string(), + )); + } + ensure_system_table_read_supported(&dynamic_options, &system_name)?; let identifier = Identifier::new(database, object.table().to_string()); match catalog.get_table(&identifier).await { Ok(mut table) => { + crate::table_loader::ensure_paimon_served(&table, &identifier)?; + ensure_system_table_read_supported(table.schema().options(), &system_name)?; if let Some(branch) = object.branch() { if !system_name.eq_ignore_ascii_case("branches") { table = table @@ -151,6 +195,12 @@ pub(crate) async fn load( .map_err(to_datafusion_error)?; } } + if system_name.eq_ignore_ascii_case("audit_log") && !dynamic_options.is_empty() { + table = table + .copy_with_time_travel(dynamic_options) + .await + .map_err(to_datafusion_error)?; + } provider_for_table(catalog, identifier, table, &system_name) } Err(paimon::Error::TableNotExist { .. }) => Err(DataFusionError::Plan(format!( @@ -189,6 +239,9 @@ mod tests { #[test] fn is_registered_is_case_insensitive() { + assert!(is_registered("audit_log")); + assert!(is_registered("Audit_Log")); + assert!(is_registered("AUDIT_LOG")); assert!(is_registered("options")); assert!(is_registered("Options")); assert!(is_registered("OPTIONS")); diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 9683bf588..f5e3e5a74 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -356,20 +356,23 @@ impl PaimonScanBuilder<'_> { self, read_fields: Vec, ) -> DFResult> { + Ok(Arc::new(self.build_scan(read_fields)?)) + } + + pub(crate) fn build_scan(self, read_fields: Vec) -> DFResult { let (projected_schema, read_type) = if let Some(indices) = self.projection { let fields: Vec = indices .iter() - .map(|&i| self.schema.field(i).clone()) + .map(|&index| self.schema.field(index).clone()) .collect(); let read_type = indices .iter() - .map(|&i| read_fields[i].clone()) - .collect::>(); + .map(|&index| read_fields[index].clone()) + .collect(); (Arc::new(Schema::new(fields)), read_type) } else { (self.schema.clone(), read_fields) }; - let splits = self.plan.into_splits(); let planned_partitions: Vec> = if splits.is_empty() { vec![Arc::from(Vec::new())] @@ -381,7 +384,7 @@ impl PaimonScanBuilder<'_> { .collect() }; - Ok(Arc::new(PaimonTableScan::try_new( + PaimonTableScan::try_new( projected_schema, self.table.clone(), read_type, @@ -392,7 +395,7 @@ impl PaimonScanBuilder<'_> { self.scan_trace, None, self.case_sensitive, - )?)) + ) } } diff --git a/crates/integrations/datafusion/tests/system_tables.rs b/crates/integrations/datafusion/tests/system_tables.rs index 1a64ccb77..dae555f74 100644 --- a/crates/integrations/datafusion/tests/system_tables.rs +++ b/crates/integrations/datafusion/tests/system_tables.rs @@ -22,7 +22,8 @@ mod common; use std::sync::Arc; use datafusion::arrow::array::{ - Array, BooleanArray, Int32Array, Int64Array, ListArray, StringArray, TimestampMillisecondArray, + Array, BooleanArray, Int32Array, Int64Array, Int8Array, ListArray, StringArray, + TimestampMillisecondArray, }; use datafusion::arrow::datatypes::{DataType, Field, TimeUnit}; use datafusion::arrow::record_batch::RecordBatch; @@ -30,6 +31,8 @@ use paimon::catalog::Identifier; use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options}; use paimon_datafusion::SQLContext; +use common::string_value; + const FIXTURE_TABLE: &str = "test_tantivy_fulltext"; fn extract_test_warehouse() -> (tempfile::TempDir, String) { @@ -85,17 +88,26 @@ async fn query_error(ctx: &SQLContext, sql: &str) -> String { } #[tokio::test] -async fn test_query_auth_table_fails_closed() { +async fn test_query_auth_system_tables_fail_closed() { let (ctx, _catalog, _tmp) = create_context().await; run_sql( &ctx, - "CREATE TABLE paimon.default.qa (id INT) WITH ('query-auth.enabled' = 'true')", + "CREATE TABLE paimon.default.qa (id INT) WITH ( + 'query-auth.enabled' = 'true', + 's3.secret-key' = 'persisted-secret' + )", ) .await; - // Data reads and data-derived system tables must all fail closed. + // Rust cannot yet apply query-auth filters or masks to table and system-table + // data. Metadata paths, including persisted options, must fail closed. for sql in [ "SELECT * FROM paimon.default.qa", + "SELECT * FROM paimon.default.qa$audit_log", + "SELECT * FROM paimon.default.qa$files", + "SELECT value FROM paimon.default.qa$options WHERE key = 's3.secret-key'", + "SELECT * FROM paimon.default.qa$schemas", + "SELECT * FROM paimon.default.qa$partitions", "SELECT * FROM paimon.default.qa$manifests", "SELECT * FROM paimon.default.qa$table_indexes", ] { @@ -105,6 +117,308 @@ async fn test_query_auth_table_fails_closed() { "`{sql}` should fail closed, got: {err}" ); } + + run_sql(&ctx, "CREATE TABLE paimon.default.qa_dynamic (id INT)").await; + run_sql(&ctx, "SET 'paimon.query-auth.enabled' = 'true'").await; + for sql in [ + "SELECT * FROM paimon.default.qa_dynamic", + "SELECT * FROM paimon.default.qa_dynamic$audit_log", + "SELECT * FROM paimon.default.qa_dynamic$files", + "SELECT * FROM paimon.default.qa_dynamic$options", + "SELECT * FROM paimon.default.qa_dynamic$schemas", + "SELECT * FROM paimon.default.qa_dynamic$partitions", + "SELECT * FROM paimon.default.qa_dynamic$manifests", + "SELECT * FROM paimon.default.qa_dynamic$table_indexes", + ] { + let err = query_error(&ctx, sql).await; + assert!( + err.contains("query-auth.enabled"), + "dynamic auth should make `{sql}` fail closed, got: {err}" + ); + } + run_sql(&ctx, "RESET 'paimon.query-auth.enabled'").await; + + run_sql(&ctx, "SET 'paimon.s3.secret-key' = 'session-secret'").await; + let batches = run_sql( + &ctx, + "SELECT value FROM paimon.default.qa_dynamic$options \ + WHERE key = 's3.secret-key'", + ) + .await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + run_sql(&ctx, "RESET 'paimon.s3.secret-key'").await; +} + +#[tokio::test] +async fn test_audit_log_rejects_dynamic_sequence_number_option() { + let (ctx, _catalog, _tmp) = create_context().await; + run_sql( + &ctx, + "CREATE TABLE paimon.default.audit_dynamic_sequence ( + id INT NOT NULL, + PRIMARY KEY (id) + ) WITH ('bucket' = '1')", + ) + .await; + + run_sql( + &ctx, + "SET 'paimon.table-read.sequence-number.enabled' = 'true'", + ) + .await; + let err = query_error( + &ctx, + "SELECT * FROM paimon.default.audit_dynamic_sequence$audit_log", + ) + .await; + assert!( + err.contains("table-read.sequence-number.enabled") + && err.contains("not supported by dynamic options"), + "unexpected error: {err}" + ); + run_sql(&ctx, "RESET 'paimon.table-read.sequence-number.enabled'").await; +} + +#[tokio::test] +async fn test_audit_log_respects_dynamic_time_travel() { + let (ctx, _catalog, _tmp) = create_context().await; + run_sql(&ctx, "CREATE TABLE paimon.default.audit_tt (id INT)").await; + run_sql(&ctx, "INSERT INTO paimon.default.audit_tt VALUES (1)").await; + run_sql(&ctx, "INSERT INTO paimon.default.audit_tt VALUES (2)").await; + + run_sql(&ctx, "SET 'paimon.scan.version' = '1'").await; + let batches = run_sql( + &ctx, + "SELECT COUNT(*) FROM paimon.default.audit_tt$audit_log", + ) + .await; + assert_eq!( + batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 1 + ); + run_sql(&ctx, "RESET 'paimon.scan.version'").await; +} + +#[tokio::test] +async fn test_audit_log_system_table_keeps_row_kinds_and_sequence_numbers() { + let (ctx, catalog, _tmp) = create_context().await; + run_sql( + &ctx, + "CREATE TABLE paimon.default.audit_rows ( + id INT NOT NULL, + value INT, + PRIMARY KEY (id) + ) WITH ( + 'bucket' = '1', + 'merge-engine' = 'deduplicate', + 'changelog-producer' = 'input', + 'table-read.sequence-number.enabled' = 'true' + )", + ) + .await; + + let table = catalog + .get_table(&Identifier::new("default", "audit_rows")) + .await + .unwrap(); + let batch = RecordBatch::try_new( + Arc::new(datafusion::arrow::datatypes::Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Int32, true), + Field::new("_VALUE_KIND", DataType::Int8, false), + ])), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 1, 2])), + Arc::new(Int32Array::from(vec![10, 20, 10, 25])), + Arc::new(Int8Array::from(vec![0, 0, 3, 2])), + ], + ) + .unwrap(); + let builder = table.new_write_builder(); + let mut write = builder.new_write().unwrap(); + write.write_arrow_batch(&batch).await.unwrap(); + let messages = write.prepare_commit().await.unwrap(); + builder.new_commit().commit(messages).await.unwrap(); + + let batches = run_sql( + &ctx, + "SELECT \"_SEQUENCE_NUMBER\", rowkind, id, value + FROM paimon.default.audit_rows$audit_log + WHERE rowkind = '-D' OR id = 2 + ORDER BY id", + ) + .await; + let mut rows = Vec::new(); + for batch in &batches { + let sequence = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let rowkind = batch.column(1); + let id = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let value = batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + rows.push(( + sequence.value(row), + string_value(rowkind.as_ref(), row).to_string(), + id.value(row), + value.value(row), + )); + } + } + assert_eq!( + rows, + vec![(2, "-D".to_string(), 1, 10), (3, "+U".to_string(), 2, 25)] + ); + + let batches = run_sql( + &ctx, + "SELECT id FROM paimon.default.audit_rows$audit_log WHERE value = 20", + ) + .await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + + let batches = run_sql( + &ctx, + "SELECT COUNT(*) FROM paimon.default.audit_rows$audit_log", + ) + .await; + assert_eq!( + batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 2 + ); + + run_sql(&ctx, "CREATE TABLE paimon.default.append_rows (id INT)").await; + run_sql( + &ctx, + "INSERT INTO paimon.default.append_rows VALUES (1), (2)", + ) + .await; + let batches = run_sql( + &ctx, + "SELECT rowkind FROM paimon.default.append_rows$audit_log", + ) + .await; + assert!(batches.iter().all(|batch| { + (0..batch.num_rows()).all(|row| string_value(batch.column(0).as_ref(), row) == "+I") + })); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + + let explain = run_sql( + &ctx, + "EXPLAIN SELECT id FROM paimon.default.audit_rows$audit_log WHERE id = 2", + ) + .await; + assert!(explain.iter().any(|batch| { + (0..batch.num_rows()).any(|row| { + let plan = string_value(batch.column(1).as_ref(), row); + plan.contains("PaimonAuditLogScan") && plan.contains("predicate=") + }) + })); +} + +#[tokio::test] +async fn test_first_row_audit_log_merges_level_zero_before_filtering() { + let (ctx, _catalog, _tmp) = create_context().await; + run_sql( + &ctx, + "CREATE TABLE paimon.default.first_row_audit ( + id INT NOT NULL, + value INT, + PRIMARY KEY (id) + ) WITH ( + 'bucket' = '1', + 'merge-engine' = 'first-row' + )", + ) + .await; + run_sql( + &ctx, + "INSERT INTO paimon.default.first_row_audit VALUES (1, 10)", + ) + .await; + run_sql( + &ctx, + "INSERT INTO paimon.default.first_row_audit VALUES (1, 20)", + ) + .await; + + let batches = run_sql( + &ctx, + "SELECT value FROM paimon.default.first_row_audit$audit_log", + ) + .await; + assert_eq!( + batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[10] + ); + + let batches = run_sql( + &ctx, + "SELECT value FROM paimon.default.first_row_audit$audit_log WHERE value = 20", + ) + .await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); +} + +#[tokio::test] +async fn test_audit_log_system_table_matches_deletion_vector_visibility() { + let (ctx, _catalog, _tmp) = create_context().await; + run_sql( + &ctx, + "CREATE TABLE paimon.default.dv_audit (id INT NOT NULL) WITH ( + 'row-tracking.enabled' = 'true', + 'data-evolution.enabled' = 'true', + 'deletion-vectors.enabled' = 'true' + )", + ) + .await; + run_sql( + &ctx, + "INSERT INTO paimon.default.dv_audit (id) VALUES (1), (2)", + ) + .await; + run_sql(&ctx, "DELETE FROM paimon.default.dv_audit WHERE id = 1").await; + + let batches = run_sql( + &ctx, + "SELECT rowkind, id FROM paimon.default.dv_audit$audit_log", + ) + .await; + assert_eq!(string_value(batches[0].column(0).as_ref(), 0), "+I"); + assert_eq!( + batches[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[2] + ); } #[tokio::test] diff --git a/crates/paimon/src/table/audit_log_table.rs b/crates/paimon/src/table/audit_log_table.rs index a6b6e5fe0..a77c7febb 100644 --- a/crates/paimon/src/table/audit_log_table.rs +++ b/crates/paimon/src/table/audit_log_table.rs @@ -16,7 +16,7 @@ // under the License. use super::incremental_scan::{IncrementalPlan, IncrementalScan, IncrementalScanMode}; -use super::{ArrowRecordBatchStream, Table}; +use super::{ArrowRecordBatchStream, AuditLogRead, DataSplit, Table, TableScan}; use crate::spec::{ BigIntType, DataField, DataType, VarCharType, ROW_KIND_FIELD_ID, ROW_KIND_FIELD_NAME, SEQUENCE_NUMBER_FIELD_ID, SEQUENCE_NUMBER_FIELD_NAME, @@ -33,8 +33,6 @@ pub struct AuditLogTable { wrapped: Table, } -const TABLE_READ_SEQUENCE_NUMBER_ENABLED: &str = "table-read.sequence-number.enabled"; - impl AuditLogTable { pub fn new(wrapped: Table) -> Self { Self { wrapped } @@ -66,9 +64,8 @@ impl AuditLogTable { fn sequence_number_enabled(&self) -> bool { self.wrapped .schema() - .options() - .get(TABLE_READ_SEQUENCE_NUMBER_ENABLED) - .is_some_and(|v| v.eq_ignore_ascii_case("true")) + .core_options() + .table_read_sequence_number_enabled() } pub fn new_incremental_scan( @@ -80,9 +77,26 @@ impl AuditLogTable { IncrementalScan::for_table(&self.wrapped, mode, start_exclusive, end_inclusive) } + /// Plan a current-state audit read for [`Self::to_arrow_for_splits`]. + pub fn new_scan(&self) -> TableScan<'_> { + self.wrapped.new_read_builder().new_audit_scan() + } + + /// Creates an audit reader using the table's configured fields and options. + pub fn new_read(&self) -> crate::Result> { + AuditLogRead::new(self.wrapped.new_read_builder().new_read()?) + } + pub fn to_arrow(&self, plan: &IncrementalPlan) -> crate::Result { plan.validate()?; - let read = self.wrapped.new_read_builder().new_read()?; - read.to_audit_log_arrow(plan) + self.new_read()?.to_arrow(plan) + } + + /// Reads the current table state, retaining retract rows for primary-key tables. + pub fn to_arrow_for_splits( + &self, + splits: &[DataSplit], + ) -> crate::Result { + self.new_read()?.to_arrow(splits) } } diff --git a/crates/paimon/src/table/kv_file_reader.rs b/crates/paimon/src/table/kv_file_reader.rs index 88b1bf806..6ff644adc 100644 --- a/crates/paimon/src/table/kv_file_reader.rs +++ b/crates/paimon/src/table/kv_file_reader.rs @@ -27,14 +27,14 @@ use super::data_file_reader::DataFileReader; use super::sort_merge::{ - AggregateMergeFunction, DeduplicateMergeFunction, PartialUpdateMergeFunction, - SortMergeReaderBuilder, + AggregateMergeFunction, ConfiguredDeduplicateMergeFunction, DeduplicateMergeFunction, + FirstRowMergeFunction, PartialUpdateMergeFunction, SortMergeReaderBuilder, }; use crate::arrow::{build_target_arrow_schema, ParquetReadBudget}; use crate::deletion_vector::DeletionVectorFactory; use crate::io::FileIO; use crate::spec::{ - BigIntType, DataField, DataFileMeta, DataType as PaimonDataType, MergeEngine, + BigIntType, CoreOptions, DataField, DataFileMeta, DataType as PaimonDataType, MergeEngine, PartialUpdateConfig, Predicate, TinyIntType, SEQUENCE_NUMBER_FIELD_ID, SEQUENCE_NUMBER_FIELD_NAME, VALUE_KIND_FIELD_ID, VALUE_KIND_FIELD_NAME, }; @@ -49,6 +49,7 @@ use std::collections::HashMap; use std::sync::Arc; /// Reads primary-key table data files using sort-merge deduplication. +#[derive(Clone)] pub(crate) struct KeyValueFileReader { file_io: FileIO, config: KeyValueReadConfig, @@ -63,6 +64,7 @@ pub(crate) struct KeyValueFileReader { /// Configuration for [`KeyValueFileReader`], grouping table schema and /// key/predicate parameters. +#[derive(Clone)] pub(crate) struct KeyValueReadConfig { pub table_name: String, pub table_options: HashMap, @@ -75,6 +77,8 @@ pub(crate) struct KeyValueReadConfig { pub merge_engine: MergeEngine, pub sequence_fields: Vec, pub read_batch_size: usize, + /// Keep a winning retract row instead of dropping it after key merge. + pub keep_delete: bool, /// Merge files from all supplied splits into one globally key-sorted stream. pub merge_splits: bool, /// Optional cap on sorted-run inputs merged concurrently by one LoserTree. @@ -282,6 +286,7 @@ impl KeyValueFileReader { self } + #[allow(clippy::too_many_arguments)] fn new_merge_function( merge_engine: MergeEngine, table_options: &HashMap, @@ -290,21 +295,28 @@ impl KeyValueFileReader { merge_output_fields: &[DataField], primary_keys: &[String], sequence_fields: &[String], + keep_delete: bool, ) -> crate::Result> { match merge_engine { + MergeEngine::Deduplicate + if keep_delete || CoreOptions::new(table_options).ignore_delete() => + { + Ok(Box::new(ConfiguredDeduplicateMergeFunction::new( + table_options, + keep_delete, + ))) + } MergeEngine::Deduplicate => Ok(Box::new(DeduplicateMergeFunction)), - MergeEngine::PartialUpdate => Ok(Box::new( - PartialUpdateMergeFunction::new_with_schema( + MergeEngine::PartialUpdate => { + Ok(Box::new(PartialUpdateMergeFunction::new_with_schema( table_options, table_name, table_fields, merge_output_fields, primary_keys, - )?, - )), - MergeEngine::FirstRow => Err(Error::Unsupported { - message: "KeyValueFileReader does not support merge-engine=first-row; first-row reads should use the non-KV path".to_string(), - }), + )?)) + } + MergeEngine::FirstRow => Ok(Box::new(FirstRowMergeFunction::new(table_options))), MergeEngine::Aggregation => Ok(Box::new(AggregateMergeFunction::new( table_options, table_name, @@ -370,11 +382,29 @@ impl KeyValueFileReader { .collect(), )) }; + let expose_sequence = self + .config + .read_type + .iter() + .any(|field| field.id() == SEQUENCE_NUMBER_FIELD_ID); + let expose_value_kind = self + .config + .read_type + .iter() + .any(|field| field.id() == VALUE_KIND_FIELD_ID); + // User columns = read_type fields + any key fields not already in read_type - // + any sequence fields not already included. + // + any sequence fields not already included. Physical system + // fields are already the first two columns of every KV file. let read_type_names: std::collections::HashSet<&str> = self.config.read_type.iter().map(|f| f.name()).collect(); - let mut user_fields: Vec = self.config.read_type.clone(); + let mut user_fields: Vec = self + .config + .read_type + .iter() + .filter(|field| !matches!(field.id(), SEQUENCE_NUMBER_FIELD_ID | VALUE_KIND_FIELD_ID)) + .cloned() + .collect(); for kf in &key_fields { if !read_type_names.contains(kf.name()) { user_fields.push(kf.clone()); @@ -423,8 +453,8 @@ impl KeyValueFileReader { // Internal read type: [_SEQ, _VK, user_fields...] let mut internal_read_type: Vec = Vec::new(); - internal_read_type.push(seq_field); - internal_read_type.push(value_kind_field); + internal_read_type.push(seq_field.clone()); + internal_read_type.push(value_kind_field.clone()); internal_read_type.extend(user_fields.clone()); let internal_schema = build_target_arrow_schema(&internal_read_type)?; @@ -447,17 +477,29 @@ impl KeyValueFileReader { .unwrap() }) .collect(); - let value_fields: Vec = user_fields - .iter() - .filter(|f| !key_names.contains(f.name())) - .cloned() - .collect(); - let value_indices: Vec = user_fields - .iter() - .enumerate() - .filter(|(_, f)| !key_names.contains(f.name())) - .map(|(i, _)| i + 2) - .collect(); + let mut value_fields = Vec::new(); + let mut value_indices = Vec::new(); + if expose_sequence { + value_fields.push(seq_field); + value_indices.push(seq_index); + } + if expose_value_kind { + value_fields.push(value_kind_field); + value_indices.push(value_kind_index); + } + value_fields.extend( + user_fields + .iter() + .filter(|field| !key_names.contains(field.name())) + .cloned(), + ); + value_indices.extend( + user_fields + .iter() + .enumerate() + .filter(|(_, field)| !key_names.contains(field.name())) + .map(|(index, _)| index + 2), + ); // If sequence.field is configured, find each field's index in the internal schema. let user_sequence_indices: Vec = self @@ -517,6 +559,7 @@ impl KeyValueFileReader { let primary_keys = self.config.primary_keys; let sequence_fields = self.config.sequence_fields; let read_batch_size = self.config.read_batch_size; + let keep_delete = self.config.keep_delete; let max_merge_input_streams = self.config.max_merge_input_streams; let parquet_read_budget = self.config.parquet_read_budget; #[cfg(test)] @@ -654,6 +697,7 @@ impl KeyValueFileReader { &merge_output_fields, &primary_keys, &sequence_fields, + keep_delete, )?, ) .build()?; @@ -1314,6 +1358,7 @@ mod tests { merge_engine: core_options.merge_engine().unwrap(), sequence_fields: Vec::new(), read_batch_size: core_options.read_batch_size().unwrap(), + keep_delete: false, merge_splits: true, max_merge_input_streams: None, parquet_read_budget: Some(budget), @@ -1433,6 +1478,7 @@ mod tests { merge_engine: core_options.merge_engine().unwrap(), sequence_fields: Vec::new(), read_batch_size: core_options.read_batch_size().unwrap(), + keep_delete: false, merge_splits: true, max_merge_input_streams: Some(256), parquet_read_budget: None, @@ -1640,6 +1686,7 @@ mod tests { .map(|field| field.to_string()) .collect(), read_batch_size: core_options.read_batch_size().unwrap(), + keep_delete: false, merge_splits: false, max_merge_input_streams: None, parquet_read_budget: None, @@ -1711,6 +1758,7 @@ mod tests { merge_engine: core_options.merge_engine().unwrap(), sequence_fields: Vec::new(), read_batch_size: core_options.read_batch_size().unwrap(), + keep_delete: false, merge_splits: false, max_merge_input_streams: None, parquet_read_budget: Some(Arc::new(ParquetReadBudget::new(2, 256 << 20).unwrap())), @@ -1905,6 +1953,7 @@ mod tests { merge_engine: core_options.merge_engine().unwrap(), sequence_fields: Vec::new(), read_batch_size: core_options.read_batch_size().unwrap(), + keep_delete: false, merge_splits, max_merge_input_streams: None, parquet_read_budget: None, @@ -1972,6 +2021,7 @@ mod tests { merge_engine: core_options.merge_engine().unwrap(), sequence_fields: Vec::new(), read_batch_size: core_options.read_batch_size().unwrap(), + keep_delete: false, merge_splits: true, max_merge_input_streams: Some(256), parquet_read_budget: None, diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 32484924a..79a490374 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -152,7 +152,7 @@ pub use source::{ merge_row_ranges, DataSplit, DataSplitBuilder, DeletionFile, PartitionBucket, Plan, RowRange, }; pub use table_commit::TableCommit; -pub use table_read::TableRead; +pub use table_read::{AuditLogInput, AuditLogRead, TableRead}; pub use table_scan::TableScan; pub use table_update::TableUpdate; pub use table_write::TableWrite; diff --git a/crates/paimon/src/table/read_builder.rs b/crates/paimon/src/table/read_builder.rs index ec8ef966e..52dc1cea9 100644 --- a/crates/paimon/src/table/read_builder.rs +++ b/crates/paimon/src/table/read_builder.rs @@ -505,10 +505,9 @@ impl<'a> PaimonReadBuilder<'a> { // `to_arrow` (e.g. an empty-splits fast path) can't bypass the guard. let core_options = self.table.schema.core_options(); core_options.ensure_read_authorized()?; - let read_type = match self.resolve_read_type()? { - None => self.table.schema.fields().to_vec(), - Some(fields) => fields, - }; + let projection = self.resolve_read_type()?; + let explicit_projection = projection.is_some(); + let read_type = projection.unwrap_or_else(|| self.table.schema.fields().to_vec()); // Pass the FULL data predicate through (including `And`/`Or`/`Not`). // Pushdown/stats skip compound nodes; the residual pass enforces the full @@ -519,6 +518,7 @@ impl<'a> PaimonReadBuilder<'a> { }; Ok( TableRead::new(self.table, read_type, self.filter.data_predicates.clone()) + .with_explicit_projection(explicit_projection) .with_parquet_read_budget(parquet_read_budget), ) } diff --git a/crates/paimon/src/table/sort_merge.rs b/crates/paimon/src/table/sort_merge.rs index a26197009..81f37c3ea 100644 --- a/crates/paimon/src/table/sort_merge.rs +++ b/crates/paimon/src/table/sort_merge.rs @@ -40,7 +40,7 @@ use futures::StreamExt; use std::cmp::Ordering; use std::collections::HashMap; use std::collections::HashSet; -use std::sync::Mutex; +use std::sync::{Arc, Mutex, OnceLock}; // --------------------------------------------------------------------------- // MergeFunction @@ -141,6 +141,39 @@ pub(crate) trait MergeFunction: Send + Sync { /// Filters out DELETE and UPDATE_BEFORE rows. pub(crate) struct DeduplicateMergeFunction; +/// Configured deduplicate merge used when deletes must be kept or ignored. +pub(crate) struct ConfiguredDeduplicateMergeFunction { + keep_delete: bool, + ignore_delete: bool, +} + +impl ConfiguredDeduplicateMergeFunction { + pub(crate) fn new(table_options: &HashMap, keep_delete: bool) -> Self { + Self { + keep_delete, + ignore_delete: CoreOptions::new(table_options).ignore_delete(), + } + } +} + +/// First-row merge used when audit reads disable the normal raw-file shortcut. +pub(crate) struct FirstRowMergeFunction { + ignore_delete: bool, +} + +impl FirstRowMergeFunction { + pub(crate) fn new(table_options: &HashMap) -> Self { + Self { + ignore_delete: CoreOptions::new(table_options).ignore_delete(), + } + } +} + +fn insert_value_kind_array() -> ArrayRef { + static INSERT: OnceLock = OnceLock::new(); + Arc::clone(INSERT.get_or_init(|| Arc::new(Int8Array::from(vec![0])))) +} + fn compare_sequence_order(lhs: &MergeRow, rhs: &MergeRow) -> Ordering { match (lhs.user_sequences.is_empty(), rhs.user_sequences.is_empty()) { (false, false) => lhs @@ -151,6 +184,32 @@ fn compare_sequence_order(lhs: &MergeRow, rhs: &MergeRow) -> Ordering { } } +fn deduplicate( + rows: &[MergeRow], + keep_delete: bool, + ignore_delete: bool, +) -> crate::Result { + let mut winner = None; + for row in rows { + if ignore_delete && !RowKind::from_value(row.value_kind)?.is_add() { + continue; + } + if winner.is_none_or(|best| compare_sequence_order(row, best).is_ge()) { + winner = Some(row); + } + } + let Some(winner) = winner else { + return Ok(MergeResult::Omit); + }; + if !keep_delete && !RowKind::from_value(winner.value_kind)?.is_add() { + return Ok(MergeResult::Omit); + } + Ok(MergeResult::SourceRow { + batch_idx: winner.batch_idx, + row_idx: winner.row_idx, + }) +} + impl MergeFunction for DeduplicateMergeFunction { fn merge( &self, @@ -159,26 +218,51 @@ impl MergeFunction for DeduplicateMergeFunction { _source_output_col_indices: &[usize], _output_schema: &SchemaRef, ) -> crate::Result { - let winner = rows - .iter() - .reduce(|best, r| { - let ord = compare_sequence_order(r, best); - // >= semantics: last-writer-wins for equal values. - if ord.is_ge() { - r - } else { - best + deduplicate(rows, false, false) + } +} + +impl MergeFunction for ConfiguredDeduplicateMergeFunction { + fn merge( + &self, + rows: &[MergeRow], + _batch_buffer: &[BufferedBatch], + _source_output_col_indices: &[usize], + _output_schema: &SchemaRef, + ) -> crate::Result { + deduplicate(rows, self.keep_delete, self.ignore_delete) + } +} + +impl MergeFunction for FirstRowMergeFunction { + fn merge( + &self, + rows: &[MergeRow], + _batch_buffer: &[BufferedBatch], + _source_output_col_indices: &[usize], + _output_schema: &SchemaRef, + ) -> crate::Result { + let mut first = None; + for row in rows { + if !RowKind::from_value(row.value_kind)?.is_add() { + if self.ignore_delete { + continue; } - }) - .expect("merge called with empty rows"); - if RowKind::from_value(winner.value_kind)?.is_add() { - Ok(MergeResult::SourceRow { - batch_idx: winner.batch_idx, - row_idx: winner.row_idx, - }) - } else { - Ok(MergeResult::Omit) + return Err(Error::Unsupported { + message: "merge-engine=first-row does not support DELETE or UPDATE_BEFORE rows; set ignore-delete=true to ignore them".to_string(), + }); + } + if first.is_none_or(|current| compare_sequence_order(row, current).is_lt()) { + first = Some(row); + } } + Ok(match first { + Some(row) => MergeResult::SourceRow { + batch_idx: row.batch_idx, + row_idx: row.row_idx, + }, + None => MergeResult::Omit, + }) } } @@ -194,6 +278,7 @@ impl MergeFunction for DeduplicateMergeFunction { #[derive(Debug)] pub(crate) struct PartialUpdateMergeFunction { ignore_delete: bool, + value_kind_index: Option, sequence_groups: Vec, grouped_fields: HashSet, aggregators: Option>, @@ -216,6 +301,7 @@ impl PartialUpdateMergeFunction { PartialUpdateConfig::new(table_options).validate_write_mode(true, table_name)?; Ok(Self { ignore_delete: CoreOptions::new(table_options).ignore_delete(), + value_kind_index: None, sequence_groups: Vec::new(), grouped_fields: HashSet::new(), aggregators: None, @@ -302,6 +388,9 @@ impl PartialUpdateMergeFunction { Ok(Self { ignore_delete: CoreOptions::new(table_options).ignore_delete(), + value_kind_index: output_fields + .iter() + .position(|field| field.id() == crate::spec::VALUE_KIND_FIELD_ID), sequence_groups, grouped_fields, aggregators: aggregators @@ -364,7 +453,9 @@ impl MergeFunction for PartialUpdateMergeFunction { saw_add = true; for (output_col_idx, selected) in selected_by_col.iter_mut().enumerate() { - if self.grouped_fields.contains(&output_col_idx) { + if self.value_kind_index == Some(output_col_idx) + || self.grouped_fields.contains(&output_col_idx) + { continue; } let source_array = batch_buffer[row.batch_idx] @@ -442,18 +533,22 @@ impl MergeFunction for PartialUpdateMergeFunction { .iter() .enumerate() .map(|(output_col_idx, field)| { - let column = match aggregators - .as_ref() - .and_then(|aggregators| aggregators.get(output_col_idx)) - .and_then(Option::as_ref) - { - Some(aggregator) => aggregator.result()?, - None => match selected_by_col[output_col_idx] { - Some((batch_idx, row_idx)) => batch_buffer[batch_idx] - .column_for_output(output_col_idx, source_output_col_indices) - .slice(row_idx, 1), - None => new_null_array(field.data_type(), 1), - }, + let column = if self.value_kind_index == Some(output_col_idx) { + insert_value_kind_array() + } else { + match aggregators + .as_ref() + .and_then(|aggregators| aggregators.get(output_col_idx)) + .and_then(Option::as_ref) + { + Some(aggregator) => aggregator.result()?, + None => match selected_by_col[output_col_idx] { + Some((batch_idx, row_idx)) => batch_buffer[batch_idx] + .column_for_output(output_col_idx, source_output_col_indices) + .slice(row_idx, 1), + None => new_null_array(field.data_type(), 1), + }, + } }; if !field.is_nullable() && column.is_null(0) { return Err(Error::DataInvalid { @@ -551,6 +646,7 @@ pub(crate) struct AggregateMergeFunction { /// One slot per output column. `None` marks primary-key columns that are /// copied through; `Some` holds the aggregator that owns the column. aggregators: Mutex>>>, + value_kind_index: Option, } impl AggregateMergeFunction { @@ -580,7 +676,12 @@ impl AggregateMergeFunction { .iter() .map(|field| -> crate::Result>> { let name = field.name(); - let agg_name: &str = if seq_set.contains(name) { + if field.id() == crate::spec::VALUE_KIND_FIELD_ID { + return Ok(None); + } + let agg_name: &str = if field.id() == crate::spec::SEQUENCE_NUMBER_FIELD_ID + || seq_set.contains(name) + { "last_value" } else if pk_set.contains(name) { return Ok(None); @@ -602,6 +703,9 @@ impl AggregateMergeFunction { Ok(Self { aggregators: Mutex::new(aggregators), + value_kind_index: output_fields + .iter() + .position(|field| field.id() == crate::spec::VALUE_KIND_FIELD_ID), }) } } @@ -673,6 +777,9 @@ impl MergeFunction for AggregateMergeFunction { .iter() .enumerate() .map(|(col_idx, slot)| -> crate::Result { + if self.value_kind_index == Some(col_idx) { + return Ok(insert_value_kind_array()); + } match slot { Some(agg) => agg.result(), None => Ok(batch_buffer[pk_source.batch_idx] diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index acaebeeb1..0fd10d871 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -25,22 +25,20 @@ use super::{ArrowRecordBatchStream, Table}; use crate::arrow::build_target_arrow_schema; use crate::arrow::ParquetReadBudget; use crate::spec::{ - BigIntType, CoreOptions, DataField, DataType, MergeEngine, Predicate, TinyIntType, - ROW_KIND_FIELD_ID, ROW_KIND_FIELD_NAME, SEQUENCE_NUMBER_FIELD_ID, SEQUENCE_NUMBER_FIELD_NAME, - VALUE_KIND_FIELD_ID, VALUE_KIND_FIELD_NAME, + CoreOptions, DataField, DataType, MergeEngine, Predicate, SEQUENCE_NUMBER_FIELD_NAME, }; use crate::DataSplit; -use arrow_array::{ - builder::StringBuilder, Array, ArrayRef, RecordBatch, RecordBatchOptions, StringArray, - UInt32Array, -}; +use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions}; use arrow_schema::Schema as ArrowSchema; -use arrow_select::concat::concat as arrow_concat; -use arrow_select::take::take; +use arrow_select::interleave::interleave; use futures::{stream, StreamExt}; use std::cmp::Ordering; +use std::collections::HashMap; use std::sync::Arc; +mod audit; +pub use audit::{AuditLogInput, AuditLogRead}; + const MAX_MERGE_INPUT_STREAMS: usize = 256; /// Table read: reads data from splits (e.g. produced by [TableScan::plan]). @@ -83,6 +81,14 @@ impl<'a> TableRead<'a> { } } + /// Preserve whether the caller explicitly selected the output columns. + pub(super) fn with_explicit_projection(mut self, explicit: bool) -> Self { + if let TableReadKind::Paimon(read) = &mut self.0 { + read.explicit_projection = explicit; + } + self + } + pub(crate) fn new_format( table: &'a Table, read_type: Vec, @@ -193,26 +199,6 @@ impl<'a> TableRead<'a> { } } - /// Returns an audit-log [`ArrowRecordBatchStream`] for an incremental plan. - /// - /// Output schema is `rowkind` (+ optional `_SEQUENCE_NUMBER`) followed by - /// the projected user columns. Primary-key Delta and Changelog rows take - /// kinds from `_VALUE_KIND`; append-only Delta rows are `+I`. Diff emits - /// `+I`/`-U`/`+U`/`-D` from before/after image comparison. - pub fn to_audit_log_arrow( - &self, - plan: &IncrementalPlan, - ) -> crate::Result { - self.ensure_query_auth_allowed()?; - plan.validate()?; - match &self.0 { - TableReadKind::Paimon(read) => read.to_audit_log_arrow(plan), - TableReadKind::Format(_) => Err(crate::Error::Unsupported { - message: "Format tables do not support audit log batch read".to_string(), - }), - } - } - fn ensure_query_auth_allowed(&self) -> crate::Result<()> { CoreOptions::new(self.table().schema().options()).ensure_read_authorized() } @@ -222,6 +208,7 @@ impl<'a> TableRead<'a> { struct PaimonTableRead<'a> { table: &'a Table, read_type: Vec, + explicit_projection: bool, data_predicates: Vec, row_filter_factory: Option>, parquet_read_budget: Option>, @@ -238,6 +225,7 @@ impl<'a> PaimonTableRead<'a> { Self { table, read_type, + explicit_projection: false, data_predicates, row_filter_factory: None, parquet_read_budget: None, @@ -357,237 +345,6 @@ impl<'a> PaimonTableRead<'a> { })) } - /// Returns an audit-log stream for a planned incremental scan. - pub fn to_audit_log_arrow( - &self, - plan: &IncrementalPlan, - ) -> crate::Result { - match plan.mode() { - IncrementalScanMode::Diff => self.audit_diff_stream(plan), - IncrementalScanMode::Delta => { - self.audit_raw_stream(plan, !self.table.schema().primary_keys().is_empty()) - } - IncrementalScanMode::Changelog => self.audit_raw_stream(plan, true), - IncrementalScanMode::Auto => Err(crate::Error::DataInvalid { - message: "Incremental plan mode Auto must be resolved before consumption" - .to_string(), - source: None, - }), - } - } - - fn audit_raw_stream( - &self, - plan: &IncrementalPlan, - has_value_kind: bool, - ) -> crate::Result { - plan.validate()?; - let core_options = self.table.schema().core_options(); - let data_splits = plan.data_splits(); - let user_read_type = self.read_type.clone(); - let include_sequence = audit_sequence_number_enabled(self.table); - let audit_schema = audit_schema_for_read_type(&user_read_type, include_sequence)?; - - let mut read_type = user_read_type.clone(); - if include_sequence { - read_type.insert( - 0, - DataField::new( - SEQUENCE_NUMBER_FIELD_ID, - SEQUENCE_NUMBER_FIELD_NAME.to_string(), - DataType::BigInt(BigIntType::new()), - ), - ); - } - if has_value_kind { - read_type.push(DataField::new( - VALUE_KIND_FIELD_ID, - VALUE_KIND_FIELD_NAME.to_string(), - DataType::TinyInt(TinyIntType::new()), - )); - } - - let reader = DataFileReader::new( - self.table.file_io.clone(), - self.table.schema_manager().clone(), - self.table.schema().id(), - self.table.schema.fields().to_vec(), - read_type, - self.data_predicates.clone(), - ) - .with_file_index_read_enabled(core_options.file_index_read_enabled()) - .with_batch_size(Some(core_options.read_batch_size()?)) - .with_parquet_read_budget(Some(self.parquet_read_budget()?)); - let raw_stream = reader.read(&data_splits)?; - - Ok(Box::pin(async_stream::try_stream! { - futures::pin_mut!(raw_stream); - while let Some(batch) = raw_stream.next().await { - let batch = batch?; - let rowkind_col: ArrayRef = if has_value_kind { - let col = batch - .column_by_name(VALUE_KIND_FIELD_NAME) - .ok_or_else(|| crate::Error::DataInvalid { - message: "Changelog audit read missing _VALUE_KIND column".to_string(), - source: None, - })?; - Arc::new(rowkind_array_from_column(col)?) - } else { - let inserts: Vec<&'static str> = (0..batch.num_rows()).map(|_| "+I").collect(); - Arc::new(StringArray::from(inserts)) - }; - - let mut columns: Vec = vec![rowkind_col]; - if include_sequence { - let seq_col = batch - .column_by_name(SEQUENCE_NUMBER_FIELD_NAME) - .ok_or_else(|| crate::Error::DataInvalid { - message: "Audit read missing _SEQUENCE_NUMBER column".to_string(), - source: None, - })?; - columns.push(seq_col.clone()); - } - for field in &user_read_type { - let col = batch - .column_by_name(field.name()) - .ok_or_else(|| crate::Error::DataInvalid { - message: format!( - "Audit read missing column '{}'", - field.name() - ), - source: None, - })?; - columns.push(col.clone()); - } - yield RecordBatch::try_new(audit_schema.clone(), columns).map_err(|e| { - crate::Error::UnexpectedError { - message: format!("Failed to build audit log batch: {e}"), - source: Some(Box::new(e)), - } - })?; - } - })) - } - - fn audit_diff_stream(&self, plan: &IncrementalPlan) -> crate::Result { - let pairs = diff_pairs(plan)?; - let parallel = CoreOptions::new(self.table.schema().options()).diff_parallelism(); - let table = self.table.clone(); - let read_type = self.read_type.clone(); - let data_predicates = self.data_predicates.clone(); - let parquet_read_budget = self.parquet_read_budget()?; - - Ok(Box::pin(async_stream::try_stream! { - let mut workers = stream::iter(pairs.into_iter().map(|(before, after)| { - let table = table.clone(); - let read_type = read_type.clone(); - let data_predicates = data_predicates.clone(); - let parquet_read_budget = Arc::clone(&parquet_read_budget); - let worker: ArrowRecordBatchStream = Box::pin(async_stream::try_stream! { - let pair_read = PaimonTableRead::new(&table, read_type, data_predicates) - .with_parquet_read_budget(parquet_read_budget); - let mut pair_stream = - pair_read.to_audit_log_arrow_for_diff(&before, &after)?; - while let Some(batch) = pair_stream.next().await { - yield batch?; - } - }); - worker - })) - .flatten_unordered(parallel); - while let Some(batch) = workers.next().await { - yield batch?; - } - })) - } - - fn to_audit_log_arrow_for_diff( - &self, - before: &[DataSplit], - after: &[DataSplit], - ) -> crate::Result { - let include_sequence = audit_sequence_number_enabled(self.table); - let audit_schema = audit_schema_for_read_type(&self.read_type, include_sequence)?; - - let mut diff_read_type = self.table.schema().fields().to_vec(); - ensure_diff_supported_read_type(&diff_read_type)?; - if include_sequence { - diff_read_type.insert( - 0, - DataField::new( - SEQUENCE_NUMBER_FIELD_ID, - SEQUENCE_NUMBER_FIELD_NAME.to_string(), - DataType::BigInt(BigIntType::new()), - ), - ); - } - - let key_indices = primary_key_indices(self.table, &diff_read_type)?; - let value_indices = value_indices_for_diff(self.table, &diff_read_type); - - let before = before.to_vec(); - let after = after.to_vec(); - let table = self.table.clone(); - let read_type_for_output = self.read_type.clone(); - let data_predicates = self.data_predicates.clone(); - let parquet_read_budget = self.parquet_read_budget()?; - - Ok(Box::pin(async_stream::try_stream! { - let core_options = CoreOptions::new(table.schema().options()); - let pair_read = PaimonTableRead::new(&table, diff_read_type.clone(), data_predicates) - .with_parquet_read_budget(parquet_read_budget); - let before_stream = - pair_read.read_pk_sorted_for_diff_with_type(&before, &core_options, &diff_read_type)?; - let after_stream = - pair_read.read_pk_sorted_for_diff_with_type(&after, &core_options, &diff_read_type)?; - let mut bc = ArrowCursor::new(before_stream).await?; - let mut ac = ArrowCursor::new(after_stream).await?; - let mut data_col_indices: Option> = None; - let mut builder = AuditBatchBuilder::new(audit_schema.clone()); - - while bc.alive() || ac.alive() { - let indices = data_col_indices.get_or_insert_with(|| { - let sample = if bc.alive() { - bc.batch() - } else { - ac.batch() - }; - diff_output_col_indices(sample, &read_type_for_output, include_sequence) - .expect("diff output column indices") - }); - if !builder.has_data_columns() { - builder.set_data_col_indices(indices.clone()); - } - match cursor_cmp(&bc, &ac, &key_indices, &value_indices)? { - CursorOrd::BeforeOnly => { - builder.push("-D", bc.batch(), bc.row()); - bc.advance().await?; - } - CursorOrd::AfterOnly => { - builder.push("+I", ac.batch(), ac.row()); - ac.advance().await?; - } - CursorOrd::EqualSame => { - bc.advance().await?; - ac.advance().await?; - } - CursorOrd::EqualDiff => { - builder.push("-U", bc.batch(), bc.row()); - builder.push("+U", ac.batch(), ac.row()); - bc.advance().await?; - ac.advance().await?; - } - } - if builder.len() >= DIFF_BATCH_SIZE { - yield builder.flush()?; - } - } - if builder.len() > 0 { - yield builder.flush()?; - } - })) - } - fn to_diff_after_image_stream( &self, before: &[DataSplit], @@ -632,8 +389,8 @@ impl<'a> PaimonTableRead<'a> { &core_options, &diff_read_type, )?; - let mut bc = ArrowCursor::new(before_stream).await?; - let mut ac = ArrowCursor::new(after_stream).await?; + let mut bc = ArrowCursor::new(before_stream, 0).await?; + let mut ac = ArrowCursor::new(after_stream, 1).await?; let mut builder = DiffAfterImageBatchBuilder::new(output_schema.clone(), output_col_indices.clone()); @@ -643,7 +400,7 @@ impl<'a> PaimonTableRead<'a> { bc.advance().await?; } CursorOrd::AfterOnly => { - builder.push(ac.batch(), ac.row()); + builder.push(ac.batch_id(), ac.batch(), ac.row()); ac.advance().await?; } CursorOrd::EqualSame => { @@ -651,7 +408,7 @@ impl<'a> PaimonTableRead<'a> { ac.advance().await?; } CursorOrd::EqualDiff => { - builder.push(ac.batch(), ac.row()); + builder.push(ac.batch_id(), ac.batch(), ac.row()); bc.advance().await?; ac.advance().await?; } @@ -703,6 +460,7 @@ impl<'a> PaimonTableRead<'a> { .map(|s| s.to_string()) .collect(), read_batch_size: core_options.read_batch_size()?, + keep_delete: false, merge_splits: true, max_merge_input_streams: Some(MAX_MERGE_INPUT_STREAMS), // Diff primes the before and after streams in sequence. Keeping @@ -844,6 +602,7 @@ impl<'a> PaimonTableRead<'a> { .map(|s| s.to_string()) .collect(), read_batch_size: core_options.read_batch_size()?, + keep_delete: false, merge_splits: false, max_merge_input_streams: (core_options.deletion_vectors_enabled() && core_options.deletion_vectors_merge_on_read()) @@ -910,70 +669,6 @@ impl<'a> PaimonTableRead<'a> { } } -fn audit_schema_for_read_type( - read_type: &[DataField], - include_sequence: bool, -) -> crate::Result> { - let mut fields = Vec::with_capacity(read_type.len() + 2); - fields.push(DataField::new( - ROW_KIND_FIELD_ID, - ROW_KIND_FIELD_NAME.to_string(), - DataType::VarChar(crate::spec::VarCharType::string_type()), - )); - if include_sequence { - fields.push(DataField::new( - SEQUENCE_NUMBER_FIELD_ID, - SEQUENCE_NUMBER_FIELD_NAME.to_string(), - DataType::BigInt(BigIntType::new()), - )); - } - fields.extend(read_type.iter().cloned()); - build_target_arrow_schema(&fields) -} - -fn audit_sequence_number_enabled(table: &Table) -> bool { - table - .schema() - .options() - .get("table-read.sequence-number.enabled") - .is_some_and(|v| v.eq_ignore_ascii_case("true")) -} - -fn rowkind_array_from_column(column: &dyn arrow_array::Array) -> crate::Result { - let values = column - .as_any() - .downcast_ref::() - .ok_or_else(|| crate::Error::DataInvalid { - message: "AuditLogTable _VALUE_KIND column must be Int8".to_string(), - source: None, - })?; - let mut strings = Vec::with_capacity(values.len()); - for idx in 0..values.len() { - if values.is_null(idx) { - return Err(crate::Error::DataInvalid { - message: format!("AuditLogTable _VALUE_KIND is null at row {idx}"), - source: None, - }); - } - let rowkind = match values.value(idx) { - 0 => "+I", - 1 => "-U", - 2 => "+U", - 3 => "-D", - value => { - return Err(crate::Error::DataInvalid { - message: format!( - "AuditLogTable _VALUE_KIND has invalid value {value} at row {idx}" - ), - source: None, - }); - } - }; - strings.push(rowkind); - } - Ok(StringArray::from(strings)) -} - const DIFF_BATCH_SIZE: usize = 8192; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -987,14 +682,18 @@ enum CursorOrd { struct ArrowCursor { stream: ArrowRecordBatchStream, batch: Option, + source_id: usize, + batch_id: usize, row: usize, } impl ArrowCursor { - async fn new(stream: ArrowRecordBatchStream) -> crate::Result { + async fn new(stream: ArrowRecordBatchStream, source_id: usize) -> crate::Result { let mut cursor = Self { stream, batch: None, + source_id, + batch_id: 0, row: 0, }; cursor.advance().await?; @@ -1013,6 +712,10 @@ impl ArrowCursor { self.row } + fn batch_id(&self) -> (usize, usize) { + (self.source_id, self.batch_id) + } + async fn advance(&mut self) -> crate::Result<()> { loop { if let Some(ref batch) = self.batch { @@ -1023,6 +726,7 @@ impl ArrowCursor { } match self.stream.next().await { Some(Ok(batch)) if batch.num_rows() > 0 => { + self.batch_id += 1; self.batch = Some(batch); self.row = 0; return Ok(()); @@ -1038,100 +742,11 @@ impl ArrowCursor { } } -struct AuditBatchBuilder { - schema: Arc, - rowkind: StringBuilder, - row_indices: Vec<(usize, usize)>, - pinned_batches: Vec, - data_col_indices: Vec, - len: usize, -} - -impl AuditBatchBuilder { - fn new(schema: Arc) -> Self { - Self { - schema, - rowkind: StringBuilder::new(), - row_indices: Vec::new(), - pinned_batches: Vec::new(), - data_col_indices: Vec::new(), - len: 0, - } - } - - fn has_data_columns(&self) -> bool { - !self.data_col_indices.is_empty() - } - - fn set_data_col_indices(&mut self, indices: Vec) { - self.data_col_indices = indices; - } - - fn len(&self) -> usize { - self.len - } - - fn push(&mut self, kind: &str, batch: &RecordBatch, row: usize) { - self.rowkind.append_value(kind); - let batch_id = self.pin_batch(batch); - self.row_indices.push((batch_id, row)); - self.len += 1; - } - - fn pin_batch(&mut self, batch: &RecordBatch) -> usize { - if let Some(last) = self.pinned_batches.last() { - if std::ptr::eq(batch, last) { - return self.pinned_batches.len() - 1; - } - } - let batch_id = self.pinned_batches.len(); - self.pinned_batches.push(batch.clone()); - batch_id - } - - fn flush(&mut self) -> crate::Result { - let mut columns: Vec = vec![Arc::new(self.rowkind.finish())]; - self.rowkind = StringBuilder::new(); - for &col_idx in &self.data_col_indices { - let taken: Vec = self - .row_indices - .iter() - .map(|(batch_id, row)| { - take( - self.pinned_batches[*batch_id].column(col_idx).as_ref(), - &UInt32Array::from(vec![*row as u32]), - None, - ) - .map_err(|e| crate::Error::UnexpectedError { - message: format!("Failed to take audit diff column: {e}"), - source: Some(Box::new(e)), - }) - }) - .collect::>>()?; - let refs: Vec<&dyn Array> = taken.iter().map(|array| array.as_ref()).collect(); - columns.push( - arrow_concat(&refs).map_err(|e| crate::Error::UnexpectedError { - message: format!("Failed to concat audit diff column: {e}"), - source: Some(Box::new(e)), - })?, - ); - } - self.row_indices.clear(); - self.pinned_batches.clear(); - self.len = 0; - RecordBatch::try_new(self.schema.clone(), columns).map_err(|e| { - crate::Error::UnexpectedError { - message: format!("Failed to build audit diff batch: {e}"), - source: Some(Box::new(e)), - } - }) - } -} - struct DiffAfterImageBatchBuilder { schema: Arc, row_indices: Vec<(usize, usize)>, pinned_batches: Vec, + pinned_batch_ids: HashMap<(usize, usize), usize>, col_indices: Vec, len: usize, } @@ -1142,6 +757,7 @@ impl DiffAfterImageBatchBuilder { schema, row_indices: Vec::new(), pinned_batches: Vec::new(), + pinned_batch_ids: HashMap::new(), col_indices, len: 0, } @@ -1151,52 +767,24 @@ impl DiffAfterImageBatchBuilder { self.len } - fn push(&mut self, batch: &RecordBatch, row: usize) { - let batch_id = self.pin_batch(batch); + fn push(&mut self, batch_id: (usize, usize), batch: &RecordBatch, row: usize) { + let batch_id = pin_batch( + &mut self.pinned_batches, + &mut self.pinned_batch_ids, + batch_id, + batch, + ); self.row_indices.push((batch_id, row)); self.len += 1; } - fn pin_batch(&mut self, batch: &RecordBatch) -> usize { - if let Some(last) = self.pinned_batches.last() { - if std::ptr::eq(batch, last) { - return self.pinned_batches.len() - 1; - } - } - let batch_id = self.pinned_batches.len(); - self.pinned_batches.push(batch.clone()); - batch_id - } - fn flush(&mut self) -> crate::Result { let row_count = self.len; - let mut columns = Vec::with_capacity(self.col_indices.len()); - for &col_idx in &self.col_indices { - let taken: Vec = self - .row_indices - .iter() - .map(|(batch_id, row)| { - take( - self.pinned_batches[*batch_id].column(col_idx).as_ref(), - &UInt32Array::from(vec![*row as u32]), - None, - ) - .map_err(|e| crate::Error::UnexpectedError { - message: format!("Failed to take diff after-image column: {e}"), - source: Some(Box::new(e)), - }) - }) - .collect::>>()?; - let refs: Vec<&dyn Array> = taken.iter().map(|array| array.as_ref()).collect(); - columns.push( - arrow_concat(&refs).map_err(|e| crate::Error::UnexpectedError { - message: format!("Failed to concat diff after-image column: {e}"), - source: Some(Box::new(e)), - })?, - ); - } + let columns = + interleave_columns(&self.pinned_batches, &self.col_indices, &self.row_indices)?; self.row_indices.clear(); self.pinned_batches.clear(); + self.pinned_batch_ids.clear(); self.len = 0; let options = RecordBatchOptions::new().with_row_count(Some(row_count)); RecordBatch::try_new_with_options(self.schema.clone(), columns, &options).map_err(|e| { @@ -1208,6 +796,41 @@ impl DiffAfterImageBatchBuilder { } } +fn pin_batch( + pinned_batches: &mut Vec, + pinned_batch_ids: &mut HashMap<(usize, usize), usize>, + batch_id: (usize, usize), + batch: &RecordBatch, +) -> usize { + if let Some(&pinned_id) = pinned_batch_ids.get(&batch_id) { + return pinned_id; + } + let pinned_id = pinned_batches.len(); + pinned_batches.push(batch.clone()); + pinned_batch_ids.insert(batch_id, pinned_id); + pinned_id +} + +fn interleave_columns( + batches: &[RecordBatch], + column_indices: &[usize], + row_indices: &[(usize, usize)], +) -> crate::Result> { + column_indices + .iter() + .map(|&column_idx| { + let arrays: Vec<&dyn Array> = batches + .iter() + .map(|batch| batch.column(column_idx).as_ref()) + .collect(); + interleave(&arrays, row_indices).map_err(|e| crate::Error::UnexpectedError { + message: format!("Failed to interleave diff column: {e}"), + source: Some(Box::new(e)), + }) + }) + .collect() +} + fn diff_pairs(plan: &IncrementalPlan) -> crate::Result, Vec)>> { plan.validate()?; if plan.mode() != IncrementalScanMode::Diff { @@ -1228,34 +851,6 @@ fn diff_pairs(plan: &IncrementalPlan) -> crate::Result, Vec< .collect() } -fn diff_output_col_indices( - batch: &RecordBatch, - read_type: &[DataField], - include_sequence: bool, -) -> crate::Result> { - let mut indices = Vec::with_capacity(read_type.len() + usize::from(include_sequence)); - if include_sequence { - indices.push( - batch - .schema() - .index_of(SEQUENCE_NUMBER_FIELD_NAME) - .map_err(|e| crate::Error::DataInvalid { - message: format!("Diff read missing _SEQUENCE_NUMBER: {e}"), - source: None, - })?, - ); - } - for field in read_type { - indices.push(batch.schema().index_of(field.name()).map_err(|e| { - crate::Error::DataInvalid { - message: format!("Diff read missing column '{}': {e}", field.name()), - source: None, - } - })?); - } - Ok(indices) -} - fn value_indices_for_diff(table: &Table, fields: &[DataField]) -> Vec { let primary_key_names = table.schema().trimmed_primary_keys(); let primary_keys: std::collections::HashSet<&str> = @@ -1495,9 +1090,51 @@ mod tests { }; use crate::table::query_auth_table; use crate::table::source::DataSplitBuilder; + use arrow_array::Int32Array; + use arrow_schema::{DataType as ArrowDataType, Field}; use futures::TryStreamExt; - fn file(name: &str, level: i32, delete_row_count: Option) -> DataFileMeta { + #[test] + fn test_diff_batch_builders_pin_each_input_batch_once() { + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "id", + ArrowDataType::Int32, + false, + )])); + let input_a = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1, 2]))]) + .unwrap(); + let input_b = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![3, 4]))]) + .unwrap(); + + let mut after = DiffAfterImageBatchBuilder::new( + Arc::new(ArrowSchema::new(vec![Field::new( + "id", + ArrowDataType::Int32, + false, + )])), + vec![0], + ); + after.push((0, 1), &input_a, 1); + after.push((1, 1), &input_b, 0); + after.push((0, 1), &input_a, 0); + after.push((1, 1), &input_b, 1); + assert_eq!(after.pinned_batches.len(), 2); + let after_batch = after.flush().unwrap(); + let after_ids = after_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + after_ids.values(), + &[2, 3, 1, 4], + "interleaved batches must preserve row order" + ); + } + + pub(super) fn file(name: &str, level: i32, delete_row_count: Option) -> DataFileMeta { DataFileMeta { file_name: name.to_string(), file_size: 128, @@ -1523,7 +1160,7 @@ mod tests { } } - fn split(files: Vec, raw_convertible: bool) -> DataSplit { + pub(super) fn split(files: Vec, raw_convertible: bool) -> DataSplit { DataSplitBuilder::new() .with_snapshot(1) .with_partition(BinaryRow::new(0)) @@ -1577,11 +1214,14 @@ mod tests { .to_vec() } - fn file_index_table(path: &str, enabled: Option) -> Table { + fn file_index_table(path: &str, enabled: Option, primary_key: bool) -> Table { let mut builder = Schema::builder().column("id", DataType::Int(IntType::new())); if let Some(enabled) = enabled { builder = builder.option("file-index.read.enabled", enabled.to_string()); } + if primary_key { + builder = builder.primary_key(["id"]).option("bucket", "1"); + } Table::new( FileIOBuilder::new("memory").build().unwrap(), Identifier::new("default", "file_index_t"), @@ -1597,7 +1237,7 @@ mod tests { indexed_file.row_count = 1; indexed_file.embedded_index = Some(embedded_bitmap_index().await); let split = split(vec![indexed_file], true); - let table = file_index_table("memory:/table_read_file_index", None); + let table = file_index_table("memory:/table_read_file_index", None, false); let fields = table.schema().fields().to_vec(); let predicate = PredicateBuilder::new(&fields) .equal("id", Datum::Int(99)) @@ -1631,8 +1271,23 @@ mod tests { .unwrap(); assert!(audit.is_empty()); + let pk_table = file_index_table("memory:/table_read_audit_file_index", None, true); + let pk_fields = pk_table.schema().fields().to_vec(); + let pk_predicate = PredicateBuilder::new(&pk_fields) + .equal("id", Datum::Int(99)) + .unwrap(); + let pk_read = TableRead::new(&pk_table, pk_fields, vec![pk_predicate]); + let splits = vec![split.clone()]; + let current_audit = pk_read + .to_audit_log_arrow(&splits) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!(current_audit.is_empty()); + let disabled_table = - file_index_table("memory:/table_read_file_index_disabled", Some(false)); + file_index_table("memory:/table_read_file_index_disabled", Some(false), false); let disabled_fields = disabled_table.schema().fields().to_vec(); let disabled_predicate = PredicateBuilder::new(&disabled_fields) .equal("id", Datum::Int(99)) @@ -1669,25 +1324,6 @@ mod tests { assert!(!pk_split_needs_merge(&dv_compacted, true)); } - #[test] - fn test_rowkind_rejects_null_value_kind() { - let values = arrow_array::Int8Array::from(vec![Some(0), None]); - assert!(matches!( - rowkind_array_from_column(&values), - Err(crate::Error::DataInvalid { ref message, .. }) if message.contains("null at row 1") - )); - } - - #[test] - fn test_rowkind_rejects_invalid_value_kind() { - let values = arrow_array::Int8Array::from(vec![4]); - assert!(matches!( - rowkind_array_from_column(&values), - Err(crate::Error::DataInvalid { ref message, .. }) - if message.contains("invalid value 4 at row 0") - )); - } - #[test] fn test_direct_table_read_fails_closed_when_query_auth_enabled() { let table = query_auth_table(); diff --git a/crates/paimon/src/table/table_read/audit.rs b/crates/paimon/src/table/table_read/audit.rs new file mode 100644 index 000000000..0ba1cf0e8 --- /dev/null +++ b/crates/paimon/src/table/table_read/audit.rs @@ -0,0 +1,940 @@ +// 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. + +//! Audit row kinds, projection and current/incremental read policy. + +use super::{ + cursor_cmp, diff_pairs, ensure_diff_supported_read_type, interleave_columns, pin_batch, + primary_key_indices, value_indices_for_diff, ArrowCursor, CursorOrd, PaimonTableRead, + TableRead, TableReadKind, DIFF_BATCH_SIZE, MAX_MERGE_INPUT_STREAMS, +}; +use crate::arrow::build_target_arrow_schema; +use crate::spec::{ + BigIntType, CoreOptions, DataField, DataType, MergeEngine, TinyIntType, ROW_KIND_FIELD_ID, + ROW_KIND_FIELD_NAME, SEQUENCE_NUMBER_FIELD_ID, SEQUENCE_NUMBER_FIELD_NAME, VALUE_KIND_FIELD_ID, + VALUE_KIND_FIELD_NAME, +}; +use crate::table::data_file_reader::DataFileReader; +use crate::table::incremental_scan::{IncrementalPlan, IncrementalScanMode}; +use crate::table::kv_file_reader::{KeyValueFileReader, KeyValueReadConfig}; +use crate::table::{ArrowRecordBatchStream, ReadBuilder, Table, TableScan}; +use crate::DataSplit; +use arrow_array::{ + builder::StringBuilder, Array, ArrayRef, RecordBatch, RecordBatchOptions, StringArray, +}; +use arrow_schema::Schema as ArrowSchema; +use futures::{stream, StreamExt}; +use std::collections::HashMap; +use std::sync::Arc; + +#[derive(Debug, Clone, Copy)] +pub enum AuditLogInput<'a> { + Current(&'a [DataSplit]), + Incremental(&'a IncrementalPlan), +} + +impl<'a> From<&'a [DataSplit]> for AuditLogInput<'a> { + fn from(splits: &'a [DataSplit]) -> Self { + Self::Current(splits) + } +} + +impl<'a, const N: usize> From<&'a [DataSplit; N]> for AuditLogInput<'a> { + fn from(splits: &'a [DataSplit; N]) -> Self { + Self::Current(splits) + } +} + +impl<'a> From<&'a Vec> for AuditLogInput<'a> { + fn from(splits: &'a Vec) -> Self { + Self::Current(splits.as_slice()) + } +} + +impl<'a> From<&'a IncrementalPlan> for AuditLogInput<'a> { + fn from(plan: &'a IncrementalPlan) -> Self { + Self::Incremental(plan) + } +} + +/// Audit reader retaining winning retract rows and exposing their physical row kind. +/// +/// Reuses the projection, predicates and Parquet budget of the supplied read. +/// Without an explicit projection, adds `rowkind` and the configured sequence column. +#[derive(Debug, Clone)] +pub struct AuditLogRead<'a> { + read: PaimonTableRead<'a>, + projection: Option>, +} + +impl<'a> AuditLogRead<'a> { + pub fn new(read: TableRead<'a>) -> crate::Result { + read.ensure_query_auth_allowed()?; + match read.0 { + TableReadKind::Paimon(read) => { + let projection = read.explicit_projection.then(|| read.read_type.clone()); + Ok(Self { read, projection }) + } + TableReadKind::Format(_) => Err(crate::Error::Unsupported { + message: "Format tables do not support audit log batch read".to_string(), + }), + } + } + + /// Reads current-state splits or a validated incremental plan. + pub fn to_arrow<'input>( + &self, + input: impl Into>, + ) -> crate::Result { + match input.into() { + AuditLogInput::Current(splits) => self.audit_current_stream(splits), + AuditLogInput::Incremental(plan) => { + plan.validate()?; + self.audit_incremental_stream(plan) + } + } + } + + fn audit_current_stream( + &self, + data_splits: &[DataSplit], + ) -> crate::Result { + let output_read_type = self.audit_read_type()?; + let include_rowkind = audit_field_requested(&output_read_type, ROW_KIND_FIELD_ID); + let include_sequence = audit_field_requested(&output_read_type, SEQUENCE_NUMBER_FIELD_ID); + let user_read_type = self.audit_user_read_type(); + let audit_schema = + audit_schema_for_read_type(&user_read_type, include_rowkind, include_sequence)?; + let has_primary_keys = !self.read.table.schema().primary_keys().is_empty(); + + let physical_stream = if has_primary_keys { + let core_options = self.read.table.schema().core_options(); + let mut read_type = Vec::with_capacity(user_read_type.len() + 2); + if include_sequence { + read_type.push(DataField::new( + SEQUENCE_NUMBER_FIELD_ID, + SEQUENCE_NUMBER_FIELD_NAME.to_string(), + DataType::BigInt(BigIntType::new()), + )); + } + if include_rowkind { + read_type.push(DataField::new( + VALUE_KIND_FIELD_ID, + VALUE_KIND_FIELD_NAME.to_string(), + DataType::TinyInt(TinyIntType::new()), + )); + } + read_type.extend(user_read_type.iter().cloned()); + + let merge_engine = core_options.merge_engine()?; + let (raw_splits, merge_splits) = partition_audit_splits(data_splits, merge_engine); + let parquet_read_budget = self.read.parquet_read_budget()?; + let raw_stream = DataFileReader::new( + self.read.table.file_io.clone(), + self.read.table.schema_manager().clone(), + self.read.table.schema().id(), + self.read.table.schema.fields().to_vec(), + read_type.clone(), + self.read.data_predicates.clone(), + ) + .with_file_index_read_enabled(core_options.file_index_read_enabled()) + .with_batch_size(Some(core_options.read_batch_size()?)) + .with_parquet_read_budget(Some(Arc::clone(&parquet_read_budget))) + .read(&raw_splits)?; + let merge_reader = KeyValueFileReader::new( + self.read.table.file_io.clone(), + KeyValueReadConfig { + table_name: self.read.table.identifier().full_name(), + table_options: self.read.table.schema().options().clone(), + schema_manager: self.read.table.schema_manager().clone(), + table_schema_id: self.read.table.schema().id(), + table_fields: self.read.table.schema.fields().to_vec(), + read_type, + predicates: self.read.data_predicates.clone(), + primary_keys: self.read.table.schema.trimmed_primary_keys(), + merge_engine, + sequence_fields: core_options + .sequence_fields() + .iter() + .map(|field| field.to_string()) + .collect(), + read_batch_size: core_options.read_batch_size()?, + keep_delete: true, + merge_splits: merge_engine == MergeEngine::FirstRow, + max_merge_input_streams: Some(MAX_MERGE_INPUT_STREAMS), + parquet_read_budget: Some(parquet_read_budget), + }, + ); + let merge_stream = if merge_engine == MergeEngine::FirstRow { + let mut groups: HashMap<(Vec, i32), Vec> = HashMap::new(); + for split in merge_splits { + groups + .entry((split.partition().to_serialized_bytes(), split.bucket())) + .or_default() + .push(split); + } + Box::pin(async_stream::try_stream! { + for splits in groups.into_values() { + let mut group_stream = merge_reader.clone().read(&splits)?; + while let Some(batch) = group_stream.next().await { + yield batch?; + } + } + }) as ArrowRecordBatchStream + } else { + merge_reader.read(&merge_splits)? + }; + Box::pin(stream::select_all([raw_stream, merge_stream])) + } else { + self.read.to_arrow(data_splits)? + }; + + let stream = audit_stream_from_physical( + physical_stream, + audit_schema, + user_read_type, + include_rowkind, + include_sequence, + has_primary_keys && include_rowkind, + ); + project_audit_stream(stream, self.projection.as_deref()) + } + + fn audit_incremental_stream( + &self, + plan: &IncrementalPlan, + ) -> crate::Result { + match plan.mode() { + IncrementalScanMode::Diff => self.audit_diff_stream(plan), + IncrementalScanMode::Delta => { + self.audit_raw_stream(plan, !self.read.table.schema().primary_keys().is_empty()) + } + IncrementalScanMode::Changelog => self.audit_raw_stream(plan, true), + IncrementalScanMode::Auto => Err(crate::Error::DataInvalid { + message: "Incremental plan mode Auto must be resolved before consumption" + .to_string(), + source: None, + }), + } + } + + fn audit_read_type(&self) -> crate::Result> { + let fields = self.projection.clone().unwrap_or_else(|| { + audit_fields_for_read_type( + &self.read.read_type, + true, + audit_sequence_number_enabled(self.read.table), + ) + }); + if audit_field_requested(&fields, SEQUENCE_NUMBER_FIELD_ID) + && !audit_sequence_number_enabled(self.read.table) + { + return Err(crate::Error::DataInvalid { + message: "Audit read requested _SEQUENCE_NUMBER but table-read.sequence-number.enabled is false".to_string(), + source: None, + }); + } + Ok(fields) + } + + fn audit_user_read_type(&self) -> Vec { + self.read + .read_type + .iter() + .filter(|field| !matches!(field.id(), ROW_KIND_FIELD_ID | SEQUENCE_NUMBER_FIELD_ID)) + .cloned() + .collect() + } + + fn audit_raw_stream( + &self, + plan: &IncrementalPlan, + has_value_kind: bool, + ) -> crate::Result { + plan.validate()?; + let core_options = self.read.table.schema().core_options(); + let data_splits = plan.data_splits(); + let output_read_type = self.audit_read_type()?; + let user_read_type = self.audit_user_read_type(); + let include_rowkind = audit_field_requested(&output_read_type, ROW_KIND_FIELD_ID); + let include_sequence = audit_field_requested(&output_read_type, SEQUENCE_NUMBER_FIELD_ID); + let audit_schema = + audit_schema_for_read_type(&user_read_type, include_rowkind, include_sequence)?; + + let mut read_type = user_read_type.clone(); + if include_sequence { + read_type.insert( + 0, + DataField::new( + SEQUENCE_NUMBER_FIELD_ID, + SEQUENCE_NUMBER_FIELD_NAME.to_string(), + DataType::BigInt(BigIntType::new()), + ), + ); + } + if has_value_kind && include_rowkind { + read_type.push(DataField::new( + VALUE_KIND_FIELD_ID, + VALUE_KIND_FIELD_NAME.to_string(), + DataType::TinyInt(TinyIntType::new()), + )); + } + + let reader = DataFileReader::new( + self.read.table.file_io.clone(), + self.read.table.schema_manager().clone(), + self.read.table.schema().id(), + self.read.table.schema.fields().to_vec(), + read_type, + self.read.data_predicates.clone(), + ) + .with_file_index_read_enabled(core_options.file_index_read_enabled()) + .with_batch_size(Some(core_options.read_batch_size()?)) + .with_parquet_read_budget(Some(self.read.parquet_read_budget()?)); + let raw_stream = reader.read(&data_splits)?; + let stream = audit_stream_from_physical( + raw_stream, + audit_schema, + user_read_type, + include_rowkind, + include_sequence, + has_value_kind && include_rowkind, + ); + project_audit_stream(stream, self.projection.as_deref()) + } + + fn audit_diff_stream(&self, plan: &IncrementalPlan) -> crate::Result { + let pairs = diff_pairs(plan)?; + let parallel = CoreOptions::new(self.read.table.schema().options()).diff_parallelism(); + let output_read_type = self.audit_read_type()?; + let include_sequence = audit_field_requested(&output_read_type, SEQUENCE_NUMBER_FIELD_ID); + let table = self.read.table.clone(); + let read_type = self.audit_user_read_type(); + let data_predicates = self.read.data_predicates.clone(); + let parquet_read_budget = self.read.parquet_read_budget()?; + + let stream: ArrowRecordBatchStream = Box::pin(async_stream::try_stream! { + let mut workers = stream::iter(pairs.into_iter().map(|(before, after)| { + let table = table.clone(); + let read_type = read_type.clone(); + let data_predicates = data_predicates.clone(); + let parquet_read_budget = Arc::clone(&parquet_read_budget); + let worker: ArrowRecordBatchStream = Box::pin(async_stream::try_stream! { + let pair_read = AuditLogRead { + read: PaimonTableRead::new(&table, read_type, data_predicates) + .with_parquet_read_budget(parquet_read_budget), + projection: None, + }; + let mut pair_stream = + pair_read.to_audit_log_arrow_for_diff( + &before, + &after, + include_sequence, + )?; + while let Some(batch) = pair_stream.next().await { + yield batch?; + } + }); + worker + })) + .flatten_unordered(parallel); + while let Some(batch) = workers.next().await { + yield batch?; + } + }); + project_audit_stream(stream, self.projection.as_deref()) + } + + fn to_audit_log_arrow_for_diff( + &self, + before: &[DataSplit], + after: &[DataSplit], + include_sequence: bool, + ) -> crate::Result { + let audit_schema = + audit_schema_for_read_type(&self.read.read_type, true, include_sequence)?; + + let mut diff_read_type = self.read.table.schema().fields().to_vec(); + ensure_diff_supported_read_type(&diff_read_type)?; + if include_sequence { + diff_read_type.insert( + 0, + DataField::new( + SEQUENCE_NUMBER_FIELD_ID, + SEQUENCE_NUMBER_FIELD_NAME.to_string(), + DataType::BigInt(BigIntType::new()), + ), + ); + } + + let key_indices = primary_key_indices(self.read.table, &diff_read_type)?; + let value_indices = value_indices_for_diff(self.read.table, &diff_read_type); + + let before = before.to_vec(); + let after = after.to_vec(); + let table = self.read.table.clone(); + let read_type_for_output = self.read.read_type.clone(); + let data_predicates = self.read.data_predicates.clone(); + let parquet_read_budget = self.read.parquet_read_budget()?; + + Ok(Box::pin(async_stream::try_stream! { + let core_options = CoreOptions::new(table.schema().options()); + let pair_read = PaimonTableRead::new(&table, diff_read_type.clone(), data_predicates) + .with_parquet_read_budget(parquet_read_budget); + let before_stream = + pair_read.read_pk_sorted_for_diff_with_type(&before, &core_options, &diff_read_type)?; + let after_stream = + pair_read.read_pk_sorted_for_diff_with_type(&after, &core_options, &diff_read_type)?; + let mut bc = ArrowCursor::new(before_stream, 0).await?; + let mut ac = ArrowCursor::new(after_stream, 1).await?; + let mut data_col_indices: Option> = None; + let mut builder = AuditBatchBuilder::new(audit_schema.clone()); + + while bc.alive() || ac.alive() { + let indices = data_col_indices.get_or_insert_with(|| { + let sample = if bc.alive() { + bc.batch() + } else { + ac.batch() + }; + diff_output_col_indices(sample, &read_type_for_output, include_sequence) + .expect("diff output column indices") + }); + if !builder.has_data_columns() { + builder.set_data_col_indices(indices.clone()); + } + match cursor_cmp(&bc, &ac, &key_indices, &value_indices)? { + CursorOrd::BeforeOnly => { + builder.push("-D", bc.batch_id(), bc.batch(), bc.row()); + bc.advance().await?; + } + CursorOrd::AfterOnly => { + builder.push("+I", ac.batch_id(), ac.batch(), ac.row()); + ac.advance().await?; + } + CursorOrd::EqualSame => { + bc.advance().await?; + ac.advance().await?; + } + CursorOrd::EqualDiff => { + builder.push("-U", bc.batch_id(), bc.batch(), bc.row()); + builder.push("+U", ac.batch_id(), ac.batch(), ac.row()); + bc.advance().await?; + ac.advance().await?; + } + } + if builder.len() >= DIFF_BATCH_SIZE { + yield builder.flush()?; + } + } + if builder.len() > 0 { + yield builder.flush()?; + } + })) + } +} + +impl TableRead<'_> { + /// Returns audit-log rows for current splits or an incremental plan. + pub fn to_audit_log_arrow<'input>( + &self, + input: impl Into>, + ) -> crate::Result { + AuditLogRead::new(self.clone())?.to_arrow(input) + } +} + +impl<'a> ReadBuilder<'a> { + /// Create a current-state audit scan that retains every visible row version. + pub fn new_audit_scan(&self) -> TableScan<'a> { + self.new_scan().with_all_versions() + } +} + +// Legacy unknown delete counts and first-row level-0 files stay on the merge path. +fn audit_raw_convertible(split: &DataSplit, merge_engine: MergeEngine) -> bool { + split.raw_convertible() + && split.data_files().iter().all(|file| { + file.delete_row_count == Some(0) + && (merge_engine != MergeEngine::FirstRow || file.level != 0) + }) +} + +fn partition_audit_splits( + data_splits: &[DataSplit], + merge_engine: MergeEngine, +) -> (Vec, Vec) { + if merge_engine != MergeEngine::FirstRow { + return data_splits + .iter() + .cloned() + .partition(|split| audit_raw_convertible(split, merge_engine)); + } + + let mut groups: HashMap<(Vec, i32), Vec> = HashMap::new(); + for split in data_splits.iter().cloned() { + groups + .entry((split.partition().to_serialized_bytes(), split.bucket())) + .or_default() + .push(split); + } + let mut raw = Vec::new(); + let mut merge = Vec::new(); + for group in groups.into_values() { + if group + .iter() + .all(|split| audit_raw_convertible(split, merge_engine)) + { + raw.extend(group); + } else { + merge.extend(group); + } + } + (raw, merge) +} + +struct AuditPhysicalProjection { + value_kind: Option, + sequence: Option, + user: Vec, +} + +fn audit_physical_projection( + schema: &ArrowSchema, + user_read_type: &[DataField], + include_rowkind: bool, + include_sequence: bool, + has_value_kind: bool, +) -> crate::Result { + let by_name: HashMap<&str, usize> = schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| (field.name().as_str(), index)) + .collect(); + let index = |name: &str| { + by_name + .get(name) + .copied() + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("Audit read missing column '{name}'"), + source: None, + }) + }; + Ok(AuditPhysicalProjection { + value_kind: (include_rowkind && has_value_kind) + .then(|| index(VALUE_KIND_FIELD_NAME)) + .transpose()?, + sequence: include_sequence + .then(|| index(SEQUENCE_NUMBER_FIELD_NAME)) + .transpose()?, + user: user_read_type + .iter() + .map(|field| index(field.name())) + .collect::>>()?, + }) +} + +fn audit_stream_from_physical( + raw_stream: ArrowRecordBatchStream, + audit_schema: Arc, + user_read_type: Vec, + include_rowkind: bool, + include_sequence: bool, + has_value_kind: bool, +) -> ArrowRecordBatchStream { + Box::pin(async_stream::try_stream! { + futures::pin_mut!(raw_stream); + let mut projection = None; + while let Some(batch) = raw_stream.next().await { + let batch = batch?; + if projection.is_none() { + projection = Some(audit_physical_projection( + batch.schema().as_ref(), + &user_read_type, + include_rowkind, + include_sequence, + has_value_kind, + )?); + } + let projection = projection.as_ref().unwrap(); + let mut columns = Vec::with_capacity(audit_schema.fields().len()); + if include_rowkind { + let rowkind_col: ArrayRef = if let Some(index) = projection.value_kind { + Arc::new(rowkind_array_from_column(batch.column(index).as_ref())?) + } else { + Arc::new(StringArray::from(vec!["+I"; batch.num_rows()])) + }; + columns.push(rowkind_col); + } + if let Some(index) = projection.sequence { + columns.push(batch.column(index).clone()); + } + columns.extend( + projection + .user + .iter() + .map(|&index| batch.column(index).clone()), + ); + let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows())); + yield RecordBatch::try_new_with_options( + audit_schema.clone(), + columns, + &options, + ) + .map_err(|error| crate::Error::UnexpectedError { + message: format!("Failed to build audit log batch: {error}"), + source: Some(Box::new(error)), + })?; + } + }) +} + +fn project_audit_stream( + stream: ArrowRecordBatchStream, + read_type: Option<&[DataField]>, +) -> crate::Result { + let Some(read_type) = read_type else { + return Ok(stream); + }; + let schema = build_target_arrow_schema(read_type)?; + let names = read_type + .iter() + .map(|field| field.name().to_string()) + .collect::>(); + Ok(Box::pin(async_stream::try_stream! { + futures::pin_mut!(stream); + let mut indices = None; + while let Some(batch) = stream.next().await { + let batch = batch?; + let indices = indices.get_or_insert_with(|| { + names + .iter() + .map(|name| batch.schema().index_of(name)) + .collect::, _>>() + }); + let indices = indices.as_ref().map_err(|error| crate::Error::DataInvalid { + message: format!("Audit read projection failed: {error}"), + source: None, + })?; + let columns = indices + .iter() + .map(|&index| batch.column(index).clone()) + .collect(); + let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows())); + yield RecordBatch::try_new_with_options(schema.clone(), columns, &options) + .map_err(|error| crate::Error::UnexpectedError { + message: format!("Failed to project audit log batch: {error}"), + source: Some(Box::new(error)), + })?; + } + })) +} + +fn audit_field_requested(read_type: &[DataField], field_id: i32) -> bool { + read_type.iter().any(|field| field.id() == field_id) +} + +fn audit_fields_for_read_type( + read_type: &[DataField], + include_rowkind: bool, + include_sequence: bool, +) -> Vec { + let mut fields = Vec::with_capacity(read_type.len() + 2); + if include_rowkind { + fields.push(DataField::new( + ROW_KIND_FIELD_ID, + ROW_KIND_FIELD_NAME.to_string(), + DataType::VarChar(crate::spec::VarCharType::string_type()), + )); + } + if include_sequence { + fields.push(DataField::new( + SEQUENCE_NUMBER_FIELD_ID, + SEQUENCE_NUMBER_FIELD_NAME.to_string(), + DataType::BigInt(BigIntType::new()), + )); + } + fields.extend(read_type.iter().cloned()); + fields +} + +fn audit_schema_for_read_type( + read_type: &[DataField], + include_rowkind: bool, + include_sequence: bool, +) -> crate::Result> { + build_target_arrow_schema(&audit_fields_for_read_type( + read_type, + include_rowkind, + include_sequence, + )) +} + +fn audit_sequence_number_enabled(table: &Table) -> bool { + table + .schema() + .core_options() + .table_read_sequence_number_enabled() +} + +fn rowkind_array_from_column(column: &dyn arrow_array::Array) -> crate::Result { + let values = column + .as_any() + .downcast_ref::() + .ok_or_else(|| crate::Error::DataInvalid { + message: "AuditLogTable _VALUE_KIND column must be Int8".to_string(), + source: None, + })?; + let mut strings = Vec::with_capacity(values.len()); + for idx in 0..values.len() { + if values.is_null(idx) { + return Err(crate::Error::DataInvalid { + message: format!("AuditLogTable _VALUE_KIND is null at row {idx}"), + source: None, + }); + } + let rowkind = match values.value(idx) { + 0 => "+I", + 1 => "-U", + 2 => "+U", + 3 => "-D", + value => { + return Err(crate::Error::DataInvalid { + message: format!( + "AuditLogTable _VALUE_KIND has invalid value {value} at row {idx}" + ), + source: None, + }); + } + }; + strings.push(rowkind); + } + Ok(StringArray::from(strings)) +} + +struct AuditBatchBuilder { + schema: Arc, + rowkind: StringBuilder, + row_indices: Vec<(usize, usize)>, + pinned_batches: Vec, + pinned_batch_ids: HashMap<(usize, usize), usize>, + data_col_indices: Vec, + len: usize, +} + +impl AuditBatchBuilder { + fn new(schema: Arc) -> Self { + Self { + schema, + rowkind: StringBuilder::new(), + row_indices: Vec::new(), + pinned_batches: Vec::new(), + pinned_batch_ids: HashMap::new(), + data_col_indices: Vec::new(), + len: 0, + } + } + + fn has_data_columns(&self) -> bool { + !self.data_col_indices.is_empty() + } + + fn set_data_col_indices(&mut self, indices: Vec) { + self.data_col_indices = indices; + } + + fn len(&self) -> usize { + self.len + } + + fn push(&mut self, kind: &str, batch_id: (usize, usize), batch: &RecordBatch, row: usize) { + self.rowkind.append_value(kind); + let batch_id = pin_batch( + &mut self.pinned_batches, + &mut self.pinned_batch_ids, + batch_id, + batch, + ); + self.row_indices.push((batch_id, row)); + self.len += 1; + } + + fn flush(&mut self) -> crate::Result { + let mut columns: Vec = vec![Arc::new(self.rowkind.finish())]; + self.rowkind = StringBuilder::new(); + columns.extend(interleave_columns( + &self.pinned_batches, + &self.data_col_indices, + &self.row_indices, + )?); + self.row_indices.clear(); + self.pinned_batches.clear(); + self.pinned_batch_ids.clear(); + self.len = 0; + RecordBatch::try_new(self.schema.clone(), columns).map_err(|e| { + crate::Error::UnexpectedError { + message: format!("Failed to build audit diff batch: {e}"), + source: Some(Box::new(e)), + } + }) + } +} + +fn diff_output_col_indices( + batch: &RecordBatch, + read_type: &[DataField], + include_sequence: bool, +) -> crate::Result> { + let mut indices = Vec::with_capacity(read_type.len() + usize::from(include_sequence)); + if include_sequence { + indices.push( + batch + .schema() + .index_of(SEQUENCE_NUMBER_FIELD_NAME) + .map_err(|e| crate::Error::DataInvalid { + message: format!("Diff read missing _SEQUENCE_NUMBER: {e}"), + source: None, + })?, + ); + } + for field in read_type { + indices.push(batch.schema().index_of(field.name()).map_err(|e| { + crate::Error::DataInvalid { + message: format!("Diff read missing column '{}': {e}", field.name()), + source: None, + } + })?); + } + Ok(indices) +} + +#[cfg(test)] +mod tests { + use super::super::tests::{file, split}; + use super::*; + use arrow_array::Int32Array; + use arrow_schema::{DataType as ArrowDataType, Field}; + use futures::TryStreamExt; + + #[tokio::test] + async fn test_default_audit_projection_bypasses_batch_rebuild() { + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "id", + ArrowDataType::Int32, + false, + )])); + let input = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1]))]) + .unwrap(); + let stream: ArrowRecordBatchStream = + Box::pin(stream::iter(vec![Ok::<_, crate::Error>(input.clone())])); + + let output = project_audit_stream(stream, None) + .unwrap() + .try_collect::>() + .await + .unwrap(); + + assert!(Arc::ptr_eq(&schema, &output[0].schema())); + + let stream: ArrowRecordBatchStream = + Box::pin(stream::iter(vec![Ok::<_, crate::Error>(input)])); + let output = project_audit_stream(stream, Some(&[])) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(output[0].num_columns(), 0); + assert_eq!(output[0].num_rows(), 1); + } + + #[test] + fn test_rowkind_rejects_null_value_kind() { + let values = arrow_array::Int8Array::from(vec![Some(0), None]); + assert!(matches!( + rowkind_array_from_column(&values), + Err(crate::Error::DataInvalid { ref message, .. }) if message.contains("null at row 1") + )); + } + + #[test] + fn test_rowkind_rejects_invalid_value_kind() { + let values = arrow_array::Int8Array::from(vec![4]); + assert!(matches!( + rowkind_array_from_column(&values), + Err(crate::Error::DataInvalid { ref message, .. }) + if message.contains("invalid value 4 at row 0") + )); + } + + #[test] + fn test_audit_batch_builder_pins_each_input_batch_once() { + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "id", + ArrowDataType::Int32, + false, + )])); + let input_a = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1, 2]))]) + .unwrap(); + let input_b = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![3, 4]))]) + .unwrap(); + + let mut audit = AuditBatchBuilder::new(Arc::new(ArrowSchema::new(vec![ + Field::new(ROW_KIND_FIELD_NAME, ArrowDataType::Utf8, false), + Field::new("id", ArrowDataType::Int32, false), + ]))); + audit.set_data_col_indices(vec![0]); + audit.push("+I", (0, 1), &input_a, 1); + audit.push("+I", (1, 1), &input_b, 0); + audit.push("+I", (0, 1), &input_a, 0); + audit.push("+I", (1, 1), &input_b, 1); + assert_eq!(audit.pinned_batches.len(), 2); + let audit_batch = audit.flush().unwrap(); + let audit_ids = audit_batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + audit_ids.values(), + &[2, 3, 1, 4], + "interleaved batches must preserve row order" + ); + } + + #[test] + fn test_audit_split_routing() { + let raw = split(vec![file("a", 5, Some(0))], true); + let merge = split(vec![file("a", 5, Some(0))], false); + let legacy = split(vec![file("a", 5, None)], true); + assert!(audit_raw_convertible(&raw, MergeEngine::Deduplicate)); + assert!(audit_raw_convertible(&raw, MergeEngine::FirstRow)); + assert!(!audit_raw_convertible(&merge, MergeEngine::Deduplicate)); + assert!(!audit_raw_convertible(&legacy, MergeEngine::Deduplicate)); + let level_zero = split(vec![file("a", 0, Some(0))], true); + assert!(audit_raw_convertible(&level_zero, MergeEngine::Deduplicate)); + assert!(!audit_raw_convertible(&level_zero, MergeEngine::FirstRow)); + let (raw_only, merge_only) = + partition_audit_splits(std::slice::from_ref(&raw), MergeEngine::FirstRow); + assert_eq!((raw_only.len(), merge_only.len()), (1, 0)); + let (raw_group, merge_group) = + partition_audit_splits(&[raw.clone(), level_zero], MergeEngine::FirstRow); + assert_eq!((raw_group.len(), merge_group.len()), (0, 2)); + } +} diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index d5de9048a..9dd8ae3d8 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -909,6 +909,15 @@ impl<'a> TableScan<'a> { } } + /// Retain all visible versions and group overlapping keys for merging, + /// preserving the read projection. + pub(super) fn with_all_versions(self) -> Self { + match self.0 { + TableScanKind::Paimon(scan) => Self(TableScanKind::Paimon(scan.with_all_versions())), + TableScanKind::Format(scan) => Self(TableScanKind::Format(scan)), + } + } + pub fn with_row_ranges(self, ranges: Vec) -> Self { match self.0 { TableScanKind::Paimon(scan) => { @@ -1021,6 +1030,8 @@ struct PaimonTableScan<'a> { /// Used by non-read paths (overwrite, truncate, writer restore) that need /// the complete file set. Normal read scans leave this as `false`. scan_all_files: bool, + /// Whether each split must contain every file whose primary-key range overlaps. + merge_key_overlaps: bool, projected_read_field_ids: Option>, } @@ -1042,6 +1053,7 @@ impl<'a> PaimonTableScan<'a> { row_ranges, row_range_optimization_disabled: false, scan_all_files: false, + merge_key_overlaps: false, projected_read_field_ids: None, } } @@ -1056,6 +1068,12 @@ impl<'a> PaimonTableScan<'a> { self } + fn with_all_versions(mut self) -> Self { + self.scan_all_files = true; + self.merge_key_overlaps = true; + self + } + /// Set row ranges for scan-time filtering. /// /// This replaces any existing row_ranges. Typically used to inject @@ -1271,7 +1289,7 @@ impl<'a> PaimonTableScan<'a> { } fn can_push_down_limit_hint(&self, row_ranges: Option<&[RowRange]>) -> bool { - can_push_down_limit_hint_for_scan(&self.data_predicates, row_ranges) + !self.scan_all_files && can_push_down_limit_hint_for_scan(&self.data_predicates, row_ranges) } fn global_index_scan_settings( @@ -1279,12 +1297,14 @@ impl<'a> PaimonTableScan<'a> { core_options: &CoreOptions, data_evolution_enabled: bool, ) -> crate::Result> { - if should_use_global_index_row_range_optimization( - self.row_range_optimization_disabled, - data_evolution_enabled, - core_options.global_index_enabled(), - !self.data_predicates.is_empty(), - ) { + if !self.scan_all_files + && should_use_global_index_row_range_optimization( + self.row_range_optimization_disabled, + data_evolution_enabled, + core_options.global_index_enabled(), + !self.data_predicates.is_empty(), + ) + { Ok(Some(GlobalIndexScanSettings { search_mode: core_options.scalar_index_search_mode()?, thread_num: core_options.global_index_thread_num()?, @@ -1419,15 +1439,14 @@ impl<'a> PaimonTableScan<'a> { /// `KeyValueFileReader`. /// /// Exempt (full predicates kept): - /// - Deletion-vector tables without merge-on-read: they read raw with + /// - Ordinary deletion-vector reads without merge-on-read: they read raw with /// per-row masks, stats are a superset of live rows, full pruning stays /// safe. With merge-on-read enabled, visible L0 versions require the /// same key-only pruning rule as an ordinary PK merge read. - /// - `merge-engine=first-row`: planned with `skip_level_zero` and read - /// via `DataFileReader` (see `TableRead::to_arrow`), no merge on the - /// read path — pruning a file drops exactly the rows the raw path's - /// exact residual filter would drop anyway. If first-row ever gains a - /// merge read path, this exemption must be revisited. + /// - Non-audit `merge-engine=first-row` reads: read via `DataFileReader` + /// without merging versions. + /// + /// Audit reads set `merge_key_overlaps` and are not exempt. fn stats_pruning_predicates(&self) -> Vec { let has_primary_keys = !self.table.schema().primary_keys().is_empty(); let core_options = CoreOptions::new(self.table.schema().options()); @@ -1440,8 +1459,8 @@ impl<'a> PaimonTableScan<'a> { Ok(crate::spec::MergeEngine::FirstRow) ); if has_primary_keys - && (!deletion_vectors_enabled || deletion_vectors_merge_on_read) - && !first_row + && (self.merge_key_overlaps + || ((!deletion_vectors_enabled || deletion_vectors_merge_on_read) && !first_row)) { retain_primary_key_conjuncts( &self.data_predicates, @@ -1916,16 +1935,17 @@ impl<'a> PaimonTableScan<'a> { // sort-merge reader sees every version of a key. The comparator decodes // the trimmed-PK min/max keys written by the kv writer. // - // Deletion-vector tables without merge-on-read and first-row tables read - // without merging (stale rows are masked by DVs / level-0 is skipped), - // so they keep plain size-based packing. DV merge-on-read includes L0 - // files and must preserve overlapping key ranges just like ordinary MOR. - let read_merges_overlapping_keys = (!core_options.deletion_vectors_enabled() - || core_options.deletion_vectors_merge_on_read()) - && !matches!( - core_options.merge_engine(), - Ok(crate::spec::MergeEngine::FirstRow) - ); + // Deletion-vector tables without merge-on-read and ordinary first-row scans + // read without merging (stale rows are masked by DVs / level-0 is skipped), + // so they keep plain size-based packing. Audit scans merge every visible + // primary-key version, so they must keep overlapping ranges together. + let read_merges_overlapping_keys = self.merge_key_overlaps + || ((!core_options.deletion_vectors_enabled() + || core_options.deletion_vectors_merge_on_read()) + && !matches!( + core_options.merge_engine(), + Ok(crate::spec::MergeEngine::FirstRow) + )); let pk_comparator = if read_merges_overlapping_keys { KeyComparator::from_table_schema(self.table.schema()) } else { @@ -2056,9 +2076,9 @@ impl<'a> PaimonTableScan<'a> { // Java MergeTreeSplitGenerator#splitForBatch). Only engines // whose writer deduplicates at flush guarantee a file never // holds two rows of one key, so only they may mark groups raw - // convertible; see merge_tree_split_for_batch. (First-row - // tables do not take this path today, but its writer dedups - // too, so keep the gate accurate.) + // convertible; see merge_tree_split_for_batch. Ordinary first-row + // scans do not take this path, but audit scans do. Its writer + // deduplicates at flush, so keep the gate accurate. let file_keys_unique = matches!( core_options.merge_engine(), Ok(crate::spec::MergeEngine::Deduplicate) @@ -2174,10 +2194,10 @@ mod tests { use crate::io::FileIOBuilder; use crate::spec::{ stats::BinaryTableStats, ArrayType, BinaryRow, BinaryRowBuilder, BucketFunctionType, - ColumnMove, CommitKind, DataField, DataFileMeta, DataType, Datum, DeletionVectorMeta, - FileKind, GlobalIndexMeta, IndexFileMeta, IndexManifestEntry, IntType, ManifestEntry, - ManifestFileMeta, Predicate, PredicateBuilder, PredicateOperator, Schema as PaimonSchema, - SchemaChange, Snapshot, TableSchema, VarCharType, + ColumnMove, CommitKind, CoreOptions, DataField, DataFileMeta, DataType, Datum, + DeletionVectorMeta, FileKind, GlobalIndexMeta, IndexFileMeta, IndexManifestEntry, IntType, + ManifestEntry, ManifestFileMeta, Predicate, PredicateBuilder, PredicateOperator, + Schema as PaimonSchema, SchemaChange, Snapshot, TableSchema, VarCharType, }; use crate::table::bucket_filter::{compute_target_buckets, extract_predicate_for_keys}; use crate::table::partition_filter::PartitionFilter; @@ -2871,6 +2891,35 @@ mod tests { )); } + #[test] + fn test_audit_scan_all_files_preserves_data_evolution_projection() { + let table = data_evolution_test_table( + "memory:/de_audit_scan_projection", + two_column_schema(0, "id", "name"), + ) + .copy_with_options(HashMap::from([( + "global-index.enabled".to_string(), + "true".to_string(), + )])); + let projected = HashSet::from([1]); + let predicate = PredicateBuilder::new(table.schema().fields()) + .equal("id", Datum::Int(1)) + .unwrap(); + let scan = PaimonTableScan::new(&table, None, vec![predicate], None, None, None) + .with_projected_read_field_ids(Some(projected.clone())) + .with_all_versions(); + + assert!(scan.scan_all_files); + assert!(scan.merge_key_overlaps); + assert_eq!(scan.projected_read_field_ids, Some(projected)); + assert!( + scan.global_index_scan_settings(&CoreOptions::new(table.schema().options()), true,) + .unwrap() + .is_none(), + "audit scans must not prune physical row versions via global indexes" + ); + } + #[test] fn test_dv_merge_on_read_controls_batch_level_zero_visibility() { assert!(should_skip_level_zero_for_scan( @@ -3719,12 +3768,68 @@ mod tests { ); } - /// `merge-engine=first-row` PK tables read raw (no merge on the read - /// path: planned with `skip_level_zero`, read via `DataFileReader`), so - /// pruning a file by a non-key conjunct cannot resurrect anything — it - /// drops exactly the rows the raw path's exact residual filter would - /// drop. The key-only gate must exempt first-row and keep full-predicate - /// stats pruning, matching the split-generation path. + #[tokio::test] + async fn test_dv_without_mor_audit_stats_pruning_ignores_non_key_conjuncts() { + let table_path = "memory:/test_dv_audit_stats_gate"; + let table = pk_stats_gate_table(table_path).copy_with_options(HashMap::from([ + ("deletion-vectors.enabled".to_string(), "true".to_string()), + ( + "deletion-vectors.merge-on-read".to_string(), + "false".to_string(), + ), + ])); + setup_scan_trace_dirs(&table).await; + + let mut old = pk_stats_file("old-version.parquet", (1, 5), (100, 200)); + old.level = 1; + let mut new = pk_stats_file("new-version.parquet", (1, 5), (10, 60)); + new.level = 1; + TableCommit::new(table.clone(), "dv-audit-gate-test".to_string()) + .commit(vec![CommitMessage::new( + BinaryRowBuilder::new(0).build_serialized(), + 0, + vec![old, new], + )]) + .await + .unwrap(); + + let fields = vec![ + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new(1, "value".to_string(), DataType::Int(IntType::new())), + ]; + let value_filter = PredicateBuilder::new(&fields) + .greater_than("value", Datum::Int(90)) + .unwrap(); + let mut reader = table.new_read_builder(); + reader.with_filter(value_filter); + + let (ordinary_plan, ordinary_trace) = reader.new_scan().plan_with_trace().await.unwrap(); + assert!(ordinary_trace.manifest_entries_pruned_by_data_stats >= 1); + assert_eq!( + ordinary_plan + .splits() + .iter() + .map(|split| split.data_files().len()) + .sum::(), + 1 + ); + + let (audit_plan, audit_trace) = reader.new_audit_scan().plan_with_trace().await.unwrap(); + assert_eq!(audit_trace.manifest_entries_pruned_by_data_stats, 0); + assert_eq!( + audit_plan + .splits() + .iter() + .map(|split| split.data_files().len()) + .sum::(), + 2, + "both key versions must reach the audit merge path" + ); + } + + /// Ordinary `merge-engine=first-row` reads skip level-0 files and read raw, + /// so full-predicate stats pruning stays safe. A scan of all files retains + /// level-0 versions for audit merging and must use key-only pruning. #[tokio::test] async fn test_first_row_table_stats_pruning_keeps_non_key_conjuncts() { let table_path = "memory:/test_first_row_stats_gate"; @@ -3782,6 +3887,18 @@ mod tests { planned_files, 1, "only the value-matching file should be planned on first-row" ); + + let (audit_plan, audit_trace) = reader.new_audit_scan().plan_with_trace().await.unwrap(); + assert_eq!(audit_trace.manifest_entries_pruned_by_data_stats, 0); + assert_eq!( + audit_plan + .splits() + .iter() + .map(|split| split.data_files().len()) + .sum::(), + 2, + "all versions must reach the first-row audit merge" + ); } #[tokio::test] diff --git a/crates/paimon/tests/audit_log_table_test.rs b/crates/paimon/tests/audit_log_table_test.rs index 662ccaa4f..6edaa0127 100644 --- a/crates/paimon/tests/audit_log_table_test.rs +++ b/crates/paimon/tests/audit_log_table_test.rs @@ -20,8 +20,8 @@ mod common; use arrow_array::{Array, Int32Array, Int64Array, RecordBatch, StringArray}; use futures::TryStreamExt; use paimon::spec::{ - DataType, IntType, Schema, TableSchema, VarCharType, ROW_KIND_FIELD_ID, ROW_KIND_FIELD_NAME, - SEQUENCE_NUMBER_FIELD_NAME, + BigIntType, DataField, DataType, IntType, Schema, TableSchema, VarCharType, ROW_KIND_FIELD_ID, + ROW_KIND_FIELD_NAME, SEQUENCE_NUMBER_FIELD_ID, SEQUENCE_NUMBER_FIELD_NAME, }; use paimon::table::{AuditLogTable, IncrementalPlan, IncrementalScanMode, IncrementalSplit}; @@ -320,6 +320,99 @@ async fn audit_log_exposes_sequence_number_when_enabled() { assert!(rows.iter().all(|(_, seq, _, _)| *seq >= 0)); } +#[tokio::test] +async fn ordinary_read_projection_keeps_sequence_number() { + let table_path = "memory:/audit_log/ordinary_sequence_projection"; + let (file_io, table) = memory_table( + table_path, + pk_schema(&[ + ("merge-engine", "deduplicate"), + ("bucket", "1"), + ("table-read.sequence-number.enabled", "true"), + ]), + ); + setup_dirs(&file_io, table_path).await; + persist_table_schema(&file_io, table_path, table.schema()).await; + write_batch(&table, &make_batch(vec![1], vec![10])).await; + + let mut builder = table.new_read_builder(); + builder.with_read_type(vec![ + DataField::new( + SEQUENCE_NUMBER_FIELD_ID, + SEQUENCE_NUMBER_FIELD_NAME.to_string(), + DataType::BigInt(BigIntType::new()), + ), + table.schema().fields()[0].clone(), + ]); + let plan = builder.new_scan().plan().await.unwrap(); + let batches: Vec = builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + + assert_eq!( + batches[0] + .schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + vec![SEQUENCE_NUMBER_FIELD_NAME, "id"] + ); + assert_eq!( + batches[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 1 + ); +} + +#[tokio::test] +async fn audit_log_current_scan_keeps_delete_and_sequence_number() { + let table_path = "memory:/audit_log/current_state"; + let (file_io, table) = memory_table( + table_path, + pk_schema(&[ + ("changelog-producer", "input"), + ("merge-engine", "deduplicate"), + ("bucket", "1"), + ("table-read.sequence-number.enabled", "true"), + ]), + ); + setup_dirs(&file_io, table_path).await; + persist_table_schema(&file_io, table_path, table.schema()).await; + + write_batch(&table, &make_batch(vec![1, 2], vec![10, 20])).await; + let builder = table.new_write_builder(); + let mut write = builder.new_write().unwrap(); + write + .write_arrow_batch(&make_batch_with_kinds(vec![1, 2], vec![10, 25], vec![3, 2])) + .await + .unwrap(); + let messages = write.prepare_commit().await.unwrap(); + builder.new_commit().commit(messages).await.unwrap(); + + let plan = table.new_read_builder().new_scan().plan().await.unwrap(); + let batches: Vec = AuditLogTable::new(table) + .to_arrow_for_splits(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + + assert_eq!( + collect_audit_rows_with_sequence(&batches), + vec![("+U".to_string(), 3, 2, 25), ("-D".to_string(), 2, 1, 10),] + ); +} + async fn audit_diff_rows( table: &paimon::table::Table, start: i64, @@ -356,6 +449,138 @@ fn assert_rows_exclude(rows: &[(String, i32, i32)], excluded: &[(&str, i32, i32) } } +#[tokio::test] +async fn audit_log_current_scan_uses_merged_rowkind() { + for merge_engine in ["partial-update", "aggregation"] { + let table_path = format!("memory:/audit_log/current_{merge_engine}"); + let (file_io, table) = memory_table( + &table_path, + pk_schema(&[("merge-engine", merge_engine), ("bucket", "1")]), + ); + setup_dirs(&file_io, &table_path).await; + persist_table_schema(&file_io, &table_path, table.schema()).await; + + write_batch(&table, &make_batch(vec![1], vec![10])).await; + let builder = table.new_write_builder(); + let mut write = builder.new_write().unwrap(); + write + .write_arrow_batch(&make_batch_with_kinds(vec![1], vec![20], vec![2])) + .await + .unwrap(); + let messages = write.prepare_commit().await.unwrap(); + builder.new_commit().commit(messages).await.unwrap(); + + let plan = table.new_read_builder().new_scan().plan().await.unwrap(); + let batches: Vec = AuditLogTable::new(table) + .to_arrow_for_splits(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!( + collect_audit_rows(&batches), + vec![("+I".to_string(), 1, 20)], + "merge-engine={merge_engine}" + ); + } +} + +#[tokio::test] +async fn audit_log_current_scan_respects_ignore_delete() { + let table_path = "memory:/audit_log/current_ignore_delete"; + let (file_io, table) = memory_table( + table_path, + pk_schema(&[ + ("merge-engine", "deduplicate"), + ("ignore-delete", "true"), + ("bucket", "1"), + ]), + ); + setup_dirs(&file_io, table_path).await; + persist_table_schema(&file_io, table_path, table.schema()).await; + + write_batch(&table, &make_batch(vec![1], vec![10])).await; + let builder = table.new_write_builder(); + let mut write = builder.new_write().unwrap(); + write + .write_arrow_batch(&make_batch_with_kinds(vec![1], vec![10], vec![3])) + .await + .unwrap(); + let messages = write.prepare_commit().await.unwrap(); + builder.new_commit().commit(messages).await.unwrap(); + + let plan = table.new_read_builder().new_scan().plan().await.unwrap(); + let batches: Vec = AuditLogTable::new(table) + .to_arrow_for_splits(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + + assert_eq!( + collect_audit_rows(&batches), + vec![("+I".to_string(), 1, 10)] + ); +} + +#[tokio::test] +async fn audit_log_current_scan_supports_first_row() { + let table_path = "memory:/audit_log/current_first_row"; + let (file_io, table) = memory_table( + table_path, + pk_schema(&[ + ("merge-engine", "first-row"), + ("bucket", "1"), + ("source.split.target-size", "1b"), + ("source.split.open-file-cost", "1b"), + ]), + ); + setup_dirs(&file_io, table_path).await; + persist_table_schema(&file_io, table_path, table.schema()).await; + + write_batch(&table, &make_batch(vec![1], vec![10])).await; + write_batch(&table, &make_batch(vec![3], vec![30])).await; + write_batch(&table, &make_batch(vec![1], vec![20])).await; + + let mut limited_reader = table.new_read_builder(); + limited_reader.with_limit(1); + let limited_plan = limited_reader.new_audit_scan().plan().await.unwrap(); + assert_eq!( + limited_plan.splits().len(), + 2, + "audit LIMIT must not discard files needed to resolve row versions" + ); + assert_eq!( + limited_plan + .splits() + .iter() + .map(|split| split.data_files().len()) + .sum::(), + 3, + "audit LIMIT must retain every physical row version" + ); + + let audit = AuditLogTable::new(table); + let plan = audit.new_scan().plan().await.unwrap(); + assert_eq!(plan.splits().len(), 2); + let mut rows = Vec::new(); + for split in plan.splits() { + let batches: Vec = audit + .to_arrow_for_splits(std::slice::from_ref(split)) + .unwrap() + .try_collect() + .await + .unwrap(); + rows.extend(collect_audit_rows(&batches)); + } + rows.sort_unstable(); + + assert_eq!( + rows, + vec![("+I".to_string(), 1, 10), ("+I".to_string(), 3, 30)] + ); +} + #[tokio::test] async fn audit_log_diff_scan_emits_row_level_delete_insert_and_updates() { let table_path = "memory:/audit_log/diff_range"; diff --git a/docs/src/sql.md b/docs/src/sql.md index c37ac51cd..6eeb46025 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -1863,7 +1863,22 @@ let df = ctx.sql("SELECT * FROM paimon.my_db.table_a JOIN paimon.my_db.table_b O ## System Tables -Access table metadata via the `$` syntax. +Access table metadata and audit rows via the `$` syntax. + +### $audit_log + +Read the current table state with each row's Paimon row kind (`+I`, `-U`, `+U`, or `-D`): + +```sql +SELECT * FROM paimon.default.my_table$audit_log; +``` + +`rowkind` is the first column, followed by the table columns. Append-only rows are +reported as `+I`; deduplicate primary-key reads retain the latest physical retract row +instead of dropping it. Other primary-key merge engines retain or reject retracts +according to their merge-engine options. As in Paimon Java, rows masked by deletion +vectors are not reconstructed as retract records. When +`table-read.sequence-number.enabled=true`, `_SEQUENCE_NUMBER` appears after `rowkind`. ### $options