From 8fc48aceb37f8c281248dd6da19a66ab1c287a32 Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 3 Sep 2026 13:04:07 +0800 Subject: [PATCH 01/15] feat: support audit log system table --- crates/integrations/datafusion/src/catalog.rs | 9 + .../datafusion/src/physical_plan/scan.rs | 143 ++++++- .../datafusion/src/system_tables/audit_log.rs | 130 ++++++ .../datafusion/src/system_tables/mod.rs | 17 + .../integrations/datafusion/src/table/mod.rs | 68 ++- .../datafusion/tests/system_tables.rs | 223 +++++++++- crates/paimon/src/table/audit_log_table.rs | 20 +- crates/paimon/src/table/kv_file_reader.rs | 100 +++-- crates/paimon/src/table/sort_merge.rs | 173 ++++++-- crates/paimon/src/table/table_read.rs | 389 +++++++++++++++--- crates/paimon/tests/audit_log_table_test.rs | 152 +++++++ docs/src/sql.md | 17 +- 12 files changed, 1290 insertions(+), 151 deletions(-) create mode 100644 crates/integrations/datafusion/src/system_tables/audit_log.rs diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index 340cf07c2..17cebb70d 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -693,11 +693,20 @@ 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(); + paimon::spec::CoreOptions::new(&dynamic_options) + .ensure_read_authorized() + .map_err(to_datafusion_error)?; 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..041bf667b 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", @@ -133,6 +137,7 @@ pub(crate) async fn load( database: String, object: ParsedObjectName, system_name: String, + dynamic_options: HashMap, ) -> DFResult>> { if !is_registered(&system_name) { return Ok(None); @@ -140,6 +145,9 @@ pub(crate) async fn load( let identifier = Identifier::new(database, object.table().to_string()); match catalog.get_table(&identifier).await { Ok(mut table) => { + paimon::spec::CoreOptions::new(table.schema().options()) + .ensure_read_authorized() + .map_err(to_datafusion_error)?; if let Some(branch) = object.branch() { if !system_name.eq_ignore_ascii_case("branches") { table = table @@ -148,6 +156,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 +200,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..92812cd55 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) { @@ -96,6 +99,7 @@ async fn test_query_auth_table_fails_closed() { // Data reads and data-derived system tables must all fail closed. for sql in [ "SELECT * FROM paimon.default.qa", + "SELECT * FROM paimon.default.qa$audit_log", "SELECT * FROM paimon.default.qa$manifests", "SELECT * FROM paimon.default.qa$table_indexes", ] { @@ -105,6 +109,223 @@ 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", + ] { + 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_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 From f50fa9a59d8fb094171242a3124bd6f089d375e9 Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 3 Sep 2026 15:01:27 +0800 Subject: [PATCH 02/15] fix(datafusion): align audit dynamic options with Java --- .../datafusion/src/system_tables/mod.rs | 8 +++++ .../datafusion/tests/system_tables.rs | 30 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/crates/integrations/datafusion/src/system_tables/mod.rs b/crates/integrations/datafusion/src/system_tables/mod.rs index 041bf667b..f3419c52d 100644 --- a/crates/integrations/datafusion/src/system_tables/mod.rs +++ b/crates/integrations/datafusion/src/system_tables/mod.rs @@ -142,6 +142,14 @@ pub(crate) async fn load( 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(), + )); + } let identifier = Identifier::new(database, object.table().to_string()); match catalog.get_table(&identifier).await { Ok(mut table) => { diff --git a/crates/integrations/datafusion/tests/system_tables.rs b/crates/integrations/datafusion/tests/system_tables.rs index 92812cd55..5c2211f62 100644 --- a/crates/integrations/datafusion/tests/system_tables.rs +++ b/crates/integrations/datafusion/tests/system_tables.rs @@ -135,6 +135,36 @@ async fn test_query_auth_table_fails_closed() { 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; From 4376087b6be7031c50944e4481bae233a7e279d6 Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 3 Sep 2026 17:13:37 +0800 Subject: [PATCH 03/15] fix(datafusion): align system table authorization with Java --- crates/integrations/datafusion/src/catalog.rs | 3 -- .../datafusion/src/system_tables/mod.rs | 35 +++++++++++++++---- .../datafusion/tests/system_tables.rs | 24 ++++++++++--- 3 files changed, 48 insertions(+), 14 deletions(-) diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index 17cebb70d..06f34c3c4 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -698,9 +698,6 @@ impl SchemaProvider for PaimonSchemaProvider { .read() .unwrap_or_else(|e| e.into_inner()) .clone(); - paimon::spec::CoreOptions::new(&dynamic_options) - .ensure_read_authorized() - .map_err(to_datafusion_error)?; return await_with_runtime(system_tables::load( Arc::clone(&self.catalog), self.database.clone(), diff --git a/crates/integrations/datafusion/src/system_tables/mod.rs b/crates/integrations/datafusion/src/system_tables/mod.rs index f3419c52d..e05682470 100644 --- a/crates/integrations/datafusion/src/system_tables/mod.rs +++ b/crates/integrations/datafusion/src/system_tables/mod.rs @@ -78,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). @@ -105,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, @@ -114,10 +139,7 @@ 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)?; + 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); } @@ -150,12 +172,11 @@ pub(crate) async fn load( .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) => { - paimon::spec::CoreOptions::new(table.schema().options()) - .ensure_read_authorized() - .map_err(to_datafusion_error)?; + 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 diff --git a/crates/integrations/datafusion/tests/system_tables.rs b/crates/integrations/datafusion/tests/system_tables.rs index 5c2211f62..ef53444c2 100644 --- a/crates/integrations/datafusion/tests/system_tables.rs +++ b/crates/integrations/datafusion/tests/system_tables.rs @@ -88,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, @@ -96,12 +96,12 @@ 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$manifests", - "SELECT * FROM paimon.default.qa$table_indexes", + "SELECT * FROM paimon.default.qa$files", ] { let err = query_error(&ctx, sql).await; assert!( @@ -110,11 +110,26 @@ async fn test_query_auth_table_fails_closed() { ); } + // 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!( @@ -122,6 +137,7 @@ async fn test_query_auth_table_fails_closed() { "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; From e16c7cba0ed2554a93b63336279bd1307f0fee0b Mon Sep 17 00:00:00 2001 From: yantian Date: Tue, 8 Sep 2026 09:37:43 +0800 Subject: [PATCH 04/15] fix(datafusion): retain system table type validation --- crates/integrations/datafusion/src/system_tables/mod.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/integrations/datafusion/src/system_tables/mod.rs b/crates/integrations/datafusion/src/system_tables/mod.rs index e05682470..a0201443f 100644 --- a/crates/integrations/datafusion/src/system_tables/mod.rs +++ b/crates/integrations/datafusion/src/system_tables/mod.rs @@ -139,6 +139,7 @@ pub(crate) fn provider_for_table( if !is_registered(system_name) { return Ok(None); } + 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); @@ -176,6 +177,7 @@ pub(crate) async fn load( 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") { From 46752b252b8305e1b68a553d4edf56d305da6f86 Mon Sep 17 00:00:00 2001 From: yantian Date: Tue, 8 Sep 2026 14:31:37 +0800 Subject: [PATCH 05/15] perf(table): batch interleave diff output --- crates/paimon/src/table/table_read.rs | 173 +++++++++++++++----------- 1 file changed, 97 insertions(+), 76 deletions(-) diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index 5309bdb80..efe486644 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -32,11 +32,9 @@ use crate::spec::{ use crate::DataSplit; use arrow_array::{ builder::StringBuilder, Array, ArrayRef, RecordBatch, RecordBatchOptions, StringArray, - UInt32Array, }; 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; @@ -1328,49 +1326,19 @@ impl AuditBatchBuilder { fn push(&mut self, kind: &str, batch: &RecordBatch, row: usize) { self.rowkind.append_value(kind); - let batch_id = self.pin_batch(batch); + let batch_id = pin_batch(&mut self.pinned_batches, 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)), - })?, - ); - } + columns.extend(interleave_columns( + &self.pinned_batches, + &self.data_col_indices, + &self.row_indices, + )?); self.row_indices.clear(); self.pinned_batches.clear(); self.len = 0; @@ -1407,49 +1375,15 @@ impl DiffAfterImageBatchBuilder { } fn push(&mut self, batch: &RecordBatch, row: usize) { - let batch_id = self.pin_batch(batch); + let batch_id = pin_batch(&mut self.pinned_batches, 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.len = 0; @@ -1463,6 +1397,43 @@ impl DiffAfterImageBatchBuilder { } } +fn pin_batch(pinned_batches: &mut Vec, batch: &RecordBatch) -> usize { + if pinned_batches.last().is_some_and(|last| { + last.num_rows() == batch.num_rows() + && last.num_columns() == batch.num_columns() + && last + .columns() + .iter() + .zip(batch.columns()) + .all(|(left, right)| Arc::ptr_eq(left, right)) + }) { + return pinned_batches.len() - 1; + } + let batch_id = pinned_batches.len(); + pinned_batches.push(batch.clone()); + batch_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 { @@ -1745,6 +1716,56 @@ mod tests { use crate::spec::{BinaryRow, DataFileMeta, DataType, IntType, Schema, TableSchema}; use crate::table::query_auth_table; use crate::table::source::DataSplitBuilder; + use arrow_array::Int32Array; + use arrow_schema::{DataType as ArrowDataType, Field}; + + #[test] + fn test_diff_batch_builders_pin_each_input_batch_once() { + let input = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![Field::new( + "id", + ArrowDataType::Int32, + false, + )])), + vec![Arc::new(Int32Array::from(vec![1, 2]))], + ) + .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", &input, 1); + audit.push("+I", &input, 0); + assert_eq!(audit.pinned_batches.len(), 1); + let audit_batch = audit.flush().unwrap(); + let audit_ids = audit_batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!((audit_ids.value(0), audit_ids.value(1)), (2, 1)); + + let mut after = DiffAfterImageBatchBuilder::new( + Arc::new(ArrowSchema::new(vec![Field::new( + "id", + ArrowDataType::Int32, + false, + )])), + vec![0], + ); + after.push(&input, 1); + after.push(&input, 0); + assert_eq!(after.pinned_batches.len(), 1); + let after_batch = after.flush().unwrap(); + let after_ids = after_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!((after_ids.value(0), after_ids.value(1)), (2, 1)); + } fn file(name: &str, level: i32, delete_row_count: Option) -> DataFileMeta { DataFileMeta { From c98dbf342743c2f1d5548a2c2617a5908cdae28a Mon Sep 17 00:00:00 2001 From: yantian Date: Tue, 8 Sep 2026 18:12:22 +0800 Subject: [PATCH 06/15] perf(table): deduplicate pinned diff batches --- crates/paimon/src/table/table_read.rs | 127 +++++++++++++++++--------- 1 file changed, 82 insertions(+), 45 deletions(-) diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index efe486644..c5d0ec637 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -652,8 +652,8 @@ impl<'a> PaimonTableRead<'a> { 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 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()); @@ -672,11 +672,11 @@ impl<'a> PaimonTableRead<'a> { } match cursor_cmp(&bc, &ac, &key_indices, &value_indices)? { CursorOrd::BeforeOnly => { - builder.push("-D", bc.batch(), bc.row()); + builder.push("-D", bc.batch_id(), bc.batch(), bc.row()); bc.advance().await?; } CursorOrd::AfterOnly => { - builder.push("+I", ac.batch(), ac.row()); + builder.push("+I", ac.batch_id(), ac.batch(), ac.row()); ac.advance().await?; } CursorOrd::EqualSame => { @@ -684,8 +684,8 @@ impl<'a> PaimonTableRead<'a> { ac.advance().await?; } CursorOrd::EqualDiff => { - builder.push("-U", bc.batch(), bc.row()); - builder.push("+U", ac.batch(), ac.row()); + 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?; } @@ -744,8 +744,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()); @@ -755,7 +755,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 => { @@ -763,7 +763,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?; } @@ -1240,14 +1240,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?; @@ -1266,6 +1270,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 { @@ -1276,6 +1284,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(()); @@ -1296,6 +1305,7 @@ struct AuditBatchBuilder { rowkind: StringBuilder, row_indices: Vec<(usize, usize)>, pinned_batches: Vec, + pinned_batch_ids: HashMap<(usize, usize), usize>, data_col_indices: Vec, len: usize, } @@ -1307,6 +1317,7 @@ impl AuditBatchBuilder { rowkind: StringBuilder::new(), row_indices: Vec::new(), pinned_batches: Vec::new(), + pinned_batch_ids: HashMap::new(), data_col_indices: Vec::new(), len: 0, } @@ -1324,9 +1335,14 @@ impl AuditBatchBuilder { self.len } - fn push(&mut self, kind: &str, batch: &RecordBatch, row: usize) { + 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, batch); + 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; } @@ -1341,6 +1357,7 @@ impl AuditBatchBuilder { )?); 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 { @@ -1355,6 +1372,7 @@ struct DiffAfterImageBatchBuilder { schema: Arc, row_indices: Vec<(usize, usize)>, pinned_batches: Vec, + pinned_batch_ids: HashMap<(usize, usize), usize>, col_indices: Vec, len: usize, } @@ -1365,6 +1383,7 @@ impl DiffAfterImageBatchBuilder { schema, row_indices: Vec::new(), pinned_batches: Vec::new(), + pinned_batch_ids: HashMap::new(), col_indices, len: 0, } @@ -1374,8 +1393,13 @@ impl DiffAfterImageBatchBuilder { self.len } - fn push(&mut self, batch: &RecordBatch, row: usize) { - let batch_id = pin_batch(&mut self.pinned_batches, 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; } @@ -1386,6 +1410,7 @@ impl DiffAfterImageBatchBuilder { 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| { @@ -1397,21 +1422,19 @@ impl DiffAfterImageBatchBuilder { } } -fn pin_batch(pinned_batches: &mut Vec, batch: &RecordBatch) -> usize { - if pinned_batches.last().is_some_and(|last| { - last.num_rows() == batch.num_rows() - && last.num_columns() == batch.num_columns() - && last - .columns() - .iter() - .zip(batch.columns()) - .all(|(left, right)| Arc::ptr_eq(left, right)) - }) { - return pinned_batches.len() - 1; +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 batch_id = pinned_batches.len(); + let pinned_id = pinned_batches.len(); pinned_batches.push(batch.clone()); - batch_id + pinned_batch_ids.insert(batch_id, pinned_id); + pinned_id } fn interleave_columns( @@ -1721,31 +1744,39 @@ mod tests { #[test] fn test_diff_batch_builders_pin_each_input_batch_once() { - let input = RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![Field::new( - "id", - ArrowDataType::Int32, - false, - )])), - vec![Arc::new(Int32Array::from(vec![1, 2]))], - ) - .unwrap(); + 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", &input, 1); - audit.push("+I", &input, 0); - assert_eq!(audit.pinned_batches.len(), 1); + 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.value(0), audit_ids.value(1)), (2, 1)); + assert_eq!( + audit_ids.values(), + &[2, 3, 1, 4], + "interleaved batches must preserve row order" + ); let mut after = DiffAfterImageBatchBuilder::new( Arc::new(ArrowSchema::new(vec![Field::new( @@ -1755,16 +1786,22 @@ mod tests { )])), vec![0], ); - after.push(&input, 1); - after.push(&input, 0); - assert_eq!(after.pinned_batches.len(), 1); + 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.value(0), after_ids.value(1)), (2, 1)); + assert_eq!( + after_ids.values(), + &[2, 3, 1, 4], + "interleaved batches must preserve row order" + ); } fn file(name: &str, level: i32, delete_row_count: Option) -> DataFileMeta { From d79327014ce96f091279ff74728cdc4474075614 Mon Sep 17 00:00:00 2001 From: yantian Date: Wed, 9 Sep 2026 16:26:26 +0800 Subject: [PATCH 07/15] fix(datafusion): harden audit log planning and auth --- .../datafusion/src/system_tables/mod.rs | 9 ++- .../integrations/datafusion/src/table/mod.rs | 73 ++++++++++++++++++- .../datafusion/tests/system_tables.rs | 36 +++++---- 3 files changed, 95 insertions(+), 23 deletions(-) diff --git a/crates/integrations/datafusion/src/system_tables/mod.rs b/crates/integrations/datafusion/src/system_tables/mod.rs index 6a832b8ef..5c23565d6 100644 --- a/crates/integrations/datafusion/src/system_tables/mod.rs +++ b/crates/integrations/datafusion/src/system_tables/mod.rs @@ -81,14 +81,19 @@ 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. +// 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. diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index d6c7dfaa6..161db1cef 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -32,7 +32,7 @@ use datafusion::logical_expr::dml::InsertOp; use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; use datafusion::physical_plan::ExecutionPlan; use paimon::spec::{ - BigIntType, CoreOptions, DataField, DataType, ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME, + BigIntType, CoreOptions, DataField, DataType, MergeEngine, ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME, }; use paimon::table::Table; @@ -393,9 +393,21 @@ impl PaimonScanBuilder<'_> { }); } + let first_row_audit = audit_log + && self + .table + .schema() + .core_options() + .merge_engine() + .map_err(to_datafusion_error)? + == MergeEngine::FirstRow; let splits = self.plan.into_splits(); let planned_partitions: Vec> = if splits.is_empty() { vec![Arc::from(Vec::new())] + } else if first_row_audit { + // ponytail: keep merge groups intact; add group-aware balancing if + // first-row audit parallelism becomes necessary. + vec![Arc::from(splits)] } else { let num_partitions = splits.len().min(self.target_partitions.max(1)); bucket_round_robin(splits, num_partitions) @@ -572,7 +584,9 @@ mod tests { use datafusion::prelude::{SessionConfig, SessionContext}; use paimon::catalog::Identifier; use paimon::spec::{ArrayType, MapType, RowType, VarCharType}; - use paimon::{Catalog, CatalogOptions, DataSplit, FileSystemCatalog, Options}; + use paimon::{ + Catalog, CatalogOptions, DataSplit, DataSplitBuilder, FileSystemCatalog, Options, + }; use crate::physical_plan::PaimonTableScan; @@ -594,6 +608,61 @@ mod tests { assert_eq!(result, vec![vec![1, 2, 3]]); } + #[test] + fn test_first_row_audit_keeps_split_group_in_one_partition() { + 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 = datafusion_read_fields(&table); + 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_audit_log(read_fields) + .unwrap(); + let scan = plan + .downcast_ref::() + .expect("Expected PaimonTableScan"); + + assert_eq!(scan.planned_partitions().len(), 1); + assert_eq!(scan.planned_partitions()[0].len(), 2); + } + fn get_test_warehouse() -> String { std::env::var("PAIMON_TEST_WAREHOUSE") .unwrap_or_else(|_| "/tmp/paimon-warehouse".to_string()) diff --git a/crates/integrations/datafusion/tests/system_tables.rs b/crates/integrations/datafusion/tests/system_tables.rs index c460e50c1..8fdd68665 100644 --- a/crates/integrations/datafusion/tests/system_tables.rs +++ b/crates/integrations/datafusion/tests/system_tables.rs @@ -88,20 +88,28 @@ async fn query_error(ctx: &SQLContext, sql: &str) -> String { } #[tokio::test] -async fn test_query_auth_system_table_policy_matches_java() { +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; - // Rust cannot yet apply row filters or masks to table data, and raw file - // statistics cannot be masked. Both paths must 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", ] { let err = query_error(&ctx, sql).await; assert!( @@ -110,26 +118,17 @@ async fn test_query_auth_system_table_policy_matches_java() { ); } - // 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", + "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!( @@ -137,7 +136,6 @@ async fn test_query_auth_system_table_policy_matches_java() { "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; From a9e4e996a8d75ce8cad730a7247a6f05282a38ba Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 10 Sep 2026 18:01:18 +0800 Subject: [PATCH 08/15] fix(table): align audit log read planning --- .../datafusion/src/physical_plan/scan.rs | 76 +---- .../datafusion/src/system_tables/audit_log.rs | 11 +- .../integrations/datafusion/src/table/mod.rs | 11 +- .../datafusion/tests/system_tables.rs | 49 ++++ crates/paimon/src/table/audit_log_table.rs | 2 +- crates/paimon/src/table/mod.rs | 2 +- crates/paimon/src/table/read_builder.rs | 20 +- crates/paimon/src/table/table_read.rs | 267 ++++++++++++------ crates/paimon/src/table/table_scan.rs | 36 ++- 9 files changed, 282 insertions(+), 192 deletions(-) diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs b/crates/integrations/datafusion/src/physical_plan/scan.rs index d1cff261c..96265f451 100644 --- a/crates/integrations/datafusion/src/physical_plan/scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/scan.rs @@ -15,13 +15,12 @@ // 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, Schema, SchemaRef as ArrowSchemaRef, TimeUnit, + DataType as ArrowDataType, SchemaRef as ArrowSchemaRef, TimeUnit, }; use datafusion::arrow::record_batch::{RecordBatch, RecordBatchOptions}; use datafusion::common::stats::Precision; @@ -51,63 +50,13 @@ 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, - ROW_KIND_FIELD_NAME, SEQUENCE_NUMBER_FIELD_NAME, -}; +use paimon::spec::{DataField, Datum, MergeEngine, Predicate, PredicateBuilder, PredicateOperator}; 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!( @@ -1195,33 +1144,14 @@ impl ExecutionPlan for PaimonTableScan { ))); } 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), - ) + read.to_audit_log_arrow(splits.as_ref()) } 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 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 diff --git a/crates/integrations/datafusion/src/system_tables/audit_log.rs b/crates/integrations/datafusion/src/system_tables/audit_log.rs index 876d413b6..3f384c31b 100644 --- a/crates/integrations/datafusion/src/system_tables/audit_log.rs +++ b/crates/integrations/datafusion/src/system_tables/audit_log.rs @@ -94,9 +94,14 @@ impl TableProvider for AuditLogTable { 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)?; + let (plan, trace) = await_with_runtime( + read_builder + .new_scan() + .with_scan_all_files() + .plan_with_trace(), + ) + .await + .map_err(to_datafusion_error)?; PaimonScanBuilder { table: &self.table, diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 161db1cef..f7bc8a574 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -371,7 +371,7 @@ impl PaimonScanBuilder<'_> { read_fields: Vec, audit_log: bool, ) -> DFResult> { - let (projected_schema, mut read_type) = if let Some(indices) = self.projection { + let (projected_schema, read_type) = if let Some(indices) = self.projection { let fields: Vec = indices .iter() .map(|&index| self.schema.field(index).clone()) @@ -384,15 +384,6 @@ impl PaimonScanBuilder<'_> { } 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 first_row_audit = audit_log && self .table diff --git a/crates/integrations/datafusion/tests/system_tables.rs b/crates/integrations/datafusion/tests/system_tables.rs index 8fdd68665..dae555f74 100644 --- a/crates/integrations/datafusion/tests/system_tables.rs +++ b/crates/integrations/datafusion/tests/system_tables.rs @@ -336,6 +336,55 @@ async fn test_audit_log_system_table_keeps_row_kinds_and_sequence_numbers() { })); } +#[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; diff --git a/crates/paimon/src/table/audit_log_table.rs b/crates/paimon/src/table/audit_log_table.rs index 099779de0..9e71edf85 100644 --- a/crates/paimon/src/table/audit_log_table.rs +++ b/crates/paimon/src/table/audit_log_table.rs @@ -91,6 +91,6 @@ impl AuditLogTable { self.wrapped .new_read_builder() .new_read()? - .to_audit_log_arrow_for_splits(splits) + .to_audit_log_arrow(splits) } } diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 32484924a..deee104c0 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, 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..6cb8f79d9 100644 --- a/crates/paimon/src/table/read_builder.rs +++ b/crates/paimon/src/table/read_builder.rs @@ -505,10 +505,17 @@ 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()? { + let audit_projection = self.resolve_read_type()?; + let mut read_type = match &audit_projection { None => self.table.schema.fields().to_vec(), - Some(fields) => fields, + Some(fields) => fields.clone(), }; + read_type.retain(|field| { + !matches!( + field.id(), + crate::spec::ROW_KIND_FIELD_ID | crate::spec::SEQUENCE_NUMBER_FIELD_ID + ) + }); // Pass the FULL data predicate through (including `And`/`Or`/`Not`). // Pushdown/stats skip compound nodes; the residual pass enforces the full @@ -517,10 +524,13 @@ impl<'a> PaimonReadBuilder<'a> { Some(budget) => Arc::clone(budget), None => configured_parquet_read_budget(self.table)?, }; - Ok( - TableRead::new(self.table, read_type, self.filter.data_predicates.clone()) - .with_parquet_read_budget(parquet_read_budget), + Ok(TableRead::new_with_audit_projection( + self.table, + read_type, + self.filter.data_predicates.clone(), + audit_projection, ) + .with_parquet_read_budget(parquet_read_budget)) } /// Resolve the effective read type, deferring projection name resolution to diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index 9eb94c826..23655a583 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -42,6 +42,30 @@ use std::sync::Arc; const MAX_MERGE_INPUT_STREAMS: usize = 256; +#[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 IncrementalPlan> for AuditLogInput<'a> { + fn from(plan: &'a IncrementalPlan) -> Self { + Self::Incremental(plan) + } +} + /// Table read: reads data from splits (e.g. produced by [TableScan::plan]). /// /// Reference: [pypaimon.read.table_read.TableRead](https://github.com/apache/paimon/blob/master/paimon-python/pypaimon/read/table_read.py) @@ -82,6 +106,18 @@ impl<'a> TableRead<'a> { } } + pub(super) fn new_with_audit_projection( + table: &'a Table, + read_type: Vec, + data_predicates: Vec, + audit_projection: Option>, + ) -> Self { + Self(TableReadKind::Paimon( + PaimonTableRead::new(table, read_type, data_predicates) + .with_audit_projection(audit_projection), + )) + } + pub(crate) fn new_format( table: &'a Table, read_type: Vec, @@ -192,54 +228,20 @@ 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( + /// Returns audit-log rows for current splits or an incremental plan. + pub fn to_audit_log_arrow<'input>( &self, - plan: &IncrementalPlan, + input: impl Into>, ) -> 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(), - }), - } - } - - /// 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::Paimon(read) => match input.into() { + AuditLogInput::Current(splits) => read.audit_current_stream(splits), + AuditLogInput::Incremental(plan) => { + plan.validate()?; + read.audit_incremental_stream(plan) + } + }, TableReadKind::Format(_) => Err(crate::Error::Unsupported { message: "Format tables do not support audit log batch read".to_string(), }), @@ -255,6 +257,7 @@ impl<'a> TableRead<'a> { struct PaimonTableRead<'a> { table: &'a Table, read_type: Vec, + audit_projection: Option>, data_predicates: Vec, row_filter_factory: Option>, parquet_read_budget: Option>, @@ -271,6 +274,7 @@ impl<'a> PaimonTableRead<'a> { Self { table, read_type, + audit_projection: None, data_predicates, row_filter_factory: None, parquet_read_budget: None, @@ -278,6 +282,11 @@ impl<'a> PaimonTableRead<'a> { } } + fn with_audit_projection(mut self, projection: Option>) -> Self { + self.audit_projection = projection; + self + } + /// Schema (fields) that this read will produce. pub fn read_type(&self) -> &[DataField] { &self.read_type @@ -390,31 +399,13 @@ 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( + fn audit_current_stream( &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 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.read_type.clone(); let audit_schema = audit_schema_for_read_type(&user_read_type, include_rowkind, include_sequence)?; @@ -450,6 +441,7 @@ impl<'a> PaimonTableRead<'a> { read_type.clone(), 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(Arc::clone(&parquet_read_budget))) .read(&raw_splits)?; @@ -501,18 +493,18 @@ impl<'a> PaimonTableRead<'a> { self.to_arrow(data_splits)? }; - Ok(audit_stream_from_physical( + 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, &output_read_type) } - /// Returns an audit-log stream for a planned incremental scan. - pub fn to_audit_log_arrow( + fn audit_incremental_stream( &self, plan: &IncrementalPlan, ) -> crate::Result { @@ -530,6 +522,25 @@ impl<'a> PaimonTableRead<'a> { } } + fn audit_read_type(&self) -> crate::Result> { + let fields = self.audit_projection.clone().unwrap_or_else(|| { + audit_fields_for_read_type( + &self.read_type, + true, + audit_sequence_number_enabled(self.table), + ) + }); + if audit_field_requested(&fields, SEQUENCE_NUMBER_FIELD_ID) + && !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, + }); + } + Ok(fields) + } + fn audit_raw_stream( &self, plan: &IncrementalPlan, @@ -538,9 +549,12 @@ impl<'a> PaimonTableRead<'a> { plan.validate()?; let core_options = self.table.schema().core_options(); let data_splits = plan.data_splits(); + let output_read_type = self.audit_read_type()?; 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, true, include_sequence)?; + 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 { @@ -553,7 +567,7 @@ impl<'a> PaimonTableRead<'a> { ), ); } - if has_value_kind { + if has_value_kind && include_rowkind { read_type.push(DataField::new( VALUE_KIND_FIELD_ID, VALUE_KIND_FIELD_NAME.to_string(), @@ -573,25 +587,28 @@ impl<'a> PaimonTableRead<'a> { .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(audit_stream_from_physical( + let stream = audit_stream_from_physical( raw_stream, audit_schema, user_read_type, - true, + include_rowkind, include_sequence, - has_value_kind, - )) + has_value_kind && include_rowkind, + ); + project_audit_stream(stream, &output_read_type) } 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 output_read_type = self.audit_read_type()?; + let include_sequence = audit_field_requested(&output_read_type, SEQUENCE_NUMBER_FIELD_ID); 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 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(); @@ -601,7 +618,11 @@ impl<'a> PaimonTableRead<'a> { 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)?; + pair_read.to_audit_log_arrow_for_diff( + &before, + &after, + include_sequence, + )?; while let Some(batch) = pair_stream.next().await { yield batch?; } @@ -612,15 +633,16 @@ impl<'a> PaimonTableRead<'a> { while let Some(batch) = workers.next().await { yield batch?; } - })) + }); + project_audit_stream(stream, &output_read_type) } fn to_audit_log_arrow_for_diff( &self, before: &[DataSplit], after: &[DataSplit], + include_sequence: bool, ) -> crate::Result { - let include_sequence = audit_sequence_number_enabled(self.table); 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(); @@ -1165,11 +1187,53 @@ fn audit_stream_from_physical( }) } -fn audit_schema_for_read_type( +fn project_audit_stream( + stream: ArrowRecordBatchStream, + read_type: &[DataField], +) -> crate::Result { + 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, -) -> crate::Result> { +) -> Vec { let mut fields = Vec::with_capacity(read_type.len() + 2); if include_rowkind { fields.push(DataField::new( @@ -1186,7 +1250,19 @@ fn audit_schema_for_read_type( )); } fields.extend(read_type.iter().cloned()); - build_target_arrow_schema(&fields) + 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 { @@ -1894,11 +1970,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"), @@ -1914,7 +1993,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)) @@ -1948,8 +2027,22 @@ 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 current_audit = pk_read + .to_audit_log_arrow(std::slice::from_ref(&split)) + .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)) diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index d5de9048a..6e289ea53 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -1423,11 +1423,9 @@ impl<'a> PaimonTableScan<'a> { /// 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. + /// - Ordinary `merge-engine=first-row` reads: planned with + /// `skip_level_zero` and read via `DataFileReader`. Audit reads use + /// `scan_all_files` and merge visible versions, so they 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()); @@ -1441,7 +1439,7 @@ impl<'a> PaimonTableScan<'a> { ); if has_primary_keys && (!deletion_vectors_enabled || deletion_vectors_merge_on_read) - && !first_row + && (!first_row || self.scan_all_files) { retain_primary_key_conjuncts( &self.data_predicates, @@ -3719,12 +3717,9 @@ 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. + /// 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 +3777,23 @@ mod tests { planned_files, 1, "only the value-matching file should be planned on first-row" ); + + let (audit_plan, audit_trace) = reader + .new_scan() + .with_scan_all_files() + .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] From cfc40aa5b8e1600dad65847d0941f008bd2072d1 Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 11 Sep 2026 09:27:21 +0800 Subject: [PATCH 09/15] fix(table): harden audit log scan planning --- .../datafusion/src/system_tables/audit_log.rs | 11 +-- crates/paimon/src/table/audit_log_table.rs | 7 +- crates/paimon/src/table/read_builder.rs | 5 ++ crates/paimon/src/table/table_read.rs | 51 +++++++++++-- crates/paimon/src/table/table_scan.rs | 73 ++++++++++++++----- crates/paimon/tests/audit_log_table_test.rs | 20 +++-- 6 files changed, 128 insertions(+), 39 deletions(-) diff --git a/crates/integrations/datafusion/src/system_tables/audit_log.rs b/crates/integrations/datafusion/src/system_tables/audit_log.rs index 3f384c31b..9f2e59e59 100644 --- a/crates/integrations/datafusion/src/system_tables/audit_log.rs +++ b/crates/integrations/datafusion/src/system_tables/audit_log.rs @@ -94,14 +94,9 @@ impl TableProvider for AuditLogTable { if let Some(limit) = pushed_limit { read_builder.with_limit(limit); } - let (plan, trace) = await_with_runtime( - read_builder - .new_scan() - .with_scan_all_files() - .plan_with_trace(), - ) - .await - .map_err(to_datafusion_error)?; + let (plan, trace) = await_with_runtime(read_builder.new_audit_scan().plan_with_trace()) + .await + .map_err(to_datafusion_error)?; PaimonScanBuilder { table: &self.table, diff --git a/crates/paimon/src/table/audit_log_table.rs b/crates/paimon/src/table/audit_log_table.rs index 9e71edf85..639c9793a 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, DataSplit, Table}; +use super::{ArrowRecordBatchStream, 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, @@ -77,6 +77,11 @@ 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() + } + pub fn to_arrow(&self, plan: &IncrementalPlan) -> crate::Result { plan.validate()?; let read = self.wrapped.new_read_builder().new_read()?; diff --git a/crates/paimon/src/table/read_builder.rs b/crates/paimon/src/table/read_builder.rs index 6cb8f79d9..02387721f 100644 --- a/crates/paimon/src/table/read_builder.rs +++ b/crates/paimon/src/table/read_builder.rs @@ -262,6 +262,11 @@ 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_scan_all_files_preserving_projection() + } + /// Create a batch incremental scan over snapshot id range /// `(start_exclusive, end_inclusive]`. /// diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index 23655a583..a162df792 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -60,6 +60,12 @@ impl<'a, const N: usize> From<&'a [DataSplit; N]> for AuditLogInput<'a> { } } +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) @@ -501,7 +507,7 @@ impl<'a> PaimonTableRead<'a> { include_sequence, has_primary_keys && include_rowkind, ); - project_audit_stream(stream, &output_read_type) + project_audit_stream(stream, self.audit_projection.as_deref()) } fn audit_incremental_stream( @@ -595,7 +601,7 @@ impl<'a> PaimonTableRead<'a> { include_sequence, has_value_kind && include_rowkind, ); - project_audit_stream(stream, &output_read_type) + project_audit_stream(stream, self.audit_projection.as_deref()) } fn audit_diff_stream(&self, plan: &IncrementalPlan) -> crate::Result { @@ -634,7 +640,7 @@ impl<'a> PaimonTableRead<'a> { yield batch?; } }); - project_audit_stream(stream, &output_read_type) + project_audit_stream(stream, self.audit_projection.as_deref()) } fn to_audit_log_arrow_for_diff( @@ -1189,8 +1195,11 @@ fn audit_stream_from_physical( fn project_audit_stream( stream: ArrowRecordBatchStream, - read_type: &[DataField], + 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() @@ -1828,6 +1837,37 @@ mod tests { 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_diff_batch_builders_pin_each_input_batch_once() { let schema = Arc::new(ArrowSchema::new(vec![Field::new( @@ -2033,8 +2073,9 @@ mod tests { .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(std::slice::from_ref(&split)) + .to_audit_log_arrow(&splits) .unwrap() .try_collect::>() .await diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index 6e289ea53..5c9f1f073 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> { } } + pub(super) fn with_scan_all_files_preserving_projection(self) -> Self { + match self.0 { + TableScanKind::Paimon(scan) => Self(TableScanKind::Paimon( + scan.with_scan_all_files_preserving_projection(), + )), + TableScanKind::Format(scan) => Self(TableScanKind::Format(scan)), + } + } + pub fn with_row_ranges(self, ranges: Vec) -> Self { match self.0 { TableScanKind::Paimon(scan) => { @@ -1056,6 +1065,11 @@ impl<'a> PaimonTableScan<'a> { self } + fn with_scan_all_files_preserving_projection(mut self) -> Self { + self.scan_all_files = true; + self + } + /// Set row ranges for scan-time filtering. /// /// This replaces any existing row_ranges. Typically used to inject @@ -1271,7 +1285,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 +1293,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()?, @@ -2172,10 +2188,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; @@ -2869,6 +2885,34 @@ 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_scan_all_files_preserving_projection(); + + assert!(scan.scan_all_files); + 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( @@ -3778,12 +3822,7 @@ mod tests { "only the value-matching file should be planned on first-row" ); - let (audit_plan, audit_trace) = reader - .new_scan() - .with_scan_all_files() - .plan_with_trace() - .await - .unwrap(); + 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 diff --git a/crates/paimon/tests/audit_log_table_test.rs b/crates/paimon/tests/audit_log_table_test.rs index 63237a6cc..76a1d82ea 100644 --- a/crates/paimon/tests/audit_log_table_test.rs +++ b/crates/paimon/tests/audit_log_table_test.rs @@ -487,15 +487,19 @@ async fn audit_log_current_scan_supports_first_row() { 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(); + 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" + ); + + let audit = AuditLogTable::new(table); + let plan = audit.new_scan().plan().await.unwrap(); assert_eq!(plan.splits().len(), 2); - let batches: Vec = AuditLogTable::new(table) + let batches: Vec = audit .to_arrow_for_splits(plan.splits()) .unwrap() .try_collect() From 0e86915ba92485685167b70447085ebcc3950de4 Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 11 Sep 2026 10:18:52 +0800 Subject: [PATCH 10/15] fix(table): make audit splits independently mergeable --- .../integrations/datafusion/src/table/mod.rs | 23 +++++--------- crates/paimon/src/table/table_scan.rs | 26 ++++++++++------ crates/paimon/tests/audit_log_table_test.rs | 31 ++++++++++++++----- 3 files changed, 46 insertions(+), 34 deletions(-) diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index f7bc8a574..afc1b8e9c 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -32,7 +32,7 @@ use datafusion::logical_expr::dml::InsertOp; use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; use datafusion::physical_plan::ExecutionPlan; use paimon::spec::{ - BigIntType, CoreOptions, DataField, DataType, MergeEngine, ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME, + BigIntType, CoreOptions, DataField, DataType, ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME, }; use paimon::table::Table; @@ -384,21 +384,9 @@ impl PaimonScanBuilder<'_> { } else { (self.schema.clone(), read_fields) }; - let first_row_audit = audit_log - && self - .table - .schema() - .core_options() - .merge_engine() - .map_err(to_datafusion_error)? - == MergeEngine::FirstRow; let splits = self.plan.into_splits(); let planned_partitions: Vec> = if splits.is_empty() { vec![Arc::from(Vec::new())] - } else if first_row_audit { - // ponytail: keep merge groups intact; add group-aware balancing if - // first-row audit parallelism becomes necessary. - vec![Arc::from(splits)] } else { let num_partitions = splits.len().min(self.target_partitions.max(1)); bucket_round_robin(splits, num_partitions) @@ -600,7 +588,7 @@ mod tests { } #[test] - fn test_first_row_audit_keeps_split_group_in_one_partition() { + fn test_first_row_audit_distributes_independent_splits() { let file_io = paimon::io::FileIOBuilder::new("memory").build().unwrap(); let schema = paimon::spec::Schema::builder() .column( @@ -650,8 +638,11 @@ mod tests { .downcast_ref::() .expect("Expected PaimonTableScan"); - assert_eq!(scan.planned_partitions().len(), 1); - assert_eq!(scan.planned_partitions()[0].len(), 2); + assert_eq!(scan.planned_partitions().len(), 2); + assert!(scan + .planned_partitions() + .iter() + .all(|splits| splits.len() == 1)); } fn get_test_warehouse() -> String { diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index 5c9f1f073..aeb4bd76d 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -1030,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>, } @@ -1051,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, } } @@ -1067,6 +1070,7 @@ impl<'a> PaimonTableScan<'a> { fn with_scan_all_files_preserving_projection(mut self) -> Self { self.scan_all_files = true; + self.merge_key_overlaps = true; self } @@ -1930,16 +1934,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 { @@ -2904,6 +2909,7 @@ mod tests { .with_scan_all_files_preserving_projection(); 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,) diff --git a/crates/paimon/tests/audit_log_table_test.rs b/crates/paimon/tests/audit_log_table_test.rs index 76a1d82ea..e2ce135d0 100644 --- a/crates/paimon/tests/audit_log_table_test.rs +++ b/crates/paimon/tests/audit_log_table_test.rs @@ -485,6 +485,7 @@ async fn audit_log_current_scan_supports_first_row() { 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(); @@ -495,20 +496,34 @@ async fn audit_log_current_scan_supports_first_row() { 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 batches: Vec = audit - .to_arrow_for_splits(plan.splits()) - .unwrap() - .try_collect() - .await - .unwrap(); + 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!( - collect_audit_rows(&batches), - vec![("+I".to_string(), 1, 10)] + rows, + vec![("+I".to_string(), 1, 10), ("+I".to_string(), 3, 30)] ); } From 7c305efe6593b719b377bf9456ef1af8884caadb Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 11 Sep 2026 10:43:35 +0800 Subject: [PATCH 11/15] docs(table): clarify first-row audit split path --- crates/paimon/src/table/table_scan.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index aeb4bd76d..ecd617d4d 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -2075,9 +2075,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) From 3af12beacbf9786c17db1fa925e0d71ad96998ee Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 11 Sep 2026 14:51:36 +0800 Subject: [PATCH 12/15] fix(table): preserve audit projection and pruning semantics --- crates/paimon/src/table/read_builder.rs | 8 +-- crates/paimon/src/table/table_read.rs | 14 +++- crates/paimon/src/table/table_scan.rs | 72 +++++++++++++++++++-- crates/paimon/tests/audit_log_table_test.rs | 58 ++++++++++++++++- 4 files changed, 134 insertions(+), 18 deletions(-) diff --git a/crates/paimon/src/table/read_builder.rs b/crates/paimon/src/table/read_builder.rs index 02387721f..6843b92d6 100644 --- a/crates/paimon/src/table/read_builder.rs +++ b/crates/paimon/src/table/read_builder.rs @@ -511,16 +511,10 @@ impl<'a> PaimonReadBuilder<'a> { let core_options = self.table.schema.core_options(); core_options.ensure_read_authorized()?; let audit_projection = self.resolve_read_type()?; - let mut read_type = match &audit_projection { + let read_type = match &audit_projection { None => self.table.schema.fields().to_vec(), Some(fields) => fields.clone(), }; - read_type.retain(|field| { - !matches!( - field.id(), - crate::spec::ROW_KIND_FIELD_ID | crate::spec::SEQUENCE_NUMBER_FIELD_ID - ) - }); // Pass the FULL data predicate through (including `And`/`Or`/`Not`). // Pushdown/stats skip compound nodes; the residual pass enforces the full diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index a162df792..2dba34d69 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -412,7 +412,7 @@ impl<'a> PaimonTableRead<'a> { 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.read_type.clone(); + 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.table.schema().primary_keys().is_empty(); @@ -547,6 +547,14 @@ impl<'a> PaimonTableRead<'a> { Ok(fields) } + fn audit_user_read_type(&self) -> Vec { + self.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, @@ -556,7 +564,7 @@ impl<'a> PaimonTableRead<'a> { let core_options = self.table.schema().core_options(); let data_splits = plan.data_splits(); let output_read_type = self.audit_read_type()?; - let user_read_type = self.read_type.clone(); + 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 = @@ -610,7 +618,7 @@ impl<'a> PaimonTableRead<'a> { let output_read_type = self.audit_read_type()?; let include_sequence = audit_field_requested(&output_read_type, SEQUENCE_NUMBER_FIELD_ID); let table = self.table.clone(); - let read_type = self.read_type.clone(); + let read_type = self.audit_user_read_type(); let data_predicates = self.data_predicates.clone(); let parquet_read_budget = self.parquet_read_budget()?; diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index ecd617d4d..1197413b3 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -1439,13 +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. - /// - Ordinary `merge-engine=first-row` reads: planned with - /// `skip_level_zero` and read via `DataFileReader`. Audit reads use - /// `scan_all_files` and merge visible versions, so they are not exempt. + /// - 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()); @@ -1458,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.scan_all_files) + && (self.merge_key_overlaps + || ((!deletion_vectors_enabled || deletion_vectors_merge_on_read) && !first_row)) { retain_primary_key_conjuncts( &self.data_predicates, @@ -3767,6 +3768,65 @@ mod tests { ); } + #[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. diff --git a/crates/paimon/tests/audit_log_table_test.rs b/crates/paimon/tests/audit_log_table_test.rs index e2ce135d0..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,60 @@ 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"; From cf5d1e2c3a5320aa35cdb6edafdf5705bb4b25d2 Mon Sep 17 00:00:00 2001 From: yantian Date: Mon, 14 Sep 2026 09:54:37 +0800 Subject: [PATCH 13/15] refactor(table): decouple audit reads from standard scans --- .../datafusion/src/physical_plan/audit_log.rs | 240 +++++ .../datafusion/src/physical_plan/mod.rs | 2 + .../datafusion/src/physical_plan/scan.rs | 227 ++--- .../datafusion/src/system_tables/audit_log.rs | 6 +- .../integrations/datafusion/src/table/mod.rs | 114 +-- crates/paimon/src/table/audit_log_table.rs | 15 +- crates/paimon/src/table/mod.rs | 2 +- crates/paimon/src/table/read_builder.rs | 23 +- crates/paimon/src/table/table_read.rs | 867 +--------------- crates/paimon/src/table/table_read/audit.rs | 940 ++++++++++++++++++ crates/paimon/src/table/table_scan.rs | 12 +- 11 files changed, 1340 insertions(+), 1108 deletions(-) create mode 100644 crates/integrations/datafusion/src/physical_plan/audit_log.rs create mode 100644 crates/paimon/src/table/table_read/audit.rs 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 96265f451..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; @@ -778,8 +778,6 @@ 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 { @@ -886,37 +884,9 @@ 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 } @@ -999,38 +969,12 @@ impl PaimonTableScan { .map(|(accumulator, field)| accumulator.finish(field.data_type(), exact_null_counts)) .collect() } -} - -impl ExecutionPlan for PaimonTableScan { - fn name(&self) -> &str { - if self.audit_log { - "PaimonAuditLogScan" - } else { - "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() @@ -1046,16 +990,7 @@ impl ExecutionPlan for PaimonTableScan { let parent_filter_handled = filters .into_iter() .map(|filter| { - 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 + if supported(&filter) && can_expr_be_pushed_down_with_schemas(&filter, schema.as_ref()) { accepted.push(filter); @@ -1092,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!( @@ -1116,7 +1053,6 @@ 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(); @@ -1143,12 +1079,7 @@ impl ExecutionPlan for PaimonTableScan { Arc::clone(&schema), ))); } - let stream = if audit_log { - read.to_audit_log_arrow(splits.as_ref()) - } else { - 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 batch = result.map_err(to_datafusion_error)?; @@ -1185,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]), @@ -1208,9 +1228,7 @@ 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 self.audit_log { - Precision::Absent - } else if all_row_counts_known + let num_rows_precision = if all_row_counts_known && self.limit.is_none() && self.filter_exact && self.runtime_filters.is_empty() @@ -1234,48 +1252,7 @@ impl DisplayAs for PaimonTableScan { _t: datafusion::physical_plan::DisplayFormatType, f: &mut std::fmt::Formatter, ) -> std::fmt::Result { - 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 - .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 index 9f2e59e59..b464afd4a 100644 --- a/crates/integrations/datafusion/src/system_tables/audit_log.rs +++ b/crates/integrations/datafusion/src/system_tables/audit_log.rs @@ -31,6 +31,7 @@ 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}; @@ -98,7 +99,7 @@ impl TableProvider for AuditLogTable { .await .map_err(to_datafusion_error)?; - PaimonScanBuilder { + let scan = PaimonScanBuilder { table: &self.table, schema: &self.schema, plan, @@ -110,7 +111,8 @@ impl TableProvider for AuditLogTable { filter_exact: false, case_sensitive: true, } - .build_audit_log(self.fields.clone()) + .build_scan(self.fields.clone())?; + Ok(Arc::new(PaimonAuditLogScan::new(scan))) } fn supports_filters_pushdown( diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index afc1b8e9c..f5e3e5a74 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -356,21 +356,10 @@ impl PaimonScanBuilder<'_> { self, read_fields: Vec, ) -> DFResult> { - self.build_scan(read_fields, false) + Ok(Arc::new(self.build_scan(read_fields)?)) } - 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> { + 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() @@ -395,31 +384,18 @@ impl PaimonScanBuilder<'_> { .collect() }; - 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, - )?)) - } + 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, + ) } } @@ -563,9 +539,7 @@ mod tests { use datafusion::prelude::{SessionConfig, SessionContext}; use paimon::catalog::Identifier; use paimon::spec::{ArrayType, MapType, RowType, VarCharType}; - use paimon::{ - Catalog, CatalogOptions, DataSplit, DataSplitBuilder, FileSystemCatalog, Options, - }; + use paimon::{Catalog, CatalogOptions, DataSplit, FileSystemCatalog, Options}; use crate::physical_plan::PaimonTableScan; @@ -587,64 +561,6 @@ mod tests { assert_eq!(result, vec![vec![1, 2, 3]]); } - #[test] - fn test_first_row_audit_distributes_independent_splits() { - 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 = datafusion_read_fields(&table); - 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_audit_log(read_fields) - .unwrap(); - let scan = plan - .downcast_ref::() - .expect("Expected PaimonTableScan"); - - assert_eq!(scan.planned_partitions().len(), 2); - assert!(scan - .planned_partitions() - .iter() - .all(|splits| splits.len() == 1)); - } - fn get_test_warehouse() -> String { std::env::var("PAIMON_TEST_WAREHOUSE") .unwrap_or_else(|_| "/tmp/paimon-warehouse".to_string()) diff --git a/crates/paimon/src/table/audit_log_table.rs b/crates/paimon/src/table/audit_log_table.rs index 639c9793a..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, DataSplit, Table, TableScan}; +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, @@ -82,10 +82,14 @@ impl AuditLogTable { 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. @@ -93,9 +97,6 @@ impl AuditLogTable { &self, splits: &[DataSplit], ) -> crate::Result { - self.wrapped - .new_read_builder() - .new_read()? - .to_audit_log_arrow(splits) + self.new_read()?.to_arrow(splits) } } diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index deee104c0..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::{AuditLogInput, 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 6843b92d6..52dc1cea9 100644 --- a/crates/paimon/src/table/read_builder.rs +++ b/crates/paimon/src/table/read_builder.rs @@ -262,11 +262,6 @@ 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_scan_all_files_preserving_projection() - } - /// Create a batch incremental scan over snapshot id range /// `(start_exclusive, end_inclusive]`. /// @@ -510,11 +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 audit_projection = self.resolve_read_type()?; - let read_type = match &audit_projection { - None => self.table.schema.fields().to_vec(), - Some(fields) => fields.clone(), - }; + 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 @@ -523,13 +516,11 @@ impl<'a> PaimonReadBuilder<'a> { Some(budget) => Arc::clone(budget), None => configured_parquet_read_budget(self.table)?, }; - Ok(TableRead::new_with_audit_projection( - self.table, - read_type, - self.filter.data_predicates.clone(), - audit_projection, + Ok( + TableRead::new(self.table, read_type, self.filter.data_predicates.clone()) + .with_explicit_projection(explicit_projection) + .with_parquet_read_budget(parquet_read_budget), ) - .with_parquet_read_budget(parquet_read_budget)) } /// Resolve the effective read type, deferring projection name resolution to diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index 2dba34d69..0fd10d871 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -25,14 +25,10 @@ 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, -}; +use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions}; use arrow_schema::Schema as ArrowSchema; use arrow_select::interleave::interleave; use futures::{stream, StreamExt}; @@ -40,37 +36,10 @@ use std::cmp::Ordering; use std::collections::HashMap; use std::sync::Arc; -const MAX_MERGE_INPUT_STREAMS: usize = 256; - -#[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()) - } -} +mod audit; +pub use audit::{AuditLogInput, AuditLogRead}; -impl<'a> From<&'a IncrementalPlan> for AuditLogInput<'a> { - fn from(plan: &'a IncrementalPlan) -> Self { - Self::Incremental(plan) - } -} +const MAX_MERGE_INPUT_STREAMS: usize = 256; /// Table read: reads data from splits (e.g. produced by [TableScan::plan]). /// @@ -112,16 +81,12 @@ impl<'a> TableRead<'a> { } } - pub(super) fn new_with_audit_projection( - table: &'a Table, - read_type: Vec, - data_predicates: Vec, - audit_projection: Option>, - ) -> Self { - Self(TableReadKind::Paimon( - PaimonTableRead::new(table, read_type, data_predicates) - .with_audit_projection(audit_projection), - )) + /// 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( @@ -234,26 +199,6 @@ impl<'a> TableRead<'a> { } } - /// Returns audit-log rows for current splits or an incremental plan. - pub fn to_audit_log_arrow<'input>( - &self, - input: impl Into>, - ) -> crate::Result { - self.ensure_query_auth_allowed()?; - match &self.0 { - TableReadKind::Paimon(read) => match input.into() { - AuditLogInput::Current(splits) => read.audit_current_stream(splits), - AuditLogInput::Incremental(plan) => { - plan.validate()?; - read.audit_incremental_stream(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() } @@ -263,7 +208,7 @@ impl<'a> TableRead<'a> { struct PaimonTableRead<'a> { table: &'a Table, read_type: Vec, - audit_projection: Option>, + explicit_projection: bool, data_predicates: Vec, row_filter_factory: Option>, parquet_read_budget: Option>, @@ -280,7 +225,7 @@ impl<'a> PaimonTableRead<'a> { Self { table, read_type, - audit_projection: None, + explicit_projection: false, data_predicates, row_filter_factory: None, parquet_read_budget: None, @@ -288,11 +233,6 @@ impl<'a> PaimonTableRead<'a> { } } - fn with_audit_projection(mut self, projection: Option>) -> Self { - self.audit_projection = projection; - self - } - /// Schema (fields) that this read will produce. pub fn read_type(&self) -> &[DataField] { &self.read_type @@ -405,339 +345,6 @@ impl<'a> PaimonTableRead<'a> { })) } - 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.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_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.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)? - }; - - 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.audit_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.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.audit_projection.clone().unwrap_or_else(|| { - audit_fields_for_read_type( - &self.read_type, - true, - audit_sequence_number_enabled(self.table), - ) - }); - if audit_field_requested(&fields, SEQUENCE_NUMBER_FIELD_ID) - && !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, - }); - } - Ok(fields) - } - - fn audit_user_read_type(&self) -> Vec { - self.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.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.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)?; - 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.audit_projection.as_deref()) - } - - 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 output_read_type = self.audit_read_type()?; - let include_sequence = audit_field_requested(&output_read_type, SEQUENCE_NUMBER_FIELD_ID); - let table = self.table.clone(); - let read_type = self.audit_user_read_type(); - let data_predicates = self.data_predicates.clone(); - let parquet_read_budget = self.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 = 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, - 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.audit_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_type, true, 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, 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()?; - } - })) - } - fn to_diff_after_image_stream( &self, before: &[DataSplit], @@ -1062,268 +669,6 @@ 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 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)) -} - const DIFF_BATCH_SIZE: usize = 8192; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -1397,74 +742,6 @@ impl ArrowCursor { } } -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)), - } - }) - } -} - struct DiffAfterImageBatchBuilder { schema: Arc, row_indices: Vec<(usize, usize)>, @@ -1574,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> = @@ -1845,37 +1094,6 @@ mod tests { 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_diff_batch_builders_pin_each_input_batch_once() { let schema = Arc::new(ArrowSchema::new(vec![Field::new( @@ -1890,28 +1108,6 @@ mod tests { 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" - ); - let mut after = DiffAfterImageBatchBuilder::new( Arc::new(ArrowSchema::new(vec![Field::new( "id", @@ -1938,7 +1134,7 @@ mod tests { ); } - fn file(name: &str, level: i32, delete_row_count: Option) -> DataFileMeta { + pub(super) fn file(name: &str, level: i32, delete_row_count: Option) -> DataFileMeta { DataFileMeta { file_name: name.to_string(), file_size: 128, @@ -1964,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)) @@ -2121,20 +1317,6 @@ 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)); @@ -2142,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 1197413b3..9dd8ae3d8 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -909,11 +909,11 @@ impl<'a> TableScan<'a> { } } - pub(super) fn with_scan_all_files_preserving_projection(self) -> Self { + /// 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_scan_all_files_preserving_projection(), - )), + TableScanKind::Paimon(scan) => Self(TableScanKind::Paimon(scan.with_all_versions())), TableScanKind::Format(scan) => Self(TableScanKind::Format(scan)), } } @@ -1068,7 +1068,7 @@ impl<'a> PaimonTableScan<'a> { self } - fn with_scan_all_files_preserving_projection(mut self) -> Self { + fn with_all_versions(mut self) -> Self { self.scan_all_files = true; self.merge_key_overlaps = true; self @@ -2907,7 +2907,7 @@ mod tests { .unwrap(); let scan = PaimonTableScan::new(&table, None, vec![predicate], None, None, None) .with_projected_read_field_ids(Some(projected.clone())) - .with_scan_all_files_preserving_projection(); + .with_all_versions(); assert!(scan.scan_all_files); assert!(scan.merge_key_overlaps); From 1781d212741d448ca206f824503b69e6fa5fc9cf Mon Sep 17 00:00:00 2001 From: yantian Date: Mon, 14 Sep 2026 11:26:16 +0800 Subject: [PATCH 14/15] refactor: simplify audit log scan and read wrappers --- .../datafusion/src/system_tables/mod.rs | 43 +- .../integrations/datafusion/src/table/mod.rs | 7 +- crates/paimon/src/table/audit_log_table.rs | 23 +- .../paimon/src/table/audit_log_table/merge.rs | 189 ++++ .../paimon/src/table/audit_log_table/read.rs | 210 ++++ .../paimon/src/table/audit_log_table/scan.rs | 265 +++++ crates/paimon/src/table/kv_file_reader.rs | 165 +-- crates/paimon/src/table/mod.rs | 4 +- crates/paimon/src/table/read_builder.rs | 8 +- crates/paimon/src/table/sort_merge.rs | 175 +--- crates/paimon/src/table/table_read.rs | 488 ++++++++- crates/paimon/src/table/table_read/audit.rs | 976 ------------------ crates/paimon/src/table/table_scan.rs | 211 +--- crates/paimon/tests/audit_log_table_test.rs | 10 +- 14 files changed, 1313 insertions(+), 1461 deletions(-) create mode 100644 crates/paimon/src/table/audit_log_table/merge.rs create mode 100644 crates/paimon/src/table/audit_log_table/read.rs create mode 100644 crates/paimon/src/table/audit_log_table/scan.rs delete mode 100644 crates/paimon/src/table/table_read/audit.rs diff --git a/crates/integrations/datafusion/src/system_tables/mod.rs b/crates/integrations/datafusion/src/system_tables/mod.rs index 09b66c2bd..22fbc5089 100644 --- a/crates/integrations/datafusion/src/system_tables/mod.rs +++ b/crates/integrations/datafusion/src/system_tables/mod.rs @@ -81,21 +81,6 @@ 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). @@ -143,21 +128,6 @@ 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, @@ -168,7 +138,10 @@ pub(crate) fn provider_for_table( return Ok(None); } crate::table_loader::ensure_paimon_served(&table, &identifier)?; - ensure_system_table_read_supported(table.schema().options(), system_name)?; + // 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)?; if system_name.eq_ignore_ascii_case("partitions") { return partitions::build(catalog, identifier, table).map(Some); } @@ -201,12 +174,16 @@ pub(crate) async fn load( .to_string(), )); } - ensure_system_table_read_supported(&dynamic_options, &system_name)?; + paimon::spec::CoreOptions::new(&dynamic_options) + .ensure_read_authorized() + .map_err(to_datafusion_error)?; 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)?; + paimon::spec::CoreOptions::new(table.schema().options()) + .ensure_read_authorized() + .map_err(to_datafusion_error)?; if let Some(branch) = object.branch() { if !system_name.eq_ignore_ascii_case("branches") { table = table diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index f5e3e5a74..f7db6c3e2 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -363,16 +363,17 @@ impl PaimonScanBuilder<'_> { let (projected_schema, read_type) = if let Some(indices) = self.projection { let fields: Vec = indices .iter() - .map(|&index| self.schema.field(index).clone()) + .map(|&i| self.schema.field(i).clone()) .collect(); let read_type = indices .iter() - .map(|&index| read_fields[index].clone()) - .collect(); + .map(|&i| read_fields[i].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())] diff --git a/crates/paimon/src/table/audit_log_table.rs b/crates/paimon/src/table/audit_log_table.rs index a77c7febb..f943188f5 100644 --- a/crates/paimon/src/table/audit_log_table.rs +++ b/crates/paimon/src/table/audit_log_table.rs @@ -16,12 +16,14 @@ // under the License. use super::incremental_scan::{IncrementalPlan, IncrementalScan, IncrementalScanMode}; -use super::{ArrowRecordBatchStream, AuditLogRead, DataSplit, Table, TableScan}; +use super::{ArrowRecordBatchStream, AuditLogRead, AuditLogScan, 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, }; +pub(super) mod merge; + /// Wrapper that exposes table rows with a leading `rowkind` audit column. /// /// Incremental reads produce: @@ -33,6 +35,8 @@ 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 } @@ -64,8 +68,9 @@ impl AuditLogTable { fn sequence_number_enabled(&self) -> bool { self.wrapped .schema() - .core_options() - .table_read_sequence_number_enabled() + .options() + .get(TABLE_READ_SEQUENCE_NUMBER_ENABLED) + .is_some_and(|v| v.eq_ignore_ascii_case("true")) } pub fn new_incremental_scan( @@ -78,18 +83,24 @@ impl AuditLogTable { } /// Plan a current-state audit read for [`Self::to_arrow_for_splits`]. - pub fn new_scan(&self) -> TableScan<'_> { + pub fn new_scan(&self) -> AuditLogScan<'_> { 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()?) + AuditLogRead::new( + self.wrapped + .new_read_builder() + .with_read_type(self.fields()?) + .new_read()?, + ) } pub fn to_arrow(&self, plan: &IncrementalPlan) -> crate::Result { plan.validate()?; - self.new_read()?.to_arrow(plan) + 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. diff --git a/crates/paimon/src/table/audit_log_table/merge.rs b/crates/paimon/src/table/audit_log_table/merge.rs new file mode 100644 index 000000000..801fa5cc7 --- /dev/null +++ b/crates/paimon/src/table/audit_log_table/merge.rs @@ -0,0 +1,189 @@ +// 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. + +//! Merge policies for current-state audit reads. + +use super::super::sort_merge::{ + compare_sequence_order, AggregateMergeFunction, BufferedBatch, MergeFunction, MergeResult, + MergeRow, PartialUpdateMergeFunction, +}; +use crate::arrow::build_target_arrow_schema; +use crate::spec::{ + CoreOptions, DataField, MergeEngine, RowKind, SEQUENCE_NUMBER_FIELD_ID, VALUE_KIND_FIELD_ID, +}; +use crate::table::kv_file_reader::KeyValueReadConfig; +use crate::Error; +use arrow_array::{Int64Array, Int8Array, RecordBatch}; +use arrow_schema::SchemaRef; +use std::sync::Arc; + +pub(in crate::table) fn new_merge_function( + config: &KeyValueReadConfig, + fields: &[DataField], +) -> crate::Result> { + let options = &config.table_options; + let engine = config.merge_engine; + if matches!(engine, MergeEngine::Deduplicate | MergeEngine::FirstRow) { + return Ok(Box::new(AuditKeyMergeFunction { + first_row: engine == MergeEngine::FirstRow, + ignore_delete: CoreOptions::new(options).ignore_delete(), + })); + } + let value_projection: Vec<_> = fields + .iter() + .enumerate() + .filter(|(_, field)| !matches!(field.id(), SEQUENCE_NUMBER_FIELD_ID | VALUE_KIND_FIELD_ID)) + .map(|(index, _)| index) + .collect(); + let value_fields: Vec<_> = value_projection + .iter() + .map(|&index| fields[index].clone()) + .collect(); + let inner: Box = match engine { + MergeEngine::PartialUpdate => Box::new(PartialUpdateMergeFunction::new_with_schema( + options, + &config.table_name, + &config.table_fields, + &value_fields, + &config.primary_keys, + )?), + MergeEngine::Aggregation => Box::new(AggregateMergeFunction::new( + options, + &config.table_name, + &value_fields, + &config.primary_keys, + &config.sequence_fields, + )?), + _ => unreachable!(), + }; + if value_projection.len() == fields.len() { + return Ok(inner); + } + Ok(Box::new(AuditValueMergeFunction { + inner, + value_projection, + value_schema: build_target_arrow_schema(&value_fields)?, + fields: fields.to_vec(), + })) +} + +// Reuse the ordinary value merge, then expose its logical INSERT and latest add sequence. +struct AuditValueMergeFunction { + inner: Box, + value_projection: Vec, + value_schema: SchemaRef, + fields: Vec, +} + +impl MergeFunction for AuditValueMergeFunction { + fn merge( + &self, + rows: &[MergeRow], + batch_buffer: &[BufferedBatch], + source_output_col_indices: &[usize], + output_schema: &SchemaRef, + ) -> crate::Result { + let indices: Vec<_> = self + .value_projection + .iter() + .map(|&index| source_output_col_indices[index]) + .collect(); + match self + .inner + .merge(rows, batch_buffer, &indices, &self.value_schema)? + { + MergeResult::MaterializedRow(batch) => { + let mut columns = batch.columns().to_vec(); + for (index, field) in self.fields.iter().enumerate() { + match field.id() { + SEQUENCE_NUMBER_FIELD_ID => { + let winner = rows + .iter() + .filter(|row| matches!(row.value_kind, 0 | 2)) + .max_by(|left, right| compare_sequence_order(left, right)) + .expect("materialized merge must contain an add row"); + columns.insert( + index, + Arc::new(Int64Array::from(vec![winner.sequence_number])), + ); + } + VALUE_KIND_FIELD_ID => { + columns.insert(index, Arc::new(Int8Array::from(vec![0]))) + } + _ => {} + } + } + Ok(MergeResult::MaterializedRow( + RecordBatch::try_new(output_schema.clone(), columns).map_err(|error| { + Error::UnexpectedError { + message: format!("Failed to build audit merge row: {error}"), + source: Some(Box::new(error)), + } + })?, + )) + } + result => Ok(result), + } + } +} + +/// Keep the winning physical row, including retracts for deduplicate tables. +struct AuditKeyMergeFunction { + first_row: bool, + ignore_delete: bool, +} + +impl MergeFunction for AuditKeyMergeFunction { + fn merge( + &self, + rows: &[MergeRow], + _batch_buffer: &[BufferedBatch], + _source_output_col_indices: &[usize], + _output_schema: &SchemaRef, + ) -> crate::Result { + let mut winner = None; + for row in rows { + if (self.first_row || self.ignore_delete) + && !RowKind::from_value(row.value_kind)?.is_add() + { + if self.ignore_delete { + continue; + } + 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 winner.is_none_or(|best| { + let order = compare_sequence_order(row, best); + if self.first_row { + order.is_lt() + } else { + order.is_ge() + } + }) { + winner = Some(row); + } + } + Ok(match winner { + Some(row) => MergeResult::SourceRow { + batch_idx: row.batch_idx, + row_idx: row.row_idx, + }, + None => MergeResult::Omit, + }) + } +} diff --git a/crates/paimon/src/table/audit_log_table/read.rs b/crates/paimon/src/table/audit_log_table/read.rs new file mode 100644 index 000000000..625d31700 --- /dev/null +++ b/crates/paimon/src/table/audit_log_table/read.rs @@ -0,0 +1,210 @@ +// 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. + +//! Current-state audit reads, including winning retract rows. + +use super::{ + audit_sequence_number_enabled, rowkind_array_from_column, PaimonTableRead, TableRead, + TableReadKind, MAX_MERGE_INPUT_STREAMS, +}; +use crate::arrow::build_target_arrow_schema; +use crate::spec::{ + DataField, DataType, MergeEngine, TinyIntType, ROW_KIND_FIELD_ID, SEQUENCE_NUMBER_FIELD_ID, + VALUE_KIND_FIELD_ID, VALUE_KIND_FIELD_NAME, +}; +use crate::table::data_file_reader::DataFileReader; +use crate::table::kv_file_reader::{KeyValueFileReader, KeyValueReadConfig}; +use crate::table::ArrowRecordBatchStream; +use crate::DataSplit; +use arrow_array::{ArrayRef, RecordBatch, RecordBatchOptions, StringArray}; +use futures::{stream, StreamExt}; +use std::sync::Arc; + +/// Reads current-state audit rows using the supplied read's exact projection, +/// predicates and Parquet budget. Include `rowkind` in the projection to expose it. +#[derive(Debug, Clone)] +pub struct AuditLogRead<'a> { + read: PaimonTableRead<'a>, +} + +impl<'a> AuditLogRead<'a> { + pub fn new(read: TableRead<'a>) -> crate::Result { + read.ensure_query_auth_allowed()?; + match read.0 { + TableReadKind::Paimon(read) => Ok(Self { read }), + TableReadKind::Format(_) => Err(crate::Error::Unsupported { + message: "Format tables do not support audit log batch read".to_string(), + }), + } + } + + /// Reads splits planned by an audit scan, retaining winning retract rows. + pub fn to_arrow(&self, data_splits: &[DataSplit]) -> crate::Result { + let output_read_type = self.read.read_type.clone(); + if output_read_type + .iter() + .any(|field| field.id() == 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, + }); + } + let audit_schema = build_target_arrow_schema(&output_read_type)?; + let has_primary_keys = !self.read.table.schema().primary_keys().is_empty(); + let mut read_type: Vec<_> = output_read_type + .iter() + .filter(|field| field.id() != ROW_KIND_FIELD_ID) + .cloned() + .collect(); + if has_primary_keys + && output_read_type + .iter() + .any(|field| field.id() == ROW_KIND_FIELD_ID) + { + read_type.push(DataField::new( + VALUE_KIND_FIELD_ID, + VALUE_KIND_FIELD_NAME.to_string(), + DataType::TinyInt(TinyIntType::new()), + )); + } + + let physical_stream = if has_primary_keys { + let core_options = self.read.table.schema().core_options(); + let merge_engine = core_options.merge_engine()?; + let (raw_splits, merge_splits): (Vec<_>, Vec<_>) = data_splits + .iter() + .cloned() + .partition(|split| audit_raw_convertible(split, 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_stream = 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()?, + merge_splits: false, + max_merge_input_streams: Some(MAX_MERGE_INPUT_STREAMS), + parquet_read_budget: Some(parquet_read_budget), + }, + ) + .read_with_merge_function( + &merge_splits, + crate::table::audit_log_table::merge::new_merge_function, + )?; + Box::pin(stream::select_all([raw_stream, merge_stream])) + } else { + let mut read = self.read.clone(); + read.read_type = read_type; + read.to_arrow(data_splits)? + }; + + Ok(Box::pin(async_stream::try_stream! { + futures::pin_mut!(physical_stream); + let mut projection = None; + while let Some(batch) = physical_stream.next().await { + let batch = batch?; + if projection.is_none() { + projection = Some(output_read_type.iter().map(|field| { + let name = if field.id() == ROW_KIND_FIELD_ID { + if !has_primary_keys { + return Ok(None); + } + VALUE_KIND_FIELD_NAME + } else { + field.name() + }; + batch.schema().index_of(name).map(Some).map_err(|error| crate::Error::DataInvalid { + message: format!("Audit read missing column '{name}': {error}"), + source: None, + }) + }).collect::>>()?); + } + let columns = output_read_type.iter().zip(projection.as_ref().unwrap()) + .map(|(field, index)| { + let column: ArrayRef = match index { + None => Arc::new(StringArray::from(vec!["+I"; batch.num_rows()])), + Some(index) if field.id() == ROW_KIND_FIELD_ID => + Arc::new(rowkind_array_from_column(batch.column(*index).as_ref())?), + Some(index) => batch.column(*index).clone(), + }; + Ok(column) + }).collect::>>()?; + 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)), + })?; + } + })) + } +} + +// 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) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::table::table_read::tests::{file, split}; + + #[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)); + } +} diff --git a/crates/paimon/src/table/audit_log_table/scan.rs b/crates/paimon/src/table/audit_log_table/scan.rs new file mode 100644 index 000000000..fa14c0f09 --- /dev/null +++ b/crates/paimon/src/table/audit_log_table/scan.rs @@ -0,0 +1,265 @@ +// 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 scan wrapper: retain row versions and only forward safe pruning. + +use super::{TableScan, TableScanKind}; +use crate::spec::MergeEngine; +use crate::table::kv_file_reader::retain_primary_key_conjuncts; +use crate::table::merge_tree_split_generator::{merge_tree_split_for_batch, KeyComparator}; +use crate::table::{DataSplit, Plan, ReadBuilder, ScanTrace}; +use std::collections::HashMap; + +/// Plans current-state audit reads over the wrapped table's files. +#[derive(Debug, Clone)] +pub struct AuditLogScan<'a> { + scan: TableScan<'a>, +} + +impl<'a> AuditLogScan<'a> { + fn new(mut scan: TableScan<'a>) -> Self { + if let TableScanKind::Paimon(inner) = &mut scan.0 { + // Preserve the read projection, including data-evolution column pruning. + inner.scan_all_files = true; + inner.limit = None; + inner.row_range_optimization_disabled = true; + inner.row_ranges = None; + if !inner.table.schema().primary_keys().is_empty() { + inner.data_predicates = retain_primary_key_conjuncts( + &inner.data_predicates, + inner.table.schema().fields(), + &inner.table.schema().trimmed_primary_keys(), + ); + } + } + Self { scan } + } + + pub async fn plan(&self) -> crate::Result { + self.plan_with_trace().await.map(|(plan, _)| plan) + } + + pub async fn plan_with_trace(&self) -> crate::Result<(Plan, ScanTrace)> { + let TableScanKind::Paimon(inner) = &self.scan.0 else { + return Err(crate::Error::Unsupported { + message: "Format tables do not support audit log batch scan".to_string(), + }); + }; + let (plan, mut trace) = self.scan.plan_with_trace().await?; + let options = inner.table.schema().core_options(); + let engine = options.merge_engine()?; + // Ordinary first-row and DV scans pack files without a key merge. Audit + // reads need the existing merge-tree planner to retain overlapping versions. + let plan = if engine == MergeEngine::FirstRow + || (options.deletion_vectors_enabled() && !options.deletion_vectors_merge_on_read()) + { + if let Some(comparator) = KeyComparator::from_table_schema(inner.table.schema()) { + let mut buckets: HashMap<(Vec, i32), Vec> = HashMap::new(); + for split in plan.into_splits() { + buckets + .entry((split.partition().to_serialized_bytes(), split.bucket())) + .or_default() + .push(split); + } + let mut splits = Vec::new(); + for bucket in buckets.into_values() { + let first = &bucket[0]; + let files = bucket + .iter() + .flat_map(|split| split.data_files().iter().cloned()) + .collect(); + let deletion_files: HashMap<_, _> = bucket + .iter() + .flat_map(|split| { + split.data_files().iter().filter_map(move |file| { + split + .deletion_file_for_data_file(file) + .map(|deletion| (file.file_name.clone(), deletion.clone())) + }) + }) + .collect(); + for group in merge_tree_split_for_batch( + files, + &comparator, + options.source_split_target_size(), + options.source_split_open_file_cost(), + matches!(engine, MergeEngine::Deduplicate | MergeEngine::FirstRow), + ) { + let mut builder = DataSplit::builder() + .with_snapshot(first.snapshot_id()) + .with_partition(first.partition().clone()) + .with_bucket(first.bucket()) + .with_bucket_path(first.bucket_path().to_string()) + .with_total_buckets(first.total_buckets()) + .with_raw_convertible(group.raw_convertible); + if !deletion_files.is_empty() { + builder = builder.with_data_deletion_files( + group + .files + .iter() + .map(|file| deletion_files.get(&file.file_name).cloned()) + .collect(), + ); + } + splits.push(builder.with_data_files(group.files).build()?); + } + } + Plan::new(splits) + } else { + plan + } + } else { + plan + }; + trace.record_final_plan( + plan.splits().len(), + plan.splits().len(), + plan.splits() + .iter() + .map(|split| split.data_files().len()) + .sum(), + ); + trace.planned_data_file_bytes = plan.planned_data_file_bytes(); + Ok((plan, trace)) + } +} + +impl<'a> ReadBuilder<'a> { + /// Create an audit scan that retains every visible row version. + pub fn new_audit_scan(&self) -> AuditLogScan<'a> { + AuditLogScan::new(self.new_scan()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::{ + BinaryRowBuilder, CoreOptions, DataField, DataType, Datum, IntType, PredicateBuilder, + }; + use crate::table::table_scan::{ + tests::{ + data_evolution_test_table, pk_stats_file, pk_stats_gate_table, setup_scan_trace_dirs, + two_column_schema, + }, + PaimonTableScan, + }; + use crate::table::{CommitMessage, TableCommit}; + use std::collections::HashSet; + #[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())); + let wrapped = AuditLogScan::new(TableScan(TableScanKind::Paimon(scan))); + let TableScanKind::Paimon(scan) = wrapped.scan.0 else { + panic!("expected Paimon scan") + }; + + assert!(scan.scan_all_files); + 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" + ); + } + + #[tokio::test] + async fn test_audit_stats_pruning_keeps_overlapping_versions() { + for option in ["deletion-vectors.enabled", "merge-engine"] { + let table_path = format!("memory:/audit_stats_{option}"); + let value = if option == "merge-engine" { + "first-row" + } else { + "true" + }; + let table = pk_stats_gate_table(&table_path).copy_with_options(HashMap::from([ + (option.to_string(), value.to_string()), + ("source.split.target-size".to_string(), "1b".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)); + // Compacted files on the same level have disjoint key ranges. + new.level = 2; + 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" + ); + assert_eq!( + audit_plan.splits().len(), + 1, + "overlapping versions must be planned together" + ); + assert!(!audit_plan.splits()[0].raw_convertible()); + } + } +} diff --git a/crates/paimon/src/table/kv_file_reader.rs b/crates/paimon/src/table/kv_file_reader.rs index 6ff644adc..3b0a31ce6 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, ConfiguredDeduplicateMergeFunction, DeduplicateMergeFunction, - FirstRowMergeFunction, PartialUpdateMergeFunction, SortMergeReaderBuilder, + AggregateMergeFunction, DeduplicateMergeFunction, MergeFunction, PartialUpdateMergeFunction, + SortMergeReaderBuilder, }; use crate::arrow::{build_target_arrow_schema, ParquetReadBudget}; use crate::deletion_vector::DeletionVectorFactory; use crate::io::FileIO; use crate::spec::{ - BigIntType, CoreOptions, DataField, DataFileMeta, DataType as PaimonDataType, MergeEngine, + BigIntType, 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,7 +49,6 @@ 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, @@ -64,7 +63,6 @@ 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, @@ -77,8 +75,6 @@ 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. @@ -286,48 +282,45 @@ impl KeyValueFileReader { self } - #[allow(clippy::too_many_arguments)] fn new_merge_function( - merge_engine: MergeEngine, - table_options: &HashMap, - table_name: &str, - table_fields: &[DataField], + config: &KeyValueReadConfig, 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, - ))) - } + ) -> crate::Result> { + match config.merge_engine { MergeEngine::Deduplicate => Ok(Box::new(DeduplicateMergeFunction)), MergeEngine::PartialUpdate => { Ok(Box::new(PartialUpdateMergeFunction::new_with_schema( - table_options, - table_name, - table_fields, + &config.table_options, + &config.table_name, + &config.table_fields, merge_output_fields, - primary_keys, + &config.primary_keys, )?)) } - MergeEngine::FirstRow => Ok(Box::new(FirstRowMergeFunction::new(table_options))), + 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::Aggregation => Ok(Box::new(AggregateMergeFunction::new( - table_options, - table_name, + &config.table_options, + &config.table_name, merge_output_fields, - primary_keys, - sequence_fields, + &config.primary_keys, + &config.sequence_fields, )?)), } } pub fn read(self, data_splits: &[DataSplit]) -> crate::Result { + self.read_with_merge_function(data_splits, Self::new_merge_function) + } + + pub(super) fn read_with_merge_function( + self, + data_splits: &[DataSplit], + merge_function: impl Fn(&KeyValueReadConfig, &[DataField]) -> crate::Result> + + Send + + 'static, + ) -> crate::Result { // A projected `_ROW_ID` is synthesized as all-nulls here, so the residual // would silently drop every row rather than hit its missing-column guard. super::row_id_predicate::reject_row_id_filter( @@ -382,17 +375,6 @@ 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. Physical system // fields are already the first two columns of every KV file. @@ -453,8 +435,8 @@ impl KeyValueFileReader { // Internal read type: [_SEQ, _VK, user_fields...] let mut internal_read_type: Vec = Vec::new(); - internal_read_type.push(seq_field.clone()); - internal_read_type.push(value_kind_field.clone()); + internal_read_type.push(seq_field); + internal_read_type.push(value_kind_field); internal_read_type.extend(user_fields.clone()); let internal_schema = build_target_arrow_schema(&internal_read_type)?; @@ -477,29 +459,20 @@ impl KeyValueFileReader { .unwrap() }) .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), - ); + let (value_indices, value_fields): (Vec<_>, Vec<_>) = internal_read_type + .iter() + .enumerate() + .filter(|(index, field)| { + !key_names.contains(field.name()) + && (*index >= 2 + || self + .config + .read_type + .iter() + .any(|requested| requested.id() == field.id())) + }) + .map(|(index, field)| (index, field.clone())) + .unzip(); // If sequence.field is configured, find each field's index in the internal schema. let user_sequence_indices: Vec = self @@ -548,24 +521,13 @@ impl KeyValueFileReader { data_splits.into_iter().map(|split| vec![split]).collect() }; let file_io = self.file_io; - let merge_engine = self.config.merge_engine; - let schema_manager = self.config.schema_manager; - let table_schema_id = self.config.table_schema_id; - let table_fields = self.config.table_fields; - let table_name = self.config.table_name; - let table_options = self.config.table_options; + let config = self.config; + let table_schema_id = config.table_schema_id; let pushdown_predicates = self.pushdown_predicates; - let residual_predicates = self.config.predicates; - 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)] let input_batch_sizes = self.input_batch_sizes; - // Build the merge output schema (keys + values, no system columns). + // Build the merge output schema (keys + projected values). let mut merge_output_fields: Vec = Vec::new(); merge_output_fields.extend(key_fields); merge_output_fields.extend(value_fields); @@ -604,13 +566,13 @@ impl KeyValueFileReader { merge_splits, ) { let input_stream_count = merge_group.len(); - ensure_merge_input_limit(input_stream_count, max_merge_input_streams)?; + ensure_merge_input_limit(input_stream_count, config.max_merge_input_streams)?; // Sort-merge must first obtain one batch from every input // stream. Keep concurrent row-group reads disabled whenever // multiple runs advance in lockstep; one run may still use // the shared budget because its files are opened serially. let group_parquet_read_budget = if input_stream_count == 1 { - parquet_read_budget.clone() + config.parquet_read_budget.clone() } else { None }; @@ -619,15 +581,15 @@ impl KeyValueFileReader { for MergeRun { files } in merge_group { let reader = DataFileReader::new( file_io.clone(), - schema_manager.clone(), + config.schema_manager.clone(), table_schema_id, - table_fields.clone(), + config.table_fields.clone(), internal_read_type.clone(), pushdown_predicates.clone(), ) - .with_batch_size(Some(read_batch_size)) + .with_batch_size(Some(config.read_batch_size)) .with_parquet_read_budget(group_parquet_read_budget.clone()); - let run_schema_manager = schema_manager.clone(); + let run_schema_manager = config.schema_manager.clone(); let run_file_io = file_io.clone(); let deletion_files_by_split = deletion_files_by_split.clone(); let run_stream: ArrowRecordBatchStream = Box::pin(try_stream! { @@ -689,16 +651,7 @@ impl KeyValueFileReader { user_sequence_indices.clone(), value_indices.clone(), merge_output_schema.clone(), - Self::new_merge_function( - merge_engine, - &table_options, - &table_name, - &table_fields, - &merge_output_fields, - &primary_keys, - &sequence_fields, - keep_delete, - )?, + merge_function(&config, &merge_output_fields)?, ) .build()?; @@ -712,13 +665,13 @@ impl KeyValueFileReader { // the merge-output batch (keys + values, including widened // predicate columns); the reorder below projects the output // back to read_type. - let batch = if residual_predicates.is_empty() { + let batch = if config.predicates.is_empty() { batch } else { match crate::arrow::residual::evaluate_predicates_mask( &batch, - &residual_predicates, - &table_fields, + &config.predicates, + &config.table_fields, &merge_output_fields, )? { Some(mask) => arrow_select::filter::filter_record_batch( @@ -1358,7 +1311,6 @@ 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), @@ -1478,7 +1430,6 @@ 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, @@ -1686,7 +1637,6 @@ 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, @@ -1758,7 +1708,6 @@ 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())), @@ -1953,7 +1902,6 @@ 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, @@ -2021,7 +1969,6 @@ 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 3de32d737..370e600df 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -171,8 +171,8 @@ pub use source::{ merge_row_ranges, DataSplit, DataSplitBuilder, DeletionFile, PartitionBucket, Plan, RowRange, }; pub use table_commit::TableCommit; -pub use table_read::{AuditLogInput, AuditLogRead, TableRead}; -pub use table_scan::TableScan; +pub use table_read::{AuditLogRead, TableRead}; +pub use table_scan::{AuditLogScan, TableScan}; pub use table_update::TableUpdate; pub use table_write::TableWrite; pub use tag_manager::TagManager; diff --git a/crates/paimon/src/table/read_builder.rs b/crates/paimon/src/table/read_builder.rs index 52dc1cea9..ec8ef966e 100644 --- a/crates/paimon/src/table/read_builder.rs +++ b/crates/paimon/src/table/read_builder.rs @@ -505,9 +505,10 @@ 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 projection = self.resolve_read_type()?; - let explicit_projection = projection.is_some(); - let read_type = projection.unwrap_or_else(|| self.table.schema.fields().to_vec()); + let read_type = match self.resolve_read_type()? { + None => self.table.schema.fields().to_vec(), + Some(fields) => fields, + }; // Pass the FULL data predicate through (including `And`/`Or`/`Not`). // Pushdown/stats skip compound nodes; the residual pass enforces the full @@ -518,7 +519,6 @@ 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 81f37c3ea..d1f2bdee8 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::{Arc, Mutex, OnceLock}; +use std::sync::Mutex; // --------------------------------------------------------------------------- // MergeFunction @@ -141,40 +141,7 @@ 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 { +pub(super) fn compare_sequence_order(lhs: &MergeRow, rhs: &MergeRow) -> Ordering { match (lhs.user_sequences.is_empty(), rhs.user_sequences.is_empty()) { (false, false) => lhs .user_sequences @@ -184,32 +151,6 @@ 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, @@ -218,51 +159,26 @@ impl MergeFunction for DeduplicateMergeFunction { _source_output_col_indices: &[usize], _output_schema: &SchemaRef, ) -> crate::Result { - 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; + 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 } - 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); - } + }) + .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) } - Ok(match first { - Some(row) => MergeResult::SourceRow { - batch_idx: row.batch_idx, - row_idx: row.row_idx, - }, - None => MergeResult::Omit, - }) } } @@ -278,7 +194,6 @@ impl MergeFunction for FirstRowMergeFunction { #[derive(Debug)] pub(crate) struct PartialUpdateMergeFunction { ignore_delete: bool, - value_kind_index: Option, sequence_groups: Vec, grouped_fields: HashSet, aggregators: Option>, @@ -301,7 +216,6 @@ 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, @@ -388,9 +302,6 @@ 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 @@ -453,9 +364,7 @@ impl MergeFunction for PartialUpdateMergeFunction { saw_add = true; for (output_col_idx, selected) in selected_by_col.iter_mut().enumerate() { - if self.value_kind_index == Some(output_col_idx) - || self.grouped_fields.contains(&output_col_idx) - { + if self.grouped_fields.contains(&output_col_idx) { continue; } let source_array = batch_buffer[row.batch_idx] @@ -533,22 +442,18 @@ impl MergeFunction for PartialUpdateMergeFunction { .iter() .enumerate() .map(|(output_col_idx, field)| { - 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), - }, - } + 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), + }, }; if !field.is_nullable() && column.is_null(0) { return Err(Error::DataInvalid { @@ -646,7 +551,6 @@ 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 { @@ -676,12 +580,7 @@ impl AggregateMergeFunction { .iter() .map(|field| -> crate::Result>> { let name = field.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) - { + let agg_name: &str = if seq_set.contains(name) { "last_value" } else if pk_set.contains(name) { return Ok(None); @@ -703,9 +602,6 @@ impl AggregateMergeFunction { Ok(Self { aggregators: Mutex::new(aggregators), - value_kind_index: output_fields - .iter() - .position(|field| field.id() == crate::spec::VALUE_KIND_FIELD_ID), }) } } @@ -777,9 +673,6 @@ 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 83cb332dc..b1926cb56 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -25,10 +25,15 @@ use super::{ArrowRecordBatchStream, Table}; use crate::arrow::build_target_arrow_schema; use crate::arrow::ParquetReadBudget; use crate::spec::{ - CoreOptions, DataField, DataType, MergeEngine, Predicate, SEQUENCE_NUMBER_FIELD_NAME, + 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, }; use crate::DataSplit; -use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, UInt32Array}; +use arrow_array::{ + builder::StringBuilder, Array, ArrayRef, RecordBatch, RecordBatchOptions, StringArray, + UInt32Array, +}; use arrow_schema::Schema as ArrowSchema; use arrow_select::concat::concat as arrow_concat; use arrow_select::take::take; @@ -36,8 +41,9 @@ use futures::{stream, StreamExt}; use std::cmp::Ordering; use std::sync::Arc; +#[path = "audit_log_table/read.rs"] mod audit; -pub use audit::{AuditLogInput, AuditLogRead}; +pub use audit::AuditLogRead; const MAX_MERGE_INPUT_STREAMS: usize = 256; @@ -81,14 +87,6 @@ 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, @@ -199,6 +197,26 @@ 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() } @@ -208,7 +226,6 @@ 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>, @@ -225,7 +242,6 @@ impl<'a> PaimonTableRead<'a> { Self { table, read_type, - explicit_projection: false, data_predicates, row_filter_factory: None, parquet_read_budget: None, @@ -345,6 +361,237 @@ 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], @@ -460,7 +707,6 @@ 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 @@ -602,7 +848,6 @@ 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()) @@ -669,6 +914,70 @@ 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)] @@ -682,7 +991,6 @@ enum CursorOrd { struct ArrowCursor { stream: ArrowRecordBatchStream, batch: Option, - batch_id: usize, row: usize, } @@ -691,7 +999,6 @@ impl ArrowCursor { let mut cursor = Self { stream, batch: None, - batch_id: 0, row: 0, }; cursor.advance().await?; @@ -710,10 +1017,6 @@ impl ArrowCursor { self.row } - fn batch_id(&self) -> usize { - self.batch_id - } - async fn advance(&mut self) -> crate::Result<()> { loop { if let Some(ref batch) = self.batch { @@ -724,7 +1027,6 @@ 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(()); @@ -740,6 +1042,96 @@ 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)>, @@ -840,6 +1232,34 @@ 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> = @@ -1225,8 +1645,9 @@ mod tests { .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) + let current_audit = AuditLogRead::new(pk_read) + .unwrap() + .to_arrow(&splits) .unwrap() .try_collect::>() .await @@ -1271,6 +1692,25 @@ 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 deleted file mode 100644 index 0336ffc17..000000000 --- a/crates/paimon/src/table/table_read/audit.rs +++ /dev/null @@ -1,976 +0,0 @@ -// 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, 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 arrow_select::interleave::interleave; -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).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", (0, bc.batch_id()), bc.batch(), bc.row()); - bc.advance().await?; - } - CursorOrd::AfterOnly => { - builder.push("+I", (1, ac.batch_id()), ac.batch(), ac.row()); - ac.advance().await?; - } - CursorOrd::EqualSame => { - bc.advance().await?; - ac.advance().await?; - } - CursorOrd::EqualDiff => { - builder.push("-U", (0, bc.batch_id()), bc.batch(), bc.row()); - builder.push("+U", (1, 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 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_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 3ed67850f..3835e4b8e 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -56,6 +56,10 @@ use indexmap::IndexMap; use std::collections::{HashMap, HashSet}; use std::sync::Arc; +#[path = "audit_log_table/scan.rs"] +mod audit; +pub use audit::AuditLogScan; + /// Path segment for manifest directory under table. const MANIFEST_DIR: &str = "manifest"; /// Path segment for index directory under table. @@ -909,15 +913,6 @@ 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) => { @@ -1030,8 +1025,6 @@ 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>, } @@ -1053,7 +1046,6 @@ impl<'a> PaimonTableScan<'a> { row_ranges, row_range_optimization_disabled: false, scan_all_files: false, - merge_key_overlaps: false, projected_read_field_ids: None, } } @@ -1068,12 +1060,6 @@ 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 @@ -1289,7 +1275,7 @@ impl<'a> PaimonTableScan<'a> { } fn can_push_down_limit_hint(&self, row_ranges: Option<&[RowRange]>) -> bool { - !self.scan_all_files && can_push_down_limit_hint_for_scan(&self.data_predicates, row_ranges) + can_push_down_limit_hint_for_scan(&self.data_predicates, row_ranges) } fn global_index_scan_settings( @@ -1297,14 +1283,12 @@ impl<'a> PaimonTableScan<'a> { core_options: &CoreOptions, data_evolution_enabled: bool, ) -> crate::Result> { - 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(), - ) - { + 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(), + ) { Ok(Some(GlobalIndexScanSettings { search_mode: core_options.scalar_index_search_mode()?, thread_num: core_options.global_index_thread_num()?, @@ -1439,14 +1423,15 @@ impl<'a> PaimonTableScan<'a> { /// `KeyValueFileReader`. /// /// Exempt (full predicates kept): - /// - Ordinary deletion-vector reads without merge-on-read: they read raw with + /// - Deletion-vector tables 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. - /// - Non-audit `merge-engine=first-row` reads: read via `DataFileReader` - /// without merging versions. - /// - /// Audit reads set `merge_key_overlaps` and are not exempt. + /// - `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. 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()); @@ -1459,8 +1444,8 @@ impl<'a> PaimonTableScan<'a> { Ok(crate::spec::MergeEngine::FirstRow) ); if has_primary_keys - && (self.merge_key_overlaps - || ((!deletion_vectors_enabled || deletion_vectors_merge_on_read) && !first_row)) + && (!deletion_vectors_enabled || deletion_vectors_merge_on_read) + && !first_row { retain_primary_key_conjuncts( &self.data_predicates, @@ -1935,17 +1920,16 @@ 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 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) - )); + // 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) + ); let pk_comparator = if read_merges_overlapping_keys { KeyComparator::from_table_schema(self.table.schema()) } else { @@ -2076,9 +2060,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. Ordinary first-row - // scans do not take this path, but audit scans do. Its writer - // deduplicates at flush, so keep the gate accurate. + // 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.) let file_keys_unique = matches!( core_options.merge_engine(), Ok(crate::spec::MergeEngine::Deduplicate) @@ -2194,10 +2178,10 @@ mod tests { use crate::io::FileIOBuilder; use crate::spec::{ stats::BinaryTableStats, ArrayType, BinaryRow, BinaryRowBuilder, BucketFunctionType, - 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, + ColumnMove, CommitKind, 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; @@ -2436,7 +2420,7 @@ mod tests { ); } - fn data_evolution_test_table(table_path: &str, schema: TableSchema) -> Table { + pub(super) fn data_evolution_test_table(table_path: &str, schema: TableSchema) -> Table { let file_io = FileIOBuilder::new("memory").build().unwrap(); let schema = schema.copy_with_options(HashMap::from([( "data-evolution.enabled".to_string(), @@ -2451,7 +2435,7 @@ mod tests { ) } - fn two_column_schema(id: i64, left: &str, right: &str) -> TableSchema { + pub(super) fn two_column_schema(id: i64, left: &str, right: &str) -> TableSchema { TableSchema::new( id, &PaimonSchema::builder() @@ -2652,7 +2636,7 @@ mod tests { ])) } - async fn setup_scan_trace_dirs(table: &Table) { + pub(super) async fn setup_scan_trace_dirs(table: &Table) { table .file_io() .mkdirs(&format!("{}/snapshot/", table.location())) @@ -2891,35 +2875,6 @@ 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( @@ -3579,7 +3534,7 @@ mod tests { ); } - fn pk_stats_gate_table(table_path: &str) -> Table { + pub(super) fn pk_stats_gate_table(table_path: &str) -> Table { let file_io = FileIOBuilder::new("memory").build().unwrap(); let schema = PaimonSchema::builder() .column("id", DataType::Int(IntType::new())) @@ -3648,7 +3603,11 @@ mod tests { builder.build_serialized() } - fn pk_stats_file(name: &str, id_range: (i32, i32), value_range: (i32, i32)) -> DataFileMeta { + pub(super) fn pk_stats_file( + name: &str, + id_range: (i32, i32), + value_range: (i32, i32), + ) -> DataFileMeta { let mut file = test_data_file_meta( two_int_stats_row(Some(id_range.0), Some(value_range.0)), two_int_stats_row(Some(id_range.1), Some(value_range.1)), @@ -3768,68 +3727,12 @@ mod tests { ); } - #[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. + /// `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_first_row_table_stats_pruning_keeps_non_key_conjuncts() { let table_path = "memory:/test_first_row_stats_gate"; @@ -3887,18 +3790,6 @@ 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 6edaa0127..3da98e1e6 100644 --- a/crates/paimon/tests/audit_log_table_test.rs +++ b/crates/paimon/tests/audit_log_table_test.rs @@ -455,7 +455,11 @@ async fn audit_log_current_scan_uses_merged_rowkind() { 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")]), + pk_schema(&[ + ("merge-engine", merge_engine), + ("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; @@ -478,8 +482,8 @@ async fn audit_log_current_scan_uses_merged_rowkind() { .await .unwrap(); assert_eq!( - collect_audit_rows(&batches), - vec![("+I".to_string(), 1, 20)], + collect_audit_rows_with_sequence(&batches), + vec![("+I".to_string(), 1, 1, 20)], "merge-engine={merge_engine}" ); } From 0b49446b97616d5a195b40fd75b542d87527f7ba Mon Sep 17 00:00:00 2001 From: yantian Date: Mon, 14 Sep 2026 11:49:22 +0800 Subject: [PATCH 15/15] refactor: remove optional audit scan hooks --- .../datafusion/src/physical_plan/audit_log.rs | 54 +---- .../datafusion/src/physical_plan/scan.rs | 222 ++++++++---------- crates/paimon/src/table/audit_log_table.rs | 12 +- crates/paimon/tests/audit_log_table_test.rs | 16 +- 4 files changed, 128 insertions(+), 176 deletions(-) diff --git a/crates/integrations/datafusion/src/physical_plan/audit_log.rs b/crates/integrations/datafusion/src/physical_plan/audit_log.rs index f98098ec8..5fdfeccb1 100644 --- a/crates/integrations/datafusion/src/physical_plan/audit_log.rs +++ b/crates/integrations/datafusion/src/physical_plan/audit_log.rs @@ -20,19 +20,14 @@ 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. +/// Retains retract rows; physical filters remain above this scan. #[derive(Debug, Clone)] pub(crate) struct PaimonAuditLogScan { inner: PaimonTableScan, @@ -64,34 +59,6 @@ impl ExecutionPlan for PaimonAuditLogScan { 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, @@ -112,10 +79,11 @@ impl ExecutionPlan for PaimonAuditLogScan { impl DisplayAs for PaimonAuditLogScan { fn fmt_as( &self, - _t: datafusion::physical_plan::DisplayFormatType, + t: datafusion::physical_plan::DisplayFormatType, f: &mut std::fmt::Formatter, ) -> std::fmt::Result { - self.inner.fmt_scan(self.name(), f) + write!(f, "{}: ", self.name())?; + self.inner.fmt_as(t, f) } } @@ -123,10 +91,13 @@ impl DisplayAs for PaimonAuditLogScan { mod tests { use super::*; use crate::table::{datafusion_arrow_schema, PaimonScanBuilder}; + use datafusion::config::ConfigOptions; 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 datafusion::physical_plan::filter_pushdown::{ + ChildFilterPushdownResult, ChildPushdownResult, FilterPushdownPhase, PushedDown, + }; use paimon::catalog::Identifier; use paimon::table::Table; use paimon::DataSplitBuilder; @@ -195,7 +166,7 @@ mod tests { } #[test] - fn test_audit_policy_survives_filter_pushdown() { + fn test_audit_physical_filters_remain_above_scan() { let scan = first_row_audit_scan(); let filters: Vec> = vec![ Arc::new(BinaryExpr::new( @@ -228,12 +199,11 @@ mod tests { assert!(matches!( result.filters.as_slice(), - [PushedDown::Yes, PushedDown::No] + [PushedDown::No, PushedDown::No] )); - let updated = result.updated_node.unwrap(); - assert!(updated.downcast_ref::().is_some()); + assert!(result.updated_node.is_none()); assert_eq!( - updated.partition_statistics(None).unwrap().num_rows, + scan.partition_statistics(None).unwrap().num_rows, Precision::Absent ); } diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs b/crates/integrations/datafusion/src/physical_plan/scan.rs index 0fc146a4f..432e3a2a0 100644 --- a/crates/integrations/datafusion/src/physical_plan/scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/scan.rs @@ -997,71 +997,6 @@ impl PaimonTableScan { .collect() } - pub(crate) fn pushdown_filters( - &self, - child_pushdown_result: ChildPushdownResult, - supported: impl Fn(&Arc) -> bool, - ) -> DFResult> { - let filters = child_pushdown_result - .parent_filters - .into_iter() - .map(|result| result.filter) - .collect::>(); - if filters.is_empty() { - return Ok(FilterPushdownPropagation::with_parent_pushdown_result( - Vec::new(), - )); - } - let schema = self.schema(); - let mut accepted = Vec::new(); - let parent_filter_handled = filters - .into_iter() - .map(|filter| { - 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. - PushedDown::Yes - } else { - PushedDown::No - } - }) - .collect::>(); - if accepted.is_empty() { - return Ok(FilterPushdownPropagation::with_parent_pushdown_result( - parent_filter_handled, - )); - } - - let mut scan = self.clone(); - for filter in accepted { - scan.decoder_filters.extend( - split_conjunction(&filter) - .into_iter() - .filter(|conjunct| { - !paimon_predicate_covers_filter( - self.pushed_predicate.as_ref(), - conjunct, - self.table.schema().fields(), - self.case_sensitive, - ) && !reads_partition_column_absent_from_files( - conjunct, - &self.table, - self.case_sensitive, - ) - }) - .cloned(), - ); - scan.runtime_filters.push(filter); - } - Ok( - FilterPushdownPropagation::with_parent_pushdown_result(parent_filter_handled) - .with_updated_node(scan), - ) - } - pub(crate) fn execute_with( &self, partition: usize, @@ -1113,8 +1048,9 @@ impl PaimonTableScan { let stream = read_splits(read, &splits).map_err(to_datafusion_error)?; let batch_schema = Arc::clone(&schema); let stream = stream.map(move |result| { - let batch = result.map_err(to_datafusion_error)?; - let mut batch = to_datafusion_batch(batch, &batch_schema)?; + let mut batch = result + .map_err(to_datafusion_error) + .and_then(|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 @@ -1146,51 +1082,6 @@ impl PaimonTableScan { futures::stream::once(fut).try_flatten(), ))) } - - 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 { @@ -1219,13 +1110,63 @@ impl ExecutionPlan for PaimonTableScan { 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), - }) + let filters = child_pushdown_result + .parent_filters + .into_iter() + .map(|result| result.filter) + .collect::>(); + if filters.is_empty() { + return Ok(FilterPushdownPropagation::with_parent_pushdown_result( + 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()) { + accepted.push(filter); + // This scan evaluates accepted expressions exactly, so the + // parent FilterExec can be removed. + PushedDown::Yes + } else { + PushedDown::No + } + }) + .collect::>(); + if accepted.is_empty() { + return Ok(FilterPushdownPropagation::with_parent_pushdown_result( + parent_filter_handled, + )); + } + + let mut scan = self.clone(); + for filter in accepted { + scan.decoder_filters.extend( + split_conjunction(&filter) + .into_iter() + .filter(|conjunct| { + !paimon_predicate_covers_filter( + self.pushed_predicate.as_ref(), + conjunct, + self.table.schema().fields(), + self.case_sensitive, + ) && !reads_partition_column_absent_from_files( + conjunct, + &self.table, + self.case_sensitive, + ) + }) + .cloned(), + ); + scan.runtime_filters.push(filter); + } + Ok( + FilterPushdownPropagation::with_parent_pushdown_result(parent_filter_handled) + .with_updated_node(Arc::new(scan)), + ) } fn execute( @@ -1283,7 +1224,48 @@ impl DisplayAs for PaimonTableScan { _t: datafusion::physical_plan::DisplayFormatType, f: &mut std::fmt::Formatter, ) -> std::fmt::Result { - self.fmt_scan(self.name(), f) + 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(()) } } diff --git a/crates/paimon/src/table/audit_log_table.rs b/crates/paimon/src/table/audit_log_table.rs index f943188f5..f3d604df3 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, AuditLogRead, AuditLogScan, DataSplit, Table}; +use super::{ArrowRecordBatchStream, AuditLogRead, AuditLogScan, Table}; use crate::spec::{ BigIntType, DataField, DataType, VarCharType, ROW_KIND_FIELD_ID, ROW_KIND_FIELD_NAME, SEQUENCE_NUMBER_FIELD_ID, SEQUENCE_NUMBER_FIELD_NAME, @@ -82,7 +82,7 @@ impl AuditLogTable { IncrementalScan::for_table(&self.wrapped, mode, start_exclusive, end_inclusive) } - /// Plan a current-state audit read for [`Self::to_arrow_for_splits`]. + /// Plan a current-state audit read for [`Self::new_read`]. pub fn new_scan(&self) -> AuditLogScan<'_> { self.wrapped.new_read_builder().new_audit_scan() } @@ -102,12 +102,4 @@ 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.new_read()?.to_arrow(splits) - } } diff --git a/crates/paimon/tests/audit_log_table_test.rs b/crates/paimon/tests/audit_log_table_test.rs index 3da98e1e6..94a31d43b 100644 --- a/crates/paimon/tests/audit_log_table_test.rs +++ b/crates/paimon/tests/audit_log_table_test.rs @@ -401,7 +401,9 @@ async fn audit_log_current_scan_keeps_delete_and_sequence_number() { let plan = table.new_read_builder().new_scan().plan().await.unwrap(); let batches: Vec = AuditLogTable::new(table) - .to_arrow_for_splits(plan.splits()) + .new_read() + .unwrap() + .to_arrow(plan.splits()) .unwrap() .try_collect() .await @@ -476,7 +478,9 @@ async fn audit_log_current_scan_uses_merged_rowkind() { let plan = table.new_read_builder().new_scan().plan().await.unwrap(); let batches: Vec = AuditLogTable::new(table) - .to_arrow_for_splits(plan.splits()) + .new_read() + .unwrap() + .to_arrow(plan.splits()) .unwrap() .try_collect() .await @@ -515,7 +519,9 @@ async fn audit_log_current_scan_respects_ignore_delete() { let plan = table.new_read_builder().new_scan().plan().await.unwrap(); let batches: Vec = AuditLogTable::new(table) - .to_arrow_for_splits(plan.splits()) + .new_read() + .unwrap() + .to_arrow(plan.splits()) .unwrap() .try_collect() .await @@ -570,7 +576,9 @@ async fn audit_log_current_scan_supports_first_row() { let mut rows = Vec::new(); for split in plan.splits() { let batches: Vec = audit - .to_arrow_for_splits(std::slice::from_ref(split)) + .new_read() + .unwrap() + .to_arrow(std::slice::from_ref(split)) .unwrap() .try_collect() .await