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/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new 51bb2726 feat(parquet/pqarrow): Read support for extension types / geo
(#969)
51bb2726 is described below
commit 51bb2726df38befb8418b590fc63719f834adcae
Author: Colton Loftus <[email protected]>
AuthorDate: Fri Jul 24 13:16:22 2026 -0400
feat(parquet/pqarrow): Read support for extension types / geo (#969)
### Rationale for this change
In https://github.com/apache/arrow-go/pull/960 I implemented write
support for extension types by mapping a logical type to its underlying
primitive/arrow type. In the case of geometry this allows the logical
geometry type to be mapped to byte array for storage.
This PR implements the read path for that by adding the parquet type to
arrow type functionality. For example, a user can read a parquet with a
geometry column and with this change, the reader will try to map it to
the designated associated arrow extension type.
### What changes are included in this PR?
An extension type can implement this interface to designate which arrow
type should be used when reading it from parquet.
```go
type ExtensionParquetLogicalType interface {
ArrowTypeFromParquet(logical schema.LogicalType, storageType
arrow.DataType) (arrow.ExtensionType, error)
}
```
The read path will try to iterate through all the extension types, see
if there are any matches where the type implements
`ArrowTypeFromParquet` and gives us something to map it to in arrow. If
so we use that for reading.
### Are these changes tested?
Yes
### Are there any user-facing changes?
Yes. Readers can now support arbitrary extension types as long as they
are registered properly in the extension type registry.
---
arrow/datatype_extension.go | 31 +++++
arrow/datatype_extension_test.go | 64 ++++++++++
parquet/pqarrow/schema.go | 57 +++++++++
parquet/pqarrow/schema_test.go | 265 +++++++++++++++++++++++++++++++++++++++
4 files changed, 417 insertions(+)
diff --git a/arrow/datatype_extension.go b/arrow/datatype_extension.go
index f0bcccdf..6ec082b5 100644
--- a/arrow/datatype_extension.go
+++ b/arrow/datatype_extension.go
@@ -19,6 +19,7 @@ package arrow
import (
"fmt"
"reflect"
+ "sort"
"sync"
)
@@ -82,6 +83,36 @@ func GetExtensionType(typName string) ExtensionType {
return nil
}
+// FindRegisteredExtensionType returns a registered extension type for
+// which the filter returns true. If no type matches, it returns nil.
+//
+// If multiple types match, the type with the lexicographically smallest
+// registered extension name is returned.
+func FindRegisteredExtensionType(filter func(ExtensionType) bool)
ExtensionType {
+ registry := getExtTypeRegistry()
+
+ // Snapshot and sort the registered names before invoking the filter so
the
+ // result is deterministic (lexicographically smallest match),
independent of
+ // sync.Map's unspecified Range order.
+ types := make(map[string]ExtensionType)
+ var names []string
+ registry.Range(func(key, value any) bool {
+ name := key.(string)
+ types[name] = value.(ExtensionType)
+ names = append(names, name)
+ return true
+ })
+
+ sort.Strings(names)
+ for _, name := range names {
+ typ := types[name]
+ if filter(typ) {
+ return typ
+ }
+ }
+ return nil
+}
+
// ExtensionType is an interface for handling user-defined types. They must be
// DataTypes and must embed arrow.ExtensionBase in them in order to work
properly
// ensuring that they always have the expected base behavior.
diff --git a/arrow/datatype_extension_test.go b/arrow/datatype_extension_test.go
index 9a3361c8..1be1eda4 100644
--- a/arrow/datatype_extension_test.go
+++ b/arrow/datatype_extension_test.go
@@ -39,6 +39,29 @@ func (BadExtensionType) Deserialize(_ arrow.DataType, _
string) (arrow.Extension
return nil, nil
}
+type findExtensionType struct {
+ arrow.ExtensionBase
+ name string
+}
+
+func newFindExtensionType(name string) *findExtensionType {
+ return &findExtensionType{
+ ExtensionBase: arrow.ExtensionBase{Storage: arrow.Null},
+ name: name,
+ }
+}
+
+func (*findExtensionType) ArrayType() reflect.Type { return nil }
+func (f *findExtensionType) ExtensionName() string { return f.name }
+func (*findExtensionType) Serialize() string { return "" }
+func (f *findExtensionType) ExtensionEquals(other arrow.ExtensionType) bool {
+ rhs, ok := other.(*findExtensionType)
+ return ok && f.name == rhs.name
+}
+func (f *findExtensionType) Deserialize(_ arrow.DataType, _ string)
(arrow.ExtensionType, error) {
+ return f, nil
+}
+
func TestMustEmbedBase(t *testing.T) {
var ext interface{} = &BadExtensionType{}
assert.Panics(t, func() {
@@ -71,6 +94,47 @@ func (e *ExtensionTypeTestSuite) TestExtensionType() {
e.False(arrow.TypeEqual(deserialized,
&arrow.FixedSizeBinaryType{ByteWidth: 16}))
}
+func (e *ExtensionTypeTestSuite) TestFindRegisteredExtensionType() {
+ found := arrow.FindRegisteredExtensionType(func(typ
arrow.ExtensionType) bool {
+ return typ.ExtensionName() == "arrow.uuid"
+ })
+ e.Same(arrow.GetExtensionType("arrow.uuid"), found)
+
+ notFound := arrow.FindRegisteredExtensionType(func(typ
arrow.ExtensionType) bool {
+ return typ.ExtensionName() == "uuid-unknown"
+ })
+ e.Nil(notFound)
+
+ calls := 0
+ found = arrow.FindRegisteredExtensionType(func(typ arrow.ExtensionType)
bool {
+ calls++
+ return true
+ })
+ e.NotNil(found)
+ e.Equal(1, calls)
+
+ // Register the higher name first so this assertion does not
accidentally
+ // pass due to registration order.
+ lowNameType := newFindExtensionType("test.find.a")
+ highNameType := newFindExtensionType("test.find.z")
+ e.Require().NoError(arrow.RegisterExtensionType(highNameType))
+ defer func() {
+
e.NoError(arrow.UnregisterExtensionType(highNameType.ExtensionName()))
+ }()
+ e.Require().NoError(arrow.RegisterExtensionType(lowNameType))
+ defer func() {
+
e.NoError(arrow.UnregisterExtensionType(lowNameType.ExtensionName()))
+ }()
+
+ found = arrow.FindRegisteredExtensionType(func(typ arrow.ExtensionType)
bool {
+ return typ.ExtensionName() == lowNameType.ExtensionName() ||
+ typ.ExtensionName() == highNameType.ExtensionName()
+ })
+ // When multiple registered types match, the lexicographically smallest
+ // registered extension name should win.
+ e.Same(lowNameType, found)
+}
+
func TestExtensionTypes(t *testing.T) {
suite.Run(t, new(ExtensionTypeTestSuite))
}
diff --git a/parquet/pqarrow/schema.go b/parquet/pqarrow/schema.go
index 66db50e1..81582303 100644
--- a/parquet/pqarrow/schema.go
+++ b/parquet/pqarrow/schema.go
@@ -129,6 +129,19 @@ type ExtensionCustomParquetType interface {
ParquetLogicalType() schema.LogicalType
}
+// ExtensionParquetLogicalType is an interface that Arrow ExtensionTypes may
+// implement to specify how a Parquet LogicalType maps back to an Arrow
+// ExtensionType when converting a Parquet schema to an Arrow schema.
+//
+// ArrowTypeFromParquet should return (nil, nil) if the logical type does not
+// map to the extension type. It should return (nil, err) if the logical type
is
+// recognized but cannot be converted into a valid extension type. If a
+// non-nil extension type is returned, that type is used and any previous
+// conversion errors from other extension types are ignored.
+type ExtensionParquetLogicalType interface {
+ ArrowTypeFromParquet(logical schema.LogicalType, storageType
arrow.DataType) (arrow.ExtensionType, error)
+}
+
func isDictionaryReadSupported(dt arrow.DataType) bool {
return arrow.IsBinaryLike(dt.ID())
}
@@ -531,6 +544,9 @@ func arrowFromByteArray(logical schema.LogicalType)
(arrow.DataType, error) {
return arrow.BinaryTypes.String, nil
case schema.DecimalLogicalType:
return arrowDecimal(logtype), nil
+ case schema.GeometryLogicalType,
+ schema.GeographyLogicalType:
+ return arrowExtensionFromParquetLogicalType(logical,
arrow.BinaryTypes.Binary)
case schema.NoLogicalType,
schema.EnumLogicalType,
schema.JSONLogicalType,
@@ -541,6 +557,47 @@ func arrowFromByteArray(logical schema.LogicalType)
(arrow.DataType, error) {
}
}
+// arrowExtensionFromParquetLogicalType asks registered extension types whether
+// they can represent the provided Parquet logical type with the given storage
+// type, falling back to the storage type when none opt in.
+func arrowExtensionFromParquetLogicalType(logical schema.LogicalType,
storageType arrow.DataType) (arrow.DataType, error) {
+ var (
+ typ arrow.ExtensionType
+ typeLookupErr error
+ )
+ matchedType := arrow.FindRegisteredExtensionType(func(extType
arrow.ExtensionType) bool {
+ converter, ok := extType.(ExtensionParquetLogicalType)
+ if !ok {
+ return false
+ }
+
+ var err error
+ typ, err = converter.ArrowTypeFromParquet(logical, storageType)
+ if err != nil {
+ // if there was a failure to match, store it as a
potential error
+ // to return if nothing else matches, but don't return
it yet.
+ // This lets one extension reject a column without
blocking another
+ // extension, while still surfacing real converter
failures when nothing
+ // handles the type.
+ if typeLookupErr == nil {
+ typeLookupErr = err
+ }
+ return false
+ }
+ return typ != nil
+ })
+ if matchedType != nil {
+ return typ, nil
+ }
+
+ if typeLookupErr != nil {
+ // If we didn't match any type and there was a type lookup
error,
+ // we can now return the error that we stored
+ return nil, typeLookupErr
+ }
+ return storageType, nil
+}
+
func arrowFromFLBA(logical schema.LogicalType, length int) (arrow.DataType,
error) {
switch logtype := logical.(type) {
case schema.DecimalLogicalType:
diff --git a/parquet/pqarrow/schema_test.go b/parquet/pqarrow/schema_test.go
index e6c7a5a3..11f969e2 100644
--- a/parquet/pqarrow/schema_test.go
+++ b/parquet/pqarrow/schema_test.go
@@ -17,15 +17,21 @@
package pqarrow_test
import (
+ "bytes"
+ "context"
"encoding/base64"
+ "errors"
+ "reflect"
"testing"
"github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
"github.com/apache/arrow-go/v18/arrow/extensions"
"github.com/apache/arrow-go/v18/arrow/flight"
"github.com/apache/arrow-go/v18/arrow/ipc"
"github.com/apache/arrow-go/v18/arrow/memory"
"github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/file"
"github.com/apache/arrow-go/v18/parquet/metadata"
"github.com/apache/arrow-go/v18/parquet/pqarrow"
"github.com/apache/arrow-go/v18/parquet/schema"
@@ -33,6 +39,100 @@ import (
"github.com/stretchr/testify/require"
)
+type testGeometryType struct {
+ arrow.ExtensionBase
+ logical schema.LogicalType
+}
+
+type testGeometryArray struct {
+ array.ExtensionArrayBase
+}
+
+type testFailingGeometryType struct {
+ arrow.ExtensionBase
+}
+
+func newTestGeometryType(logical schema.LogicalType) *testGeometryType {
+ return &testGeometryType{
+ ExtensionBase: arrow.ExtensionBase{Storage:
arrow.BinaryTypes.Binary},
+ logical: logical,
+ }
+}
+
+func newTestFailingGeometryType() *testFailingGeometryType {
+ return &testFailingGeometryType{
+ ExtensionBase: arrow.ExtensionBase{Storage:
arrow.BinaryTypes.Binary},
+ }
+}
+
+func (*testGeometryType) ArrayType() reflect.Type {
+ return reflect.TypeFor[testGeometryArray]()
+}
+
+func (*testFailingGeometryType) ArrayType() reflect.Type {
+ return reflect.TypeFor[testGeometryArray]()
+}
+
+func (*testGeometryType) ExtensionName() string {
+ return "test.geospatial"
+}
+
+func (*testFailingGeometryType) ExtensionName() string {
+ return "test.failing_geospatial"
+}
+
+func (t *testGeometryType) ExtensionEquals(other arrow.ExtensionType) bool {
+ rhs, ok := other.(*testGeometryType)
+ return ok && t.ExtensionName() == rhs.ExtensionName() &&
t.logical.Equals(rhs.logical)
+}
+
+func (t *testFailingGeometryType) ExtensionEquals(other arrow.ExtensionType)
bool {
+ _, ok := other.(*testFailingGeometryType)
+ return ok
+}
+
+func (*testGeometryType) Serialize() string {
+ return ""
+}
+
+func (*testFailingGeometryType) Serialize() string {
+ return ""
+}
+
+func (*testGeometryType) Deserialize(storageType arrow.DataType, data string)
(arrow.ExtensionType, error) {
+ return newTestGeometryType(schema.GeometryLogicalType{}), nil
+}
+
+func (*testFailingGeometryType) Deserialize(storageType arrow.DataType, data
string) (arrow.ExtensionType, error) {
+ return newTestFailingGeometryType(), nil
+}
+
+func (*testGeometryType) ArrowTypeFromParquet(logical schema.LogicalType,
storageType arrow.DataType) (arrow.ExtensionType, error) {
+ if !arrow.TypeEqual(storageType, arrow.BinaryTypes.Binary) {
+ return nil, nil
+ }
+ switch logical.(type) {
+ case schema.GeometryLogicalType, schema.GeographyLogicalType:
+ return newTestGeometryType(logical), nil
+ default:
+ return nil, nil
+ }
+}
+
+func (*testFailingGeometryType) ArrowTypeFromParquet(logical
schema.LogicalType, storageType arrow.DataType) (arrow.ExtensionType, error) {
+ return nil, errors.New("not my geospatial logical type")
+}
+
+func (t *testGeometryType) ParquetLogicalType() schema.LogicalType {
+ return t.logical
+}
+
+var (
+ _ pqarrow.ExtensionCustomParquetType = (*testGeometryType)(nil)
+ _ pqarrow.ExtensionParquetLogicalType = (*testGeometryType)(nil)
+ _ pqarrow.ExtensionParquetLogicalType = (*testFailingGeometryType)(nil)
+)
+
func TestGetOriginSchemaBase64(t *testing.T) {
uuidType := extensions.NewUUIDType()
md := arrow.NewMetadata([]string{"PARQUET:field_id"}, []string{"-1"})
@@ -97,6 +197,171 @@ func TestGetOriginSchemaUnregisteredExtension(t
*testing.T) {
assert.Truef(t, expArrSc.Equal(arrsc), "expected: %s\ngot: %s",
expArrSc, arrsc)
}
+// TestFromParquetGeospatialRegisteredExtension verifies that Parquet Geometry
and
+// Geography logical byte-array columns are mapped to a registered Arrow
extension
+// type when the extension opts into Parquet logical type conversion.
+func TestFromParquetGeospatialRegisteredExtension(t *testing.T) {
+ geoType := newTestGeometryType(schema.GeometryLogicalType{})
+ require.NoError(t, arrow.RegisterExtensionType(geoType))
+ defer func() {
+ require.NoError(t,
arrow.UnregisterExtensionType(geoType.ExtensionName()))
+ }()
+
+ geometryLogical := schema.GeometryLogicalType{Crs: "EPSG:4326"}
+ geographyLogical := schema.GeographyLogicalType{Crs: "OGC:CRS84",
Algorithm: schema.GeographyEdgeKarney}
+ geomNode := schema.Must(schema.NewPrimitiveNodeLogical("geometry",
parquet.Repetitions.Optional,
+ geometryLogical, parquet.Types.ByteArray, -1, -1))
+ geogNode := schema.Must(schema.NewPrimitiveNodeLogical("geography",
parquet.Repetitions.Optional,
+ geographyLogical, parquet.Types.ByteArray, -1, -1))
+ pqschema :=
schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema",
parquet.Repetitions.Required,
+ schema.FieldList{geomNode, geogNode}, -1)))
+
+ arrsc, err := pqarrow.FromParquet(pqschema, nil, nil)
+ require.NoError(t, err)
+ require.Equal(t, 2, arrsc.NumFields())
+
+ for i, logical := range []schema.LogicalType{geometryLogical,
geographyLogical} {
+ require.Equal(t, arrow.EXTENSION, arrsc.Field(i).Type.ID())
+ extType, ok := arrsc.Field(i).Type.(*testGeometryType)
+ require.True(t, ok)
+ assert.True(t, logical.Equals(extType.logical))
+ assert.True(t, arrow.TypeEqual(arrow.BinaryTypes.Binary,
extType.StorageType()))
+ }
+}
+
+func TestFromParquetGeospatialNoRegisteredExtensionFallsBackToBinary(t
*testing.T) {
+ geometryLogical := schema.GeometryLogicalType{Crs: "EPSG:4326"}
+ geomNode := schema.Must(schema.NewPrimitiveNodeLogical("geometry",
parquet.Repetitions.Optional,
+ geometryLogical, parquet.Types.ByteArray, -1, -1))
+ pqschema :=
schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema",
parquet.Repetitions.Required,
+ schema.FieldList{geomNode}, -1)))
+
+ arrsc, err := pqarrow.FromParquet(pqschema, nil, nil)
+ require.NoError(t, err)
+ require.Equal(t, 1, arrsc.NumFields())
+ assert.True(t, arrow.TypeEqual(arrow.BinaryTypes.Binary,
arrsc.Field(0).Type))
+}
+
+func
TestFromParquetGeospatialRegisteredExtensionContinuesAfterConverterError(t
*testing.T) {
+ failingType := newTestFailingGeometryType()
+ require.NoError(t, arrow.RegisterExtensionType(failingType))
+ defer func() {
+ require.NoError(t,
arrow.UnregisterExtensionType(failingType.ExtensionName()))
+ }()
+
+ geoType := newTestGeometryType(schema.GeometryLogicalType{})
+ require.NoError(t, arrow.RegisterExtensionType(geoType))
+ defer func() {
+ require.NoError(t,
arrow.UnregisterExtensionType(geoType.ExtensionName()))
+ }()
+
+ geometryLogical := schema.GeometryLogicalType{Crs: "EPSG:4326"}
+ geomNode := schema.Must(schema.NewPrimitiveNodeLogical("geometry",
parquet.Repetitions.Optional,
+ geometryLogical, parquet.Types.ByteArray, -1, -1))
+ pqschema :=
schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema",
parquet.Repetitions.Required,
+ schema.FieldList{geomNode}, -1)))
+
+ arrsc, err := pqarrow.FromParquet(pqschema, nil, nil)
+ require.NoError(t, err)
+ require.Equal(t, 1, arrsc.NumFields())
+
+ extType, ok := arrsc.Field(0).Type.(*testGeometryType)
+ require.True(t, ok)
+ assert.True(t, geometryLogical.Equals(extType.logical))
+}
+
+func
TestFromParquetGeospatialRegisteredExtensionReturnsConverterErrorWhenNoMatch(t
*testing.T) {
+ failingType := newTestFailingGeometryType()
+ require.NoError(t, arrow.RegisterExtensionType(failingType))
+ defer func() {
+ require.NoError(t,
arrow.UnregisterExtensionType(failingType.ExtensionName()))
+ }()
+
+ geometryLogical := schema.GeometryLogicalType{Crs: "EPSG:4326"}
+ geomNode := schema.Must(schema.NewPrimitiveNodeLogical("geometry",
parquet.Repetitions.Optional,
+ geometryLogical, parquet.Types.ByteArray, -1, -1))
+ pqschema :=
schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema",
parquet.Repetitions.Required,
+ schema.FieldList{geomNode}, -1)))
+
+ _, err := pqarrow.FromParquet(pqschema, nil, nil)
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), "not my geospatial logical type")
+}
+
+// TestReadWriteGeospatialRegisteredExtensionWithoutStoredSchema verifies that
a
+// registered extension type can round trip using only the Parquet logical
type.
+// Keeping StoreSchema disabled ensures the reader reconstructs the extension
via
+// ExtensionParquetLogicalType interface logic instead of restoring it from
+// ARROW:schema metadata.
+func TestReadWriteGeospatialRegisteredExtensionWithoutStoredSchema(t
*testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ // Register the extension type so the reader can discover its Parquet
mapping.
+ logical := schema.GeometryLogicalType{Crs: "EPSG:4326"}
+ geoType := newTestGeometryType(logical)
+ require.NoError(t, arrow.RegisterExtensionType(geoType))
+ defer func() {
+ require.NoError(t,
arrow.UnregisterExtensionType(geoType.ExtensionName()))
+ }()
+
+ // Build an extension array whose storage is the byte array written to
Parquet.
+ bldr := array.NewExtensionBuilder(mem, geoType)
+ defer bldr.Release()
+ binaryBldr := bldr.StorageBuilder().(*array.BinaryBuilder)
+ binaryBldr.AppendValues([][]byte{
+ {1, 2, 3},
+ nil,
+ {4, 5, 6, 7},
+ }, []bool{true, false, true})
+
+ arr := bldr.NewArray()
+ defer arr.Release()
+
+ field := arrow.Field{Name: "geometry", Type: geoType, Nullable: true}
+ col := arrow.NewColumnFromArr(field, arr)
+ defer col.Release()
+ tbl := array.NewTable(arrow.NewSchema([]arrow.Field{field}, nil),
[]arrow.Column{col}, -1)
+ defer tbl.Release()
+
+ var buf bytes.Buffer
+ arrowProps :=
pqarrow.NewArrowWriterProperties(pqarrow.WithAllocator(mem))
+ // Write without using WithStoreSchema so the file does not embed the
original Arrow
+ // extension type in ARROW:schema metadata. This forces the read path
to infer the
+ // extension from the Parquet logical type and the registered extension
interface.
+ require.NoError(t, pqarrow.WriteTable(tbl, &buf, tbl.NumRows(),
+ parquet.NewWriterProperties(parquet.WithAllocator(mem)),
arrowProps))
+
+ pf, err := file.NewParquetReader(bytes.NewReader(buf.Bytes()),
file.WithReadProps(parquet.NewReaderProperties(mem)))
+ require.NoError(t, err)
+ defer pf.Close()
+
+ // Confirm the extension type can be recovered from the Parquet logical
type.
+ require.Nil(t,
pf.MetaData().KeyValueMetadata().FindValue("ARROW:schema"))
+ require.True(t,
pf.MetaData().Schema.Column(0).LogicalType().Equals(logical))
+
+ reader, err := pqarrow.NewFileReader(pf, pqarrow.ArrowReadProperties{},
mem)
+ require.NoError(t, err)
+
+ readTbl, err := reader.ReadTable(context.Background())
+ require.NoError(t, err)
+ defer readTbl.Release()
+
+ // Validate that the read path restored the extension type and
preserved values.
+ require.Equal(t, tbl.NumRows(), readTbl.NumRows())
+ require.Equal(t, tbl.NumCols(), readTbl.NumCols())
+ require.Equal(t, arrow.EXTENSION, readTbl.Column(0).DataType().ID())
+
+ readType, ok := readTbl.Column(0).DataType().(*testGeometryType)
+ require.True(t, ok)
+ assert.True(t, logical.Equals(readType.logical))
+ assert.True(t, arrow.TypeEqual(arrow.BinaryTypes.Binary,
readType.StorageType()))
+
+ expected := tbl.Column(0).Data().Chunk(0)
+ actual := readTbl.Column(0).Data().Chunk(0)
+ assert.Truef(t, array.Equal(expected, actual), "expected: %T %s\ngot:
%T %s", expected, expected, actual, actual)
+}
+
func TestToParquetWriterConfig(t *testing.T) {
origSc := arrow.NewSchema([]arrow.Field{
{Name: "f1", Type: arrow.BinaryTypes.String},
