laskoviymishka commented on code in PR #2002:
URL: https://github.com/apache/iceberg-go/pull/2002#discussion_r4024419800
##########
table/arrow_utils.go:
##########
@@ -436,7 +436,7 @@ func (c convertToIceberg) Primitive(dt arrow.DataType)
(result iceberg.NestedFie
switch dt.ExtensionName() {
case "arrow.uuid":
result.Type = iceberg.PrimitiveTypes.UUID
- case "parquet.variant":
+ case extensions.VariantExtensionName,
extensions.LegacyVariantExtensionName:
Review Comment:
Taking both names on the read side is the right call. The one thing I'd
still note: after the 18.8 bump every variant column we write embeds
`arrow.parquet.variant`, and a reader still on 18.7 only matches the literal
`parquet.variant`, falls through, and hits the `default: panic`. So a
mixed-version deployment reading a freshly-written variant column crashes
during the upgrade window.
Not a blocker, but I'd land the CHANGELOG/upgrade note I flagged last round,
calling out that variant columns written after this need 18.8+ readers. If we
wanted to be extra safe we could keep emitting the legacy name in
`VisitVariant` for a release, though I don't think that's necessary. wdyt?
##########
table/internal/parquet_files.go:
##########
@@ -711,9 +711,91 @@ func getWriteProperties(writeProps any, arrowSchema
*arrow.Schema) (*parquet.Wri
wp = append(wp, parquet.WithStoreDecimalAsInteger(true))
}
+ // Match Iceberg Java: apply parquet-mr's cost-based dictionary
fallback to every leaf
+ // column so high-cardinality columns fall back to PLAIN rather than
keeping a dictionary.
+ // arrow-go otherwise enables it only for uncompressed columns, so zstd
(our default) would
+ // retain dictionaries on all-distinct columns and roughly double their
size.
+ costFallback, err := dictCostFallbackProps(arrowSchema, wp)
+ if err != nil {
+ return nil, err
+ }
+ wp = append(wp, costFallback...)
+
return parquet.NewWriterProperties(wp...), nil
}
+// dictCostFallbackProps returns a WithDictionaryCostFallbackFor(true)
property per leaf, walking the arrow schema directly (extensions unwrapped) and
falling back to pqarrow.ToParquet for list/map schemas.
+func dictCostFallbackProps(arrowSchema *arrow.Schema, base
[]parquet.WriterProperty) ([]parquet.WriterProperty, error) {
+ if schemaHasListOrMap(arrowSchema) {
+ return dictCostFallbackViaParquet(arrowSchema, base)
Review Comment:
Small gap I'd close before this is done: this list/map branch is the one
path none of the new tests reach. `TestDictCostFallbackWalkMatchesToParquet`
asserts `require.False(schemaHasListOrMap(sc))` on every case, so
`dictCostFallbackViaParquet` never actually runs in CI, and the one list-schema
test in `parquet_files_test.go` passes `WriteProps` directly and skips
`getWriteProperties`. Since this is a real size fix, a wrong leaf path here
would silently keep inflating list/map files. A single `arrow.ListOfField` case
that asserts `schemaHasListOrMap` is true and round-trips the leaf paths
through `DictionaryCostFallbackEnabledFor` would lock it.
Nearby and purely defensive: the direct walk treats `*arrow.DictionaryType`
as a leaf, and `schemaHasListOrMap` doesn't look through it either, so a
dict-of-struct column would take the walk path and set fallback on the wrong
leaf. Doesn't arise from our normal writes, but unwrapping `DictionaryType` (or
adding it to `schemaHasListOrMap`) closes it.
##########
table/variant_residual.go:
##########
@@ -68,7 +68,112 @@ func buildExtractColumn(col iceberg.VariantExtractColumn,
rec arrow.RecordBatch,
return nil, arrow.Field{}, fmt.Errorf("%w: variant extract
column %q is not a VariantArray (got %T)", iceberg.ErrInvalidArgument, varName,
arr)
}
- for i := range n {
+ out, err := extractColumnValues(ctx, varr, col, typ, dt, mem)
+ if err != nil {
+ return nil, arrow.Field{}, err
+ }
+
+ field := arrow.Field{
+ Name: col.Name,
+ Type: dt,
+ Nullable: true,
+ Metadata: arrow.NewMetadata([]string{ArrowParquetFieldIDKey},
[]string{strconv.Itoa(col.FieldID)}),
+ }
+
+ return out, field, nil
+}
+
+// variantPathOf returns the extract term's member-name path (kept off the
public BoundExtract interface); false if the term has none.
+func variantPathOf(t iceberg.BoundExtract) (variant.VariantPath, bool) {
+ vp, ok := t.(interface {
+ VariantPath() variant.VariantPath
+ })
+ if !ok {
+ return variant.VariantPath{}, false
+ }
+
+ return vp.VariantPath(), true
+}
+
+// extractColumnValues navigates columnarly then casts each leaf with
iceberg's cast; a VariantGet
+// navigation error (arrow-rs errors where the residual filter wants
null-on-miss) falls back to per-row.
+func extractColumnValues(ctx context.Context, varr *extensions.VariantArray,
col iceberg.VariantExtractColumn, typ iceberg.PrimitiveType, dt arrow.DataType,
mem memory.Allocator) (arrow.Array, error) {
+ if err := ctx.Err(); err != nil {
+ return nil, err
+ }
+ path, ok := variantPathOf(col.Term)
+ if !ok {
+ return extractColumnValuesPerRow(ctx, varr, col, dt, mem)
+ }
+ if fast := tryShreddedTypedColumn(varr, path, dt, mem); fast != nil {
+ return fast, nil
+ }
+ if !varr.IsShredded() {
+ return extractColumnValuesPerRow(ctx, varr, col, dt, mem)
+ }
+
+ extracted, err := compute.VariantGet(ctx, varr,
compute.VariantGetOptions{Path: path})
+ if err != nil {
+ if errors.Is(err, context.Canceled) || errors.Is(err,
context.DeadlineExceeded) {
+ return nil, err
+ }
+
+ return extractColumnValuesPerRow(ctx, varr, col, dt, mem)
+ }
+ defer extracted.Release()
+
+ leaves, ok := extracted.(*extensions.VariantArray)
+ if !ok {
+ return extractColumnValuesPerRow(ctx, varr, col, dt, mem)
+ }
+
+ bldr := array.NewBuilder(mem, dt)
+ defer bldr.Release()
+
+ varName := col.Term.Ref().Field().Name
+ for i := range leaves.Len() {
Review Comment:
Nice, the blanket fallback inspects the error now (`errors.Is(err,
context.Canceled)` above) instead of swallowing it, thanks.
One nit on this loop: it's the one leg that still doesn't poll, so on a 128K
batch a cancel after `VariantGet` returns isn't noticed until we've walked
every row, where the per-row path checks every 4096. Dropping the same `if
i%4096 == 0 { if err := ctx.Err(); err != nil { return nil, err } }` in here
lines all three tiers up.
##########
table/variant_residual.go:
##########
@@ -91,18 +196,138 @@ func buildExtractColumn(col iceberg.VariantExtractColumn,
rec arrow.RecordBatch,
}
if aerr := appendExtractLiteral(bldr, lit); aerr != nil {
- return nil, arrow.Field{}, aerr
+ return nil, aerr
}
}
- field := arrow.Field{
- Name: col.Name,
- Type: dt,
- Nullable: true,
- Metadata: arrow.NewMetadata([]string{ArrowParquetFieldIDKey},
[]string{strconv.Itoa(col.FieldID)}),
+ return bldr.NewArray(), nil
+}
+
+// tryShreddedTypedColumn returns the field's typed leaf column when it is
shredded to exactly dt, else nil.
+func tryShreddedTypedColumn(varr *extensions.VariantArray, path
variant.VariantPath, dt arrow.DataType, mem memory.Allocator) arrow.Array {
+ if path.Len() == 0 || varr.Data().Offset() != 0 {
+ return nil
+ }
+ tv := varr.Shredded()
+ if tv == nil || rootResidualHidesRows(varr, tv) {
+ return nil
+ }
+ n := varr.Len()
+
+ var mask *memory.Buffer
+ badOffset := false
+ mergeValidity := func(arr arrow.Array) {
+ if arr.Data().Offset() != 0 {
+ badOffset = true // child at a non-zero offset: our
offset-0 bit indexing would be wrong
+
+ return
+ }
+ if arr.NullN() == 0 {
+ return
+ }
+ vb := arr.Data().Buffers()[0]
+ if vb == nil {
+ return
+ }
+ if mask == nil {
+ mask = memory.NewResizableBuffer(mem)
+ mask.Resize(int(bitutil.BytesForBits(int64(n))))
+ copy(mask.Bytes(), vb.Bytes())
+
+ return
+ }
+ merged := bitutil.BitmapAndAlloc(mem, mask.Bytes(), vb.Bytes(),
0, 0, int64(n), 0)
+ mask.Release()
+ mask = merged
+ }
+ bail := func() arrow.Array {
+ if mask != nil {
+ mask.Release()
+ }
+
+ return nil
+ }
+
+ mergeValidity(varr.Storage())
+
+ cur := tv
+ for i := range path.Len() {
+ name, _, isField := path.StepAt(i)
+ if !isField {
+ return bail()
+ }
+ st, ok := cur.(*array.Struct)
+ if !ok {
+ return bail()
+ }
+ mergeValidity(st)
+ idx, ok := st.DataType().(*arrow.StructType).FieldIdx(name)
+ if !ok {
+ return bail()
+ }
+ field, ok := st.Field(idx).(*array.Struct)
+ if !ok {
+ return bail()
+ }
+ fty := field.DataType().(*arrow.StructType)
+ if vIdx, ok := fty.FieldIdx("value"); ok {
+ if v := field.Field(vIdx); v.NullN() != v.Len() {
+ return bail()
+ }
+ }
+ tvIdx, ok := fty.FieldIdx("typed_value")
+ if !ok {
+ return bail()
+ }
+ cur = field.Field(tvIdx)
Review Comment:
This is the wrapper-validity case from last round.
`TestFastPathWrapperFieldNullMatchesPerRow` landed and the fast and per-row
paths agree, so we're consistent either way. The open question is just which
behavior is right: we check the wrapper's `value` child is all-null but never
fold the wrapper struct's own validity into `mask`, so a row where the wrapper
is null but `typed_value` is live returns the live value. The shredding spec
reads a null wrapper as an absent key, i.e. the strict result is null.
Either fold `mergeValidity(field)` in here for strict-null semantics, or
keep the permissive read and drop the `spec-undefined shape` label in the test
(the spec does define this, it's malformed writer output we're choosing to
tolerate). I lean slightly toward the guard, but I'm fine either way as long as
it's documented. wdyt?
##########
table/variant_residual.go:
##########
@@ -91,18 +196,138 @@ func buildExtractColumn(col iceberg.VariantExtractColumn,
rec arrow.RecordBatch,
}
if aerr := appendExtractLiteral(bldr, lit); aerr != nil {
- return nil, arrow.Field{}, aerr
+ return nil, aerr
}
}
- field := arrow.Field{
- Name: col.Name,
- Type: dt,
- Nullable: true,
- Metadata: arrow.NewMetadata([]string{ArrowParquetFieldIDKey},
[]string{strconv.Itoa(col.FieldID)}),
+ return bldr.NewArray(), nil
+}
+
+// tryShreddedTypedColumn returns the field's typed leaf column when it is
shredded to exactly dt, else nil.
+func tryShreddedTypedColumn(varr *extensions.VariantArray, path
variant.VariantPath, dt arrow.DataType, mem memory.Allocator) arrow.Array {
+ if path.Len() == 0 || varr.Data().Offset() != 0 {
+ return nil
+ }
+ tv := varr.Shredded()
+ if tv == nil || rootResidualHidesRows(varr, tv) {
+ return nil
+ }
+ n := varr.Len()
+
+ var mask *memory.Buffer
+ badOffset := false
+ mergeValidity := func(arr arrow.Array) {
+ if arr.Data().Offset() != 0 {
+ badOffset = true // child at a non-zero offset: our
offset-0 bit indexing would be wrong
+
+ return
+ }
+ if arr.NullN() == 0 {
+ return
+ }
+ vb := arr.Data().Buffers()[0]
+ if vb == nil {
+ return
+ }
+ if mask == nil {
+ mask = memory.NewResizableBuffer(mem)
+ mask.Resize(int(bitutil.BytesForBits(int64(n))))
+ copy(mask.Bytes(), vb.Bytes())
+
+ return
+ }
+ merged := bitutil.BitmapAndAlloc(mem, mask.Bytes(), vb.Bytes(),
0, 0, int64(n), 0)
+ mask.Release()
+ mask = merged
+ }
+ bail := func() arrow.Array {
+ if mask != nil {
+ mask.Release()
+ }
+
+ return nil
+ }
+
+ mergeValidity(varr.Storage())
+
+ cur := tv
+ for i := range path.Len() {
+ name, _, isField := path.StepAt(i)
+ if !isField {
+ return bail()
+ }
+ st, ok := cur.(*array.Struct)
+ if !ok {
+ return bail()
+ }
+ mergeValidity(st)
+ idx, ok := st.DataType().(*arrow.StructType).FieldIdx(name)
+ if !ok {
+ return bail()
+ }
+ field, ok := st.Field(idx).(*array.Struct)
+ if !ok {
+ return bail()
+ }
+ fty := field.DataType().(*arrow.StructType)
+ if vIdx, ok := fty.FieldIdx("value"); ok {
+ if v := field.Field(vIdx); v.NullN() != v.Len() {
+ return bail()
+ }
+ }
+ tvIdx, ok := fty.FieldIdx("typed_value")
+ if !ok {
+ return bail()
+ }
+ cur = field.Field(tvIdx)
+ }
+
+ if badOffset || !arrow.TypeEqual(cur.DataType(), dt) {
+ return bail()
+ }
+
+ // no ancestor/row nulls: the leaf's own validity already describes the
result, return it zero-copy
+ if mask == nil {
+ cur.Retain()
+
+ return cur
+ }
+
+ mergeValidity(cur)
+ if badOffset {
+ return bail()
+ }
+ curData := cur.Data()
+ buffers := append([]*memory.Buffer(nil), curData.Buffers()...)
+ buffers[0] = mask
+ nullCount := n - bitutil.CountSetBits(mask.Bytes(), 0, n)
+ d := array.NewData(cur.DataType(), n, buffers, curData.Children(),
nullCount, 0)
+ mask.Release()
+ out := array.MakeFromData(d)
+ d.Release()
+
+ return out
+}
+
+// rootResidualHidesRows reports whether any row's whole object lives in the
root residual (value present, typed_value null) - not representable by the
typed tree, so the fast path must fall back.
+func rootResidualHidesRows(varr *extensions.VariantArray, tv arrow.Array) bool
{
+ uv := varr.UntypedValues()
+ if uv == nil || tv.NullN() == 0 || uv.NullN() == uv.Len() {
+ return false
+ }
+ if tv.Data().Offset() != 0 || uv.Data().Offset() != 0 {
+ return true // non-zero child offset: can't safely bit-index;
presume residual so the caller bails
+ }
+ tvb := tv.Data().Buffers()[0]
Review Comment:
Tiny asymmetry: `uvb` gets the `uvb == nil ||` guard on the next line, but
`tvb` is dereferenced directly in `bitutil.BitIsSet(tvb.Bytes(), i)`. We only
reach here with `tv.NullN() > 0`, so on canonical Arrow it's non-nil, but
`mergeValidity` guards the same invariant with `if vb == nil`, so a
non-canonical array with nulls and a nil validity buffer would panic here where
it bails there. Either mirror the nil-guard on `tvb` or drop a one-line comment
on why it can't be nil.
--
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]
