leaves12138 commented on code in PR #494: URL: https://github.com/apache/paimon-rust/pull/494#discussion_r3557533997
########## crates/integrations/datafusion/src/sql_function.rs: ########## @@ -0,0 +1,409 @@ +// 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. + +use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; +use std::ops::ControlFlow; +use std::sync::Arc; + +use datafusion::error::{DataFusionError, Result as DFResult}; +use datafusion::sql::sqlparser::ast::{ + visit_expressions, visit_expressions_mut, Expr as SqlExpr, FunctionArg, FunctionArgExpr, + FunctionArguments, Ident, ObjectName, Statement, +}; +use datafusion::sql::sqlparser::dialect::GenericDialect; +use datafusion::sql::sqlparser::parser::Parser; +use paimon::catalog::{Catalog, Function, Identifier}; + +const MAX_EXPANSION_DEPTH: usize = 32; +const MAX_FUNCTION_REFERENCES: usize = 1024; +type FunctionReference = (String, Identifier); + +pub(crate) async fn expand_sql( + sql: &str, + catalogs: &HashMap<String, Arc<dyn Catalog>>, + current_catalog: &str, + current_database: &str, +) -> DFResult<String> { + let mut statements = Parser::parse_sql(&GenericDialect {}, sql) + .map_err(|error| DataFusionError::Plan(format!("Invalid REST SQL: {error}")))?; + if statements.len() != 1 { + return Err(DataFusionError::Plan(format!( + "REST SQL must contain exactly one statement, found {}", + statements.len() + ))); + } + let statement = expand_statement( + statements.remove(0), + catalogs, + current_catalog, + current_database, + ) + .await?; + Ok(statement.to_string()) +} + +pub(crate) async fn expand_statement( + mut statement: Statement, + catalogs: &HashMap<String, Arc<dyn Catalog>>, + current_catalog: &str, + current_database: &str, +) -> DFResult<Statement> { + let mut functions = HashMap::new(); + let mut dependencies = HashMap::new(); + for _ in 0..MAX_EXPANSION_DEPTH { + let mut references = BTreeMap::new(); + let _: ControlFlow<()> = visit_expressions(&statement, |expr| { + if let SqlExpr::Function(function) = expr { + if let Some(reference) = + function_reference(function, current_catalog, current_database) + { + references.insert(reference, ()); + } + } + ControlFlow::Continue(()) + }); + + let mut pending = references.into_keys().collect::<VecDeque<_>>(); + while let Some(reference) = pending.pop_front() { + if functions.contains_key(&reference) { + continue; + } + if functions.len() >= MAX_FUNCTION_REFERENCES { + return Err(DataFusionError::Plan(format!( + "REST SQL function expansion exceeded the reference limit of {MAX_FUNCTION_REFERENCES}" + ))); + } + let catalog = catalogs.get(&reference.0).ok_or_else(|| { Review Comment: This lookup regresses `SQLContext` usages that have no registered Paimon catalog. Every function call is treated as a REST-function candidate, `current_catalog` is then DataFusion's default `datafusion`, and this branch returns `Unknown catalog 'datafusion'` before registered/built-in UDF resolution. The current Python integration CI fails `video_snapshot` and `vector_from_json` tests for exactly this reason. Please make expansion a no-op/fallback when the current catalog is not in the Paimon catalog map (while retaining an appropriate error for explicitly invalid qualification), and add a no-catalog built-in regression test. ########## crates/integrations/datafusion/src/sql_function.rs: ########## @@ -0,0 +1,409 @@ +// 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. + +use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; +use std::ops::ControlFlow; +use std::sync::Arc; + +use datafusion::error::{DataFusionError, Result as DFResult}; +use datafusion::sql::sqlparser::ast::{ + visit_expressions, visit_expressions_mut, Expr as SqlExpr, FunctionArg, FunctionArgExpr, + FunctionArguments, Ident, ObjectName, Statement, +}; +use datafusion::sql::sqlparser::dialect::GenericDialect; +use datafusion::sql::sqlparser::parser::Parser; +use paimon::catalog::{Catalog, Function, Identifier}; + +const MAX_EXPANSION_DEPTH: usize = 32; +const MAX_FUNCTION_REFERENCES: usize = 1024; +type FunctionReference = (String, Identifier); + +pub(crate) async fn expand_sql( + sql: &str, + catalogs: &HashMap<String, Arc<dyn Catalog>>, + current_catalog: &str, + current_database: &str, +) -> DFResult<String> { + let mut statements = Parser::parse_sql(&GenericDialect {}, sql) + .map_err(|error| DataFusionError::Plan(format!("Invalid REST SQL: {error}")))?; + if statements.len() != 1 { + return Err(DataFusionError::Plan(format!( + "REST SQL must contain exactly one statement, found {}", + statements.len() + ))); + } + let statement = expand_statement( + statements.remove(0), + catalogs, + current_catalog, + current_database, + ) + .await?; + Ok(statement.to_string()) +} + +pub(crate) async fn expand_statement( + mut statement: Statement, + catalogs: &HashMap<String, Arc<dyn Catalog>>, + current_catalog: &str, + current_database: &str, +) -> DFResult<Statement> { + let mut functions = HashMap::new(); + let mut dependencies = HashMap::new(); + for _ in 0..MAX_EXPANSION_DEPTH { + let mut references = BTreeMap::new(); + let _: ControlFlow<()> = visit_expressions(&statement, |expr| { + if let SqlExpr::Function(function) = expr { + if let Some(reference) = + function_reference(function, current_catalog, current_database) + { + references.insert(reference, ()); + } + } + ControlFlow::Continue(()) + }); + + let mut pending = references.into_keys().collect::<VecDeque<_>>(); + while let Some(reference) = pending.pop_front() { + if functions.contains_key(&reference) { + continue; + } + if functions.len() >= MAX_FUNCTION_REFERENCES { + return Err(DataFusionError::Plan(format!( + "REST SQL function expansion exceeded the reference limit of {MAX_FUNCTION_REFERENCES}" + ))); + } + let catalog = catalogs.get(&reference.0).ok_or_else(|| { + DataFusionError::Plan(format!("Unknown catalog '{}'", reference.0)) + })?; + let function = match catalog.get_function(&reference.1).await { + Ok(function) => Some(function), + Err(paimon::Error::FunctionNotExist { .. }) + | Err(paimon::Error::Unsupported { .. }) => None, + Err(error) => return Err(crate::to_datafusion_error(error)), Review Comment: Please preserve compatibility with REST servers that do not implement the new function endpoint. This code probes `get_function` for every scalar/aggregate call, including built-ins such as `abs` and `count`; `RESTCatalog::get_function` currently leaves HTTP 501 as `Error::RestApi(NotImplemented)`, so `SELECT abs(-1)` fails instead of falling through to DataFusion. I reproduced this with a catalog returning `RestError::NotImplemented`. Please map 501 to `Unsupported` at the REST catalog boundary or handle it as a non-persistent function here, and add a regression test. The same compatibility consideration applies to the new view endpoint. ########## crates/integrations/datafusion/src/sql_function.rs: ########## @@ -0,0 +1,409 @@ +// 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. + +use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; +use std::ops::ControlFlow; +use std::sync::Arc; + +use datafusion::error::{DataFusionError, Result as DFResult}; +use datafusion::sql::sqlparser::ast::{ + visit_expressions, visit_expressions_mut, Expr as SqlExpr, FunctionArg, FunctionArgExpr, + FunctionArguments, Ident, ObjectName, Statement, +}; +use datafusion::sql::sqlparser::dialect::GenericDialect; +use datafusion::sql::sqlparser::parser::Parser; +use paimon::catalog::{Catalog, Function, Identifier}; + +const MAX_EXPANSION_DEPTH: usize = 32; +const MAX_FUNCTION_REFERENCES: usize = 1024; +type FunctionReference = (String, Identifier); + +pub(crate) async fn expand_sql( + sql: &str, + catalogs: &HashMap<String, Arc<dyn Catalog>>, + current_catalog: &str, + current_database: &str, +) -> DFResult<String> { + let mut statements = Parser::parse_sql(&GenericDialect {}, sql) + .map_err(|error| DataFusionError::Plan(format!("Invalid REST SQL: {error}")))?; + if statements.len() != 1 { + return Err(DataFusionError::Plan(format!( + "REST SQL must contain exactly one statement, found {}", + statements.len() + ))); + } + let statement = expand_statement( + statements.remove(0), + catalogs, + current_catalog, + current_database, + ) + .await?; + Ok(statement.to_string()) +} + +pub(crate) async fn expand_statement( + mut statement: Statement, + catalogs: &HashMap<String, Arc<dyn Catalog>>, + current_catalog: &str, + current_database: &str, +) -> DFResult<Statement> { + let mut functions = HashMap::new(); + let mut dependencies = HashMap::new(); + for _ in 0..MAX_EXPANSION_DEPTH { + let mut references = BTreeMap::new(); + let _: ControlFlow<()> = visit_expressions(&statement, |expr| { + if let SqlExpr::Function(function) = expr { + if let Some(reference) = + function_reference(function, current_catalog, current_database) + { + references.insert(reference, ()); + } + } + ControlFlow::Continue(()) + }); + + let mut pending = references.into_keys().collect::<VecDeque<_>>(); + while let Some(reference) = pending.pop_front() { + if functions.contains_key(&reference) { + continue; + } + if functions.len() >= MAX_FUNCTION_REFERENCES { + return Err(DataFusionError::Plan(format!( + "REST SQL function expansion exceeded the reference limit of {MAX_FUNCTION_REFERENCES}" + ))); + } + let catalog = catalogs.get(&reference.0).ok_or_else(|| { + DataFusionError::Plan(format!("Unknown catalog '{}'", reference.0)) + })?; + let function = match catalog.get_function(&reference.1).await { + Ok(function) => Some(function), + Err(paimon::Error::FunctionNotExist { .. }) + | Err(paimon::Error::Unsupported { .. }) => None, + Err(error) => return Err(crate::to_datafusion_error(error)), + }; + if let Some(function) = &function { + let nested = function_dependencies(function, &reference.0, reference.1.database()); + pending.extend(nested.iter().cloned()); + dependencies.insert(reference.clone(), nested); + } + functions.insert(reference, function); + } + + if let Some(cycle) = find_dependency_cycle(&dependencies) { + let path = cycle + .iter() + .map(|(catalog, identifier)| format!("{catalog}.{}", identifier.full_name())) + .collect::<Vec<_>>() + .join(" -> "); + return Err(DataFusionError::Plan(format!( + "recursive REST SQL function dependency detected: {path}" + ))); + } + + let mut expanded_count = 0; + let flow = visit_expressions_mut(&mut statement, |expr| { + let SqlExpr::Function(call) = expr else { + return ControlFlow::Continue(()); + }; + let Some(reference) = function_reference(call, current_catalog, current_database) + else { + return ControlFlow::Continue(()); + }; + let Some(Some(function)) = functions.get(&reference) else { + return ControlFlow::Continue(()); + }; + + match expand_call(function, call, &reference.0, &functions) { + Ok(expanded) => { + *expr = expanded; + expanded_count += 1; Review Comment: The current limits bound nesting depth and the number of distinct function definitions, but not the number of expanded call sites. A chain such as `f0(x) = f1(x) + f1(x)`, ..., `f31(x) = x` creates roughly `2^31` subexpressions while using only 32 references and 32 rounds, so the process can exhaust memory before the depth guard fires. Please add a total expanded-call/AST-size budget checked before replacement, with a branching-chain regression test using a small configured limit. ########## crates/integrations/datafusion/src/catalog.rs: ########## @@ -501,3 +589,182 @@ impl SchemaProvider for PaimonSchemaProvider { ) } } + +fn enforce_view_schema(plan: LogicalPlan, view: &View) -> DFResult<LogicalPlan> { + let declared_fields = view.schema().fields(); + let actual_fields = plan.schema().fields(); + if actual_fields.len() != declared_fields.len() { + return Err(plan_datafusion_err!( + "REST catalog view '{}' declares {} fields but its query returns {}", + view.full_name(), + declared_fields.len(), + actual_fields.len() + )); + } + + let expressions = declared_fields + .iter() + .enumerate() + .map(|(index, declared)| { + let (qualifier, actual) = plan.schema().qualified_field(index); + let column = match qualifier { + Some(qualifier) => Column::new(Some(qualifier.clone()), actual.name()), + None => Column::new_unqualified(actual.name()), + }; + let target_type = paimon::arrow::paimon_type_to_arrow(declared.data_type()) + .map_err(to_datafusion_error)?; + Ok(cast(Expr::Column(column), target_type).alias(declared.name())) + }) + .collect::<DFResult<Vec<_>>>()?; + + LogicalPlanBuilder::from(plan).project(expressions)?.build() +} + +const MAX_VIEW_DEPENDENCIES: usize = 64; + +async fn validate_view_dependencies( + catalog: &Arc<dyn Catalog>, + catalog_name: &str, + root: &View, +) -> DFResult<()> { + let mut queue = VecDeque::from([root.clone()]); + let mut loaded = HashSet::from([root.identifier().clone()]); + let mut dependencies = HashMap::<Identifier, Vec<Identifier>>::new(); + + while let Some(view) = queue.pop_front() { + let candidates = view_relation_identifiers(&view, catalog_name)?; + let mut view_dependencies = Vec::new(); + for identifier in candidates { + match catalog.get_table(&identifier).await { + Ok(_) => continue, + Err(paimon::Error::TableNotExist { .. }) + | Err(paimon::Error::Unsupported { .. }) => {} + Err(error) => return Err(to_datafusion_error(error)), + } + + let dependency = match catalog.get_view(&identifier).await { + Ok(view) => view, + Err(paimon::Error::ViewNotExist { .. }) + | Err(paimon::Error::Unsupported { .. }) => continue, + Err(error) => return Err(to_datafusion_error(error)), + }; + view_dependencies.push(identifier.clone()); + if loaded.insert(identifier) { + if loaded.len() > MAX_VIEW_DEPENDENCIES { + return Err(plan_datafusion_err!( + "REST catalog view '{}' exceeds the dependency limit of {}", + root.full_name(), + MAX_VIEW_DEPENDENCIES + )); + } + queue.push_back(dependency); + } + } + dependencies.insert(view.identifier().clone(), view_dependencies); + + if let Some(cycle) = find_view_dependency_cycle(&dependencies) { + let path = cycle + .iter() + .map(Identifier::full_name) + .collect::<Vec<_>>() + .join(" -> "); + return Err(plan_datafusion_err!( + "recursive REST catalog view dependency detected: {path}" + )); + } + } + Ok(()) +} + +fn view_relation_identifiers(view: &View, catalog_name: &str) -> DFResult<Vec<Identifier>> { + let statements = + Parser::parse_sql(&GenericDialect {}, view.query_for("datafusion")).map_err(|error| { + plan_datafusion_err!( + "Invalid SQL for REST catalog view '{}': {error}", + view.full_name() + ) + })?; + if statements.len() != 1 { + return Err(plan_datafusion_err!( + "REST catalog view '{}' must contain exactly one SQL statement", + view.full_name() + )); + } + if !matches!(statements.first(), Some(Statement::Query(_))) { + return Err(plan_datafusion_err!( + "REST catalog view '{}' must contain a single read-only query", + view.full_name() + )); + } + + let mut identifiers = Vec::new(); + let _: std::ops::ControlFlow<()> = visit_relations(&statements, |relation| { Review Comment: `visit_relations` also visits references to local CTE names, but this code treats every unqualified relation as a catalog object. For example, a persistent view named `cte_view` with `WITH cte_view AS (SELECT CAST(42 AS BIGINT) AS answer) SELECT * FROM cte_view` is rejected as `default.cte_view -> default.cte_view` recursion; I reproduced this locally. Please exclude CTE aliases according to their SQL scope (including nested CTEs) before catalog lookup, and add a regression test. -- 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]
