leaves12138 commented on code in PR #460: URL: https://github.com/apache/paimon-rust/pull/460#discussion_r3527356005
########## crates/integrations/datafusion/src/variant_pushdown.rs: ########## @@ -0,0 +1,608 @@ +// 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::any::Any; +use std::cmp::Ordering; +use std::collections::{HashMap, HashSet}; +use std::fmt; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; +use datafusion::catalog::default_table_source::source_as_provider; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::{ + internal_err, plan_err, Column, DFSchema, DFSchemaRef, Result as DFResult, ScalarValue, +}; +use datafusion::execution::context::SessionState; +use datafusion::functions::core::expr_fn::get_field; +use datafusion::logical_expr::expr::{InList, ScalarFunction}; +use datafusion::logical_expr::{ + Between, BinaryExpr, Case, Cast, Expr, Extension, Filter, Like, LogicalPlan, Projection, + TableScan, TryCast, UserDefinedLogicalNode, +}; +use datafusion::optimizer::{ApplyOrder, OptimizerConfig, OptimizerRule}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner}; +use paimon::spec::{ + variant_extraction_row, BigIntType, BooleanType, DataField, DataType, DateType, DecimalType, + DoubleType, FloatType, IntType, SmallIntType, TinyIntType, VarCharType, +}; +use paimon::table::Table; + +use crate::error::to_datafusion_error; +use crate::filter_pushdown::analyze_filters; +use crate::physical_plan::PaimonTableScan; +use crate::runtime::await_with_runtime; +use crate::table::{bucket_round_robin, datafusion_read_fields, PaimonTableProvider}; + +#[derive(Debug)] +pub(crate) struct RewriteVariantExtractions; + +impl RewriteVariantExtractions { + pub(crate) fn new() -> Self { + Self + } +} + +impl OptimizerRule for RewriteVariantExtractions { + fn name(&self) -> &str { + "rewrite_variant_extractions" + } + + fn apply_order(&self) -> Option<ApplyOrder> { + Some(ApplyOrder::BottomUp) + } + + fn rewrite( + &self, + plan: LogicalPlan, + _config: &dyn OptimizerConfig, + ) -> DFResult<Transformed<LogicalPlan>> { + let LogicalPlan::Projection(projection) = plan else { + return Ok(Transformed::no(plan)); + }; + let rewrite = match projection.input.as_ref() { + LogicalPlan::TableScan(scan) => build_projection_rewrite(&projection, scan, None)?, + LogicalPlan::Filter(filter) => match filter.input.as_ref() { + LogicalPlan::TableScan(scan) => { + build_projection_rewrite(&projection, scan, Some(filter))? + } + _ => None, + }, + _ => None, + }; + let Some(rewrite) = rewrite else { + return Ok(Transformed::no(LogicalPlan::Projection(projection))); + }; + Ok(Transformed::yes(LogicalPlan::Projection( + Projection::try_new_with_schema( + rewrite.exprs, + Arc::new(rewrite.input), + projection.schema, + )?, + ))) + } +} + +#[derive(Debug)] +struct ProjectionRewrite { + input: LogicalPlan, + exprs: Vec<Expr>, +} + +#[derive(Debug, Clone)] +struct VariantExtractionScanNode { + table: Table, + schema: DFSchemaRef, + arrow_schema: ArrowSchemaRef, + read_type: Vec<DataField>, + filters: Vec<Expr>, + fetch: Option<usize>, + pushed_variants: String, +} + +impl UserDefinedLogicalNode for VariantExtractionScanNode { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "PaimonVariantExtractionScan" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn check_invariants(&self, _check: datafusion::logical_expr::InvariantLevel) -> DFResult<()> { + Ok(()) + } + + fn expressions(&self) -> Vec<Expr> { + vec![] + } + + fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "PaimonVariantExtractionScan: PushedVariants=[{}]", + self.pushed_variants + ) + } + + fn with_exprs_and_inputs( + &self, + exprs: Vec<Expr>, + inputs: Vec<LogicalPlan>, + ) -> DFResult<Arc<dyn UserDefinedLogicalNode>> { + if !exprs.is_empty() || !inputs.is_empty() { + return internal_err!("PaimonVariantExtractionScan expects no expressions or inputs"); + } + Ok(Arc::new(self.clone())) + } + + fn dyn_hash(&self, mut state: &mut dyn Hasher) { + self.name().hash(&mut state); + self.table.location().hash(&mut state); + self.read_type.hash(&mut state); + self.filters.hash(&mut state); + self.fetch.hash(&mut state); + } + + fn dyn_eq(&self, other: &dyn UserDefinedLogicalNode) -> bool { + other.as_any().downcast_ref::<Self>().is_some_and(|other| { + self.table.location() == other.table.location() + && self.read_type == other.read_type + && self.filters == other.filters + && self.fetch == other.fetch + }) + } + + fn dyn_ord(&self, other: &dyn UserDefinedLogicalNode) -> Option<Ordering> { + let other = other.as_any().downcast_ref::<Self>()?; + if self.dyn_eq(other) { + Some(Ordering::Equal) + } else { + Some(format!("{self:?}").cmp(&format!("{other:?}"))) + } + } +} + +#[derive(Debug)] +pub(crate) struct VariantExtractionExtensionPlanner; + +#[async_trait] +impl ExtensionPlanner for VariantExtractionExtensionPlanner { + async fn plan_extension( + &self, + _planner: &dyn PhysicalPlanner, + node: &dyn UserDefinedLogicalNode, + logical_inputs: &[&LogicalPlan], + physical_inputs: &[Arc<dyn ExecutionPlan>], + session_state: &SessionState, + ) -> DFResult<Option<Arc<dyn ExecutionPlan>>> { + let Some(node) = node.as_any().downcast_ref::<VariantExtractionScanNode>() else { + return Ok(None); + }; + if !logical_inputs.is_empty() || !physical_inputs.is_empty() { + return internal_err!("PaimonVariantExtractionScan physical planning expects no input"); + } + + let filter_analysis = analyze_filters(&node.filters, node.table.schema().fields()); + let mut read_builder = node.table.new_read_builder(); + read_builder.with_read_type(node.read_type.clone()); + if let Some(filter) = filter_analysis.pushed_predicate.clone() { + read_builder.with_filter(filter); + } + let pushed_limit = node + .fetch + .filter(|_| !filter_analysis.has_untranslated_residual); + if let Some(limit) = pushed_limit { + read_builder.with_limit(limit); + } + let scan = read_builder.new_scan(); + let (plan, scan_trace) = await_with_runtime(scan.plan_with_trace()) + .await + .map_err(to_datafusion_error)?; + + let splits = plan.splits().to_vec(); + let target = session_state.config_options().execution.target_partitions; + let planned_partitions: Vec<Arc<[_]>> = if splits.is_empty() { + vec![Arc::from(Vec::new())] + } else { + let num_partitions = splits.len().min(target.max(1)); + bucket_round_robin(splits, num_partitions) + .into_iter() + .map(Arc::from) + .collect() + }; + let filter_exact = !filter_analysis.has_untranslated_residual + && filter_analysis + .pushed_predicate + .as_ref() + .is_none_or(|p| read_builder.is_exact_filter_pushdown(p)); + + Ok(Some(Arc::new(PaimonTableScan::new( + Arc::clone(&node.arrow_schema), + node.table.clone(), + node.read_type.clone(), + filter_analysis.pushed_predicate, + planned_partitions, + pushed_limit, + filter_exact, + Some(scan_trace), + Some(node.pushed_variants.clone()), + )))) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct VariantGetCall { + column: Column, + path: String, + type_name: String, + data_type: DataType, + fail_on_error: bool, +} + +#[derive(Debug, Clone)] +struct AcceptedExtraction { + call: VariantGetCall, + field_name: String, +} + +fn build_projection_rewrite( + projection: &Projection, + scan: &TableScan, + filter: Option<&Filter>, +) -> DFResult<Option<ProjectionRewrite>> { + let provider = source_as_provider(&scan.source)?; + let Some(provider) = provider.downcast_ref::<PaimonTableProvider>() else { + return Ok(None); + }; + + let mut calls = Vec::new(); + let mut full_columns = HashSet::new(); + for expr in &projection.expr { + collect_variant_usage(expr, &mut calls, &mut full_columns)?; + } + for expr in &scan.filters { + collect_variant_usage(expr, &mut calls, &mut full_columns)?; + } + if let Some(filter) = filter { + collect_variant_usage(&filter.predicate, &mut calls, &mut full_columns)?; + } + if calls.is_empty() { + return Ok(None); + } + + let read_fields = datafusion_read_fields(provider.table()); + let mut by_column: HashMap<String, Vec<VariantGetCall>> = HashMap::new(); + for call in calls { + if full_columns.contains(&call.column.name) { + continue; + } + let Some(table_field) = read_fields + .iter() + .find(|field| field.name() == call.column.name) + else { + continue; + }; + if !matches!(table_field.data_type(), DataType::Variant(_)) { + continue; + } + let entries = by_column.entry(call.column.name.clone()).or_default(); + if !entries.iter().any(|existing| existing == &call) { + entries.push(call); + } + } + if by_column.is_empty() { + return Ok(None); + } + + let read_indices = scan + .projection + .clone() + .unwrap_or_else(|| (0..read_fields.len()).collect()); + let mut read_type = Vec::with_capacity(read_indices.len()); + for idx in read_indices { + let Some(field) = read_fields.get(idx).cloned() else { + return plan_err!("Paimon TableScan projection index is out of bounds"); + }; + if let Some(extractions) = by_column.get(field.name()) { + let extraction_row = variant_extraction_row( + field.data_type().is_nullable(), + extractions.iter().map(|call| { + ( + call.data_type.clone(), + call.path.clone(), + call.fail_on_error, + "UTC".to_string(), + ) + }), + ); + read_type.push( + DataField::new( + field.id(), + field.name().to_string(), + DataType::Row(extraction_row), + ) + .with_description(field.description().map(ToString::to_string)), + ); + } else { + read_type.push(field); + } + } + + let accepted = by_column + .values() + .flat_map(|calls| { + calls + .iter() + .enumerate() + .map(|(idx, call)| AcceptedExtraction { + call: call.clone(), + field_name: idx.to_string(), + }) + }) + .collect::<Vec<_>>(); + let exprs = projection + .expr + .iter() + .cloned() + .map(|expr| rewrite_variant_gets(expr, &accepted)) + .collect::<DFResult<Vec<_>>>()?; + let filter_predicate = filter + .map(|filter| rewrite_variant_gets(filter.predicate.clone(), &accepted)) + .transpose()?; + + let arrow_schema = + paimon::arrow::build_target_arrow_schema(&read_type).map_err(to_datafusion_error)?; + let schema = Arc::new(DFSchema::try_from_qualified_schema( + scan.table_name.clone(), + arrow_schema.as_ref(), + )?); + let pushed_variants = describe_pushed_variants(&by_column); + let mut filters = scan.filters.clone(); + if let Some(filter) = filter { + filters.push(filter.predicate.clone()); + } + + let scan_input = LogicalPlan::Extension(Extension { + node: Arc::new(VariantExtractionScanNode { + table: provider.table().clone(), + schema, + arrow_schema, + read_type, + filters, + fetch: scan.fetch, + pushed_variants, + }), + }); + let input = if let Some(predicate) = filter_predicate { + LogicalPlan::Filter(Filter::try_new(predicate, Arc::new(scan_input))?) + } else { + scan_input + }; + + Ok(Some(ProjectionRewrite { input, exprs })) +} + +fn collect_variant_usage( + expr: &Expr, + calls: &mut Vec<VariantGetCall>, + full_columns: &mut HashSet<String>, +) -> DFResult<()> { + match parse_variant_get(expr)? { + VariantGetParse::Scalar(call) => { + calls.push(call); + return Ok(()); + } + VariantGetParse::FullVariant(column) => { + full_columns.insert(column.name); + return Ok(()); + } + VariantGetParse::NotVariantGet => {} + } + + match expr { + Expr::Alias(alias) => collect_variant_usage(alias.expr.as_ref(), calls, full_columns), + Expr::Column(column) => { + full_columns.insert(column.name.clone()); + Ok(()) + } + Expr::BinaryExpr(BinaryExpr { left, right, .. }) => { + collect_variant_usage(left, calls, full_columns)?; + collect_variant_usage(right, calls, full_columns) + } + Expr::Like(Like { expr, pattern, .. }) | Expr::SimilarTo(Like { expr, pattern, .. }) => { + collect_variant_usage(expr, calls, full_columns)?; + collect_variant_usage(pattern, calls, full_columns) + } + Expr::Not(inner) + | Expr::IsNotNull(inner) + | Expr::IsNull(inner) + | Expr::IsTrue(inner) + | Expr::IsFalse(inner) + | Expr::IsUnknown(inner) + | Expr::IsNotTrue(inner) + | Expr::IsNotFalse(inner) + | Expr::IsNotUnknown(inner) + | Expr::Negative(inner) => collect_variant_usage(inner, calls, full_columns), + Expr::Between(Between { + expr, low, high, .. + }) => { + collect_variant_usage(expr, calls, full_columns)?; + collect_variant_usage(low, calls, full_columns)?; + collect_variant_usage(high, calls, full_columns) + } + Expr::Case(Case { + expr, + when_then_expr, + else_expr, + }) => { + if let Some(expr) = expr { + collect_variant_usage(expr, calls, full_columns)?; + } + for (when, then) in when_then_expr { + collect_variant_usage(when, calls, full_columns)?; + collect_variant_usage(then, calls, full_columns)?; + } + if let Some(expr) = else_expr { + collect_variant_usage(expr, calls, full_columns)?; + } + Ok(()) + } + Expr::Cast(Cast { expr, .. }) | Expr::TryCast(TryCast { expr, .. }) => { + collect_variant_usage(expr, calls, full_columns) + } + Expr::ScalarFunction(ScalarFunction { args, .. }) => { + for arg in args { + collect_variant_usage(arg, calls, full_columns)?; + } + Ok(()) + } + Expr::InList(InList { expr, list, .. }) => { + collect_variant_usage(expr, calls, full_columns)?; + for expr in list { + collect_variant_usage(expr, calls, full_columns)?; + } + Ok(()) + } + _ => Ok(()), + } +} + +fn rewrite_variant_gets(expr: Expr, accepted: &[AcceptedExtraction]) -> DFResult<Expr> { + expr.transform(|expr| { + if let VariantGetParse::Scalar(call) = parse_variant_get(&expr)? { + if let Some(accepted) = accepted.iter().find(|accepted| accepted.call == call) { + return Ok(Transformed::yes(get_field( + Expr::Column(call.column), + accepted.field_name.as_str(), + ))); + } + } + Ok(Transformed::no(expr)) + }) + .map(|transformed| transformed.data) +} + +enum VariantGetParse { + Scalar(VariantGetCall), + FullVariant(Column), + NotVariantGet, +} + +fn parse_variant_get(expr: &Expr) -> DFResult<VariantGetParse> { + let Expr::ScalarFunction(func) = expr else { + return Ok(VariantGetParse::NotVariantGet); + }; + let name = func.name(); + if name != "variant_get" && name != "try_variant_get" { + return Ok(VariantGetParse::NotVariantGet); + } + if func.args.len() != 2 && func.args.len() != 3 { + return Ok(VariantGetParse::NotVariantGet); + } + let Expr::Column(column) = &func.args[0] else { + return Ok(VariantGetParse::NotVariantGet); + }; + let Some(path) = string_literal(&func.args[1]) else { + return Ok(VariantGetParse::NotVariantGet); + }; + let Some(type_name) = func.args.get(2).and_then(string_literal) else { + return Ok(VariantGetParse::FullVariant(column.clone())); + }; + let Some(data_type) = paimon_type_for_variant_get(&type_name)? else { + return Ok(VariantGetParse::FullVariant(column.clone())); + }; + Ok(VariantGetParse::Scalar(VariantGetCall { + column: column.clone(), + path, + type_name: type_name.trim().to_ascii_lowercase(), + data_type, + fail_on_error: name == "variant_get", + })) +} + +fn string_literal(expr: &Expr) -> Option<String> { + match expr { + Expr::Literal(ScalarValue::Utf8(Some(value)), _) + | Expr::Literal(ScalarValue::LargeUtf8(Some(value)), _) + | Expr::Literal(ScalarValue::Utf8View(Some(value)), _) => Some(value.clone()), + _ => None, + } +} + +fn paimon_type_for_variant_get(type_name: &str) -> DFResult<Option<DataType>> { + let normalized = type_name.trim().to_ascii_lowercase(); + Ok(match normalized.as_str() { + "variant" => None, + "boolean" | "bool" => Some(DataType::Boolean(BooleanType::new())), + "byte" | "tinyint" => Some(DataType::TinyInt(TinyIntType::new())), + "short" | "smallint" => Some(DataType::SmallInt(SmallIntType::new())), + "int" | "integer" => Some(DataType::Int(IntType::new())), + "long" | "bigint" => Some(DataType::BigInt(BigIntType::new())), + "float" | "real" => Some(DataType::Float(FloatType::new())), + "double" => Some(DataType::Double(DoubleType::new())), + "string" | "varchar" | "text" => Some(DataType::VarChar(VarCharType::string_type())), + "date" => Some(DataType::Date(DateType::new())), Review Comment: One small consistency issue: this pushdown type parser accepts `date`, but the public `variant_get` / `try_variant_get` UDF type parser in `variant_functions.rs` does not support `date` (and the UDF cast path has no Arrow Date output branch). This makes the supported type list diverge from the non-pushdown path. Could we either remove `date` here, or add matching UDF support and tests for `variant_get(..., 'date')`? -- 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]
