From 6169d4e42de8a9b9d5e5a19e1302d1ece0c98128 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20P=C3=BCtz?= Date: Wed, 29 Jul 2026 19:15:20 +0200 Subject: [PATCH 1/2] fix(table): rewrap extension arrays instead of casting when projecting Arrow has no extension-to-extension cast kernel. Projecting a column whose file and requested types are the same named extension with differing parameters fails with "cast_extension has no kernel matching input types", e.g. geoarrow.wkb spelling the same CRS as an SRID vs. an authority code. The values already carry the extension's storage encoding, so rewrapping that storage in the requested type is zero-copy and correct; storage is cast only on a width mismatch, and the requested type stays authoritative. --- table/arrow_utils.go | 33 +++++++ table/arrow_utils_test.go | 197 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 230 insertions(+) diff --git a/table/arrow_utils.go b/table/arrow_utils.go index 6e9e1e38c..d0beb8199 100644 --- a/table/arrow_utils.go +++ b/table/arrow_utils.go @@ -1089,6 +1089,10 @@ func (a *arrowProjectionVisitor) castIfNeeded(field iceberg.NestedField, vals ar targetType := a.typeToArrowType(field.Type) if !arrow.TypeEqual(targetType, vals.DataType()) { + if out, ok := a.rewrapExtension(targetType, vals); ok { + return out + } + switch field.Type.(type) { case iceberg.TimestampType: tt, tgtok := targetType.(*arrow.TimestampType) @@ -1131,6 +1135,35 @@ func (a *arrowProjectionVisitor) castIfNeeded(field iceberg.NestedField, vals ar return vals } +// rewrapExtension re-wraps vals in targetType when both are extension types +// with the same extension name but differing extension parameters, returning +// false when the types are not such a pair. Arrow's compute layer has no +// extension-to-extension cast kernel, and the values are already in the +// extension's storage encoding, so only the storage array is cast and that only +// when the storage types differ. +func (a *arrowProjectionVisitor) rewrapExtension(targetType arrow.DataType, vals arrow.Array) (arrow.Array, bool) { + tgtExt, ok := targetType.(arrow.ExtensionType) + if !ok { + return nil, false + } + + srcExt, ok := vals.DataType().(arrow.ExtensionType) + if !ok || srcExt.ExtensionName() != tgtExt.ExtensionName() { + return nil, false + } + + if arrow.TypeEqual(srcExt.StorageType(), tgtExt.StorageType()) { + return array.NewExtensionArrayWithStorage(tgtExt, + vals.(array.ExtensionArray).Storage()), true + } + + storage := retOrPanic(compute.CastArray(a.ctx, vals, + compute.SafeCastOptions(tgtExt.StorageType()))) + defer storage.Release() + + return array.NewExtensionArrayWithStorage(tgtExt, storage), true +} + func canDowncastTimestampPrecision(fileType, readType iceberg.Type, enabled bool) bool { if !enabled { return false diff --git a/table/arrow_utils_test.go b/table/arrow_utils_test.go index efc8367eb..48f603dd3 100644 --- a/table/arrow_utils_test.go +++ b/table/arrow_utils_test.go @@ -2555,3 +2555,200 @@ func TestToRequestedSchemaMissingNestedFieldID(t *testing.T) { }, nil) require.True(t, targetSchema.Equal(rec2.Schema()), "Schema is not perfectly equal") } + +// Some writers tag a non-PROJJSON CRS as an SRID, while the Iceberg schema +// resolves the same CRS as an authority code. Both describe the same WKB +// payload, so projection must not attempt a value cast. +func sridWKBType(t *testing.T, storage arrow.DataType, crs string, edges geoarrow.EdgeInterpolation) arrow.ExtensionType { + t.Helper() + + meta := geoarrow.Metadata{ + CRS: jsonCRS(crs), + CRSType: geoarrow.CRSTypeSRID, + Edges: edges, + } + + switch storage.ID() { + case arrow.BINARY: + return geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage(), geoarrow.WKBWithMetadata(meta)) + case arrow.LARGE_BINARY: + return geoarrow.NewWKBType(geoarrow.WKBWithLargeBinaryStorage(), geoarrow.WKBWithMetadata(meta)) + default: + t.Fatalf("unsupported WKB storage type %s", storage) + + return nil + } +} + +func newExtensionArrayOverBinary(t *testing.T, mem memory.Allocator, dt arrow.ExtensionType, values [][]byte) arrow.Array { + t.Helper() + + bldr := array.NewBinaryBuilder(mem, dt.StorageType().(arrow.BinaryDataType)) + defer bldr.Release() + + for _, v := range values { + if v == nil { + bldr.AppendNull() + + continue + } + bldr.Append(v) + } + + storage := bldr.NewArray() + defer storage.Release() + + return array.NewExtensionArrayWithStorage(dt, storage) +} + +func binaryValueAt(t *testing.T, arr arrow.Array, i int) []byte { + t.Helper() + + switch a := arr.(type) { + case *array.Binary: + return a.Value(i) + case *array.LargeBinary: + return a.Value(i) + default: + t.Fatalf("expected binary-like array, got %T", arr) + + return nil + } +} + +func singleColumnRecord(field arrow.Field, arr arrow.Array) arrow.RecordBatch { + sc := arrow.NewSchema([]arrow.Field{field}, nil) + + return array.NewRecordBatch(sc, []arrow.Array{arr}, int64(arr.Len())) +} + +func TestToRequestedSchemaGeoExtensionRewrap(t *testing.T) { + // POINT(1 2) and POINT(3 4) little-endian WKB. + point12 := []byte{ + 0x01, 0x01, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xf0, 0x3f, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x40, + } + point34 := []byte{ + 0x01, 0x01, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x08, 0x40, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x40, + } + values := [][]byte{point12, nil, point34} + + geography, err := iceberg.GeographyTypeOf("OGC:CRS84", "spherical") + require.NoError(t, err) + + tests := []struct { + name string + icebergType iceberg.Type + sourceStorage arrow.DataType + sourceEdges geoarrow.EdgeInterpolation + useLargeTypes bool + }{ + { + name: "geometry_metadata_only", + icebergType: iceberg.GeometryType{}, + sourceStorage: arrow.BinaryTypes.Binary, + }, + { + name: "geometry_storage_width_mismatch", + icebergType: iceberg.GeometryType{}, + sourceStorage: arrow.BinaryTypes.Binary, + useLargeTypes: true, + }, + { + name: "geography_spherical_metadata_only", + icebergType: geography, + sourceStorage: arrow.BinaryTypes.Binary, + sourceEdges: geoarrow.EdgeSpherical, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + sourceType := sridWKBType(t, tt.sourceStorage, "OGC:CRS84", tt.sourceEdges) + source := newExtensionArrayOverBinary(t, mem, sourceType, values) + defer source.Release() + + rec := singleColumnRecord(arrow.Field{ + Name: "geom", Type: sourceType, Nullable: true, Metadata: fieldIDMeta("1"), + }, source) + defer rec.Release() + + sc := iceberg.NewSchema(0, iceberg.NestedField{ + ID: 1, Name: "geom", Type: tt.icebergType, Required: false, + }) + + opts := table.SchemaOptions{IncludeFieldIDs: true, UseLargeTypes: tt.useLargeTypes} + + ctx := compute.WithAllocator(context.Background(), mem) + out, err := table.ToRequestedSchema(ctx, sc, sc, rec, opts) + require.NoError(t, err) + defer out.Release() + + want, err := table.TypeToArrowTypeWithOptions(tt.icebergType, table.ArrowSchemaOptions{ + IncludeFieldIDs: opts.IncludeFieldIDs, + UseLargeTypes: opts.UseLargeTypes, + }) + require.NoError(t, err) + + got, ok := out.Column(0).(array.ExtensionArray) + require.True(t, ok, "expected extension array, got %T", out.Column(0)) + assert.True(t, arrow.TypeEqual(want, got.DataType()), "expected: %s\ngot: %s", want, got.DataType()) + + gotStorage := got.Storage() + require.Equal(t, len(values), gotStorage.Len()) + + for i, v := range values { + if v == nil { + assert.True(t, gotStorage.IsNull(i), "row %d should be null", i) + + continue + } + require.False(t, gotStorage.IsNull(i), "row %d should not be null", i) + assert.Equal(t, v, binaryValueAt(t, gotStorage, i)) + } + + if arrow.TypeEqual(tt.sourceStorage, got.DataType().(arrow.ExtensionType).StorageType()) { + srcBufs := source.(array.ExtensionArray).Storage().Data().Buffers() + gotBufs := gotStorage.Data().Buffers() + require.Equal(t, len(srcBufs), len(gotBufs)) + + for i := range srcBufs { + assert.Same(t, srcBufs[i], gotBufs[i], "buffer %d should be shared", i) + } + } + }) + } +} + +func TestToRequestedSchemaMismatchedExtensionNameNotRewrapped(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + sourceType := extensions.NewUUIDType() + bldr := array.NewExtensionBuilder(mem, sourceType) + defer bldr.Release() + bldr.AppendNull() + + source := bldr.NewArray() + defer source.Release() + + rec := singleColumnRecord(arrow.Field{ + Name: "geom", Type: sourceType, Nullable: true, Metadata: fieldIDMeta("1"), + }, source) + defer rec.Release() + + sc := iceberg.NewSchema(0, iceberg.NestedField{ + ID: 1, Name: "geom", Type: iceberg.GeometryType{}, Required: false, + }) + + ctx := compute.WithAllocator(context.Background(), mem) + _, err := table.ToRequestedSchema(ctx, sc, sc, rec, table.SchemaOptions{IncludeFieldIDs: true}) + require.Error(t, err) + assert.ErrorContains(t, err, "cast_extension") +} From e6bb3f9b638b0dbb8a3597842fc26fb57f4d3484 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20P=C3=BCtz?= Date: Thu, 30 Jul 2026 08:50:04 +0200 Subject: [PATCH 2/2] fix(table): address review feedback on extension rewrap Cast the source storage array directly on a width mismatch instead of relying on Arrow's extension-input cast registration, document the invariant the rewrap depends on (a shared extension name implies the storage payload is interpretation-compatible), make the geometry test cases state planar edges explicitly, and stop matching Arrow's internal kernel name in the negative test. --- table/arrow_utils.go | 7 +++++-- table/arrow_utils_test.go | 6 +++++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/table/arrow_utils.go b/table/arrow_utils.go index d0beb8199..335672e8e 100644 --- a/table/arrow_utils.go +++ b/table/arrow_utils.go @@ -1140,7 +1140,9 @@ func (a *arrowProjectionVisitor) castIfNeeded(field iceberg.NestedField, vals ar // false when the types are not such a pair. Arrow's compute layer has no // extension-to-extension cast kernel, and the values are already in the // extension's storage encoding, so only the storage array is cast and that only -// when the storage types differ. +// when the storage types differ. This is safe only because a shared extension +// name implies the storage payload is interpretation-compatible: parameters may +// differ, the encoded values do not. func (a *arrowProjectionVisitor) rewrapExtension(targetType arrow.DataType, vals arrow.Array) (arrow.Array, bool) { tgtExt, ok := targetType.(arrow.ExtensionType) if !ok { @@ -1157,7 +1159,8 @@ func (a *arrowProjectionVisitor) rewrapExtension(targetType arrow.DataType, vals vals.(array.ExtensionArray).Storage()), true } - storage := retOrPanic(compute.CastArray(a.ctx, vals, + storage := retOrPanic(compute.CastArray(a.ctx, + vals.(array.ExtensionArray).Storage(), compute.SafeCastOptions(tgtExt.StorageType()))) defer storage.Release() diff --git a/table/arrow_utils_test.go b/table/arrow_utils_test.go index 48f603dd3..4b616ee35 100644 --- a/table/arrow_utils_test.go +++ b/table/arrow_utils_test.go @@ -2650,11 +2650,13 @@ func TestToRequestedSchemaGeoExtensionRewrap(t *testing.T) { name: "geometry_metadata_only", icebergType: iceberg.GeometryType{}, sourceStorage: arrow.BinaryTypes.Binary, + sourceEdges: geoarrow.EdgePlanar, }, { name: "geometry_storage_width_mismatch", icebergType: iceberg.GeometryType{}, sourceStorage: arrow.BinaryTypes.Binary, + sourceEdges: geoarrow.EdgePlanar, useLargeTypes: true, }, { @@ -2749,6 +2751,8 @@ func TestToRequestedSchemaMismatchedExtensionNameNotRewrapped(t *testing.T) { ctx := compute.WithAllocator(context.Background(), mem) _, err := table.ToRequestedSchema(ctx, sc, sc, rec, table.SchemaOptions{IncludeFieldIDs: true}) + // The failure comes from Arrow's cast internals, whose message is not a + // stable surface; what matters is that mismatched extension names are not + // silently rewrapped. require.Error(t, err) - assert.ErrorContains(t, err, "cast_extension") }