diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index e5093f140..e619d2170 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -440,6 +440,22 @@ impl PaimonCatalogProvider { self } + pub(crate) fn catalog(&self) -> &Arc { + &self.catalog + } + + /// The session's dynamic options as they stand now. + pub(crate) fn session_options(&self) -> HashMap { + self.dynamic_options + .read() + .unwrap_or_else(|e| e.into_inner()) + .clone() + } + + pub(crate) fn blob_reader_registry(&self) -> &BlobReaderRegistry { + &self.blob_reader_registry + } + /// Register an engine for a table type the Paimon reader cannot serve /// (e.g. [`PaimonTableType::IcebergTable`]); everything else takes the /// Paimon path unchanged. Kept inside the provider so the registered diff --git a/crates/integrations/datafusion/src/incremental_query.rs b/crates/integrations/datafusion/src/incremental_query.rs new file mode 100644 index 000000000..3ac1a40d6 --- /dev/null +++ b/crates/integrations/datafusion/src/incremental_query.rs @@ -0,0 +1,511 @@ +// 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. + +//! The `paimon_incremental_query` table function reads a snapshot-id range +//! `(start_exclusive, end_inclusive]` as a batch. Snapshot planning stays in +//! paimon core; this parses arguments, picks the core planner that matches +//! Java's `incremental-between` for the table, and streams its output. +//! +//! ```sql +//! SELECT * FROM paimon_incremental_query('db.t', 0, 5); +//! SELECT * FROM paimon_incremental_query('db.t', 0, 5, 'delta'); +//! SELECT * FROM paimon_incremental_query('db.t$audit_log', 0, 5, 'changelog'); +//! ``` +//! +//! A `$audit_log` table-name suffix prepends the `rowkind` audit column (and the +//! optional `_SEQUENCE_NUMBER`), mirroring the core audit-log read. +//! +//! The table name resolves as in a `FROM` clause, in the calling session, and +//! the session's dynamic options apply as they do to an ordinary read of the +//! table, except those that choose another range to read. + +use std::collections::HashMap; +use std::fmt; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; +use datafusion::catalog::{CatalogProviderList, Session, TableFunctionArgs, TableFunctionImpl}; +use datafusion::common::stats::Precision; +use datafusion::common::{internal_err, project_schema, Statistics}; +use datafusion::datasource::{TableProvider, TableType}; +use datafusion::error::{DataFusionError, Result as DFResult}; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::logical_expr::Expr; +use datafusion::physical_expr::EquivalenceProperties; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, +}; +use datafusion::prelude::SessionContext; +use futures::{stream, StreamExt, TryStreamExt}; +use paimon::catalog::DEFAULT_MAIN_BRANCH; +use paimon::resource::ResourceContext; +use paimon::spec::{CoreOptions, DataField}; +use paimon::table::{ + ArrowRecordBatchStream, AuditLogTable, IncrementalPlan, IncrementalScanMode, IncrementalSplit, + Table, +}; + +use crate::catalog::PaimonCatalogProvider; +use crate::error::to_datafusion_error; +use crate::memory::reader_resources; +use crate::physical_plan::scan::to_datafusion_batch; +use crate::runtime::{await_with_runtime, block_on_with_runtime}; +use crate::system_tables::ensure_audit_log_dynamic_options; +use crate::table::datafusion_arrow_schema; +use crate::table_function_args::{ + extract_int_literal, extract_string_literal, resolve_table_argument, +}; +use crate::table_loader::load_data_table_for_read; + +const FUNCTION_NAME: &str = "paimon_incremental_query"; +const AUDIT_LOG_SUFFIX: &str = "$audit_log"; + +/// Registers `paimon_incremental_query`, which looks its table up among the +/// catalogs of `ctx` when it is called. +pub(crate) fn register_incremental_query(ctx: &SessionContext) { + ctx.register_udtf( + FUNCTION_NAME, + Arc::new(IncrementalQueryFunction { + catalogs: Arc::clone(ctx.state().catalog_list()), + }), + ); +} + +pub(crate) struct IncrementalQueryFunction { + /// The catalog list `register_catalog` adds to, as ordinary queries see it. + catalogs: Arc, +} + +impl fmt::Debug for IncrementalQueryFunction { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("IncrementalQueryFunction") + .finish_non_exhaustive() + } +} + +/// Parses the optional mode argument. `Auto` resolves from `changelog-producer` +/// during planning; `Diff` is handled by core's diff planning. +fn parse_mode(raw: &str) -> DFResult { + match raw.to_ascii_lowercase().as_str() { + "auto" => Ok(IncrementalScanMode::Auto), + "delta" => Ok(IncrementalScanMode::Delta), + "changelog" => Ok(IncrementalScanMode::Changelog), + "diff" => Ok(IncrementalScanMode::Diff), + other => Err(DataFusionError::Plan(format!( + "{FUNCTION_NAME}: unknown scan mode '{other}', expected auto, delta, changelog, or diff" + ))), + } +} + +impl TableFunctionImpl for IncrementalQueryFunction { + fn call_with_args(&self, args: TableFunctionArgs) -> DFResult> { + let exprs = args.exprs(); + if exprs.len() < 3 || exprs.len() > 4 { + return Err(DataFusionError::Plan(format!( + "{FUNCTION_NAME} requires 3 or 4 arguments: (table_name, start_snapshot_exclusive, end_snapshot_inclusive [, mode])" + ))); + } + let table_name = extract_string_literal(FUNCTION_NAME, &exprs[0], "table_name")?; + let start_exclusive = + extract_int_literal(FUNCTION_NAME, &exprs[1], "start_snapshot_exclusive")?; + let end_inclusive = + extract_int_literal(FUNCTION_NAME, &exprs[2], "end_snapshot_inclusive")?; + let mode = match exprs.get(3) { + Some(expr) => parse_mode(&extract_string_literal(FUNCTION_NAME, expr, "mode")?)?, + None => IncrementalScanMode::Auto, + }; + if end_inclusive < start_exclusive { + return Err(DataFusionError::Plan(format!( + "{FUNCTION_NAME}: end_snapshot_inclusive ({end_inclusive}) must be >= start_snapshot_exclusive ({start_exclusive})" + ))); + } + + // `db.t$audit_log` selects the audit-log projection of `db.t`. + let (base_name, audit_log) = match table_name.strip_suffix(AUDIT_LOG_SUFFIX) { + Some(base) => (base.to_string(), true), + None => (table_name.clone(), false), + }; + // The calling session resolves the name, so `USE` applies here as it + // does to a `FROM` clause. + let (catalog_name, identifier) = resolve_table_argument( + FUNCTION_NAME, + &base_name, + &args.session().config_options().catalog, + )?; + let registered = self + .catalogs + .catalog(&catalog_name) + .ok_or_else(|| DataFusionError::Plan(format!("Unknown catalog '{catalog_name}'")))?; + let provider = registered + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Plan(format!("Catalog '{catalog_name}' is not a Paimon catalog")) + })?; + + // Session options apply as they do to an ordinary read of the table; + // those choosing another range are refused before anything is read. + let session_options = provider.session_options(); + ensure_no_range_options(&session_options, |key| { + format!( + "{FUNCTION_NAME} reads the range its arguments give and cannot run while \ + scan option '{key}' is set; RESET 'paimon.{key}' first" + ) + })?; + if audit_log { + ensure_audit_log_dynamic_options(&session_options)?; + } + let catalog = Arc::clone(provider.catalog()); + let table = block_on_with_runtime( + async move { + let table = load_data_table_for_read(&catalog, &identifier, FUNCTION_NAME).await?; + let table = if session_options.is_empty() { + table + } else { + table + .copy_with_time_travel(session_options) + .await + .map_err(to_datafusion_error)? + }; + // Checked on the table read: a session `branch` reads that + // branch, with that branch's options. + ensure_no_range_options(table.schema().options(), |key| { + let read = match table.branch() { + DEFAULT_MAIN_BRANCH => format!("table '{}'", identifier.full_name()), + branch => { + format!("branch '{branch}' of table '{}'", identifier.full_name()) + } + }; + format!( + "{FUNCTION_NAME} reads the range its arguments give and cannot run on \ + {read}, which sets scan option '{key}'" + ) + })?; + Ok::<_, DataFusionError>(table) + }, + "paimon_incremental_query: catalog access thread panicked", + )?; + // As an ordinary read does, so BLOB descriptors it returns can be read. + provider + .blob_reader_registry() + .register_if_absent(table.location().to_string(), table.file_io().clone()); + + Ok(Arc::new(IncrementalQueryTableProvider::try_new( + table, + audit_log, + mode, + start_exclusive, + end_inclusive, + )?)) + } +} + +/// Fails with `message(key)` when `options` choose what to read, which the +/// function's arguments do. +fn ensure_no_range_options( + options: &HashMap, + message: impl FnOnce(&str) -> String, +) -> DFResult<()> { + match CoreOptions::new(options) + .configured_scan_range_options() + .first() + { + Some(key) => Err(DataFusionError::Plan(message(key))), + None => Ok(()), + } +} + +#[derive(Debug, Clone)] +struct IncrementalQueryTableProvider { + table: Table, + audit_log: bool, + mode: IncrementalScanMode, + start_exclusive: i64, + end_inclusive: i64, + fields: Vec, + schema: ArrowSchemaRef, +} + +impl IncrementalQueryTableProvider { + fn try_new( + table: Table, + audit_log: bool, + mode: IncrementalScanMode, + start_exclusive: i64, + end_inclusive: i64, + ) -> DFResult { + // Audit-log output leads with `rowkind` (+ optional `_SEQUENCE_NUMBER`); + // the plain projection is the table's own fields. + let fields = if audit_log { + AuditLogTable::new(table.clone()) + .fields() + .map_err(to_datafusion_error)? + } else { + table.schema().fields().to_vec() + }; + // Core streams raw Arrow (`Utf8`, not `Utf8View`), so declare the schema + // the reader actually yields rather than the view-forced scan schema. + let schema = datafusion_arrow_schema(&fields, false)?; + Ok(Self { + table, + audit_log, + mode, + start_exclusive, + end_inclusive, + fields, + schema, + }) + } + + /// Splits a projection into the table fields core reads and the column of + /// core's batches each output column comes from. Core's audit-log batches + /// lead with the same audit columns as `fields`, then the read fields. + fn read_projection(&self, indices: &[usize]) -> ReadProjection { + let audit_columns = self.fields.len() - self.table.schema().fields().len(); + let mut read_type = Vec::new(); + let columns = indices + .iter() + .map(|&index| { + if index < audit_columns { + return index; + } + read_type.push(self.fields[index].clone()); + audit_columns + read_type.len() - 1 + }) + .collect(); + ReadProjection { read_type, columns } + } +} + +#[async_trait] +impl TableProvider for IncrementalQueryTableProvider { + fn schema(&self) -> ArrowSchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + limit: Option, + ) -> DFResult> { + let projected_schema = project_schema(&self.schema, projection)?; + // An outer `LIMIT 0` needs no rows. + if limit == Some(0) { + return Ok(Arc::new(EmptyExec::new(projected_schema))); + } + Ok(Arc::new(IncrementalQueryExec::new( + self.table.clone(), + self.audit_log, + self.mode, + self.start_exclusive, + self.end_inclusive, + projection.map(|indices| self.read_projection(indices)), + projected_schema, + ))) + } +} + +/// The table fields to read, and per output column the column of core's +/// batches it comes from. +#[derive(Debug, Clone)] +struct ReadProjection { + read_type: Vec, + columns: Vec, +} + +/// Execution-time plan: builds the incremental plan in paimon core and streams +/// its Arrow output when polled, so planning and `EXPLAIN` stay cheap. +#[derive(Debug, Clone)] +struct IncrementalQueryExec { + table: Table, + audit_log: bool, + mode: IncrementalScanMode, + start_exclusive: i64, + end_inclusive: i64, + projection: Option, + output_schema: ArrowSchemaRef, + plan_properties: Arc, +} + +impl IncrementalQueryExec { + fn new( + table: Table, + audit_log: bool, + mode: IncrementalScanMode, + start_exclusive: i64, + end_inclusive: i64, + projection: Option, + output_schema: ArrowSchemaRef, + ) -> Self { + let plan_properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(output_schema.clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + table, + audit_log, + mode, + start_exclusive, + end_inclusive, + projection, + output_schema, + plan_properties, + } + } + + /// Plans the range in paimon core and opens its batch stream. Reader memory + /// comes from the executing task's pool, as in the ordinary scan. The + /// projection is pushed into the read: Delta and Changelog decode the + /// projected fields, while Diff decodes full rows to compare them and + /// projects its output. The stream is `'static`, so it outlives the + /// transient read builder it is built from. + async fn open(&self, resources: ResourceContext) -> DFResult { + let mut read_builder = self.table.new_read_builder(); + read_builder.with_resources(resources); + if let Some(projection) = &self.projection { + read_builder.with_read_type(projection.read_type.clone()); + } + let scan = + read_builder.new_incremental_scan(self.mode, self.start_exclusive, self.end_inclusive); + let plan = if self.plans_range_as_one_batch() { + let combined = scan + .plan_combined_delta() + .await + .map_err(to_datafusion_error)?; + IncrementalPlan::try_new( + IncrementalScanMode::Delta, + combined + .into_splits() + .into_iter() + .map(IncrementalSplit::Data) + .collect(), + ) + } else { + scan.plan().await + } + .map_err(to_datafusion_error)?; + let read = read_builder.new_read().map_err(to_datafusion_error)?; + if self.audit_log { + read.to_audit_log_arrow(&plan) + } else { + read.to_incremental_arrow(&plan) + } + .map_err(to_datafusion_error) + } + + /// Java's `incremental-between` scan plans a data-evolution Delta range as + /// one batch, so a column update merges into the rows it updates when they + /// were written in the same range; planned per commit, it would come back + /// as rows of its own. `Auto` is Delta here, since such a table has no + /// primary key and so no changelog producer. An empty range is left to + /// `plan`; within the available snapshots it returns no splits, as Java's + /// empty-range scanner does. + fn plans_range_as_one_batch(&self) -> bool { + CoreOptions::new(self.table.schema().options()).data_evolution_enabled() + && matches!( + self.mode, + IncrementalScanMode::Auto | IncrementalScanMode::Delta + ) + && self.start_exclusive < self.end_inclusive + } +} +impl DisplayAs for IncrementalQueryExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "IncrementalQueryExec: mode={:?}, range=({}, {}], audit_log={}", + self.mode, self.start_exclusive, self.end_inclusive, self.audit_log + ) + } +} + +impl ExecutionPlan for IncrementalQueryExec { + fn name(&self) -> &str { + "IncrementalQueryExec" + } + + fn properties(&self) -> &Arc { + &self.plan_properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DFResult> { + if !children.is_empty() { + return internal_err!("IncrementalQueryExec is a leaf and takes no children"); + } + Ok(self) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DFResult { + if partition != 0 { + return internal_err!( + "IncrementalQueryExec has a single partition, got partition {partition}" + ); + } + let exec = self.clone(); + let stream = stream::once(async move { + let resources = reader_resources(&context, partition).map_err(to_datafusion_error)?; + let batches = await_with_runtime(exec.open(resources)).await?; + // Batches are forwarded as core emits them, projected, and rebuilt + // against the declared schema as the ordinary scan does. + let columns = exec.projection.map(|projection| projection.columns); + let schema = exec.output_schema; + Ok::<_, DataFusionError>(batches.map(move |batch| { + let batch = batch.map_err(to_datafusion_error)?; + let batch = match &columns { + Some(columns) => batch.project(columns)?, + None => batch, + }; + to_datafusion_batch(batch, &schema) + })) + }) + .try_flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.output_schema.clone(), + stream, + ))) + } + + fn partition_statistics(&self, _partition: Option) -> DFResult> { + Ok(Arc::new(Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Absent, + column_statistics: Statistics::unknown_column(&self.output_schema), + })) + } +} diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index cfab74c0e..ebdbf39c8 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -50,6 +50,7 @@ mod format_table_truncate; #[cfg(feature = "fulltext")] mod full_text_search; mod hybrid_search; +mod incremental_query; mod lateral_vector_search; mod memory; mod merge_condition; diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs b/crates/integrations/datafusion/src/physical_plan/scan.rs index e7edacbaa..5e7bcf38d 100644 --- a/crates/integrations/datafusion/src/physical_plan/scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/scan.rs @@ -59,7 +59,10 @@ use paimon::DataSplit; use crate::error::to_datafusion_error; use crate::filter_pushdown::scalar_to_datum; -fn to_datafusion_batch(batch: RecordBatch, schema: &ArrowSchemaRef) -> DFResult { +pub(crate) fn to_datafusion_batch( + batch: RecordBatch, + schema: &ArrowSchemaRef, +) -> DFResult { if batch.num_columns() != schema.fields().len() { return Err(datafusion::error::DataFusionError::Execution(format!( "Paimon reader returned {} columns for DataFusion schema with {} fields", diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 450a0c19c..f7dc7588a 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -170,6 +170,7 @@ impl SQLContextBuilder { let ctx = SessionContext::new_with_state(state); crate::blob_descriptor_functions::register_blob_descriptor_functions(&ctx); crate::variant_functions::register_variant_functions(&ctx); + crate::incremental_query::register_incremental_query(&ctx); SQLContext { ctx, catalogs: HashMap::new(), diff --git a/crates/integrations/datafusion/src/system_tables/mod.rs b/crates/integrations/datafusion/src/system_tables/mod.rs index 6b7d36d2c..dbd474b6f 100644 --- a/crates/integrations/datafusion/src/system_tables/mod.rs +++ b/crates/integrations/datafusion/src/system_tables/mod.rs @@ -145,6 +145,19 @@ pub(crate) fn provider_for_table( .map(Some) } +/// Rejects session options a `$audit_log` read does not take. +pub(crate) fn ensure_audit_log_dynamic_options( + dynamic_options: &HashMap, +) -> DFResult<()> { + if 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(), + )); + } + Ok(()) +} + /// Builds a system-table provider from a loaded Paimon table. pub(crate) async fn load( catalog: Arc, @@ -157,13 +170,8 @@ 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(), - )); + if system_name.eq_ignore_ascii_case("audit_log") { + ensure_audit_log_dynamic_options(&dynamic_options)?; } paimon::spec::CoreOptions::new(&dynamic_options) .ensure_read_authorized() diff --git a/crates/integrations/datafusion/src/table_function_args.rs b/crates/integrations/datafusion/src/table_function_args.rs index c0a6c833e..4f7f1d5a4 100644 --- a/crates/integrations/datafusion/src/table_function_args.rs +++ b/crates/integrations/datafusion/src/table_function_args.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use datafusion::common::config::CatalogOptions; use datafusion::common::ScalarValue; use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::logical_expr::Expr; @@ -70,11 +71,35 @@ pub(crate) fn parse_table_identifier( name: &str, default_database: &str, ) -> DFResult { - let parts: Vec<&str> = name.split('.').collect(); - match parts.len() { - 1 => Ok(Identifier::new(default_database, parts[0])), - 2 => Ok(Identifier::new(parts[0], parts[1])), - 3 => Ok(Identifier::new(parts[1], parts[2])), + let (_, database, table) = split_table_name(function_name, name)?; + Ok(Identifier::new(database.unwrap_or(default_database), table)) +} + +/// Resolves a table-name argument the way a `FROM` clause resolves a table +/// name: `table` and `database.table` take the session's current catalog (and +/// database), and `catalog.database.table` names the catalog. Returns the +/// catalog's name and the table's identifier in that catalog. +pub(crate) fn resolve_table_argument( + function_name: &str, + name: &str, + session: &CatalogOptions, +) -> DFResult<(String, Identifier)> { + let (catalog, database, table) = split_table_name(function_name, name)?; + Ok(( + catalog.unwrap_or(&session.default_catalog).to_string(), + Identifier::new(database.unwrap_or(&session.default_schema), table), + )) +} + +/// Splits `table`, `database.table` or `catalog.database.table`. +fn split_table_name<'a>( + function_name: &str, + name: &'a str, +) -> DFResult<(Option<&'a str>, Option<&'a str>, &'a str)> { + match name.split('.').collect::>().as_slice() { + [table] => Ok((None, None, *table)), + [database, table] => Ok((None, Some(*database), *table)), + [catalog, database, table] => Ok((Some(*catalog), Some(*database), *table)), _ => Err(DataFusionError::Plan(format!( "{function_name}: invalid table name '{name}', expected 'table', 'database.table', or 'catalog.database.table'" ))), diff --git a/crates/integrations/datafusion/tests/incremental_query.rs b/crates/integrations/datafusion/tests/incremental_query.rs new file mode 100644 index 000000000..5c2dfd3e9 --- /dev/null +++ b/crates/integrations/datafusion/tests/incremental_query.rs @@ -0,0 +1,1020 @@ +// 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. + +//! Integration tests for the `paimon_incremental_query` table function. + +use std::sync::Arc; + +use datafusion::arrow::array::RecordBatch; +use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use futures::StreamExt; +use paimon::catalog::{Catalog, Identifier}; +use paimon::table::{IncrementalScanMode, IncrementalSplit}; +use paimon::{CatalogOptions, FileSystemCatalog, Options}; +use paimon_datafusion::SQLContext; +use tempfile::TempDir; + +async fn run(ctx: &SQLContext, sql: &str) -> Vec { + ctx.sql(sql).await.unwrap().collect().await.unwrap() +} + +fn total_rows(batches: &[RecordBatch]) -> usize { + batches.iter().map(RecordBatch::num_rows).sum() +} + +fn catalog(warehouse: &TempDir) -> Arc { + let mut options = Options::new(); + options.set( + CatalogOptions::WAREHOUSE, + format!("file://{}", warehouse.path().display()), + ); + Arc::new(FileSystemCatalog::new(options).unwrap()) +} + +async fn context(warehouse: &TempDir) -> SQLContext { + let mut ctx = SQLContext::new(); + ctx.register_catalog("paimon", catalog(warehouse)) + .await + .unwrap(); + ctx +} + +/// A session over `warehouse` whose memory pool admits `limit` bytes. +async fn context_with_memory_limit( + warehouse: &TempDir, + limit: usize, +) -> (SQLContext, Arc) { + let pool = Arc::new(GreedyMemoryPool::new(limit)); + let runtime_env = RuntimeEnvBuilder::new() + .with_memory_pool(pool.clone()) + .build_arc() + .unwrap(); + let mut ctx = SQLContext::builder().with_runtime_env(runtime_env).build(); + ctx.register_catalog("paimon", catalog(warehouse)) + .await + .unwrap(); + (ctx, pool) +} + +/// Append-only table with one row per snapshot: snapshot 1 = (1,'a'), +/// snapshot 2 = (2,'b'), snapshot 3 = (3,'c'). +async fn setup_three_snapshots() -> (TempDir, SQLContext) { + let temp_dir = TempDir::new().expect("temp dir"); + let ctx = context(&temp_dir).await; + + run(&ctx, "CREATE TABLE paimon.default.t (id INT, name STRING)").await; + run(&ctx, "INSERT INTO paimon.default.t VALUES (1, 'a')").await; + run(&ctx, "INSERT INTO paimon.default.t VALUES (2, 'b')").await; + run(&ctx, "INSERT INTO paimon.default.t VALUES (3, 'c')").await; + (temp_dir, ctx) +} +/// Collects the `id` column across batches into a sorted vector. +fn sorted_ids(batches: &[RecordBatch]) -> Vec { + use datafusion::arrow::array::{Array, Int32Array}; + let mut ids = Vec::new(); + for batch in batches { + let col = batch + .column_by_name("id") + .expect("id column") + .as_any() + .downcast_ref::() + .expect("id is Int32"); + for i in 0..col.len() { + ids.push(col.value(i)); + } + } + ids.sort_unstable(); + ids +} + +#[tokio::test] +async fn test_auto_range_is_start_exclusive_end_inclusive() { + let (_tmp, ctx) = setup_three_snapshots().await; + + // (1, 3]: snapshots 2 and 3 only — start is exclusive, end inclusive. + let batches = run( + &ctx, + "SELECT * FROM paimon_incremental_query('default.t', 1, 3)", + ) + .await; + assert_eq!(total_rows(&batches), 2); + assert_eq!(sorted_ids(&batches), vec![2, 3]); + + // (0, 3]: every snapshot. + let all = run( + &ctx, + "SELECT * FROM paimon_incremental_query('default.t', 0, 3)", + ) + .await; + assert_eq!(sorted_ids(&all), vec![1, 2, 3]); + + // (3, 3]: empty range. + let empty = run( + &ctx, + "SELECT * FROM paimon_incremental_query('default.t', 3, 3)", + ) + .await; + assert_eq!(total_rows(&empty), 0); +} + +#[tokio::test] +async fn test_explicit_delta_mode_matches_auto() { + let (_tmp, ctx) = setup_three_snapshots().await; + let batches = run( + &ctx, + "SELECT * FROM paimon_incremental_query('default.t', 0, 2, 'delta')", + ) + .await; + assert_eq!(sorted_ids(&batches), vec![1, 2]); +} + +#[tokio::test] +async fn test_projection_selects_single_column() { + let (_tmp, ctx) = setup_three_snapshots().await; + let batches = run( + &ctx, + "SELECT id FROM paimon_incremental_query('default.t', 0, 3) ORDER BY id", + ) + .await; + assert_eq!(batches[0].num_columns(), 1); + assert_eq!(sorted_ids(&batches), vec![1, 2, 3]); +} +#[tokio::test] +async fn test_audit_log_suffix_prepends_rowkind() { + use datafusion::arrow::array::{Array, StringArray}; + let (_tmp, ctx) = setup_three_snapshots().await; + let batches = run( + &ctx, + "SELECT * FROM paimon_incremental_query('default.t$audit_log', 0, 3, 'delta')", + ) + .await; + assert_eq!(total_rows(&batches), 3); + let batch = &batches[0]; + assert_eq!( + batch.schema().field(0).name(), + "rowkind", + "audit-log output must lead with the rowkind column" + ); + let rowkind = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("rowkind is Utf8"); + // Append-only rows are all inserts. + for i in 0..rowkind.len() { + assert_eq!(rowkind.value(i), "+I"); + } +} + +async fn expect_error(ctx: &SQLContext, sql: &str) -> String { + match ctx.sql(sql).await { + Err(e) => e.to_string(), + Ok(df) => df + .collect() + .await + .expect_err("expected the query to fail") + .to_string(), + } +} + +#[tokio::test] +async fn test_rejects_unknown_mode() { + let (_tmp, ctx) = setup_three_snapshots().await; + let err = expect_error( + &ctx, + "SELECT * FROM paimon_incremental_query('default.t', 0, 1, 'bogus')", + ) + .await; + assert!( + err.contains("unknown scan mode"), + "error must name the bad mode: {err}" + ); +} + +#[tokio::test] +async fn test_rejects_bad_arity() { + let (_tmp, ctx) = setup_three_snapshots().await; + let err = expect_error( + &ctx, + "SELECT * FROM paimon_incremental_query('default.t', 0)", + ) + .await; + assert!( + err.contains("requires 3 or 4 arguments"), + "error must explain the arity: {err}" + ); +} + +#[tokio::test] +async fn test_rejects_end_before_start() { + let (_tmp, ctx) = setup_three_snapshots().await; + let err = expect_error( + &ctx, + "SELECT * FROM paimon_incremental_query('default.t', 5, 2)", + ) + .await; + assert!( + err.contains("must be >="), + "error must explain the range order: {err}" + ); +} + +#[tokio::test] +async fn test_reads_reserve_memory_from_the_session_pool() { + let (tmp, _ctx) = setup_three_snapshots().await; + let (ctx, _pool) = context_with_memory_limit(&tmp, 0).await; + // A pool that admits nothing refuses the ordinary scan, and must refuse + // the incremental read the same way. + for sql in [ + "SELECT * FROM paimon.default.t", + "SELECT * FROM paimon_incremental_query('default.t', 0, 3)", + ] { + let err = expect_error(&ctx, sql).await; + assert!(err.contains("Resources exhausted"), "{sql}: {err}"); + } +} + +#[tokio::test] +async fn test_projection_skips_unselected_columns() { + let tmp = TempDir::new().unwrap(); + let ctx = context(&tmp).await; + run( + &ctx, + "CREATE TABLE paimon.default.wide (id INT, payload STRING)", + ) + .await; + run( + &ctx, + "INSERT INTO paimon.default.wide VALUES (1, repeat('x', 1048576)), (2, repeat('y', 1048576))", + ) + .await; + + // The pool admits the `id` column but not the 1 MiB payloads, so only a + // read that never decodes `payload` can succeed. + let (ctx, pool) = context_with_memory_limit(&tmp, 64 * 1024).await; + let ids = run( + &ctx, + "SELECT id FROM paimon_incremental_query('default.wide', 0, 1)", + ) + .await; + assert_eq!(sorted_ids(&ids), vec![1, 2]); + // The reader releases its reservation; emitted batches belong to the consumer. + assert_eq!(pool.reserved(), 0); + let err = expect_error( + &ctx, + "SELECT * FROM paimon_incremental_query('default.wide', 0, 1)", + ) + .await; + assert!(err.contains("Resources exhausted"), "{err}"); +} + +#[tokio::test] +async fn test_batches_stream_before_the_whole_range_is_read() { + let (tmp, ctx) = setup_three_snapshots().await; + // Lose the range's last data file. Reading the whole range before emitting + // anything would fail before the first batch. + let table = catalog(&tmp) + .get_table(&Identifier::new("default", "t")) + .await + .unwrap(); + let plan = table + .new_read_builder() + .new_incremental_scan(IncrementalScanMode::Delta, 0, 3) + .plan() + .await + .unwrap(); + let Some(IncrementalSplit::Data(split)) = plan.splits().last() else { + panic!("expected a data split"); + }; + let file = split.data_files().last().expect("a data file"); + table + .file_io() + .delete_file(&format!("{}/{}", split.bucket_path(), file.file_name)) + .await + .unwrap(); + + let mut stream = ctx + .sql("SELECT * FROM paimon_incremental_query('default.t', 0, 3)") + .await + .unwrap() + .execute_stream() + .await + .unwrap(); + assert!(stream.next().await.unwrap().is_ok()); + let mut rest = Vec::new(); + while let Some(batch) = stream.next().await { + rest.push(batch); + } + assert!( + rest.iter().any(Result::is_err), + "the missing file still fails the read" + ); +} + +#[tokio::test] +async fn test_audit_log_projection_keeps_column_offsets() { + use datafusion::arrow::array::{Array, Int32Array, Int8Array, StringArray}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + // With sequence numbers, core's audit batches lead with two audit columns, + // so a projected data column sits two positions further right. + let tmp = TempDir::new().unwrap(); + let ctx = context(&tmp).await; + run( + &ctx, + "CREATE TABLE paimon.default.pk (id INT NOT NULL, name STRING, PRIMARY KEY (id)) WITH ( + 'bucket' = '1', + 'changelog-producer' = 'input', + 'table-read.sequence-number.enabled' = 'true' + )", + ) + .await; + run( + &ctx, + "INSERT INTO paimon.default.pk VALUES (1, 'a'), (2, 'b')", + ) + .await; + // Snapshot 2 updates id 1 and deletes id 2. + let table = catalog(&tmp) + .get_table(&Identifier::new("default", "pk")) + .await + .unwrap(); + let changes = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + Field::new("_VALUE_KIND", DataType::Int8, false), + ])), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2])), + Arc::new(StringArray::from(vec!["a", "a2", "b"])), + // -U, +U, -D + Arc::new(Int8Array::from(vec![1, 2, 3])), + ], + ) + .unwrap(); + let builder = table.new_write_builder(); + let mut write = builder.new_write().unwrap(); + write.write_arrow_batch(&changes).await.unwrap(); + let messages = write.prepare_commit().await.unwrap(); + builder.new_commit().commit(messages).await.unwrap(); + + for mode in ["changelog", "diff"] { + let batches = run( + &ctx, + &format!( + "SELECT name, rowkind, id FROM paimon_incremental_query('default.pk$audit_log', 1, 2, '{mode}') ORDER BY id, rowkind" + ), + ) + .await; + let mut rows = Vec::new(); + for batch in &batches { + let name = batch.column(0).as_any().downcast_ref::(); + let rowkind = batch.column(1).as_any().downcast_ref::(); + let id = batch.column(2).as_any().downcast_ref::(); + let (name, rowkind, id) = (name.unwrap(), rowkind.unwrap(), id.unwrap()); + for row in 0..batch.num_rows() { + rows.push(( + name.value(row).to_string(), + rowkind.value(row).to_string(), + id.value(row), + )); + } + } + let expected = [("a2", "+U", 1), ("a", "-U", 1), ("b", "-D", 2)] + .map(|(name, rowkind, id)| (name.to_string(), rowkind.to_string(), id)); + assert_eq!(rows, expected, "{mode}"); + + // Skipping `id` makes `name` the only read field, so its column comes + // right after the audit columns rather than at its table position. + let batches = run( + &ctx, + &format!( + "SELECT rowkind, name FROM paimon_incremental_query('default.pk$audit_log', 1, 2, '{mode}') ORDER BY name, rowkind" + ), + ) + .await; + let mut rows = Vec::new(); + for batch in &batches { + let rowkind = batch.column(0).as_any().downcast_ref::(); + let name = batch.column(1).as_any().downcast_ref::(); + let (rowkind, name) = (rowkind.unwrap(), name.unwrap()); + for row in 0..batch.num_rows() { + rows.push((rowkind.value(row).to_string(), name.value(row).to_string())); + } + } + let expected = [("-U", "a"), ("+U", "a2"), ("-D", "b")] + .map(|(rowkind, name)| (rowkind.to_string(), name.to_string())); + assert_eq!(rows, expected, "{mode}"); + } +} + +#[tokio::test] +async fn test_batches_match_the_declared_schema() { + // `DataFrame::cache` keeps the batches only if the plan declares their schema. + let (_tmp, ctx) = setup_three_snapshots().await; + for sql in [ + "SELECT * FROM paimon_incremental_query('default.t', 0, 3)", + "SELECT name, rowkind FROM paimon_incremental_query('default.t$audit_log', 0, 3)", + ] { + let cached = ctx.sql(sql).await.unwrap().cache().await.unwrap(); + assert_eq!(total_rows(&cached.collect().await.unwrap()), 3, "{sql}"); + } +} + +#[tokio::test] +async fn test_blob_rows_match_the_ordinary_read() { + use datafusion::arrow::util::pretty::pretty_format_batches; + // BLOB values live apart from the row's other columns: in a file of their + // own on a data-evolution table, behind a stored descriptor on a + // primary-key table. Reading the range must return the ordinary read. + let tmp = TempDir::new().unwrap(); + let ctx = context(&tmp).await; + run( + &ctx, + "CREATE TABLE paimon.default.de_assets (id INT, picture BLOB) WITH ( + 'row-tracking.enabled' = 'true', + 'data-evolution.enabled' = 'true' + )", + ) + .await; + run( + &ctx, + "CREATE TABLE paimon.default.pk_assets (id INT NOT NULL, picture BLOB, PRIMARY KEY (id)) + WITH ('bucket' = '1')", + ) + .await; + for table in ["de_assets", "pk_assets"] { + for values in ["(1, X'414141'), (2, X'4242')", "(3, X'43')"] { + run( + &ctx, + &format!("INSERT INTO paimon.default.{table} (id, picture) VALUES {values}"), + ) + .await; + } + let format = + |batches: Vec| pretty_format_batches(&batches).unwrap().to_string(); + let expected = format( + run( + &ctx, + &format!("SELECT id, picture FROM paimon.default.{table} ORDER BY id"), + ) + .await, + ); + for name in [table.to_string(), format!("{table}$audit_log")] { + let actual = format( + run( + &ctx, + &format!( + "SELECT id, picture FROM paimon_incremental_query('default.{name}', 0, 2) ORDER BY id" + ), + ) + .await, + ); + assert_eq!(actual, expected, "{name}"); + } + for (start, rows) in [(0, 3), (1, 1)] { + let batches = run( + &ctx, + &format!( + "SELECT count(*) AS n FROM paimon_incremental_query('default.{table}', {start}, 2)" + ), + ) + .await; + let count = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, rows, "{table} from {start}"); + } + } +} + +#[tokio::test] +async fn test_data_evolution_range_merges_column_updates() { + use datafusion::arrow::array::{Int64Array, StringArray}; + use datafusion::arrow::util::pretty::pretty_format_batches; + // MERGE INTO on a data-evolution table writes only the updated `name` to a + // file of its own. A range that also holds the rows it updates reads both + // files together, so the update merges into them; read per commit, the + // update would come back as rows whose NOT NULL `id` is missing. + let tmp = TempDir::new().unwrap(); + let ctx = context(&tmp).await; + run( + &ctx, + "CREATE TABLE paimon.default.de (id INT NOT NULL, name STRING, value INT) WITH ( + 'row-tracking.enabled' = 'true', + 'data-evolution.enabled' = 'true' + )", + ) + .await; + run( + &ctx, + "INSERT INTO paimon.default.de (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20)", + ) + .await; + run( + &ctx, + "CREATE TEMPORARY TABLE paimon.default.renames AS + SELECT * FROM (VALUES (1, 'ALICE')) AS t(id, name)", + ) + .await; + run( + &ctx, + "MERGE INTO paimon.default.de t USING paimon.default.renames s ON t.id = s.id + WHEN MATCHED THEN UPDATE SET name = s.name", + ) + .await; + let format = |batches: Vec| pretty_format_batches(&batches).unwrap().to_string(); + let expected = format( + run( + &ctx, + "SELECT id, name, value FROM paimon.default.de ORDER BY id", + ) + .await, + ); + // Auto is Delta on a data-evolution table, so both read the range as one batch. + for mode in ["", ", 'delta'"] { + for name in ["de", "de$audit_log"] { + let actual = format( + run( + &ctx, + &format!( + "SELECT id, name, value FROM paimon_incremental_query('default.{name}', 0, 2{mode}) ORDER BY id" + ), + ) + .await, + ); + assert_eq!(actual, expected, "{name}{mode}"); + } + let batches = run( + &ctx, + &format!( + "SELECT rowkind FROM paimon_incremental_query('default.de$audit_log', 0, 2{mode})" + ), + ) + .await; + let kinds: Vec<&str> = batches + .iter() + .flat_map(|batch| { + let kinds = batch.column(0).as_any().downcast_ref::(); + kinds.unwrap().iter().map(Option::unwrap) + }) + .collect(); + assert_eq!(kinds, ["+I", "+I"], "{mode}"); + } + // An empty range plans nothing, as on other tables. + for (start, end, rows) in [(0, 2, 2), (0, 0, 0)] { + let batches = run( + &ctx, + &format!( + "SELECT count(*) AS n FROM paimon_incremental_query('default.de', {start}, {end})" + ), + ) + .await; + let count = batches[0].column(0).as_any().downcast_ref::(); + assert_eq!(count.unwrap().value(0), rows, "({start}, {end}]"); + } +} + +#[tokio::test] +async fn test_data_evolution_range_of_column_updates_reads_unwritten_columns_as_null() { + use datafusion::arrow::array::{Array, Int32Array, StringArray}; + // Snapshot 2 updates `name` and snapshot 3 updates `value`, so (1, 3] + // holds two column files of the same rows and no file with `id`. + let tmp = TempDir::new().unwrap(); + let ctx = context(&tmp).await; + run( + &ctx, + "CREATE TABLE paimon.default.de (id INT NOT NULL, name STRING, value INT) WITH ( + 'row-tracking.enabled' = 'true', + 'data-evolution.enabled' = 'true' + )", + ) + .await; + run( + &ctx, + "INSERT INTO paimon.default.de (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20)", + ) + .await; + for (source, values, column) in [ + ("renames", "(1, 'ALICE')", "name"), + ("revalues", "(2, 99)", "value"), + ] { + run( + &ctx, + &format!( + "CREATE TEMPORARY TABLE paimon.default.{source} AS + SELECT * FROM (VALUES {values}) AS t(id, {column})" + ), + ) + .await; + run( + &ctx, + &format!( + "MERGE INTO paimon.default.de t USING paimon.default.{source} s ON t.id = s.id + WHEN MATCHED THEN UPDATE SET {column} = s.{column}" + ), + ) + .await; + } + for mode in ["", ", 'delta'"] { + for name in ["de", "de$audit_log"] { + let batches = run( + &ctx, + &format!( + "SELECT id, name, value FROM paimon_incremental_query('default.{name}', 1, 3{mode})" + ), + ) + .await; + let mut rows = Vec::new(); + for batch in &batches { + let id = batch.column(0).as_any().downcast_ref::(); + let name = batch.column(1).as_any().downcast_ref::(); + let value = batch.column(2).as_any().downcast_ref::(); + let (id, name, value) = (id.unwrap(), name.unwrap(), value.unwrap()); + for row in 0..batch.num_rows() { + rows.push(( + id.is_valid(row).then(|| id.value(row)), + name.value(row).to_string(), + value.value(row), + )); + } + } + rows.sort(); + let expected = [ + (None, "ALICE".to_string(), 10), + (None, "bob".to_string(), 99), + ]; + assert_eq!(rows, expected, "{name}{mode}"); + } + } +} + +/// Asserts that, for `name`, the function reads the rows a `FROM` clause in +/// the same session reads. +async fn assert_reads_as_from(ctx: &SQLContext, name: &str) { + let expected = sorted_ids(&run(ctx, &format!("SELECT id FROM {name}")).await); + let actual = sorted_ids( + &run( + ctx, + &format!("SELECT id FROM paimon_incremental_query('{name}', 0, 1)"), + ) + .await, + ); + assert_eq!(actual, expected, "{name}"); +} + +#[tokio::test] +async fn test_table_name_resolves_as_in_a_from_clause() { + use datafusion::catalog::MemoryCatalogProvider; + // Two warehouses with a `default.t` each, and a database in cat_a named + // like the catalog cat_b. + let (a, b) = (TempDir::new().unwrap(), TempDir::new().unwrap()); + let mut ctx = SQLContext::new(); + ctx.register_catalog("cat_a", catalog(&a)).await.unwrap(); + ctx.register_catalog("cat_b", catalog(&b)).await.unwrap(); + run(&ctx, "CREATE SCHEMA cat_a.cat_b").await; + for (table, id) in [ + ("cat_a.default.t", 1), + ("cat_b.default.t", 2), + ("cat_a.cat_b.t", 3), + ] { + run(&ctx, &format!("CREATE TABLE {table} (id INT)")).await; + run(&ctx, &format!("INSERT INTO {table} VALUES ({id})")).await; + } + + // cat_a, registered first, is the current catalog. + for name in [ + "cat_a.default.t", + "cat_b.default.t", + "default.t", + "t", + "cat_b.t", + ] { + assert_reads_as_from(&ctx, name).await; + } + run(&ctx, "USE cat_b.default").await; + for name in ["t", "default.t", "cat_a.default.t"] { + assert_reads_as_from(&ctx, name).await; + } + run(&ctx, "USE cat_a.cat_b").await; + assert_reads_as_from(&ctx, "t").await; + + let err = expect_error( + &ctx, + "SELECT * FROM paimon_incremental_query('nope.default.t', 0, 1)", + ) + .await; + assert!(err.contains("Unknown catalog 'nope'"), "{err}"); + ctx.ctx() + .register_catalog("mem", Arc::new(MemoryCatalogProvider::new())); + let err = expect_error( + &ctx, + "SELECT * FROM paimon_incremental_query('mem.public.t', 0, 1)", + ) + .await; + assert!( + err.contains("Catalog 'mem' is not a Paimon catalog"), + "{err}" + ); +} +#[tokio::test] +async fn test_session_read_options_apply_as_in_an_ordinary_read() { + use datafusion::arrow::util::pretty::pretty_format_batches; + let tmp = TempDir::new().unwrap(); + let ctx = context(&tmp).await; + run( + &ctx, + "CREATE TABLE paimon.default.de_assets (id INT, picture BLOB) WITH ( + 'row-tracking.enabled' = 'true', + 'data-evolution.enabled' = 'true' + )", + ) + .await; + run( + &ctx, + "CREATE TABLE paimon.default.pk_assets (id INT NOT NULL, picture BLOB, PRIMARY KEY (id)) + WITH ('bucket' = '1')", + ) + .await; + let format = |batches: Vec| pretty_format_batches(&batches).unwrap().to_string(); + for table in ["de_assets", "pk_assets"] { + run( + &ctx, + &format!("INSERT INTO paimon.default.{table} (id, picture) VALUES (1, X'414141')"), + ) + .await; + let ordinary = format!("SELECT id, picture FROM paimon.default.{table}"); + let payloads = format(run(&ctx, &ordinary).await); + run(&ctx, "SET 'paimon.blob-as-descriptor' = 'true'").await; + let descriptors = format(run(&ctx, &ordinary).await); + assert_ne!( + descriptors, payloads, + "{table}: the option changes the read" + ); + for (setting, expected) in [ + ("SET 'paimon.blob-as-descriptor' = 'true'", &descriptors), + ("RESET 'paimon.blob-as-descriptor'", &payloads), + ] { + run(&ctx, setting).await; + for name in [table.to_string(), format!("{table}$audit_log")] { + let sql = format!( + "SELECT id, picture FROM paimon_incremental_query('default.{name}', 0, 1)" + ); + assert_eq!( + &format(run(&ctx, &sql).await), + expected, + "{name}: {setting}" + ); + } + } + } + + // Descriptors the function returns can be read back through the session, + // as those of an ordinary read can. + let fresh = context(&tmp).await; + run(&fresh, "SET 'paimon.blob-as-descriptor' = 'true'").await; + run( + &fresh, + "SELECT * FROM paimon_incremental_query('default.pk_assets', 0, 1)", + ) + .await; + let location = catalog(&tmp) + .get_table(&Identifier::new("default", "pk_assets")) + .await + .unwrap() + .location() + .to_string(); + assert!(fresh + .blob_reader_registry() + .resolve(&format!("{location}/bucket-0/data.blob")) + .is_some()); + + run(&ctx, "CREATE TABLE paimon.default.t (id INT)").await; + run( + &ctx, + "INSERT INTO paimon.default.t VALUES (1), (2), (3), (4), (5)", + ) + .await; + run(&ctx, "SET 'paimon.read.batch-size' = '2'").await; + let sizes = |batches: Vec| { + batches + .iter() + .map(RecordBatch::num_rows) + .filter(|rows| *rows > 0) + .collect::>() + }; + let ordinary = sizes(run(&ctx, "SELECT id FROM paimon.default.t").await); + assert_eq!(ordinary, [2, 2, 1]); + let incremental = sizes( + run( + &ctx, + "SELECT id FROM paimon_incremental_query('default.t', 0, 1)", + ) + .await, + ); + assert_eq!(incremental, ordinary); +} + +#[tokio::test] +async fn test_options_that_choose_another_range_are_refused() { + let (_tmp, ctx) = setup_three_snapshots().await; + for (key, value) in [ + ("scan.snapshot-id", "1"), + // A tag that does not exist: the option is refused before any lookup. + ("scan.tag-name", "not_a_tag"), + ("scan.version", "1"), + ("scan.timestamp-millis", "0"), + ("scan.timestamp", "2024-01-01 00:00:00"), + ("scan.watermark", "0"), + ("incremental-between", "1,2"), + ("incremental-between-timestamp", "0,1"), + ("incremental-between-scan-mode", "delta"), + ("scan.file-creation-time-millis", "0"), + ("scan.creation-time-millis", "0"), + ("scan.mode", "latest"), + ] { + run(&ctx, &format!("SET 'paimon.{key}' = '{value}'")).await; + for name in ["t", "t$audit_log"] { + let err = expect_error( + &ctx, + &format!("SELECT * FROM paimon_incremental_query('default.{name}', 1, 3)"), + ) + .await; + assert!( + err.contains(&format!( + "scan option '{key}' is set; RESET 'paimon.{key}' first" + )), + "{key} on {name}: {err}" + ); + } + run(&ctx, &format!("RESET 'paimon.{key}'")).await; + } + run(&ctx, "SET 'paimon.scan.mode' = 'default'").await; + let batches = run( + &ctx, + "SELECT * FROM paimon_incremental_query('default.t', 1, 3)", + ) + .await; + assert_eq!(sorted_ids(&batches), vec![2, 3]); + run(&ctx, "RESET 'paimon.scan.mode'").await; + + // The table's own options choose no range either. + run( + &ctx, + "CREATE TABLE paimon.default.latest (id INT) WITH ('scan.mode' = 'latest')", + ) + .await; + run(&ctx, "INSERT INTO paimon.default.latest VALUES (1)").await; + let err = expect_error( + &ctx, + "SELECT * FROM paimon_incremental_query('default.latest', 0, 1)", + ) + .await; + assert!( + err.contains("table 'default.latest', which sets scan option 'scan.mode'"), + "{err}" + ); + // A stored time-travel selector is refused too, with or without other + // session options. + run(&ctx, "CREATE TABLE paimon.default.tagged (id INT)").await; + run(&ctx, "INSERT INTO paimon.default.tagged VALUES (1)").await; + run( + &ctx, + "ALTER TABLE paimon.default.tagged SET TBLPROPERTIES('scan.tag-name' = 'not_a_tag')", + ) + .await; + for setting in [None, Some("read.batch-size")] { + if let Some(key) = setting { + run(&ctx, &format!("SET 'paimon.{key}' = '2'")).await; + } + let err = expect_error( + &ctx, + "SELECT * FROM paimon_incremental_query('default.tagged', 0, 1)", + ) + .await; + assert!( + err.contains("table 'default.tagged', which sets scan option 'scan.tag-name'"), + "{setting:?}: {err}" + ); + if let Some(key) = setting { + run(&ctx, &format!("RESET 'paimon.{key}'")).await; + } + } +} + +#[tokio::test] +async fn test_range_options_are_checked_on_the_branch_read() { + use paimon::table::{BranchManager, SnapshotManager, TagManager}; + // A session `branch` reads that branch, so its options decide: `on_branch` + // stores `scan.mode=latest` only on b1, `on_main` only on the main branch. + let tmp = TempDir::new().unwrap(); + let ctx = context(&tmp).await; + let warehouse = catalog(&tmp); + for (table, created, altered) in [ + ("on_branch", "latest", "default"), + ("on_main", "default", "latest"), + ] { + run( + &ctx, + &format!( + "CREATE TABLE paimon.default.{table} (id INT) WITH ('scan.mode' = '{created}')" + ), + ) + .await; + run( + &ctx, + &format!("INSERT INTO paimon.default.{table} VALUES (1)"), + ) + .await; + let loaded = warehouse + .get_table(&Identifier::new("default", table)) + .await + .unwrap(); + let (io, location) = (loaded.file_io().clone(), loaded.location().to_string()); + let snapshot = SnapshotManager::new(io.clone(), location.clone()) + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + TagManager::new(io.clone(), location.clone()) + .create("base", &snapshot) + .await + .unwrap(); + BranchManager::new(io, location) + .create_branch_from_tag("b1", "base") + .await + .unwrap(); + run( + &ctx, + &format!( + "ALTER TABLE paimon.default.{table} SET TBLPROPERTIES('scan.mode' = '{altered}')" + ), + ) + .await; + } + + let query = + |table: &str| format!("SELECT id FROM paimon_incremental_query('default.{table}', 0, 1)"); + assert_eq!(sorted_ids(&run(&ctx, &query("on_branch")).await), vec![1]); + let err = expect_error(&ctx, &query("on_main")).await; + assert!( + err.contains("table 'default.on_main', which sets scan option 'scan.mode'"), + "{err}" + ); + run(&ctx, "SET 'paimon.branch' = 'b1'").await; + let err = expect_error(&ctx, &query("on_branch")).await; + assert!( + err.contains( + "branch 'b1' of table 'default.on_branch', which sets scan option 'scan.mode'" + ), + "{err}" + ); + assert_eq!(sorted_ids(&run(&ctx, &query("on_main")).await), vec![1]); +} + +#[tokio::test] +async fn test_audit_log_refuses_the_session_options_its_ordinary_read_refuses() { + let tmp = TempDir::new().unwrap(); + let ctx = context(&tmp).await; + run( + &ctx, + "CREATE TABLE paimon.default.pk (id INT NOT NULL, PRIMARY KEY (id)) WITH ('bucket' = '1')", + ) + .await; + run(&ctx, "INSERT INTO paimon.default.pk VALUES (1)").await; + run( + &ctx, + "SET 'paimon.table-read.sequence-number.enabled' = 'true'", + ) + .await; + for sql in [ + "SELECT * FROM paimon.default.pk$audit_log", + "SELECT * FROM paimon_incremental_query('default.pk$audit_log', 0, 1)", + ] { + let err = expect_error(&ctx, sql).await; + assert!( + err.contains( + "table-read.sequence-number.enabled is not supported by dynamic options for $audit_log" + ), + "{sql}: {err}" + ); + } +} diff --git a/crates/paimon/src/spec/core_options.rs b/crates/paimon/src/spec/core_options.rs index c98509752..9c1b2ff50 100644 --- a/crates/paimon/src/spec/core_options.rs +++ b/crates/paimon/src/spec/core_options.rs @@ -1391,6 +1391,34 @@ impl<'a> CoreOptions<'a> { selectors } + /// Options that choose which snapshots or files a scan reads: the + /// time-travel selectors, an incremental range or timestamp window, a file + /// creation-time filter, and a `scan.mode` other than `default`. A read + /// whose range its caller sets, such as an incremental query, cannot also + /// honor them. + pub fn configured_scan_range_options(&self) -> Vec<&'static str> { + let mut options = self.configured_time_travel_selectors(); + options.extend( + [ + INCREMENTAL_BETWEEN_OPTION, + INCREMENTAL_BETWEEN_TIMESTAMP_OPTION, + INCREMENTAL_BETWEEN_SCAN_MODE_OPTION, + "scan.file-creation-time-millis", + "scan.creation-time-millis", + ] + .into_iter() + .filter(|key| self.options.contains_key(*key)), + ); + if self + .options + .get(SCAN_MODE_OPTION) + .is_some_and(|mode| !mode.eq_ignore_ascii_case("default")) + { + options.push(SCAN_MODE_OPTION); + } + options + } + /// Validates and normalizes the internal time-travel selector. /// /// This is the semantic owner for selector mutual exclusion and strict @@ -3786,6 +3814,38 @@ mod tests { .is_err()); } + #[test] + fn test_scan_range_options_list_each_snapshot_or_file_selector() { + for (key, value) in [ + (SCAN_SNAPSHOT_ID_OPTION, "1"), + (SCAN_TAG_NAME_OPTION, "t"), + (SCAN_VERSION_OPTION, "1"), + (SCAN_TIMESTAMP_MILLIS_OPTION, "1"), + (SCAN_TIMESTAMP_OPTION, "2024-01-01 00:00:00"), + (SCAN_WATERMARK_OPTION, "1"), + (INCREMENTAL_BETWEEN_OPTION, "1,2"), + (INCREMENTAL_BETWEEN_TIMESTAMP_OPTION, "1,2"), + (INCREMENTAL_BETWEEN_SCAN_MODE_OPTION, "delta"), + ("scan.file-creation-time-millis", "1"), + ("scan.creation-time-millis", "1"), + (SCAN_MODE_OPTION, "latest"), + ] { + let options = HashMap::from([(key.to_string(), value.to_string())]); + assert_eq!( + CoreOptions::new(&options).configured_scan_range_options(), + vec![key], + "{key}" + ); + } + let options = HashMap::from([ + (SCAN_MODE_OPTION.to_string(), "DEFAULT".to_string()), + ("read.batch-size".to_string(), "2".to_string()), + ]); + assert!(CoreOptions::new(&options) + .configured_scan_range_options() + .is_empty()); + } + #[test] fn test_watermark_conflicts_with_other_selectors() { let options = HashMap::from([ diff --git a/crates/paimon/src/table/data_evolution_nested.rs b/crates/paimon/src/table/data_evolution_nested.rs index c75203f64..9254b1c52 100644 --- a/crates/paimon/src/table/data_evolution_nested.rs +++ b/crates/paimon/src/table/data_evolution_nested.rs @@ -43,10 +43,13 @@ pub(super) struct NestedSelection { pub whole_source: Option, } -/// Return `None` only when no file supplies any leaf of this ROW. +/// Return `None` only when no file supplies any leaf of this ROW. In change +/// events (`change_events`), a child no file wrote reads as NULL even when it +/// is NOT NULL. pub(super) fn select_nested_row( read_field: &DataField, sources: &[PhysicalRowSource<'_>], + change_events: bool, ) -> Result> { let DataType::Row(read_row) = read_field.data_type() else { return Err(Error::DataInvalid { @@ -123,7 +126,7 @@ pub(super) fn select_nested_row( ), }); } - if providers.is_empty() && !child.data_type().is_nullable() { + if providers.is_empty() && !change_events && !child.data_type().is_nullable() { return Err(Error::DataInvalid { message: format!( "Cannot read non-nullable nested field '{}.{}' without a provider", @@ -516,7 +519,7 @@ mod tests { fields: &whole, }]; assert_eq!( - select_nested_row(&requested, &sources) + select_nested_row(&requested, &sources, false) .unwrap() .unwrap() .whole_source, @@ -535,7 +538,9 @@ mod tests { fields: &older, }, ]; - let selection = select_nested_row(&requested, &sources).unwrap().unwrap(); + let selection = select_nested_row(&requested, &sources, false) + .unwrap() + .unwrap(); assert_eq!(selection.children, vec![Some(1), Some(1)]); assert_eq!(selection.whole_source, None); let old_request = source_read_field(&requested, ¤t, 1, &older, &selection).unwrap(); @@ -559,7 +564,7 @@ mod tests { fields: &base, }, ]; - let selected = select_nested_row(&row(vec![name(), age()]), &sources) + let selected = select_nested_row(&row(vec![name(), age()]), &sources, false) .unwrap() .unwrap(); assert_eq!(selected.anchors, vec![0, 1]); @@ -581,7 +586,7 @@ mod tests { fields: &base, }, ]; - let selected = select_nested_row(&row(vec![name()]), &sources) + let selected = select_nested_row(&row(vec![name()]), &sources, false) .unwrap() .unwrap(); assert_eq!(selected.anchors, vec![0, 1]); @@ -605,7 +610,9 @@ mod tests { fields: &base_file, }, ]; - let selection = select_nested_row(&requested, &sources).unwrap().unwrap(); + let selection = select_nested_row(&requested, &sources, false) + .unwrap() + .unwrap(); let anchor = source_read_field(&requested, ¤t, 0, &age_file, &selection).unwrap(); let provider = source_read_field(&requested, ¤t, 1, &base_file, &selection).unwrap(); @@ -631,7 +638,9 @@ mod tests { source_index: 0, fields: &physical, }]; - let selection = select_nested_row(&requested, &sources).unwrap().unwrap(); + let selection = select_nested_row(&requested, &sources, false) + .unwrap() + .unwrap(); let read = source_read_field(&requested, ¤t, 0, &physical, &selection).unwrap(); let DataType::Row(read_row) = read.data_type() else { panic!("read field must be ROW") @@ -656,7 +665,7 @@ mod tests { fields: &older, }, ]; - let selected = select_nested_row(&row(vec![name(), age()]), &sources) + let selected = select_nested_row(&row(vec![name(), age()]), &sources, false) .unwrap() .unwrap(); assert_eq!(selected.anchors, vec![0]); @@ -683,7 +692,7 @@ mod tests { }, ]; assert!(matches!( - select_nested_row(&read, &sources), + select_nested_row(&read, &sources, false), Err(Error::Unsupported { message }) if message.contains("profile.address") )); } @@ -697,11 +706,27 @@ mod tests { fields: &available, }]; assert!(matches!( - select_nested_row(&row(vec![required]), &sources), + select_nested_row(&row(vec![required]), &sources, false), Err(Error::DataInvalid { message, .. }) if message.contains("profile.required") )); } + #[test] + fn change_events_read_a_missing_non_nullable_subfield_as_null() { + let required = field(4, "required", DataType::Int(IntType::with_nullable(false))); + let available = vec![row(vec![name()])]; + let sources = [PhysicalRowSource { + source_index: 0, + fields: &available, + }]; + let selected = select_nested_row(&row(vec![name(), required]), &sources, true) + .unwrap() + .unwrap(); + assert_eq!(selected.anchors, vec![0]); + assert_eq!(selected.children, vec![Some(0), None]); + assert_eq!(selected.whole_source, None); + } + fn struct_batch( names: Vec>, ages: Vec>, diff --git a/crates/paimon/src/table/data_evolution_reader.rs b/crates/paimon/src/table/data_evolution_reader.rs index 6b0bff0c7..20696535f 100644 --- a/crates/paimon/src/table/data_evolution_reader.rs +++ b/crates/paimon/src/table/data_evolution_reader.rs @@ -1142,6 +1142,7 @@ impl DataEvolutionReader { partition_fields: &partition_fields, row_id_pushdown: row_ranges.is_some(), nested_enabled, + change_events: split.is_streaming(), }, )?; @@ -2518,6 +2519,11 @@ struct SourcePlanReadContext<'a> { partition_fields: &'a [DataField], row_id_pushdown: bool, nested_enabled: bool, + /// The split holds change events rather than a materialized snapshot + /// ([`DataSplit::is_streaming`]). A field no file of the group wrote was + /// not written in the planned range, so it reads as NULL, NOT NULL fields + /// included. + change_events: bool, } #[cfg(test)] @@ -2537,6 +2543,7 @@ fn build_source_plan( partition_fields: &[], row_id_pushdown: false, nested_enabled: false, + change_events: false, }, ) } @@ -2553,6 +2560,7 @@ fn build_source_plan_with_row_id_pushdown( partition_fields, row_id_pushdown, nested_enabled, + change_events, } = context; let mut sources = Vec::new(); let mut normal_providers: HashMap = HashMap::new(); // field_id -> source_idx @@ -2697,9 +2705,11 @@ fn build_source_plan_with_row_id_pushdown( }, ) .collect::>(); - if let Some(selection) = - super::data_evolution_nested::select_nested_row(field, &physical_sources)? - { + if let Some(selection) = super::data_evolution_nested::select_nested_row( + field, + &physical_sources, + change_events, + )? { if selection.whole_source.is_none() { let current_field = table_fields .iter() @@ -2771,6 +2781,7 @@ fn build_source_plan_with_row_id_pushdown( column_plan.push(Some((source_idx, field_offset))); nested_plan.push(None); } else if !field.data_type().is_nullable() + && !change_events && !partition_fields .iter() .any(|partition| partition.id() == field.id()) @@ -4602,6 +4613,7 @@ mod tests { partition_fields: &[], row_id_pushdown: true, nested_enabled: false, + change_events: false, }, ) .unwrap_err(); @@ -4610,6 +4622,52 @@ mod tests { if message.contains("non-nullable field 'payload'"))); } + #[test] + fn change_events_read_a_field_no_file_wrote_as_null() { + // Two column updates of one file, neither of which wrote `id`. + let files = vec![ + data_file("a-update.parquet", 0, 4, 2, Some(vec!["a"])), + data_file("b-update.parquet", 0, 4, 3, Some(vec!["b"])), + ]; + let prepared_group = PreparedMergeGroup::new(&files).unwrap(); + // Newest first: the `b` update, then the `a` update. + let file_infos = vec![resolved_info(vec![3]), resolved_info(vec![2])]; + let read_type = vec![ + DataField::new( + 1, + "id".to_string(), + DataType::Int(IntType::with_nullable(false)), + ), + DataField::new(2, "a".to_string(), DataType::Int(IntType::new())), + DataField::new(3, "b".to_string(), DataType::Int(IntType::new())), + ]; + let plan = |change_events| { + build_source_plan_with_row_id_pushdown( + &prepared_group, + &file_infos, + &read_type, + &HashSet::new(), + SourcePlanReadContext { + table_fields: &read_type, + partition_fields: &[], + row_id_pushdown: false, + nested_enabled: false, + change_events, + }, + ) + }; + + assert!( + matches!(plan(false), Err(Error::DataInvalid { message, .. }) + if message.contains("non-nullable field 'id'")) + ); + let source_plan = plan(true).unwrap(); + assert_eq!( + source_plan.column_plan, + vec![None, Some((1, 0)), Some((0, 0))] + ); + } + #[test] fn test_build_source_plan_validates_complete_range_only_for_read_blob_source() { let files = vec![ @@ -4694,6 +4752,7 @@ mod tests { partition_fields: &[], row_id_pushdown: false, nested_enabled: false, + change_events: false, }, ) .unwrap(); diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index 47d38e5e1..ac761a18e 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -776,10 +776,21 @@ impl<'a> PaimonTableRead<'a> { } } } - // Delta / Changelog rows are read as-is from planned files (no full-table - // merge against historical base versions). + // Delta / Changelog rows are read as-is from planned files, without a merge + // against historical base versions. Files are still decoded as an ordinary + // read decodes them: primary-key BLOB descriptors resolve to payloads, and + // data-evolution files in a split that hold column slices of the same rows + // merge by row id. let data_splits = self.matching_splits(&data_splits)?; - self.new_data_file_reader()?.read(&data_splits) + let core_options = self.table.schema().core_options(); + if !self.table.schema().primary_keys().is_empty() { + return self + .resolving_primary_key_blobs(&core_options, |read| read.read_raw(&data_splits)); + } + if core_options.data_evolution_enabled() { + return self.read_with_evolution(&data_splits, &core_options); + } + self.read_raw(&data_splits) } fn to_incremental_diff_arrow( @@ -828,6 +839,14 @@ impl<'a> PaimonTableRead<'a> { ) -> crate::Result { match plan.mode() { IncrementalScanMode::Diff => self.audit_diff_stream(plan), + // A data-evolution table has no primary key, so its rows are inserts, + // read as `to_incremental_arrow` reads them. + IncrementalScanMode::Delta + if self.table.schema().core_options().data_evolution_enabled() => + { + let schema = audit_schema_for_read_type(&self.read_type, false)?; + prepend_insert_row_kind_stream(self.to_incremental_arrow(plan)?, schema) + } IncrementalScanMode::Delta => { self.audit_raw_stream(plan, !self.table.schema().primary_keys().is_empty()) } @@ -1900,6 +1919,18 @@ impl<'a> PaimonTableRead<'a> { &self, data_splits: &[DataSplit], core_options: &CoreOptions<'_>, + ) -> crate::Result { + self.resolving_primary_key_blobs(core_options, |read| { + read.read_pk(data_splits, core_options) + }) + } + + /// Runs `read` over the stored BLOB descriptors and resolves primary-key + /// managed BLOB payloads in its output. + fn resolving_primary_key_blobs( + &self, + core_options: &CoreOptions<'_>, + read: impl FnOnce(&Self) -> crate::Result, ) -> crate::Result { use super::managed_blob_reader::{resolve_primary_key_blob_stream, ManagedBlobReadPlan}; @@ -1912,7 +1943,7 @@ impl<'a> PaimonTableRead<'a> { let mut inner = self.clone(); inner.read_type = plan.scan_fields().to_vec(); inner.data_predicates.clear(); - let stream = inner.read_pk(data_splits, core_options)?; + let stream = read(&inner)?; return Ok(plan.finish( stream, core_options, @@ -1922,7 +1953,7 @@ impl<'a> PaimonTableRead<'a> { )); } - let stream = self.read_pk(data_splits, core_options)?; + let stream = read(self)?; Ok(resolve_primary_key_blob_stream( stream, self.read_type(), diff --git a/crates/paimon/tests/incremental_batch_scan_test.rs b/crates/paimon/tests/incremental_batch_scan_test.rs index 193afb123..d4a4b87e4 100644 --- a/crates/paimon/tests/incremental_batch_scan_test.rs +++ b/crates/paimon/tests/incremental_batch_scan_test.rs @@ -1811,3 +1811,453 @@ async fn delta_keeps_physical_retractions_and_changelog_row_kinds() { } assert!(read_current_pairs(&table).await.is_empty()); } + +type BlobRow = (Option, Option>); + +/// `(id, picture)` of every row, sorted, located by column name so a leading +/// `rowkind` column is skipped. +fn blob_rows(batches: &[RecordBatch]) -> Vec { + use arrow_array::LargeBinaryArray; + let mut rows = Vec::new(); + for batch in batches { + let ids = batch.column_by_name("id").unwrap(); + let ids = ids.as_any().downcast_ref::().unwrap(); + let pictures = batch.column_by_name("picture").unwrap(); + let pictures = pictures + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + rows.push(( + ids.is_valid(row).then(|| ids.value(row)), + pictures.is_valid(row).then(|| pictures.value(row).to_vec()), + )); + } + } + rows.sort(); + rows +} + +/// Reads `(start, end]` per commit, as rows and as audit rows. +async fn read_delta_both_ways( + table: &paimon::table::Table, + start: i64, + end: i64, +) -> (Vec, Vec) { + let builder = table.new_read_builder(); + let plan = builder + .new_incremental_scan(IncrementalScanMode::Delta, start, end) + .plan() + .await + .unwrap(); + let read = builder.new_read().unwrap(); + let rows = read + .to_incremental_arrow(&plan) + .unwrap() + .try_collect() + .await + .unwrap(); + let audit = read + .to_audit_log_arrow(&plan) + .unwrap() + .try_collect() + .await + .unwrap(); + (rows, audit) +} + +/// A data-evolution commit writes BLOB and dedicated VECTOR columns to files of +/// their own beside the row's other columns. Each commit must still read back +/// as whole rows, as in an ordinary read. +#[tokio::test] +async fn delta_merges_data_evolution_column_files_into_rows() { + use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; + use arrow_array::{ArrayRef, FixedSizeListArray, Float32Array, LargeBinaryArray, StringArray}; + use arrow_schema::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; + use paimon::spec::{BlobType, DataType, FloatType, IntType, Schema, TableSchema, VectorType}; + use std::sync::Arc; + + let path = "memory:/incremental_batch/data_evolution_files"; + let vector = VectorType::try_new(true, 2, DataType::Float(FloatType::new())).unwrap(); + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("picture", DataType::Blob(BlobType::new())) + .column("embedding", DataType::Vector(vector)) + .option("row-tracking.enabled", "true") + .option("data-evolution.enabled", "true") + .option("vector.file.format", "parquet") + .build() + .unwrap(); + let (io, table) = memory_table(path, TableSchema::new(0, &schema)); + setup_dirs(&io, path).await; + persist_table_schema(&io, path, table.schema()).await; + let element = Arc::new(Field::new("element", ArrowDataType::Float32, true)); + let batch = |ids: Vec, pictures: Vec<&[u8]>, embeddings: Vec<[f32; 2]>| { + let mut builder = + FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(element.clone()); + for embedding in embeddings { + builder.values().append_slice(&embedding); + builder.append(true); + } + RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + Field::new("id", ArrowDataType::Int32, true), + Field::new("picture", ArrowDataType::LargeBinary, true), + Field::new( + "embedding", + ArrowDataType::FixedSizeList(element.clone(), 2), + true, + ), + ])), + vec![ + Arc::new(Int32Array::from(ids)) as ArrayRef, + Arc::new(LargeBinaryArray::from(pictures)), + Arc::new(builder.finish()), + ], + ) + .unwrap() + }; + write_batch( + &table, + &batch( + vec![1, 2], + vec![b"aaa", b"bb"], + vec![[1.0, 0.0], [0.0, 1.0]], + ), + ) + .await; + write_batch(&table, &batch(vec![3], vec![b"c"], vec![[0.5, 0.5]])).await; + + let embeddings = |batches: &[RecordBatch]| { + let mut rows = Vec::new(); + for batch in batches { + let ids = batch.column_by_name("id").unwrap(); + let ids = ids.as_any().downcast_ref::().unwrap(); + let lists = batch.column_by_name("embedding").unwrap(); + let lists = lists.as_any().downcast_ref::().unwrap(); + for row in 0..batch.num_rows() { + let values = lists.is_valid(row).then(|| { + let values = lists.value(row); + let values = values.as_any().downcast_ref::().unwrap(); + values.values().to_vec() + }); + rows.push((ids.is_valid(row).then(|| ids.value(row)), values)); + } + } + rows.sort_by_key(|row| row.0); + rows + }; + let picture = |id: i32, bytes: &[u8]| (Some(id), Some(bytes.to_vec())); + for (start, ids) in [(0, vec![1, 2, 3]), (1, vec![3])] { + let (rows, audit) = read_delta_both_ways(&table, start, 2).await; + let all = [picture(1, b"aaa"), picture(2, b"bb"), picture(3, b"c")]; + let expected: Vec = all + .into_iter() + .filter(|row| ids.contains(&row.0.unwrap())) + .collect(); + assert_eq!(blob_rows(&rows), expected, "rows from {start}"); + assert_eq!(blob_rows(&audit), expected, "audit rows from {start}"); + let vectors = [ + (Some(1), Some(vec![1.0, 0.0])), + (Some(2), Some(vec![0.0, 1.0])), + (Some(3), Some(vec![0.5, 0.5])), + ]; + let expected: Vec<_> = vectors + .into_iter() + .filter(|row| ids.contains(&row.0.unwrap())) + .collect(); + assert_eq!(embeddings(&rows), expected, "embeddings from {start}"); + let kinds: Vec<_> = audit + .iter() + .flat_map(|batch| { + let kinds = batch.column(0).as_any().downcast_ref::(); + kinds.unwrap().iter().map(|kind| kind.unwrap().to_owned()) + }) + .collect(); + assert_eq!(kinds, vec!["+I"; ids.len()], "row kinds from {start}"); + } +} + +/// A primary-key table keeps managed BLOB descriptors in its data files. Each +/// commit must read back with the payloads, as in an ordinary read. +#[tokio::test] +async fn delta_resolves_primary_key_blob_payloads() { + use arrow_array::{ArrayRef, LargeBinaryArray}; + use arrow_schema::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; + use paimon::spec::{BlobType, DataType, IntType, Schema, TableSchema}; + use std::sync::Arc; + + let path = "memory:/incremental_batch/primary_key_blobs"; + let schema = Schema::builder() + .column("id", DataType::Int(IntType::with_nullable(false))) + .column("picture", DataType::Blob(BlobType::new())) + .primary_key(["id"]) + .option("bucket", "1") + .build() + .unwrap(); + let (io, table) = memory_table(path, TableSchema::new(0, &schema)); + setup_dirs(&io, path).await; + persist_table_schema(&io, path, table.schema()).await; + let batch = |ids: Vec, pictures: Vec<&[u8]>| { + RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + Field::new("id", ArrowDataType::Int32, false), + Field::new("picture", ArrowDataType::LargeBinary, true), + ])), + vec![ + Arc::new(Int32Array::from(ids)) as ArrayRef, + Arc::new(LargeBinaryArray::from(pictures)), + ], + ) + .unwrap() + }; + write_batch(&table, &batch(vec![1, 2], vec![b"aaa", b"bb"])).await; + write_batch(&table, &batch(vec![2], vec![b"c"])).await; + + let (rows, audit) = read_delta_both_ways(&table, 0, 2).await; + let expected = vec![ + (Some(1), Some(b"aaa".to_vec())), + (Some(2), Some(b"bb".to_vec())), + (Some(2), Some(b"c".to_vec())), + ]; + assert_eq!(blob_rows(&rows), expected); + assert_eq!(blob_rows(&audit), expected); +} + +/// Each commit reads back as a streaming read of that snapshot does, including +/// a data-evolution update that rewrites a column of existing rows. +#[tokio::test] +async fn delta_reads_each_commit_as_its_stream_follow_up() { + use arrow_array::{ArrayRef, Int64Array, LargeBinaryArray}; + use paimon::spec::{BlobType, DataType, IntType, Schema, TableSchema}; + use std::sync::Arc; + + let path = "memory:/incremental_batch/stream_follow_up"; + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("picture", DataType::Blob(BlobType::new())) + .option("row-tracking.enabled", "true") + .option("data-evolution.enabled", "true") + .build() + .unwrap(); + let (io, table) = memory_table(path, TableSchema::new(0, &schema)); + setup_dirs(&io, path).await; + persist_table_schema(&io, path, table.schema()).await; + let batch = |ids: Vec, pictures: Vec<&[u8]>| { + RecordBatch::try_from_iter([ + ("id", Arc::new(Int32Array::from(ids)) as ArrayRef), + ("picture", Arc::new(LargeBinaryArray::from(pictures))), + ]) + .unwrap() + }; + write_batch(&table, &batch(vec![1, 2], vec![b"aaa", b"bb"])).await; + write_batch(&table, &batch(vec![3], vec![b"c"])).await; + // Snapshot 3 rewrites `picture` of row 0. + let builder = table.new_write_builder(); + let mut update = builder + .new_data_evolution_writer(vec!["picture".into()]) + .unwrap(); + update + .add_matched_batch( + RecordBatch::try_from_iter([ + ("_ROW_ID", Arc::new(Int64Array::from(vec![0])) as ArrayRef), + ( + "picture", + Arc::new(LargeBinaryArray::from(vec![b"new".as_slice()])), + ), + ]) + .unwrap(), + ) + .unwrap(); + builder + .new_commit() + .commit(update.prepare_commit().await.unwrap()) + .await + .unwrap(); + + let read_builder = table.new_read_builder(); + let mut stream = read_builder.new_stream_scan().unwrap(); + for snapshot in 1..=3 { + stream.restore(Some(snapshot)).unwrap(); + let plan = stream.plan().await.unwrap().unwrap(); + let streamed: Vec = read_builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + let (rows, audit) = read_delta_both_ways(&table, snapshot - 1, snapshot).await; + assert!(!streamed.is_empty(), "snapshot {snapshot}"); + assert_eq!( + blob_rows(&rows), + blob_rows(&streamed), + "snapshot {snapshot}" + ); + assert_eq!( + blob_rows(&audit), + blob_rows(&streamed), + "snapshot {snapshot}" + ); + } +} + +/// A range of two column updates of the same rows, without the rows they +/// update: every reader of the combined plan returns the updated columns and +/// reads the column no file of the range wrote as NULL, NOT NULL `id` +/// included. +#[tokio::test] +async fn combined_delta_reads_the_columns_a_range_did_not_write_as_null() { + use arrow_array::{ArrayRef, Int64Array, StringArray}; + use paimon::spec::{DataType, IntType, Schema, TableSchema, VarCharType}; + use paimon::table::{IncrementalPlan, IncrementalSplit}; + use std::collections::HashMap; + use std::sync::Arc; + + let path = "memory:/incremental_batch/column_updates_only"; + let schema = Schema::builder() + .column("id", DataType::Int(IntType::with_nullable(false))) + .column("name", DataType::VarChar(VarCharType::string_type())) + .column("value", DataType::Int(IntType::new())) + .option("row-tracking.enabled", "true") + .option("data-evolution.enabled", "true") + .build() + .unwrap(); + let (io, table) = memory_table(path, TableSchema::new(0, &schema)); + setup_dirs(&io, path).await; + persist_table_schema(&io, path, table.schema()).await; + let rows = RecordBatch::try_from_iter_with_nullable([ + ( + "id", + Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef, + false, + ), + ( + "name", + Arc::new(StringArray::from(vec!["alice", "bob"])), + true, + ), + ("value", Arc::new(Int32Array::from(vec![10, 20])), true), + ]) + .unwrap(); + write_batch(&table, &rows).await; + // Snapshot 2 rewrites `name` of row 0 and snapshot 3 `value` of row 1. + for (column, row_id, values) in [ + ( + "name", + 0, + Arc::new(StringArray::from(vec!["ALICE"])) as ArrayRef, + ), + ("value", 1, Arc::new(Int32Array::from(vec![99]))), + ] { + // Distinct commit times, for the timestamp window below. + std::thread::sleep(std::time::Duration::from_millis(5)); + let builder = table.new_write_builder(); + let mut update = builder + .new_data_evolution_writer(vec![column.into()]) + .unwrap(); + update + .add_matched_batch( + RecordBatch::try_from_iter([ + ( + "_ROW_ID", + Arc::new(Int64Array::from(vec![row_id])) as ArrayRef, + ), + (column, values), + ]) + .unwrap(), + ) + .unwrap(); + builder + .new_commit() + .commit(update.prepare_commit().await.unwrap()) + .await + .unwrap(); + } + type Row = (Option, Option, Option); + let collect = |batches: &[RecordBatch]| { + let mut rows: Vec = Vec::new(); + for batch in batches { + let column = |name: &str| batch.column_by_name(name).unwrap().clone(); + let (id, name, value) = (column("id"), column("name"), column("value")); + let id = id.as_any().downcast_ref::().unwrap(); + let name = name.as_any().downcast_ref::().unwrap(); + let value = value.as_any().downcast_ref::().unwrap(); + for row in 0..batch.num_rows() { + rows.push(( + id.is_valid(row).then(|| id.value(row)), + name.is_valid(row).then(|| name.value(row).to_string()), + value.is_valid(row).then(|| value.value(row)), + )); + } + } + rows.sort(); + rows + }; + let updates: Vec = vec![ + (None, Some("ALICE".into()), Some(10)), + (None, Some("bob".into()), Some(99)), + ]; + + let builder = table.new_read_builder(); + let splits = builder + .new_incremental_scan(IncrementalScanMode::Delta, 1, 3) + .plan_combined_delta() + .await + .unwrap() + .into_splits(); + let read = builder.new_read().unwrap(); + let batches: Vec = read.to_arrow(&splits).unwrap().try_collect().await.unwrap(); + assert_eq!(collect(&batches), updates, "to_arrow"); + let plan = IncrementalPlan::try_new( + IncrementalScanMode::Delta, + splits.into_iter().map(IncrementalSplit::Data).collect(), + ) + .unwrap(); + for (reader, stream) in [ + ("to_incremental_arrow", read.to_incremental_arrow(&plan)), + ("to_audit_log_arrow", read.to_audit_log_arrow(&plan)), + ] { + let batches: Vec = stream.unwrap().try_collect().await.unwrap(); + assert_eq!(collect(&batches), updates, "{reader}"); + } + + // An ordinary scan given the same range as a timestamp window plans it + // the same way. + let snapshots = table.snapshot_manager(); + let start = snapshots.get_snapshot(1).await.unwrap().time_millis(); + let end = snapshots.get_snapshot(3).await.unwrap().time_millis(); + let windowed = table.copy_with_options(HashMap::from([( + "incremental-between-timestamp".to_string(), + format!("{start},{end}"), + )])); + let builder = windowed.new_read_builder(); + 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!(collect(&batches), updates, "incremental-between-timestamp"); + + // The table itself still reads whole rows. + let builder = table.new_read_builder(); + let plan = builder.new_scan().plan().await.unwrap(); + let batches: Vec = builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + let current: Vec = vec![ + (Some(1), Some("ALICE".into()), Some(10)), + (Some(2), Some("bob".into()), Some(99)), + ]; + assert_eq!(collect(&batches), current, "snapshot read"); +} diff --git a/crates/paimon/tests/partition_filter_test.rs b/crates/paimon/tests/partition_filter_test.rs index 0ee258e05..c55a61bf4 100644 --- a/crates/paimon/tests/partition_filter_test.rs +++ b/crates/paimon/tests/partition_filter_test.rs @@ -380,3 +380,38 @@ async fn ordinary_partition_filter_survives_diff_workers_and_audit_images() { .all(|id| *id == 1)); } } + +#[tokio::test] +async fn ordinary_partition_filter_survives_incremental_delta_reads() { + // Delta reads take a separate path per table kind, plain and audit-log + // alike; the read's partition conjuncts must restrict every one of them. + for mode in ["append", "pk", "de"] { + let table = partitioned(mode).await; + let plan = table + .new_read_builder() + .new_incremental_scan(IncrementalScanMode::Delta, 0, 1) + .plan() + .await + .unwrap(); + let mut builder = table.new_read_builder(); + builder + .with_projection(&["id"]) + .unwrap() + .with_filter(equal(&table, "p", "a")); + let read = builder.new_read().unwrap(); + for audit in [false, true] { + let stream = if audit { + read.to_audit_log_arrow(&plan) + } else { + read.to_incremental_arrow(&plan) + }; + let mut actual: Vec = Vec::new(); + for batch in stream.unwrap().try_collect::>().await.unwrap() { + let column = batch.column_by_name("id").unwrap().as_any(); + actual.extend(column.downcast_ref::().unwrap().values()); + } + actual.sort_unstable(); + assert_eq!(actual, [1, 3], "{mode} audit={audit}"); + } + } +} diff --git a/docs/src/sql.md b/docs/src/sql.md index b00cdfafe..f921ecf8f 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -69,9 +69,10 @@ async fn example() -> Result<(), Box> { } ``` -`SQLContext::new` creates a session context with the Paimon relation planner and +`SQLContext::new` creates a session context with the Paimon relation planner, the catalog-independent `path_to_descriptor` and `descriptor_to_string` scalar -functions pre-registered. Use `register_catalog(...).await` to add one or more +functions and the `paimon_incremental_query` table-valued function +pre-registered. Use `register_catalog(...).await` to add one or more Paimon catalogs; registering a catalog also registers the built-in scalar function `blob_view` (alias `sys.blob_view`) and the built-in table-valued functions (`vector_search`, `hybrid_search`, and `full_text_search` when the @@ -2310,6 +2311,44 @@ given value (snapshots without a watermark are skipped). It is mutually exclusive with the other time-travel selectors. If no matching snapshot exists, scan planning fails. +## Incremental Query + +`paimon_incremental_query` reads the changes between two snapshots. The range +excludes the start snapshot and includes the end snapshot: + +```sql +SELECT * FROM paimon_incremental_query('my_db.my_table', 1, 5); +SELECT * FROM paimon_incremental_query('my_db.my_table$audit_log', 1, 5, 'changelog'); +``` + +The optional fourth argument selects the mode: + +- `delta` returns the rows each APPEND commit wrote, every version of a key included. + A data-evolution table's range is read as one batch, as Java's `incremental-between` + scan reads it, so a column update merges into the rows it updates when they were + written in the range. For rows written before the range it returns every row of the + files the updates rewrote, unchanged rows included, with only the updated columns + set and the others NULL, NOT NULL columns included. (Java's reader instead fails + when two such updates of the same rows leave a NOT NULL column unwritten.) +- `changelog` returns the changelog rows of each commit. +- `diff` compares the table at the two snapshots and returns the added and changed + rows (`$audit_log` adds `-U` and `-D` rows). It requires a primary-key table with + `merge-engine=deduplicate` and no BLOB or VARIANT columns. +- `auto`, the default, is `delta` when `changelog-producer` is `none` and + `changelog` otherwise. + +A `$audit_log` suffix puts `rowkind` first, followed by `_SEQUENCE_NUMBER` when +`table-read.sequence-number.enabled=true`. + +The table name resolves as in a `FROM` clause of the same session: `my_table` and +`my_db.my_table` use the current catalog (and database), and +`my_catalog.my_db.my_table` names the catalog. Dynamic options set with +`SET 'paimon.…'` apply as they do to a query on the table, except options that +choose another range to read (time-travel selectors, `incremental-between*`, file +creation-time filters and a `scan.mode` other than `default`), which are rejected +whether the session or the table read sets them. With `SET 'paimon.branch'`, the +table read is that branch. + ## Dynamic Options (SET / RESET) Use `SET` to configure session-scoped Paimon dynamic options that apply to subsequent table loads: