nssalian commented on code in PR #2002:
URL: https://github.com/apache/iceberg-go/pull/2002#discussion_r4020195505
##########
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:
This is done as well. `mergeValidity` sets a `badOffset` flag on any
non-zero child offset and the fast path bails; `rootResidualHidesRows` bails on
non-zero `tv`/`uv` offset. Test: `TestFastPathSlicedOffsetFallsBack` (slices a
`VariantArray` to a non-zero offset, asserts the fast path bails and matches
per-row).
##########
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:
Done. `defer checked.AssertSize(t, 0)` (runs last), `defer tbl.Release()`
after the ReadTable check, `defer out.Release()`.
##########
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:
Done. `TestMiddleTierColumnarCast` shreds `a` as int32 and extracts int64,
so the fast path bails on `TypeEqual` and the array stays shredded, driving
`compute.VariantGet` + the cast; it asserts the promoted non-null values (`1`,
`2`), confirming the promotions Java's `castTo` does.
--
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]
