mbutrovich commented on code in PR #2668: URL: https://github.com/apache/iceberg-rust/pull/2668#discussion_r3500804903
########## crates/iceberg/src/partitioning.rs: ########## @@ -0,0 +1,94 @@ +// 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 type utilities for Iceberg tables. + +use crate::spec::{NestedField, NestedFieldRef, PartitionSpec, Schema, StructType, Transform}; +use crate::{Error, ErrorKind, Result}; + +/// Computes the unified partition type across all partition specs in the table. +/// +/// This is equivalent to Java's `Partitioning.partitionType(table)`. The result is a +/// StructType containing all partition fields ever used across all specs, enabling correct +/// representation of the `_partition` metadata column when partition evolution has occurred. +/// +/// Matches Java's behavior: +/// - Specs are sorted by spec_id in descending order (newer specs first), so newer field +/// names take precedence when deduplicating by field_id. +/// - Void and unknown transform fields are skipped. +/// - Fields are deduplicated by field_id — each unique field_id appears exactly once. +/// +/// # Arguments +/// * `partition_specs` - Iterator over all partition specs in the table +/// * `schema` - The current table schema (needed to determine result types of transforms) +pub fn compute_unified_partition_type<'a>( + partition_specs: impl Iterator<Item = &'a PartitionSpec>, + schema: &Schema, +) -> Result<StructType> { + let mut seen_field_ids = std::collections::HashSet::new(); + let mut struct_fields: Vec<NestedFieldRef> = Vec::new(); + + // Sort specs by spec_id descending (newer first) to match Java's behavior: + // newer field names take precedence when deduplicating by field_id. + let mut specs: Vec<&PartitionSpec> = partition_specs.collect(); + specs.sort_by_key(|s| std::cmp::Reverse(s.spec_id())); + + for spec in specs { + for field in spec.fields() { + if seen_field_ids.contains(&field.field_id) { + continue; + } + + // Skip void transforms (dropped partition columns) + if matches!(field.transform, Transform::Void) { + continue; + } + + // Reject unknown transforms — the table uses a spec feature that + // this version of iceberg-rust doesn't support. Matching Java's + // behavior where getResultType() throws for unknown transforms. + if matches!(field.transform, Transform::Unknown) { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + format!( + "Partition field '{}' uses an unknown transform that is not \ + supported by this version of iceberg-rust", + field.name + ), + )); + } + + seen_field_ids.insert(field.field_id); + + let source_field = schema.field_by_id(field.source_id).ok_or_else(|| { Review Comment: Java's `Partitioning.partitionType` only considers `allActiveFieldIds`, which filters out fields whose source column is gone (`schema.findField(field.sourceId()) != null` in `Partitioning.java`). Here we return `ErrorKind::Unexpected` instead. If I'm reading both right, a table that dropped a column it used to partition on (allowed under https://iceberg.apache.org/spec/#partition-evolution) would error the whole scan when `_partition` is projected. Would `continue`-ing on a missing source column match Java's intent better? A test that drops a partition source column would settle it either way. ########## crates/iceberg/src/arrow/value.rs: ########## @@ -909,6 +909,38 @@ pub(crate) fn create_primitive_array_repeated( let vals: Vec<Option<&[u8]>> = vec![None; num_rows]; Arc::new(BinaryArray::from_opt_vec(vals)) } + (DataType::LargeBinary, Some(PrimitiveLiteral::Binary(value))) => { Review Comment: @advancedxy already questioned whether these arms are related; you explained they're needed for `_partition` struct children and logged the Binary-to-LargeBinary mapping as #2698. This is a separate, reuse-only angle on the same lines, not a relevance question. I'd push for using arrow-rs here rather than growing our own array-construction match. The hand-rolled arms are duplicate code we own forever: every Arrow type a partition column can take is another arm, and a missing one is a runtime error (the `(dt, _)` catch-all), not a compile error. arrow-rs already maintains correct, exhaustive versions: - `arrow_array::new_null_array(dt, len)` builds an all-null array for any `DataType`, recursively for `Struct` (`arrow-array/src/array/mod.rs:1020`; struct path `struct_array.rs:192`). This replaces every `(dt, None) => ...` arm and the `(DataType::Struct(fields), None)` arm. - `PrimitiveArray::from_value(v, n)` (`arrow-array/src/array/primitive_array.rs:808`) for the primitive `Some` arms. - `GenericByteArray::new_repeated(v, n)` (`arrow-array/src/array/byte_array.rs:198`) for String/Binary/LargeBinary. Delegating to these shrinks the match substantially and means new partition types are covered by arrow-rs instead of by us adding (and testing) another arm. If there's a reason to keep the explicit arms (e.g. a type arrow-rs handles differently), worth a comment saying so; otherwise I'd lean on the library. A test driving `_partition` through decimal/timestamptz/uuid would also turn any remaining gap into a test failure rather than a read-time error. (Could fold into #2698.) ########## crates/iceberg/src/partitioning.rs: ########## @@ -0,0 +1,94 @@ +// 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 type utilities for Iceberg tables. + +use crate::spec::{NestedField, NestedFieldRef, PartitionSpec, Schema, StructType, Transform}; +use crate::{Error, ErrorKind, Result}; + +/// Computes the unified partition type across all partition specs in the table. +/// +/// This is equivalent to Java's `Partitioning.partitionType(table)`. The result is a +/// StructType containing all partition fields ever used across all specs, enabling correct +/// representation of the `_partition` metadata column when partition evolution has occurred. +/// +/// Matches Java's behavior: +/// - Specs are sorted by spec_id in descending order (newer specs first), so newer field +/// names take precedence when deduplicating by field_id. +/// - Void and unknown transform fields are skipped. +/// - Fields are deduplicated by field_id — each unique field_id appears exactly once. +/// +/// # Arguments +/// * `partition_specs` - Iterator over all partition specs in the table +/// * `schema` - The current table schema (needed to determine result types of transforms) +pub fn compute_unified_partition_type<'a>( + partition_specs: impl Iterator<Item = &'a PartitionSpec>, + schema: &Schema, +) -> Result<StructType> { + let mut seen_field_ids = std::collections::HashSet::new(); + let mut struct_fields: Vec<NestedFieldRef> = Vec::new(); + + // Sort specs by spec_id descending (newer first) to match Java's behavior: + // newer field names take precedence when deduplicating by field_id. + let mut specs: Vec<&PartitionSpec> = partition_specs.collect(); + specs.sort_by_key(|s| std::cmp::Reverse(s.spec_id())); + + for spec in specs { + for field in spec.fields() { + if seen_field_ids.contains(&field.field_id) { + continue; + } + + // Skip void transforms (dropped partition columns) + if matches!(field.transform, Transform::Void) { Review Comment: Void/unknown handling was already discussed and implemented, so I won't re-litigate it. One residual nuance only if you're interested: because we `continue` on Void before inserting into `seen_field_ids`, if the newest spec marks a field Void and an older spec has it non-Void, the field is still included (good, matches Java's type-upgrade intent), but its name then comes from the older spec, whereas Java prefers the newest spec's name (`Partitioning.java` populates `nameMap` from the first/newest spec that defines the id). Minor and possibly fine as-is; a void-in-newest test would pin whichever behavior you intend. ########## crates/iceberg/src/arrow/record_batch_transformer.rs: ########## @@ -240,11 +274,46 @@ impl RecordBatchTransformerBuilder { Ok(self) } + /// Set the _partition metadata column constant. + /// + /// This builds the struct constant for the _partition column from the unified partition + /// type (across all specs) and the current file's partition data. + /// + /// # Arguments + /// * `unified_partition_type` - The unified partition type across all specs + /// * `partition_spec` - The partition spec for this specific file + /// * `partition_data` - The partition values for this file + #[cfg(test)] Review Comment: Since `build_partition_column_constant` is already callable, tests could do `with_partition_column_precomputed(build_partition_column_constant(...)?)`, letting us drop the `#[cfg(test)]` variant and the `_precomputed` suffix. One setter, no test-only API. Reasonable? ########## crates/iceberg/src/arrow/record_batch_transformer.rs: ########## @@ -193,6 +216,16 @@ pub(crate) struct RecordBatchTransformerBuilder { snapshot_schema: Arc<IcebergSchema>, projected_iceberg_field_ids: Vec<i32>, constant_fields: HashMap<i32, Datum>, + partition_column: Option<PartitionColumnConstant>, Review Comment: Related to the #2699 follow-up (constant map holds only primitive `Datum`, no struct variant). `_file`/`_spec_id`/identity partitions already synthesize a scalar constant (via `constant_fields`, ending in `ColumnSource::Add`), and `_partition` synthesizes a struct constant (`AddStructConstant`). Rather than adding struct support to `Datum`, one option is to converge the two synthesized paths into a small map: ``` field_id -> MetadataColumnSource { ScalarConstant(Datum), // _file, _spec_id, identity partitions StructConstant { fields, child_values }, // _partition } ``` with `ColumnSource` collapsing toward `FromMetadata { field_id }`. Both arms are test-covered (existing `_file` test plus this PR's `_partition` tests). I'd leave `_pos` (#2746) out: arrow-rs `RowNumber` makes it a real source column that is passed through, not synthesized, so its `virtual_fields` set is a different axis. (@advancedxy floated "use a `Struct` type for the child values" on the related thread; the enum is a way to do that without widening `Datum`.) Not a blocker for this PR; feel free to fold into #2699 if that's a better home. ########## crates/iceberg/src/arrow/record_batch_transformer.rs: ########## @@ -359,6 +432,21 @@ impl RecordBatchTransformer { let fields: Result<Vec<_>> = projected_iceberg_field_ids .iter() .map(|field_id| { + // Handle _partition struct column + if *field_id == RESERVED_FIELD_ID_PARTITION + && let Some(pc) = partition_column + { + let struct_type = DataType::Struct(pc.fields.clone()); + let nullable = pc.fields.is_empty(); + let arrow_field = + Field::new(RESERVED_COL_NAME_PARTITION, struct_type, nullable) Review Comment: This exact block recurs several times across these PRs. A small `fn field_with_id(name, dt, nullable, field_id) -> Arc<Field>` helper would remove the repetition and keep the field-id metadata format in one place. ########## crates/iceberg/src/scan/context.rs: ########## @@ -138,6 +142,7 @@ impl ManifestEntryContext { // TODO: Pass actual PartitionSpec through context chain for native flow .with_partition_spec(None) Review Comment: `build_partition_column_constant` needs `task.partition_spec` to be `Some` for partitioned tables, but this PR keeps `partition_spec(None)`; #2695 is the one that wires the real spec through `context.rs`. Is the plan to rebase on #2695? If so, an end-to-end test asserting non-null partition values would confirm the wiring once rebased. ########## crates/iceberg/src/scan/mod.rs: ########## @@ -300,6 +303,20 @@ impl<'a> TableScanBuilder<'a> { .transpose()? .map(Arc::new); + // Compute unified partition type if _partition is projected Review Comment: Gating on "`_partition` projected" is a nice touch. Java builds the type from `table.specs().values()` (all specs). Could we confirm `partition_specs_iter()` enumerates every historical spec, not just current, and ideally pin it with a multi-spec test? That is what makes evolution visible in `_partition`. ########## crates/iceberg/src/arrow/record_batch_transformer.rs: ########## @@ -193,6 +216,16 @@ pub(crate) struct RecordBatchTransformerBuilder { snapshot_schema: Arc<IcebergSchema>, projected_iceberg_field_ids: Vec<i32>, constant_fields: HashMap<i32, Datum>, + partition_column: Option<PartitionColumnConstant>, +} + +/// Pre-computed data for the _partition struct constant. +#[derive(Debug, Clone, PartialEq)] +pub struct PartitionColumnConstant { + /// Arrow struct fields (names, types, nullability) for the _partition column. + pub fields: Fields, Review Comment: These two are assumed equal-length (the `zip` in `create_struct_column` silently truncates otherwise). A `new(fields, child_values)` constructor plus accessors would let us assert that invariant in one place and keeps with the crate's "no public fields" convention. Minor, but it closes a small footgun. ########## crates/iceberg/src/arrow/record_batch_transformer.rs: ########## @@ -656,6 +761,84 @@ impl RecordBatchTransformer { create_primitive_array_repeated(target_type, prim_lit, num_rows) } } + + fn create_struct_column( Review Comment: `create_struct_column` is hand-rolled struct assembly that arrow-rs already provides: the empty/all-null path is `StructArray::new_null(fields, len)` / `new_null_array` (`arrow-array/src/array/struct_array.rs:192`), and the populated path is `StructArray::try_new(fields, arrays, nulls)`, which validates `fields.len() == arrays.len()` and child-length agreement (`struct_array.rs:106`). I'd reuse those rather than maintain our own constructor: `try_new` also enforces the `fields`/`child_values` length invariant that the current `zip` silently truncates. Same principle as the `value.rs` note: prefer the library's validated constructors over duplicating them here. ########## crates/iceberg/src/arrow/reader/pipeline.rs: ########## @@ -255,6 +260,30 @@ impl FileScanTaskReader { record_batch_transformer_builder.with_partition(partition_spec, partition_data)?; } + // Add the _partition metadata struct column if it's in the projected fields. + // Computed lazily here at read time from the unified partition type + task's spec + data. + if task + .project_field_ids() + .contains(&RESERVED_FIELD_ID_PARTITION) + && let Some(unified_type) = &task.unified_partition_type + { + if unified_type.fields().is_empty() { Review Comment: The "unpartitioned becomes empty struct becomes null" rule (which matches Java `PartitionUtil.constantsMap`, where an empty `partitionType` stores `null`) currently lives in the orchestration layer. If `build_partition_column_constant` handled the empty case internally and returned the empty constant, `pipeline.rs` could always call one entry point and drop the branch, and a direct caller couldn't forget the null case. Thoughts? ########## crates/iceberg/src/partitioning.rs: ########## @@ -0,0 +1,94 @@ +// 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 type utilities for Iceberg tables. + +use crate::spec::{NestedField, NestedFieldRef, PartitionSpec, Schema, StructType, Transform}; +use crate::{Error, ErrorKind, Result}; + +/// Computes the unified partition type across all partition specs in the table. +/// +/// This is equivalent to Java's `Partitioning.partitionType(table)`. The result is a +/// StructType containing all partition fields ever used across all specs, enabling correct +/// representation of the `_partition` metadata column when partition evolution has occurred. +/// +/// Matches Java's behavior: +/// - Specs are sorted by spec_id in descending order (newer specs first), so newer field +/// names take precedence when deduplicating by field_id. +/// - Void and unknown transform fields are skipped. +/// - Fields are deduplicated by field_id — each unique field_id appears exactly once. +/// +/// # Arguments +/// * `partition_specs` - Iterator over all partition specs in the table +/// * `schema` - The current table schema (needed to determine result types of transforms) +pub fn compute_unified_partition_type<'a>( + partition_specs: impl Iterator<Item = &'a PartitionSpec>, + schema: &Schema, +) -> Result<StructType> { + let mut seen_field_ids = std::collections::HashSet::new(); Review Comment: Small consistency thing: `use std::collections::HashSet;` / `use std::cmp::Reverse;` up top matches the surrounding style. ########## crates/iceberg/src/partitioning.rs: ########## @@ -0,0 +1,94 @@ +// 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 type utilities for Iceberg tables. + +use crate::spec::{NestedField, NestedFieldRef, PartitionSpec, Schema, StructType, Transform}; +use crate::{Error, ErrorKind, Result}; + +/// Computes the unified partition type across all partition specs in the table. +/// +/// This is equivalent to Java's `Partitioning.partitionType(table)`. The result is a +/// StructType containing all partition fields ever used across all specs, enabling correct +/// representation of the `_partition` metadata column when partition evolution has occurred. +/// +/// Matches Java's behavior: +/// - Specs are sorted by spec_id in descending order (newer specs first), so newer field +/// names take precedence when deduplicating by field_id. +/// - Void and unknown transform fields are skipped. +/// - Fields are deduplicated by field_id — each unique field_id appears exactly once. +/// +/// # Arguments +/// * `partition_specs` - Iterator over all partition specs in the table +/// * `schema` - The current table schema (needed to determine result types of transforms) +pub fn compute_unified_partition_type<'a>( + partition_specs: impl Iterator<Item = &'a PartitionSpec>, + schema: &Schema, +) -> Result<StructType> { + let mut seen_field_ids = std::collections::HashSet::new(); + let mut struct_fields: Vec<NestedFieldRef> = Vec::new(); + + // Sort specs by spec_id descending (newer first) to match Java's behavior: + // newer field names take precedence when deduplicating by field_id. + let mut specs: Vec<&PartitionSpec> = partition_specs.collect(); + specs.sort_by_key(|s| std::cmp::Reverse(s.spec_id())); + + for spec in specs { + for field in spec.fields() { + if seen_field_ids.contains(&field.field_id) { + continue; + } + + // Skip void transforms (dropped partition columns) + if matches!(field.transform, Transform::Void) { + continue; + } + + // Reject unknown transforms — the table uses a spec feature that + // this version of iceberg-rust doesn't support. Matching Java's + // behavior where getResultType() throws for unknown transforms. + if matches!(field.transform, Transform::Unknown) { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + format!( + "Partition field '{}' uses an unknown transform that is not \ + supported by this version of iceberg-rust", + field.name + ), + )); + } + + seen_field_ids.insert(field.field_id); + + let source_field = schema.field_by_id(field.source_id).ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + format!( + "No column with source column id {} in schema for partition field {}", + field.source_id, field.name + ), + ) + })?; + + let res_type = field.transform.result_type(&source_field.field_type)?; + let nested = NestedField::optional(field.field_id, &field.name, res_type).into(); + struct_fields.push(nested); + } + } + + Ok(StructType::new(struct_fields)) Review Comment: `buildPartitionProjectionType` sorts its output fields by id ascending (`fieldMap.keySet().stream().sorted(Comparator.naturalOrder())` in `Partitioning.java`). Here we emit them in spec-descending, first-seen order. In the common case these coincide, but under partition evolution I think they can differ, and since the `_partition` struct schema is consumed by engines (this is the Comet path), matching Java's ordering seems safer. Worth a sort plus an evolution test that would otherwise reorder? Curious whether you already considered this and decided it's fine. ########## crates/iceberg/src/scan/task.rs: ########## @@ -116,6 +116,17 @@ pub struct FileScanTask { #[builder(default)] pub name_mapping: Option<Arc<NameMapping>>, + /// The unified partition type across all specs in the table. + /// When `RESERVED_FIELD_ID_PARTITION` is in the projected field IDs, the reader + /// uses this type along with the task's partition_spec and partition data to + /// materialize the `_partition` struct column at read time. + #[serde(default)] + #[serde(skip_serializing_if = "Option::is_none")] + #[serde(serialize_with = "serialize_not_implemented")] Review Comment: This connects to your own thread on the PR (Comet doesn't get `unified_partition_type` from Iceberg Java, and parthchandra's plan is for Comet to compute the equivalent, per the linked planner.rs). Independent of that: the serde attrs here match the existing `partition` / `partition_spec` / `name_mapping` fields, so it's consistent, not a regression. The implication worth one line in the PR text: this field (like the others) does not round-trip through serde, so the `_partition` machinery works in the native scan flow only. Since Comet builds tasks natively that's fine, just easy for a future reader to assume otherwise. -- 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]
