laskoviymishka commented on code in PR #2015:
URL: https://github.com/apache/iceberg-go/pull/2015#discussion_r4030782779
##########
table/partitioned_fanout_writer.go:
##########
@@ -397,6 +397,66 @@ func newPartitionExtractionPlan(spec
iceberg.PartitionSpec, schema *iceberg.Sche
}, nil
}
+// resolveArrowColumnPath resolves a dot-joined nested field name into a path
of Arrow field indices.
+func resolveArrowColumnPath(recordSchema *arrow.Schema, colName string)
([]int, bool) {
+ segments := strings.Split(colName, ".")
Review Comment:
Splitting on `.` before trying an exact match regresses top-level fields
whose name contains a literal dot. `FindColumnName` joins nested segments with
`.` but returns top-level names verbatim, so a field genuinely named
`user.name` comes back as `user.name`, gets split into `[user, name]`, and we
go looking for a struct field `user` that isn't there, and
`resolveArrowColumnPath` returns false so the plan build fails. Pre-PR,
`recordSchema.FieldIndices(colName)` matched that name exactly and the write
succeeded.
I'd try the exact top-level name first and only fall back to dot-splitting
when it misses, which is what PyIceberg's `_get_field_from_arrow_table` does.
Cleaner still: `Schema.columnPathSegments` (schema.go:528) already walks the
field-ID map and yields each raw name as its own segment without ever joining
on `.`, so it can't reintroduce this ambiguity. It's unexported today, but
exporting it and using it here sidesteps the split entirely. wdyt?
##########
table/partitioned_fanout_writer.go:
##########
@@ -411,17 +471,38 @@ func (p *partitionExtractionPlan)
getRecordPartitions(record arrow.RecordBatch)
partitionMap := newPartitionMapNode()
partitionRec := make(partitionRecord, len(p.fields))
- partitionColumns := make([]arrow.Array, len(p.fields))
+
+ // we track these separately to avoid traversing the chain for top
level fields.
+ topLevelColumns := make([]arrow.Array, len(p.fields))
+ columnChains := make([][]arrow.Array, len(p.fields))
for i, fieldInfo := range p.fields {
- if fieldInfo.columnIndex >= 0 {
- partitionColumns[i] =
record.Column(fieldInfo.columnIndex)
+ switch len(fieldInfo.columnPath) {
+ case 0:
Review Comment:
An empty `case 0:` with no comment is easy to misread as a mistake. A
one-liner noting the intent (source field isn't present in the record schema,
so leave both arrays nil and let the row loop treat it as absent) would save
the next reader a double-take.
##########
table/partitioned_fanout_writer_test.go:
##########
@@ -1471,6 +1471,59 @@ func (s *FanoutWriterTestSuite)
TestPartitionExtractionPlanHandlesReorderedRecor
s.ElementsMatch([]int32{7, 8}, values)
}
+func (s *FanoutWriterTestSuite) TestPartitionExtractionPlanNestedSourceField()
{
+ // "payload" is a struct column; the partition source field
("event_time",
+ // iceberg field ID 2) lives inside it rather than at the top level.
+ icebergSchema := iceberg.NewSchema(0,
+ iceberg.NestedField{ID: 1, Name: "payload", Type:
&iceberg.StructType{
+ FieldList: []iceberg.NestedField{
+ {ID: 2, Name: "event_time", Type:
iceberg.PrimitiveTypes.Int32},
+ },
+ }},
+ )
+ spec := iceberg.NewPartitionSpec(iceberg.PartitionField{
+ SourceIDs: []int{2}, FieldID: 1000, Transform:
iceberg.IdentityTransform{}, Name: "event_time",
+ })
+
+ structType := arrow.StructOf(arrow.Field{Name: "event_time", Type:
arrow.PrimitiveTypes.Int32, Nullable: true})
+ arrowSchema := arrow.NewSchema([]arrow.Field{
+ {Name: "payload", Type: structType, Nullable: true},
+ }, nil)
+
+ plan, err := newPartitionExtractionPlan(spec, icebergSchema,
arrowSchema)
+ s.Require().NoError(err)
+ s.Equal([]int{0, 0}, plan.fields[0].columnPath)
+
+ bldr := array.NewRecordBuilder(s.mem, arrowSchema)
+ defer bldr.Release()
+ structBldr := bldr.Field(0).(*array.StructBuilder)
+ eventTimeBldr := structBldr.FieldBuilder(0).(*array.Int32Builder)
+
+ structBldr.Append(true)
+ eventTimeBldr.Append(7)
+
+ structBldr.Append(true)
+ eventTimeBldr.Append(8)
+
+ structBldr.Append(false) // null struct row -> partition value should
be nil
+ eventTimeBldr.AppendNull()
Review Comment:
The new test only exercises one shape: two-level nesting, identity
transform, with a fully-null struct at the leaf. The most load-bearing gap is a
non-null struct with a null leaf, since that's the only case that reaches the
`leaf.IsNull` guard in `leafColumnAt`, and right now nothing would fail if that
check were inverted or dropped. A fourth row with `structBldr.Append(true)` and
`eventTimeBldr.AppendNull()` asserting a nil partition would lock it.
While we're here, the missing-segment error path, 3+ levels of nesting, a
spec mixing top-level and nested fields, and the dotted-name case from the
regression above are all worth a case each so this can't creep back in.
##########
table/partitioned_fanout_writer.go:
##########
@@ -397,6 +397,66 @@ func newPartitionExtractionPlan(spec
iceberg.PartitionSpec, schema *iceberg.Sche
}, nil
}
+// resolveArrowColumnPath resolves a dot-joined nested field name into a path
of Arrow field indices.
+func resolveArrowColumnPath(recordSchema *arrow.Schema, colName string)
([]int, bool) {
+ segments := strings.Split(colName, ".")
+
+ indices := recordSchema.FieldIndices(segments[0])
+ if len(indices) == 0 {
+ return nil, false
+ }
+
+ path := make([]int, 1, len(segments))
+ path[0] = indices[0]
+ fieldType := recordSchema.Field(indices[0]).Type
+
+ for _, segment := range segments[1:] {
+ structType, ok := fieldType.(*arrow.StructType)
+ if !ok {
+ return nil, false
+ }
+ fieldIdx, ok := structType.FieldIdx(segment)
+ if !ok {
+ return nil, false
+ }
+ path = append(path, fieldIdx)
+ fieldType = structType.Field(fieldIdx).Type
+ }
+
+ return path, true
+}
+
+// resolveColumnChain walks a resolved column path from the record's top-level
column
+// down to the leaf, returning the array at each step of the path.
+func resolveColumnChain(record arrow.RecordBatch, path []int) []arrow.Array {
+ chain := make([]arrow.Array, len(path))
+ col := record.Column(path[0])
+ chain[0] = col
+ for i, fieldIdx := range path[1:] {
+ col = col.(*array.Struct).Field(fieldIdx)
Review Comment:
This assertion panics instead of erroring when the runtime array isn't a
struct. The schema-level guard in `resolveArrowColumnPath` checks the Arrow
field type is a struct, but a dictionary-encoded struct column satisfies that
and still arrives as `*array.Dictionary` here, so `col.(*array.Struct)` blows
up. A panic in a fanout worker goroutine isn't caught by the errgroup, so it
takes down the process rather than surfacing through `errCh`.
I'd make this a comma-ok assertion and return an error, probably changing
`resolveColumnChain` to `([]arrow.Array, error)` and threading it through the
two call sites. wdyt?
--
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]
