diff --git a/arrow/internal/arrjson/arrjson.go b/arrow/internal/arrjson/arrjson.go index afd34f503..b6b242451 100644 --- a/arrow/internal/arrjson/arrjson.go +++ b/arrow/internal/arrjson/arrjson.go @@ -99,6 +99,10 @@ type FieldDict struct { Type json.RawMessage `json:"indexType"` idxType arrow.DataType `json:"-"` Ordered bool `json:"isOrdered"` + // Metadata holds the metadata of the dictionary value type, such as + // the extension type keys of a dictionary whose values are an + // extension type. + Metadata []metaKV `json:"metadata,omitempty"` } type Field struct { @@ -269,6 +273,16 @@ func (f FieldWrapper) MarshalJSON() ([]byte, error) { if f.Dictionary.Type, err = typeToJSON(f.Dictionary.idxType); err != nil { return nil, err } + // the extension type keys of a dictionary value type go in the + // dictionary's own metadata rather than the field's. + f.Dictionary.Metadata = nil + if exType, ok := f.arrowType.(arrow.ExtensionType); ok { + f.Dictionary.Metadata = []metaKV{{Key: ipc.ExtensionTypeKeyName, Value: exType.ExtensionName()}} + if serializedData := exType.Serialize(); len(serializedData) > 0 { + f.Dictionary.Metadata = append(f.Dictionary.Metadata, metaKV{Key: ipc.ExtensionMetadataKeyName, Value: serializedData}) + } + f.arrowType = exType.StorageType() + } } if f.Type, err = typeToJSON(f.arrowType); err != nil { @@ -564,68 +578,77 @@ func (f *FieldWrapper) UnmarshalJSON(data []byte) error { if f.Dictionary.idxType, err = typeFromJSON(f.Dictionary.Type, nil); err != nil { return err } - f.arrowType = &arrow.DictionaryType{IndexType: f.Dictionary.idxType, ValueType: f.arrowType} + // metadata of the dictionary value type other than a registered + // extension type has nowhere to go and is dropped. + valueType, _, err := extensionTypeFromJSON(f.arrowType, f.Dictionary.Metadata) + if err != nil { + return err + } + f.arrowType = &arrow.DictionaryType{ + IndexType: f.Dictionary.idxType, + ValueType: valueType, + Ordered: f.Dictionary.Ordered, + } } if len(f.Metadata) > 0 { // unmarshal the key/value metadata pairs - var ( - mdkeys = make([]string, 0, len(f.Metadata)) - mdvals = make([]string, 0, len(f.Metadata)) - extKeyIdx = -1 - extDataIdx = -1 - ) - - for i, kv := range f.Metadata { - switch kv.Key { - case ipc.ExtensionTypeKeyName: - extKeyIdx = i - case ipc.ExtensionMetadataKeyName: - extDataIdx = i - } - mdkeys = append(mdkeys, kv.Key) - mdvals = append(mdvals, kv.Value) + var rest []metaKV + if f.arrowType, rest, err = extensionTypeFromJSON(f.arrowType, f.Metadata); err != nil { + return err } - if extKeyIdx == -1 { // no extension metadata just create the metadata - f.arrowMeta = arrow.NewMetadata(mdkeys, mdvals) - return nil + mdkeys := make([]string, 0, len(rest)) + mdvals := make([]string, 0, len(rest)) + for _, kv := range rest { + mdkeys = append(mdkeys, kv.Key) + mdvals = append(mdvals, kv.Value) } + f.arrowMeta = arrow.NewMetadata(mdkeys, mdvals) + } - extType := arrow.GetExtensionType(mdvals[extKeyIdx]) - if extType == nil { // unregistered extension type, just keep the metadata - f.arrowMeta = arrow.NewMetadata(mdkeys, mdvals) - return nil - } + return nil +} - var extData string - if extDataIdx > -1 { - extData = mdvals[extDataIdx] - // if both extension type and extension type metadata exist - // filter out both keys - newkeys := make([]string, 0, len(mdkeys)-2) - newvals := make([]string, 0, len(mdvals)-2) - for i := range mdkeys { - if i != extKeyIdx && i != extDataIdx { - newkeys = append(newkeys, mdkeys[i]) - newvals = append(newvals, mdvals[i]) - } - } - mdkeys = newkeys - mdvals = newvals - } else { - // if only extension type key is present, we can simplify filtering it out - mdkeys = append(mdkeys[:extKeyIdx], mdkeys[extKeyIdx+1:]...) - mdvals = append(mdvals[:extKeyIdx], mdvals[extKeyIdx+1:]...) +// extensionTypeFromJSON returns the extension type named by the +// ARROW:extension:name key of md, with the given storage type, along with +// the rest of md without the extension type keys. If md does not name an +// extension type, or names one that is not registered, the storage type and +// md are returned unchanged. +func extensionTypeFromJSON(storage arrow.DataType, md []metaKV) (arrow.DataType, []metaKV, error) { + extKeyIdx, extDataIdx := -1, -1 + for i, kv := range md { + switch kv.Key { + case ipc.ExtensionTypeKeyName: + extKeyIdx = i + case ipc.ExtensionMetadataKeyName: + extDataIdx = i } + } + if extKeyIdx == -1 { + return storage, md, nil + } - if f.arrowType, err = extType.Deserialize(f.arrowType, extData); err != nil { - return err - } + extType := arrow.GetExtensionType(md[extKeyIdx].Value) + if extType == nil { // unregistered extension type, keep the storage type + return storage, md, nil + } - f.arrowMeta = arrow.NewMetadata(mdkeys, mdvals) + var extData string + if extDataIdx > -1 { + extData = md[extDataIdx].Value + } + dt, err := extType.Deserialize(storage, extData) + if err != nil { + return nil, nil, err } - return err + rest := make([]metaKV, 0, len(md)) + for i, kv := range md { + if i != extKeyIdx && i != extDataIdx { + rest = append(rest, kv) + } + } + return dt, rest, nil } // the structs below represent various configurations of the Type @@ -739,6 +762,9 @@ func fieldsToJSON(fields []arrow.Field, parentPos dictutils.FieldPos, mapper *di if typ.ID() == arrow.DICTIONARY { dictType := typ.(*arrow.DictionaryType) typ = dictType.ValueType + if typ.ID() == arrow.EXTENSION { + typ = typ.(arrow.ExtensionType).StorageType() + } dictID, err := mapper.GetFieldID(pos.Path()) if err != nil { panic(err) diff --git a/arrow/internal/arrjson/arrjson_test.go b/arrow/internal/arrjson/arrjson_test.go index 232270b9d..59dad25b3 100644 --- a/arrow/internal/arrjson/arrjson_test.go +++ b/arrow/internal/arrjson/arrjson_test.go @@ -17,15 +17,23 @@ package arrjson import ( + "bytes" "errors" "io" "os" + "strings" "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/internal/arrdata" + "github.com/apache/arrow-go/v18/arrow/ipc" "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/internal/json" + "github.com/apache/arrow-go/v18/internal/types" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestReadWrite(t *testing.T) { @@ -6382,3 +6390,172 @@ func makeViewTypesWantJSONs() string { ] }` } + +func TestDictionaryOfExtensionMetadata(t *testing.T) { + for _, extType := range []arrow.ExtensionType{types.NewParametric1Type(1), types.NewDictExtensionType()} { + if arrow.GetExtensionType(extType.ExtensionName()) == nil { + require.NoError(t, arrow.RegisterExtensionType(extType)) + defer arrow.UnregisterExtensionType(extType.ExtensionName()) + } + } + + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + uuidType := extensions.NewUUIDType() + p1Type := types.NewParametric1Type(6) + + uuidValues, _, err := array.FromJSON(mem, uuidType, strings.NewReader( + `["00000000-0000-0000-0000-000000000001", null, "00000000-0000-0000-0000-000000000002"]`)) + require.NoError(t, err) + defer uuidValues.Release() + p1Storage, _, err := array.FromJSON(mem, arrow.PrimitiveTypes.Int32, strings.NewReader(`[7, null, 9]`)) + require.NoError(t, err) + defer p1Storage.Release() + p1Values := array.NewExtensionArrayWithStorage(p1Type, p1Storage) + defer p1Values.Release() + indices, _, err := array.FromJSON(mem, arrow.PrimitiveTypes.Int16, strings.NewReader(`[2, 0, null, 1]`)) + require.NoError(t, err) + defer indices.Release() + + dictOfUUID := array.NewDictionaryArray( + &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int16, ValueType: uuidType}, indices, uuidValues) + defer dictOfUUID.Release() + dictOfP1 := array.NewDictionaryArray( + &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int16, ValueType: p1Type}, indices, p1Values) + defer dictOfP1.Release() + dictExt, _, err := array.FromJSON(mem, types.NewDictExtensionType(), strings.NewReader(`["a", "b", null, "a"]`)) + require.NoError(t, err) + defer dictExt.Release() + + schema := arrow.NewSchema([]arrow.Field{ + {Name: "uuid", Type: dictOfUUID.DataType(), Nullable: true}, + {Name: "p1", Type: dictOfP1.DataType(), Nullable: true}, + {Name: "dict_ext", Type: dictExt.DataType(), Nullable: true}, + }, nil) + rec := array.NewRecordBatch(schema, []arrow.Array{dictOfUUID, dictOfP1, dictExt}, 4) + defer rec.Release() + + var buf bytes.Buffer + w, err := NewWriter(&buf, schema) + require.NoError(t, err) + require.NoError(t, w.Write(rec)) + require.NoError(t, w.Close()) + + // the extension keys of a dictionary value type are in the "dictionary" + // object, and those of an extension type with dictionary storage are in + // the field's own "metadata". + var raw struct { + Schema struct { + Fields []struct { + Metadata []metaKV `json:"metadata"` + Dictionary struct { + Metadata []metaKV `json:"metadata"` + } `json:"dictionary"` + } `json:"fields"` + } `json:"schema"` + } + require.NoError(t, json.Unmarshal(buf.Bytes(), &raw)) + fields := raw.Schema.Fields + require.Len(t, fields, 3) + assert.Nil(t, fields[0].Metadata) + assert.Equal(t, []metaKV{{Key: ipc.ExtensionTypeKeyName, Value: "arrow.uuid"}}, fields[0].Dictionary.Metadata) + assert.Nil(t, fields[1].Metadata) + assert.Equal(t, []metaKV{ + {Key: ipc.ExtensionTypeKeyName, Value: "parametric-type-1"}, + {Key: ipc.ExtensionMetadataKeyName, Value: p1Type.Serialize()}, + }, fields[1].Dictionary.Metadata) + assert.Equal(t, []metaKV{ + {Key: ipc.ExtensionTypeKeyName, Value: "dict-extension"}, + {Key: ipc.ExtensionMetadataKeyName, Value: "dict-extension-serialized"}, + }, fields[2].Metadata) + assert.Nil(t, fields[2].Dictionary.Metadata) + + r, err := NewReader(bytes.NewReader(buf.Bytes()), WithAllocator(mem)) + require.NoError(t, err) + defer r.Release() + assert.Truef(t, r.Schema().Equal(schema), "got=%s\nwant=%s", r.Schema(), schema) + got, err := r.Read() + require.NoError(t, err) + assert.Truef(t, array.RecordEqual(rec, got), "got=%s\nwant=%s", got, rec) +} + +func TestDictionaryOfUnregisteredExtension(t *testing.T) { + const data = `{ + "schema": {"fields": [{ + "name": "f", "nullable": true, "children": [], + "type": {"name": "int", "isSigned": true, "bitWidth": 32}, + "dictionary": {"id": 0, "indexType": {"name": "int", "isSigned": true, "bitWidth": 8}, "isOrdered": false, + "metadata": [{"key": "ARROW:extension:name", "value": "unregistered"}, + {"key": "ARROW:extension:metadata", "value": "data"}]} + }]}, + "batches": [], + "dictionaries": [] +}` + r, err := NewReader(strings.NewReader(data)) + require.NoError(t, err) + defer r.Release() + + want := arrow.NewSchema([]arrow.Field{{Name: "f", Nullable: true, + Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8, ValueType: arrow.PrimitiveTypes.Int32}}}, nil) + assert.Truef(t, r.Schema().Equal(want), "got=%s\nwant=%s", r.Schema(), want) + assert.Zero(t, r.Schema().Field(0).Metadata.Len()) +} + +func TestDictionaryRoundTrip(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + uuidType := extensions.NewUUIDType() + uuidDictType := &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int16, ValueType: uuidType} + uuidDict, err := array.DictArrayFromJSON(mem, uuidDictType, `[1, null, 0, 1, 2]`, + `["00000000-0000-0000-0000-000000000001", "00000000-0000-0000-0000-000000000002", null]`) + require.NoError(t, err) + defer uuidDict.Release() + uuidDictData := uuidDict.Data().(*array.Data) + + ordered, err := array.DictArrayFromJSON(mem, + &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8, ValueType: arrow.BinaryTypes.String, Ordered: true}, + `[0, 1, null, 1, 0]`, `["low", "high"]`) + require.NoError(t, err) + defer ordered.Release() + + // extension type whose storage is a dictionary of an extension type + opaqueType := extensions.NewOpaqueType(uuidDictType, "type", "vendor") + opaqueData := array.NewDataWithDictionary(opaqueType, uuidDict.Len(), uuidDictData.Buffers(), + uuidDict.NullN(), 0, uuidDictData.Dictionary().(*array.Data)) + defer opaqueData.Release() + opaqueOfDict := array.MakeFromData(opaqueData) + defer opaqueOfDict.Release() + + // list of dictionary of an extension type + offsets := memory.NewBufferBytes(arrow.Int32Traits.CastToBytes([]int32{0, 2, 2, 3, 5, 5})) + listData := array.NewData(arrow.ListOf(uuidDictType), 5, []*memory.Buffer{nil, offsets}, + []arrow.ArrayData{uuidDictData}, 0, 0) + defer listData.Release() + listOfDict := array.MakeFromData(listData) + defer listOfDict.Release() + + schema := arrow.NewSchema([]arrow.Field{ + {Name: "ordered", Type: ordered.DataType(), Nullable: true}, + {Name: "opaque_of_dict", Type: opaqueOfDict.DataType(), Nullable: true}, + {Name: "list_of_dict", Type: listOfDict.DataType(), Nullable: true}, + }, nil) + rec := array.NewRecordBatch(schema, []arrow.Array{ordered, opaqueOfDict, listOfDict}, 5) + defer rec.Release() + + var buf bytes.Buffer + w, err := NewWriter(&buf, schema) + require.NoError(t, err) + require.NoError(t, w.Write(rec)) + require.NoError(t, w.Close()) + + r, err := NewReader(bytes.NewReader(buf.Bytes()), WithAllocator(mem)) + require.NoError(t, err) + defer r.Release() + assert.Truef(t, r.Schema().Equal(schema), "got=%s\nwant=%s", r.Schema(), schema) + assert.True(t, r.Schema().Field(0).Type.(*arrow.DictionaryType).Ordered) + got, err := r.Read() + require.NoError(t, err) + assert.Truef(t, array.RecordEqual(rec, got), "got=%s\nwant=%s", got, rec) +} diff --git a/arrow/internal/dictutils/dict.go b/arrow/internal/dictutils/dict.go index a90058f13..b274cda6b 100644 --- a/arrow/internal/dictutils/dict.go +++ b/arrow/internal/dictutils/dict.go @@ -114,7 +114,11 @@ func (d *Mapper) ImportField(pos FieldPos, field arrow.Field) { if dt.ID() == arrow.DICTIONARY { d.InsertPath(pos) // import nested dicts - if nested, ok := dt.(*arrow.DictionaryType).ValueType.(arrow.NestedType); ok { + valueType := dt.(*arrow.DictionaryType).ValueType + if valueType.ID() == arrow.EXTENSION { + valueType = valueType.(arrow.ExtensionType).StorageType() + } + if nested, ok := valueType.(arrow.NestedType); ok { d.ImportFields(pos, nested.Fields()) } return diff --git a/arrow/internal/flatbuf/DictionaryEncoding.go b/arrow/internal/flatbuf/DictionaryEncoding.go index 691678027..457c4f349 100644 --- a/arrow/internal/flatbuf/DictionaryEncoding.go +++ b/arrow/internal/flatbuf/DictionaryEncoding.go @@ -99,7 +99,7 @@ func (rcv *DictionaryEncoding) IndexType(obj *Int) *Int { /// recommended to prefer signed integer types over unsigned integer types /// and to avoid uint64 indices unless they are required by an application. /// By default, dictionaries are not ordered, or the order does not have -/// semantic meaning. In some statistical, applications, dictionary-encoding +/// semantic meaning. In some statistical applications, dictionary-encoding /// is used to represent ordered categorical data, and we provide a way to /// preserve that metadata here func (rcv *DictionaryEncoding) IsOrdered() bool { @@ -111,7 +111,7 @@ func (rcv *DictionaryEncoding) IsOrdered() bool { } /// By default, dictionaries are not ordered, or the order does not have -/// semantic meaning. In some statistical, applications, dictionary-encoding +/// semantic meaning. In some statistical applications, dictionary-encoding /// is used to represent ordered categorical data, and we provide a way to /// preserve that metadata here func (rcv *DictionaryEncoding) MutateIsOrdered(n bool) bool { @@ -130,8 +130,30 @@ func (rcv *DictionaryEncoding) MutateDictionaryKind(n DictionaryKind) bool { return rcv._tab.MutateInt16Slot(10, int16(n)) } +/// Metadata for the dictionary value type, e.g. the extension type of the values +func (rcv *DictionaryEncoding) CustomMetadata(obj *KeyValue, j int) bool { + o := flatbuffers.UOffsetT(rcv._tab.Offset(12)) + if o != 0 { + x := rcv._tab.Vector(o) + x += flatbuffers.UOffsetT(j) * 4 + x = rcv._tab.Indirect(x) + obj.Init(rcv._tab.Bytes, x) + return true + } + return false +} + +func (rcv *DictionaryEncoding) CustomMetadataLength() int { + o := flatbuffers.UOffsetT(rcv._tab.Offset(12)) + if o != 0 { + return rcv._tab.VectorLen(o) + } + return 0 +} + +/// Metadata for the dictionary value type, e.g. the extension type of the values func DictionaryEncodingStart(builder *flatbuffers.Builder) { - builder.StartObject(4) + builder.StartObject(5) } func DictionaryEncodingAddId(builder *flatbuffers.Builder, id int64) { builder.PrependInt64Slot(0, id, 0) @@ -145,6 +167,12 @@ func DictionaryEncodingAddIsOrdered(builder *flatbuffers.Builder, isOrdered bool func DictionaryEncodingAddDictionaryKind(builder *flatbuffers.Builder, dictionaryKind DictionaryKind) { builder.PrependInt16Slot(3, int16(dictionaryKind), 0) } +func DictionaryEncodingAddCustomMetadata(builder *flatbuffers.Builder, customMetadata flatbuffers.UOffsetT) { + builder.PrependUOffsetTSlot(4, flatbuffers.UOffsetT(customMetadata), 0) +} +func DictionaryEncodingStartCustomMetadataVector(builder *flatbuffers.Builder, numElems int) flatbuffers.UOffsetT { + return builder.StartVector(4, numElems, 4) +} func DictionaryEncodingEnd(builder *flatbuffers.Builder) flatbuffers.UOffsetT { return builder.EndObject() } diff --git a/arrow/internal/flight_integration/scenario.go b/arrow/internal/flight_integration/scenario.go index b95184430..ac6d48af0 100644 --- a/arrow/internal/flight_integration/scenario.go +++ b/arrow/internal/flight_integration/scenario.go @@ -45,6 +45,8 @@ import ( "google.golang.org/grpc/status" "google.golang.org/protobuf/proto" timestamppb "google.golang.org/protobuf/types/known/timestamppb" + + _ "github.com/apache/arrow-go/v18/arrow/internal/integration" ) type Scenario interface { diff --git a/arrow/internal/integration/integration.go b/arrow/internal/integration/integration.go new file mode 100644 index 000000000..b3461dd45 --- /dev/null +++ b/arrow/internal/integration/integration.go @@ -0,0 +1,31 @@ +// 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 integration holds setup shared by the programs used for +// cross-implementation integration testing (the IPC JSON integration and +// file/stream conversion programs and the Flight integration client and +// server). Importing it registers the canonical extension types, so that +// those programs read and write them, including as dictionary value types, +// like other implementations do. +// +// The C Data integration library does not import it: the C Data importer +// does not consult the extension type registry. +package integration + +import ( + // register the canonical extension types + _ "github.com/apache/arrow-go/v18/arrow/extensions" +) diff --git a/arrow/ipc/cmd/arrow-file-to-stream/main.go b/arrow/ipc/cmd/arrow-file-to-stream/main.go index f7fc29379..5b0df340f 100644 --- a/arrow/ipc/cmd/arrow-file-to-stream/main.go +++ b/arrow/ipc/cmd/arrow-file-to-stream/main.go @@ -27,6 +27,8 @@ import ( "github.com/apache/arrow-go/v18/arrow/arrio" "github.com/apache/arrow-go/v18/arrow/ipc" "github.com/apache/arrow-go/v18/arrow/memory" + + _ "github.com/apache/arrow-go/v18/arrow/internal/integration" ) func main() { diff --git a/arrow/ipc/cmd/arrow-json-integration-test/main.go b/arrow/ipc/cmd/arrow-json-integration-test/main.go index 44535116b..e264ab919 100644 --- a/arrow/ipc/cmd/arrow-json-integration-test/main.go +++ b/arrow/ipc/cmd/arrow-json-integration-test/main.go @@ -26,6 +26,8 @@ import ( "github.com/apache/arrow-go/v18/arrow/arrio" "github.com/apache/arrow-go/v18/arrow/internal/arrjson" "github.com/apache/arrow-go/v18/arrow/ipc" + + _ "github.com/apache/arrow-go/v18/arrow/internal/integration" ) func main() { diff --git a/arrow/ipc/cmd/arrow-stream-to-file/main.go b/arrow/ipc/cmd/arrow-stream-to-file/main.go index feb111aa5..3328122cd 100644 --- a/arrow/ipc/cmd/arrow-stream-to-file/main.go +++ b/arrow/ipc/cmd/arrow-stream-to-file/main.go @@ -27,6 +27,8 @@ import ( "github.com/apache/arrow-go/v18/arrow/arrio" "github.com/apache/arrow-go/v18/arrow/ipc" "github.com/apache/arrow-go/v18/arrow/memory" + + _ "github.com/apache/arrow-go/v18/arrow/internal/integration" ) func main() { diff --git a/arrow/ipc/metadata.go b/arrow/ipc/metadata.go index 7a3a2321d..e355d6b2b 100644 --- a/arrow/ipc/metadata.go +++ b/arrow/ipc/metadata.go @@ -259,6 +259,9 @@ type fieldVisitor struct { offset flatbuffers.UOffsetT kids []flatbuffers.UOffsetT meta map[string]string + // dictMeta is the metadata of the value type of a dictionary-encoded + // field, written to DictionaryEncoding.custom_metadata. + dictMeta arrow.Metadata } func (fv *fieldVisitor) visit(field arrow.Field) { @@ -497,6 +500,15 @@ func (fv *fieldVisitor) visit(field arrow.Field) { case *arrow.DictionaryType: field.Type = dt.ValueType + if extType, ok := dt.ValueType.(arrow.ExtensionType); ok { + // the extension metadata of a dictionary's value type goes in + // DictionaryEncoding.custom_metadata so that it can't be confused + // with an extension type whose storage type is a dictionary. + field.Type = extType.StorageType() + fv.dictMeta = arrow.NewMetadata( + []string{ExtensionTypeKeyName, ExtensionMetadataKeyName}, + []string{extType.ExtensionName(), extType.Serialize()}) + } fv.visit(field) case arrow.UnionType: @@ -550,7 +562,8 @@ func (fv *fieldVisitor) result(field arrow.Field) flatbuffers.UOffsetT { var dictFB flatbuffers.UOffsetT if storageType.ID() == arrow.DICTIONARY { - idxType := field.Type.(*arrow.DictionaryType).IndexType.(arrow.FixedWidthDataType) + dictType := storageType.(*arrow.DictionaryType) + idxType := dictType.IndexType.(arrow.FixedWidthDataType) dictID, err := fv.memo.GetFieldID(fv.pos.Path()) if err != nil { @@ -564,10 +577,12 @@ func (fv *fieldVisitor) result(field arrow.Field) flatbuffers.UOffsetT { signed = true } indexTypeOffset := intToFB(fv.b, int32(idxType.BitWidth()), signed) + dictMetaFB := metadataToFB(fv.b, fv.dictMeta, flatbuf.DictionaryEncodingStartCustomMetadataVector) flatbuf.DictionaryEncodingStart(fv.b) flatbuf.DictionaryEncodingAddId(fv.b, dictID) flatbuf.DictionaryEncodingAddIndexType(fv.b, indexTypeOffset) - flatbuf.DictionaryEncodingAddIsOrdered(fv.b, field.Type.(*arrow.DictionaryType).Ordered) + flatbuf.DictionaryEncodingAddIsOrdered(fv.b, dictType.Ordered) + flatbuf.DictionaryEncodingAddCustomMetadata(fv.b, dictMetaFB) dictFB = flatbuf.DictionaryEncodingEnd(fv.b) } @@ -647,6 +662,16 @@ func typeFromFB(field *flatbuf.Field, pos dictutils.FieldPos, children []arrow.F } dictValueType = dt + // the dictionary value type may be an extension type, described by + // the DictionaryEncoding's own custom metadata. Other keys have + // nowhere to go and are dropped. + dictMeta, err := metadataFromFB(encoding) + if err != nil { + return nil, err + } + if dictValueType, err = extensionTypeFromMetadata(dictValueType, &dictMeta); err != nil { + return nil, err + } dt = &arrow.DictionaryType{IndexType: idxType, ValueType: dictValueType, Ordered: encoding.IsOrdered()} dictID = encoding.Id() @@ -660,56 +685,58 @@ func typeFromFB(field *flatbuf.Field, pos dictutils.FieldPos, children []arrow.F } // look for extension metadata in custom metadata field. - if md.Len() > 0 { - i := md.FindKey(ExtensionTypeKeyName) - if i < 0 { - return dt, err - } + return extensionTypeFromMetadata(dt, md) +} - extType := arrow.GetExtensionType(md.Values()[i]) - if extType == nil { - // if the extension type is unknown, we do not error here. - // simply return the storage type. - return dt, err - } +// extensionTypeFromMetadata returns the extension type described by the +// ARROW:extension:name and ARROW:extension:metadata keys of md, with storage +// type dt, and removes those keys from md. If md does not name an extension +// type, or names one that is not registered, dt and md are returned unchanged. +func extensionTypeFromMetadata(dt arrow.DataType, md *arrow.Metadata) (arrow.DataType, error) { + if md.Len() == 0 { + return dt, nil + } - var ( - data string - dataIdx int - ) + i := md.FindKey(ExtensionTypeKeyName) + if i < 0 { + return dt, nil + } - if dataIdx = md.FindKey(ExtensionMetadataKeyName); dataIdx >= 0 { - data = md.Values()[dataIdx] - } + extType := arrow.GetExtensionType(md.Values()[i]) + if extType == nil { + // if the extension type is unknown, we do not error here. + // simply return the storage type. + return dt, nil + } - dt, err = extType.Deserialize(dt, data) - if err != nil { - return dt, err - } + var ( + data string + dataIdx int + ) - mdkeys := md.Keys() - mdvals := md.Values() - if dataIdx < 0 { - // if there was no extension metadata, just the name, we only have to - // remove the extension name metadata key/value to ensure roundtrip - // metadata consistency - *md = arrow.NewMetadata(append(mdkeys[:i], mdkeys[i+1:]...), append(mdvals[:i], mdvals[i+1:]...)) - } else { - // if there was extension metadata, we need to remove both the type name - // and the extension metadata keys and values. - newkeys := make([]string, 0, md.Len()-2) - newvals := make([]string, 0, md.Len()-2) - for j := range mdkeys { - if j != i && j != dataIdx { // copy everything except the extension metadata keys/values - newkeys = append(newkeys, mdkeys[j]) - newvals = append(newvals, mdvals[j]) - } - } - *md = arrow.NewMetadata(newkeys, newvals) + if dataIdx = md.FindKey(ExtensionMetadataKeyName); dataIdx >= 0 { + data = md.Values()[dataIdx] + } + + dt, err := extType.Deserialize(dt, data) + if err != nil { + return dt, err + } + + // remove the extension type keys from the metadata, building new slices + // since those returned by Keys and Values may be shared with other copies + // of md. + newkeys := make([]string, 0, md.Len()) + newvals := make([]string, 0, md.Len()) + for j, k := range md.Keys() { + if j != i && j != dataIdx { + newkeys = append(newkeys, k) + newvals = append(newvals, md.Values()[j]) } } + *md = arrow.NewMetadata(newkeys, newvals) - return dt, err + return dt, nil } func concreteTypeFromFB(typ flatbuf.Type, data flatbuffers.Table, children []arrow.Field) (arrow.DataType, error) { diff --git a/arrow/ipc/metadata_test.go b/arrow/ipc/metadata_test.go index ac8820dcc..d37ba47aa 100644 --- a/arrow/ipc/metadata_test.go +++ b/arrow/ipc/metadata_test.go @@ -18,7 +18,10 @@ package ipc import ( "bytes" + "errors" + "io" "reflect" + "strings" "testing" "github.com/apache/arrow-go/v18/arrow" @@ -27,8 +30,10 @@ import ( "github.com/apache/arrow-go/v18/arrow/internal/dictutils" "github.com/apache/arrow-go/v18/arrow/internal/flatbuf" "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/internal/types" flatbuffers "github.com/google/flatbuffers/go" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestRWSchema(t *testing.T) { @@ -266,3 +271,538 @@ func TestUnrecognizedExtensionType(t *testing.T) { assert.Truef(t, array.RecordEqual(rec, batchNoExt), "expected: %s\ngot: %s\n", batchNoExt, rec) } + +// registerExtensionType registers extType for the duration of the test if no +// extension type of that name is registered yet. +func registerExtensionType(t *testing.T, extType arrow.ExtensionType) { + t.Helper() + if arrow.GetExtensionType(extType.ExtensionName()) != nil { + return + } + if err := arrow.RegisterExtensionType(extType); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { arrow.UnregisterExtensionType(extType.ExtensionName()) }) +} + +// fieldFBRoundTrip serializes field to a flatbuffers Field and returns that +// Field along with the arrow.Field read back from it. +func fieldFBRoundTrip(t *testing.T, field arrow.Field) (*flatbuf.Field, arrow.Field) { + t.Helper() + schema := arrow.NewSchema([]arrow.Field{field}, nil) + var memo dictutils.Memo + memo.Mapper.ImportSchema(schema) + + b := flatbuffers.NewBuilder(0) + b.Finish(schemaToFB(b, schema, &memo.Mapper)) + fbSchema := flatbuf.GetRootAsSchema(b.FinishedBytes(), 0) + + var fbField flatbuf.Field + if !fbSchema.Fields(&fbField, 0) { + t.Fatal("could not read field from flatbuffer") + } + + readMemo := dictutils.NewMemo() + defer readMemo.Clear() + got, err := schemaFromFB(fbSchema, &readMemo) + if err != nil { + t.Fatal(err) + } + return &fbField, got.Field(0) +} + +func fbMetadata(t *testing.T, md customMetadataer) map[string]string { + t.Helper() + meta, err := metadataFromFB(md) + if err != nil { + t.Fatal(err) + } + return meta.ToMap() +} + +func TestDictionaryOfExtensionMetadata(t *testing.T) { + registerExtensionType(t, types.NewParametric1Type(1)) + registerExtensionType(t, types.NewDictExtensionType()) + + uuidType := extensions.NewUUIDType() + p1Type := types.NewParametric1Type(6) + dictExtType := types.NewDictExtensionType() + opaqueOfDictOfUUID := extensions.NewOpaqueType( + &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int16, ValueType: uuidType}, + "type", "vendor") + dictOfNested := &arrow.DictionaryType{ + IndexType: arrow.PrimitiveTypes.Int32, + ValueType: extensions.NewOpaqueType(arrow.StructOf(arrow.Field{ + Name: "d", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8, ValueType: uuidType}, + }), "type", "vendor"), + } + + tests := []struct { + name string + typ arrow.DataType + fieldMeta map[string]string + dictMeta map[string]string + }{ + { + name: "dictionary of uuid", + typ: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: uuidType}, + dictMeta: map[string]string{ExtensionTypeKeyName: "arrow.uuid", ExtensionMetadataKeyName: ""}, + }, + { + name: "dictionary of extension type with metadata", + typ: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Uint8, ValueType: p1Type, Ordered: true}, + dictMeta: map[string]string{ExtensionTypeKeyName: "parametric-type-1", ExtensionMetadataKeyName: p1Type.Serialize()}, + }, + { + name: "extension type with dictionary storage", + typ: dictExtType, + fieldMeta: map[string]string{ExtensionTypeKeyName: "dict-extension", ExtensionMetadataKeyName: "dict-extension-serialized"}, + }, + { + name: "extension type with dictionary of extension type storage", + typ: opaqueOfDictOfUUID, + fieldMeta: map[string]string{ExtensionTypeKeyName: "arrow.opaque", ExtensionMetadataKeyName: opaqueOfDictOfUUID.Serialize()}, + dictMeta: map[string]string{ExtensionTypeKeyName: "arrow.uuid", ExtensionMetadataKeyName: ""}, + }, + { + name: "dictionary of extension type with nested dictionary storage", + typ: dictOfNested, + dictMeta: map[string]string{ExtensionTypeKeyName: "arrow.opaque", ExtensionMetadataKeyName: dictOfNested.ValueType.(arrow.ExtensionType).Serialize()}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + field := arrow.Field{Name: "f", Type: tc.typ, Nullable: true, + Metadata: arrow.NewMetadata([]string{"k"}, []string{"v"})} + fbField, got := fieldFBRoundTrip(t, field) + + wantFieldMeta := map[string]string{"k": "v"} + for k, v := range tc.fieldMeta { + wantFieldMeta[k] = v + } + assert.Equal(t, wantFieldMeta, fbMetadata(t, fbField)) + + encoding := fbField.Dictionary(nil) + if assert.NotNil(t, encoding) { + if tc.dictMeta == nil { + assert.Zero(t, encoding.CustomMetadataLength()) + } else { + assert.Equal(t, tc.dictMeta, fbMetadata(t, encoding)) + } + } + + assert.Truef(t, got.Equal(field), "got=%s\nwant=%s", got, field) + }) + } +} + +func TestDictionaryOfExtensionRoundTrip(t *testing.T) { + registerExtensionType(t, types.NewParametric1Type(1)) + registerExtensionType(t, types.NewDictExtensionType()) + + pool := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer pool.AssertSize(t, 0) + + uuidValues := exampleUUID(pool) + defer uuidValues.Release() + + p1Storage, _, err := array.FromJSON(pool, arrow.PrimitiveTypes.Int32, strings.NewReader(`[1, 2, null, 4]`)) + if err != nil { + t.Fatal(err) + } + defer p1Storage.Release() + p1Values := array.NewExtensionArrayWithStorage(types.NewParametric1Type(6), p1Storage) + defer p1Values.Release() + + dictExt, _, err := array.FromJSON(pool, types.NewDictExtensionType(), strings.NewReader(`["a", "b", null, "a", "c"]`)) + if err != nil { + t.Fatal(err) + } + defer dictExt.Release() + + indices, _, err := array.FromJSON(pool, arrow.PrimitiveTypes.Int32, strings.NewReader(`[3, 0, 1, null, 3]`)) + if err != nil { + t.Fatal(err) + } + defer indices.Release() + + dictOfUUID := array.NewDictionaryArray( + &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: uuidValues.DataType()}, indices, uuidValues) + defer dictOfUUID.Release() + dictOfP1 := array.NewDictionaryArray( + &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: p1Values.DataType()}, indices, p1Values) + defer dictOfP1.Release() + + schema := arrow.NewSchema([]arrow.Field{ + {Name: "uuid", Type: dictOfUUID.DataType(), Nullable: true}, + {Name: "p1", Type: dictOfP1.DataType(), Nullable: true}, + {Name: "dict_ext", Type: dictExt.DataType(), Nullable: true}, + }, nil) + batch := array.NewRecordBatch(schema, []arrow.Array{dictOfUUID, dictOfP1, dictExt}, 5) + defer batch.Release() + + for _, format := range []string{"stream", "file"} { + t.Run(format, func(t *testing.T) { + var ( + buf bytes.Buffer + recs []arrow.RecordBatch + ) + switch format { + case "stream": + w := NewWriter(&buf, WithAllocator(pool), WithSchema(schema)) + assert.NoError(t, w.Write(batch)) + assert.NoError(t, w.Close()) + + r, err := NewReader(&buf, WithAllocator(pool)) + if err != nil { + t.Fatal(err) + } + defer r.Release() + assert.Truef(t, r.Schema().Equal(schema), "got=%s\nwant=%s", r.Schema(), schema) + for r.Next() { + rec := r.RecordBatch() + rec.Retain() + recs = append(recs, rec) + } + assert.NoError(t, r.Err()) + case "file": + w, err := NewFileWriter(&buf, WithAllocator(pool), WithSchema(schema)) + if err != nil { + t.Fatal(err) + } + assert.NoError(t, w.Write(batch)) + assert.NoError(t, w.Close()) + + r, err := NewFileReader(bytes.NewReader(buf.Bytes()), WithAllocator(pool)) + if err != nil { + t.Fatal(err) + } + defer r.Close() + assert.Truef(t, r.Schema().Equal(schema), "got=%s\nwant=%s", r.Schema(), schema) + for i := 0; i < r.NumRecords(); i++ { + rec, err := r.RecordBatch(i) + if err != nil { + t.Fatal(err) + } + rec.Retain() + recs = append(recs, rec) + } + } + + if assert.Len(t, recs, 1) { + assert.Truef(t, array.RecordEqual(batch, recs[0]), "got=%s\nwant=%s", recs[0], batch) + assert.IsType(t, extensions.NewUUIDType(), + recs[0].Column(0).(*array.Dictionary).Dictionary().DataType()) + } + for _, rec := range recs { + rec.Release() + } + }) + } +} + +func TestDictionaryOfUnregisteredExtension(t *testing.T) { + pool := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer pool.AssertSize(t, 0) + + // an extension type whose name is never registered + extType := types.NewParametric2Type(424242) + require.Nil(t, arrow.GetExtensionType(extType.ExtensionName())) + + storage, _, err := array.FromJSON(pool, arrow.PrimitiveTypes.Int32, strings.NewReader(`[1, null, 3, 4]`)) + require.NoError(t, err) + defer storage.Release() + values := array.NewExtensionArrayWithStorage(extType, storage) + defer values.Release() + indices, _, err := array.FromJSON(pool, arrow.PrimitiveTypes.Int32, strings.NewReader(`[1, 0, null, 3]`)) + require.NoError(t, err) + defer indices.Release() + dictArr := array.NewDictionaryArray( + &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: extType}, indices, values) + defer dictArr.Release() + + schema := arrow.NewSchema([]arrow.Field{{Name: "f0", Type: dictArr.DataType(), Nullable: true}}, nil) + batch := array.NewRecordBatch(schema, []arrow.Array{dictArr}, 4) + defer batch.Release() + + var buf bytes.Buffer + wr := NewWriter(&buf, WithAllocator(pool), WithSchema(schema)) + require.NoError(t, wr.Write(batch)) + require.NoError(t, wr.Close()) + + rdr, err := NewReader(&buf, WithAllocator(pool)) + require.NoError(t, err) + defer rdr.Release() + + // the dictionary keeps its storage value type and the field gets no + // extension metadata + storageDictType := &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.PrimitiveTypes.Int32} + want := arrow.NewSchema([]arrow.Field{{Name: "f0", Type: storageDictType, Nullable: true}}, nil) + assert.Truef(t, rdr.Schema().Equal(want), "got=%s\nwant=%s", rdr.Schema(), want) + assert.Zero(t, rdr.Schema().Field(0).Metadata.Len()) + + require.True(t, rdr.Next()) + storageDict := array.NewDictionaryArray(storageDictType, indices, storage) + defer storageDict.Release() + assert.Truef(t, array.Equal(storageDict, rdr.RecordBatch().Column(0)), + "got=%s\nwant=%s", rdr.RecordBatch().Column(0), storageDict) +} + +// dictFieldSchemaFB builds a flatbuffers Schema with a single dictionary +// encoded field of the given value type, whose DictionaryEncoding has the +// given custom metadata. +func dictFieldSchemaFB(t *testing.T, valueType arrow.DataType, dictMeta arrow.Metadata) *flatbuf.Schema { + t.Helper() + b := flatbuffers.NewBuilder(0) + + fv := fieldVisitor{b: b, meta: make(map[string]string)} + fv.visit(arrow.Field{Name: "f", Type: valueType}) + nameFB := b.CreateString("f") + flatbuf.FieldStartChildrenVector(b, 0) + kidsFB := b.EndVector(0) + indexTypeFB := intToFB(b, 8, true) + dictMetaFB := metadataToFB(b, dictMeta, flatbuf.DictionaryEncodingStartCustomMetadataVector) + flatbuf.DictionaryEncodingStart(b) + flatbuf.DictionaryEncodingAddId(b, 0) + flatbuf.DictionaryEncodingAddIndexType(b, indexTypeFB) + flatbuf.DictionaryEncodingAddCustomMetadata(b, dictMetaFB) + dictFB := flatbuf.DictionaryEncodingEnd(b) + + flatbuf.FieldStart(b) + flatbuf.FieldAddName(b, nameFB) + flatbuf.FieldAddNullable(b, true) + flatbuf.FieldAddTypeType(b, fv.dtype) + flatbuf.FieldAddType(b, fv.offset) + flatbuf.FieldAddDictionary(b, dictFB) + flatbuf.FieldAddChildren(b, kidsFB) + fieldFB := flatbuf.FieldEnd(b) + + flatbuf.SchemaStartFieldsVector(b, 1) + b.PrependUOffsetT(fieldFB) + fieldsFB := b.EndVector(1) + flatbuf.SchemaStart(b) + flatbuf.SchemaAddFields(b, fieldsFB) + b.Finish(flatbuf.SchemaEnd(b)) + return flatbuf.GetRootAsSchema(b.FinishedBytes(), 0) +} + +func TestDictionaryEncodingCustomMetadataOnRead(t *testing.T) { + uuidStorage := &arrow.FixedSizeBinaryType{ByteWidth: 16} + tests := []struct { + name string + valueType arrow.DataType + dictMeta map[string]string + want arrow.DataType + wantErr string + }{ + { + name: "non-extension keys are dropped", + valueType: arrow.PrimitiveTypes.Int32, + dictMeta: map[string]string{"k": "v"}, + want: arrow.PrimitiveTypes.Int32, + }, + { + name: "unregistered extension name is dropped", + valueType: arrow.PrimitiveTypes.Int32, + dictMeta: map[string]string{ExtensionTypeKeyName: "arrow-go.test.never-registered", + ExtensionMetadataKeyName: "data", "k": "v"}, + want: arrow.PrimitiveTypes.Int32, + }, + { + name: "registered extension with other keys", + valueType: uuidStorage, + dictMeta: map[string]string{ExtensionTypeKeyName: "arrow.uuid", "k": "v"}, + want: extensions.NewUUIDType(), + }, + { + name: "registered extension over the wrong storage type", + valueType: arrow.PrimitiveTypes.Int32, + dictMeta: map[string]string{ExtensionTypeKeyName: "arrow.uuid"}, + wantErr: "invalid storage type for UUIDType: int32", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + memo := dictutils.NewMemo() + defer memo.Clear() + got, err := schemaFromFB(dictFieldSchemaFB(t, tc.valueType, arrow.MetadataFrom(tc.dictMeta)), &memo) + if tc.wantErr != "" { + assert.ErrorContains(t, err, tc.wantErr) + return + } + require.NoError(t, err) + + want := &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8, ValueType: tc.want} + assert.Truef(t, arrow.TypeEqual(want, got.Field(0).Type), "got=%s\nwant=%s", got.Field(0).Type, want) + assert.Zero(t, got.Field(0).Metadata.Len()) + }) + } +} + +func TestExtensionTypeFromMetadataDoesNotMutateInput(t *testing.T) { + keys := []string{"a", ExtensionTypeKeyName, "b"} + vals := []string{"1", "arrow.uuid", "2"} + md := arrow.NewMetadata(keys, vals) + orig := md + + dt, err := extensionTypeFromMetadata(&arrow.FixedSizeBinaryType{ByteWidth: 16}, &md) + require.NoError(t, err) + assert.Equal(t, "arrow.uuid", dt.(arrow.ExtensionType).ExtensionName()) + assert.Equal(t, []string{"a", "b"}, md.Keys()) + assert.Equal(t, []string{"1", "2"}, md.Values()) + assert.Equal(t, keys, orig.Keys()) + assert.Equal(t, vals, orig.Values()) +} + +func TestDictionaryOfNestedDictionaryExtensionRoundTrip(t *testing.T) { + pool := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer pool.AssertSize(t, 0) + + // dictionary>>> + uuidValues := exampleUUID(pool) + defer uuidValues.Release() + innerIdx, _, err := array.FromJSON(pool, arrow.PrimitiveTypes.Int8, strings.NewReader(`[3, 1, null, 2, 1]`)) + require.NoError(t, err) + defer innerIdx.Release() + inner := array.NewDictionaryArray( + &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8, ValueType: uuidValues.DataType()}, innerIdx, uuidValues) + defer inner.Release() + + structArr, err := array.NewStructArrayWithFields([]arrow.Array{inner}, + []arrow.Field{{Name: "d", Type: inner.DataType(), Nullable: true}}) + require.NoError(t, err) + defer structArr.Release() + opaqueType := extensions.NewOpaqueType(structArr.DataType(), "type", "vendor") + opaqueValues := array.NewExtensionArrayWithStorage(opaqueType, structArr) + defer opaqueValues.Release() + + outerIdx, _, err := array.FromJSON(pool, arrow.PrimitiveTypes.Int32, strings.NewReader(`[4, 0, null, 2, 2, 1]`)) + require.NoError(t, err) + defer outerIdx.Release() + outer := array.NewDictionaryArray( + &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: opaqueType}, outerIdx, opaqueValues) + defer outer.Release() + + schema := arrow.NewSchema([]arrow.Field{{Name: "f", Type: outer.DataType(), Nullable: true}}, nil) + batch := array.NewRecordBatch(schema, []arrow.Array{outer}, int64(outer.Len())) + defer batch.Release() + + var buf bytes.Buffer + w := NewWriter(&buf, WithAllocator(pool), WithSchema(schema)) + require.NoError(t, w.Write(batch)) + require.NoError(t, w.Close()) + + r, err := NewReader(&buf, WithAllocator(pool)) + require.NoError(t, err) + defer r.Release() + assert.Truef(t, r.Schema().Equal(schema), "got=%s\nwant=%s", r.Schema(), schema) + require.True(t, r.Next()) + assert.Truef(t, array.RecordEqual(batch, r.RecordBatch()), "got=%s\nwant=%s", r.RecordBatch(), batch) + assert.False(t, r.Next()) + assert.NoError(t, r.Err()) +} + +func TestDictionaryOfExtensionDeltas(t *testing.T) { + pool := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer pool.AssertSize(t, 0) + + uuidType := extensions.NewUUIDType() + dictType := &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: uuidType} + schema := arrow.NewSchema([]arrow.Field{{Name: "f", Type: dictType, Nullable: true}}, nil) + + makeBatch := func(values, indices string) arrow.RecordBatch { + dict, _, err := array.FromJSON(pool, uuidType, strings.NewReader(values)) + require.NoError(t, err) + defer dict.Release() + idx, _, err := array.FromJSON(pool, arrow.PrimitiveTypes.Int32, strings.NewReader(indices)) + require.NoError(t, err) + defer idx.Release() + arr := array.NewDictionaryArray(dictType, idx, dict) + defer arr.Release() + return array.NewRecordBatch(schema, []arrow.Array{arr}, int64(arr.Len())) + } + + const ( + uuid1 = `"00000000-0000-0000-0000-000000000001"` + uuid2 = `"00000000-0000-0000-0000-000000000002"` + uuid3 = `"00000000-0000-0000-0000-000000000003"` + ) + batches := []arrow.RecordBatch{ + makeBatch(`[`+uuid1+`, `+uuid2+`]`, `[0, 1, null, 0]`), + makeBatch(`[`+uuid1+`, `+uuid2+`, `+uuid3+`]`, `[2, 0, 1, 2]`), + } + defer func() { + for _, b := range batches { + b.Release() + } + }() + + var buf bytes.Buffer + w := NewWriter(&buf, WithAllocator(pool), WithSchema(schema), WithDictionaryDeltas(true)) + for _, b := range batches { + require.NoError(t, w.Write(b)) + } + require.NoError(t, w.Close()) + + // the second dictionary is written as a delta + var deltas []bool + mr := NewMessageReader(bytes.NewReader(buf.Bytes()), WithAllocator(pool)) + for { + msg, err := mr.Message() + if errors.Is(err, io.EOF) { + break + } + require.NoError(t, err) + if msg.Type() == MessageDictionaryBatch { + var md flatbuf.DictionaryBatch + initFB(&md, msg.msg.Header) + deltas = append(deltas, md.IsDelta()) + } + } + mr.Release() + assert.Equal(t, []bool{false, true}, deltas) + + r, err := NewReader(bytes.NewReader(buf.Bytes()), WithAllocator(pool)) + require.NoError(t, err) + defer r.Release() + assert.Truef(t, r.Schema().Equal(schema), "got=%s\nwant=%s", r.Schema(), schema) + for i, want := range batches { + require.True(t, r.Next(), "batch %d", i) + got := r.RecordBatch() + assert.Truef(t, array.RecordEqual(want, got), "batch %d: got=%s\nwant=%s", i, got, want) + assert.IsType(t, uuidType, got.Column(0).(*array.Dictionary).Dictionary().DataType()) + } + assert.False(t, r.Next()) + assert.NoError(t, r.Err()) +} + +func TestDictionaryLegacyFieldExtensionMetadata(t *testing.T) { + // Writers that predate DictionaryEncoding.custom_metadata put the + // extension keys of a dictionary's value type in Field.custom_metadata, + // which reads as an extension type with dictionary storage. + storageDictType := &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: &arrow.FixedSizeBinaryType{ByteWidth: 16}} + + t.Run("registered", func(t *testing.T) { + field := arrow.Field{Name: "f", Type: storageDictType, Nullable: true, + Metadata: arrow.NewMetadata([]string{ExtensionMetadataKeyName, ExtensionTypeKeyName}, []string{"", "arrow.uuid"})} + schema := arrow.NewSchema([]arrow.Field{field}, nil) + var memo dictutils.Memo + memo.Mapper.ImportSchema(schema) + b := flatbuffers.NewBuilder(0) + b.Finish(schemaToFB(b, schema, &memo.Mapper)) + + readMemo := dictutils.NewMemo() + _, err := schemaFromFB(flatbuf.GetRootAsSchema(b.FinishedBytes(), 0), &readMemo) + assert.ErrorContains(t, err, "invalid storage type for UUIDType: dictionary") + }) + + t.Run("unregistered", func(t *testing.T) { + field := arrow.Field{Name: "f", Type: storageDictType, Nullable: true, + Metadata: arrow.NewMetadata([]string{ExtensionTypeKeyName}, []string{"unregistered"})} + fbField, got := fieldFBRoundTrip(t, field) + assert.Zero(t, fbField.Dictionary(nil).CustomMetadataLength()) + assert.Truef(t, got.Equal(field), "got=%s\nwant=%s", got, field) + }) +}