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/scan.rs b/crates/integrations/datafusion/src/physical_plan/scan.rs index fe0a38590..d1cff261c 100644 --- a/crates/integrations/datafusion/src/physical_plan/scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/scan.rs @@ -15,12 +15,13 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::sync::Arc; use datafusion::arrow::array::BooleanArray; use datafusion::arrow::compute::{cast, filter_record_batch}; use datafusion::arrow::datatypes::{ - DataType as ArrowDataType, SchemaRef as ArrowSchemaRef, TimeUnit, + DataType as ArrowDataType, Schema, SchemaRef as ArrowSchemaRef, TimeUnit, }; use datafusion::arrow::record_batch::{RecordBatch, RecordBatchOptions}; use datafusion::common::stats::Precision; @@ -50,13 +51,63 @@ use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{DisplayAs, ExecutionPlan, Partitioning, PlanProperties}; use futures::{FutureExt, StreamExt, TryStreamExt}; use paimon::arrow::ParquetReadBudget; -use paimon::spec::{DataField, Datum, MergeEngine, Predicate, PredicateBuilder, PredicateOperator}; +use paimon::spec::{ + DataField, Datum, MergeEngine, Predicate, PredicateBuilder, PredicateOperator, + ROW_KIND_FIELD_NAME, SEQUENCE_NUMBER_FIELD_NAME, +}; use paimon::table::{ScanTrace, Table}; use paimon::DataSplit; use crate::error::to_datafusion_error; use crate::filter_pushdown::scalar_to_datum; +struct AuditProjection { + indices: Vec, + schema: ArrowSchemaRef, +} + +fn audit_projection(batch: &RecordBatch, schema: &ArrowSchemaRef) -> DFResult { + let batch_schema = batch.schema(); + let by_name: HashMap<&str, usize> = batch_schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| (field.name().as_str(), index)) + .collect(); + let indices = schema + .fields() + .iter() + .map(|field| { + by_name.get(field.name().as_str()).copied().ok_or_else(|| { + datafusion::error::DataFusionError::Execution(format!( + "Audit log reader did not return projected column '{}'", + field.name() + )) + }) + }) + .collect::>>()?; + let fields = indices + .iter() + .map(|&index| batch.schema().field(index).clone()) + .collect::>(); + Ok(AuditProjection { + indices, + schema: Arc::new(Schema::new(fields)), + }) +} + +fn project_audit_batch(batch: RecordBatch, projection: &AuditProjection) -> DFResult { + let row_count = batch.num_rows(); + let columns = projection + .indices + .iter() + .map(|&index| batch.column(index).clone()) + .collect(); + let options = RecordBatchOptions::new().with_row_count(Some(row_count)); + RecordBatch::try_new_with_options(projection.schema.clone(), columns, &options) + .map_err(Into::into) +} + fn to_datafusion_batch(batch: RecordBatch, schema: &ArrowSchemaRef) -> DFResult { if batch.num_columns() != schema.fields().len() { return Err(datafusion::error::DataFusionError::Execution(format!( @@ -778,6 +829,8 @@ pub struct PaimonTableScan { decoder_filters: Vec>, /// Query-wide budget shared by every DataFusion scan partition. parquet_read_budget: Arc, + /// Retain retract rows and expose their row kind through `$audit_log`. + audit_log: bool, } impl PaimonTableScan { @@ -884,9 +937,37 @@ impl PaimonTableScan { runtime_filters: Vec::new(), decoder_filters: Vec::new(), parquet_read_budget, + audit_log: false, } } + #[allow(clippy::too_many_arguments)] + pub(crate) fn try_new_audit_log( + schema: ArrowSchemaRef, + table: Table, + read_type: Vec, + pushed_predicate: Option, + planned_partitions: Vec>, + limit: Option, + scan_trace: Option, + case_sensitive: bool, + ) -> DFResult { + let mut scan = Self::try_new( + schema, + table, + read_type, + pushed_predicate, + planned_partitions, + limit, + false, + scan_trace, + None, + case_sensitive, + )?; + scan.audit_log = true; + Ok(scan) + } + pub fn table(&self) -> &Table { &self.table } @@ -973,7 +1054,11 @@ impl PaimonTableScan { impl ExecutionPlan for PaimonTableScan { fn name(&self) -> &str { - "PaimonTableScan" + if self.audit_log { + "PaimonAuditLogScan" + } else { + "PaimonTableScan" + } } fn properties(&self) -> &Arc { @@ -1007,13 +1092,23 @@ 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()) { + let physical_columns_available = !self.audit_log + || collect_columns(&filter).iter().all(|column| { + resolve_physical_field( + column.name(), + self.table.schema().fields(), + self.case_sensitive, + ) + .is_some() + }); + if physical_columns_available + && 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. @@ -1072,6 +1167,7 @@ impl ExecutionPlan for PaimonTableScan { let runtime_filters = self.runtime_filters.clone(); let decoder_filters = self.decoder_filters.clone(); let parquet_read_budget = Arc::clone(&self.parquet_read_budget); + let audit_log = self.audit_log; let fut = async move { let mut read_builder = table.new_read_builder(); @@ -1098,12 +1194,35 @@ impl ExecutionPlan for PaimonTableScan { Arc::clone(&schema), ))); } - let stream = read.to_arrow(&splits).map_err(to_datafusion_error)?; + let stream = if audit_log { + read.to_projected_audit_log_arrow_for_splits( + &splits, + schema + .fields() + .iter() + .any(|field| field.name() == ROW_KIND_FIELD_NAME), + schema + .fields() + .iter() + .any(|field| field.name() == SEQUENCE_NUMBER_FIELD_NAME), + ) + } else { + read.to_arrow(&splits) + } + .map_err(to_datafusion_error)?; let batch_schema = Arc::clone(&schema); + let mut cached_audit_projection = None; 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 batch = if audit_log { + if cached_audit_projection.is_none() { + cached_audit_projection = Some(audit_projection(&batch, &batch_schema)?); + } + project_audit_batch(batch, cached_audit_projection.as_ref().unwrap())? + } else { + batch + }; + 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 @@ -1159,7 +1278,9 @@ impl ExecutionPlan for PaimonTableScan { // 1. All splits have known merged_row_count (no deletion files with unknown cardinality) // 2. No limit is applied (limit would make row count inexact) // 3. Filter is exact (no residual filtering needed above the scan) - let num_rows_precision = if all_row_counts_known + let num_rows_precision = if self.audit_log { + Precision::Absent + } else if all_row_counts_known && self.limit.is_none() && self.filter_exact && self.runtime_filters.is_empty() @@ -1183,7 +1304,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())?; + write!(f, "{}: table={}", self.name(), self.table.identifier())?; let total_splits: usize = self.planned_partitions.iter().map(|p| p.len()).sum(); let total_files: usize = self 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..876d413b6 --- /dev/null +++ b/crates/integrations/datafusion/src/system_tables/audit_log.rs @@ -0,0 +1,130 @@ +// 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::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_scan().plan_with_trace()) + .await + .map_err(to_datafusion_error)?; + + 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_audit_log(self.fields.clone()) + } + + 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 392db7da5..a0201443f 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 files; mod manifests; @@ -48,6 +50,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), ("files", files::build), ("manifests", manifests::build), @@ -61,6 +64,7 @@ const TABLES: &[(&str, Builder)] = &[ ]; const SYSTEM_TABLE_NAMES: &[&str] = &[ + "audit_log", "branches", "files", "manifests", @@ -74,6 +78,16 @@ const SYSTEM_TABLE_NAMES: &[&str] = &[ "tags", ]; +// Matches Java SystemTableLoader's physical-metadata restriction. Audit log is +// also rejected until Rust can apply the base table's row filters and masks. +const QUERY_AUTH_UNSUPPORTED_TABLES: &[&str] = &[ + "audit_log", + "files", + "file_key_ranges", + "binlog", + "statistics", +]; + /// 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). @@ -101,6 +115,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, @@ -110,10 +139,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); } @@ -133,13 +160,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 @@ -148,6 +187,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!( @@ -186,6 +231,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..d6c7dfaa6 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -356,19 +356,42 @@ impl PaimonScanBuilder<'_> { self, read_fields: Vec, ) -> DFResult> { - let (projected_schema, read_type) = if let Some(indices) = self.projection { + self.build_scan(read_fields, false) + } + + pub(crate) fn build_audit_log( + self, + audit_fields: Vec, + ) -> DFResult> { + self.build_scan(audit_fields, true) + } + + fn build_scan( + self, + read_fields: Vec, + audit_log: bool, + ) -> DFResult> { + let (projected_schema, mut 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) }; + if audit_log { + read_type.retain(|field| { + !matches!( + field.id(), + paimon::spec::ROW_KIND_FIELD_ID | paimon::spec::SEQUENCE_NUMBER_FIELD_ID + ) + }); + } let splits = self.plan.into_splits(); let planned_partitions: Vec> = if splits.is_empty() { @@ -381,18 +404,31 @@ impl PaimonScanBuilder<'_> { .collect() }; - Ok(Arc::new(PaimonTableScan::try_new( - projected_schema, - self.table.clone(), - read_type, - self.pushed_predicate, - planned_partitions, - self.limit, - self.filter_exact, - self.scan_trace, - None, - self.case_sensitive, - )?)) + if audit_log { + Ok(Arc::new(PaimonTableScan::try_new_audit_log( + projected_schema, + self.table.clone(), + read_type, + self.pushed_predicate, + planned_partitions, + self.limit, + self.scan_trace, + self.case_sensitive, + )?)) + } else { + Ok(Arc::new(PaimonTableScan::try_new( + projected_schema, + self.table.clone(), + read_type, + self.pushed_predicate, + planned_partitions, + self.limit, + self.filter_exact, + self.scan_trace, + None, + self.case_sensitive, + )?)) + } } } diff --git a/crates/integrations/datafusion/tests/system_tables.rs b/crates/integrations/datafusion/tests/system_tables.rs index 582d6ef10..ef53444c2 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,7 +88,7 @@ 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_table_policy_matches_java() { let (ctx, _catalog, _tmp) = create_context().await; run_sql( &ctx, @@ -93,18 +96,282 @@ async fn test_query_auth_table_fails_closed() { ) .await; - // Data reads and data-derived system tables must all fail closed. + // Rust cannot yet apply row filters or masks to table data, and raw file + // statistics cannot be masked. Both paths must fail closed. for sql in [ "SELECT * FROM paimon.default.qa", + "SELECT * FROM paimon.default.qa$audit_log", + "SELECT * FROM paimon.default.qa$files", + ] { + let err = query_error(&ctx, sql).await; + assert!( + err.contains("query-auth.enabled"), + "`{sql}` should fail closed, got: {err}" + ); + } + + // Match Java SystemTableLoader: schema and non-physical metadata remain readable. + let batches = run_sql( + &ctx, + "SELECT value FROM paimon.default.qa$options WHERE key = 'query-auth.enabled'", + ) + .await; + assert_eq!(string_value(batches[0].column(0).as_ref(), 0), "true"); + for sql in [ "SELECT * FROM paimon.default.qa$manifests", "SELECT * FROM paimon.default.qa$table_indexes", + ] { + run_sql(&ctx, sql).await; + } + + 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", ] { let err = query_error(&ctx, sql).await; assert!( err.contains("query-auth.enabled"), - "`{sql}` should fail closed, got: {err}" + "dynamic auth should make `{sql}` fail closed, got: {err}" ); } + run_sql(&ctx, "SELECT * FROM paimon.default.qa_dynamic$options").await; + 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_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..099779de0 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, DataSplit, Table}; 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( @@ -85,4 +82,15 @@ impl AuditLogTable { let read = self.wrapped.new_read_builder().new_read()?; read.to_audit_log_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.wrapped + .new_read_builder() + .new_read()? + .to_audit_log_arrow_for_splits(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/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 99c01f7d9..5309bdb80 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -39,6 +39,7 @@ use arrow_select::concat::concat as arrow_concat; use arrow_select::take::take; use futures::{stream, StreamExt}; use std::cmp::Ordering; +use std::collections::HashMap; use std::sync::Arc; const MAX_MERGE_INPUT_STREAMS: usize = 256; @@ -213,6 +214,40 @@ impl<'a> TableRead<'a> { } } + /// Returns the current table state as audit-log rows for planned data splits. + pub fn to_audit_log_arrow_for_splits( + &self, + data_splits: &[DataSplit], + ) -> crate::Result { + self.ensure_query_auth_allowed()?; + match &self.0 { + TableReadKind::Paimon(read) => read.to_audit_log_arrow_for_splits(data_splits), + TableReadKind::Format(_) => Err(crate::Error::Unsupported { + message: "Format tables do not support audit log batch read".to_string(), + }), + } + } + + /// As [`Self::to_audit_log_arrow_for_splits`], omitting unrequested system columns. + pub fn to_projected_audit_log_arrow_for_splits( + &self, + data_splits: &[DataSplit], + include_rowkind: bool, + include_sequence: bool, + ) -> crate::Result { + self.ensure_query_auth_allowed()?; + match &self.0 { + TableReadKind::Paimon(read) => read.to_projected_audit_log_arrow_for_splits( + data_splits, + include_rowkind, + include_sequence, + ), + 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() } @@ -357,6 +392,127 @@ impl<'a> PaimonTableRead<'a> { })) } + /// Returns the current table state as audit-log rows. + pub fn to_audit_log_arrow_for_splits( + &self, + data_splits: &[DataSplit], + ) -> crate::Result { + self.to_projected_audit_log_arrow_for_splits( + data_splits, + true, + audit_sequence_number_enabled(self.table), + ) + } + + /// Returns projected current-state audit rows without materializing omitted system columns. + pub fn to_projected_audit_log_arrow_for_splits( + &self, + data_splits: &[DataSplit], + include_rowkind: bool, + include_sequence: bool, + ) -> crate::Result { + if include_sequence && !audit_sequence_number_enabled(self.table) { + return Err(crate::Error::DataInvalid { + message: "Audit read requested _SEQUENCE_NUMBER but table-read.sequence-number.enabled is false".to_string(), + source: None, + }); + } + let user_read_type = self.read_type.clone(); + let audit_schema = + audit_schema_for_read_type(&user_read_type, include_rowkind, include_sequence)?; + let has_primary_keys = !self.table.schema().primary_keys().is_empty(); + + let physical_stream = if has_primary_keys { + let core_options = self.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.parquet_read_budget()?; + let raw_stream = DataFileReader::new( + self.table.file_io.clone(), + self.table.schema_manager().clone(), + self.table.schema().id(), + self.table.schema.fields().to_vec(), + read_type.clone(), + self.data_predicates.clone(), + ) + .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.table.file_io.clone(), + KeyValueReadConfig { + table_name: self.table.identifier().full_name(), + table_options: self.table.schema().options().clone(), + schema_manager: self.table.schema_manager().clone(), + table_schema_id: self.table.schema().id(), + table_fields: self.table.schema.fields().to_vec(), + read_type, + predicates: self.data_predicates.clone(), + primary_keys: self.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.to_arrow(data_splits)? + }; + + Ok(audit_stream_from_physical( + physical_stream, + audit_schema, + user_read_type, + include_rowkind, + include_sequence, + has_primary_keys && include_rowkind, + )) + } + /// Returns an audit-log stream for a planned incremental scan. pub fn to_audit_log_arrow( &self, @@ -385,7 +541,7 @@ impl<'a> PaimonTableRead<'a> { 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 audit_schema = audit_schema_for_read_type(&user_read_type, true, include_sequence)?; let mut read_type = user_read_type.clone(); if include_sequence { @@ -417,54 +573,14 @@ impl<'a> PaimonTableRead<'a> { .with_batch_size(Some(self.table.schema().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)), - } - })?; - } - })) + Ok(audit_stream_from_physical( + raw_stream, + audit_schema, + user_read_type, + true, + include_sequence, + has_value_kind, + )) } fn audit_diff_stream(&self, plan: &IncrementalPlan) -> crate::Result { @@ -505,7 +621,7 @@ impl<'a> PaimonTableRead<'a> { 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 audit_schema = audit_schema_for_read_type(&self.read_type, true, include_sequence)?; let mut diff_read_type = self.table.schema().fields().to_vec(); ensure_diff_supported_read_type(&diff_read_type)?; @@ -701,6 +817,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 @@ -842,6 +959,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()) @@ -906,16 +1024,158 @@ impl<'a> PaimonTableRead<'a> { } } +// 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 audit_schema_for_read_type( read_type: &[DataField], + include_rowkind: bool, 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_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, @@ -930,9 +1190,8 @@ fn audit_schema_for_read_type( 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")) + .core_options() + .table_read_sequence_number_enabled() } fn rowkind_array_from_column(column: &dyn arrow_array::Array) -> crate::Result { @@ -1557,6 +1816,20 @@ mod tests { let legacy = split(vec![file("a", 5, None)], true); assert!(pk_split_needs_merge(&legacy, false)); + 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)); + // Deletion-vector tables dispatch on level 0 only. let dv_l0 = split(vec![file("a", 0, None)], false); assert!(pk_split_needs_merge(&dv_l0, true)); diff --git a/crates/paimon/tests/audit_log_table_test.rs b/crates/paimon/tests/audit_log_table_test.rs index 662ccaa4f..63237a6cc 100644 --- a/crates/paimon/tests/audit_log_table_test.rs +++ b/crates/paimon/tests/audit_log_table_test.rs @@ -320,6 +320,45 @@ async fn audit_log_exposes_sequence_number_when_enabled() { assert!(rows.iter().all(|(_, seq, _, _)| *seq >= 0)); } +#[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 +395,119 @@ 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![1], vec![20])).await; + + let plan = table + .new_read_builder() + .new_scan() + .with_scan_all_files() + .plan() + .await + .unwrap(); + assert_eq!(plan.splits().len(), 2); + 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_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 c53bacbf5..a680c1cf6 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