liurenjie1024 commented on code in PR #1602: URL: https://github.com/apache/iceberg-rust/pull/1602#discussion_r2381620650
########## crates/integrations/datafusion/src/physical_plan/project.rs: ########## @@ -0,0 +1,494 @@ +// 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. + +//! Partition value projection for Iceberg tables. + +use std::sync::{Arc, Mutex}; + +use datafusion::arrow::array::{ArrayRef, RecordBatch, StructArray}; +use datafusion::arrow::datatypes::{DataType, Schema as ArrowSchema}; +use datafusion::common::Result as DFResult; +use datafusion::error::DataFusionError; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::{ColumnarValue, ExecutionPlan}; +use iceberg::arrow::record_batch_projector::RecordBatchProjector; +use iceberg::spec::{PartitionSpec, Schema}; +use iceberg::table::Table; + +use crate::to_datafusion_error; + +/// Column name for the combined partition values struct +const PARTITION_VALUES_COLUMN: &str = "_partition"; + +/// Extends an ExecutionPlan with partition value calculations for Iceberg tables. +/// +/// This function takes an input ExecutionPlan and extends it with an additional column +/// containing calculated partition values based on the table's partition specification. +/// For unpartitioned tables, returns the original plan unchanged. +/// +/// # Arguments +/// * `input` - The input ExecutionPlan to extend +/// * `table` - The Iceberg table with partition specification +/// +/// # Returns +/// * `Ok(Arc<dyn ExecutionPlan>)` - Extended plan with partition values column +/// * `Err` - If partition spec is not found or transformation fails +pub fn project_with_partition( + input: Arc<dyn ExecutionPlan>, + table: &Table, +) -> DFResult<Arc<dyn ExecutionPlan>> { + let metadata = table.metadata(); + let partition_spec = metadata.default_partition_spec(); + let table_schema = metadata.current_schema(); + + if partition_spec.is_unpartitioned() { + return Ok(input); + } + + let input_schema = input.schema(); + let partition_type = build_partition_type(partition_spec, table_schema.as_ref())?; + let calculator = PartitionValueCalculator::new( + partition_spec.as_ref().clone(), + table_schema.as_ref().clone(), + partition_type, + )?; + + let mut projection_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = + Vec::with_capacity(input_schema.fields().len() + 1); + + for (index, field) in input_schema.fields().iter().enumerate() { + let column_expr = Arc::new(Column::new(field.name(), index)); + projection_exprs.push((column_expr, field.name().clone())); + } + + let partition_expr = Arc::new(PartitionExpr::new(calculator)); + projection_exprs.push((partition_expr, PARTITION_VALUES_COLUMN.to_string())); + + let projection = ProjectionExec::try_new(projection_exprs, input)?; + Ok(Arc::new(projection)) +} + +/// PhysicalExpr implementation for partition value calculation +#[derive(Debug, Clone)] +struct PartitionExpr { + calculator: Arc<Mutex<PartitionValueCalculator>>, +} + +impl PartitionExpr { + fn new(calculator: PartitionValueCalculator) -> Self { + Self { + calculator: Arc::new(Mutex::new(calculator)), + } + } +} + +// Manual PartialEq/Eq implementations for pointer-based equality +// (two PartitionExpr are equal if they share the same calculator instance) +impl PartialEq for PartitionExpr { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.calculator, &other.calculator) + } +} + +impl Eq for PartitionExpr {} + +impl PhysicalExpr for PartitionExpr { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn data_type(&self, _input_schema: &ArrowSchema) -> DFResult<DataType> { + let calculator = self + .calculator + .lock() + .map_err(|e| DataFusionError::Internal(format!("Failed to lock calculator: {}", e)))?; + Ok(calculator.partition_type.clone()) + } + + fn nullable(&self, _input_schema: &ArrowSchema) -> DFResult<bool> { + Ok(false) + } + + fn evaluate(&self, batch: &RecordBatch) -> DFResult<ColumnarValue> { + let mut calculator = self + .calculator + .lock() + .map_err(|e| DataFusionError::Internal(format!("Failed to lock calculator: {}", e)))?; + let array = calculator.calculate(batch)?; + Ok(ColumnarValue::Array(array)) + } + + fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> { + vec![] + } + + fn with_new_children( + self: Arc<Self>, + _children: Vec<Arc<dyn PhysicalExpr>>, + ) -> DFResult<Arc<dyn PhysicalExpr>> { + Ok(self) + } + + fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if let Ok(calculator) = self.calculator.lock() { + let field_names: Vec<String> = calculator + .partition_spec + .fields() + .iter() + .map(|pf| format!("{}({})", pf.transform, pf.name)) + .collect(); + write!(f, "iceberg_partition_values[{}]", field_names.join(", ")) + } else { + write!(f, "iceberg_partition_values") + } + } +} + +impl std::fmt::Display for PartitionExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if let Ok(calculator) = self.calculator.lock() { + let field_names: Vec<&str> = calculator + .partition_spec + .fields() + .iter() + .map(|pf| pf.name.as_str()) + .collect(); + write!(f, "iceberg_partition_values({})", field_names.join(", ")) + } else { + write!(f, "iceberg_partition_values") + } + } +} + +impl std::hash::Hash for PartitionExpr { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + std::any::TypeId::of::<Self>().hash(state); + } +} + +/// Calculator for partition values in Iceberg tables +#[derive(Debug, Clone)] +struct PartitionValueCalculator { + partition_spec: PartitionSpec, + table_schema: Schema, + partition_type: DataType, + projector: Option<RecordBatchProjector>, +} + +impl PartitionValueCalculator { + fn new( + partition_spec: PartitionSpec, + table_schema: Schema, + partition_type: DataType, + ) -> DFResult<Self> { + if partition_spec.is_unpartitioned() { + return Err(DataFusionError::Internal( + "Cannot create partition calculator for unpartitioned table".to_string(), + )); + } + + Ok(Self { + partition_spec, + table_schema, + partition_type, + projector: None, + }) + } + + fn calculate(&mut self, batch: &RecordBatch) -> DFResult<ArrayRef> { + if self.projector.is_none() { + let source_field_ids: Vec<i32> = self + .partition_spec + .fields() + .iter() + .map(|pf| pf.source_id) + .collect(); + + let projector = RecordBatchProjector::from_iceberg_schema_mapping( Review Comment: We don't need first batch to get input's schema, see: https://github.com/apache/datafusion/blob/921f4a028409f71b68bed7d05a348255bb6f0fba/datafusion/physical-plan/src/execution_plan.rs#L106 ########## crates/iceberg/src/arrow/record_batch_projector.rs: ########## @@ -77,6 +79,80 @@ impl RecordBatchProjector { }) } + /// Create RecordBatchProjector using Iceberg schema for field mapping. + /// + /// This constructor is more flexible and works with any Arrow schema by using + /// the Iceberg schema to map field names to field IDs. + /// + /// # Arguments + /// * `original_schema` - The original Arrow schema (doesn't need field ID metadata) + /// * `iceberg_schema` - The Iceberg schema for field ID mapping + /// * `target_field_ids` - The field IDs to project + pub fn from_iceberg_schema_mapping( + original_schema: SchemaRef, + iceberg_schema: Arc<IcebergSchema>, + target_field_ids: &[i32], + ) -> Result<Self> { + let field_id_fetch_func = |field: &Field| -> Result<Option<i64>> { + // First try to get field ID from metadata (Parquet case) + if let Some(value) = field.metadata().get(PARQUET_FIELD_ID_META_KEY) { + let field_id = value.parse::<i32>().map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + "Failed to parse field id".to_string(), + ) + .with_context("value", value) + .with_source(e) + })?; + return Ok(Some(field_id as i64)); + } + + // Fallback: use Iceberg schema's built-in field lookup + if let Some(iceberg_field) = iceberg_schema.field_by_name(field.name()) { Review Comment: This is incorrect. Iceberg schema's `field_by_name` uses full path(including parent name), while this is not name. ########## crates/integrations/datafusion/src/physical_plan/project.rs: ########## @@ -15,125 +15,203 @@ // specific language governing permissions and limitations // under the License. -//! Utilities for calculating partition values for Iceberg tables. -//! -//! This module provides functions to calculate partition values from record batches -//! based on Iceberg partition specifications. These utilities are used when writing -//! data to partitioned Iceberg tables. +//! Partition value projection for Iceberg tables. use std::sync::Arc; use datafusion::arrow::array::{ArrayRef, RecordBatch, StructArray}; -use datafusion::arrow::datatypes::{ - DataType, Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, -}; +use datafusion::arrow::datatypes::{DataType, Schema as ArrowSchema}; use datafusion::common::Result as DFResult; use datafusion::error::DataFusionError; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::{ColumnarValue, ExecutionPlan}; use iceberg::spec::{PartitionSpec, Schema}; +use iceberg::table::Table; use crate::to_datafusion_error; /// Column name for the combined partition values struct -#[allow(dead_code)] -pub(crate) const PARTITION_VALUES_COLUMN: &str = "_iceberg_partition_values"; +const PARTITION_VALUES_COLUMN: &str = "_partition"; -/// Create an output schema by adding a single partition values struct column to the input schema. -/// Returns the original schema unchanged if the table is unpartitioned. +/// Extends an ExecutionPlan with partition value calculations for Iceberg tables. +/// +/// This function takes an input ExecutionPlan and extends it with an additional column +/// containing calculated partition values based on the table's partition specification. +/// For unpartitioned tables, returns the original plan unchanged. +/// +/// # Arguments +/// * `input` - The input ExecutionPlan to extend +/// * `table` - The Iceberg table with partition specification +/// +/// # Returns +/// * `Ok(Arc<dyn ExecutionPlan>)` - Extended plan with partition values column +/// * `Err` - If partition spec is not found or transformation fails #[allow(dead_code)] -pub(crate) fn create_schema_with_partition_columns( - input_schema: &ArrowSchema, - partition_spec: &PartitionSpec, - table_schema: &Schema, -) -> DFResult<ArrowSchemaRef> { +pub fn project_with_partition( + input: Arc<dyn ExecutionPlan>, + table: &Table, +) -> DFResult<Arc<dyn ExecutionPlan>> { + let metadata = table.metadata(); + let partition_spec = metadata + .partition_spec_by_id(metadata.default_partition_spec_id()) + .ok_or_else(|| DataFusionError::Internal("Default partition spec not found".to_string()))?; + let table_schema = metadata.current_schema(); + if partition_spec.is_unpartitioned() { - return Ok(Arc::new(input_schema.clone())); + return Ok(input); } - let mut fields: Vec<Arc<Field>> = input_schema.fields().to_vec(); + let input_schema = input.schema(); + let partition_type = build_partition_type(partition_spec, table_schema.as_ref())?; + let calculator = PartitionValueCalculator::new( + partition_spec.as_ref().clone(), + table_schema.as_ref().clone(), + partition_type, + ); - let partition_struct_type = partition_spec - .partition_type(table_schema) - .map_err(to_datafusion_error)?; + let mut projection_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = Vec::new(); - let arrow_struct_type = - iceberg::arrow::type_to_arrow_type(&iceberg::spec::Type::Struct(partition_struct_type)) - .map_err(to_datafusion_error)?; + for (index, field) in input_schema.fields().iter().enumerate() { + let column_expr = Arc::new(Column::new(field.name(), index)); + projection_exprs.push((column_expr, field.name().clone())); + } - fields.push(Arc::new(Field::new( - PARTITION_VALUES_COLUMN, - arrow_struct_type, - false, // Partition values are generally not null - ))); + let partition_expr = Arc::new(PartitionExpr::new(calculator)); + projection_exprs.push((partition_expr, PARTITION_VALUES_COLUMN.to_string())); - Ok(Arc::new(ArrowSchema::new(fields))) + let projection = ProjectionExec::try_new(projection_exprs, input)?; + Ok(Arc::new(projection)) } -/// Calculate partition values for a record batch and return as a single struct array. -/// Returns None if the table is unpartitioned. -/// -/// # Arguments -/// * `batch` - The record batch to calculate partition values for -/// * `partition_spec` - The partition specification defining the partition fields -/// * `table_schema` - The Iceberg table schema -/// * `expected_partition_type` - The expected Arrow struct type for the partition values -#[allow(dead_code)] -pub(crate) fn calculate_partition_values( - batch: &RecordBatch, - partition_spec: &PartitionSpec, - table_schema: &Schema, - expected_partition_type: &DataType, -) -> DFResult<Option<ArrayRef>> { - if partition_spec.is_unpartitioned() { - return Ok(None); +/// PhysicalExpr implementation for partition value calculation +#[derive(Debug, Clone, PartialEq, Eq)] +struct PartitionExpr { + calculator: PartitionValueCalculator, +} + +impl PartitionExpr { + fn new(calculator: PartitionValueCalculator) -> Self { + Self { calculator } + } +} + +impl PhysicalExpr for PartitionExpr { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn data_type(&self, _input_schema: &ArrowSchema) -> DFResult<DataType> { + Ok(self.calculator.partition_type.clone()) + } + + fn nullable(&self, _input_schema: &ArrowSchema) -> DFResult<bool> { + Ok(false) + } + + fn evaluate(&self, batch: &RecordBatch) -> DFResult<ColumnarValue> { + let array = self.calculator.calculate(batch)?; + Ok(ColumnarValue::Array(array)) + } + + fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> { + vec![] } - let batch_schema = batch.schema(); - let mut partition_values = Vec::with_capacity(partition_spec.fields().len()); + fn with_new_children( + self: Arc<Self>, + _children: Vec<Arc<dyn PhysicalExpr>>, + ) -> DFResult<Arc<dyn PhysicalExpr>> { + Ok(self) + } + + fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "partition_values") + } +} + +impl std::fmt::Display for PartitionExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "partition_values") + } +} + +impl std::hash::Hash for PartitionExpr { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + std::any::TypeId::of::<Self>().hash(state); + } +} + +/// Calculator for partition values in Iceberg tables +#[derive(Debug, Clone, PartialEq, Eq)] +struct PartitionValueCalculator { + partition_spec: PartitionSpec, + table_schema: Schema, + partition_type: DataType, +} - let expected_struct_fields = match expected_partition_type { - DataType::Struct(fields) => fields.clone(), - _ => { +impl PartitionValueCalculator { + fn new(partition_spec: PartitionSpec, table_schema: Schema, partition_type: DataType) -> Self { + Self { + partition_spec, + table_schema, + partition_type, + } + } + + fn calculate(&self, batch: &RecordBatch) -> DFResult<ArrayRef> { + if self.partition_spec.is_unpartitioned() { return Err(DataFusionError::Internal( - "Expected partition type must be a struct".to_string(), + "Cannot calculate partition values for unpartitioned table".to_string(), )); } - }; - for pf in partition_spec.fields() { - let source_field = table_schema.field_by_id(pf.source_id).ok_or_else(|| { - DataFusionError::Internal(format!( - "Source field not found with id {} when calculating partition values", - pf.source_id - )) - })?; + let batch_schema = batch.schema(); + let mut partition_values = Vec::with_capacity(self.partition_spec.fields().len()); + + let expected_struct_fields = match &self.partition_type { + DataType::Struct(fields) => fields.clone(), + _ => { + return Err(DataFusionError::Internal( + "Expected partition type must be a struct".to_string(), + )); + } + }; + + for pf in self.partition_spec.fields() { + let source_field = self.table_schema.field_by_id(pf.source_id).ok_or_else(|| { + DataFusionError::Internal(format!( + "Source field not found with id {} when calculating partition values", + pf.source_id + )) + })?; - let field_path = find_field_path(table_schema, source_field.id)?; - let index_path = resolve_arrow_index_path(batch_schema.as_ref(), &field_path)?; + let field_path = find_field_path(&self.table_schema, source_field.id)?; + let index_path = resolve_arrow_index_path(batch_schema.as_ref(), &field_path)?; - let source_column = extract_column_by_index_path(batch, &index_path)?; + let source_column = extract_column_by_index_path(batch, &index_path)?; - let transform_fn = iceberg::transform::create_transform_function(&pf.transform) - .map_err(to_datafusion_error)?; - let partition_value = transform_fn - .transform(source_column) - .map_err(to_datafusion_error)?; + let transform_fn = iceberg::transform::create_transform_function(&pf.transform) Review Comment: This is not resolved, we could create trnasform functions in constructor. ########## crates/iceberg/src/arrow/record_batch_projector.rs: ########## @@ -77,6 +79,80 @@ impl RecordBatchProjector { }) } + /// Create RecordBatchProjector using Iceberg schema for field mapping. + /// + /// This constructor is more flexible and works with any Arrow schema by using + /// the Iceberg schema to map field names to field IDs. + /// + /// # Arguments + /// * `original_schema` - The original Arrow schema (doesn't need field ID metadata) + /// * `iceberg_schema` - The Iceberg schema for field ID mapping + /// * `target_field_ids` - The field IDs to project + pub fn from_iceberg_schema_mapping( + original_schema: SchemaRef, + iceberg_schema: Arc<IcebergSchema>, + target_field_ids: &[i32], + ) -> Result<Self> { + let field_id_fetch_func = |field: &Field| -> Result<Option<i64>> { Review Comment: This method is unnecessarily to be too complicated, following approach could simplify this: 1. Create an arrow schema using iceberg schema. 2. Prune arrow schema created in step 1 by matching top level field name. Then we can pass pruned arrow schema to original constructor, and search field by `PARQUET_FIELD_ID_META_KEY` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
