From 5509e0622db36b66179f251c7088902ea1026800 Mon Sep 17 00:00:00 2001 From: pandas886 <123344357+pandas886@users.noreply.github.com> Date: Sat, 11 Jul 2026 22:44:37 +0800 Subject: [PATCH 1/2] feat: support DataFusion paimon_incremental_query --- .../datafusion/src/incremental_query.rs | 217 ++++++++++ crates/integrations/datafusion/src/lib.rs | 2 + .../src/physical_plan/incremental_scan.rs | 187 +++++++++ .../datafusion/src/physical_plan/mod.rs | 2 + .../datafusion/src/sql_context.rs | 22 +- .../datafusion/src/table_function_args.rs | 180 +++++++- .../datafusion/tests/incremental_query.rs | 387 ++++++++++++++++++ docs/src/sql.md | 72 +++- 8 files changed, 1057 insertions(+), 12 deletions(-) create mode 100644 crates/integrations/datafusion/src/incremental_query.rs create mode 100644 crates/integrations/datafusion/src/physical_plan/incremental_scan.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..d21f46055 --- /dev/null +++ b/crates/integrations/datafusion/src/incremental_query.rs @@ -0,0 +1,217 @@ +// 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. + +//! `paimon_incremental_query` table-valued function for DataFusion. +//! +//! Usage: +//! ```sql +//! SELECT * FROM paimon_incremental_query('table_name', start_snapshot_exclusive, end_snapshot_inclusive); +//! SELECT * FROM paimon_incremental_query('table_name$audit_log', 0, 2, 'diff'); +//! ``` +//! +//! This module only parses SQL arguments and adapts the resulting +//! [`IncrementalPlan`](paimon::table::IncrementalPlan) into a DataFusion +//! execution plan. Snapshot planning lives in paimon-core +//! (`IncrementalScan` / `AuditLogTable`). + +use std::fmt::Debug; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; +use datafusion::catalog::Session; +use datafusion::catalog::TableFunctionImpl; +use datafusion::common::project_schema; +use datafusion::datasource::{TableProvider, TableType}; +use datafusion::error::Result as DFResult; +use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionContext; +use paimon::catalog::Catalog; +use paimon::table::{AuditLogTable, IncrementalScanMode, Table}; + +use crate::error::to_datafusion_error; +use crate::physical_plan::PaimonIncrementalScan; +use crate::runtime::{await_with_runtime, block_on_with_runtime}; +use crate::table::PaimonTableProvider; +use crate::table_function_args::{ + extract_snapshot_bound_literal, extract_string_literal, parse_incremental_scan_mode, + parse_incremental_table_ref, +}; + +const FUNCTION_NAME: &str = "paimon_incremental_query"; + +/// Register the `paimon_incremental_query` table-valued function on a [`SessionContext`]. +pub fn register_incremental_query( + ctx: &SessionContext, + catalog: Arc, + default_database: &str, +) { + ctx.register_udtf( + FUNCTION_NAME, + Arc::new(IncrementalQueryFunction::new(catalog, default_database)), + ); +} + +/// Table function that performs batch incremental reads between snapshot IDs. +/// +/// Arguments: +/// `(table_name STRING, start_snapshot_exclusive BIGINT, end_snapshot_inclusive BIGINT [, scan_mode STRING])` +pub struct IncrementalQueryFunction { + catalog: Arc, + default_database: String, +} + +impl Debug for IncrementalQueryFunction { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("IncrementalQueryFunction") + .field("default_database", &self.default_database) + .finish() + } +} + +impl IncrementalQueryFunction { + pub fn new(catalog: Arc, default_database: &str) -> Self { + Self { + catalog, + default_database: default_database.to_string(), + } + } +} + +impl TableFunctionImpl for IncrementalQueryFunction { + fn call(&self, args: &[Expr]) -> DFResult> { + if args.len() != 3 && args.len() != 4 { + return Err(datafusion::error::DataFusionError::Plan(format!( + "{FUNCTION_NAME} requires 3 or 4 arguments: \ + (table_name, start_snapshot_exclusive, end_snapshot_inclusive [, scan_mode])" + ))); + } + + let table_name = extract_string_literal(FUNCTION_NAME, &args[0], "table_name")?; + let start_exclusive = + extract_snapshot_bound_literal(FUNCTION_NAME, &args[1], "start_snapshot_exclusive")?; + let end_inclusive = + extract_snapshot_bound_literal(FUNCTION_NAME, &args[2], "end_snapshot_inclusive")?; + let scan_mode = if args.len() == 4 { + parse_incremental_scan_mode(FUNCTION_NAME, &args[3])? + } else { + IncrementalScanMode::Auto + }; + + let table_ref = + parse_incremental_table_ref(FUNCTION_NAME, &table_name, &self.default_database)?; + + let catalog = Arc::clone(&self.catalog); + let identifier = table_ref.identifier.clone(); + let table = block_on_with_runtime( + async move { catalog.get_table(&identifier).await }, + "paimon_incremental_query: catalog access thread panicked", + ) + .map_err(to_datafusion_error)?; + + let schema = if table_ref.audit_log { + let fields = AuditLogTable::new(table.clone()) + .fields() + .map_err(to_datafusion_error)?; + paimon::arrow::build_target_arrow_schema(&fields).map_err(to_datafusion_error)? + } else { + PaimonTableProvider::try_new(table.clone())?.schema() + }; + + Ok(Arc::new(IncrementalQueryTableProvider { + table, + schema, + start_exclusive, + end_inclusive, + scan_mode, + audit_log: table_ref.audit_log, + })) + } +} + +#[derive(Debug)] +struct IncrementalQueryTableProvider { + table: Table, + schema: ArrowSchemaRef, + start_exclusive: i64, + end_inclusive: i64, + scan_mode: IncrementalScanMode, + audit_log: bool, +} + +#[async_trait] +impl TableProvider for IncrementalQueryTableProvider { + fn schema(&self) -> ArrowSchemaRef { + self.schema.clone() + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> DFResult> { + let table = self.table.clone(); + let start_exclusive = self.start_exclusive; + let end_inclusive = self.end_inclusive; + let scan_mode = self.scan_mode; + + // Planning is delegated entirely to paimon-core IncrementalScan. + let incremental_plan = await_with_runtime(async move { + table + .new_read_builder() + .new_incremental_scan(scan_mode, start_exclusive, end_inclusive) + .plan() + .await + .map_err(to_datafusion_error) + }) + .await?; + + let projected_schema = project_schema(&self.schema, projection)?; + let projected_columns = projection.map(|indices| { + indices + .iter() + .map(|&i| self.schema.field(i).name().clone()) + .collect::>() + }); + + Ok(Arc::new(PaimonIncrementalScan::new( + projected_schema, + self.table.clone(), + incremental_plan, + self.audit_log, + projected_columns, + ))) + } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> DFResult> { + // Residual filter evaluation is performed by DataFusion above this plan. + Ok(vec![ + TableProviderFilterPushDown::Unsupported; + filters.len() + ]) + } +} diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index 05e4037a2..3ddb65a94 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -46,6 +46,7 @@ mod filter_pushdown; #[cfg(feature = "fulltext")] mod full_text_search; mod hybrid_search; +mod incremental_query; mod lateral_vector_search; mod merge_into; mod physical_plan; @@ -79,6 +80,7 @@ pub use error::to_datafusion_error; #[cfg(feature = "fulltext")] pub use full_text_search::{register_full_text_search, FullTextSearchFunction}; pub use hybrid_search::{register_hybrid_search, HybridSearchFunction}; +pub use incremental_query::{register_incremental_query, IncrementalQueryFunction}; pub use physical_plan::PaimonTableScan; pub use relation_planner::PaimonRelationPlanner; pub use sql_context::SQLContext; diff --git a/crates/integrations/datafusion/src/physical_plan/incremental_scan.rs b/crates/integrations/datafusion/src/physical_plan/incremental_scan.rs new file mode 100644 index 000000000..b31038769 --- /dev/null +++ b/crates/integrations/datafusion/src/physical_plan/incremental_scan.rs @@ -0,0 +1,187 @@ +// 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. + +//! Physical plan for batch incremental reads via paimon-core +//! [`IncrementalScan`](paimon::table::IncrementalScan) / +//! [`AuditLogTable`](paimon::table::AuditLogTable). + +use std::sync::Arc; + +use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; +use datafusion::error::Result as DFResult; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr::EquivalenceProperties; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{DisplayAs, ExecutionPlan, Partitioning, PlanProperties}; +use futures::{StreamExt, TryStreamExt}; +use paimon::table::{AuditLogTable, IncrementalPlan, Table}; + +use crate::error::to_datafusion_error; + +pub(crate) const PLAN_NAME: &str = "PaimonIncrementalScan"; + +/// Execution plan that streams incremental snapshot-range rows from a Paimon table. +/// +/// Planning of snapshot ranges is performed eagerly in +/// [`crate::incremental_query::IncrementalQueryFunction`]; this plan only executes +/// the precomputed [`IncrementalPlan`]. +#[derive(Debug)] +pub(crate) struct PaimonIncrementalScan { + table: Table, + incremental_plan: IncrementalPlan, + audit_log: bool, + projected_columns: Option>, + schema: ArrowSchemaRef, + plan_properties: Arc, +} + +impl PaimonIncrementalScan { + pub(crate) fn new( + schema: ArrowSchemaRef, + table: Table, + incremental_plan: IncrementalPlan, + audit_log: bool, + projected_columns: Option>, + ) -> Self { + let plan_properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + table, + incremental_plan, + audit_log, + projected_columns, + schema, + plan_properties, + } + } +} + +impl ExecutionPlan for PaimonIncrementalScan { + fn name(&self) -> &str { + PLAN_NAME + } + + fn properties(&self) -> &Arc { + &self.plan_properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DFResult> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> DFResult { + let table = self.table.clone(); + let plan = self.incremental_plan.clone(); + let audit_log = self.audit_log; + let projected_columns = self.projected_columns.clone(); + let schema = self.schema.clone(); + + let fut = async move { + let stream = if audit_log { + // AuditLogTable always materialises the full audit schema + // (rowkind + table fields). Column projection is applied below. + AuditLogTable::new(table) + .to_arrow(&plan) + .map_err(to_datafusion_error)? + } else { + let mut read_builder = table.new_read_builder(); + if let Some(ref columns) = projected_columns { + let col_refs: Vec<&str> = columns.iter().map(|s| s.as_str()).collect(); + read_builder + .with_projection(&col_refs) + .map_err(to_datafusion_error)?; + } + let read = read_builder.new_read().map_err(to_datafusion_error)?; + read.to_incremental_arrow(&plan) + .map_err(to_datafusion_error)? + }; + + let stream = stream.map(move |result| { + let batch = result.map_err(to_datafusion_error)?; + if audit_log { + if let Some(ref columns) = projected_columns { + let indices = columns + .iter() + .map(|name| { + batch.schema().index_of(name).map_err(|e| { + datafusion::error::DataFusionError::Internal(format!( + "PaimonIncrementalScan: projection column '{name}': {e}" + )) + }) + }) + .collect::>>()?; + return batch.project(&indices).map_err(|e| { + datafusion::error::DataFusionError::ArrowError(Box::new(e), None) + }); + } + } + Ok(batch) + }); + + Ok::<_, datafusion::error::DataFusionError>(RecordBatchStreamAdapter::new( + schema, + Box::pin(stream), + )) + }; + + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + futures::stream::once(fut).try_flatten(), + ))) + } +} + +impl DisplayAs for PaimonIncrementalScan { + fn fmt_as( + &self, + _format: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!( + f, + "PaimonIncrementalScan(audit_log={}, splits={})", + self.audit_log, + self.incremental_plan.splits().len() + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn incremental_scan_plan_name_constant() { + assert_eq!(PLAN_NAME, "PaimonIncrementalScan"); + } +} diff --git a/crates/integrations/datafusion/src/physical_plan/mod.rs b/crates/integrations/datafusion/src/physical_plan/mod.rs index 2d1905035..fa57847ee 100644 --- a/crates/integrations/datafusion/src/physical_plan/mod.rs +++ b/crates/integrations/datafusion/src/physical_plan/mod.rs @@ -15,10 +15,12 @@ // specific language governing permissions and limitations // under the License. +pub(crate) mod incremental_scan; pub(crate) mod scan; mod search_score; pub(crate) mod sink; +pub(crate) use incremental_scan::PaimonIncrementalScan; pub use scan::PaimonTableScan; pub(crate) use search_score::{SearchScoreExec, SearchScoreOutputColumn}; pub use sink::PaimonDataSink; diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 1803892b2..865744dc9 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -160,11 +160,11 @@ impl SQLContext { /// /// `default_db = Some("")` is rejected — pass `None` to opt out instead. /// - /// **Note on built-in TVFs (`vector_search`, `full_text_search`):** when - /// `default_db = None`, bare table names inside these functions still resolve - /// against the literal namespace `"default"` (the fallback in - /// [`register_table_functions`]). Callers using `None` must qualify table names - /// (`'db.table'` or `'catalog.db.table'`) in those calls. + /// **Note on built-in TVFs (`paimon_incremental_query`, `vector_search`, + /// `full_text_search`):** when `default_db = None`, bare table names inside + /// these functions still resolve against the literal namespace `"default"` + /// (the fallback in [`register_table_functions`]). Callers using `None` must + /// qualify table names (`'db.table'` or `'catalog.db.table'`) in those calls. pub async fn register_catalog_with_default_db( &mut self, catalog_name: impl Into, @@ -3483,6 +3483,11 @@ fn register_table_functions( dynamic_options: DynamicOptions, ) { crate::blob_view::register_blob_view(ctx, Arc::clone(catalog), default_database); + crate::incremental_query::register_incremental_query( + ctx, + Arc::clone(catalog), + default_database, + ); crate::vector_search::register_vector_search_with_dynamic_options( ctx, Arc::clone(catalog), @@ -5683,9 +5688,10 @@ mod tests { async fn register_catalog_with_none_table_function_resolves_bare_name_to_literal_default() { // Documents the fallback in register_table_functions: `default_db.unwrap_or("default")`. // When the caller opts out of default-db init, bare table names inside built-in TVFs - // (vector_search / full_text_search) still resolve against the literal namespace - // `"default"` — so a caller using `None` MUST use fully-qualified names with these - // functions or they'll hit a `default.` lookup that may not exist / be readable. + // (paimon_incremental_query / vector_search / full_text_search) still resolve + // against the literal namespace `"default"` — so a caller using `None` MUST use + // fully-qualified names with these functions or they'll hit a `default.` + // lookup that may not exist / be readable. let catalog = Arc::new(ProbeTrackingCatalog::new()); let mut ctx = SQLContext::new(); ctx.register_catalog_with_default_db("paimon", catalog, None) diff --git a/crates/integrations/datafusion/src/table_function_args.rs b/crates/integrations/datafusion/src/table_function_args.rs index c0a6c833e..d637dfab6 100644 --- a/crates/integrations/datafusion/src/table_function_args.rs +++ b/crates/integrations/datafusion/src/table_function_args.rs @@ -18,7 +18,109 @@ use datafusion::common::ScalarValue; use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::logical_expr::Expr; -use paimon::catalog::Identifier; +use paimon::catalog::{Identifier, SYSTEM_TABLE_SPLITTER}; +use paimon::table::IncrementalScanMode; + +/// Parsed target of `paimon_incremental_query`: base table plus optional `$audit_log`. +#[derive(Debug, Clone)] +pub(crate) struct IncrementalTableRef { + pub identifier: Identifier, + pub audit_log: bool, +} + +/// Parse table name for incremental query TVF. +/// +/// Accepts `table`, `database.table`, or `catalog.database.table`, with optional +/// `$audit_log` suffix on the object name. +pub(crate) fn parse_incremental_table_ref( + function_name: &str, + name: &str, + default_database: &str, +) -> DFResult { + let parts: Vec<&str> = name.split('.').collect(); + let (database, object_name) = match parts.len() { + 1 => (default_database, parts[0]), + 2 => (parts[0], parts[1]), + 3 => (parts[1], parts[2]), + _ => { + return Err(DataFusionError::Plan(format!( + "{function_name}: invalid table name '{name}', expected 'table', 'database.table', or 'catalog.database.table'" + ))); + } + }; + + let mut object_parts = object_name.splitn(2, SYSTEM_TABLE_SPLITTER); + let base = object_parts.next().unwrap_or(object_name); + let suffix = object_parts.next(); + let audit_log = match suffix { + None => false, + Some("audit_log") => true, + Some(other) => { + return Err(DataFusionError::Plan(format!( + "{function_name}: unsupported system table suffix '${other}' in '{name}'" + ))); + } + }; + + Ok(IncrementalTableRef { + identifier: Identifier::new(database.to_string(), base.to_string()), + audit_log, + }) +} + +/// Accept integer literals or decimal strings (Spark `'0'`, `'1'` compatibility). +pub(crate) fn extract_snapshot_bound_literal( + function_name: &str, + expr: &Expr, + name: &str, +) -> DFResult { + match expr { + Expr::Literal(scalar, _) => match scalar { + ScalarValue::Int8(Some(v)) => Ok(*v as i64), + ScalarValue::Int16(Some(v)) => Ok(*v as i64), + ScalarValue::Int32(Some(v)) => Ok(*v as i64), + ScalarValue::Int64(Some(v)) => Ok(*v), + ScalarValue::UInt8(Some(v)) => Ok(*v as i64), + ScalarValue::UInt16(Some(v)) => Ok(*v as i64), + ScalarValue::UInt32(Some(v)) => Ok(*v as i64), + ScalarValue::UInt64(Some(v)) => i64::try_from(*v).map_err(|_| { + DataFusionError::Plan(format!( + "{function_name}: {name} value {v} exceeds i64 range" + )) + }), + ScalarValue::Utf8(Some(s)) | ScalarValue::LargeUtf8(Some(s)) => { + s.parse::().map_err(|_| { + DataFusionError::Plan(format!( + "{function_name}: {name} must be an integer snapshot id, got string '{s}'" + )) + }) + } + _ => Err(DataFusionError::Plan(format!( + "{function_name}: {name} must be an integer or numeric string literal, got: {expr}" + ))), + }, + _ => Err(DataFusionError::Plan(format!( + "{function_name}: {name} must be a literal, got: {expr}" + ))), + } +} + +/// Parse optional TVF 4th argument; aligns with Java `incremental-between-scan-mode`. +pub(crate) fn parse_incremental_scan_mode( + function_name: &str, + expr: &Expr, +) -> DFResult { + let mode = extract_string_literal(function_name, expr, "scan_mode")?; + match mode.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}: unsupported scan_mode '{other}', expected auto, delta, changelog, or diff" + ))), + } +} pub(crate) fn extract_string_literal( function_name: &str, @@ -80,3 +182,79 @@ pub(crate) fn parse_table_identifier( ))), } } + +#[cfg(test)] +mod incremental_query_args_tests { + use super::*; + + #[test] + fn parse_incremental_table_ref_plain_name() { + let parsed = + parse_incremental_table_ref("paimon_incremental_query", "orders", "default").unwrap(); + assert_eq!(parsed.identifier.database(), "default"); + assert_eq!(parsed.identifier.object(), "orders"); + assert!(!parsed.audit_log); + } + + #[test] + fn parse_incremental_table_ref_with_audit_log_suffix() { + let parsed = parse_incremental_table_ref( + "paimon_incremental_query", + "paimon.test_db.orders$audit_log", + "default", + ) + .unwrap(); + assert_eq!(parsed.identifier.database(), "test_db"); + assert_eq!(parsed.identifier.object(), "orders"); + assert!(parsed.audit_log); + } + + #[test] + fn extract_snapshot_bound_accepts_int_and_numeric_string() { + let int_expr = Expr::Literal(ScalarValue::Int64(Some(2)), None); + let str_expr = Expr::Literal(ScalarValue::Utf8(Some("2".to_string())), None); + assert_eq!( + extract_snapshot_bound_literal("paimon_incremental_query", &int_expr, "end").unwrap(), + 2 + ); + assert_eq!( + extract_snapshot_bound_literal("paimon_incremental_query", &str_expr, "end").unwrap(), + 2 + ); + } + + #[test] + fn parse_incremental_table_ref_rejects_unknown_system_suffix() { + let err = + parse_incremental_table_ref("paimon_incremental_query", "orders$snapshots", "default") + .unwrap_err() + .to_string(); + assert!(err.contains("snapshots"), "unexpected error: {err}"); + } + + #[test] + fn parse_incremental_scan_mode_accepts_aliases() { + let lit = |s: &str| Expr::Literal(ScalarValue::Utf8(Some(s.to_string())), None); + assert_eq!( + parse_incremental_scan_mode("paimon_incremental_query", &lit("auto")).unwrap(), + IncrementalScanMode::Auto + ); + assert_eq!( + parse_incremental_scan_mode("paimon_incremental_query", &lit("DIFF")).unwrap(), + IncrementalScanMode::Diff + ); + assert_eq!( + parse_incremental_scan_mode("paimon_incremental_query", &lit("changelog")).unwrap(), + IncrementalScanMode::Changelog + ); + } + + #[test] + fn parse_incremental_scan_mode_rejects_unknown() { + let lit = Expr::Literal(ScalarValue::Utf8(Some("nope".to_string())), None); + let err = parse_incremental_scan_mode("paimon_incremental_query", &lit) + .unwrap_err() + .to_string(); + assert!(err.contains("nope"), "unexpected error: {err}"); + } +} diff --git a/crates/integrations/datafusion/tests/incremental_query.rs b/crates/integrations/datafusion/tests/incremental_query.rs new file mode 100644 index 000000000..f7d22d060 --- /dev/null +++ b/crates/integrations/datafusion/tests/incremental_query.rs @@ -0,0 +1,387 @@ +// 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 `paimon_incremental_query` TVF. +//! +//! Uses pure Rust `changelog-producer=input` tables (no compact / lookup). + +mod common; + +use std::sync::Arc; + +use arrow_array::{Array, Int32Array, Int8Array, RecordBatch}; +use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; +use datafusion::arrow::array::StringArray; +use paimon::catalog::Identifier; +use paimon::spec::VALUE_KIND_FIELD_NAME; +use paimon::{Catalog, FileSystemCatalog}; + +async fn setup_table_with_three_snapshots(sql: &paimon_datafusion::SQLContext) { + common::exec( + sql, + "CREATE TABLE paimon.test_db.inc_t ( + id INT NOT NULL, + value INT, + PRIMARY KEY (id) + ) WITH ('bucket' = '1', 'changelog-producer' = 'input')", + ) + .await; + common::exec(sql, "INSERT INTO paimon.test_db.inc_t VALUES (1, 10)").await; + common::exec(sql, "INSERT INTO paimon.test_db.inc_t VALUES (2, 20)").await; + common::exec(sql, "INSERT INTO paimon.test_db.inc_t VALUES (3, 30)").await; +} + +fn collect_id_value(batches: &[datafusion::arrow::record_batch::RecordBatch]) -> Vec<(i32, i32)> { + let mut rows = Vec::new(); + for batch in batches { + let ids = batch + .column_by_name("id") + .and_then(|c| c.as_any().downcast_ref::()) + .expect("id column"); + let values = batch + .column_by_name("value") + .and_then(|c| c.as_any().downcast_ref::()) + .expect("value column"); + for row in 0..batch.num_rows() { + rows.push((ids.value(row), values.value(row))); + } + } + rows.sort_unstable(); + rows +} + +fn collect_audit_rows( + batches: &[datafusion::arrow::record_batch::RecordBatch], +) -> Vec<(String, i32, i32)> { + let mut rows = Vec::new(); + for batch in batches { + let kinds = batch + .column_by_name("rowkind") + .and_then(|c| c.as_any().downcast_ref::()) + .expect("rowkind"); + let ids = batch + .column_by_name("id") + .and_then(|c| c.as_any().downcast_ref::()) + .expect("id"); + let values = batch + .column_by_name("value") + .and_then(|c| c.as_any().downcast_ref::()) + .expect("value"); + for row in 0..batch.num_rows() { + rows.push(( + kinds.value(row).to_string(), + ids.value(row), + values.value(row), + )); + } + } + rows.sort_unstable(); + rows +} + +/// Commit a single-row input-changelog delete (`RowKind::Delete` = 3). +async fn commit_input_changelog_delete( + catalog: &Arc, + table_name: &str, + id: i32, +) { + let table = catalog + .get_table(&Identifier::new("test_db", table_name)) + .await + .expect("table should exist"); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new(VALUE_KIND_FIELD_NAME, ArrowDataType::Int8, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![id])), + Arc::new(Int8Array::from(vec![3])), // RowKind::Delete + ], + ) + .unwrap(); + let wb = table.new_write_builder(); + let mut writer = wb.new_write().expect("writer"); + writer + .write_arrow_batch(&batch) + .await + .expect("write changelog delete"); + let messages = writer.prepare_commit().await.expect("prepare"); + wb.new_commit() + .commit(messages) + .await + .expect("commit changelog delete"); +} + +/// 3-arg form defaults to Auto (input producer → changelog semantics). +#[tokio::test] +async fn incremental_query_without_audit_log_returns_delta_rows() { + let (_tmp, ctx) = common::setup_sql_context().await; + setup_table_with_three_snapshots(&ctx).await; + + let batches = ctx + .sql("SELECT id, value FROM paimon_incremental_query('paimon.test_db.inc_t', 0, 2)") + .await + .expect("plan") + .collect() + .await + .expect("execute"); + + assert_eq!(collect_id_value(&batches), vec![(1, 10), (2, 20)]); +} + +#[tokio::test] +async fn incremental_query_with_audit_log_suffix_exposes_rowkind() { + let (_tmp, ctx) = common::setup_sql_context().await; + setup_table_with_three_snapshots(&ctx).await; + + let batches = ctx + .sql( + "SELECT rowkind, id, value \ + FROM paimon_incremental_query('paimon.test_db.inc_t$audit_log', 0, 1)", + ) + .await + .expect("plan") + .collect() + .await + .expect("execute"); + + let batch = &batches[0]; + assert_eq!(batch.schema().field(0).name(), "rowkind"); + let kinds = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(kinds.iter().all(|k| k == Some("+I"))); +} + +#[tokio::test] +async fn incremental_query_explicit_auto_and_changelog_match_default() { + let (_tmp, ctx) = common::setup_sql_context().await; + setup_table_with_three_snapshots(&ctx).await; + + let default_rows = ctx + .sql("SELECT id, value FROM paimon_incremental_query('paimon.test_db.inc_t', 0, 1)") + .await + .unwrap() + .collect() + .await + .unwrap(); + let auto_rows = ctx + .sql("SELECT id, value FROM paimon_incremental_query('paimon.test_db.inc_t', 0, 1, 'auto')") + .await + .unwrap() + .collect() + .await + .unwrap(); + let changelog_rows = ctx + .sql( + "SELECT id, value FROM paimon_incremental_query('paimon.test_db.inc_t', 0, 1, 'changelog')", + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + + assert_eq!( + collect_id_value(&default_rows), + collect_id_value(&auto_rows) + ); + assert_eq!( + collect_id_value(&default_rows), + collect_id_value(&changelog_rows) + ); +} + +#[tokio::test] +async fn incremental_query_explicit_delta_mode() { + let (_tmp, ctx) = common::setup_sql_context().await; + setup_table_with_three_snapshots(&ctx).await; + + let batches = ctx + .sql( + "SELECT id, value FROM paimon_incremental_query('paimon.test_db.inc_t', 0, 2, 'delta')", + ) + .await + .expect("plan") + .collect() + .await + .expect("execute"); + + // Delta reads data files from APPEND snapshots in (0, 2] → rows from snapshots 1 and 2. + assert_eq!(collect_id_value(&batches), vec![(1, 10), (2, 20)]); +} + +#[tokio::test] +async fn incremental_query_diff_mode_on_pk_table_with_audit_log() { + let (_tmp, ctx) = common::setup_sql_context().await; + common::exec( + &ctx, + "CREATE TABLE paimon.test_db.inc_pk ( + id INT NOT NULL, + value INT, + PRIMARY KEY (id) + ) WITH ('bucket' = '1', 'merge-engine' = 'deduplicate')", + ) + .await; + common::exec(&ctx, "INSERT INTO paimon.test_db.inc_pk VALUES (1, 10)").await; + common::exec(&ctx, "INSERT INTO paimon.test_db.inc_pk VALUES (2, 20)").await; + // PK deduplicate merge treats a later insert with the same key as an update. + common::exec(&ctx, "INSERT INTO paimon.test_db.inc_pk VALUES (1, 11)").await; + + let batches = ctx + .sql( + "SELECT rowkind, id, value \ + FROM paimon_incremental_query('paimon.test_db.inc_pk$audit_log', 1, 3, 'diff')", + ) + .await + .expect("plan") + .collect() + .await + .expect("execute"); + + let rows = collect_audit_rows(&batches); + assert!( + rows.iter() + .any(|(k, id, v)| k == "-U" && *id == 1 && *v == 10), + "diff audit_log should include -U before image, got {rows:?}" + ); + assert!( + rows.iter() + .any(|(k, id, v)| k == "+U" && *id == 1 && *v == 11), + "diff audit_log should include +U after image, got {rows:?}" + ); +} + +#[tokio::test] +async fn incremental_query_rejects_invalid_scan_mode() { + let (_tmp, ctx) = common::setup_sql_context().await; + setup_table_with_three_snapshots(&ctx).await; + + let err = ctx + .sql("SELECT * FROM paimon_incremental_query('paimon.test_db.inc_t', 0, 1, 'nope')") + .await + .expect_err("invalid scan_mode should fail at planning"); + let msg = err.to_string(); + assert!( + msg.to_lowercase().contains("scan_mode") || msg.contains("nope"), + "unexpected error: {msg}" + ); +} + +#[tokio::test] +async fn incremental_query_rejects_illegal_snapshot_range() { + let (_tmp, ctx) = common::setup_sql_context().await; + setup_table_with_three_snapshots(&ctx).await; + + let err = ctx + .sql("SELECT * FROM paimon_incremental_query('paimon.test_db.inc_t', 0, 99)") + .await + .expect("plan may succeed; execution/planning of scan fails") + .collect() + .await + .expect_err("out-of-range end snapshot should fail"); + let msg = err.to_string().to_lowercase(); + assert!( + msg.contains("snapshot") || msg.contains("99") || msg.contains("range"), + "unexpected error: {err}" + ); +} + +#[tokio::test] +async fn incremental_query_supports_projection_and_filter() { + let (_tmp, ctx) = common::setup_sql_context().await; + setup_table_with_three_snapshots(&ctx).await; + + let batches = ctx + .sql("SELECT id FROM paimon_incremental_query('paimon.test_db.inc_t', 0, 2) WHERE id = 2") + .await + .expect("plan") + .collect() + .await + .expect("execute"); + + let mut ids = Vec::new(); + for batch in &batches { + assert_eq!(batch.num_columns(), 1); + assert_eq!(batch.schema().field(0).name(), "id"); + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + ids.push(col.value(row)); + } + } + assert_eq!(ids, vec![2]); +} + +#[tokio::test] +async fn incremental_query_audit_log_sees_delete_rowkind() { + let (tmp, catalog) = common::create_test_env(); + let ctx = common::create_sql_context(Arc::clone(&catalog)).await; + common::exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + + common::exec( + &ctx, + "CREATE TABLE paimon.test_db.inc_del ( + id INT NOT NULL, + PRIMARY KEY (id) + ) WITH ('bucket' = '1', 'changelog-producer' = 'input')", + ) + .await; + common::exec(&ctx, "INSERT INTO paimon.test_db.inc_del VALUES (1)").await; + common::exec(&ctx, "INSERT INTO paimon.test_db.inc_del VALUES (2)").await; + commit_input_changelog_delete(&catalog, "inc_del", 1).await; + let _ = tmp; + + let batches = ctx + .sql( + "SELECT rowkind, id FROM paimon_incremental_query('paimon.test_db.inc_del$audit_log', 1, 3)", + ) + .await + .expect("plan") + .collect() + .await + .expect("execute"); + + let mut kinds = Vec::new(); + for batch in &batches { + let kind_col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let id_col = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + kinds.push((kind_col.value(row).to_string(), id_col.value(row))); + } + } + kinds.sort_unstable(); + assert!( + kinds.iter().any(|(k, id)| k == "-D" && *id == 1), + "audit_log incremental query should expose -D, got {kinds:?}" + ); +} diff --git a/docs/src/sql.md b/docs/src/sql.md index b0da06bf8..9dc379922 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -73,9 +73,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 (`paimon_incremental_query`, `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. ### REST Catalog Views and SQL Functions @@ -1169,6 +1169,72 @@ CROSS JOIN LATERAL vector_search( ) AS r; ``` +## Incremental Query + +Paimon supports batch incremental reads between snapshot IDs via the +`paimon_incremental_query` table-valued function. This aligns with Spark's +`paimon_incremental_query` TVF: results are returned as `RecordBatch` streams +within a single SQL query (not a long-running CDC stream). + +### Registration + +When you use a `SQLContext`, `paimon_incremental_query` is registered +automatically for every catalog you register — no extra setup is needed. + +With a raw DataFusion `SessionContext`, register it explicitly: + +```rust +use paimon_datafusion::register_incremental_query; + +register_incremental_query(&ctx, catalog.clone(), "default"); +``` + +### Usage + +```sql +SELECT * FROM paimon_incremental_query('table_name', start_snapshot_id, end_snapshot_id); +SELECT * FROM paimon_incremental_query('table_name$audit_log', start_snapshot_id, end_snapshot_id); +SELECT * FROM paimon_incremental_query('table_name$audit_log', start_snapshot_id, end_snapshot_id, 'diff'); +``` + +| Argument | Type | Description | +|---|---|---| +| `table_name` | STRING | Table name, fully qualified (`catalog.db.table`) or short form. Append `$audit_log` to read the audit-log view with a `rowkind` column (`+I`, `-U`, `+U`, `-D`). | +| `start_snapshot_id` | BIGINT | Exclusive lower bound snapshot ID | +| `end_snapshot_id` | BIGINT | Inclusive upper bound snapshot ID | +| `scan_mode` | STRING | Optional. One of `auto` (default), `delta`, `changelog`, or `diff`. | + +Examples: + +```sql +-- Delta / auto rows between snapshots 0 and 2 (exclusive start, inclusive end) +SELECT * FROM paimon_incremental_query('paimon.my_db.orders', 0, 2); + +-- Audit log with row kinds +SELECT rowkind, * FROM paimon_incremental_query('paimon.my_db.orders$audit_log', 0, 2); + +-- Diff mode on a primary-key table (before/after images) +SELECT rowkind, * FROM paimon_incremental_query('paimon.my_db.orders$audit_log', 0, 2, 'diff'); +``` + +The snapshot interval is **(start, end]** — rows committed after +`start_snapshot_id` through `end_snapshot_id` inclusive. Omitting `scan_mode` +uses `auto`, which picks an appropriate mode based on table options (for +example `changelog-producer = input` tables often behave like changelog reads). + +Modes: + +| Mode | Behaviour | +|---|---| +| `auto` | Resolve to `delta` when `changelog-producer=none`, otherwise `changelog`. | +| `delta` | Read data files from APPEND snapshots in the range. | +| `changelog` | Read existing changelog manifest files in the range (skips OVERWRITE and snapshots without changelog manifests). Does not generate changelogs. | +| `diff` | Compare before/after snapshot states for primary-key tables (`merge-engine=deduplicate`). | + +Illegal `scan_mode` values and out-of-range snapshot bounds fail with clear +errors. Projection and residual filters are supported via the normal DataFusion +plan (column projection is applied by the incremental scan; filters are residual). + ## Vector Search Paimon supports approximate nearest neighbor (ANN) vector search through global From eb99a8175592ed24df0a57d1118543ed8bb60637 Mon Sep 17 00:00:00 2001 From: pandas886 <123344357+pandas886@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:30:15 +0800 Subject: [PATCH 2/2] fix(datafusion): resolve incremental query catalogs correctly --- .../datafusion/src/incremental_query.rs | 117 ++++++++++++++++-- .../datafusion/src/sql_context.rs | 22 +++- .../datafusion/src/table_function_args.rs | 46 ++++--- .../datafusion/tests/incremental_query.rs | 84 +++++++++++++ 4 files changed, 234 insertions(+), 35 deletions(-) diff --git a/crates/integrations/datafusion/src/incremental_query.rs b/crates/integrations/datafusion/src/incremental_query.rs index d21f46055..f430a41d4 100644 --- a/crates/integrations/datafusion/src/incremental_query.rs +++ b/crates/integrations/datafusion/src/incremental_query.rs @@ -28,8 +28,9 @@ //! execution plan. Snapshot planning lives in paimon-core //! (`IncrementalScan` / `AuditLogTable`). +use std::collections::HashMap; use std::fmt::Debug; -use std::sync::Arc; +use std::sync::{Arc, RwLock}; use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; @@ -55,6 +56,70 @@ use crate::table_function_args::{ const FUNCTION_NAME: &str = "paimon_incremental_query"; +#[derive(Clone)] +struct CatalogBinding { + catalog: Arc, + default_database: String, +} + +#[derive(Clone, Default)] +pub(crate) struct IncrementalQueryCatalogs { + catalogs: Arc>>, + current_catalog: Arc>>, + current_database: Arc>>, +} + +impl IncrementalQueryCatalogs { + pub(crate) fn register(&self, name: String, catalog: Arc, default_database: &str) { + self.catalogs.write().unwrap().insert( + name, + CatalogBinding { + catalog, + default_database: default_database.to_string(), + }, + ); + } + + pub(crate) fn set_current(&self, name: String) { + *self.current_catalog.write().unwrap() = Some(name); + } + + pub(crate) fn set_current_database(&self, name: String) { + *self.current_database.write().unwrap() = Some(name); + } + + fn resolve(&self, name: Option<&str>) -> DFResult { + let name = match name { + Some(name) => name.to_string(), + None => self + .current_catalog + .read() + .unwrap() + .clone() + .ok_or_else(|| { + datafusion::error::DataFusionError::Plan( + "paimon_incremental_query: no current catalog is configured".to_string(), + ) + })?, + }; + let mut binding = self + .catalogs + .read() + .unwrap() + .get(&name) + .cloned() + .ok_or_else(|| { + datafusion::error::DataFusionError::Plan(format!( + "paimon_incremental_query: unknown catalog '{name}'" + )) + })?; + if let Some(current_database) = self.current_database.read().unwrap().as_ref() { + binding.default_database.clone_from(current_database); + } + Ok(binding) + } +} + /// Register the `paimon_incremental_query` table-valued function on a [`SessionContext`]. pub fn register_incremental_query( ctx: &SessionContext, @@ -67,28 +132,56 @@ pub fn register_incremental_query( ); } +pub(crate) fn register_incremental_query_catalogs( + ctx: &SessionContext, + catalogs: IncrementalQueryCatalogs, +) { + ctx.register_udtf( + FUNCTION_NAME, + Arc::new(IncrementalQueryFunction::with_catalogs(catalogs)), + ); +} + /// Table function that performs batch incremental reads between snapshot IDs. /// /// Arguments: /// `(table_name STRING, start_snapshot_exclusive BIGINT, end_snapshot_inclusive BIGINT [, scan_mode STRING])` pub struct IncrementalQueryFunction { - catalog: Arc, - default_database: String, + source: CatalogSource, +} + +#[derive(Clone)] +enum CatalogSource { + Fixed(CatalogBinding), + Registered(IncrementalQueryCatalogs), } impl Debug for IncrementalQueryFunction { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("IncrementalQueryFunction") - .field("default_database", &self.default_database) - .finish() + f.debug_struct("IncrementalQueryFunction").finish() } } impl IncrementalQueryFunction { pub fn new(catalog: Arc, default_database: &str) -> Self { Self { - catalog, - default_database: default_database.to_string(), + source: CatalogSource::Fixed(CatalogBinding { + catalog, + default_database: default_database.to_string(), + }), + } + } + + fn with_catalogs(catalogs: IncrementalQueryCatalogs) -> Self { + Self { + source: CatalogSource::Registered(catalogs), + } + } + + fn resolve_catalog(&self, name: Option<&str>) -> DFResult { + match &self.source { + CatalogSource::Fixed(binding) => Ok(binding.clone()), + CatalogSource::Registered(catalogs) => catalogs.resolve(name), } } } @@ -113,11 +206,11 @@ impl TableFunctionImpl for IncrementalQueryFunction { IncrementalScanMode::Auto }; - let table_ref = - parse_incremental_table_ref(FUNCTION_NAME, &table_name, &self.default_database)?; + let table_ref = parse_incremental_table_ref(FUNCTION_NAME, &table_name)?; + let binding = self.resolve_catalog(table_ref.catalog.as_deref())?; - let catalog = Arc::clone(&self.catalog); - let identifier = table_ref.identifier.clone(); + let catalog = Arc::clone(&binding.catalog); + let identifier = table_ref.identifier(&binding.default_database); let table = block_on_with_runtime( async move { catalog.get_table(&identifier).await }, "paimon_incremental_query: catalog access thread panicked", diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 865744dc9..9e5c4e827 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -96,6 +96,7 @@ use crate::{BlobReaderRegistry, DynamicOptions}; pub struct SQLContext { ctx: SessionContext, catalogs: HashMap>, + incremental_query_catalogs: crate::incremental_query::IncrementalQueryCatalogs, /// Session-scoped dynamic options set via `SET 'paimon.key' = 'value'`. dynamic_options: DynamicOptions, blob_reader_registry: BlobReaderRegistry, @@ -122,11 +123,18 @@ impl SQLContext { )) .build(); let ctx = SessionContext::new_with_state(state); + let incremental_query_catalogs = + crate::incremental_query::IncrementalQueryCatalogs::default(); + crate::incremental_query::register_incremental_query_catalogs( + &ctx, + incremental_query_catalogs.clone(), + ); crate::blob_descriptor_functions::register_blob_descriptor_functions(&ctx); crate::variant_functions::register_variant_functions(&ctx); Self { ctx, catalogs: HashMap::new(), + incremental_query_catalogs, dynamic_options: Default::default(), blob_reader_registry: BlobReaderRegistry::default(), } @@ -203,6 +211,11 @@ impl SQLContext { Some(session_state), )), ); + self.incremental_query_catalogs.register( + catalog_name.clone(), + catalog.clone(), + default_db.unwrap_or("default"), + ); register_table_functions( &self.ctx, &catalog, @@ -241,6 +254,8 @@ impl SQLContext { "SET datafusion.catalog.default_catalog = '{catalog_name}'" )) .await?; + self.incremental_query_catalogs + .set_current(catalog_name.clone()); Ok(()) } @@ -256,6 +271,8 @@ impl SQLContext { "SET datafusion.catalog.default_schema = '{database_name}'" )) .await?; + self.incremental_query_catalogs + .set_current_database(database_name.to_string()); Ok(()) } @@ -3483,11 +3500,6 @@ fn register_table_functions( dynamic_options: DynamicOptions, ) { crate::blob_view::register_blob_view(ctx, Arc::clone(catalog), default_database); - crate::incremental_query::register_incremental_query( - ctx, - Arc::clone(catalog), - default_database, - ); crate::vector_search::register_vector_search_with_dynamic_options( ctx, Arc::clone(catalog), diff --git a/crates/integrations/datafusion/src/table_function_args.rs b/crates/integrations/datafusion/src/table_function_args.rs index d637dfab6..05fced08f 100644 --- a/crates/integrations/datafusion/src/table_function_args.rs +++ b/crates/integrations/datafusion/src/table_function_args.rs @@ -24,10 +24,21 @@ use paimon::table::IncrementalScanMode; /// Parsed target of `paimon_incremental_query`: base table plus optional `$audit_log`. #[derive(Debug, Clone)] pub(crate) struct IncrementalTableRef { - pub identifier: Identifier, + pub catalog: Option, + pub database: Option, + pub object: String, pub audit_log: bool, } +impl IncrementalTableRef { + pub fn identifier(&self, default_database: &str) -> Identifier { + Identifier::new( + self.database.as_deref().unwrap_or(default_database), + &self.object, + ) + } +} + /// Parse table name for incremental query TVF. /// /// Accepts `table`, `database.table`, or `catalog.database.table`, with optional @@ -35,13 +46,12 @@ pub(crate) struct IncrementalTableRef { pub(crate) fn parse_incremental_table_ref( function_name: &str, name: &str, - default_database: &str, ) -> DFResult { let parts: Vec<&str> = name.split('.').collect(); - let (database, object_name) = match parts.len() { - 1 => (default_database, parts[0]), - 2 => (parts[0], parts[1]), - 3 => (parts[1], parts[2]), + let (catalog, database, object_name) = match parts.len() { + 1 => (None, None, parts[0]), + 2 => (None, Some(parts[0]), parts[1]), + 3 => (Some(parts[0]), Some(parts[1]), parts[2]), _ => { return Err(DataFusionError::Plan(format!( "{function_name}: invalid table name '{name}', expected 'table', 'database.table', or 'catalog.database.table'" @@ -63,7 +73,9 @@ pub(crate) fn parse_incremental_table_ref( }; Ok(IncrementalTableRef { - identifier: Identifier::new(database.to_string(), base.to_string()), + catalog: catalog.map(str::to_string), + database: database.map(str::to_string), + object: base.to_string(), audit_log, }) } @@ -189,10 +201,9 @@ mod incremental_query_args_tests { #[test] fn parse_incremental_table_ref_plain_name() { - let parsed = - parse_incremental_table_ref("paimon_incremental_query", "orders", "default").unwrap(); - assert_eq!(parsed.identifier.database(), "default"); - assert_eq!(parsed.identifier.object(), "orders"); + let parsed = parse_incremental_table_ref("paimon_incremental_query", "orders").unwrap(); + assert_eq!(parsed.identifier("default").database(), "default"); + assert_eq!(parsed.identifier("default").object(), "orders"); assert!(!parsed.audit_log); } @@ -201,11 +212,11 @@ mod incremental_query_args_tests { let parsed = parse_incremental_table_ref( "paimon_incremental_query", "paimon.test_db.orders$audit_log", - "default", ) .unwrap(); - assert_eq!(parsed.identifier.database(), "test_db"); - assert_eq!(parsed.identifier.object(), "orders"); + assert_eq!(parsed.catalog.as_deref(), Some("paimon")); + assert_eq!(parsed.identifier("default").database(), "test_db"); + assert_eq!(parsed.identifier("default").object(), "orders"); assert!(parsed.audit_log); } @@ -225,10 +236,9 @@ mod incremental_query_args_tests { #[test] fn parse_incremental_table_ref_rejects_unknown_system_suffix() { - let err = - parse_incremental_table_ref("paimon_incremental_query", "orders$snapshots", "default") - .unwrap_err() - .to_string(); + let err = parse_incremental_table_ref("paimon_incremental_query", "orders$snapshots") + .unwrap_err() + .to_string(); assert!(err.contains("snapshots"), "unexpected error: {err}"); } diff --git a/crates/integrations/datafusion/tests/incremental_query.rs b/crates/integrations/datafusion/tests/incremental_query.rs index f7d22d060..8491394b4 100644 --- a/crates/integrations/datafusion/tests/incremental_query.rs +++ b/crates/integrations/datafusion/tests/incremental_query.rs @@ -145,6 +145,90 @@ async fn incremental_query_without_audit_log_returns_delta_rows() { assert_eq!(collect_id_value(&batches), vec![(1, 10), (2, 20)]); } +#[tokio::test] +async fn incremental_query_resolves_the_qualified_catalog() { + let (_first_tmp, first_catalog) = common::create_test_env(); + let (_second_tmp, second_catalog) = common::create_test_env(); + let mut ctx = paimon_datafusion::SQLContext::new(); + ctx.register_catalog("first", first_catalog).await.unwrap(); + ctx.register_catalog("second", second_catalog) + .await + .unwrap(); + + common::exec( + &ctx, + "CREATE TABLE first.default.inc_t (id INT) WITH ('bucket' = '1', 'bucket-key' = 'id')", + ) + .await; + common::exec( + &ctx, + "CREATE TABLE second.default.inc_t (id INT) WITH ('bucket' = '1', 'bucket-key' = 'id')", + ) + .await; + common::exec(&ctx, "CREATE SCHEMA second.analytics").await; + common::exec( + &ctx, + "CREATE TABLE second.analytics.inc_t (id INT) WITH ('bucket' = '1', 'bucket-key' = 'id')", + ) + .await; + common::exec(&ctx, "INSERT INTO first.default.inc_t VALUES (1)").await; + common::exec(&ctx, "INSERT INTO second.default.inc_t VALUES (2)").await; + common::exec(&ctx, "INSERT INTO second.analytics.inc_t VALUES (3)").await; + + let first = ctx + .sql("SELECT id FROM paimon_incremental_query('first.default.inc_t', 0, 1)") + .await + .unwrap() + .collect() + .await + .unwrap(); + let second = ctx + .sql("SELECT id FROM paimon_incremental_query('second.default.inc_t', 0, 1)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + assert_eq!(collect_id_value_only(&first), vec![1]); + assert_eq!(collect_id_value_only(&second), vec![2]); + + ctx.set_current_catalog("second").await.unwrap(); + let current = ctx + .sql("SELECT id FROM paimon_incremental_query('inc_t', 0, 1)") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(collect_id_value_only(¤t), vec![2]); + + ctx.set_current_database("analytics").await.unwrap(); + let current_database = ctx + .sql("SELECT id FROM paimon_incremental_query('inc_t', 0, 1)") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(collect_id_value_only(¤t_database), vec![3]); +} + +fn collect_id_value_only(batches: &[datafusion::arrow::record_batch::RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .and_then(|column| column.as_any().downcast_ref::()) + .expect("id column") + .values() + .iter() + .copied() + }) + .collect() +} + #[tokio::test] async fn incremental_query_with_audit_log_suffix_exposes_rowkind() { let (_tmp, ctx) = common::setup_sql_context().await;