QuakeWang commented on code in PR #460: URL: https://github.com/apache/paimon-rust/pull/460#discussion_r3526984345
########## 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, 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 table_fields = provider.table().schema().fields(); Review Comment: This uses the table schema fields directly, so the variant extraction scan drops DataFusion-only fields such as `_ROW_ID` that the normal `PaimonTableProvider::scan` exposes via `datafusion_read_fields` when data evolution is enabled. A query selecting `_ROW_ID` together with a pushed `variant_get` can build an extension scan schema without `_ROW_ID`, causing later projection/field resolution to fail. This rewrite should use the same read-field/projection mapping as the normal scan and keep system fields untouched; please add a regression for data-evolution + `_ROW_ID` + `variant_get`. ########## crates/paimon/src/variant.rs: ########## @@ -2601,6 +2601,29 @@ pub(crate) fn rebuild_shredded( builder.result() } +pub(crate) fn cast_variant_to_shredded_value( + variant: VariantRef<'_>, + data_type: &DataType, + fail_on_error: bool, +) -> Result<Option<ShreddedValue>> { + if variant.is_null()? { + return Ok(None); + } + let Some(scalar) = scalar_schema_for_type(data_type)? else { + return Err(Error::Unsupported { + message: format!("Unsupported Variant extraction target type: {data_type:?}"), + }); + }; + match try_typed_shred(variant, &scalar)? { Review Comment: This cast path does not preserve the existing `variant_get` / `try_variant_get` UDF semantics. `try_typed_shred` only handles exact VariantKind-to-target cases plus a few narrow numeric conversions, while the current UDF also casts strings to numeric/bool/decimal, long/decimal/string to double, and object/array to string. For example, after pushdown `try_variant_get(payload, '$.age_text', 'int')` returns NULL for `{"age_text":"27"}` instead of 27. We should either reuse the same cast semantics here or only push down cases that are provably equivalent, with regression tests for these casts. -- 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]
