leaves12138 commented on code in PR #494: URL: https://github.com/apache/paimon-rust/pull/494#discussion_r3558354059
########## crates/integrations/datafusion/src/sql_function.rs: ########## @@ -0,0 +1,453 @@ +// 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; +const MAX_EXPANDED_CALLS: 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( + statement: Statement, + catalogs: &HashMap<String, Arc<dyn Catalog>>, + current_catalog: &str, + current_database: &str, +) -> DFResult<Statement> { + expand_statement_with_budget( + statement, + catalogs, + current_catalog, + current_database, + MAX_EXPANDED_CALLS, + ) + .await +} + +pub(crate) async fn expand_statement_with_budget( + mut statement: Statement, + catalogs: &HashMap<String, Arc<dyn Catalog>>, + current_catalog: &str, + current_database: &str, + max_expanded_calls: usize, +) -> DFResult<Statement> { + let mut functions = HashMap::new(); + let mut dependencies = HashMap::new(); + let mut total_expanded_calls = 0; + 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 Some(catalog) = catalogs.get(&reference.0) else { + functions.insert(reference, None); + continue; + }; + 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(()); + }; + if total_expanded_calls >= max_expanded_calls { + return ControlFlow::Break(DataFusionError::Plan(format!( + "REST SQL function expansion budget of {max_expanded_calls} calls exceeded" + ))); + } + total_expanded_calls += 1; + + match expand_call(function, call, &reference.0, &functions) { + Ok(expanded) => { + *expr = expanded; + expanded_count += 1; + ControlFlow::Continue(()) + } + Err(error) => ControlFlow::Break(error), + } + }); + if let ControlFlow::Break(error) = flow { + return Err(error); + } + if expanded_count == 0 { + return Ok(statement); + } + } + + Err(DataFusionError::Plan(format!( + "REST SQL function expansion exceeded maximum depth of {MAX_EXPANSION_DEPTH}" + ))) +} + +fn function_dependencies( + function: &Function, + current_catalog: &str, + current_database: &str, +) -> Vec<FunctionReference> { + let Some(definition) = function + .definition("datafusion") + .and_then(paimon::catalog::FunctionDefinition::sql) + else { + return Vec::new(); + }; + let Ok(mut parser) = Parser::new(&GenericDialect {}).try_with_sql(definition) else { + return Vec::new(); + }; + let Ok(body) = parser.parse_expr() else { + return Vec::new(); + }; + let mut dependencies = Vec::new(); + let _: ControlFlow<()> = visit_expressions(&body, |expr| { + if let SqlExpr::Function(call) = expr { + if let Some(reference) = function_reference(call, current_catalog, current_database) { + dependencies.push(reference); + } + } + ControlFlow::Continue(()) + }); + dependencies +} + +fn find_dependency_cycle( + dependencies: &HashMap<FunctionReference, Vec<FunctionReference>>, +) -> Option<Vec<FunctionReference>> { + fn visit( + reference: &FunctionReference, + dependencies: &HashMap<FunctionReference, Vec<FunctionReference>>, + finished: &mut HashSet<FunctionReference>, + path: &mut Vec<FunctionReference>, + ) -> Option<Vec<FunctionReference>> { + if let Some(start) = path.iter().position(|entry| entry == reference) { + let mut cycle = path[start..].to_vec(); + cycle.push(reference.clone()); + return Some(cycle); + } + if finished.contains(reference) { + return None; + } + + path.push(reference.clone()); + if let Some(next_references) = dependencies.get(reference) { + for next in next_references { + if let Some(cycle) = visit(next, dependencies, finished, path) { + return Some(cycle); + } + } + } + path.pop(); + finished.insert(reference.clone()); + None + } + + let mut finished = HashSet::new(); + for reference in dependencies.keys() { + if let Some(cycle) = visit(reference, dependencies, &mut finished, &mut Vec::new()) { + return Some(cycle); + } + } + None +} + +fn function_reference( + function: &datafusion::sql::sqlparser::ast::Function, + current_catalog: &str, + current_database: &str, +) -> Option<(String, Identifier)> { + let identifiers = function + .name + .0 + .iter() + .map(|part| part.as_ident().map(|ident| ident.value.clone())) Review Comment: Please apply DataFusion identifier normalization before REST function lookup and parameter substitution. A catalog function stored as `plus_one` is not expanded for `SELECT PLUS_ONE(41)` because the raw `Ident.value` is used; DataFusion lowercases unquoted identifiers by default, so the query should resolve the same function. I reproduced the current behavior as `Invalid function 'plus_one'`. Catalog/database components and unquoted parameter identifiers inside stored definitions need the same treatment. ########## crates/integrations/datafusion/src/catalog.rs: ########## @@ -501,3 +589,283 @@ 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 visitor = ViewRelationVisitor::new(catalog_name, view.identifier().database()); + let _: std::ops::ControlFlow<()> = statements.visit(&mut visitor); + Ok(visitor.identifiers) +} + +type SqlIdentifierKey = (Option<char>, String); + +struct QueryCteScope { + visible: HashSet<SqlIdentifierKey>, + cte_query_visibility: HashMap<usize, HashSet<SqlIdentifierKey>>, +} + +struct ViewRelationVisitor<'a> { + catalog_name: &'a str, + current_database: &'a str, + scopes: Vec<QueryCteScope>, + identifiers: Vec<Identifier>, +} + +impl<'a> ViewRelationVisitor<'a> { + fn new(catalog_name: &'a str, current_database: &'a str) -> Self { + Self { + catalog_name, + current_database, + scopes: Vec::new(), + identifiers: Vec::new(), + } + } +} + +impl Visitor for ViewRelationVisitor<'_> { + type Break = (); + + fn pre_visit_query(&mut self, query: &Query) -> std::ops::ControlFlow<Self::Break> { + let query_address = query as *const Query as usize; + let inherited = self + .scopes + .last() + .map(|scope| { + scope + .cte_query_visibility + .get(&query_address) + .unwrap_or(&scope.visible) + .clone() + }) + .unwrap_or_default(); + let mut visible = inherited.clone(); + let mut cte_query_visibility = HashMap::new(); + + if let Some(with) = &query.with { + let local_ctes = with + .cte_tables + .iter() + .map(|cte| sql_identifier_key(&cte.alias.name)) + .collect::<Vec<_>>(); + if with.recursive { + visible.extend(local_ctes); + for cte in &with.cte_tables { + cte_query_visibility + .insert(cte.query.as_ref() as *const Query as usize, visible.clone()); + } + } else { + for (cte, alias) in with.cte_tables.iter().zip(local_ctes) { + cte_query_visibility + .insert(cte.query.as_ref() as *const Query as usize, visible.clone()); + visible.insert(alias); + } + } + } + + self.scopes.push(QueryCteScope { + visible, + cte_query_visibility, + }); + std::ops::ControlFlow::Continue(()) + } + + fn post_visit_query(&mut self, _query: &Query) -> std::ops::ControlFlow<Self::Break> { + self.scopes.pop(); + std::ops::ControlFlow::Continue(()) + } + + fn pre_visit_relation(&mut self, relation: &ObjectName) -> std::ops::ControlFlow<Self::Break> { + let is_cte = match relation.0.as_slice() { + [part] => part.as_ident().is_some_and(|identifier| { + self.scopes + .last() + .is_some_and(|scope| scope.visible.contains(&sql_identifier_key(identifier))) + }), + _ => false, + }; + if !is_cte { + if let Some(identifier) = + relation_identifier(relation, self.catalog_name, self.current_database) + { + self.identifiers.push(identifier); + } + } + std::ops::ControlFlow::Continue(()) + } +} + +fn sql_identifier_key(identifier: &Ident) -> SqlIdentifierKey { Review Comment: Please match DataFusion's identifier normalization rather than treating `quote_style` as part of the CTE identity. With default normalization, `WITH cte_view AS (...) SELECT * FROM "cte_view"` references the same CTE, but these keys differ (`None` versus `Some('\"')`). If the persistent view itself is named `cte_view`, dependency validation reports a false self-cycle. I reproduced this locally: `recursive REST catalog view dependency detected: default.cte_view -> default.cte_view`. The same normalization should be applied when relation parts are converted to catalog/database/object names. ########## crates/integrations/datafusion/src/sql_context.rs: ########## @@ -475,6 +480,24 @@ impl SQLContext { ) .await } + Statement::Query(_) => { Review Comment: Please ensure REST function expansion is used by the other query paths as well. In particular, `sql()` returns through `handle_time_travel_query` before reaching this arm, and that handler calls the raw `SessionContext`, so a query such as `SELECT plus_one(id) FROM t VERSION AS OF 1` bypasses expansion and fails function resolution. `EXPLAIN SELECT plus_one(1)` also does not enter the `Statement::Query` arm. The PR states that queries executed through `SQLContext` can call these functions, so these existing query forms should not silently lose the feature (or the limitation should be explicit and intentional). ########## crates/integrations/datafusion/src/catalog.rs: ########## @@ -65,28 +74,20 @@ impl Debug for PaimonCatalogProvider { impl PaimonCatalogProvider { /// Creates a new [`PaimonCatalogProvider`]. /// - /// For standalone use without `SET`/`RESET` support. - /// When used via [`SQLContext`], the handler creates the provider - /// internally with shared dynamic options. - pub fn new(catalog: Arc<dyn Catalog>) -> Self { - PaimonCatalogProvider { - catalog, - dynamic_options: Default::default(), - temp_tables: Arc::new(RwLock::new(HashMap::new())), - blob_reader_registry: BlobReaderRegistry::default(), - } - } - - pub(crate) fn with_dynamic_options( + pub fn new( Review Comment: Please preserve the existing public `PaimonCatalogProvider::new(Arc<dyn Catalog>)` API. This changes it to a five-argument constructor, so every downstream crate using the provider directly stops compiling; I confirmed the old call now fails with E0061. The extra session state is only needed for the `SQLContext` integration, so keeping `new(catalog)` and adding a crate-private/session-aware constructor would avoid an unrelated source-breaking change. `PaimonSchemaProvider::new` is changed similarly. -- 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]
