From 4fad44b28836ed1c1494ee1ba92f05f03c1bedff Mon Sep 17 00:00:00 2001 From: jackylee-ch Date: Mon, 5 Oct 2026 20:48:05 +0800 Subject: [PATCH 1/2] feat(datafusion): support the paimon_incremental_query table function --- .../datafusion/src/incremental_query.rs | 365 ++++++++++++++++++ crates/integrations/datafusion/src/lib.rs | 1 + .../datafusion/src/sql_context.rs | 5 + .../datafusion/tests/incremental_query.rs | 201 ++++++++++ 4 files changed, 572 insertions(+) create mode 100644 crates/integrations/datafusion/src/incremental_query.rs create mode 100644 crates/integrations/datafusion/tests/incremental_query.rs diff --git a/crates/integrations/datafusion/src/incremental_query.rs b/crates/integrations/datafusion/src/incremental_query.rs new file mode 100644 index 000000000..420b96df0 --- /dev/null +++ b/crates/integrations/datafusion/src/incremental_query.rs @@ -0,0 +1,365 @@ +// 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 only parses arguments and streams the core plan's 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. + +use std::fmt; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::catalog::{Session, 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, TryStreamExt}; +use paimon::catalog::Catalog; +use paimon::table::{AuditLogTable, IncrementalScanMode, Table}; + +use crate::error::to_datafusion_error; +use crate::runtime::{await_with_runtime, block_on_with_runtime}; +use crate::table::datafusion_arrow_schema; +use crate::table_function_args::{ + extract_int_literal, extract_string_literal, parse_table_identifier, +}; +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` against `catalog`. +pub(crate) fn register_incremental_query( + ctx: &SessionContext, + catalog: Arc, + default_database: &str, +) { + ctx.register_udtf( + FUNCTION_NAME, + Arc::new(IncrementalQueryFunction { + catalog, + default_database: default_database.to_string(), + }), + ); +} + +pub(crate) struct IncrementalQueryFunction { + catalog: Arc, + default_database: String, +} + +impl fmt::Debug for IncrementalQueryFunction { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("IncrementalQueryFunction") + .field("default_database", &self.default_database) + .finish() + } +} + +/// 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(&self, args: &[Expr]) -> DFResult> { + if args.len() < 3 || args.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, &args[0], "table_name")?; + let start_exclusive = + extract_int_literal(FUNCTION_NAME, &args[1], "start_snapshot_exclusive")?; + let end_inclusive = extract_int_literal(FUNCTION_NAME, &args[2], "end_snapshot_inclusive")?; + let mode = match args.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), + }; + let identifier = parse_table_identifier(FUNCTION_NAME, &base_name, &self.default_database)?; + + let catalog = Arc::clone(&self.catalog); + let table = block_on_with_runtime( + async move { load_data_table_for_read(&catalog, &identifier, FUNCTION_NAME).await }, + "paimon_incremental_query: catalog access thread panicked", + )?; + + Ok(Arc::new(IncrementalQueryTableProvider::try_new( + table, + audit_log, + mode, + start_exclusive, + end_inclusive, + )?)) + } +} +#[derive(Debug, Clone)] +struct IncrementalQueryTableProvider { + table: Table, + audit_log: bool, + mode: IncrementalScanMode, + start_exclusive: i64, + end_inclusive: i64, + 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, + schema, + }) + } +} + +#[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.cloned(), + projected_schema, + ))) + } +} +/// 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, + } + } + + async fn compute_batches(&self) -> DFResult> { + // Core owns snapshot planning; its stream is `'static`, so it outlives + // the transient read builder it is built from. + let batches = await_with_runtime(async { + let read_builder = self.table.new_read_builder(); + let scan = read_builder.new_incremental_scan( + self.mode, + self.start_exclusive, + self.end_inclusive, + ); + let plan = scan.plan().await.map_err(to_datafusion_error)?; + let read = read_builder.new_read().map_err(to_datafusion_error)?; + let mut stream = if self.audit_log { + read.to_audit_log_arrow(&plan) + } else { + read.to_incremental_arrow(&plan) + } + .map_err(to_datafusion_error)?; + let mut batches = Vec::new(); + while let Some(batch) = stream.try_next().await.map_err(to_datafusion_error)? { + batches.push(batch); + } + Ok::<_, DataFusionError>(batches) + }) + .await?; + + match &self.projection { + Some(indices) => batches + .iter() + .map(|batch| batch.project(indices).map_err(DataFusionError::from)) + .collect(), + None => Ok(batches), + } + } +} +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 batches = exec.compute_batches().await?; + Ok::<_, DataFusionError>(stream::iter(batches.into_iter().map(Ok))) + }) + .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 7043303e0..bb17cb732 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_into; diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index f5f27276b..0236417be 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -3842,6 +3842,11 @@ fn register_table_functions( #[cfg(feature = "fulltext")] crate::full_text_search::register_full_text_search(ctx, Arc::clone(catalog), default_database); crate::hybrid_search::register_hybrid_search(ctx, Arc::clone(catalog), default_database); + crate::incremental_query::register_incremental_query( + ctx, + Arc::clone(catalog), + default_database, + ); } #[cfg(test)] diff --git a/crates/integrations/datafusion/tests/incremental_query.rs b/crates/integrations/datafusion/tests/incremental_query.rs new file mode 100644 index 000000000..8510d4b80 --- /dev/null +++ b/crates/integrations/datafusion/tests/incremental_query.rs @@ -0,0 +1,201 @@ +// 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 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() +} + +/// 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 warehouse = format!("file://{}", temp_dir.path().display()); + let mut options = Options::new(); + options.set(CatalogOptions::WAREHOUSE, warehouse); + let catalog = Arc::new(FileSystemCatalog::new(options).unwrap()); + let mut ctx = SQLContext::new(); + ctx.register_catalog("paimon", catalog).await.unwrap(); + + 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}" + ); +} From 95681813d97441e5e8d12d27ee6186fae06dc8d1 Mon Sep 17 00:00:00 2001 From: jackylee-ch Date: Thu, 15 Oct 2026 19:43:22 +0800 Subject: [PATCH 2/2] fix: stream incremental reads and decode them as ordinary reads `IncrementalQueryExec` collected the whole snapshot range into a `Vec` before emitting a batch, decoded every column even for a narrow projection, and ignored the executing `TaskContext`, so the DataFusion memory pool never saw the reader's working memory. Forward core's batch stream lazily, push the projection into the read (`$audit_log` columns are mapped by position, one read field per projected column), attach the task's reader resources, and rebuild each batch against the declared (projected) schema, as the ordinary scan does. Core's `to_incremental_arrow` decoded each planned file raw, so a data-evolution BLOB or VECTOR column file came back as rows of its own and a primary-key BLOB column as its stored descriptor. Decode those files as `to_arrow` does, read limit included, and plan a data-evolution Delta range as one batch, as Java's `incremental-between` scan does, so a column update merges into the rows it updates when both are in the range. Document `paimon_incremental_query` in the SQL guide. --- .../datafusion/src/incremental_query.rs | 154 +++++-- .../datafusion/src/physical_plan/scan.rs | 5 +- .../datafusion/tests/incremental_query.rs | 405 +++++++++++++++++- crates/paimon/src/table/table_read.rs | 41 +- .../tests/incremental_batch_scan_test.rs | 291 +++++++++++++ crates/paimon/tests/partition_filter_test.rs | 35 ++ docs/src/sql.md | 34 +- 7 files changed, 908 insertions(+), 57 deletions(-) diff --git a/crates/integrations/datafusion/src/incremental_query.rs b/crates/integrations/datafusion/src/incremental_query.rs index 420b96df0..cdf2f0ab0 100644 --- a/crates/integrations/datafusion/src/incremental_query.rs +++ b/crates/integrations/datafusion/src/incremental_query.rs @@ -17,7 +17,8 @@ //! 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 only parses arguments and streams the core plan's output. +//! 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); @@ -33,7 +34,6 @@ use std::sync::Arc; use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; -use datafusion::arrow::record_batch::RecordBatch; use datafusion::catalog::{Session, TableFunctionImpl}; use datafusion::common::stats::Precision; use datafusion::common::{internal_err, project_schema, Statistics}; @@ -49,11 +49,18 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, }; use datafusion::prelude::SessionContext; -use futures::{stream, TryStreamExt}; +use futures::{stream, StreamExt, TryStreamExt}; use paimon::catalog::Catalog; -use paimon::table::{AuditLogTable, IncrementalScanMode, Table}; +use paimon::resource::ResourceContext; +use paimon::spec::{CoreOptions, DataField}; +use paimon::table::{ + ArrowRecordBatchStream, AuditLogTable, IncrementalPlan, IncrementalScanMode, IncrementalSplit, + Table, +}; 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::table::datafusion_arrow_schema; use crate::table_function_args::{ @@ -156,6 +163,7 @@ struct IncrementalQueryTableProvider { mode: IncrementalScanMode, start_exclusive: i64, end_inclusive: i64, + fields: Vec, schema: ArrowSchemaRef, } @@ -185,9 +193,29 @@ impl IncrementalQueryTableProvider { 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] @@ -218,11 +246,20 @@ impl TableProvider for IncrementalQueryTableProvider { self.mode, self.start_exclusive, self.end_inclusive, - projection.cloned(), + 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)] @@ -232,7 +269,7 @@ struct IncrementalQueryExec { mode: IncrementalScanMode, start_exclusive: i64, end_inclusive: i64, - projection: Option>, + projection: Option, output_schema: ArrowSchemaRef, plan_properties: Arc, } @@ -244,7 +281,7 @@ impl IncrementalQueryExec { mode: IncrementalScanMode, start_exclusive: i64, end_inclusive: i64, - projection: Option>, + projection: Option, output_schema: ArrowSchemaRef, ) -> Self { let plan_properties = Arc::new(PlanProperties::new( @@ -265,39 +302,60 @@ impl IncrementalQueryExec { } } - async fn compute_batches(&self) -> DFResult> { - // Core owns snapshot planning; its stream is `'static`, so it outlives - // the transient read builder it is built from. - let batches = await_with_runtime(async { - let read_builder = self.table.new_read_builder(); - let scan = read_builder.new_incremental_scan( - self.mode, - self.start_exclusive, - self.end_inclusive, - ); - let plan = scan.plan().await.map_err(to_datafusion_error)?; - let read = read_builder.new_read().map_err(to_datafusion_error)?; - let mut stream = if self.audit_log { - read.to_audit_log_arrow(&plan) - } else { - read.to_incremental_arrow(&plan) - } - .map_err(to_datafusion_error)?; - let mut batches = Vec::new(); - while let Some(batch) = stream.try_next().await.map_err(to_datafusion_error)? { - batches.push(batch); - } - Ok::<_, DataFusionError>(batches) - }) - .await?; - - match &self.projection { - Some(indices) => batches - .iter() - .map(|batch| batch.project(indices).map_err(DataFusionError::from)) - .collect(), - None => Ok(batches), + /// 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 { @@ -336,7 +394,7 @@ impl ExecutionPlan for IncrementalQueryExec { fn execute( &self, partition: usize, - _context: Arc, + context: Arc, ) -> DFResult { if partition != 0 { return internal_err!( @@ -345,8 +403,20 @@ impl ExecutionPlan for IncrementalQueryExec { } let exec = self.clone(); let stream = stream::once(async move { - let batches = exec.compute_batches().await?; - Ok::<_, DataFusionError>(stream::iter(batches.into_iter().map(Ok))) + 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( diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs b/crates/integrations/datafusion/src/physical_plan/scan.rs index 0a9eb5ab6..e3ec6d115 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/tests/incremental_query.rs b/crates/integrations/datafusion/tests/incremental_query.rs index 8510d4b80..e55da99a0 100644 --- a/crates/integrations/datafusion/tests/incremental_query.rs +++ b/crates/integrations/datafusion/tests/incremental_query.rs @@ -20,6 +20,11 @@ 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; @@ -32,16 +37,45 @@ 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 warehouse = format!("file://{}", temp_dir.path().display()); - let mut options = Options::new(); - options.set(CatalogOptions::WAREHOUSE, warehouse); - let catalog = Arc::new(FileSystemCatalog::new(options).unwrap()); - let mut ctx = SQLContext::new(); - ctx.register_catalog("paimon", catalog).await.unwrap(); + 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; @@ -199,3 +233,362 @@ async fn test_rejects_end_before_start() { "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}]"); + } +} diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index 4ae7b7112..951f868a7 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -466,10 +466,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( @@ -518,6 +529,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()) } @@ -1519,6 +1538,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}; @@ -1531,7 +1562,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, @@ -1541,7 +1572,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..60eb2846f 100644 --- a/crates/paimon/tests/incremental_batch_scan_test.rs +++ b/crates/paimon/tests/incremental_batch_scan_test.rs @@ -1811,3 +1811,294 @@ 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}" + ); + } +} 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 120b8a333..89c6d3565 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -74,9 +74,9 @@ the catalog-independent `path_to_descriptor` and `descriptor_to_string` scalar functions 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 -`fulltext` feature is enabled) against it. It also manages session-scoped -dynamic options internally for `SET`/`RESET` support. +functions (`vector_search`, `hybrid_search`, `paimon_incremental_query`, and +`full_text_search` when the `fulltext` feature is enabled) against it. It also +manages session-scoped dynamic options internally for `SET`/`RESET` support. ### REST Catalog Views and SQL Functions @@ -2231,6 +2231,34 @@ 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 update rewrote, unchanged rows included, with only the updated columns + set and the others NULL, NOT NULL columns included. +- `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`. + ## Dynamic Options (SET / RESET) Use `SET` to configure session-scoped Paimon dynamic options that apply to subsequent table loads: