feat: push down variant extractions in DataFusion (#460)

Author: JingsongLi

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

bindings/python/src/read.rs

@@ -69,17 +69,18 @@ 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);
     }
     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.

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

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");
 

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()

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;

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,19 +58,22 @@ 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 {
     #[allow(clippy::too_many_arguments)]
     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();