jackylee-ch commented on code in PR #1042: URL: https://github.com/apache/paimon-rust/pull/1042#discussion_r4343697620
########## crates/integrations/datafusion/src/incremental_query.rs: ########## @@ -0,0 +1,435 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The `paimon_incremental_query` table function reads a snapshot-id range +//! `(start_exclusive, end_inclusive]` as a batch. Snapshot planning stays in +//! paimon core; this parses arguments, picks the core planner that matches +//! Java's `incremental-between` for the table, and streams its output. +//! +//! ```sql +//! SELECT * FROM paimon_incremental_query('db.t', 0, 5); +//! SELECT * FROM paimon_incremental_query('db.t', 0, 5, 'delta'); +//! SELECT * FROM paimon_incremental_query('db.t$audit_log', 0, 5, 'changelog'); +//! ``` +//! +//! A `$audit_log` table-name suffix prepends the `rowkind` audit column (and the +//! optional `_SEQUENCE_NUMBER`), mirroring the core audit-log read. + +use std::fmt; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; +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, StreamExt, TryStreamExt}; +use paimon::catalog::Catalog; +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::{ + 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<dyn Catalog>, + default_database: &str, +) { + ctx.register_udtf( + FUNCTION_NAME, + Arc::new(IncrementalQueryFunction { + catalog, + default_database: default_database.to_string(), + }), + ); +} + +pub(crate) struct IncrementalQueryFunction { + catalog: Arc<dyn Catalog>, + 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<IncrementalScanMode> { + 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<Arc<dyn TableProvider>> { + 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); Review Comment: Fixed in f8016e1: the function now resolves the name in the calling session as a `FROM` clause does (`t` and `db.t` take the current catalog and database, `cat.db.t` names the catalog) and looks the catalog up in the session's catalog list; an unknown or non-Paimon catalog is an error. Added the two-warehouse regression, checked against ordinary SELECT for 1-, 2- and 3-part names, after `USE`, and with a database named like the other catalog. `vector_search`, `full_text_search`, `hybrid_search` and `blob_view` still bind the catalog registered last; this PR leaves them unchanged. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
