You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iceberg.apache.org by ze...@apache.org on 2026/08/17 20:01:01 UTC
(iceberg-go) branch main updated: fix(table): normalize EWKB geometry values on write (#1670)
This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-go.git
The following commit(s) were added to refs/heads/main by this push:
new e15f6e061 fix(table): normalize EWKB geometry values on write (#1670)
e15f6e061 is described below
commit e15f6e06169973d498fe7c818cad53f805a0eece
Author: Minh Vu <vu...@gmail.com>
AuthorDate: Mon Aug 17 22:00:55 2026 +0200
fix(table): normalize EWKB geometry values on write (#1670)
* fix(table): normalize EWKB geometry values on write
* fix(table): use writer allocator for EWKB normalization
* fix(table): normalize nested EWKB values on write
* style(table): satisfy lint in nested geo normalization
* style(table): satisfy nested geo lint
* fix(table): preserve nested geo slice semantics
* fix(table): skip unreachable nested geo values
* fix(table): cache geo normalization columns
* test(table): cover cached geo normalization
* style(table): satisfy integer range lint
---
table/internal/geo_codec.go | 29 +-
table/internal/geo_codec_internal_test.go | 110 ++++-
table/internal/geo_normalize_internal_test.go | 603 ++++++++++++++++++++++++++
table/internal/parquet_files.go | 437 +++++++++++++++++--
table/internal/parquet_files_internal_test.go | 28 ++
table/internal/parquet_files_test.go | 88 ++++
6 files changed, 1258 insertions(+), 37 deletions(-)
diff --git a/table/internal/geo_codec.go b/table/internal/geo_codec.go
index dd400f5a9..48de8c119 100644
--- a/table/internal/geo_codec.go
+++ b/table/internal/geo_codec.go
@@ -26,6 +26,7 @@ import (
"github.com/twpayne/go-geom"
"github.com/twpayne/go-geom/encoding/ewkb"
"github.com/twpayne/go-geom/encoding/wkb"
+ "github.com/twpayne/go-geom/encoding/wkbcommon"
)
// Geo bounding box dimension indices. X is longitude/easting, Y is
@@ -39,6 +40,8 @@ const (
geoNumDims
)
+var wkbEmptyPointOption = wkbcommon.WKBOptionEmptyPointHandling(wkbcommon.EmptyPointHandlingNaN)
+
// geoBoundsAccumulator computes a geospatial bounding box from a stream of WKB
// values. Iceberg stores geometry/geography column bounds using the single-value
// serialization for geospatial types (spec Appendix D): the concatenation of
@@ -107,13 +110,35 @@ const (
// dimension in the type value itself (PointZ = 1001); some writers instead emit
// EWKB, which flags Z/M in the high bits of the type word and may embed an SRID
// after it. Both yield the same coordinates, so bounds do not depend on the
-// encoding. Stored values are never rewritten - this decode is read-only.
+// encoding.
func decodeWKB(data []byte) (geom.T, error) {
if isEWKB(data) {
return ewkb.Unmarshal(data)
}
- return wkb.Unmarshal(data)
+ // ISO empty points use NaN coordinates; keep them as empty geometries so the
+ // bounds accumulator can skip them instead of rejecting the value.
+ return wkb.Unmarshal(data, wkbEmptyPointOption)
+}
+
+func normalizeWKB(data []byte) ([]byte, error) {
+ if !isEWKB(data) {
+ return data, nil
+ }
+
+ // EWKB decoding already preserves empty-point NaNs; the explicit option is
+ // needed on the ISO marshal side so those coordinates remain an empty point.
+ g, err := ewkb.Unmarshal(data)
+ if err != nil {
+ return nil, err
+ }
+
+ var order binary.ByteOrder = binary.LittleEndian
+ if data[0] == wkbBigEndian {
+ order = binary.BigEndian
+ }
+
+ return wkb.Marshal(g, order, wkbEmptyPointOption)
}
// isEWKB reports whether data's type word carries EWKB flags. The two decoders
diff --git a/table/internal/geo_codec_internal_test.go b/table/internal/geo_codec_internal_test.go
index 8845b569f..990ac9885 100644
--- a/table/internal/geo_codec_internal_test.go
+++ b/table/internal/geo_codec_internal_test.go
@@ -242,14 +242,16 @@ func TestGeoBoundsAccumulatorInvalidWKB(t *testing.T) {
// of the type word (ewkbFlagZ etc., shared with geo_codec.go) and optionally
// embeds an SRID after it.
const (
- wkbPoint = 1
- wkbLineString = 2
- wkbGeometryCollection = 7
- wkbPointZ = 1001
- wkbPointM = 2001
- wkbPointZM = 3001
- wkbLineStringZ = 1002
- wkbGeometryCollectionZ = 1007
+ wkbPoint = 1
+ wkbLineString = 2
+ wkbGeometryCollection = 7
+ wkbPointZ = 1001
+ wkbPointM = 2001
+ wkbPointZM = 3001
+ wkbLineStringZ = 1002
+ wkbLineStringZM = 3002
+ wkbGeometryCollectionZ = 1007
+ wkbGeometryCollectionZM = 3007
)
// wkbBuilder assembles a WKB value byte by byte: the byte-order marker, then
@@ -523,6 +525,98 @@ func TestGeoBoundsAccumulatorEWKBMatchesISO(t *testing.T) {
}
}
+func TestNormalizeWKB(t *testing.T) {
+ tests := []struct {
+ name string
+ in []byte
+ want []byte
+ }{
+ {
+ name: "little endian z",
+ in: newWKBBuilder(wkbPoint|ewkbFlagZ).f64(1, 2, 3).bytes(),
+ want: newWKBBuilder(wkbPointZ).f64(1, 2, 3).bytes(),
+ },
+ {
+ name: "big endian z",
+ in: newXDRWKBBuilder(wkbPoint|ewkbFlagZ).f64(1, 2, 3).bytes(),
+ want: newXDRWKBBuilder(wkbPointZ).f64(1, 2, 3).bytes(),
+ },
+ {
+ name: "embedded srid",
+ in: newWKBBuilder(wkbPoint|ewkbFlagSRID).u32(4326).f64(1, 2).bytes(),
+ want: newWKBBuilder(wkbPoint).f64(1, 2).bytes(),
+ },
+ {
+ name: "iso unchanged",
+ in: newWKBBuilder(wkbPointZ).f64(1, 2, 3).bytes(),
+ want: newWKBBuilder(wkbPointZ).f64(1, 2, 3).bytes(),
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got, err := normalizeWKB(tt.in)
+ require.NoError(t, err)
+ require.Equal(t, tt.want, got)
+ if tt.name == "iso unchanged" {
+ require.Same(t, &tt.in[0], &got[0])
+ }
+ })
+ }
+}
+
+func TestNormalizeWKBEmptyPoints(t *testing.T) {
+ tests := []struct {
+ name string
+ flags uint32
+ isoType uint32
+ coords []float64
+ }{
+ {name: "xy", flags: 0, isoType: wkbPoint, coords: []float64{geom.PointEmptyCoord(), geom.PointEmptyCoord()}},
+ {name: "z", flags: ewkbFlagZ, isoType: wkbPointZ, coords: []float64{geom.PointEmptyCoord(), geom.PointEmptyCoord(), geom.PointEmptyCoord()}},
+ {name: "m", flags: ewkbFlagM, isoType: wkbPointM, coords: []float64{geom.PointEmptyCoord(), geom.PointEmptyCoord(), geom.PointEmptyCoord()}},
+ {name: "zm", flags: ewkbFlagZ | ewkbFlagM, isoType: wkbPointZM, coords: []float64{geom.PointEmptyCoord(), geom.PointEmptyCoord(), geom.PointEmptyCoord(), geom.PointEmptyCoord()}},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ in := newWKBBuilder(wkbPoint | tt.flags | ewkbFlagSRID).u32(4326).f64(tt.coords...).bytes()
+ want := newWKBBuilder(tt.isoType).f64(tt.coords...).bytes()
+
+ got, err := normalizeWKB(in)
+ require.NoError(t, err)
+ require.Equal(t, want, got)
+
+ decoded, err := decodeWKB(got)
+ require.NoError(t, err)
+ require.True(t, decoded.Empty())
+ acc := newGeoBoundsAccumulator(false)
+ require.NoError(t, acc.AddWKB(got))
+ })
+ }
+}
+
+func TestNormalizeWKBGeometryCollection(t *testing.T) {
+ in := newXDRWKBBuilder(wkbGeometryCollection|ewkbFlagZ|ewkbFlagM|ewkbFlagSRID).
+ u32(4326).
+ u32(2).
+ nested(
+ newXDRWKBBuilder(wkbPoint|ewkbFlagZ|ewkbFlagM|ewkbFlagSRID).u32(4326).f64(1, 2, 3, 4).bytes(),
+ newXDRWKBBuilder(wkbLineString|ewkbFlagZ|ewkbFlagM|ewkbFlagSRID).u32(4326).u32(2).f64(5, 6, 7, 8, 9, 10, 11, 12).bytes(),
+ ).bytes()
+ want := newXDRWKBBuilder(wkbGeometryCollectionZM).
+ u32(2).
+ nested(
+ newXDRWKBBuilder(wkbPointZM).f64(1, 2, 3, 4).bytes(),
+ newXDRWKBBuilder(wkbLineStringZM).u32(2).f64(5, 6, 7, 8, 9, 10, 11, 12).bytes(),
+ ).bytes()
+
+ got, err := normalizeWKB(in)
+ require.NoError(t, err)
+ require.Equal(t, want, got)
+ require.False(t, isEWKB(got))
+}
+
// TestGeoBoundsAccumulatorRejectsInvalidWKB verifies that malformed values still
// error rather than panicking or silently contributing no coordinates.
//
diff --git a/table/internal/geo_normalize_internal_test.go b/table/internal/geo_normalize_internal_test.go
new file mode 100644
index 000000000..ef0bad28f
--- /dev/null
+++ b/table/internal/geo_normalize_internal_test.go
@@ -0,0 +1,603 @@
+// 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 internal
+
+import (
+ "fmt"
+ "slices"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/geoarrow/geoarrow-go"
+ "github.com/stretchr/testify/require"
+)
+
+func TestNormalizeWKBArray(t *testing.T) {
+ typeDef := geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())
+ builder := array.NewBinaryBuilder(memory.DefaultAllocator, arrow.BinaryTypes.Binary)
+ builder.Append(newWKBBuilder(wkbPoint|ewkbFlagSRID).u32(4326).f64(1, 2).bytes())
+ builder.Append(newWKBBuilder(wkbPointZ).f64(3, 4, 5).bytes())
+ builder.AppendNull()
+ storage := builder.NewArray()
+ builder.Release()
+ ext := array.NewExtensionArrayWithStorage(typeDef, storage).(array.ExtensionArray)
+ storage.Release()
+ defer ext.Release()
+
+ normalized, changed, err := normalizeWKBArray(ext, memory.DefaultAllocator)
+ require.NoError(t, err)
+ require.True(t, changed)
+ defer normalized.Release()
+
+ arr := normalized.(*geoarrow.WKBArray)
+ require.Equal(t, newWKBBuilder(wkbPoint).f64(1, 2).bytes(), []byte(arr.Value(0)))
+ require.Equal(t, newWKBBuilder(wkbPointZ).f64(3, 4, 5).bytes(), []byte(arr.Value(1)))
+ require.True(t, arr.IsNull(2))
+}
+
+func TestNormalizeWKBArrayUsesProvidedAllocator(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ typeDef := geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())
+ builder := array.NewBinaryBuilder(mem, arrow.BinaryTypes.Binary)
+ builder.Append(newWKBBuilder(wkbPoint|ewkbFlagSRID).u32(4326).f64(1, 2).bytes())
+ storage := builder.NewArray()
+ builder.Release()
+ ext := array.NewExtensionArrayWithStorage(typeDef, storage).(array.ExtensionArray)
+ storage.Release()
+ defer ext.Release()
+
+ normalized, changed, err := normalizeWKBArray(ext, mem)
+ require.NoError(t, err)
+ require.True(t, changed)
+ require.NotZero(t, mem.CurrentAlloc())
+ normalized.Release()
+}
+
+func TestNormalizeNestedArray(t *testing.T) {
+ mem := memory.DefaultAllocator
+ typeDef := geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())
+ ewkb := newWKBBuilder(wkbPoint|ewkbFlagSRID).u32(4326).f64(1, 2).bytes()
+ iso := newWKBBuilder(wkbPoint).f64(3, 4).bytes()
+
+ geo := newGeoWKBArray(t, mem, typeDef, ewkb, iso)
+ structArray, err := array.NewStructArrayWithFields(
+ []arrow.Array{geo},
+ []arrow.Field{{Name: "geom", Type: typeDef, Nullable: true}},
+ )
+ require.NoError(t, err)
+ geo.Release()
+
+ listArray := newGeoListArray(t, mem, typeDef, false, ewkb, iso)
+ largeListArray := newGeoListArray(t, mem, typeDef, true, ewkb, iso)
+ fixedListArray := newGeoFixedSizeListArray(t, mem, typeDef, ewkb, iso)
+ mapArray := newGeoMapArray(t, mem, typeDef, ewkb, iso)
+ deepArray, err := array.NewStructArrayWithFields(
+ []arrow.Array{listArray},
+ []arrow.Field{{Name: "nested", Type: listArray.DataType(), Nullable: true}},
+ )
+ require.NoError(t, err)
+
+ tests := []struct {
+ name string
+ arr arrow.Array
+ }{
+ {name: "struct", arr: structArray},
+ {name: "list", arr: listArray},
+ {name: "large list", arr: largeListArray},
+ {name: "fixed size list", arr: fixedListArray},
+ {name: "map", arr: mapArray},
+ {name: "deep nesting", arr: deepArray},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ normalized, changed, err := normalizeNestedArray(tt.arr, mem)
+ require.NoError(t, err)
+ require.True(t, changed)
+ defer normalized.Release()
+ assertNoEWKB(t, normalized)
+ })
+ }
+
+ structArray.Release()
+ listArray.Release()
+ largeListArray.Release()
+ fixedListArray.Release()
+ mapArray.Release()
+ deepArray.Release()
+}
+
+func TestNormalizeNestedArrayPreservesSliceAndNulls(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ typeDef := geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())
+ ewkb := newWKBBuilder(wkbPoint|ewkbFlagSRID).u32(4326).f64(1, 2).bytes()
+ iso := newWKBBuilder(wkbPoint).f64(3, 4).bytes()
+
+ builder := array.NewListBuilder(mem, arrow.BinaryTypes.Binary)
+ values := builder.ValueBuilder().(*array.BinaryBuilder)
+ builder.Append(true)
+ values.Append(iso)
+ builder.AppendNull()
+ builder.Append(true)
+ values.Append(ewkb)
+ raw := builder.NewArray()
+ builder.Release()
+ geo := newGeoWKBArray(t, mem, typeDef, iso, ewkb)
+ data := array.NewData(arrow.ListOf(typeDef), raw.Len(), raw.Data().Buffers(),
+ []arrow.ArrayData{geo.Data()}, raw.Data().NullN(), raw.Data().Offset())
+ base := array.NewListData(data)
+ data.Release()
+ raw.Release()
+ geo.Release()
+ defer base.Release()
+
+ sliced := array.NewSlice(base, 1, 3)
+ defer sliced.Release()
+ normalized, changed, err := normalizeNestedArray(sliced, mem)
+ require.NoError(t, err)
+ require.True(t, changed)
+ defer normalized.Release()
+
+ require.Zero(t, normalized.Data().Offset())
+ require.True(t, normalized.IsNull(0))
+ require.True(t, normalized.IsValid(1))
+ storage := normalized.(*array.List).ListValues().(array.ExtensionArray).Storage().(wkbStorage)
+ require.Equal(t, newWKBBuilder(wkbPoint).f64(1, 2).bytes(), storage.Value(0))
+ assertNoEWKB(t, normalized)
+}
+
+func TestNormalizeNestedArrayIgnoresValuesOutsideListSlice(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ typeDef := geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())
+ malformedEWKB := newWKBBuilder(wkbPoint | ewkbFlagSRID).u32(4326).bytes()
+ iso := newWKBBuilder(wkbPoint).f64(3, 4).bytes()
+
+ builder := array.NewListBuilder(mem, arrow.BinaryTypes.Binary)
+ values := builder.ValueBuilder().(*array.BinaryBuilder)
+ builder.Append(true)
+ values.Append(malformedEWKB)
+ builder.Append(true)
+ values.Append(iso)
+ raw := builder.NewArray()
+ builder.Release()
+ geo := newGeoWKBArray(t, mem, typeDef, malformedEWKB, iso)
+ data := array.NewData(arrow.ListOf(typeDef), raw.Len(), raw.Data().Buffers(),
+ []arrow.ArrayData{geo.Data()}, raw.Data().NullN(), raw.Data().Offset())
+ base := array.NewListData(data)
+ data.Release()
+ raw.Release()
+ geo.Release()
+ defer base.Release()
+
+ sliced := array.NewSlice(base, 1, 2)
+ defer sliced.Release()
+ normalized, changed, err := normalizeNestedArray(sliced, mem)
+ require.NoError(t, err)
+ require.False(t, changed)
+ require.Nil(t, normalized)
+}
+
+func TestNormalizeNestedArrayPreservesStructSlice(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ typeDef := geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())
+ ewkb := newWKBBuilder(wkbPoint|ewkbFlagSRID).u32(4326).f64(1, 2).bytes()
+ iso := newWKBBuilder(wkbPoint).f64(3, 4).bytes()
+ geo := newGeoWKBArray(t, mem, typeDef, iso, ewkb, iso)
+ valuesBuilder := array.NewInt32Builder(mem)
+ valuesBuilder.AppendValues([]int32{1, 2, 3}, nil)
+ values := valuesBuilder.NewArray()
+ valuesBuilder.Release()
+
+ input, err := array.NewStructArrayWithFields(
+ []arrow.Array{geo, values},
+ []arrow.Field{
+ {Name: "geom", Type: typeDef, Nullable: true},
+ {Name: "value", Type: arrow.PrimitiveTypes.Int32, Nullable: false},
+ },
+ )
+ require.NoError(t, err)
+ geo.Release()
+ values.Release()
+ defer input.Release()
+
+ sliced := array.NewSlice(input, 1, 3)
+ defer sliced.Release()
+ normalized, changed, err := normalizeNestedArray(sliced, mem)
+ require.NoError(t, err)
+ require.True(t, changed)
+ defer normalized.Release()
+
+ result := normalized.(*array.Struct)
+ require.Equal(t, 2, result.Len())
+ require.Zero(t, result.Data().Offset())
+ require.True(t, result.IsValid(0))
+ require.True(t, result.IsValid(1))
+ storage := result.Field(0).(array.ExtensionArray).Storage().(wkbStorage)
+ require.Equal(t, newWKBBuilder(wkbPoint).f64(1, 2).bytes(), storage.Value(0))
+ require.Equal(t, iso, storage.Value(1))
+ assertNoEWKB(t, result)
+}
+
+func TestNormalizeNestedArrayIgnoresStructNullChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ typeDef := geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())
+ malformedEWKB := newWKBBuilder(wkbPoint | ewkbFlagSRID).u32(4326).bytes()
+ ewkb := newWKBBuilder(wkbPoint|ewkbFlagSRID).u32(4326).f64(1, 2).bytes()
+ geo := newGeoWKBArray(t, mem, typeDef, malformedEWKB, ewkb)
+ validity := memory.NewBufferBytes([]byte{0b00000010})
+
+ input, err := array.NewStructArrayWithFieldsAndNulls(
+ []arrow.Array{geo},
+ []arrow.Field{{Name: "geom", Type: typeDef, Nullable: true}},
+ validity, 1, 0,
+ )
+ require.NoError(t, err)
+ validity.Release()
+ geo.Release()
+ defer input.Release()
+
+ normalized, changed, err := normalizeNestedArray(input, mem)
+ require.NoError(t, err)
+ require.True(t, changed)
+ defer normalized.Release()
+
+ result := normalized.(*array.Struct)
+ require.True(t, result.IsNull(0))
+ require.True(t, result.IsValid(1))
+ storage := result.Field(0).(array.ExtensionArray).Storage().(wkbStorage)
+ require.Equal(t, malformedEWKB, storage.Value(0))
+ require.Equal(t, newWKBBuilder(wkbPoint).f64(1, 2).bytes(), storage.Value(1))
+}
+
+func TestNormalizeNestedArrayIgnoresNullListChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ typeDef := geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())
+ malformedEWKB := newWKBBuilder(wkbPoint | ewkbFlagSRID).u32(4326).bytes()
+ ewkb := newWKBBuilder(wkbPoint|ewkbFlagSRID).u32(4326).f64(1, 2).bytes()
+ geo := newGeoWKBArray(t, mem, typeDef, malformedEWKB, ewkb)
+ validity := memory.NewBufferBytes([]byte{0b00000010})
+ offsets := memory.NewBufferBytes(arrow.Int32Traits.CastToBytes([]int32{0, 1, 2}))
+ data := array.NewData(arrow.ListOf(typeDef), 2,
+ []*memory.Buffer{validity, offsets}, []arrow.ArrayData{geo.Data()}, 1, 0)
+ base := array.NewListData(data)
+ data.Release()
+ validity.Release()
+ offsets.Release()
+ geo.Release()
+ defer base.Release()
+
+ normalized, changed, err := normalizeNestedArray(base, mem)
+ require.NoError(t, err)
+ require.True(t, changed)
+ defer normalized.Release()
+
+ result := normalized.(*array.List)
+ require.Zero(t, result.Data().Offset())
+ require.True(t, result.IsNull(0))
+ require.True(t, result.IsValid(1))
+ start, end := result.ValueOffsets(0)
+ require.Equal(t, int64(0), start)
+ require.Equal(t, int64(1), end)
+ start, end = result.ValueOffsets(1)
+ require.Equal(t, int64(1), start)
+ require.Equal(t, int64(2), end)
+ storage := result.ListValues().(array.ExtensionArray).Storage().(wkbStorage)
+ require.Equal(t, malformedEWKB, storage.Value(0))
+ require.Equal(t, newWKBBuilder(wkbPoint).f64(1, 2).bytes(), storage.Value(1))
+}
+
+func TestNormalizeNestedArrayIgnoresAncestorNullChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ typeDef := geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())
+ malformedEWKB := newWKBBuilder(wkbPoint | ewkbFlagSRID).u32(4326).bytes()
+ ewkb := newWKBBuilder(wkbPoint|ewkbFlagSRID).u32(4326).f64(1, 2).bytes()
+ geo := newGeoWKBArray(t, mem, typeDef, malformedEWKB, ewkb)
+ offsets := memory.NewBufferBytes(arrow.Int32Traits.CastToBytes([]int32{0, 1, 2}))
+ listData := array.NewData(arrow.ListOf(typeDef), 2,
+ []*memory.Buffer{nil, offsets}, []arrow.ArrayData{geo.Data()}, 0, 0)
+ list := array.NewListData(listData)
+ listData.Release()
+ offsets.Release()
+ geo.Release()
+
+ validity := memory.NewBufferBytes([]byte{0b00000010})
+ input, err := array.NewStructArrayWithFieldsAndNulls(
+ []arrow.Array{list},
+ []arrow.Field{{Name: "items", Type: list.DataType(), Nullable: true}},
+ validity, 1, 0,
+ )
+ require.NoError(t, err)
+ validity.Release()
+ list.Release()
+ defer input.Release()
+
+ normalized, changed, err := normalizeNestedArray(input, mem)
+ require.NoError(t, err)
+ require.True(t, changed)
+ defer normalized.Release()
+
+ result := normalized.(*array.Struct)
+ require.True(t, result.IsNull(0))
+ require.True(t, result.IsValid(1))
+ values := result.Field(0).(*array.List).ListValues().(array.ExtensionArray).Storage().(wkbStorage)
+ require.Equal(t, malformedEWKB, values.Value(0))
+ require.Equal(t, newWKBBuilder(wkbPoint).f64(1, 2).bytes(), values.Value(1))
+}
+
+func TestNormalizeNestedArrayReusesUnchangedChildren(t *testing.T) {
+ mem := memory.DefaultAllocator
+ typeDef := geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())
+ ewkb := newWKBBuilder(wkbPoint|ewkbFlagSRID).u32(4326).f64(1, 2).bytes()
+ geo := newGeoWKBArray(t, mem, typeDef, ewkb)
+ intsBuilder := array.NewInt32Builder(mem)
+ intsBuilder.AppendValues([]int32{1}, nil)
+ ints := intsBuilder.NewArray()
+ intsBuilder.Release()
+
+ input, err := array.NewStructArrayWithFields(
+ []arrow.Array{geo, ints},
+ []arrow.Field{
+ {Name: "geom", Type: typeDef, Nullable: true},
+ {Name: "value", Type: arrow.PrimitiveTypes.Int32, Nullable: false},
+ },
+ )
+ require.NoError(t, err)
+ geo.Release()
+ ints.Release()
+ defer input.Release()
+
+ normalized, changed, err := normalizeNestedArray(input, mem)
+ require.NoError(t, err)
+ require.True(t, changed)
+ defer normalized.Release()
+
+ result := normalized.(*array.Struct)
+ require.Same(t, input.Field(1).Data(), result.Field(1).Data())
+ assertNoEWKB(t, result)
+}
+
+func TestNormalizeGeoBatchNormalizesNestedEWKB(t *testing.T) {
+ mem := memory.DefaultAllocator
+ typeDef := geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())
+ ewkb := newWKBBuilder(wkbPoint|ewkbFlagSRID).u32(4326).f64(1, 2).bytes()
+ list := newGeoListArray(t, mem, typeDef, false, ewkb)
+ defer list.Release()
+
+ valuesBuilder := array.NewInt32Builder(mem)
+ valuesBuilder.Append(7)
+ values := valuesBuilder.NewArray()
+ valuesBuilder.Release()
+ defer values.Release()
+
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "value", Type: arrow.PrimitiveTypes.Int32, Nullable: false},
+ {Name: "nested", Type: list.DataType(), Nullable: true},
+ }, nil)
+ batch := array.NewRecordBatch(schema, []arrow.Array{values, list}, int64(list.Len()))
+ defer batch.Release()
+
+ originalStorage := list.(*array.List).ListValues().(array.ExtensionArray).Storage().(wkbStorage)
+ require.True(t, isEWKB(originalStorage.Value(0)))
+ writer := &ParquetFileWriter{mem: mem, geoNormalizeCols: collectGeoNormalizationColumns(schema)}
+ normalized, err := writer.normalizeGeoBatch(batch)
+ require.NoError(t, err)
+ require.NotSame(t, batch, normalized)
+ defer normalized.Release()
+ require.Same(t, batch.Column(0).Data(), normalized.Column(0).Data())
+ assertNoEWKB(t, normalized.Column(1))
+ require.True(t, isEWKB(originalStorage.Value(0)), "normalization must not mutate the input batch")
+}
+
+func TestCollectGeoNormalizationColumns(t *testing.T) {
+ typeDef := geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "plain", Type: arrow.PrimitiveTypes.Int32},
+ {Name: "direct", Type: typeDef},
+ {Name: "struct", Type: arrow.StructOf(arrow.Field{Name: "geom", Type: typeDef})},
+ {Name: "list", Type: arrow.ListOf(typeDef)},
+ {Name: "large list", Type: arrow.LargeListOf(typeDef)},
+ {Name: "fixed size list", Type: arrow.FixedSizeListOf(2, typeDef)},
+ {Name: "map", Type: arrow.MapOf(arrow.BinaryTypes.String, typeDef)},
+ {Name: "deep", Type: arrow.StructOf(arrow.Field{
+ Name: "items", Type: arrow.ListOf(arrow.StructOf(arrow.Field{Name: "geom", Type: typeDef})),
+ })},
+ {Name: "nested plain", Type: arrow.StructOf(arrow.Field{Name: "value", Type: arrow.ListOf(arrow.PrimitiveTypes.Int32)})},
+ }, nil)
+
+ require.Equal(t, []int{1, 2, 3, 4, 5, 6, 7}, collectGeoNormalizationColumns(schema))
+}
+
+func TestNormalizeGeoBatchSkipsNonGeoColumns(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ builder := array.NewInt32Builder(mem)
+ builder.Append(1)
+ values := builder.NewArray()
+ builder.Release()
+ defer values.Release()
+
+ schema := arrow.NewSchema([]arrow.Field{{Name: "value", Type: arrow.PrimitiveTypes.Int32}}, nil)
+ batch := array.NewRecordBatch(schema, []arrow.Array{values}, 1)
+ defer batch.Release()
+
+ writer := &ParquetFileWriter{mem: mem}
+ normalized, err := writer.normalizeGeoBatch(batch)
+ require.NoError(t, err)
+ require.Same(t, batch, normalized)
+}
+
+func TestNormalizeGeoBatchRejectsReachableMalformedEWKB(t *testing.T) {
+ mem := memory.DefaultAllocator
+ typeDef := geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())
+ malformedEWKB := newWKBBuilder(wkbPoint | ewkbFlagSRID).u32(4326).bytes()
+ geo := newGeoWKBArray(t, mem, typeDef, malformedEWKB)
+ defer geo.Release()
+
+ schema := arrow.NewSchema([]arrow.Field{{Name: "geom", Type: typeDef, Nullable: true}}, nil)
+ batch := array.NewRecordBatch(schema, []arrow.Array{geo}, 1)
+ defer batch.Release()
+
+ writer := &ParquetFileWriter{mem: mem, geoNormalizeCols: []int{0}}
+ _, err := writer.normalizeGeoBatch(batch)
+ require.ErrorContains(t, err, "normalizing geo values for column 0")
+}
+
+func newGeoWKBArray(t *testing.T, mem memory.Allocator, typeDef *geoarrow.WKBType, values ...[]byte) arrow.Array {
+ builder := array.NewBinaryBuilder(mem, arrow.BinaryTypes.Binary)
+ for _, value := range values {
+ if value == nil {
+ builder.AppendNull()
+ } else {
+ builder.Append(value)
+ }
+ }
+ storage := builder.NewArray()
+ builder.Release()
+ ext := array.NewExtensionArrayWithStorage(typeDef, storage)
+ storage.Release()
+
+ return ext
+}
+
+func newGeoListArray(t *testing.T, mem memory.Allocator, typeDef *geoarrow.WKBType, large bool, values ...[]byte) arrow.Array {
+ if large {
+ builder := array.NewLargeListBuilder(mem, arrow.BinaryTypes.Binary)
+ items := builder.ValueBuilder().(*array.BinaryBuilder)
+ for _, value := range values {
+ builder.Append(true)
+ items.Append(value)
+ }
+ raw := builder.NewArray()
+ builder.Release()
+
+ return replaceListValues(t, raw, typeDef, values, true)
+ }
+
+ builder := array.NewListBuilder(mem, arrow.BinaryTypes.Binary)
+ items := builder.ValueBuilder().(*array.BinaryBuilder)
+ for _, value := range values {
+ builder.Append(true)
+ items.Append(value)
+ }
+ raw := builder.NewArray()
+ builder.Release()
+
+ return replaceListValues(t, raw, typeDef, values, false)
+}
+
+func newGeoFixedSizeListArray(t *testing.T, mem memory.Allocator, typeDef *geoarrow.WKBType, values ...[]byte) arrow.Array {
+ builder := array.NewFixedSizeListBuilder(mem, 1, arrow.BinaryTypes.Binary)
+ items := builder.ValueBuilder().(*array.BinaryBuilder)
+ for _, value := range values {
+ builder.Append(true)
+ items.Append(value)
+ }
+ raw := builder.NewArray()
+ builder.Release()
+
+ return replaceListValues(t, raw, typeDef, values, false)
+}
+
+func newGeoMapArray(t *testing.T, mem memory.Allocator, typeDef *geoarrow.WKBType, values ...[]byte) arrow.Array {
+ builder := array.NewMapBuilder(mem, arrow.BinaryTypes.String, arrow.BinaryTypes.Binary, false)
+ keys := builder.KeyBuilder().(*array.StringBuilder)
+ items := builder.ItemBuilder().(*array.BinaryBuilder)
+ for idx, value := range values {
+ builder.Append(true)
+ keys.Append(fmt.Sprintf("key-%d", idx))
+ items.Append(value)
+ }
+ raw := builder.NewMapArray()
+ builder.Release()
+ geo := newGeoWKBArray(t, mem, typeDef, values...)
+ entryData := raw.Data().Children()[0]
+ entryChildren := slices.Clone(entryData.Children())
+ entryChildren[1] = geo.Data()
+ newEntryData := array.NewData(entryData.DataType(), entryData.Len(), entryData.Buffers(),
+ entryChildren, entryData.NullN(), entryData.Offset())
+ mapData := array.NewData(raw.DataType(), raw.Len(), raw.Data().Buffers(),
+ []arrow.ArrayData{newEntryData}, raw.Data().NullN(), raw.Data().Offset())
+ result := array.NewMapData(mapData)
+ mapData.Release()
+ newEntryData.Release()
+ geo.Release()
+ raw.Release()
+
+ return result
+}
+
+func replaceListValues(t *testing.T, raw arrow.Array, typeDef *geoarrow.WKBType, values [][]byte, large bool) arrow.Array {
+ t.Helper()
+ mem := memory.DefaultAllocator
+ geo := newGeoWKBArray(t, mem, typeDef, values...)
+ var listType arrow.DataType
+ var newArray func(arrow.ArrayData) arrow.Array
+ if large {
+ listType = arrow.LargeListOf(typeDef)
+ newArray = func(data arrow.ArrayData) arrow.Array { return array.NewLargeListData(data) }
+ } else if _, ok := raw.(*array.FixedSizeList); ok {
+ listType = arrow.FixedSizeListOf(1, typeDef)
+ newArray = func(data arrow.ArrayData) arrow.Array { return array.NewFixedSizeListData(data) }
+ } else {
+ listType = arrow.ListOf(typeDef)
+ newArray = func(data arrow.ArrayData) arrow.Array { return array.NewListData(data) }
+ }
+ data := array.NewData(listType, raw.Len(), raw.Data().Buffers(),
+ []arrow.ArrayData{geo.Data()}, raw.Data().NullN(), raw.Data().Offset())
+ result := newArray(data)
+ data.Release()
+ geo.Release()
+ raw.Release()
+
+ return result
+}
+
+func assertNoEWKB(t *testing.T, arr arrow.Array) {
+ t.Helper()
+ switch nested := arr.(type) {
+ case array.ExtensionArray:
+ storage := nested.Storage().(wkbStorage)
+ for idx := range storage.Len() {
+ if storage.IsNull(idx) {
+ continue
+ }
+ require.False(t, isEWKB(storage.Value(idx)))
+ }
+ case *array.Struct:
+ for idx := range nested.NumField() {
+ assertNoEWKB(t, nested.Field(idx))
+ }
+ case *array.Map:
+ assertNoEWKB(t, nested.ListValues())
+ case *array.List:
+ assertNoEWKB(t, nested.ListValues())
+ case *array.LargeList:
+ assertNoEWKB(t, nested.ListValues())
+ case *array.FixedSizeList:
+ assertNoEWKB(t, nested.ListValues())
+ }
+}
diff --git a/table/internal/parquet_files.go b/table/internal/parquet_files.go
index f0546e30d..a0332d9fd 100644
--- a/table/internal/parquet_files.go
+++ b/table/internal/parquet_files.go
@@ -33,6 +33,7 @@ import (
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/bitutil"
"github.com/apache/arrow-go/v18/arrow/compute"
"github.com/apache/arrow-go/v18/arrow/decimal128"
"github.com/apache/arrow-go/v18/arrow/extensions"
@@ -497,18 +498,20 @@ func writeDataFileBatches(w FileWriter, batches []arrow.RecordBatch) (iceberg.Da
// lifecycle. It writes Arrow record batches to a Parquet file and tracks bytes
// written for rolling file decisions.
type ParquetFileWriter struct {
- pqWriter *pqarrow.FileWriter
- counter *internal.CountingWriter
- fileCloser io.Closer
- fs iceio.WriteFileIO
- format parquetFormat
- info WriteFileInfo
- partition map[int]any
- colMapping map[string]int
- geoCols []geoColumn
- geoAccs map[int]*geoBoundsAccumulator
- arrowSchema *arrow.Schema
- rowGroupBytes int64
+ pqWriter *pqarrow.FileWriter
+ counter *internal.CountingWriter
+ fileCloser io.Closer
+ mem memory.Allocator
+ fs iceio.WriteFileIO
+ format parquetFormat
+ info WriteFileInfo
+ partition map[int]any
+ colMapping map[string]int
+ geoCols []geoColumn
+ geoNormalizeCols []int
+ geoAccs map[int]*geoBoundsAccumulator
+ arrowSchema *arrow.Schema
+ rowGroupBytes int64
}
// geoColumn locates a top-level geometry/geography column: its position in the
@@ -538,6 +541,40 @@ func collectGeoColumns(sc *arrow.Schema, colMapping map[string]int) []geoColumn
return result
}
+// collectGeoNormalizationColumns finds top-level columns whose supported Arrow
+// container types can contain a GeoArrow WKB value. The result is computed once
+// per writer so batches without geo columns take a constant-time fast path.
+func collectGeoNormalizationColumns(sc *arrow.Schema) []int {
+ var result []int
+ for i, f := range sc.Fields() {
+ if arrowTypeContainsWKB(f.Type) {
+ result = append(result, i)
+ }
+ }
+
+ return result
+}
+
+func arrowTypeContainsWKB(dt arrow.DataType) bool {
+ // Keep this traversal in sync with normalizeNestedArrayReachable. The table
+ // write path does not currently pass dictionary- or run-end-encoded WKB
+ // arrays; supporting those inputs requires unwrapping them in both places.
+ switch t := dt.(type) {
+ case *geoarrow.WKBType:
+ return true
+ case *arrow.StructType:
+ for _, field := range t.Fields() {
+ if arrowTypeContainsWKB(field.Type) {
+ return true
+ }
+ }
+ case *arrow.ListType, *arrow.LargeListType, *arrow.FixedSizeListType, *arrow.MapType:
+ return arrowTypeContainsWKB(dt.(arrow.ListLikeType).Elem())
+ }
+
+ return false
+}
+
// NewFileWriter creates a ParquetFileWriter that writes batches to a single
// Parquet file. Call Write to append batches, BytesWritten to check actual
// compressed file size, and Close to finalize and get the resulting DataFile.
@@ -583,6 +620,7 @@ func (p parquetFormat) NewFileWriter(ctx context.Context, fs iceio.WriteFileIO,
}
geoCols := collectGeoColumns(arrowSchema, colMapping)
+ geoNormalizeCols := collectGeoNormalizationColumns(arrowSchema)
geoAccs := make(map[int]*geoBoundsAccumulator, len(geoCols))
for _, gc := range geoCols {
// The accumulator must know whether the column is geometry or geography
@@ -595,18 +633,20 @@ func (p parquetFormat) NewFileWriter(ctx context.Context, fs iceio.WriteFileIO,
}
return &ParquetFileWriter{
- pqWriter: writer,
- counter: counter,
- fileCloser: fw,
- fs: fs,
- format: p,
- info: info,
- partition: partitionValues,
- colMapping: colMapping,
- geoCols: geoCols,
- geoAccs: geoAccs,
- arrowSchema: arrowSchema,
- rowGroupBytes: rowGroupTargetSizeBytes,
+ pqWriter: writer,
+ counter: counter,
+ fileCloser: fw,
+ mem: mem,
+ fs: fs,
+ format: p,
+ info: info,
+ partition: partitionValues,
+ colMapping: colMapping,
+ geoCols: geoCols,
+ geoNormalizeCols: geoNormalizeCols,
+ geoAccs: geoAccs,
+ arrowSchema: arrowSchema,
+ rowGroupBytes: rowGroupTargetSizeBytes,
}, nil
}
@@ -666,21 +706,29 @@ func getWriteProperties(writeProps any, arrowSchema *arrow.Schema) (*parquet.Wri
// Write appends a record batch to the Parquet file.
func (w *ParquetFileWriter) Write(batch arrow.RecordBatch) error {
- if err := w.accumulateGeoBounds(batch); err != nil {
+ writeBatch, err := w.normalizeGeoBatch(batch)
+ if err != nil {
+ return err
+ }
+ if writeBatch != batch {
+ defer writeBatch.Release()
+ }
+
+ if err := w.accumulateGeoBounds(writeBatch); err != nil {
return err
}
// Rotate before writing the next non-empty batch. Rotating immediately after
// crossing the target would leave an empty trailing row group when the file
// is closed without another write.
- if batch.NumRows() > 0 &&
+ if writeBatch.NumRows() > 0 &&
w.pqWriter.RowGroupTotalBytesWritten() >= w.rowGroupBytes {
if err := w.pqWriter.NewBufferedRowGroupChecked(); err != nil {
return err
}
}
- return w.pqWriter.WriteBuffered(batch)
+ return w.pqWriter.WriteBuffered(writeBatch)
}
// wkbStorage is the subset of the binary Arrow arrays that back a geoarrow WKB
@@ -690,6 +738,341 @@ type wkbStorage interface {
Value(int) []byte
}
+func (w *ParquetFileWriter) normalizeGeoBatch(batch arrow.RecordBatch) (arrow.RecordBatch, error) {
+ if len(w.geoNormalizeCols) == 0 {
+ return batch, nil
+ }
+
+ columns := make(map[int]arrow.Array, len(w.geoNormalizeCols))
+ for _, idx := range w.geoNormalizeCols {
+ if idx >= int(batch.NumCols()) {
+ continue
+ }
+
+ normalized, changed, err := normalizeNestedArray(batch.Column(idx), w.mem)
+ if err != nil {
+ for _, col := range columns {
+ col.Release()
+ }
+
+ return nil, fmt.Errorf("normalizing geo values for column %d: %w", idx, err)
+ }
+ if changed {
+ columns[idx] = normalized
+ }
+ }
+
+ if len(columns) == 0 {
+ return batch, nil
+ }
+
+ resultColumns := slices.Clone(batch.Columns())
+ for idx, col := range columns {
+ resultColumns[idx] = col
+ }
+ result := array.NewRecordBatch(batch.Schema(), resultColumns, batch.NumRows())
+ for _, col := range columns {
+ col.Release()
+ }
+
+ return result, nil
+}
+
+// normalizeNestedArray walks an Arrow array and replaces only nested WKB values
+// that use EWKB encoding. Changed nested containers are rebuilt from their
+// logical values with offset zero, so materialized children cannot be sliced a
+// second time by an enclosing Arrow container. Recursive calls carry an
+// optional reachability mask so values hidden by nullable ancestors are copied
+// without being decoded. Such unreachable child slots may retain EWKB or
+// malformed bytes, which is safe because they cannot be observed through the
+// containing null value.
+func normalizeNestedArray(arr arrow.Array, mem memory.Allocator) (arrow.Array, bool, error) {
+ return normalizeNestedArrayReachable(arr, nil, mem)
+}
+
+func normalizeNestedArrayReachable(arr arrow.Array, active []bool, mem memory.Allocator) (arrow.Array, bool, error) {
+ if active != nil && !hasReachableValues(active) {
+ return nil, false, nil
+ }
+
+ if ext, ok := arr.(array.ExtensionArray); ok {
+ if _, ok := ext.ExtensionType().(*geoarrow.WKBType); !ok {
+ return nil, false, nil
+ }
+
+ return normalizeWKBArrayReachable(ext, active, mem)
+ }
+
+ // The table write path currently supplies unencoded nested arrays.
+ switch nested := arr.(type) {
+ case *array.Struct:
+ return normalizeStruct(nested, active, mem)
+
+ case *array.Map:
+ return normalizeListChild(nested, active, mem, func(data arrow.ArrayData) arrow.Array {
+ return array.NewMapData(data)
+ })
+
+ case *array.List:
+ return normalizeListChild(nested, active, mem, func(data arrow.ArrayData) arrow.Array {
+ return array.NewListData(data)
+ })
+
+ case *array.LargeList:
+ return normalizeListChild(nested, active, mem, func(data arrow.ArrayData) arrow.Array {
+ return array.NewLargeListData(data)
+ })
+
+ case *array.FixedSizeList:
+ return normalizeListChild(nested, active, mem, func(data arrow.ArrayData) arrow.Array {
+ return array.NewFixedSizeListData(data)
+ })
+ }
+
+ return nil, false, nil
+}
+
+func isReachable(active []bool, idx int) bool {
+ return active == nil || active[idx]
+}
+
+func hasReachableValues(active []bool) bool {
+ for _, reachable := range active {
+ if reachable {
+ return true
+ }
+ }
+
+ return false
+}
+
+func reachableWithValidity(active []bool, arr arrow.Array) []bool {
+ if arr.NullN() == 0 {
+ return active
+ }
+
+ reachable := make([]bool, arr.Len())
+ for i := range reachable {
+ reachable[i] = isReachable(active, i) && arr.IsValid(i)
+ }
+
+ return reachable
+}
+
+func normalizeStruct(nested *array.Struct, active []bool, mem memory.Allocator) (arrow.Array, bool, error) {
+ children := make([]arrow.ArrayData, nested.NumField())
+ changedChildren := make([]arrow.Array, 0, len(children))
+ childActive := reachableWithValidity(active, nested)
+ for idx := range children {
+ field := nested.Field(idx)
+ children[idx] = field.Data()
+ normalized, changed, err := normalizeNestedArrayReachable(field, childActive, mem)
+ if err != nil {
+ for _, changedChild := range changedChildren {
+ changedChild.Release()
+ }
+
+ return nil, false, err
+ }
+ if changed {
+ children[idx] = normalized.Data()
+ changedChildren = append(changedChildren, normalized)
+ }
+ }
+ if len(changedChildren) == 0 {
+ return nil, false, nil
+ }
+ defer func() {
+ for _, changedChild := range changedChildren {
+ changedChild.Release()
+ }
+ }()
+
+ validity, nulls := rebasedValidity(nested, mem)
+ if validity != nil {
+ defer validity.Release()
+ }
+ buffers := slices.Clone(nested.Data().Buffers())
+ if len(buffers) == 0 {
+ buffers = []*memory.Buffer{validity}
+ } else {
+ buffers[0] = validity
+ }
+ data := array.NewData(nested.DataType(), nested.Len(), buffers, children, nulls, 0)
+ result := array.NewStructData(data)
+ data.Release()
+
+ return result, true, nil
+}
+
+func normalizeListChild(list array.ListLike, active []bool, mem memory.Allocator,
+ newArray func(arrow.ArrayData) arrow.Array,
+) (arrow.Array, bool, error) {
+ if list.Len() == 0 {
+ return nil, false, nil
+ }
+
+ start, _ := list.ValueOffsets(0)
+ _, end := list.ValueOffsets(list.Len() - 1)
+ if start == end {
+ return nil, false, nil
+ }
+
+ childActive := make([]bool, int(end-start))
+ for i := range list.Len() {
+ rowStart, rowEnd := list.ValueOffsets(i)
+ if !isReachable(active, i) || list.IsNull(i) {
+ continue
+ }
+ for j := rowStart; j < rowEnd; j++ {
+ childActive[int(j-start)] = true
+ }
+ }
+ if !hasReachableValues(childActive) {
+ return nil, false, nil
+ }
+
+ values := array.NewSlice(list.ListValues(), start, end)
+ defer values.Release()
+ normalized, changed, err := normalizeNestedArrayReachable(values, childActive, mem)
+ if err != nil || !changed {
+ return nil, false, err
+ }
+ defer normalized.Release()
+
+ children := slices.Clone(list.Data().Children())
+ children[0] = normalized.Data()
+
+ validity, nulls := rebasedValidity(list, mem)
+ if validity != nil {
+ defer validity.Release()
+ }
+ buffers := slices.Clone(list.Data().Buffers())
+ if len(buffers) == 0 {
+ buffers = []*memory.Buffer{validity}
+ }
+ buffers[0] = validity
+ if list.DataType().ID() != arrow.FIXED_SIZE_LIST {
+ offsets, err := rebasedListOffsets(list, start, mem)
+ if err != nil {
+ return nil, false, err
+ }
+ defer offsets.Release()
+ if len(buffers) < 2 {
+ return nil, false, errors.New("list array has no offsets buffer")
+ }
+ buffers[1] = offsets
+ }
+
+ newData := array.NewData(list.DataType(), list.Len(), buffers, children, nulls, 0)
+ result := newArray(newData)
+ newData.Release()
+
+ return result, true, nil
+}
+
+func rebasedValidity(arr arrow.Array, mem memory.Allocator) (*memory.Buffer, int) {
+ nulls := arr.NullN()
+ if nulls == 0 {
+ return nil, 0
+ }
+
+ data := mem.Allocate(int(bitutil.BytesForBits(int64(arr.Len()))))
+ for i := range data {
+ data[i] = 0
+ }
+ for i := range arr.Len() {
+ if arr.IsValid(i) {
+ bitutil.SetBit(data, i)
+ }
+ }
+
+ return memory.NewBufferWithAllocator(data, mem), nulls
+}
+
+func rebasedListOffsets(list array.ListLike, first int64, mem memory.Allocator) (*memory.Buffer, error) {
+ var (
+ data []byte
+ set func(int, int64)
+ )
+ switch list.DataType().ID() {
+ case arrow.LIST, arrow.MAP:
+ data = mem.Allocate((list.Len() + 1) * 4)
+ offsets := arrow.Int32Traits.CastFromBytes(data)
+ set = func(idx int, value int64) {
+ offsets[idx] = int32(value)
+ }
+ case arrow.LARGE_LIST:
+ data = mem.Allocate((list.Len() + 1) * 8)
+ offsets := arrow.Int64Traits.CastFromBytes(data)
+ set = func(idx int, value int64) {
+ offsets[idx] = value
+ }
+ default:
+ return nil, fmt.Errorf("unsupported list type %s", list.DataType())
+ }
+
+ buffer := memory.NewBufferWithAllocator(data, mem)
+ for i := range list.Len() {
+ start, end := list.ValueOffsets(i)
+ set(i, start-first)
+ set(i+1, end-first)
+ }
+
+ return buffer, nil
+}
+
+func normalizeWKBArray(ext array.ExtensionArray, mem memory.Allocator) (arrow.Array, bool, error) {
+ return normalizeWKBArrayReachable(ext, nil, mem)
+}
+
+func normalizeWKBArrayReachable(ext array.ExtensionArray, active []bool, mem memory.Allocator) (arrow.Array, bool, error) {
+ storage, ok := ext.Storage().(wkbStorage)
+ if !ok {
+ return nil, false, nil
+ }
+
+ needsNormalization := false
+ for i := range storage.Len() {
+ if isReachable(active, i) && !storage.IsNull(i) && isEWKB(storage.Value(i)) {
+ needsNormalization = true
+
+ break
+ }
+ }
+ if !needsNormalization {
+ return nil, false, nil
+ }
+
+ builder := array.NewBinaryBuilder(mem, storage.DataType().(arrow.BinaryDataType))
+ defer builder.Release()
+
+ for i := range storage.Len() {
+ if storage.IsNull(i) {
+ builder.AppendNull()
+
+ continue
+ }
+ if !isReachable(active, i) || !isEWKB(storage.Value(i)) {
+ builder.Append(storage.Value(i))
+
+ continue
+ }
+
+ value, err := normalizeWKB(storage.Value(i))
+ if err != nil {
+ return nil, false, err
+ }
+ builder.Append(value)
+ }
+
+ normalizedStorage := builder.NewArray()
+ normalized := array.NewExtensionArrayWithStorage(ext.ExtensionType(), normalizedStorage)
+ normalizedStorage.Release()
+
+ return normalized, true, nil
+}
+
// accumulateGeoBounds extends the per-field bounding boxes with the WKB values
// in this batch. Null rows are skipped; a malformed WKB value fails the write.
func (w *ParquetFileWriter) accumulateGeoBounds(batch arrow.RecordBatch) error {
diff --git a/table/internal/parquet_files_internal_test.go b/table/internal/parquet_files_internal_test.go
index 78b3725bf..273edfe2c 100644
--- a/table/internal/parquet_files_internal_test.go
+++ b/table/internal/parquet_files_internal_test.go
@@ -18,11 +18,14 @@
package internal
import (
+ "context"
"errors"
"testing"
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/iceberg-go"
+ iceio "github.com/apache/iceberg-go/io"
+ "github.com/geoarrow/geoarrow-go"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@@ -88,3 +91,28 @@ func TestWriteDataFileBatchesAbortsOnWriteError(t *testing.T) {
})
}
}
+
+func TestNewFileWriterCachesGeoNormalizationColumns(t *testing.T) {
+ typeDef := geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())
+ arrowSchema := arrow.NewSchema([]arrow.Field{
+ {Name: "id", Type: arrow.PrimitiveTypes.Int32, Nullable: false},
+ {Name: "geom", Type: typeDef, Nullable: true},
+ }, nil)
+ fileSchema := iceberg.NewSchema(0,
+ iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int32, Required: true},
+ iceberg.NestedField{ID: 2, Name: "geom", Type: iceberg.GeometryType{}, Required: false},
+ )
+ format := parquetFormat{}
+ writer, err := format.NewFileWriter(context.Background(), iceio.NewMemFS(), nil, WriteFileInfo{
+ FileSchema: fileSchema,
+ Spec: *iceberg.UnpartitionedSpec,
+ FileName: "geo-cache.parquet",
+ WriteProps: format.GetWriteProperties(iceberg.Properties{}),
+ }, arrowSchema)
+ require.NoError(t, err)
+ defer func() { require.NoError(t, writer.Abort()) }()
+
+ parquetWriter, ok := writer.(*ParquetFileWriter)
+ require.True(t, ok)
+ require.Equal(t, []int{1}, parquetWriter.geoNormalizeCols)
+}
diff --git a/table/internal/parquet_files_test.go b/table/internal/parquet_files_test.go
index 25d62abcf..ff7150f7d 100644
--- a/table/internal/parquet_files_test.go
+++ b/table/internal/parquet_files_test.go
@@ -35,6 +35,7 @@ import (
"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"
@@ -53,6 +54,8 @@ import (
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
+ "github.com/twpayne/go-geom/encoding/wkb"
+ "github.com/twpayne/go-geom/encoding/wkt"
)
type abortTestFile struct {
@@ -1113,6 +1116,91 @@ func TestParquetGeoArrowExtensionMetadataRoundTrip(t *testing.T) {
assertWKBValues(t, batch.Column(4), defaultGeogWKB, nil)
}
+func TestWriteDataFileNormalizesEWKBOnDisk(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ ctx := compute.WithAllocator(context.Background(), mem)
+ geomType, err := iceberg.GeometryTypeOf("srid:4326")
+ require.NoError(t, err)
+
+ iceSchema := iceberg.NewSchema(0,
+ iceberg.NestedField{ID: 1, Name: "geom", Type: geomType, Required: false},
+ )
+ arrowSchema, err := table.SchemaToArrowSchema(iceSchema, nil, true, false)
+ require.NoError(t, err)
+
+ isoWKB := pointWKB(t, 1, 2)
+ builder := array.NewBinaryBuilder(mem, arrow.BinaryTypes.Binary)
+ builder.Append(ewkbWithSRID(isoWKB, 4326))
+ builder.AppendNull()
+ storage := builder.NewArray()
+ builder.Release()
+ ext := array.NewExtensionArrayWithStorage(
+ arrowSchema.Field(0).Type.(arrow.ExtensionType), storage,
+ ).(array.ExtensionArray)
+ storage.Release()
+ record := array.NewRecordBatch(arrowSchema, []arrow.Array{ext}, 2)
+ ext.Release()
+ defer record.Release()
+
+ fs := iceio.NewMemFS()
+ fm := internal.GetFileFormat(iceberg.ParquetFile)
+ _, err = fm.WriteDataFile(ctx, fs, nil, internal.WriteFileInfo{
+ FileSchema: iceSchema,
+ Spec: *iceberg.UnpartitionedSpec,
+ FileName: "geo.parquet",
+ StatsCols: map[int]internal.StatisticsCollector{
+ 1: {FieldID: 1, IcebergTyp: geomType, ColName: "geom", Mode: internal.MetricsMode{Typ: internal.MetricModeFull}},
+ },
+ WriteProps: fm.GetWriteProperties(iceberg.Properties{}),
+ Content: iceberg.EntryContentData,
+ }, []arrow.RecordBatch{record})
+ require.NoError(t, err)
+
+ output, err := fs.Open("geo.parquet")
+ require.NoError(t, err)
+ data, err := io.ReadAll(output)
+ closeErr := output.Close()
+ require.NoError(t, err)
+ require.NoError(t, closeErr)
+ pqReader, err := file.NewParquetReader(bytes.NewReader(data))
+ require.NoError(t, err)
+ defer pqReader.Close()
+
+ arrowReader, err := pqarrow.NewFileReader(pqReader, pqarrow.ArrowReadProperties{}, mem)
+ require.NoError(t, err)
+ reader := internal.WrapParquetFileReader(arrowReader)
+ records, err := reader.GetRecords(ctx, nil, nil)
+ require.NoError(t, err)
+ defer records.Release()
+ require.True(t, records.Next())
+
+ batch := records.RecordBatch()
+ defer batch.Release()
+ assertWKBValues(t, batch.Column(0), geoarrow.WKBBytes(isoWKB), nil)
+}
+
+func pointWKB(t *testing.T, x, y float64) []byte {
+ t.Helper()
+ geom, err := wkt.Unmarshal(fmt.Sprintf("POINT (%g %g)", x, y))
+ require.NoError(t, err)
+ wkbBytes, err := wkb.Marshal(geom, wkb.NDR)
+ require.NoError(t, err)
+
+ return wkbBytes
+}
+
+func ewkbWithSRID(isoWKB []byte, srid uint32) []byte {
+ result := make([]byte, len(isoWKB)+4)
+ copy(result[:5], isoWKB[:5])
+ binary.LittleEndian.PutUint32(result[1:5], binary.LittleEndian.Uint32(isoWKB[1:5])|0x20000000)
+ binary.LittleEndian.PutUint32(result[5:9], srid)
+ copy(result[9:], isoWKB[5:])
+
+ return result
+}
+
func TestWriteDataFileErrOnClose(t *testing.T) {
ctx := context.Background()
fm := internal.GetFileFormat(iceberg.ParquetFile)