liurenjie1024 commented on code in PR #1602:
URL: https://github.com/apache/iceberg-rust/pull/1602#discussion_r2378331183
##########
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)]
Review Comment:
Do we still need this for pub function?
##########
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() {
Review Comment:
Should we move this to constructor?
##########
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")
Review Comment:
```suggestion
write!(f, "iceberg_partition_values")
```
Should we consider adding partition spec?
##########
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
Review Comment:
How about using
https://github.com/apache/iceberg-rust/blob/55cc6c32106e1551e09d178e0ca050aab1e6f736/crates/iceberg/src/spec/table_metadata.rs#L342
?
##########
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();
Review Comment:
nit: `Vec::with_capacity`
##########
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)?;
Review Comment:
We don't need to to them for every batch.
##########
crates/integrations/datafusion/src/physical_plan/project.rs:
##########
@@ -0,0 +1,511 @@
+// 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.
+
+//! 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.
+
+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::common::Result as DFResult;
+use datafusion::error::DataFusionError;
+use iceberg::spec::{PartitionSpec, Schema};
+
+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";
+
+/// 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.
+#[allow(dead_code)]
+pub(crate) fn create_schema_with_partition_columns(
+ input_schema: &ArrowSchema,
+ partition_spec: &PartitionSpec,
+ table_schema: &Schema,
+) -> DFResult<ArrowSchemaRef> {
+ if partition_spec.is_unpartitioned() {
+ return Ok(Arc::new(input_schema.clone()));
+ }
+
+ let mut fields: Vec<Arc<Field>> = input_schema.fields().to_vec();
+
+ let partition_struct_type = partition_spec
+ .partition_type(table_schema)
+ .map_err(to_datafusion_error)?;
+
+ let arrow_struct_type =
+
iceberg::arrow::type_to_arrow_type(&iceberg::spec::Type::Struct(partition_struct_type))
+ .map_err(to_datafusion_error)?;
+
+ fields.push(Arc::new(Field::new(
+ PARTITION_VALUES_COLUMN,
+ arrow_struct_type,
+ false, // Partition values are generally not null
+ )));
+
+ Ok(Arc::new(ArrowSchema::new(fields)))
+}
+
+/// 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);
+ }
+
+ let batch_schema = batch.schema();
+ let mut partition_values =
Vec::with_capacity(partition_spec.fields().len());
+
+ let expected_struct_fields = match expected_partition_type {
+ DataType::Struct(fields) => fields.clone(),
+ _ => {
+ return Err(DataFusionError::Internal(
+ "Expected partition type must be a struct".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 field_path = find_field_path(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 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)?;
+
+ partition_values.push(partition_value);
+ }
+
+ let struct_array = StructArray::try_new(
+ expected_struct_fields,
+ partition_values,
+ None, // No null buffer for the struct array itself
+ )
+ .map_err(|e| DataFusionError::ArrowError(e, None))?;
+
+ Ok(Some(Arc::new(struct_array)))
+}
+
+/// Extract a column from a record batch by following an index path.
+/// The index path specifies the column indices to traverse for nested
structures.
+#[allow(dead_code)]
+fn extract_column_by_index_path(batch: &RecordBatch, index_path: &[usize]) ->
DFResult<ArrayRef> {
Review Comment:
I'm not convinced. There are two ways to solve your issue:
1. Add a constructor in `RecordBatchProjector` to accept iceberg schema and
target field ids.
2. Convert iceberg schema to arrow schema, the converter will add `field_id`
metadata.
Personally I prefer approach 1, but I don't have a strong opinion about.
After using `RecordBatchProjector`, the whole pr could be simplified a lot.
##########
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:
Ditto, this only needs to be done once.
--
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]
