This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new b7f573b feat: push down variant extractions in DataFusion (#460)
b7f573b is described below
commit b7f573b2657aef990c08270e29f6df9b948c7039
Author: Jingsong Lee <[email protected]>
AuthorDate: Mon Jul 6 16:47:06 2026 +0800
feat: push down variant extractions in DataFusion (#460)
---
bindings/c/src/table.rs | 15 +-
bindings/python/src/read.rs | 9 +-
crates/integration_tests/tests/append_tables.rs | 3 +-
crates/integration_tests/tests/read_tables.rs | 33 +-
.../datafusion/src/lateral_vector_search.rs | 14 +-
crates/integrations/datafusion/src/lib.rs | 1 +
.../datafusion/src/physical_plan/scan.rs | 54 +-
crates/integrations/datafusion/src/table/mod.rs | 60 +-
.../datafusion/src/variant_pushdown.rs | 607 +++++++++++++++++++++
.../datafusion/tests/variant_pushdown.rs | 361 ++++++++++++
crates/paimon/src/arrow/format/shredding.rs | 9 +-
crates/paimon/src/arrow/shredding.rs | 271 ++++++++-
crates/paimon/src/spec/mod.rs | 3 +
crates/paimon/src/spec/variant_metadata.rs | 217 ++++++++
crates/paimon/src/table/bucket_assigner_cross.rs | 2 +-
crates/paimon/src/table/data_evolution_reader.rs | 14 +-
crates/paimon/src/table/data_evolution_writer.rs | 2 +-
crates/paimon/src/table/data_file_reader.rs | 121 +++-
.../paimon/src/table/full_text_search_builder.rs | 2 +-
.../paimon/src/table/lumina_index_build_builder.rs | 2 +-
crates/paimon/src/table/read_builder.rs | 132 ++---
crates/paimon/src/table/table_scan.rs | 19 +-
crates/paimon/src/table/vector_search_builder.rs | 2 +-
crates/paimon/src/variant.rs | 178 ++++++
docs/src/sql.md | 2 +-
25 files changed, 1961 insertions(+), 172 deletions(-)
diff --git a/bindings/c/src/table.rs b/bindings/c/src/table.rs
index c6496f6..5c1622d 100644
--- a/bindings/c/src/table.rs
+++ b/bindings/c/src/table.rs
@@ -112,8 +112,7 @@ pub unsafe extern "C" fn paimon_read_builder_free(rb: *mut
paimon_read_builder)
///
/// The `columns` parameter is a null-terminated array of null-terminated C
strings.
/// Output order follows the caller-specified order. Unknown or duplicate names
-/// cause `paimon_read_builder_new_read()` to fail; an empty list is a valid
-/// zero-column projection.
+/// are validated immediately; an empty list is a valid zero-column projection.
///
/// # Safety
/// `rb` must be a valid pointer from `paimon_table_new_read_builder`, or null
(returns error).
@@ -149,6 +148,11 @@ pub unsafe extern "C" fn
paimon_read_builder_with_projection(
ptr = ptr.add(1);
}
+ let col_refs: Vec<&str> = col_names.iter().map(String::as_str).collect();
+ if let Err(e) = state.table.new_read_builder().with_projection(&col_refs) {
+ return paimon_error::from_paimon(e);
+ }
+
state.projected_columns = Some(col_names);
std::ptr::null_mut()
}
@@ -229,7 +233,12 @@ pub unsafe extern "C" fn paimon_read_builder_new_read(
// Apply projection if set
if let Some(ref columns) = state.projected_columns {
let col_refs: Vec<&str> = columns.iter().map(|s| s.as_str()).collect();
- rb_rust.with_projection(&col_refs);
+ if let Err(e) = rb_rust.with_projection(&col_refs) {
+ return paimon_result_new_read {
+ read: std::ptr::null_mut(),
+ error: paimon_error::from_paimon(e),
+ };
+ }
}
// Apply filter if set
diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs
index 3cd662c..d49f8fd 100644
--- a/bindings/python/src/read.rs
+++ b/bindings/python/src/read.rs
@@ -69,10 +69,10 @@ fn apply_read_config(
projection: &Option<Vec<String>>,
limit: Option<usize>,
filter: &Option<Predicate>,
-) {
+) -> PyResult<()> {
if let Some(projection) = projection {
let cols: Vec<&str> = projection.iter().map(String::as_str).collect();
- builder.with_projection(&cols);
+ builder.with_projection(&cols).map_err(to_py_err)?;
}
if let Some(limit) = limit {
builder.with_limit(limit);
@@ -80,6 +80,7 @@ fn apply_read_config(
if let Some(filter) = filter {
builder.with_filter(filter.clone());
}
+ Ok(())
}
/// Extract a sequence of Python `Split` objects into core `DataSplit`s.
Accepts
@@ -220,7 +221,7 @@ impl PyTableScan {
let splits = py.detach(|| {
rt.block_on(async {
let mut builder = self.table.new_read_builder();
- apply_read_config(&mut builder, &self.projection, self.limit,
&self.filter);
+ apply_read_config(&mut builder, &self.projection, self.limit,
&self.filter)?;
let plan = builder.new_scan().plan().await.map_err(to_py_err)?;
Ok::<_, PyErr>(plan.splits().to_vec())
})
@@ -246,7 +247,7 @@ impl PyTableRead {
let batches = py.detach(|| {
rt.block_on(async {
let mut builder = self.table.new_read_builder();
- apply_read_config(&mut builder, &self.projection, self.limit,
&self.filter);
+ apply_read_config(&mut builder, &self.projection, self.limit,
&self.filter)?;
// Validate config (e.g. projection) before the empty-splits
fast
// path so an invalid projection fails consistently regardless
of
// how many splits are passed.
diff --git a/crates/integration_tests/tests/append_tables.rs
b/crates/integration_tests/tests/append_tables.rs
index 981c44d..6d22d5e 100644
--- a/crates/integration_tests/tests/append_tables.rs
+++ b/crates/integration_tests/tests/append_tables.rs
@@ -253,7 +253,8 @@ async fn test_unpartitioned_projection() {
// Read with projection
let mut rb = table.new_read_builder();
- rb.with_projection(&["value"]);
+ rb.with_projection(&["value"])
+ .expect("Projection should succeed");
let plan = rb.new_scan().plan().await.unwrap();
let read = rb.new_read().unwrap();
let result: Vec<RecordBatch> = read
diff --git a/crates/integration_tests/tests/read_tables.rs
b/crates/integration_tests/tests/read_tables.rs
index ac45170..c4b7677 100644
--- a/crates/integration_tests/tests/read_tables.rs
+++ b/crates/integration_tests/tests/read_tables.rs
@@ -46,7 +46,9 @@ async fn scan_and_read<C: Catalog + ?Sized>(
let mut read_builder = table.new_read_builder();
if let Some(cols) = projection {
- read_builder.with_projection(cols);
+ read_builder
+ .with_projection(cols)
+ .expect("Invalid projection");
}
let scan = read_builder.new_scan();
let plan = scan.plan().await.expect("Failed to plan scan");
@@ -107,7 +109,9 @@ async fn scan_and_read_with_projection_and_filter(
) -> (Plan, Vec<RecordBatch>) {
let mut read_builder = table.new_read_builder();
if let Some(cols) = projection {
- read_builder.with_projection(cols);
+ read_builder
+ .with_projection(cols)
+ .expect("Invalid projection");
}
read_builder.with_filter(filter);
let scan = read_builder.new_scan();
@@ -486,7 +490,9 @@ async fn test_read_projection_empty() {
let table = get_table_from_catalog(&catalog, "simple_log_table").await;
let mut read_builder = table.new_read_builder();
- read_builder.with_projection(&[]);
+ read_builder
+ .with_projection(&[])
+ .expect("Empty projection should succeed");
let read = read_builder
.new_read()
.expect("Empty projection should succeed");
@@ -528,9 +534,8 @@ async fn test_read_projection_unknown_column() {
let table = get_table_from_catalog(&catalog, "simple_log_table").await;
let mut read_builder = table.new_read_builder();
- read_builder.with_projection(&["id", "nonexistent_column"]);
let err = read_builder
- .new_read()
+ .with_projection(&["id", "nonexistent_column"])
.expect_err("Unknown columns should fail");
assert!(
@@ -551,9 +556,8 @@ async fn test_read_projection_all_invalid() {
let table = get_table_from_catalog(&catalog, "simple_log_table").await;
let mut read_builder = table.new_read_builder();
- read_builder.with_projection(&["nonexistent_a", "nonexistent_b"]);
let err = read_builder
- .new_read()
+ .with_projection(&["nonexistent_a", "nonexistent_b"])
.expect_err("All-invalid projection should fail");
assert!(
@@ -574,9 +578,8 @@ async fn test_read_projection_duplicate_column() {
let table = get_table_from_catalog(&catalog, "simple_log_table").await;
let mut read_builder = table.new_read_builder();
- read_builder.with_projection(&["id", "id"]);
let err = read_builder
- .new_read()
+ .with_projection(&["id", "id"])
.expect_err("Duplicate projection should fail");
assert!(
@@ -3164,7 +3167,9 @@ async fn
test_read_data_evolution_table_with_row_id_projection() {
// Project _ROW_ID along with regular columns
let mut read_builder = table.new_read_builder();
- read_builder.with_projection(&["_ROW_ID", "id", "name"]);
+ read_builder
+ .with_projection(&["_ROW_ID", "id", "name"])
+ .expect("Row ID projection should succeed");
let scan = read_builder.new_scan();
let plan = scan.plan().await.expect("Failed to plan scan");
@@ -3233,7 +3238,9 @@ async fn
test_read_data_evolution_table_only_row_id_with_row_ranges() {
// Project only _ROW_ID with a partial row range
let mid = min_row_id + (max_row_id - min_row_id) / 2;
let mut read_builder = table.new_read_builder();
- read_builder.with_projection(&["_ROW_ID"]);
+ read_builder
+ .with_projection(&["_ROW_ID"])
+ .expect("Row ID projection should succeed");
read_builder.with_row_ranges(vec![RowRange::new(min_row_id, mid)]);
let scan = read_builder.new_scan();
let plan = scan.plan().await.expect("plan");
@@ -3267,7 +3274,9 @@ async fn
test_read_data_evolution_mixed_format_row_id_projection() {
let table = get_table_from_catalog(&catalog,
"data_evolution_mixed_format_add_column").await;
let mut read_builder = table.new_read_builder();
- read_builder.with_projection(&["_ROW_ID", "id"]);
+ read_builder
+ .with_projection(&["_ROW_ID", "id"])
+ .expect("Row ID projection should succeed");
let scan = read_builder.new_scan();
let plan = scan.plan().await.expect("Failed to plan scan");
diff --git a/crates/integrations/datafusion/src/lateral_vector_search.rs
b/crates/integrations/datafusion/src/lateral_vector_search.rs
index 5a44244..30342b1 100644
--- a/crates/integrations/datafusion/src/lateral_vector_search.rs
+++ b/crates/integrations/datafusion/src/lateral_vector_search.rs
@@ -72,9 +72,10 @@ impl QueryPlanner for PaimonQueryPlanner {
logical_plan: &LogicalPlan,
session_state: &SessionState,
) -> DFResult<Arc<dyn ExecutionPlan>> {
- let planner =
DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(
- LateralVectorSearchExtensionPlanner,
- )]);
+ let planner = DefaultPhysicalPlanner::with_extension_planners(vec![
+ Arc::new(LateralVectorSearchExtensionPlanner),
+
Arc::new(crate::variant_pushdown::VariantExtractionExtensionPlanner),
+ ]);
planner
.create_physical_plan(logical_plan, session_state)
.await
@@ -91,8 +92,10 @@ impl RewriteLateralVectorSearch {
}
pub(crate) fn optimizer_rules() -> Vec<Arc<dyn OptimizerRule + Send + Sync>> {
- let mut rules: Vec<Arc<dyn OptimizerRule + Send + Sync>> =
- vec![Arc::new(RewriteLateralVectorSearch::new())];
+ let mut rules: Vec<Arc<dyn OptimizerRule + Send + Sync>> = vec![
+ Arc::new(crate::variant_pushdown::RewriteVariantExtractions::new()),
+ Arc::new(RewriteLateralVectorSearch::new()),
+ ];
rules.extend(Optimizer::default().rules);
rules
}
@@ -611,6 +614,7 @@ async fn read_target_rows(
let mut read_builder = table.new_read_builder();
read_builder
.with_projection(&projection_refs)
+ .map_err(to_datafusion_error)?
.with_row_ranges(row_ranges);
let plan = read_builder
.new_scan()
diff --git a/crates/integrations/datafusion/src/lib.rs
b/crates/integrations/datafusion/src/lib.rs
index 8ede2a0..ec838c0 100644
--- a/crates/integrations/datafusion/src/lib.rs
+++ b/crates/integrations/datafusion/src/lib.rs
@@ -55,6 +55,7 @@ mod table;
mod table_function_args;
mod update;
mod variant_functions;
+mod variant_pushdown;
mod vector_search;
use std::collections::HashMap;
diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs
b/crates/integrations/datafusion/src/physical_plan/scan.rs
index 1c234b2..bcde017 100644
--- a/crates/integrations/datafusion/src/physical_plan/scan.rs
+++ b/crates/integrations/datafusion/src/physical_plan/scan.rs
@@ -27,7 +27,7 @@ use datafusion::physical_plan::execution_plan::{Boundedness,
EmissionType};
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{DisplayAs, ExecutionPlan, Partitioning,
PlanProperties};
use futures::{StreamExt, TryStreamExt};
-use paimon::spec::Predicate;
+use paimon::spec::{DataField, Predicate};
use paimon::table::{ScanTrace, Table};
use paimon::DataSplit;
@@ -41,8 +41,8 @@ use crate::error::to_datafusion_error;
#[derive(Debug)]
pub struct PaimonTableScan {
table: Table,
- /// Projected column names (if None, reads all columns).
- projected_columns: Option<Vec<String>>,
+ /// Full Paimon read type for nested or connector-defined projections.
+ read_type: Vec<DataField>,
/// Filter translated from DataFusion expressions and reused during
execute()
/// so reader-side pruning reaches the actual read path.
pushed_predicate: Option<Predicate>,
@@ -58,6 +58,8 @@ pub struct PaimonTableScan {
filter_exact: bool,
/// Metadata-pruning trace captured during eager scan planning.
scan_trace: Option<ScanTrace>,
+ /// Human-readable Variant extraction summary for explain output.
+ pushed_variants: Option<String>,
}
impl PaimonTableScan {
@@ -65,12 +67,13 @@ impl PaimonTableScan {
pub(crate) fn new(
schema: ArrowSchemaRef,
table: Table,
- projected_columns: Option<Vec<String>>,
+ read_type: Vec<DataField>,
pushed_predicate: Option<Predicate>,
planned_partitions: Vec<Arc<[DataSplit]>>,
limit: Option<usize>,
filter_exact: bool,
scan_trace: Option<ScanTrace>,
+ pushed_variants: Option<String>,
) -> Self {
let plan_properties = Arc::new(PlanProperties::new(
EquivalenceProperties::new(schema.clone()),
@@ -80,13 +83,14 @@ impl PaimonTableScan {
));
Self {
table,
- projected_columns,
+ read_type,
pushed_predicate,
planned_partitions,
plan_properties,
limit,
filter_exact,
scan_trace,
+ pushed_variants,
}
}
@@ -143,16 +147,13 @@ impl ExecutionPlan for PaimonTableScan {
let table = self.table.clone();
let schema = self.schema();
- let projected_columns = self.projected_columns.clone();
+ let read_type = self.read_type.clone();
let pushed_predicate = self.pushed_predicate.clone();
let fut = async move {
let mut read_builder = table.new_read_builder();
- if let Some(ref columns) = projected_columns {
- let col_refs: Vec<&str> = columns.iter().map(|s|
s.as_str()).collect();
- read_builder.with_projection(&col_refs);
- }
+ read_builder.with_read_type(read_type);
if let Some(filter) = pushed_predicate {
read_builder.with_filter(filter);
}
@@ -236,9 +237,12 @@ impl DisplayAs for PaimonTableScan {
self.planned_partitions.len()
)?;
- if let Some(ref columns) = self.projected_columns {
- write!(f, ", projection=[{}]", columns.join(", "))?;
- }
+ let columns = self
+ .read_type
+ .iter()
+ .map(|field| field.name())
+ .collect::<Vec<_>>();
+ write!(f, ", projection=[{}]", columns.join(", "))?;
if let Some(ref predicate) = self.pushed_predicate {
write!(f, ", predicate={predicate}")?;
}
@@ -248,6 +252,9 @@ impl DisplayAs for PaimonTableScan {
if let Some(ref trace) = self.scan_trace {
write!(f, ", trace={trace}")?;
}
+ if let Some(ref pushed_variants) = self.pushed_variants {
+ write!(f, ", PushedVariants=[{pushed_variants}]")?;
+ }
Ok(())
}
}
@@ -282,18 +289,27 @@ mod tests {
)]))
}
+ fn test_read_type() -> Vec<DataField> {
+ vec![DataField::new(
+ 0,
+ "id".to_string(),
+ DataType::Int(IntType::new()),
+ )]
+ }
+
#[test]
fn test_partition_count_empty_plan() {
let schema = test_schema();
let scan = PaimonTableScan::new(
schema,
dummy_table(),
- None,
+ test_read_type(),
None,
vec![Arc::from(Vec::new())],
None,
false,
None,
+ None,
);
assert_eq!(scan.properties().output_partitioning().partition_count(),
1);
}
@@ -309,12 +325,13 @@ mod tests {
let scan = PaimonTableScan::new(
schema,
dummy_table(),
- None,
+ test_read_type(),
None,
planned_partitions,
None,
false,
None,
+ None,
);
assert_eq!(scan.properties().output_partitioning().partition_count(),
3);
}
@@ -387,12 +404,17 @@ mod tests {
let scan = PaimonTableScan::new(
schema,
table,
- Some(vec!["id".to_string()]),
+ vec![DataField::new(
+ 0,
+ "id".to_string(),
+ DataType::Int(IntType::new()),
+ )],
Some(pushed_predicate),
vec![Arc::from(vec![split])],
None,
false,
None,
+ None,
);
let ctx = SessionContext::new();
diff --git a/crates/integrations/datafusion/src/table/mod.rs
b/crates/integrations/datafusion/src/table/mod.rs
index 1b8e349..c0a9457 100644
--- a/crates/integrations/datafusion/src/table/mod.rs
+++ b/crates/integrations/datafusion/src/table/mod.rs
@@ -28,6 +28,9 @@ use datafusion::error::Result as DFResult;
use datafusion::logical_expr::dml::InsertOp;
use datafusion::logical_expr::{Expr, TableProviderFilterPushDown};
use datafusion::physical_plan::ExecutionPlan;
+use paimon::spec::{
+ BigIntType, CoreOptions, DataField, DataType, ROW_ID_FIELD_ID,
ROW_ID_FIELD_NAME,
+};
use paimon::table::Table;
use crate::physical_plan::PaimonDataSink;
@@ -40,6 +43,18 @@ use crate::filter_pushdown::{analyze_filters,
classify_filter_pushdown};
use crate::physical_plan::PaimonTableScan;
use crate::runtime::await_with_runtime;
+pub(crate) fn datafusion_read_fields(table: &Table) -> Vec<DataField> {
+ let mut fields = table.schema().fields().to_vec();
+ if CoreOptions::new(table.schema().options()).data_evolution_enabled() {
+ fields.push(DataField::new(
+ ROW_ID_FIELD_ID,
+ ROW_ID_FIELD_NAME.to_string(),
+ DataType::BigInt(BigIntType::with_nullable(true)),
+ ));
+ }
+ fields
+}
+
/// Read-only table provider for a Paimon table.
///
/// Supports full table scan, column projection, and predicate pushdown for
@@ -60,15 +75,7 @@ impl PaimonTableProvider {
///
/// Loads the table schema and converts it to Arrow for DataFusion.
pub fn try_new(table: Table) -> DFResult<Self> {
- let mut fields = table.schema().fields().to_vec();
- let core_options =
paimon::spec::CoreOptions::new(table.schema().options());
- if core_options.data_evolution_enabled() {
- fields.push(paimon::spec::DataField::new(
- paimon::spec::ROW_ID_FIELD_ID,
- paimon::spec::ROW_ID_FIELD_NAME.to_string(),
-
paimon::spec::DataType::BigInt(paimon::spec::BigIntType::with_nullable(true)),
- ));
- }
+ let fields = datafusion_read_fields(&table);
let schema =
paimon::arrow::build_target_arrow_schema(&fields).map_err(to_datafusion_error)?;
Ok(Self { table, schema })
@@ -113,21 +120,19 @@ pub(crate) struct PaimonScanBuilder<'a> {
impl PaimonScanBuilder<'_> {
/// Build a [`PaimonTableScan`] from the configured parameters.
pub(crate) fn build(self) -> DFResult<Arc<dyn ExecutionPlan>> {
- let (projected_schema, projected_columns) = if let Some(indices) =
self.projection {
+ let read_fields = datafusion_read_fields(self.table);
+ let (projected_schema, read_type) = if let Some(indices) =
self.projection {
let fields: Vec<Field> = indices
.iter()
.map(|&i| self.schema.field(i).clone())
.collect();
- let column_names: Vec<String> = fields.iter().map(|f|
f.name().clone()).collect();
- (Arc::new(Schema::new(fields)), Some(column_names))
- } else {
- let column_names: Vec<String> = self
- .schema
- .fields()
+ let read_type = indices
.iter()
- .map(|f| f.name().clone())
- .collect();
- (self.schema.clone(), Some(column_names))
+ .map(|&i| read_fields[i].clone())
+ .collect::<Vec<_>>();
+ (Arc::new(Schema::new(fields)), read_type)
+ } else {
+ (self.schema.clone(), read_fields)
};
let splits = self.plan.splits().to_vec();
@@ -144,12 +149,13 @@ impl PaimonScanBuilder<'_> {
Ok(Arc::new(PaimonTableScan::new(
projected_schema,
self.table.clone(),
- projected_columns,
+ read_type,
self.pushed_predicate,
planned_partitions,
self.limit,
self.filter_exact,
self.scan_trace,
+ None,
)))
}
}
@@ -174,15 +180,13 @@ impl TableProvider for PaimonTableProvider {
// Plan splits eagerly so we know partition count upfront.
let filter_analysis = analyze_filters(filters,
self.table.schema().fields());
let mut read_builder = self.table.new_read_builder();
- let projected_columns = projection.map(|indices| {
- indices
+ if let Some(indices) = projection {
+ let read_fields = datafusion_read_fields(&self.table);
+ let read_type = indices
.iter()
- .map(|&i| self.schema.field(i).name().clone())
- .collect::<Vec<_>>()
- });
- if let Some(ref columns) = projected_columns {
- let col_refs: Vec<&str> = columns.iter().map(|s|
s.as_str()).collect();
- read_builder.with_projection(&col_refs);
+ .map(|&i| read_fields[i].clone())
+ .collect::<Vec<_>>();
+ read_builder.with_read_type(read_type);
}
if let Some(filter) = filter_analysis.pushed_predicate.clone() {
read_builder.with_filter(filter);
diff --git a/crates/integrations/datafusion/src/variant_pushdown.rs
b/crates/integrations/datafusion/src/variant_pushdown.rs
new file mode 100644
index 0000000..1aff558
--- /dev/null
+++ b/crates/integrations/datafusion/src/variant_pushdown.rs
@@ -0,0 +1,607 @@
+// 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,
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())),
+ "decimal" => Some(DataType::Decimal(
+ DecimalType::new(10, 0).map_err(to_datafusion_error)?,
+ )),
+ _ if normalized.starts_with("decimal(") && normalized.ends_with(')')
=> {
+ let inner = &normalized["decimal(".len()..normalized.len() - 1];
+ let Some((precision, scale)) = inner.split_once(',') else {
+ return plan_err!("Invalid decimal type for variant_get:
{type_name}");
+ };
+ let precision = precision.trim().parse::<u32>().map_err(|e| {
+ datafusion::error::DataFusionError::Plan(format!("Invalid
decimal precision: {e}"))
+ })?;
+ let scale = scale.trim().parse::<u32>().map_err(|e| {
+ datafusion::error::DataFusionError::Plan(format!("Invalid
decimal scale: {e}"))
+ })?;
+ Some(DataType::Decimal(
+ DecimalType::new(precision,
scale).map_err(to_datafusion_error)?,
+ ))
+ }
+ _ => return plan_err!("Unsupported variant_get type for pushdown:
{type_name}"),
+ })
+}
+
+fn describe_pushed_variants(by_column: &HashMap<String, Vec<VariantGetCall>>)
-> String {
+ let mut columns = by_column.keys().cloned().collect::<Vec<_>>();
+ columns.sort();
+ columns
+ .into_iter()
+ .map(|column| {
+ let paths = by_column[&column]
+ .iter()
+ .map(|call| call.path.as_str())
+ .collect::<Vec<_>>()
+ .join(",");
+ format!("{column}=[{paths}]")
+ })
+ .collect::<Vec<_>>()
+ .join(", ")
+}
diff --git a/crates/integrations/datafusion/tests/variant_pushdown.rs
b/crates/integrations/datafusion/tests/variant_pushdown.rs
new file mode 100644
index 0000000..f5dc32f
--- /dev/null
+++ b/crates/integrations/datafusion/tests/variant_pushdown.rs
@@ -0,0 +1,361 @@
+// 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.
+
+mod common;
+
+use datafusion::arrow::array::{Array, Float64Array, Int32Array, Int64Array,
StringArray};
+use datafusion::physical_plan::displayable;
+use paimon_datafusion::SQLContext;
+
+async fn setup_shredded_variant_table() -> (tempfile::TempDir, SQLContext) {
+ let (tmp, sql_context) = common::setup_sql_context().await;
+ common::exec(
+ &sql_context,
+ r#"
+ CREATE TABLE paimon.test_db.t (
+ id INT,
+ payload VARIANT
+ ) WITH (
+ 'file.format' = 'parquet',
+ 'variant.shreddingSchema' =
+
'{"type":"ROW","fields":[{"name":"payload","type":{"type":"ROW","fields":[{"name":"age","type":"INT"},{"name":"city","type":"STRING"}]}}]}'
+ )
+ "#,
+ )
+ .await;
+ (tmp, sql_context)
+}
+
+async fn setup_shredded_variant_table_with_rows() -> (tempfile::TempDir,
SQLContext) {
+ let (tmp, sql_context) = setup_shredded_variant_table().await;
+ common::exec(
+ &sql_context,
+ r#"
+ INSERT INTO paimon.test_db.t
+ SELECT 1,
parse_json('{"age":27,"age_text":"27","city":"Beijing","a;b":11,"unused":"large"}')
+ UNION ALL
+ SELECT 2,
parse_json('{"age":32,"age_text":"32","city":"Hangzhou","a;b":22,"unused":"large"}')
+ "#,
+ )
+ .await;
+ (tmp, sql_context)
+}
+
+async fn setup_data_evolution_shredded_variant_table_with_rows() ->
(tempfile::TempDir, SQLContext)
+{
+ let (tmp, sql_context) = common::setup_sql_context().await;
+ common::exec(
+ &sql_context,
+ r#"
+ CREATE TABLE paimon.test_db.de_t (
+ id INT,
+ payload VARIANT
+ ) WITH (
+ 'file.format' = 'parquet',
+ 'data-evolution.enabled' = 'true',
+ 'variant.shreddingSchema' =
+
'{"type":"ROW","fields":[{"name":"payload","type":{"type":"ROW","fields":[{"name":"age","type":"INT"},{"name":"city","type":"STRING"}]}}]}'
+ )
+ "#,
+ )
+ .await;
+ common::exec(
+ &sql_context,
+ r#"
+ INSERT INTO paimon.test_db.de_t (id, payload)
+ SELECT 1, parse_json('{"age":27,"city":"Beijing"}')
+ UNION ALL
+ SELECT 2, parse_json('{"age":32,"city":"Hangzhou"}')
+ "#,
+ )
+ .await;
+ (tmp, sql_context)
+}
+
+#[tokio::test]
+async fn variant_get_projection_pushes_extractions_into_scan() {
+ let (_tmp, sql_context) = setup_shredded_variant_table_with_rows().await;
+ let sql = r#"
+ SELECT
+ id,
+ variant_get(payload, '$.age', 'int') AS age,
+ variant_get(payload, '$.city', 'string') AS city
+ FROM paimon.test_db.t
+ ORDER BY id
+ "#;
+ let df = sql_context.sql(sql).await.unwrap();
+ let plan = df.create_physical_plan().await.unwrap();
+ let plan_text = displayable(plan.as_ref()).indent(true).to_string();
+ assert!(
+ plan_text.contains("PushedVariants=[payload=[$.age,$.city]]"),
+ "plan should push variant extractions, got:\n{plan_text}"
+ );
+
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
+ assert_eq!(batches.len(), 1);
+ let ids = batches[0]
+ .column_by_name("id")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ let ages = batches[0]
+ .column_by_name("age")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ let cities = batches[0]
+ .column_by_name("city")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+
+ assert_eq!(ids.values(), &[1, 2]);
+ assert_eq!(ages.values(), &[27, 32]);
+ assert_eq!(cities.value(0), "Beijing");
+ assert_eq!(cities.value(1), "Hangzhou");
+}
+
+#[tokio::test]
+async fn variant_get_filter_pushes_extraction_into_scan() {
+ let (_tmp, sql_context) = setup_shredded_variant_table_with_rows().await;
+ let sql = r#"
+ SELECT id
+ FROM paimon.test_db.t
+ WHERE variant_get(payload, '$.age', 'int') >= 30
+ ORDER BY id
+ "#;
+
+ let df = sql_context.sql(sql).await.unwrap();
+ let plan = df.create_physical_plan().await.unwrap();
+ let plan_text = displayable(plan.as_ref()).indent(true).to_string();
+ assert!(
+ plan_text.contains("PushedVariants=[payload=[$.age]]"),
+ "plan should push filter variant extraction, got:\n{plan_text}"
+ );
+
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
+ assert_eq!(batches.len(), 1);
+ let ids = batches[0]
+ .column_by_name("id")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(ids.values(), &[2]);
+}
+
+#[tokio::test]
+async fn variant_get_long_to_double_pushdown_matches_udf_cast() {
+ let (_tmp, sql_context) = setup_shredded_variant_table_with_rows().await;
+ let sql = r#"
+ SELECT variant_get(payload, '$.age', 'double') AS value
+ FROM paimon.test_db.t
+ ORDER BY id
+ "#;
+
+ let df = sql_context.sql(sql).await.unwrap();
+ let plan = df.create_physical_plan().await.unwrap();
+ let plan_text = displayable(plan.as_ref()).indent(true).to_string();
+ assert!(
+ plan_text.contains("PushedVariants=[payload=[$.age]]"),
+ "plan should push long-to-double variant extraction, got:\n{plan_text}"
+ );
+
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
+ assert_eq!(batches.len(), 1);
+ let values = batches[0]
+ .column_by_name("value")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Float64Array>()
+ .unwrap();
+ assert_eq!(values.values(), &[27.0, 32.0]);
+}
+
+#[tokio::test]
+async fn try_variant_get_string_to_int_pushdown_matches_udf_cast() {
+ let (_tmp, sql_context) = setup_shredded_variant_table_with_rows().await;
+ let sql = r#"
+ SELECT try_variant_get(payload, '$.age_text', 'int') AS value
+ FROM paimon.test_db.t
+ ORDER BY id
+ "#;
+
+ let df = sql_context.sql(sql).await.unwrap();
+ let plan = df.create_physical_plan().await.unwrap();
+ let plan_text = displayable(plan.as_ref()).indent(true).to_string();
+ assert!(
+ plan_text.contains("PushedVariants=[payload=[$.age_text]]"),
+ "plan should push string-to-int try_variant_get extraction,
got:\n{plan_text}"
+ );
+
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
+ assert_eq!(batches.len(), 1);
+ let values = batches[0]
+ .column_by_name("value")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(values.values(), &[27, 32]);
+}
+
+#[tokio::test]
+async fn variant_get_date_type_remains_unsupported() {
+ let (_tmp, sql_context) = setup_shredded_variant_table_with_rows().await;
+ let sql = r#"
+ SELECT variant_get(payload, '$.age', 'date') AS value
+ FROM paimon.test_db.t
+ "#;
+
+ let err = match sql_context.sql(sql).await {
+ Ok(df) => df.collect().await.unwrap_err(),
+ Err(err) => err,
+ };
+ assert!(
+ err.to_string()
+ .contains("Unsupported variant_get type: date"),
+ "expected public variant_get type parser to reject date, got: {err:?}"
+ );
+}
+
+#[tokio::test]
+async fn full_variant_projection_prevents_extraction_pushdown() {
+ let (_tmp, sql_context) = setup_shredded_variant_table_with_rows().await;
+ let sql = r#"
+ SELECT
+ payload,
+ variant_get(payload, '$.age', 'int') AS age
+ FROM paimon.test_db.t
+ ORDER BY age
+ "#;
+
+ let df = sql_context.sql(sql).await.unwrap();
+ let plan = df.create_physical_plan().await.unwrap();
+ let plan_text = displayable(plan.as_ref()).indent(true).to_string();
+ assert!(
+ !plan_text.contains("PushedVariants="),
+ "plan should keep full Variant reads when the query also projects
payload, got:\n{plan_text}"
+ );
+
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
+ assert_eq!(batches.len(), 1);
+ let ages = batches[0]
+ .column_by_name("age")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(ages.values(), &[27, 32]);
+}
+
+#[tokio::test]
+async fn data_evolution_row_id_survives_variant_extraction_pushdown() {
+ let (_tmp, sql_context) =
setup_data_evolution_shredded_variant_table_with_rows().await;
+ let sql = r#"
+ SELECT "_ROW_ID", variant_get(payload, '$.age', 'int') AS age
+ FROM paimon.test_db.de_t
+ ORDER BY "_ROW_ID"
+ "#;
+
+ let df = sql_context.sql(sql).await.unwrap();
+ let plan = df.create_physical_plan().await.unwrap();
+ let plan_text = displayable(plan.as_ref()).indent(true).to_string();
+ assert!(
+ plan_text.contains("PushedVariants=[payload=[$.age]]"),
+ "plan should push variant extraction while preserving _ROW_ID,
got:\n{plan_text}"
+ );
+
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
+ assert_eq!(batches.len(), 1);
+ let row_ids = batches[0]
+ .column_by_name("_ROW_ID")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ let ages = batches[0]
+ .column_by_name("age")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+
+ assert_eq!(row_ids.len(), 2);
+ assert_eq!(ages.values(), &[27, 32]);
+}
+
+#[tokio::test]
+async fn try_variant_get_invalid_path_pushdown_returns_null() {
+ let (_tmp, sql_context) = setup_shredded_variant_table_with_rows().await;
+ let sql = r#"
+ SELECT try_variant_get(payload, 'invalid_path', 'int') AS value
+ FROM paimon.test_db.t
+ ORDER BY id
+ "#;
+
+ let df = sql_context.sql(sql).await.unwrap();
+ let plan = df.create_physical_plan().await.unwrap();
+ let plan_text = displayable(plan.as_ref()).indent(true).to_string();
+ assert!(
+ plan_text.contains("PushedVariants=[payload=[invalid_path]]"),
+ "plan should push invalid-path try_variant_get extraction,
got:\n{plan_text}"
+ );
+
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
+ assert_eq!(batches.len(), 1);
+ let values = batches[0]
+ .column_by_name("value")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(values.len(), 2);
+ assert!(values.is_null(0));
+ assert!(values.is_null(1));
+}
+
+#[tokio::test]
+async fn variant_get_path_with_semicolon_pushes_extraction() {
+ let (_tmp, sql_context) = setup_shredded_variant_table_with_rows().await;
+ let sql = r#"
+ SELECT variant_get(payload, '$.a;b', 'int') AS value
+ FROM paimon.test_db.t
+ ORDER BY id
+ "#;
+
+ let df = sql_context.sql(sql).await.unwrap();
+ let plan = df.create_physical_plan().await.unwrap();
+ let plan_text = displayable(plan.as_ref()).indent(true).to_string();
+ assert!(
+ plan_text.contains("PushedVariants=[payload=[$.a;b]]"),
+ "plan should push semicolon-path variant extraction, got:\n{plan_text}"
+ );
+
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
+ assert_eq!(batches.len(), 1);
+ let values = batches[0]
+ .column_by_name("value")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(values.values(), &[11, 22]);
+}
diff --git a/crates/paimon/src/arrow/format/shredding.rs
b/crates/paimon/src/arrow/format/shredding.rs
index 2e236c6..d33e98a 100644
--- a/crates/paimon/src/arrow/format/shredding.rs
+++ b/crates/paimon/src/arrow/format/shredding.rs
@@ -19,8 +19,9 @@ use super::{FilePredicates, FormatFileReader,
FormatFileWriter};
use crate::arrow::build_target_arrow_schema;
use crate::arrow::shredding::{
assemble_shredded_variant_batch, batch_to_shredded_physical,
- configured_variant_shredding_fields, contains_variant_fields,
infer_variant_shredding_fields,
- should_infer_variant_shredding_fields,
variant_shredding_infer_buffer_row_count,
+ configured_variant_shredding_fields, contains_variant_fields,
contains_variant_read_fields,
+ infer_variant_shredding_fields, should_infer_variant_shredding_fields,
+ variant_shredding_infer_buffer_row_count,
};
use crate::io::FileRead;
use crate::spec::DataField;
@@ -48,7 +49,7 @@ pub(crate) fn maybe_wrap_reader(
reader: Box<dyn FormatFileReader>,
read_fields: &[DataField],
) -> Box<dyn FormatFileReader> {
- if contains_variant_fields(read_fields) {
+ if contains_variant_read_fields(read_fields) {
Box::new(ShreddingFormatReader::new(reader))
} else {
reader
@@ -83,7 +84,7 @@ impl FormatFileReader for ShreddingFormatReader {
row_selection,
)
.await?;
- if !contains_variant_fields(read_fields) {
+ if !contains_variant_read_fields(read_fields) {
return Ok(stream);
}
let read_fields = read_fields.to_vec();
diff --git a/crates/paimon/src/arrow/shredding.rs
b/crates/paimon/src/arrow/shredding.rs
index 3556c1b..a5d0356 100644
--- a/crates/paimon/src/arrow/shredding.rs
+++ b/crates/paimon/src/arrow/shredding.rs
@@ -18,12 +18,14 @@
use crate::arrow::{
arrow_to_paimon_type, build_target_arrow_schema, is_variant_arrow_fields,
paimon_type_to_arrow,
};
-use crate::spec::{ArrayType, DataField, DataType, RowType};
+use crate::spec::{
+ is_variant_extraction_row, parse_variant_metadata, ArrayType, DataField,
DataType, RowType,
+};
use crate::variant::{
- build_variant_schema, cast_shredded, infer_variant_shredding_schema,
rebuild_shredded,
- variant_shredding_type, GenericVariant, ShreddedRow, ShreddedValue,
- VariantShreddingInferConfig, VARIANT_METADATA_FIELD_NAME,
VARIANT_TYPED_VALUE_FIELD_NAME,
- VARIANT_VALUE_FIELD_NAME,
+ build_variant_schema, cast_shredded, cast_variant_to_shredded_value,
+ infer_variant_shredding_schema, rebuild_shredded, variant_shredding_type,
GenericVariant,
+ ShreddedRow, ShreddedValue, VariantShreddingInferConfig,
VARIANT_METADATA_FIELD_NAME,
+ VARIANT_TYPED_VALUE_FIELD_NAME, VARIANT_VALUE_FIELD_NAME,
};
use crate::{Error, Result};
use arrow_array::{
@@ -258,6 +260,15 @@ pub(crate) fn contains_variant_fields(fields:
&[DataField]) -> bool {
})
}
+pub(crate) fn contains_variant_read_fields(fields: &[DataField]) -> bool {
+ fields.iter().any(|field| match field.data_type() {
+ DataType::Variant(_) => true,
+ DataType::Row(row_type) if is_variant_extraction_row(row_type) => true,
+ DataType::Row(row_type) =>
contains_variant_read_fields(row_type.fields()),
+ _ => false,
+ })
+}
+
fn paths_to_variant(fields: &[DataField]) -> Vec<Vec<usize>> {
let mut result = Vec::new();
for (idx, field) in fields.iter().enumerate() {
@@ -680,11 +691,130 @@ fn assemble_array_to_logical(
DataType::Variant(_) if is_variant_storage_array(array) => {
Ok(Some(assemble_shredded_variant_array(array)?))
}
+ DataType::Row(row_type)
+ if is_variant_extraction_row(row_type)
+ && (is_variant_storage_array(array) ||
is_plain_variant_array(array)) =>
+ {
+ Ok(Some(assemble_variant_extraction_array(array, row_type)?))
+ }
DataType::Row(row_type) => assemble_row_array_to_logical(array,
row_type),
_ => Ok(None),
}
}
+fn assemble_variant_extraction_array(array: &dyn Array, row_type: &RowType) ->
Result<ArrayRef> {
+ let input = array
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .ok_or_else(|| Error::DataInvalid {
+ message: "Variant extraction column must be
StructArray".to_string(),
+ source: None,
+ })?;
+ let fields = row_type.fields();
+ let metadata = fields
+ .iter()
+ .map(|field| {
+ let Some(description) = field.description() else {
+ return Err(Error::DataInvalid {
+ message: "Variant extraction field is missing
metadata".to_string(),
+ source: None,
+ });
+ };
+ parse_variant_metadata(description)
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ let mut values_by_field = vec![Vec::with_capacity(input.len());
fields.len()];
+ let mut validities = Vec::with_capacity(input.len());
+ for row in 0..input.len() {
+ if input.is_null(row) {
+ validities.push(false);
+ for values in &mut values_by_field {
+ values.push(None);
+ }
+ continue;
+ }
+
+ validities.push(true);
+ let variant = variant_from_storage_row(input, row)?;
+ for (field_idx, field) in fields.iter().enumerate() {
+ let field_metadata = &metadata[field_idx];
+ let value = match &variant {
+ Some(variant) => match variant.get_path(field_metadata.path())
{
+ Ok(Some(extracted)) => cast_variant_to_shredded_value(
+ extracted,
+ field.data_type(),
+ field_metadata.fail_on_error(),
+ )?,
+ Ok(None) => None,
+ Err(e) if !field_metadata.fail_on_error() => {
+ let _ = e;
+ None
+ }
+ Err(e) => return Err(e),
+ },
+ None => None,
+ };
+ values_by_field[field_idx].push(value);
+ }
+ }
+
+ let arrow_fields: Fields = fields
+ .iter()
+ .map(|field| {
+ Ok(ArrowField::new(
+ field.name(),
+ paimon_type_to_arrow(field.data_type())?,
+ field.data_type().is_nullable(),
+ ))
+ })
+ .collect::<Result<Vec<_>>>()?
+ .into();
+ let columns = values_by_field
+ .iter()
+ .zip(fields)
+ .map(|(values, field)| array_from_values(values, field.data_type()))
+ .collect::<Result<Vec<_>>>()?;
+
+ Ok(Arc::new(
+ StructArray::try_new(arrow_fields, columns,
Some(null_buffer(validities))).map_err(
+ |e| Error::UnexpectedError {
+ message: format!("Failed to build Variant extraction
StructArray: {e}"),
+ source: Some(Box::new(e)),
+ },
+ )?,
+ ))
+}
+
+fn variant_from_storage_row(input: &StructArray, row: usize) ->
Result<Option<GenericVariant>> {
+ if input.is_null(row) {
+ return Ok(None);
+ }
+ if input
+ .column_by_name(VARIANT_TYPED_VALUE_FIELD_NAME)
+ .is_none()
+ {
+ return variant_from_array(input, row);
+ }
+
+ let row_type = match arrow_to_paimon_type(input.data_type(), true)? {
+ DataType::Row(row_type) => row_type,
+ DataType::Variant(_) => return variant_from_array(input, row),
+ other => {
+ return Err(Error::DataInvalid {
+ message: format!("Variant storage physical type must be ROW,
got {other:?}"),
+ source: None,
+ })
+ }
+ };
+ let schema = build_variant_schema(&row_type)?;
+ let shredded = struct_row_at(input, row, &row_type)?.ok_or_else(||
Error::DataInvalid {
+ message: "Variant extraction storage row is null".to_string(),
+ source: None,
+ })?;
+ rebuild_shredded(&shredded, &schema).map(Some)
+}
+
fn assemble_row_array_to_logical(
array: &dyn Array,
row_type: &RowType,
@@ -763,6 +893,13 @@ fn is_variant_storage_array(array: &dyn Array) -> bool {
&& (!is_variant_arrow_fields(fields) ||
is_shredded_variant_array(array))
}
+fn is_plain_variant_array(array: &dyn Array) -> bool {
+ let ArrowDataType::Struct(fields) = array.data_type() else {
+ return false;
+ };
+ is_variant_arrow_fields(fields)
+}
+
fn shredded_rows_to_struct_array(
rows: Vec<Option<ShreddedRow>>,
row_type: &RowType,
@@ -1186,7 +1323,7 @@ fn null_buffer(validities: Vec<bool>) -> NullBuffer {
mod tests {
use super::*;
use crate::arrow::variant_arrow_type;
- use crate::spec::{IntType, VariantType};
+ use crate::spec::{variant_extraction_row, IntType, VarCharType,
VariantType};
fn variant_array_for_test(values: &[GenericVariant]) -> ArrayRef {
let value_items = values
@@ -1289,6 +1426,128 @@ mod tests {
);
}
+ #[test]
+ fn variant_extraction_row_reads_plain_variant_array() {
+ let variants = vec![
+
GenericVariant::parse_json(r#"{"age":27,"name":"Alice"}"#).unwrap(),
+ GenericVariant::parse_json(r#"{"age":"old"}"#).unwrap(),
+ ];
+ let array = variant_array_for_test(&variants);
+ let row_type = variant_extraction_row(
+ true,
+ vec![
+ (
+ DataType::Int(IntType::new()),
+ "$.age".to_string(),
+ false,
+ "UTC".to_string(),
+ ),
+ (
+ DataType::VarChar(VarCharType::string_type()),
+ "$.name".to_string(),
+ true,
+ "UTC".to_string(),
+ ),
+ ],
+ );
+
+ let extracted = assemble_array_to_logical(array.as_ref(),
&DataType::Row(row_type))
+ .unwrap()
+ .expect("extracted");
+ let extracted =
extracted.as_any().downcast_ref::<StructArray>().unwrap();
+ let ages = extracted
+ .column_by_name("0")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ let names = extracted
+ .column_by_name("1")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+
+ assert_eq!(ages.value(0), 27);
+ assert!(ages.is_null(1));
+ assert_eq!(names.value(0), "Alice");
+ assert!(names.is_null(1));
+ }
+
+ #[test]
+ fn variant_extraction_row_reads_shredded_variant_array() {
+ let logical_fields = vec![DataField::new(
+ 1,
+ "v".to_string(),
+ DataType::Variant(VariantType::new()),
+ )];
+ let options = HashMap::from([(
+ "variant.shreddingSchema".to_string(),
+
r#"{"type":"ROW","fields":[{"name":"v","type":{"type":"ROW","fields":[{"name":"age","type":"INT"},{"name":"name","type":"STRING"}]}}]}"#.to_string(),
+ )]);
+ let physical_fields =
configured_variant_shredding_fields(&logical_fields, &options)
+ .unwrap()
+ .expect("shredding fields");
+ let variants = vec![
+
GenericVariant::parse_json(r#"{"age":27,"name":"Alice"}"#).unwrap(),
+
GenericVariant::parse_json(r#"{"age":"old","name":"Bob"}"#).unwrap(),
+ ];
+ let batch = RecordBatch::try_new(
+ build_target_arrow_schema(&logical_fields).unwrap(),
+ vec![variant_array_for_test(&variants)],
+ )
+ .unwrap();
+ let physical =
+ batch_to_shredded_physical(&batch, &logical_fields,
&physical_fields).unwrap();
+ let extraction_row_type = variant_extraction_row(
+ true,
+ vec![
+ (
+ DataType::Int(IntType::new()),
+ "$.age".to_string(),
+ false,
+ "UTC".to_string(),
+ ),
+ (
+ DataType::VarChar(VarCharType::string_type()),
+ "$.name".to_string(),
+ true,
+ "UTC".to_string(),
+ ),
+ ],
+ );
+ let read_fields = vec![DataField::new(
+ 1,
+ "v".to_string(),
+ DataType::Row(extraction_row_type),
+ )];
+
+ let extracted_batch = assemble_shredded_variant_batch(physical,
&read_fields).unwrap();
+ let extracted = extracted_batch
+ .column_by_name("v")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap();
+ let ages = extracted
+ .column_by_name("0")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ let names = extracted
+ .column_by_name("1")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+
+ assert_eq!(ages.value(0), 27);
+ assert!(ages.is_null(1));
+ assert_eq!(names.value(0), "Alice");
+ assert_eq!(names.value(1), "Bob");
+ }
+
#[test]
fn inferred_schema_transforms_and_reassembles_variant_batch() {
let logical_fields = vec![
diff --git a/crates/paimon/src/spec/mod.rs b/crates/paimon/src/spec/mod.rs
index 79bdc7c..12a0cc6 100644
--- a/crates/paimon/src/spec/mod.rs
+++ b/crates/paimon/src/spec/mod.rs
@@ -49,6 +49,9 @@ pub use schema::*;
mod schema_change;
pub use schema_change::*;
+mod variant_metadata;
+pub use variant_metadata::*;
+
mod snapshot;
pub use snapshot::*;
diff --git a/crates/paimon/src/spec/variant_metadata.rs
b/crates/paimon/src/spec/variant_metadata.rs
new file mode 100644
index 0000000..b172931
--- /dev/null
+++ b/crates/paimon/src/spec/variant_metadata.rs
@@ -0,0 +1,217 @@
+// 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.
+
+//! Java-compatible markers for Variant extraction read types.
+
+use crate::spec::{DataField, DataType, RowType};
+use crate::{Error, Result};
+use serde::{Deserialize, Serialize};
+
+pub const VARIANT_METADATA_KEY: &str = "__VARIANT_METADATA";
+const VARIANT_METADATA_DELIMITER: &str = ";";
+
+#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct VariantFieldMetadata {
+ path: String,
+ fail_on_error: bool,
+ time_zone_id: String,
+}
+
+impl VariantFieldMetadata {
+ pub fn new(
+ path: impl Into<String>,
+ fail_on_error: bool,
+ time_zone_id: impl Into<String>,
+ ) -> Self {
+ Self {
+ path: path.into(),
+ fail_on_error,
+ time_zone_id: time_zone_id.into(),
+ }
+ }
+
+ pub fn path(&self) -> &str {
+ &self.path
+ }
+
+ pub fn fail_on_error(&self) -> bool {
+ self.fail_on_error
+ }
+
+ pub fn time_zone_id(&self) -> &str {
+ &self.time_zone_id
+ }
+}
+
+pub fn build_variant_metadata(path: &str, fail_on_error: bool, time_zone_id:
&str) -> String {
+ let metadata = VariantFieldMetadata::new(path, fail_on_error,
time_zone_id);
+ let json = serde_json::to_string(&metadata)
+ .expect("Variant metadata serialization should not fail for string and
bool fields");
+ format!("{VARIANT_METADATA_KEY}{json}")
+}
+
+pub fn parse_variant_metadata(description: &str) ->
Result<VariantFieldMetadata> {
+ let Some(raw) = description.strip_prefix(VARIANT_METADATA_KEY) else {
+ return Err(Error::DataInvalid {
+ message: "Variant metadata description is missing
marker".to_string(),
+ source: None,
+ });
+ };
+ if raw.trim_start().starts_with('{') {
+ let metadata =
+ serde_json::from_str::<VariantFieldMetadata>(raw).map_err(|e|
Error::DataInvalid {
+ message: "Malformed Variant metadata JSON
description".to_string(),
+ source: Some(Box::new(e)),
+ })?;
+ validate_variant_metadata(&metadata)?;
+ return Ok(metadata);
+ }
+
+ parse_legacy_variant_metadata(raw)
+}
+
+fn parse_legacy_variant_metadata(raw: &str) -> Result<VariantFieldMetadata> {
+ let mut parts = raw.split(VARIANT_METADATA_DELIMITER);
+ let path = parts.next().unwrap_or_default();
+ let fail_on_error = parts.next().ok_or_else(|| Error::DataInvalid {
+ message: "Variant metadata description is missing
failOnError".to_string(),
+ source: None,
+ })?;
+ let time_zone_id = parts.next().ok_or_else(|| Error::DataInvalid {
+ message: "Variant metadata description is missing
timeZoneId".to_string(),
+ source: None,
+ })?;
+ if parts.next().is_some() || path.is_empty() {
+ return Err(Error::DataInvalid {
+ message: "Malformed Variant metadata description".to_string(),
+ source: None,
+ });
+ }
+ let fail_on_error = fail_on_error
+ .parse::<bool>()
+ .map_err(|e| Error::DataInvalid {
+ message: format!("Invalid Variant metadata failOnError value
'{fail_on_error}'"),
+ source: Some(Box::new(e)),
+ })?;
+ Ok(VariantFieldMetadata::new(
+ path.to_string(),
+ fail_on_error,
+ time_zone_id.to_string(),
+ ))
+}
+
+fn validate_variant_metadata(metadata: &VariantFieldMetadata) -> Result<()> {
+ if metadata.path().is_empty() {
+ return Err(Error::DataInvalid {
+ message: "Malformed Variant metadata description".to_string(),
+ source: None,
+ });
+ }
+ Ok(())
+}
+
+pub fn is_variant_metadata_description(description: Option<&str>) -> bool {
+ description.is_some_and(|description|
description.starts_with(VARIANT_METADATA_KEY))
+}
+
+pub fn is_variant_extraction_row_type(data_type: &DataType) -> bool {
+ let DataType::Row(row_type) = data_type else {
+ return false;
+ };
+ is_variant_extraction_row(row_type)
+}
+
+pub fn is_variant_extraction_row(row_type: &RowType) -> bool {
+ !row_type.fields().is_empty()
+ && row_type
+ .fields()
+ .iter()
+ .all(|field| is_variant_metadata_description(field.description()))
+}
+
+pub fn variant_extraction_row(
+ nullable: bool,
+ extractions: impl IntoIterator<Item = (DataType, String, bool, String)>,
+) -> RowType {
+ let fields = extractions
+ .into_iter()
+ .enumerate()
+ .map(|(idx, (data_type, path, fail_on_error, time_zone_id))| {
+ DataField::new(idx as i32, idx.to_string(),
data_type).with_description(Some(
+ build_variant_metadata(&path, fail_on_error, &time_zone_id),
+ ))
+ })
+ .collect();
+ RowType::with_nullable(nullable, fields)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::spec::{DataType, IntType, VarCharType};
+
+ #[test]
+ fn parses_variant_metadata_description() {
+ let description = build_variant_metadata("$.a", false, "UTC");
+ let metadata = parse_variant_metadata(&description).unwrap();
+ assert_eq!(metadata.path(), "$.a");
+ assert!(!metadata.fail_on_error());
+ assert_eq!(metadata.time_zone_id(), "UTC");
+ }
+
+ #[test]
+ fn parses_variant_metadata_description_with_delimiters() {
+ let description = build_variant_metadata("$.a;b", true, "UTC;8");
+ let metadata = parse_variant_metadata(&description).unwrap();
+ assert_eq!(metadata.path(), "$.a;b");
+ assert!(metadata.fail_on_error());
+ assert_eq!(metadata.time_zone_id(), "UTC;8");
+ }
+
+ #[test]
+ fn parses_legacy_variant_metadata_description() {
+ let description = format!("{VARIANT_METADATA_KEY}$.a;false;UTC");
+ let metadata = parse_variant_metadata(&description).unwrap();
+ assert_eq!(metadata.path(), "$.a");
+ assert!(!metadata.fail_on_error());
+ assert_eq!(metadata.time_zone_id(), "UTC");
+ }
+
+ #[test]
+ fn identifies_variant_extraction_row_type() {
+ let row = variant_extraction_row(
+ true,
+ vec![
+ (
+ DataType::Int(IntType::new()),
+ "$.age".to_string(),
+ true,
+ "UTC".to_string(),
+ ),
+ (
+ DataType::VarChar(VarCharType::string_type()),
+ "$.name".to_string(),
+ false,
+ "UTC".to_string(),
+ ),
+ ],
+ );
+ assert!(is_variant_extraction_row(&row));
+ assert!(is_variant_extraction_row_type(&DataType::Row(row)));
+ }
+}
diff --git a/crates/paimon/src/table/bucket_assigner_cross.rs
b/crates/paimon/src/table/bucket_assigner_cross.rs
index c3f1acc..f0e711b 100644
--- a/crates/paimon/src/table/bucket_assigner_cross.rs
+++ b/crates/paimon/src/table/bucket_assigner_cross.rs
@@ -86,7 +86,7 @@ impl GlobalPartitionIndex {
let projected_pk_indices: Vec<usize> = (0..pk_fields.len()).collect();
let mut rb = table.new_read_builder();
- rb.with_projection(&pk_field_names);
+ rb.with_projection(&pk_field_names)?;
let scan = rb.new_scan().with_scan_all_files();
let plan = scan.plan().await?;
let read = rb.new_read()?;
diff --git a/crates/paimon/src/table/data_evolution_reader.rs
b/crates/paimon/src/table/data_evolution_reader.rs
index 1a1d86d..7bebb5c 100644
--- a/crates/paimon/src/table/data_evolution_reader.rs
+++ b/crates/paimon/src/table/data_evolution_reader.rs
@@ -3799,7 +3799,10 @@ mod tests {
]);
let mut builder = table.new_read_builder();
- builder.with_projection(&["id"]).with_filter(predicate);
+ builder
+ .with_projection(&["id"])
+ .unwrap()
+ .with_filter(predicate);
let read = builder.new_read().unwrap();
let batches = read
.to_arrow(&[split])
@@ -3864,7 +3867,10 @@ mod tests {
let predicate = pb.equal("value", Datum::Int(20)).unwrap();
let mut builder = table.new_read_builder();
- builder.with_projection(&["id"]).with_filter(predicate);
+ builder
+ .with_projection(&["id"])
+ .unwrap()
+ .with_filter(predicate);
let read = builder.new_read().unwrap();
let batches = read
.to_arrow(&[split])
@@ -3917,6 +3923,7 @@ mod tests {
let mut builder = table.new_read_builder();
builder
.with_projection(&["id", crate::spec::ROW_ID_FIELD_NAME])
+ .unwrap()
.with_filter(predicate);
let read = builder.new_read().unwrap();
let batches = read
@@ -3981,6 +3988,7 @@ mod tests {
let mut builder = table.new_read_builder();
builder
.with_projection(&["id", crate::spec::ROW_ID_FIELD_NAME])
+ .unwrap()
.with_filter(predicate);
let read = builder.new_read().unwrap();
let batches = read
@@ -4066,6 +4074,7 @@ mod tests {
let mut builder = table.new_read_builder();
builder
.with_projection(&["id"])
+ .unwrap()
.with_filter(pb.equal("value", Datum::Int(10)).unwrap());
let read = builder.new_read().unwrap();
let batches = read
@@ -4080,6 +4089,7 @@ mod tests {
let mut builder = table.new_read_builder();
builder
.with_projection(&["id"])
+ .unwrap()
.with_filter(pb.is_null("value").unwrap());
let read = builder.new_read().unwrap();
let batches = read
diff --git a/crates/paimon/src/table/data_evolution_writer.rs
b/crates/paimon/src/table/data_evolution_writer.rs
index a319d77..2228b78 100644
--- a/crates/paimon/src/table/data_evolution_writer.rs
+++ b/crates/paimon/src/table/data_evolution_writer.rs
@@ -234,7 +234,7 @@ impl DataEvolutionWriter {
// Read original columns from the entire file group (base +
partial-column files).
let col_refs: Vec<&str> = self.update_columns.iter().map(|s|
s.as_str()).collect();
let mut rb = self.table.new_read_builder();
- rb.with_projection(&col_refs);
+ rb.with_projection(&col_refs)?;
let read = rb.new_read()?;
// Base + partial-column files share row-id ranges, so physical
diff --git a/crates/paimon/src/table/data_file_reader.rs
b/crates/paimon/src/table/data_file_reader.rs
index 6a4457f..5a7e0bf 100644
--- a/crates/paimon/src/table/data_file_reader.rs
+++ b/crates/paimon/src/table/data_file_reader.rs
@@ -20,7 +20,9 @@ use crate::arrow::format::create_format_reader;
use crate::arrow::schema_evolution::{create_index_mapping, NULL_FIELD_INDEX};
use crate::deletion_vector::{DeletionVector, DeletionVectorFactory};
use crate::io::FileIO;
-use crate::spec::{DataField, DataFileMeta, Predicate, ROW_ID_FIELD_NAME};
+use crate::spec::{
+ is_variant_extraction_row_type, DataField, DataFileMeta, DataType,
Predicate, ROW_ID_FIELD_NAME,
+};
use crate::table::schema_manager::SchemaManager;
use crate::table::ArrowRecordBatchStream;
use crate::table::RowRange;
@@ -199,18 +201,8 @@ impl DataFileReader {
// Compute index mapping and determine which columns to read from the
file.
let (projected_read_fields, index_mapping) = if let Some(ref df) =
data_fields {
let mapping = create_index_mapping(&read_type, df);
- match mapping {
- Some(ref idx_map) => {
- let mut seen = std::collections::HashSet::new();
- let fields_to_read: Vec<DataField> = idx_map
- .iter()
- .filter(|&&idx| idx != NULL_FIELD_INDEX &&
seen.insert(idx))
- .map(|&idx| df[idx as usize].clone())
- .collect();
- (fields_to_read, Some(idx_map.clone()))
- }
- None => (df.clone(), None),
- }
+ let fields_to_read = read_data_fields(df, &read_type)?;
+ (fields_to_read, mapping)
} else {
(
read_type
@@ -367,6 +359,67 @@ impl DataFileReader {
}
}
+fn read_data_fields(
+ all_data_fields: &[DataField],
+ expected_fields: &[DataField],
+) -> crate::Result<Vec<DataField>> {
+ let mut read_fields = Vec::new();
+ for data_field in all_data_fields {
+ if let Some(expected) = expected_fields
+ .iter()
+ .find(|field| field.id() == data_field.id())
+ {
+ if let Some(pruned_type) =
+ prune_data_type(expected.data_type(), data_field.data_type())?
+ {
+ read_fields.push(data_field_with_type(data_field,
pruned_type));
+ }
+ }
+ }
+ Ok(read_fields)
+}
+
+fn prune_data_type(read_type: &DataType, data_type: &DataType) ->
crate::Result<Option<DataType>> {
+ match read_type {
+ DataType::Row(read_row) if is_variant_extraction_row_type(read_type)
=> {
+ Ok(Some(DataType::Row(read_row.clone())))
+ }
+ DataType::Row(read_row) => {
+ let DataType::Row(data_row) = data_type else {
+ return Ok(Some(data_type.clone()));
+ };
+ let mut fields = Vec::new();
+ for read_field in read_row.fields() {
+ if let Some(data_field) = data_row
+ .fields()
+ .iter()
+ .find(|field| field.id() == read_field.id())
+ {
+ if let Some(pruned_type) =
+ prune_data_type(read_field.data_type(),
data_field.data_type())?
+ {
+ fields.push(data_field_with_type(data_field,
pruned_type));
+ }
+ }
+ }
+ if fields.is_empty() {
+ Ok(None)
+ } else {
+ Ok(Some(DataType::Row(crate::spec::RowType::with_nullable(
+ read_type.is_nullable(),
+ fields,
+ ))))
+ }
+ }
+ _ => Ok(Some(data_type.clone())),
+ }
+}
+
+fn data_field_with_type(field: &DataField, data_type: DataType) -> DataField {
+ DataField::new(field.id(), field.name().to_string(), data_type)
+ .with_description(field.description().map(ToString::to_string))
+}
+
fn is_row_file(file_meta: &DataFileMeta) -> bool {
file_meta.file_name.to_ascii_lowercase().ends_with(".row")
|| file_meta
@@ -596,9 +649,11 @@ mod row_tests {
use crate::io::FileIOBuilder;
use crate::spec::stats::BinaryTableStats;
use crate::spec::{
- BinaryRow, DataFileMeta, DataType, Datum, IntType, Predicate,
PredicateBuilder, VarCharType,
+ is_variant_extraction_row_type, variant_extraction_row, BigIntType,
BinaryRow,
+ DataFileMeta, DataType, Datum, IntType, Predicate, PredicateBuilder,
RowType, VarCharType,
};
use crate::table::source::DataSplitBuilder;
+ use crate::variant::variant_shredding_type;
use arrow_array::{Int32Array, StringArray};
use futures::TryStreamExt;
@@ -631,6 +686,44 @@ mod row_tests {
}
}
+ #[test]
+ fn read_data_fields_preserves_variant_extraction_row_type() {
+ let configured = DataType::Row(RowType::new(vec![field(
+ 0,
+ "age",
+ DataType::Int(IntType::new()),
+ )]));
+ let physical_type = variant_shredding_type(&configured).unwrap();
+ let data_field = field(1, "v", physical_type);
+ let extraction_type = DataType::Row(variant_extraction_row(
+ true,
+ vec![(
+ DataType::Int(IntType::new()),
+ "$.age".to_string(),
+ true,
+ "UTC".to_string(),
+ )],
+ ));
+ let expected_field = field(1, "v", extraction_type.clone());
+
+ let read_fields = read_data_fields(&[data_field],
&[expected_field]).unwrap();
+
+ assert_eq!(read_fields.len(), 1);
+ assert!(is_variant_extraction_row_type(read_fields[0].data_type()));
+ assert_eq!(read_fields[0].data_type(), &extraction_type);
+ }
+
+ #[test]
+ fn read_data_fields_uses_physical_type_for_castable_field() {
+ let data_field = field(1, "n", DataType::Int(IntType::new()));
+ let expected_field = field(1, "n",
DataType::BigInt(BigIntType::new()));
+
+ let read_fields = read_data_fields(&[data_field],
&[expected_field]).unwrap();
+
+ assert_eq!(read_fields.len(), 1);
+ assert_eq!(read_fields[0].data_type(), &DataType::Int(IntType::new()));
+ }
+
#[tokio::test]
async fn row_projection_reads_full_file_schema_before_projecting() {
let fields = vec![
diff --git a/crates/paimon/src/table/full_text_search_builder.rs
b/crates/paimon/src/table/full_text_search_builder.rs
index 616d1ed..0b92674 100644
--- a/crates/paimon/src/table/full_text_search_builder.rs
+++ b/crates/paimon/src/table/full_text_search_builder.rs
@@ -317,7 +317,7 @@ async fn read_raw_full_text_search(
let mut read_builder = table.new_read_builder();
read_builder
- .with_projection(&[search.field_name.as_str(), ROW_ID_FIELD_NAME])
+ .with_projection(&[search.field_name.as_str(), ROW_ID_FIELD_NAME])?
.with_row_ranges(raw_ranges.to_vec());
let plan = read_builder.new_scan().plan().await?;
if plan.splits().is_empty() {
diff --git a/crates/paimon/src/table/lumina_index_build_builder.rs
b/crates/paimon/src/table/lumina_index_build_builder.rs
index ea5c126..b6c06c1 100644
--- a/crates/paimon/src/table/lumina_index_build_builder.rs
+++ b/crates/paimon/src/table/lumina_index_build_builder.rs
@@ -600,7 +600,7 @@ async fn extract_vectors(
.build()?;
let mut read_builder = table.new_read_builder();
- read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME]);
+ read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?;
let read = read_builder.new_read()?;
let batches = read.to_arrow(&[split])?.try_collect::<Vec<_>>().await?;
extract_vectors_from_batches(
diff --git a/crates/paimon/src/table/read_builder.rs
b/crates/paimon/src/table/read_builder.rs
index 8cc66e3..1d54b8d 100644
--- a/crates/paimon/src/table/read_builder.rs
+++ b/crates/paimon/src/table/read_builder.rs
@@ -110,7 +110,7 @@ fn normalize_filter(table: &Table, filter: Predicate) ->
NormalizedFilter {
#[derive(Debug, Clone)]
pub struct ReadBuilder<'a> {
table: &'a Table,
- projected_fields: Option<Vec<String>>,
+ read_type: Option<Vec<DataField>>,
filter: NormalizedFilter,
limit: Option<usize>,
row_ranges: Option<Vec<RowRange>>,
@@ -120,7 +120,7 @@ impl<'a> ReadBuilder<'a> {
pub(crate) fn new(table: &'a Table) -> Self {
Self {
table,
- projected_fields: None,
+ read_type: None,
filter: NormalizedFilter::default(),
limit: None,
row_ranges: None,
@@ -128,10 +128,18 @@ impl<'a> ReadBuilder<'a> {
}
/// Set column projection by name. Output order follows the
caller-specified order.
- /// Unknown or duplicate names cause `new_read()` to fail; an empty list
is a valid
+ /// Unknown or duplicate names cause this method to fail; an empty list is
a valid
/// zero-column projection.
- pub fn with_projection(&mut self, columns: &[&str]) -> &mut Self {
- self.projected_fields = Some(columns.iter().map(|c|
(*c).to_string()).collect());
+ pub fn with_projection(&mut self, columns: &[&str]) -> Result<&mut Self> {
+ let projection_names = columns.iter().map(|c|
(*c).to_string()).collect::<Vec<_>>();
+ self.read_type =
Some(self.resolve_projected_fields(&projection_names)?);
+ Ok(self)
+ }
+
+ /// Set the full read type, including nested field pruning or
connector-defined
+ /// logical read types such as Variant extractions.
+ pub fn with_read_type(&mut self, read_type: Vec<DataField>) -> &mut Self {
+ self.read_type = Some(read_type);
self
}
@@ -224,7 +232,7 @@ impl<'a> ReadBuilder<'a> {
self.limit,
self.row_ranges.clone(),
)
- .with_projection(self.projected_fields.clone())
+
.with_projected_read_field_ids(projected_read_field_ids(&self.read_type))
}
/// Create a table read for consuming splits (e.g. from a scan plan).
@@ -232,9 +240,9 @@ impl<'a> ReadBuilder<'a> {
// Fail closed at read construction so bindings that short-circuit
before
// `to_arrow` (e.g. an empty-splits fast path) can't bypass the guard.
CoreOptions::new(self.table.schema.options()).ensure_read_authorized()?;
- let read_type = match &self.projected_fields {
+ let read_type = match &self.read_type {
None => self.table.schema.fields().to_vec(),
- Some(projected) => self.resolve_projected_fields(projected)?,
+ Some(fields) => fields.clone(),
};
// Pass the FULL data predicate through (including `And`/`Or`/`Not`).
@@ -249,12 +257,12 @@ impl<'a> ReadBuilder<'a> {
pub(super) fn resolve_projected_fields(
&self,
- projected_fields: &[String],
+ projection_names: &[String],
) -> Result<Vec<DataField>> {
resolve_projected_fields(
self.table.identifier().full_name(),
self.table.schema.fields(),
- projected_fields,
+ projection_names,
)
}
}
@@ -262,19 +270,19 @@ impl<'a> ReadBuilder<'a> {
pub(super) fn resolve_projected_fields(
full_name: String,
fields: &[DataField],
- projected_fields: &[String],
+ projection_names: &[String],
) -> Result<Vec<DataField>> {
- if projected_fields.is_empty() {
+ if projection_names.is_empty() {
return Ok(Vec::new());
}
let field_map: HashMap<&str, &DataField> =
fields.iter().map(|field| (field.name(), field)).collect();
- let mut seen = HashSet::with_capacity(projected_fields.len());
- let mut resolved = Vec::with_capacity(projected_fields.len());
+ let mut seen = HashSet::with_capacity(projection_names.len());
+ let mut resolved = Vec::with_capacity(projection_names.len());
- for name in projected_fields {
+ for name in projection_names {
if !seen.insert(name.as_str()) {
return Err(Error::ConfigInvalid {
message: format!("Duplicate projection column '{name}' for
table {full_name}"),
@@ -302,21 +310,18 @@ pub(super) fn resolve_projected_fields(
Ok(resolved)
}
-pub(super) fn projected_read_field_ids(
- full_name: String,
- fields: &[DataField],
- projected_fields: Option<&Vec<String>>,
-) -> Result<Option<HashSet<i32>>> {
- let Some(projected) = projected_fields else {
- return Ok(None);
- };
- let fields = resolve_projected_fields(full_name, fields, projected)?;
- let field_ids = fields
- .into_iter()
+pub(super) fn projected_read_field_ids_from_fields(fields: &[DataField]) ->
HashSet<i32> {
+ fields
+ .iter()
.filter(|field| !is_system_projection_field(field.id()))
.map(|field| field.id())
- .collect::<HashSet<_>>();
- Ok(Some(field_ids))
+ .collect::<HashSet<_>>()
+}
+
+fn projected_read_field_ids(read_type: &Option<Vec<DataField>>) ->
Option<HashSet<i32>> {
+ read_type
+ .as_ref()
+ .map(|fields| projected_read_field_ids_from_fields(fields))
}
pub(super) fn is_system_projection_field(field_id: i32) -> bool {
@@ -339,7 +344,8 @@ mod tests {
use crate::catalog::Identifier;
use crate::io::FileIOBuilder;
use crate::spec::{
- BinaryRow, DataType, IntType, Predicate, PredicateBuilder, Schema,
TableSchema, VarCharType,
+ BinaryRow, DataField, DataType, IntType, Predicate, PredicateBuilder,
Schema, TableSchema,
+ VarCharType,
};
use crate::table::{query_auth_table, DataSplitBuilder, Table};
use arrow_array::{Int32Array, RecordBatch};
@@ -432,43 +438,37 @@ mod tests {
#[test]
fn test_projected_read_field_ids_uses_projection_ids() {
- let table = simple_table();
- let projected = vec!["id".to_string()];
+ let read_type = vec![DataField::new(
+ 1,
+ "id".to_string(),
+ DataType::Int(IntType::new()),
+ )];
assert_eq!(
- super::projected_read_field_ids(
- table.identifier().full_name(),
- table.schema().fields(),
- Some(&projected),
- )
- .unwrap(),
- Some(HashSet::from([1]))
+ super::projected_read_field_ids_from_fields(&read_type),
+ HashSet::from([1])
);
}
#[test]
fn test_projected_read_field_ids_ignores_system_only_projection() {
- let table = simple_table();
- let projected = vec![crate::spec::ROW_ID_FIELD_NAME.to_string()];
+ let read_type = vec![DataField::new(
+ crate::spec::ROW_ID_FIELD_ID,
+ crate::spec::ROW_ID_FIELD_NAME.to_string(),
+ DataType::Int(IntType::new()),
+ )];
assert_eq!(
- super::projected_read_field_ids(
- table.identifier().full_name(),
- table.schema().fields(),
- Some(&projected),
- )
- .unwrap(),
- Some(HashSet::new())
+ super::projected_read_field_ids_from_fields(&read_type),
+ HashSet::new()
);
}
- #[tokio::test]
- async fn test_new_scan_validates_unknown_projection() {
+ #[test]
+ fn test_with_projection_validates_unknown_projection() {
let table = simple_table();
let mut builder = ReadBuilder::new(&table);
- builder.with_projection(&["missing"]);
-
- let err = builder.new_scan().plan().await.unwrap_err();
+ let err = builder.with_projection(&["missing"]).unwrap_err();
assert!(matches!(
err,
@@ -479,13 +479,11 @@ mod tests {
));
}
- #[tokio::test]
- async fn test_new_scan_validates_duplicate_projection() {
+ #[test]
+ fn test_with_projection_validates_duplicate_projection() {
let table = simple_table();
let mut builder = ReadBuilder::new(&table);
- builder.with_projection(&["id", "id"]);
-
- let err = builder.new_scan().plan().await.unwrap_err();
+ let err = builder.with_projection(&["id", "id"]).unwrap_err();
assert!(matches!(
err,
@@ -565,7 +563,10 @@ mod tests {
.unwrap();
let mut builder = table.new_read_builder();
- builder.with_projection(&["id"]).with_filter(predicate);
+ builder
+ .with_projection(&["id"])
+ .unwrap()
+ .with_filter(predicate);
let read = builder.new_read().unwrap();
let batches = read
.to_arrow(&[split])
@@ -681,7 +682,10 @@ mod tests {
.unwrap();
let mut builder = table.new_read_builder();
- builder.with_projection(&["id"]).with_filter(predicate);
+ builder
+ .with_projection(&["id"])
+ .unwrap()
+ .with_filter(predicate);
let read = builder.new_read().unwrap();
let batches = read
.to_arrow(&[split])
@@ -747,7 +751,10 @@ mod tests {
]);
let mut builder = table.new_read_builder();
- builder.with_projection(&["id"]).with_filter(predicate);
+ builder
+ .with_projection(&["id"])
+ .unwrap()
+ .with_filter(predicate);
let read = builder.new_read().unwrap();
let batches = read
.to_arrow(&[split])
@@ -812,7 +819,10 @@ mod tests {
]);
let mut builder = table.new_read_builder();
- builder.with_projection(&["id"]).with_filter(predicate);
+ builder
+ .with_projection(&["id"])
+ .unwrap()
+ .with_filter(predicate);
let read = builder.new_read().unwrap();
let batches = read
.to_arrow(&[split])
diff --git a/crates/paimon/src/table/table_scan.rs
b/crates/paimon/src/table/table_scan.rs
index df59954..cab4781 100644
--- a/crates/paimon/src/table/table_scan.rs
+++ b/crates/paimon/src/table/table_scan.rs
@@ -572,7 +572,7 @@ pub struct TableScan<'a> {
/// Used by non-read paths (overwrite, truncate, writer restore) that need
/// the complete file set. Normal read scans leave this as `false`.
scan_all_files: bool,
- projected_fields: Option<Vec<String>>,
+ projected_read_field_ids: Option<HashSet<i32>>,
}
impl<'a> TableScan<'a> {
@@ -592,7 +592,7 @@ impl<'a> TableScan<'a> {
limit,
row_ranges,
scan_all_files: false,
- projected_fields: None,
+ projected_read_field_ids: None,
}
}
@@ -602,7 +602,7 @@ impl<'a> TableScan<'a> {
/// the complete file set regardless of merge engine or DV settings.
pub fn with_scan_all_files(mut self) -> Self {
self.scan_all_files = true;
- self.projected_fields = None;
+ self.projected_read_field_ids = None;
self
}
@@ -619,8 +619,11 @@ impl<'a> TableScan<'a> {
self
}
- pub(crate) fn with_projection(mut self, projected_fields:
Option<Vec<String>>) -> Self {
- self.projected_fields = projected_fields;
+ pub(super) fn with_projected_read_field_ids(
+ mut self,
+ projected_read_field_ids: Option<HashSet<i32>>,
+ ) -> Self {
+ self.projected_read_field_ids = projected_read_field_ids;
self
}
@@ -681,11 +684,7 @@ impl<'a> TableScan<'a> {
}
fn projected_read_field_ids(&self) -> crate::Result<Option<HashSet<i32>>> {
- super::read_builder::projected_read_field_ids(
- self.table.identifier().full_name(),
- self.table.schema().fields(),
- self.projected_fields.as_ref(),
- )
+ Ok(self.projected_read_field_ids.clone())
}
async fn resolve_snapshot(&self) -> crate::Result<Option<Snapshot>> {
diff --git a/crates/paimon/src/table/vector_search_builder.rs
b/crates/paimon/src/table/vector_search_builder.rs
index fc0d2b7..33003c4 100644
--- a/crates/paimon/src/table/vector_search_builder.rs
+++ b/crates/paimon/src/table/vector_search_builder.rs
@@ -652,7 +652,7 @@ async fn read_raw_batch_vector_search(
let mut read_builder = table.new_read_builder();
read_builder
- .with_projection(&[field_name.as_str(), ROW_ID_FIELD_NAME])
+ .with_projection(&[field_name.as_str(), ROW_ID_FIELD_NAME])?
.with_row_ranges(raw_ranges.to_vec());
let plan = read_builder.new_scan().plan().await?;
if plan.splits().is_empty() {
diff --git a/crates/paimon/src/variant.rs b/crates/paimon/src/variant.rs
index 7d4d869..00f442e 100644
--- a/crates/paimon/src/variant.rs
+++ b/crates/paimon/src/variant.rs
@@ -2601,6 +2601,184 @@ 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 cast_variant_to_extraction_value(variant, &scalar) {
+ Some(value) => Ok(Some(value)),
+ None if fail_on_error => Err(Error::DataInvalid {
+ message: format!("Cannot cast Variant value to {data_type:?}"),
+ source: None,
+ }),
+ None => Ok(None),
+ }
+}
+
+fn cast_variant_to_extraction_value(
+ variant: VariantRef<'_>,
+ scalar: &VariantScalarSchema,
+) -> Option<ShreddedValue> {
+ match scalar {
+ VariantScalarSchema::Boolean =>
cast_variant_to_bool(variant).map(ShreddedValue::Boolean),
+ VariantScalarSchema::Int8 => cast_variant_to_i64(variant)
+ .and_then(|value| i8::try_from(value).ok())
+ .map(ShreddedValue::Int8),
+ VariantScalarSchema::Int16 => cast_variant_to_i64(variant)
+ .and_then(|value| i16::try_from(value).ok())
+ .map(ShreddedValue::Int16),
+ VariantScalarSchema::Int32 => cast_variant_to_i64(variant)
+ .and_then(|value| i32::try_from(value).ok())
+ .map(ShreddedValue::Int32),
+ VariantScalarSchema::Int64 =>
cast_variant_to_i64(variant).map(ShreddedValue::Int64),
+ VariantScalarSchema::Float32 => {
+ cast_variant_to_f64(variant).map(|value|
ShreddedValue::Float32(value as f32))
+ }
+ VariantScalarSchema::Float64 =>
cast_variant_to_f64(variant).map(ShreddedValue::Float64),
+ VariantScalarSchema::Decimal { precision, scale } => {
+ cast_variant_to_decimal(variant, *precision,
*scale).map(ShreddedValue::Decimal128)
+ }
+ VariantScalarSchema::String =>
cast_variant_to_string(variant).map(ShreddedValue::String),
+ VariantScalarSchema::Binary
+ | VariantScalarSchema::Date32
+ | VariantScalarSchema::Timestamp
+ | VariantScalarSchema::TimestampNtz => try_typed_shred(variant,
scalar).ok().flatten(),
+ }
+}
+
+fn cast_variant_to_bool(variant: VariantRef<'_>) -> Option<bool> {
+ match variant.kind().ok()? {
+ VariantKind::Boolean => variant.get_boolean().ok(),
+ VariantKind::String => match
variant.get_string().ok()?.to_ascii_lowercase().as_str() {
+ "true" => Some(true),
+ "false" => Some(false),
+ _ => None,
+ },
+ _ => None,
+ }
+}
+
+fn cast_variant_to_i64(variant: VariantRef<'_>) -> Option<i64> {
+ match variant.kind().ok()? {
+ VariantKind::Long
+ | VariantKind::Date
+ | VariantKind::Timestamp
+ | VariantKind::TimestampNtz => variant.get_long().ok(),
+ VariantKind::String => variant.get_string().ok()?.parse::<i64>().ok(),
+ VariantKind::Decimal => {
+ let decimal = variant.get_decimal().ok()?;
+ rescale_decimal_exact(decimal.unscaled, decimal.scale, 0)
+ .and_then(|value| i64::try_from(value).ok())
+ }
+ _ => None,
+ }
+}
+
+fn cast_variant_to_f64(variant: VariantRef<'_>) -> Option<f64> {
+ match variant.kind().ok()? {
+ VariantKind::Long
+ | VariantKind::Date
+ | VariantKind::Timestamp
+ | VariantKind::TimestampNtz => Some(variant.get_long().ok()? as f64),
+ VariantKind::Double => variant.get_double().ok(),
+ VariantKind::Float => Some(variant.get_float().ok()? as f64),
+ VariantKind::Decimal => {
+ let decimal = variant.get_decimal().ok()?;
+ Some(decimal.unscaled as f64 / 10f64.powi(decimal.scale as i32))
+ }
+ VariantKind::String => variant.get_string().ok()?.parse::<f64>().ok(),
+ _ => None,
+ }
+}
+
+fn cast_variant_to_string(variant: VariantRef<'_>) -> Option<String> {
+ match variant.kind().ok()? {
+ VariantKind::Object | VariantKind::Array => variant.to_json().ok(),
+ VariantKind::Boolean => Some(variant.get_boolean().ok()?.to_string()),
+ VariantKind::Long
+ | VariantKind::Date
+ | VariantKind::Timestamp
+ | VariantKind::TimestampNtz =>
Some(variant.get_long().ok()?.to_string()),
+ VariantKind::String => variant.get_string().ok(),
+ VariantKind::Double => Some(variant.get_double().ok()?.to_string()),
+ VariantKind::Decimal =>
Some(variant.get_decimal().ok()?.to_plain_string()),
+ VariantKind::Float => Some(variant.get_float().ok()?.to_string()),
+ _ => variant.to_json().ok(),
+ }
+}
+
+fn cast_variant_to_decimal(variant: VariantRef<'_>, precision: u8, scale: i8)
-> Option<i128> {
+ let unscaled = match variant.kind().ok()? {
+ VariantKind::Long
+ | VariantKind::Date
+ | VariantKind::Timestamp
+ | VariantKind::TimestampNtz => {
+ rescale_decimal_exact(variant.get_long().ok()? as i128, 0, scale)?
+ }
+ VariantKind::Decimal => {
+ let decimal = variant.get_decimal().ok()?;
+ rescale_decimal_exact(decimal.unscaled, decimal.scale, scale)?
+ }
+ VariantKind::String => {
+ let parsed = parse_decimal_string(&variant.get_string().ok()?)?;
+ rescale_decimal_exact(parsed.unscaled, parsed.scale, scale)?
+ }
+ _ => return None,
+ };
+ (decimal_precision(unscaled) <= precision).then_some(unscaled)
+}
+
+fn parse_decimal_string(input: &str) -> Option<VariantDecimal> {
+ let input = input.trim();
+ if input.is_empty() || input.contains(['e', 'E']) {
+ return None;
+ }
+ let negative = input.starts_with('-');
+ let unsigned = input.strip_prefix('-').unwrap_or(input);
+ if unsigned.is_empty()
+ || unsigned.matches('.').count() > 1
+ || !unsigned.bytes().all(|ch| ch == b'.' || ch.is_ascii_digit())
+ {
+ return None;
+ }
+ let scale = unsigned
+ .split_once('.')
+ .map(|(_, fraction)| fraction.len())
+ .unwrap_or(0);
+ let digits: String = unsigned
+ .bytes()
+ .filter(|ch| *ch != b'.')
+ .map(char::from)
+ .collect();
+ let significant = digits.trim_start_matches('0');
+ let precision = if significant.is_empty() {
+ 1
+ } else {
+ significant.len()
+ };
+ if precision > 38 || scale > 38 {
+ return None;
+ }
+ let mut unscaled = digits.parse::<i128>().ok()?;
+ if negative {
+ unscaled = -unscaled;
+ }
+ Some(VariantDecimal {
+ unscaled,
+ precision: precision as u8,
+ scale: scale as i8,
+ })
+}
+
fn rebuild_shredded_into(
row: &ShreddedRow,
metadata: &[u8],
diff --git a/docs/src/sql.md b/docs/src/sql.md
index 70c282b..6ae3b62 100644
--- a/docs/src/sql.md
+++ b/docs/src/sql.md
@@ -260,7 +260,7 @@ Current limitations:
- `schema_of_variant`, `schema_of_variant_agg`, `to_variant_object`,
`variant_explode`, and `variant_explode_outer` are not implemented yet.
- `variant_get` currently casts to scalar types and `VARIANT`. It does not yet
cast directly to `ARRAY`, `MAP`, or `STRUCT`.
-- Predicate pushdown is not applied through `variant_get`; DataFusion
evaluates Variant filters after reading rows.
+- Simple `variant_get` and `try_variant_get` expressions over a `VARIANT`
column, a literal path, and a scalar literal type can be pushed into scans as
Variant extraction fields for projections and filters. Predicate translation
through `variant_get` is still not applied to Paimon/Parquet statistics;
DataFusion evaluates those filters after reading the extracted field.
With a raw DataFusion `SessionContext`, register these scalar functions
explicitly: