laskoviymishka commented on code in PR #2002:
URL: https://github.com/apache/iceberg-go/pull/2002#discussion_r4044956123
##########
table/internal/parquet_files.go:
##########
@@ -711,9 +711,97 @@ 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)
+ }
+
+ var props []parquet.WriterProperty
+ var walk func(prefix string, dt arrow.DataType)
+ walk = func(prefix string, dt arrow.DataType) {
+ if ext, ok := dt.(arrow.ExtensionType); ok {
+ dt = ext.StorageType()
+ }
+ if d, ok := dt.(*arrow.DictionaryType); ok {
+ dt = d.ValueType
+ }
+ if st, ok := dt.(*arrow.StructType); ok {
+ for _, f := range st.Fields() {
+ walk(prefix+"."+f.Name, f.Type)
+ }
+
+ return
+ }
+ props = append(props,
parquet.WithDictionaryCostFallbackFor(prefix, true))
+ }
+ for _, f := range arrowSchema.Fields() {
+ walk(f.Name, f.Type)
+ }
+
+ return props, nil
+}
+
+// dictCostFallbackViaParquet is the authoritative path for list/map schemas,
whose parquet leaf naming the direct walk does not reproduce.
+func dictCostFallbackViaParquet(arrowSchema *arrow.Schema, base
[]parquet.WriterProperty) ([]parquet.WriterProperty, error) {
+ parquetSchema, err := pqarrow.ToParquet(arrowSchema,
parquet.NewWriterProperties(base...), pqarrow.DefaultWriterProps())
+ if err != nil {
+ return nil, err
+ }
+
+ props := make([]parquet.WriterProperty, 0, parquetSchema.NumColumns())
+ for i := range parquetSchema.NumColumns() {
+ props = append(props,
parquet.WithDictionaryCostFallbackFor(parquetSchema.Column(i).Path(), true))
+ }
+
+ return props, nil
+}
+
+func schemaHasListOrMap(sc *arrow.Schema) bool {
+ var has func(dt arrow.DataType) bool
+ has = func(dt arrow.DataType) bool {
+ if ext, ok := dt.(arrow.ExtensionType); ok {
+ dt = ext.StorageType()
+ }
+ if d, ok := dt.(*arrow.DictionaryType); ok {
+ dt = d.ValueType
+ }
+ switch t := dt.(type) {
+ case *arrow.ListType, *arrow.LargeListType,
*arrow.FixedSizeListType,
Review Comment:
We unwrap `DictionaryType` here but not `*arrow.RunEndEncodedType`, so a
run-end column over a list/map slips past this check and takes the direct-walk
branch, which then registers the RLE node as a single leaf where
`pqarrow.ToParquet` would produce group children, and the cost-fallback
property becomes a no-op for those leaves.
Purely latent today, since Iceberg-generated schemas never carry RLE, so
this only bites a user-provided Arrow schema, and the downside is a dictionary
that should've fallen back (file size, not correctness). A
`*arrow.RunEndEncodedType` arm recursing into `t.Encoded()` here, plus the
parallel unwrap in `dictCostFallbackProps`'s walk, would close it.
##########
table/variant_residual.go:
##########
@@ -91,18 +201,142 @@ 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 {
Review Comment:
Small inconsistency with `rootResidualHidesRows`: there, `tvb == nil` with
`NullN() > 0` bails conservatively (returns true), but here `mergeValidity`
bare-returns on that same shape, so `mask` stays nil, we take the zero-copy
branch, and an ancestor that should null out rows reads as live.
It's non-canonical (nil buffer plus a non-zero null count needs a
non-conformant producer or a hand-built array), and the per-row path already
handles it via `IsNull(i)`, but the fast path doesn't. I'd set `badOffset =
true` here so it bails the same way. And since `badOffset` is really doing
double duty as a general "can't trust this, bail" flag now, `shouldBail` might
read truer. wdyt?
##########
table/variant_residual_fastpath_test.go:
##########
@@ -0,0 +1,592 @@
+// 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.
+
+package table
+
+import (
+ "context"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/compute"
+ "github.com/apache/arrow-go/v18/arrow/decimal128"
+ "github.com/apache/arrow-go/v18/arrow/extensions"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/parquet/variant"
+ "github.com/apache/iceberg-go"
+ "github.com/stretchr/testify/require"
+)
+
+// buildVariantExtractRec builds a "payload" variant column of shreddedType
(nil => unshredded) from rows, plus the bound extract term for path/typ.
+func buildVariantExtractRec(t testing.TB, mem memory.Allocator, shreddedType
*extensions.VariantType, path string, typ iceberg.PrimitiveType, rows
[]map[string]any) (arrow.RecordBatch, iceberg.VariantExtractColumn) {
+ t.Helper()
+ iceSchema := iceberg.NewSchema(0, iceberg.NestedField{ID: 2, Name:
"payload", Type: iceberg.VariantType{}})
+
+ vt := shreddedType
+ if vt == nil {
+ vt = extensions.NewDefaultVariantType()
+ }
+ vb := extensions.NewVariantBuilder(mem, vt)
+ for _, row := range rows {
+ if row == nil {
+ vb.AppendNull()
+
+ continue
+ }
+ var b variant.Builder
+ require.NoError(t, b.Append(row))
+ v, err := b.Build()
+ require.NoError(t, err)
+ vb.Append(v)
+ }
+ pArr := vb.NewArray()
+ vb.Release()
+
+ md := arrow.NewMetadata([]string{ArrowParquetFieldIDKey}, []string{"2"})
+ arrSchema := arrow.NewSchema([]arrow.Field{{Name: "payload", Type:
pArr.DataType(), Nullable: true, Metadata: md}}, nil)
+ rec := array.NewRecordBatch(arrSchema, []arrow.Array{pArr},
int64(pArr.Len()))
+ pArr.Release()
+
+ term, err := iceberg.Extract("payload", path, typ).Bind(iceSchema, true)
+ require.NoError(t, err)
+ col := iceberg.VariantExtractColumn{Term: term.(iceberg.BoundExtract),
FieldID: 100, Name: "_x", SourcePath: []string{"payload"}}
+
+ return rec, col
+}
+
+func shredStruct(fields ...arrow.Field) *extensions.VariantType {
+ return extensions.NewShreddedVariantType(arrow.StructOf(fields...))
+}
+
+// vpath fetches the extract term's member-name path (now off the public
BoundExtract interface).
+func vpath(tb testing.TB, term iceberg.BoundExtract) variant.VariantPath {
+ tb.Helper()
+ p, ok := variantPathOf(term)
+ require.True(tb, ok)
+
+ return p
+}
+
+// TestExtractFastPathParity: fast-path output must match the per-row walk on
every branch; wantFast asserts whether it fires.
+func TestExtractFastPathParity(t *testing.T) {
+ i64 := arrow.PrimitiveTypes.Int64
+
+ cases := []struct {
+ name string
+ shred *extensions.VariantType
+ path string
+ typ iceberg.PrimitiveType
+ rows []map[string]any
+ wantFast bool
+ }{
+ {
+ name: "exact-match int64 shredded",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, {"a":
int64(2)}, {"a": int64(3), "city": "x"}},
+ wantFast: true,
+ },
+ {
+ name: "exact-match with absent field yields null",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, {"b":
int64(9)}, {"a": int64(3)}},
+ wantFast: true,
+ },
+ {
+ name: "nested exact-match int64 shredded",
+ shred: shredStruct(arrow.Field{Name: "a", Type:
arrow.StructOf(
+ arrow.Field{Name: "b", Type: i64},
+ )}),
+ path: "$.a.b", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{
+ {"a": map[string]any{"b": int64(7)}},
+ {"a": map[string]any{"b": int64(8)}},
+ },
+ wantFast: true,
+ },
+ {
+ name: "promotion int32->int64 skips fast path",
+ shred: shredStruct(arrow.Field{Name: "a", Type:
arrow.PrimitiveTypes.Int32}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int32(1)}, {"a":
int32(2)}},
+ wantFast: false,
+ },
+ {
+ name: "int64 extracted as float64 skips fast path
(iceberg nulls it)",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Float64,
+ rows: []map[string]any{{"a": int64(5)}},
+ wantFast: false,
+ },
+ {
+ name: "unshredded skips fast path",
+ shred: nil,
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, {"a":
int64(2)}},
+ wantFast: false,
+ },
+ {
+ name: "field-level residual (mixed types) skips fast
path",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, {"a":
"not-an-int"}, {"a": int64(3)}},
+ wantFast: false,
+ },
+ {
+ name: "null row folds into validity",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, nil, {"a":
int64(3)}},
+ wantFast: true,
+ },
+ {
+ name: "null row and absent field both null in merged
mask",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, nil, {"b":
int64(9)}, {"a": int64(4)}},
+ wantFast: true,
+ },
+ {
+ name: "nested absent intermediate folds into validity",
+ shred: shredStruct(arrow.Field{Name: "a", Type:
arrow.StructOf(
+ arrow.Field{Name: "b", Type: i64},
+ )}),
+ path: "$.a.b", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{
+ {"a": map[string]any{"b": int64(7)}},
+ {"b": int64(9)},
+ {"a": map[string]any{"b": int64(3)}},
+ },
+ wantFast: true,
+ },
+ {
+ name: "field absent from shredded schema skips fast
path",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.c", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, {"a":
int64(2)}},
+ wantFast: false,
+ },
+ {
+ name: "string clean shredded",
+ shred: shredStruct(arrow.Field{Name: "a", Type:
arrow.BinaryTypes.String}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.String,
+ rows: []map[string]any{{"a": "x"}, {"a": "yy"},
{"a": "zzz"}},
+ wantFast: true,
+ },
+ {
+ name: "string with null row folds into validity",
+ shred: shredStruct(arrow.Field{Name: "a", Type:
arrow.BinaryTypes.String}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.String,
+ rows: []map[string]any{{"a": "x"}, nil, {"a":
"zzz"}},
+ wantFast: true,
+ },
+ {
+ name: "mask allocated then bail on field residual",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, nil, {"a":
"str"}},
+ wantFast: false,
+ },
+ }
+
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(t.Context(), mem)
+
+ rec, col := buildVariantExtractRec(t, mem, tc.shred,
tc.path, tc.typ, tc.rows)
+ defer rec.Release()
+ varr := resolveVariantSource(rec,
col.Term.Ref().Field().ID, col.SourcePath).(*extensions.VariantArray)
+ dt, err := TypeToArrowType(tc.typ, false, false)
+ require.NoError(t, err)
+
+ ref := tryShreddedTypedColumn(varr, vpath(t, col.Term),
dt, mem)
+ if ref != nil {
+ defer ref.Release()
+ }
+ require.Equal(t, tc.wantFast, ref != nil, "fast path
firing")
+
+ got, _, err := buildExtractColumn(ctx, col, rec, mem)
+ require.NoError(t, err)
+ defer got.Release()
+
+ want, err := extractColumnValuesPerRow(ctx, varr, col,
dt, mem)
+ require.NoError(t, err)
+ defer want.Release()
+
+ require.Truef(t, array.Equal(got, want),
+ "fast/columnar output diverges from per-row
reference\n got=%v\nwant=%v", got, want)
+
+ if ref != nil {
+ require.True(t, sharesDataBuffers(got, ref),
"fast path must share the typed column's data (no per-row rebuild)")
+ }
+ })
+ }
+}
+
+// TestFastPathRootResidualObjectFallsBack: a row whose whole object is in the
root residual (value present, typed_value null) must fall back, not be nulled.
+func TestFastPathRootResidualObjectFallsBack(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(t.Context(), mem)
+
+ vt :=
extensions.NewShreddedVariantType(arrow.StructOf(arrow.Field{Name: "a", Type:
arrow.PrimitiveTypes.Int64}))
+ build := func(obj map[string]any) ([]byte, []byte) {
+ var b variant.Builder
+ require.NoError(t, b.Append(obj))
+ v, err := b.Build()
+ require.NoError(t, err)
+
+ return v.Metadata().Bytes(), v.Bytes()
+ }
+ m0, _ := build(map[string]any{"a": int64(1)})
+ m1, v1 := build(map[string]any{"a": int64(5)})
+
+ sb := array.NewStructBuilder(mem, vt.StorageType().(*arrow.StructType))
+ defer sb.Release()
+ metaB := sb.FieldBuilder(0).(*array.BinaryBuilder)
+ valB := sb.FieldBuilder(1).(*array.BinaryBuilder)
+ tvB := sb.FieldBuilder(2).(*array.StructBuilder)
+ aB := tvB.FieldBuilder(0).(*array.StructBuilder)
+ aValB := aB.FieldBuilder(0).(*array.BinaryBuilder)
+ aTypedB := aB.FieldBuilder(1).(*array.Int64Builder)
+
+ // row 0: shredded {a:1}
+ sb.Append(true)
+ metaB.Append(m0)
+ valB.AppendNull()
+ tvB.Append(true)
+ aB.Append(true)
+ aValB.AppendNull()
+ aTypedB.Append(1)
+ // row 1: whole object {a:5} in the root residual, typed_value null
+ sb.Append(true)
+ metaB.Append(m1)
+ valB.Append(v1)
+ tvB.AppendNulls(1)
Review Comment:
Row 1 appends a null to `tvB` but never appends to `aB`/`aValB`/`aTypedB`,
so the struct ends up with `tvB` at length 2 and its children at length 1, a
length mismatch that violates the Arrow invariant. It passes only because
`rootResidualHidesRows` returns true for this row and `tryShreddedTypedColumn`
bails before indexing the children, and the per-row fallback reads
`varr.Value(i)` rather than the subtree.
If the bail guard ever changes, this fixture surfaces as a bounds error
instead of the case it's meant to test. Building the storage directly with
`array.NewData` (like `TestFastPathWrapperFieldNullIsAbsentKey` does) would
keep the child lengths honest.
##########
table/variant_residual_fastpath_test.go:
##########
@@ -0,0 +1,592 @@
+// 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.
+
+package table
+
+import (
+ "context"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/compute"
+ "github.com/apache/arrow-go/v18/arrow/decimal128"
+ "github.com/apache/arrow-go/v18/arrow/extensions"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/parquet/variant"
+ "github.com/apache/iceberg-go"
+ "github.com/stretchr/testify/require"
+)
+
+// buildVariantExtractRec builds a "payload" variant column of shreddedType
(nil => unshredded) from rows, plus the bound extract term for path/typ.
+func buildVariantExtractRec(t testing.TB, mem memory.Allocator, shreddedType
*extensions.VariantType, path string, typ iceberg.PrimitiveType, rows
[]map[string]any) (arrow.RecordBatch, iceberg.VariantExtractColumn) {
+ t.Helper()
+ iceSchema := iceberg.NewSchema(0, iceberg.NestedField{ID: 2, Name:
"payload", Type: iceberg.VariantType{}})
+
+ vt := shreddedType
+ if vt == nil {
+ vt = extensions.NewDefaultVariantType()
+ }
+ vb := extensions.NewVariantBuilder(mem, vt)
+ for _, row := range rows {
+ if row == nil {
+ vb.AppendNull()
+
+ continue
+ }
+ var b variant.Builder
+ require.NoError(t, b.Append(row))
+ v, err := b.Build()
+ require.NoError(t, err)
+ vb.Append(v)
+ }
+ pArr := vb.NewArray()
+ vb.Release()
+
+ md := arrow.NewMetadata([]string{ArrowParquetFieldIDKey}, []string{"2"})
+ arrSchema := arrow.NewSchema([]arrow.Field{{Name: "payload", Type:
pArr.DataType(), Nullable: true, Metadata: md}}, nil)
+ rec := array.NewRecordBatch(arrSchema, []arrow.Array{pArr},
int64(pArr.Len()))
+ pArr.Release()
+
+ term, err := iceberg.Extract("payload", path, typ).Bind(iceSchema, true)
+ require.NoError(t, err)
+ col := iceberg.VariantExtractColumn{Term: term.(iceberg.BoundExtract),
FieldID: 100, Name: "_x", SourcePath: []string{"payload"}}
+
+ return rec, col
+}
+
+func shredStruct(fields ...arrow.Field) *extensions.VariantType {
+ return extensions.NewShreddedVariantType(arrow.StructOf(fields...))
+}
+
+// vpath fetches the extract term's member-name path (now off the public
BoundExtract interface).
+func vpath(tb testing.TB, term iceberg.BoundExtract) variant.VariantPath {
+ tb.Helper()
+ p, ok := variantPathOf(term)
+ require.True(tb, ok)
+
+ return p
+}
+
+// TestExtractFastPathParity: fast-path output must match the per-row walk on
every branch; wantFast asserts whether it fires.
+func TestExtractFastPathParity(t *testing.T) {
+ i64 := arrow.PrimitiveTypes.Int64
+
+ cases := []struct {
+ name string
+ shred *extensions.VariantType
+ path string
+ typ iceberg.PrimitiveType
+ rows []map[string]any
+ wantFast bool
+ }{
+ {
+ name: "exact-match int64 shredded",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, {"a":
int64(2)}, {"a": int64(3), "city": "x"}},
+ wantFast: true,
+ },
+ {
+ name: "exact-match with absent field yields null",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, {"b":
int64(9)}, {"a": int64(3)}},
+ wantFast: true,
+ },
+ {
+ name: "nested exact-match int64 shredded",
+ shred: shredStruct(arrow.Field{Name: "a", Type:
arrow.StructOf(
+ arrow.Field{Name: "b", Type: i64},
+ )}),
+ path: "$.a.b", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{
+ {"a": map[string]any{"b": int64(7)}},
+ {"a": map[string]any{"b": int64(8)}},
+ },
+ wantFast: true,
+ },
+ {
+ name: "promotion int32->int64 skips fast path",
+ shred: shredStruct(arrow.Field{Name: "a", Type:
arrow.PrimitiveTypes.Int32}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int32(1)}, {"a":
int32(2)}},
+ wantFast: false,
+ },
+ {
+ name: "int64 extracted as float64 skips fast path
(iceberg nulls it)",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Float64,
+ rows: []map[string]any{{"a": int64(5)}},
+ wantFast: false,
+ },
+ {
+ name: "unshredded skips fast path",
+ shred: nil,
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, {"a":
int64(2)}},
+ wantFast: false,
+ },
+ {
+ name: "field-level residual (mixed types) skips fast
path",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, {"a":
"not-an-int"}, {"a": int64(3)}},
+ wantFast: false,
+ },
+ {
+ name: "null row folds into validity",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, nil, {"a":
int64(3)}},
+ wantFast: true,
+ },
+ {
+ name: "null row and absent field both null in merged
mask",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, nil, {"b":
int64(9)}, {"a": int64(4)}},
+ wantFast: true,
+ },
+ {
+ name: "nested absent intermediate folds into validity",
+ shred: shredStruct(arrow.Field{Name: "a", Type:
arrow.StructOf(
+ arrow.Field{Name: "b", Type: i64},
+ )}),
+ path: "$.a.b", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{
+ {"a": map[string]any{"b": int64(7)}},
+ {"b": int64(9)},
+ {"a": map[string]any{"b": int64(3)}},
+ },
+ wantFast: true,
+ },
+ {
+ name: "field absent from shredded schema skips fast
path",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.c", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, {"a":
int64(2)}},
+ wantFast: false,
+ },
+ {
+ name: "string clean shredded",
+ shred: shredStruct(arrow.Field{Name: "a", Type:
arrow.BinaryTypes.String}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.String,
+ rows: []map[string]any{{"a": "x"}, {"a": "yy"},
{"a": "zzz"}},
+ wantFast: true,
+ },
+ {
+ name: "string with null row folds into validity",
+ shred: shredStruct(arrow.Field{Name: "a", Type:
arrow.BinaryTypes.String}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.String,
+ rows: []map[string]any{{"a": "x"}, nil, {"a":
"zzz"}},
+ wantFast: true,
+ },
+ {
+ name: "mask allocated then bail on field residual",
+ shred: shredStruct(arrow.Field{Name: "a", Type: i64}),
+ path: "$.a", typ: iceberg.PrimitiveTypes.Int64,
+ rows: []map[string]any{{"a": int64(1)}, nil, {"a":
"str"}},
+ wantFast: false,
+ },
+ }
+
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(t.Context(), mem)
+
+ rec, col := buildVariantExtractRec(t, mem, tc.shred,
tc.path, tc.typ, tc.rows)
+ defer rec.Release()
+ varr := resolveVariantSource(rec,
col.Term.Ref().Field().ID, col.SourcePath).(*extensions.VariantArray)
+ dt, err := TypeToArrowType(tc.typ, false, false)
+ require.NoError(t, err)
+
+ ref := tryShreddedTypedColumn(varr, vpath(t, col.Term),
dt, mem)
+ if ref != nil {
+ defer ref.Release()
+ }
+ require.Equal(t, tc.wantFast, ref != nil, "fast path
firing")
+
+ got, _, err := buildExtractColumn(ctx, col, rec, mem)
+ require.NoError(t, err)
+ defer got.Release()
+
+ want, err := extractColumnValuesPerRow(ctx, varr, col,
dt, mem)
+ require.NoError(t, err)
+ defer want.Release()
+
+ require.Truef(t, array.Equal(got, want),
+ "fast/columnar output diverges from per-row
reference\n got=%v\nwant=%v", got, want)
+
+ if ref != nil {
+ require.True(t, sharesDataBuffers(got, ref),
"fast path must share the typed column's data (no per-row rebuild)")
+ }
+ })
+ }
+}
+
+// TestFastPathRootResidualObjectFallsBack: a row whose whole object is in the
root residual (value present, typed_value null) must fall back, not be nulled.
+func TestFastPathRootResidualObjectFallsBack(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(t.Context(), mem)
+
+ vt :=
extensions.NewShreddedVariantType(arrow.StructOf(arrow.Field{Name: "a", Type:
arrow.PrimitiveTypes.Int64}))
+ build := func(obj map[string]any) ([]byte, []byte) {
+ var b variant.Builder
+ require.NoError(t, b.Append(obj))
+ v, err := b.Build()
+ require.NoError(t, err)
+
+ return v.Metadata().Bytes(), v.Bytes()
+ }
+ m0, _ := build(map[string]any{"a": int64(1)})
+ m1, v1 := build(map[string]any{"a": int64(5)})
+
+ sb := array.NewStructBuilder(mem, vt.StorageType().(*arrow.StructType))
+ defer sb.Release()
+ metaB := sb.FieldBuilder(0).(*array.BinaryBuilder)
+ valB := sb.FieldBuilder(1).(*array.BinaryBuilder)
+ tvB := sb.FieldBuilder(2).(*array.StructBuilder)
+ aB := tvB.FieldBuilder(0).(*array.StructBuilder)
+ aValB := aB.FieldBuilder(0).(*array.BinaryBuilder)
+ aTypedB := aB.FieldBuilder(1).(*array.Int64Builder)
+
+ // row 0: shredded {a:1}
+ sb.Append(true)
+ metaB.Append(m0)
+ valB.AppendNull()
+ tvB.Append(true)
+ aB.Append(true)
+ aValB.AppendNull()
+ aTypedB.Append(1)
+ // row 1: whole object {a:5} in the root residual, typed_value null
+ sb.Append(true)
+ metaB.Append(m1)
+ valB.Append(v1)
+ tvB.AppendNulls(1)
+
+ storage := sb.NewStructArray()
+ defer storage.Release()
+ varr := array.NewExtensionArrayWithStorage(vt,
storage).(*extensions.VariantArray)
+ defer varr.Release()
+
+ md := arrow.NewMetadata([]string{ArrowParquetFieldIDKey}, []string{"2"})
+ rec := array.NewRecordBatch(arrow.NewSchema([]arrow.Field{{Name:
"payload", Type: vt, Nullable: true, Metadata: md}}, nil), []arrow.Array{varr},
2)
+ defer rec.Release()
+
+ iceSchema := iceberg.NewSchema(0, iceberg.NestedField{ID: 2, Name:
"payload", Type: iceberg.VariantType{}})
+ term, err := iceberg.Extract("payload", "$.a",
iceberg.PrimitiveTypes.Int64).Bind(iceSchema, true)
+ require.NoError(t, err)
+ col := iceberg.VariantExtractColumn{Term: term.(iceberg.BoundExtract),
FieldID: 100, Name: "_x", SourcePath: []string{"payload"}}
+ dt, err := TypeToArrowType(iceberg.PrimitiveTypes.Int64, false, false)
+ require.NoError(t, err)
+
+ require.Nil(t, tryShreddedTypedColumn(varr, vpath(t, col.Term), dt,
mem), "must not fast-path a root-residual object row")
+
+ got, _, err := buildExtractColumn(ctx, col, rec, mem)
+ require.NoError(t, err)
+ defer got.Release()
+ want, err := extractColumnValuesPerRow(ctx, varr, col, dt, mem)
+ require.NoError(t, err)
+ defer want.Release()
+
+ require.Truef(t, array.Equal(got, want), "got=%v want=%v", got, want)
+ require.False(t, got.IsNull(1), "root-residual row must be extracted,
not nulled")
+ require.EqualValues(t, 5, got.(*array.Int64).Value(1))
+}
+
+// TestFastPathDecimalScaleNearMissFallsBack: a shredded decimal whose scale
differs from the target (arrow.TypeEqual near-miss) must fall back, not return
the mistyped column.
+func TestFastPathDecimalScaleNearMissFallsBack(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(t.Context(), mem)
+
+ shredType := &arrow.Decimal128Type{Precision: 10, Scale: 2}
+ vt :=
extensions.NewShreddedVariantType(arrow.StructOf(arrow.Field{Name: "a", Type:
shredType}))
+
+ var b variant.Builder
+ require.NoError(t, b.Append(map[string]any{"a": int64(1)}))
+ v, err := b.Build()
+ require.NoError(t, err)
+ meta := v.Metadata().Bytes()
+
+ sb := array.NewStructBuilder(mem, vt.StorageType().(*arrow.StructType))
+ defer sb.Release()
+ metaB := sb.FieldBuilder(0).(*array.BinaryBuilder)
+ valB := sb.FieldBuilder(1).(*array.BinaryBuilder)
+ tvB := sb.FieldBuilder(2).(*array.StructBuilder)
+ aB := tvB.FieldBuilder(0).(*array.StructBuilder)
+ aValB := aB.FieldBuilder(0).(*array.BinaryBuilder)
+ aDecB := aB.FieldBuilder(1).(*array.Decimal128Builder)
+
+ sb.Append(true)
+ metaB.Append(meta)
+ valB.AppendNull()
+ tvB.Append(true)
+ aB.Append(true)
+ aValB.AppendNull()
+ aDecB.Append(decimal128.FromI64(150)) // 1.50 at scale 2
+
+ storage := sb.NewStructArray()
+ defer storage.Release()
+ varr := array.NewExtensionArrayWithStorage(vt,
storage).(*extensions.VariantArray)
+ defer varr.Release()
+
+ md := arrow.NewMetadata([]string{ArrowParquetFieldIDKey}, []string{"2"})
+ rec := array.NewRecordBatch(arrow.NewSchema([]arrow.Field{{Name:
"payload", Type: vt, Nullable: true, Metadata: md}}, nil), []arrow.Array{varr},
1)
+ defer rec.Release()
+
+ iceSchema := iceberg.NewSchema(0, iceberg.NestedField{ID: 2, Name:
"payload", Type: iceberg.VariantType{}})
+ term, err := iceberg.Extract("payload", "$.a",
iceberg.DecimalTypeOf(10, 4)).Bind(iceSchema, true)
+ require.NoError(t, err)
+ col := iceberg.VariantExtractColumn{Term: term.(iceberg.BoundExtract),
FieldID: 100, Name: "_x", SourcePath: []string{"payload"}}
+ dt, err := TypeToArrowType(iceberg.DecimalTypeOf(10, 4), false, false)
+ require.NoError(t, err)
+
+ require.Nil(t, tryShreddedTypedColumn(varr, vpath(t, col.Term), dt,
mem), "scale near-miss must not fast-path")
+
+ got, _, err := buildExtractColumn(ctx, col, rec, mem)
+ require.NoError(t, err)
+ defer got.Release()
+ want, err := extractColumnValuesPerRow(ctx, varr, col, dt, mem)
+ require.NoError(t, err)
+ defer want.Release()
+ require.Truef(t, array.Equal(got, want), "got=%v want=%v", got, want)
+}
+
+// TestFastPathTimestampTzNearMissFallsBack: a shredded tz-aware timestamp
extracted as a zoneless timestamp (arrow.TypeEqual near-miss on tz) must fall
back.
+func TestFastPathTimestampTzNearMissFallsBack(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(t.Context(), mem)
+
+ shredType := &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone:
"UTC"}
+ vt :=
extensions.NewShreddedVariantType(arrow.StructOf(arrow.Field{Name: "a", Type:
shredType}))
+
+ var b variant.Builder
+ require.NoError(t, b.Append(map[string]any{"a": int64(1)}))
+ v, err := b.Build()
+ require.NoError(t, err)
+ meta := v.Metadata().Bytes()
+
+ sb := array.NewStructBuilder(mem, vt.StorageType().(*arrow.StructType))
+ defer sb.Release()
+ metaB := sb.FieldBuilder(0).(*array.BinaryBuilder)
+ valB := sb.FieldBuilder(1).(*array.BinaryBuilder)
+ tvB := sb.FieldBuilder(2).(*array.StructBuilder)
+ aB := tvB.FieldBuilder(0).(*array.StructBuilder)
+ aValB := aB.FieldBuilder(0).(*array.BinaryBuilder)
+ aTsB := aB.FieldBuilder(1).(*array.TimestampBuilder)
+
+ sb.Append(true)
+ metaB.Append(meta)
+ valB.AppendNull()
+ tvB.Append(true)
+ aB.Append(true)
+ aValB.AppendNull()
+ aTsB.Append(arrow.Timestamp(1_000_000))
+
+ storage := sb.NewStructArray()
+ defer storage.Release()
+ varr := array.NewExtensionArrayWithStorage(vt,
storage).(*extensions.VariantArray)
+ defer varr.Release()
+
+ md := arrow.NewMetadata([]string{ArrowParquetFieldIDKey}, []string{"2"})
+ rec := array.NewRecordBatch(arrow.NewSchema([]arrow.Field{{Name:
"payload", Type: vt, Nullable: true, Metadata: md}}, nil), []arrow.Array{varr},
1)
+ defer rec.Release()
+
+ iceSchema := iceberg.NewSchema(0, iceberg.NestedField{ID: 2, Name:
"payload", Type: iceberg.VariantType{}})
+ term, err := iceberg.Extract("payload", "$.a",
iceberg.PrimitiveTypes.Timestamp).Bind(iceSchema, true)
+ require.NoError(t, err)
+ col := iceberg.VariantExtractColumn{Term: term.(iceberg.BoundExtract),
FieldID: 100, Name: "_x", SourcePath: []string{"payload"}}
+ dt, err := TypeToArrowType(iceberg.PrimitiveTypes.Timestamp, false,
false)
+ require.NoError(t, err)
+
+ require.Nil(t, tryShreddedTypedColumn(varr, vpath(t, col.Term), dt,
mem), "tz near-miss must not fast-path")
+
+ got, _, err := buildExtractColumn(ctx, col, rec, mem)
+ require.NoError(t, err)
+ defer got.Release()
+ want, err := extractColumnValuesPerRow(ctx, varr, col, dt, mem)
+ require.NoError(t, err)
+ defer want.Release()
+ require.Truef(t, array.Equal(got, want), "got=%v want=%v", got, want)
+}
+
+// sharesDataBuffers reports whether a and b share their non-validity buffers
(validity may differ after a mask merge).
+func sharesDataBuffers(a, b arrow.Array) bool {
+ ba, bb := a.Data().Buffers(), b.Data().Buffers()
+ if len(ba) != len(bb) || len(ba) < 2 {
+ return false
+ }
+ for i := 1; i < len(ba); i++ {
+ if ba[i] != bb[i] {
+ return false
+ }
+ }
+
+ return true
+}
+
+// TestFastPathWrapperFieldNullIsAbsentKey: a null per-key wrapper is an
absent key (shredding spec), so the fast path reads null even if the
typed_value child is left live by a non-conformant writer.
+func TestFastPathWrapperFieldNullIsAbsentKey(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(t.Context(), mem)
+
+ vt :=
extensions.NewShreddedVariantType(arrow.StructOf(arrow.Field{Name: "a", Type:
arrow.PrimitiveTypes.Int64}))
+ var b variant.Builder
+ require.NoError(t, b.Append(map[string]any{"a": int64(1)}))
+ v, err := b.Build()
+ require.NoError(t, err)
+ meta := v.Metadata().Bytes()
+
+ storageType := vt.StorageType().(*arrow.StructType)
+ tvType := storageType.Field(2).Type.(*arrow.StructType) // struct<a>
+ aType := tvType.Field(0).Type.(*arrow.StructType) // a:
struct<value, typed_value>
+
+ binArr := func(vals [][]byte) arrow.Array {
+ bb := array.NewBinaryBuilder(mem, arrow.BinaryTypes.Binary)
+ defer bb.Release()
+ for _, x := range vals {
+ if x == nil {
+ bb.AppendNull()
+ } else {
+ bb.Append(x)
+ }
+ }
+
+ return bb.NewArray()
+ }
+
+ i64b := array.NewInt64Builder(mem)
+ defer i64b.Release()
+ i64b.Append(1)
+ i64b.Append(5) // row 1 typed_value stays live under the (about-to-be)
null wrapper
+ aTyped := i64b.NewArray()
+ defer aTyped.Release()
+ aVal := binArr([][]byte{nil, nil})
+ defer aVal.Release()
+
+ // wrapper "a": row0 valid, row1 NULL (bit1=0), children left untouched
-> non-cascaded shape
+ wrapValidity := memory.NewBufferBytes([]byte{0x01})
+ aData := array.NewData(aType, 2, []*memory.Buffer{wrapValidity},
[]arrow.ArrayData{aVal.Data(), aTyped.Data()}, 1, 0)
+ defer aData.Release()
+ aStruct := array.NewStructData(aData)
+ defer aStruct.Release()
+
+ tvData := array.NewData(tvType, 2, []*memory.Buffer{nil},
[]arrow.ArrayData{aStruct.Data()}, 0, 0)
+ defer tvData.Release()
+ tvStruct := array.NewStructData(tvData)
+ defer tvStruct.Release()
+
+ metaArr := binArr([][]byte{meta, meta})
+ defer metaArr.Release()
+ rootVal := binArr([][]byte{nil, nil})
+ defer rootVal.Release()
+
+ storageData := array.NewData(storageType, 2, []*memory.Buffer{nil},
[]arrow.ArrayData{metaArr.Data(), rootVal.Data(), tvStruct.Data()}, 0, 0)
+ defer storageData.Release()
+ storage := array.NewStructData(storageData)
+ defer storage.Release()
+ varr := array.NewExtensionArrayWithStorage(vt,
storage).(*extensions.VariantArray)
+ defer varr.Release()
+
+ md := arrow.NewMetadata([]string{ArrowParquetFieldIDKey}, []string{"2"})
+ rec := array.NewRecordBatch(arrow.NewSchema([]arrow.Field{{Name:
"payload", Type: vt, Nullable: true, Metadata: md}}, nil), []arrow.Array{varr},
2)
+ defer rec.Release()
+
+ iceSchema := iceberg.NewSchema(0, iceberg.NestedField{ID: 2, Name:
"payload", Type: iceberg.VariantType{}})
+ term, err := iceberg.Extract("payload", "$.a",
iceberg.PrimitiveTypes.Int64).Bind(iceSchema, true)
+ require.NoError(t, err)
+ col := iceberg.VariantExtractColumn{Term: term.(iceberg.BoundExtract),
FieldID: 100, Name: "_x", SourcePath: []string{"payload"}}
+
+ got, _, err := buildExtractColumn(ctx, col, rec, mem)
+ require.NoError(t, err)
+ defer got.Release()
+ require.Truef(t, got.IsNull(1), "null wrapper is an absent key (spec):
must read null, got=%v", got)
+ require.EqualValues(t, 1, got.(*array.Int64).Value(0))
+}
+
+// TestFastPathSlicedOffsetFallsBack: a VariantArray sliced to a non-zero
offset must bail from the fast path and still match the per-row reference.
+func TestFastPathSlicedOffsetFallsBack(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(t.Context(), mem)
+
+ rec, col := buildVariantExtractRec(t, mem,
shredStruct(arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}),
+ "$.a", iceberg.PrimitiveTypes.Int64, []map[string]any{{"a":
int64(1)}, {"a": int64(2)}, {"a": int64(3)}})
+ defer rec.Release()
+ full := resolveVariantSource(rec, col.Term.Ref().Field().ID,
col.SourcePath).(*extensions.VariantArray)
+ sliced := array.NewSlice(full, 1,
int64(full.Len())).(*extensions.VariantArray)
+ defer sliced.Release()
+ dt, err := TypeToArrowType(iceberg.PrimitiveTypes.Int64, false, false)
+ require.NoError(t, err)
+
+ require.Nil(t, tryShreddedTypedColumn(sliced, vpath(t, col.Term), dt,
mem), "sliced (offset!=0) array must bail")
+
+ got, err := extractColumnValues(ctx, sliced, col,
iceberg.PrimitiveTypes.Int64, dt, mem)
+ require.NoError(t, err)
+ defer got.Release()
+ want, err := extractColumnValuesPerRow(ctx, sliced, col, dt, mem)
+ require.NoError(t, err)
+ defer want.Release()
+ require.Truef(t, array.Equal(got, want), "got=%v want=%v", got, want)
+}
+
+// TestExtractColumnValuesContextCancelled: a cancelled context returns its
error instead of a silent partial result.
+func TestExtractColumnValuesContextCancelled(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ base := compute.WithAllocator(context.Background(), mem)
+ ctx, cancel := context.WithCancel(base)
+ cancel()
+
+ rec, col := buildVariantExtractRec(t, mem,
shredStruct(arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}),
+ "$.a", iceberg.PrimitiveTypes.Int64, []map[string]any{{"a":
int64(1)}})
+ defer rec.Release()
+
+ _, _, err := buildExtractColumn(ctx, col, rec, mem)
Review Comment:
This pre-cancels and runs a single row, so it returns at the top-level
`ctx.Err()` guard and never reaches the `i%4096 == 0` polls in the middle-tier
loop or the per-row walk, which are exactly the checks the round-2 ask added.
If either poll were off-by-one or dropped in a refactor, this test would stay
green.
A subtest with >4096 rows, cancelling via a goroutine partway and asserting
`errors.Is(err, context.Canceled)`, would exercise the running loops. Not
blocking, but it's the coverage that matches the fix.
--
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]
