laskoviymishka commented on code in PR #2002:
URL: https://github.com/apache/iceberg-go/pull/2002#discussion_r4008128725
##########
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:
not a blocker, but worth a note. after this bump we write
`arrow.parquet.variant` into the Arrow IPC schema, so an older iceberg-go
(still on v18.7, matching only `parquet.variant`) won't recognize the field and
falls back to treating it as a plain struct. the Parquet-level VARIANT logical
type is still there, so Java/pqarrow and iceberg-rust read it fine; it's
specifically old iceberg-go readers in a mixed-version fleet during a rolling
upgrade.
reading is already handled on both names here, so there's nothing to change
in the code. I'd just add a CHANGELOG line noting that variant files written
after this release need iceberg-go >= this version to be read back as variant
by iceberg-go itself.
##########
table/variant_residual.go:
##########
@@ -91,18 +167,124 @@ 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
+ mergeValidity := func(arr arrow.Array) {
+ if arr.NullN() == 0 {
+ return
+ }
+ vb := arr.Data().Buffers()[0]
+ if vb == nil {
+ return
+ }
+ if mask == nil {
+ mask = bitutil.BitmapAndAlloc(mem, vb.Bytes(),
vb.Bytes(), 0, 0, int64(n), 0)
+
+ 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)
Review Comment:
I think there's a subtle correctness hole here. we merge validity for the
intermediate typed_value struct (`mergeValidity(st)`) and for the leaf
(`mergeValidity(cur)` after the loop), but never for `field` itself, the
per-key `{value, typed_value}` wrapper struct.
for a row where the key is absent from the object, `field.IsNull(i)` is
true, but its children's null bits aren't guaranteed set: the Arrow spec says
child data under a null struct entry is undefined. we only get away with it
because `VariantBuilder.AppendNull` and the Parquet reader cascade the null
into every child, which is exactly why CI is green. an externally-built
shredded array (or a third-party pqarrow writer) where `field.IsNull(i)` is
true but `field.Field(tvIdx).IsNull(i)` is false would return a live value for
a row that should be null, which flips residual eval so a row that `a IS NULL`
should exclude passes.
adding `mergeValidity(field)` right after the `fty` line closes it with one
extra bitmap AND. wdyt?
##########
table/variant_residual.go:
##########
@@ -68,7 +67,84 @@ 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
+}
+
+// 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 fast := tryShreddedTypedColumn(varr, col.Term.VariantPath(), dt,
mem); fast != nil {
+ return fast, nil
+ }
+ if !varr.IsShredded() {
+ return extractColumnValuesPerRow(varr, col, dt, mem)
+ }
+
+ extracted, err := compute.VariantGet(ctx, varr,
compute.VariantGetOptions{Path: col.Term.VariantPath()})
Review Comment:
I think there's a real issue here. `err != nil` treats every `VariantGet`
failure as a navigation miss and drops to the per-row walk, but the comment
only means the null-on-miss case. `context.Canceled`/`DeadlineExceeded` and
allocator OOM all land in this branch too, and `extractColumnValuesPerRow`
doesn't check the context either, so a cancelled or timed-out scan quietly
finishes this batch (and the next ones) and looks like success to the caller.
I'd inspect the error before falling back: propagate anything wrapping
`context.Canceled`/`DeadlineExceeded` (or just check `ctx.Err()` at the top of
`extractColumnValues`), and fall through to per-row only on an actual
navigation miss.
##########
table/internal/parquet_files.go:
##########
@@ -711,9 +711,35 @@ 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 for every leaf
+// column of arrowSchema, resolving parquet leaf paths through arrow-go's own
schema conversion.
+func dictCostFallbackProps(arrowSchema *arrow.Schema, base
[]parquet.WriterProperty) ([]parquet.WriterProperty, error) {
+ parquetSchema, err := pqarrow.ToParquet(arrowSchema,
parquet.NewWriterProperties(base...), pqarrow.DefaultWriterProps())
Review Comment:
this runs a full `pqarrow.ToParquet` conversion on every writer init just to
enumerate leaf column paths, and the caller in `getWriteProperties` rebuilds
essentially the same schema moments later. for wide schemas and short-lived
microbatch (one-file) writers that fixed cost gets paid per file and roughly
doubles the schema-conversion work.
not a blocker. could we walk the `arrow.Schema` leaves directly to get the
paths instead of materializing a whole parquet.Schema? wdyt?
##########
table/variant_shredded_write_test.go:
##########
@@ -1354,6 +1354,7 @@ func TestShreddedVariantExtractResidualNoLeak(t
*testing.T) {
require.NoError(t, err)
out.Release()
+ tbl.Release()
Review Comment:
the whole point of reading `tbl` on `checked` is to catch a fast-path
zero-copy leak, but `tbl.Release()` and `checked.AssertSize` here aren't
deferred, and there are several `require`s above them. if any of those fails,
`FailNow` skips both: the tracked memory leaks and the assertion never fires,
so the leak check passes silently on the exact failure it exists to catch.
`defer checked.AssertSize(t, 0)` first (so it runs last) and `defer
tbl.Release()` right after the `ReadTable` NoError check would make the guard
actually hold.
##########
table/variant_residual.go:
##########
@@ -91,18 +167,124 @@ 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
+ mergeValidity := func(arr arrow.Array) {
+ if arr.NullN() == 0 {
+ return
+ }
+ vb := arr.Data().Buffers()[0]
+ if vb == nil {
+ return
+ }
+ if mask == nil {
+ mask = bitutil.BitmapAndAlloc(mem, vb.Bytes(),
vb.Bytes(), 0, 0, int64(n), 0)
Review Comment:
tiny one: on first init `mask == nil`, so this ANDs `vb` with itself, which
is a copy done the expensive way. could we clone the bitmap on init and keep
the AND only for the merge path? it's the hot path, so worth the two lines.
##########
table/variant_residual.go:
##########
@@ -91,18 +167,124 @@ 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
+ mergeValidity := func(arr arrow.Array) {
+ if arr.NullN() == 0 {
+ return
+ }
+ vb := arr.Data().Buffers()[0]
+ if vb == nil {
+ return
+ }
+ if mask == nil {
+ mask = bitutil.BitmapAndAlloc(mem, vb.Bytes(),
vb.Bytes(), 0, 0, int64(n), 0)
+
+ 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 !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)
+ 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
+ }
+ tvb := tv.Data().Buffers()[0]
+ uvb := uv.Data().Buffers()[0]
+ for i := range varr.Len() {
+ valuePresent := uvb == nil || bitutil.BitIsSet(uvb.Bytes(), i)
+ if valuePresent && !bitutil.BitIsSet(tvb.Bytes(), i) {
Review Comment:
related fragility. the `varr.Data().Offset() != 0` guard up in
`tryShreddedTypedColumn` only covers the top-level array, but here we index
`tvb`/`uvb` with the raw `i`, and `mergeValidity` calls `BitmapAndAlloc(..., 0,
0, ...)` with zero offsets too. both assume the child arrays have offset 0.
`tv`/`uv` are direct children off the Parquet read path so that holds today,
but there's no guard, and a sliced child would silently read the wrong rows.
simplest fail-safe is to bail (presume residual present, skip the fast path)
if any child offset is non-zero, or thread `Data().Offset()+i` into the bit
index. either is fine, I'd just make it explicit rather than leaning on the
outer guard.
there's also no test that slices a VariantArray to a non-zero offset and
checks the fast path bails and matches the per-row reference; that fixture is
what would lock this down once it's guarded.
##########
table/variant_residual_fastpath_test.go:
##########
@@ -0,0 +1,434 @@
+// 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 (
+ "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...))
+}
+
+// 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,
col.Term.VariantPath(), dt, mem)
Review Comment:
`ref` isn't deferred, so if the `require.NoError` below fires `t.FailNow()`
while `ref != nil`, the `ref.Release()` at the bottom of the subtest never runs
and the deferred `mem.AssertSize(t, 0)` then reports a non-zero size, masking
the real failure with a spurious leak. a conditional `defer ref.Release()`
right after the probe (dropping the manual release at the end) keeps the
failure readable.
##########
variant_extract.go:
##########
@@ -29,6 +29,8 @@ type BoundExtract interface {
BoundTerm
Path() string
+ // VariantPath returns the term's member-name path for columnar
extraction via compute.VariantGet.
+ VariantPath() variant.VariantPath
Review Comment:
I'd pull `VariantPath()` off the public `BoundExtract` interface. it's added
to an exported interface in the root package, so it's a breaking change: any
downstream type implementing `BoundExtract` (mocks, engine adapters, the
transferia CDC layer) stops compiling until it adds the method. and the
`variant.VariantPath` return type leaks arrow-go's path representation into our
public contract for what's really an internal fast-path optimization.
since the only caller is `table/`, I'd keep it internal: a narrow unexported
`variantPathProvider interface { VariantPath() variant.VariantPath }` that
`table/` type-asserts against, or expose `[]string` and rebuild the path inside
`table/`. either keeps `BoundExtract` free of the arrow-go coupling. wdyt?
##########
table/variant_shredded_write_test.go:
##########
@@ -2068,10 +2069,124 @@ func TestBuildExtractColumnWrongType(t *testing.T) {
require.NoError(t, err)
col := iceberg.VariantExtractColumn{Term: term.(iceberg.BoundExtract),
FieldID: 100, Name: "_x"}
- _, _, err = buildExtractColumn(col, rec, mem)
+ _, _, err = buildExtractColumn(context.Background(), col, rec, mem)
require.ErrorIs(t, err, iceberg.ErrInvalidArgument)
}
+// TestBuildExtractColumnShreddedColumnar exercises the compute.VariantGet
fast path over a
+// shredded column and asserts per-row values match.
+func TestBuildExtractColumnShreddedColumnar(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(t.Context(), mem)
+
+ iceSchema := iceberg.NewSchema(0, iceberg.NestedField{ID: 2, Name:
"payload", Type: iceberg.VariantType{}})
+
+ shredded := extensions.NewShreddedVariantType(arrow.StructOf(
+ arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64},
+ ))
+ vb := extensions.NewVariantBuilder(mem, shredded)
+ for i := range 4 {
+ var b variant.Builder
+ require.NoError(t, b.Append(map[string]any{"a": int64(i * 10),
"city": "NYC"}))
+ 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()
+ defer rec.Release()
+
+ 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"}}
+
+ arr, _, err := buildExtractColumn(ctx, col, rec, mem)
+ require.NoError(t, err)
+ defer arr.Release()
+ require.Equal(t, 4, arr.Len())
+ got := arr.(*array.Int64)
+ for i := range 4 {
+ require.EqualValues(t, i*10, got.Value(i), "row %d", i)
+ }
+}
+
+// TestBuildExtractColumnKeepsIcebergCast guards that the columnar path casts
with iceberg's
+// restrictive CastVariantLiteral, not arrow-go's permissive cast: an int64
extracted as float64
+// is null under iceberg's cast (no int->float coercion), whereas arrow-go's
cast would yield 5.0.
+func TestBuildExtractColumnKeepsIcebergCast(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(t.Context(), mem)
+
+ iceSchema := iceberg.NewSchema(0, iceberg.NestedField{ID: 2, Name:
"payload", Type: iceberg.VariantType{}})
+
+ vb := extensions.NewVariantBuilder(mem,
extensions.NewDefaultVariantType())
Review Comment:
this reads well but I don't think it reaches the tier it's aimed at.
`TestBuildExtractColumnKeepsIcebergCast` builds an unshredded variant
(`NewDefaultVariantType`), so `extractColumnValues` takes the
`!varr.IsShredded()` branch straight to `extractColumnValuesPerRow` and never
touches `compute.VariantGet` + `CastVariantLiteral`. and
`TestBuildExtractColumnShreddedColumnar` just above, despite its comment saying
it 'exercises the compute.VariantGet fast path', shreds `a` to exactly Int64
and extracts `$.a` as Int64, so `tryShreddedTypedColumn` matches on `TypeEqual`
and returns the leaf zero-copy before `VariantGet` is ever called. so the
middle columnar tier, and iceberg's cast on it, is still unexercised; both of
these would pass with that whole branch deleted.
to actually drive it we'd want a shredded array where the fast path bails
but `VariantGet` resolves: shred `a` to a type that differs from the extract
target so `tryShreddedTypedColumn` fails the `TypeEqual` check, then assert the
columnar output matches the per-row reference. a promotion case (say an int32
typed leaf extracted as int64) would also confirm the cast on that path covers
the promotions Java's `castTo` does rather than returning null. 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]
