feat(arrow/ipc): store dictionary value type metadata in DictionaryEncoding

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)

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)
+}

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