From d24193dc0b423f1ed1e046586fc031233c4d7a2b Mon Sep 17 00:00:00 2001 From: Curt Hagenlocher Date: Sun, 18 Oct 2026 18:28:49 -0700 Subject: [PATCH] Support DictionaryEncoding.custom_metadata for dictionaries of extension types Implements the proposed format 1.6 change (apache/arrow#53994): the extension metadata of a dictionary's value type is written to and read from DictionaryEncoding.custom_metadata. Field.custom_metadata on a dictionary-encoded field now describes the whole dictionary type, as in Arrow C++ and arrow-go. Co-Authored-By: Claude Opus 5.5 --- format/Schema.fbs | 4 + src/Apache.Arrow/C/CArrowSchemaExporter.cs | 5 +- .../Flatbuf/DictionaryEncoding.cs | 18 +- src/Apache.Arrow/Ipc/ArrowStreamWriter.cs | 40 +++- src/Apache.Arrow/Ipc/DictionaryMemo.cs | 5 +- src/Apache.Arrow/Ipc/MessageSerializer.cs | 57 +++-- test/Apache.Arrow.Tests/ExtensionTypeTests.cs | 226 +++++++++++++++++- 7 files changed, 322 insertions(+), 33 deletions(-) diff --git a/format/Schema.fbs b/format/Schema.fbs index f902b6bc..9eb5de1e 100644 --- a/format/Schema.fbs +++ b/format/Schema.fbs @@ -25,6 +25,7 @@ /// Version 1.4 - Add BinaryView, Utf8View, variadicBufferCounts, ListView, and /// LargeListView. /// Version 1.5 - Add 32-bit and 64-bit as allowed bit widths for Decimal +/// Version 1.6 - Add custom_metadata to DictionaryEncoding. namespace org.apache.arrow.flatbuf; @@ -503,6 +504,9 @@ table DictionaryEncoding { isOrdered: bool; dictionaryKind: DictionaryKind; + + /// Metadata for the dictionary value type, e.g. the extension type of the values + custom_metadata: [ KeyValue ]; } /// ---------------------------------------------------------------------- diff --git a/src/Apache.Arrow/C/CArrowSchemaExporter.cs b/src/Apache.Arrow/C/CArrowSchemaExporter.cs index da875c01..011134ca 100644 --- a/src/Apache.Arrow/C/CArrowSchemaExporter.cs +++ b/src/Apache.Arrow/C/CArrowSchemaExporter.cs @@ -69,8 +69,9 @@ private static unsafe void ExportType(IArrowType datatype, CArrowSchema* schema, // The C Data Interface puts the extension metadata of a dictionary's value type on // the dictionary child; on the field itself, it means the extension's storage is the - // whole dictionary. The IPC reader leaves a value type's extension metadata on the - // field, so move it only when it names the value type. Unresolved metadata stays put. + // whole dictionary. Earlier versions of this library's IPC reader left a value type's + // extension metadata on the field, so move it only when it names the value type. + // Unresolved metadata stays put. IReadOnlyDictionary dictionaryMetadata = null; if (datatype is DictionaryType dictionaryType && dictionaryType.ValueType is ExtensionType valueType && diff --git a/src/Apache.Arrow/Flatbuf/DictionaryEncoding.cs b/src/Apache.Arrow/Flatbuf/DictionaryEncoding.cs index ecf1dd78..69d79229 100644 --- a/src/Apache.Arrow/Flatbuf/DictionaryEncoding.cs +++ b/src/Apache.Arrow/Flatbuf/DictionaryEncoding.cs @@ -35,25 +35,36 @@ internal struct DictionaryEncoding : IFlatbufferObject /// preserve that metadata here public bool IsOrdered { get { int o = __p.__offset(8); return o != 0 ? 0!=__p.bb.Get(o + __p.bb_pos) : (bool)false; } } public DictionaryKind DictionaryKind { get { int o = __p.__offset(10); return o != 0 ? (DictionaryKind)__p.bb.GetShort(o + __p.bb_pos) : DictionaryKind.DenseArray; } } + /// Metadata for the dictionary value type, e.g. the extension type of the values + public KeyValue? CustomMetadata(int j) { int o = __p.__offset(12); return o != 0 ? (KeyValue?)(new KeyValue()).__assign(__p.__indirect(__p.__vector(o) + j * 4), __p.bb) : null; } + public int CustomMetadataLength { get { int o = __p.__offset(12); return o != 0 ? __p.__vector_len(o) : 0; } } public static Offset CreateDictionaryEncoding(FlatBufferBuilder builder, long id = 0, Offset indexTypeOffset = default(Offset), bool isOrdered = false, - DictionaryKind dictionaryKind = DictionaryKind.DenseArray) { - builder.StartTable(4); + DictionaryKind dictionaryKind = DictionaryKind.DenseArray, + VectorOffset custom_metadataOffset = default(VectorOffset)) { + builder.StartTable(5); DictionaryEncoding.AddId(builder, id); + DictionaryEncoding.AddCustomMetadata(builder, custom_metadataOffset); DictionaryEncoding.AddIndexType(builder, indexTypeOffset); DictionaryEncoding.AddDictionaryKind(builder, dictionaryKind); DictionaryEncoding.AddIsOrdered(builder, isOrdered); return DictionaryEncoding.EndDictionaryEncoding(builder); } - public static void StartDictionaryEncoding(FlatBufferBuilder builder) { builder.StartTable(4); } + public static void StartDictionaryEncoding(FlatBufferBuilder builder) { builder.StartTable(5); } public static void AddId(FlatBufferBuilder builder, long id) { builder.AddLong(0, id, 0); } public static void AddIndexType(FlatBufferBuilder builder, Offset indexTypeOffset) { builder.AddOffset(1, indexTypeOffset.Value, 0); } public static void AddIsOrdered(FlatBufferBuilder builder, bool isOrdered) { builder.AddBool(2, isOrdered, false); } public static void AddDictionaryKind(FlatBufferBuilder builder, DictionaryKind dictionaryKind) { builder.AddShort(3, (short)dictionaryKind, 0); } + public static void AddCustomMetadata(FlatBufferBuilder builder, VectorOffset customMetadataOffset) { builder.AddOffset(4, customMetadataOffset.Value, 0); } + public static VectorOffset CreateCustomMetadataVector(FlatBufferBuilder builder, Offset[] data) { builder.StartVector(4, data.Length, 4); for (int i = data.Length - 1; i >= 0; i--) builder.AddOffset(data[i].Value); return builder.EndVector(); } + public static VectorOffset CreateCustomMetadataVectorBlock(FlatBufferBuilder builder, Offset[] data) { builder.StartVector(4, data.Length, 4); builder.Add(data); return builder.EndVector(); } + public static VectorOffset CreateCustomMetadataVectorBlock(FlatBufferBuilder builder, ArraySegment> data) { builder.StartVector(4, data.Count, 4); builder.Add(data); return builder.EndVector(); } + public static VectorOffset CreateCustomMetadataVectorBlock(FlatBufferBuilder builder, IntPtr dataPtr, int sizeInBytes) { builder.StartVector(1, sizeInBytes, 1); builder.Add>(dataPtr, sizeInBytes); return builder.EndVector(); } + public static void StartCustomMetadataVector(FlatBufferBuilder builder, int numElems) { builder.StartVector(4, numElems, 4); } public static Offset EndDictionaryEncoding(FlatBufferBuilder builder) { int o = builder.EndTable(); return new Offset(o); @@ -70,6 +81,7 @@ static public bool Verify(Google.FlatBuffers.Verifier verifier, uint tablePos) && verifier.VerifyTable(tablePos, 6 /*IndexType*/, IntVerify.Verify, false) && verifier.VerifyField(tablePos, 8 /*IsOrdered*/, 1 /*bool*/, 1, false) && verifier.VerifyField(tablePos, 10 /*DictionaryKind*/, 2 /*DictionaryKind*/, 2, false) + && verifier.VerifyVectorOfTables(tablePos, 12 /*CustomMetadata*/, KeyValueVerify.Verify, false) && verifier.VerifyTableEnd(tablePos); } } diff --git a/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs index aa63985e..53b676d2 100644 --- a/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs +++ b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs @@ -1248,6 +1248,8 @@ private VectorOffset GetChildrenFieldOffset(Field field) targetDataType = extType.StorageType; if (targetDataType is DictionaryType dictionaryType) targetDataType = dictionaryType.ValueType; + if (targetDataType is ExtensionType valueExtType) + targetDataType = valueExtType.StorageType; if (!(targetDataType is NestedType type)) { @@ -1279,12 +1281,24 @@ private VectorOffset GetFieldMetadataOffset(Field field) { IReadOnlyDictionary metadata = field.Metadata; - // Inject extension type metadata if the field type is/wraps an ExtensionType - ExtensionType extType = field.DataType as ExtensionType; - if (extType == null && field.DataType is DictionaryType dt && dt.ValueType is ExtensionType dext) - extType = dext; + // A dictionary value type's extension metadata goes in DictionaryEncoding.custom_metadata; + // in Field.custom_metadata, it would mean the extension's storage is the whole dictionary. + // Field metadata that names the value type, which is how earlier versions of this + // library's IPC reader left it, is dropped rather than written there. + if (field.DataType is DictionaryType dt && + dt.ValueType is ExtensionType valueType && + metadata != null && + metadata.TryGetValue("ARROW:extension:name", out string extName) && + extName == valueType.Name) + { + var filtered = new Dictionary((IDictionary)metadata); + filtered.Remove("ARROW:extension:name"); + filtered.Remove("ARROW:extension:metadata"); + metadata = filtered; + } - if (extType != null) + // Inject extension type metadata if the field type is an ExtensionType + if (field.DataType is ExtensionType extType) { // Merge extension metadata with existing field metadata var merged = metadata != null @@ -1319,7 +1333,21 @@ private VectorOffset GetFieldMetadataOffset(Field field) var indexType = dicType.IndexType as NumberType; Offset indexOffset = Flatbuf.Int.CreateInt(Builder, indexType.BitWidth, indexType.IsSigned); - return Flatbuf.DictionaryEncoding.CreateDictionaryEncoding(Builder, id, indexOffset, dicType.Ordered); + + // DictionaryEncoding.custom_metadata (format 1.6) holds the value type's extension metadata + VectorOffset valueMetadataOffset = default; + if (dicType.ValueType is ExtensionType valueType) + { + Offset[] valueMetadataOffsets = GetMetadataOffsets(new Dictionary + { + ["ARROW:extension:name"] = valueType.Name, + ["ARROW:extension:metadata"] = valueType.ExtensionMetadata ?? "", + }); + valueMetadataOffset = Flatbuf.DictionaryEncoding.CreateCustomMetadataVector(Builder, valueMetadataOffsets); + } + + return Flatbuf.DictionaryEncoding.CreateDictionaryEncoding(Builder, id, indexOffset, dicType.Ordered, + Flatbuf.DictionaryKind.DenseArray, valueMetadataOffset); } /// diff --git a/src/Apache.Arrow/Ipc/DictionaryMemo.cs b/src/Apache.Arrow/Ipc/DictionaryMemo.cs index 0e151b69..59e2c112 100644 --- a/src/Apache.Arrow/Ipc/DictionaryMemo.cs +++ b/src/Apache.Arrow/Ipc/DictionaryMemo.cs @@ -61,12 +61,13 @@ public void AddField(long id, Field field) throw new ArgumentException($"Field {field.Name} is already in Memo"); } - if (field.DataType.TypeId != ArrowTypeId.Dictionary) + // The field may be an extension type whose storage is a dictionary + if (!(Types.ArrowTypeExtensions.GetStorageType(field.DataType) is DictionaryType dictionaryType)) { throw new ArgumentException($"Field type is not DictionaryType: Name={field.Name}, {field.DataType.Name}"); } - IArrowType valueType = ((DictionaryType)field.DataType).ValueType; + IArrowType valueType = dictionaryType.ValueType; if (_idToValueType.TryGetValue(id, out IArrowType valueTypeInDic)) { diff --git a/src/Apache.Arrow/Ipc/MessageSerializer.cs b/src/Apache.Arrow/Ipc/MessageSerializer.cs index 675f9468..082ba6b1 100644 --- a/src/Apache.Arrow/Ipc/MessageSerializer.cs +++ b/src/Apache.Arrow/Ipc/MessageSerializer.cs @@ -90,29 +90,20 @@ private static Field FieldFromFlatbuffer(Flatbuf.Field flatbufField, ref Diction Flatbuf.DictionaryEncoding? dictionaryEncoding = flatbufField.Dictionary; IArrowType type = GetFieldArrowType(flatbufField, childFields); - // Build metadata first (needed for extension type resolution) - Dictionary metadata = flatbufField.CustomMetadataLength > 0 ? new Dictionary() : null; - for (int i = 0; i < flatbufField.CustomMetadataLength; i++) - { - Flatbuf.KeyValue keyValue = flatbufField.CustomMetadata(i).GetValueOrDefault(); - metadata[keyValue.Key] = keyValue.Value; - } - - // Resolve extension type from metadata (before dictionary wrapping) - if (metadata != null && - metadata.TryGetValue("ARROW:extension:name", out string extName) && - extensionRegistry != null && - extensionRegistry.TryGetDefinition(extName, out ExtensionDefinition extDef)) + if (dictionaryEncoding.HasValue) { - metadata.TryGetValue("ARROW:extension:metadata", out string extMeta); - if (extDef.TryCreateType(type, extMeta, out ExtensionType extType)) + // DictionaryEncoding.custom_metadata (format 1.6) describes the dictionary value + // type. DictionaryType has nowhere to keep it, so, as in Arrow C++, only extension + // metadata is used and an unresolved value type's metadata is dropped. + Dictionary valueMetadata = null; + for (int i = 0; i < dictionaryEncoding.Value.CustomMetadataLength; i++) { - type = extType; + Flatbuf.KeyValue keyValue = dictionaryEncoding.Value.CustomMetadata(i).GetValueOrDefault(); + valueMetadata ??= new Dictionary(); + valueMetadata[keyValue.Key] = keyValue.Value; } - } + type = ResolveExtensionType(type, valueMetadata, extensionRegistry); - if (dictionaryEncoding.HasValue) - { Flatbuf.Int? indexTypeAsInt = dictionaryEncoding.Value.IndexType; IArrowType indexType = indexTypeAsInt.HasValue ? GetNumberType(indexTypeAsInt.Value.BitWidth, indexTypeAsInt.Value.IsSigned) : @@ -121,6 +112,17 @@ private static Field FieldFromFlatbuffer(Flatbuf.Field flatbufField, ref Diction type = new DictionaryType(indexType, type, dictionaryEncoding.Value.IsOrdered); } + Dictionary metadata = flatbufField.CustomMetadataLength > 0 ? new Dictionary() : null; + for (int i = 0; i < flatbufField.CustomMetadataLength; i++) + { + Flatbuf.KeyValue keyValue = flatbufField.CustomMetadata(i).GetValueOrDefault(); + metadata[keyValue.Key] = keyValue.Value; + } + + // Field.custom_metadata describes the field's whole type, so on a dictionary-encoded + // field its extension metadata names an extension type whose storage is the dictionary. + type = ResolveExtensionType(type, metadata, extensionRegistry); + var arrowField = new Field(flatbufField.Name, type, flatbufField.Nullable, metadata, copyCollections: false); if (dictionaryEncoding.HasValue) @@ -132,6 +134,23 @@ private static Field FieldFromFlatbuffer(Flatbuf.Field flatbufField, ref Diction return arrowField; } + private static IArrowType ResolveExtensionType(IArrowType type, Dictionary metadata, ExtensionTypeRegistry extensionRegistry) + { + if (metadata != null && + metadata.TryGetValue("ARROW:extension:name", out string extName) && + extensionRegistry != null && + extensionRegistry.TryGetDefinition(extName, out ExtensionDefinition extDef)) + { + metadata.TryGetValue("ARROW:extension:metadata", out string extMeta); + if (extDef.TryCreateType(type, extMeta, out ExtensionType extType)) + { + return extType; + } + } + + return type; + } + private static Types.IArrowType GetFieldArrowType(Flatbuf.Field field, Field[] childFields = null) { switch (field.TypeType) diff --git a/test/Apache.Arrow.Tests/ExtensionTypeTests.cs b/test/Apache.Arrow.Tests/ExtensionTypeTests.cs index d87823f6..41e7a577 100644 --- a/test/Apache.Arrow.Tests/ExtensionTypeTests.cs +++ b/test/Apache.Arrow.Tests/ExtensionTypeTests.cs @@ -566,7 +566,8 @@ public unsafe void CDataExportDictionaryOfExtensionPutsMetadataOnDictionary() [Fact] public unsafe void CDataExportMovesValueExtensionMetadataFromDictionaryField() { - // The IPC reader leaves a dictionary value type's extension metadata on the field. + // Earlier versions of the IPC reader left a dictionary value type's extension metadata + // on the field. var field = new Field("f", s_guidDictionaryType, true, new Dictionary { @@ -745,6 +746,229 @@ public unsafe void CDataDictionaryOfExtensionArrayRoundTrip() Assert.Equal(new int?[] { 1, null, 0, 1 }, (Int32Array)importedArray.Indices); } + [Fact] + public void IpcDictionaryOfExtensionWritesValueMetadataToDictionaryEncoding() + { + // Earlier versions of the IPC reader left a dictionary value type's extension metadata + // on the field. + var field = new Field("f", s_guidDictionaryType, true, + new Dictionary + { + ["ARROW:extension:name"] = "arrow.uuid", + ["ARROW:extension:metadata"] = "", + ["other"] = "value", + }); + + Flatbuf.Field written = WriteSchemaField(field); + + Assert.Equal(Flatbuf.Type.FixedSizeBinary, written.TypeType); + Assert.Equal(new Dictionary { ["other"] = "value" }, ReadFieldMetadata(written)); + Assert.Equal( + new Dictionary { ["ARROW:extension:name"] = "arrow.uuid", ["ARROW:extension:metadata"] = "" }, + ReadDictionaryMetadata(written)); + } + + [Fact] + public void IpcDictionaryOfExtensionRoundTrip() + { + var values = new Guid?[] { Guid.NewGuid(), Guid.NewGuid() }; + RecordBatch batch = BuildGuidDictionaryRecordBatch(values); + + RecordBatch readBatch = IpcRoundTrip(batch, GuidRegistry()); + + Field field = readBatch.Schema.GetFieldByIndex(0); + DictionaryType dictionaryType = Assert.IsType(field.DataType); + Assert.IsType(dictionaryType.IndexType); + Assert.IsType(dictionaryType.ValueType); + Assert.Null(field.Metadata); + + DictionaryArray readArray = Assert.IsType(readBatch.Column(0)); + Assert.Equal(values, Assert.IsType(readArray.Dictionary)); + Assert.Equal(new int?[] { 1, null, 0, 1 }, (Int32Array)readArray.Indices); + } + + [Fact] + public void IpcDictionaryOfExtensionRoundTripWithoutRegistry() + { + RecordBatch batch = BuildGuidDictionaryRecordBatch(new Guid?[] { Guid.NewGuid(), Guid.NewGuid() }); + + RecordBatch readBatch = IpcRoundTrip(batch, new ExtensionTypeRegistry()); + + // As in Arrow C++, an unresolved value type's metadata is dropped rather than moved + // to the field, where it would describe the whole dictionary. + Field field = readBatch.Schema.GetFieldByIndex(0); + DictionaryType dictionaryType = Assert.IsType(field.DataType); + Assert.IsType(dictionaryType.ValueType); + Assert.Null(field.Metadata); + Assert.IsType(((DictionaryArray)readBatch.Column(0)).Dictionary); + } + + [Fact] + public void IpcNestedDictionaryOfExtensionRoundTrip() + { + var values = new Guid?[] { Guid.NewGuid(), Guid.NewGuid() }; + DictionaryArray dictionaryArray = (DictionaryArray)BuildGuidDictionaryRecordBatch(values).Column(0); + var listType = new ListType(new Field("item", s_guidDictionaryType, true)); + ArrowBuffer offsets = new ArrowBuffer.Builder().Append(0).Append(1).Append(4).Build(); + var listArray = new ListArray(listType, 2, offsets, dictionaryArray, ArrowBuffer.Empty); + var schema = new Schema(new[] { new Field("f", listType, true) }, null); + var batch = new RecordBatch(schema, new[] { listArray }, listArray.Length); + + RecordBatch readBatch = IpcRoundTrip(batch, GuidRegistry()); + + ListType readListType = Assert.IsType(readBatch.Schema.GetFieldByIndex(0).DataType); + DictionaryType dictionaryType = Assert.IsType(readListType.ValueDataType); + Assert.IsType(dictionaryType.ValueType); + Assert.Null(readListType.ValueField.Metadata); + + ListArray readList = Assert.IsType(readBatch.Column(0)); + DictionaryArray readValues = Assert.IsType(readList.Values); + Assert.Equal(values, Assert.IsType(readValues.Dictionary)); + Assert.Equal(new int?[] { 1, null, 0, 1 }, (Int32Array)readValues.Indices); + } + + [Fact] + public void IpcDictionaryFieldExtensionMetadataDescribesWholeDictionary() + { + // Writers before format 1.6, including earlier versions of this library, put a + // dictionary value type's extension metadata on the field. Like Arrow C++, it is + // resolved against the whole dictionary type, not the values. + var metadata = new Dictionary + { + ["ARROW:extension:name"] = "arrow.uuid", + ["ARROW:extension:metadata"] = "", + }; + var dictionaryType = new DictionaryType(Int32Type.Default, new FixedSizeBinaryType(16), false); + var field = new Field("f", dictionaryType, true, metadata); + var storage = (FixedSizeBinaryArray)((GuidArray)BuildGuidRecordBatch(new Guid?[] { Guid.NewGuid() }).Column(0)).Storage; + Int32Array indices = new Int32Array.Builder().Append(0).AppendNull().Build(); + var array = new DictionaryArray(dictionaryType, indices, storage); + var batch = new RecordBatch(new Schema(new[] { field }, null), new[] { array }, array.Length); + + Flatbuf.Field written = WriteSchemaField(field); + Assert.Equal(metadata, ReadFieldMetadata(written)); + Assert.Empty(ReadDictionaryMetadata(written)); + + RecordBatch readBatch = IpcRoundTrip(batch, GuidRegistry()); + + Field readField = readBatch.Schema.GetFieldByIndex(0); + DictionaryType readType = Assert.IsType(readField.DataType); + Assert.IsType(readType.ValueType); + Assert.Equal(metadata, readField.Metadata); + } + + [Fact] + public void IpcDictionaryStorageExtensionRoundTrip() + { + var dictionaryType = new DictionaryType(Int32Type.Default, StringType.Default, false); + var extensionType = new DictionaryStorageType(dictionaryType); + var field = new Field("f", extensionType, true); + StringArray dictionary = new StringArray.Builder().Append("a").Append("b").Build(); + Int32Array indices = new Int32Array.Builder().Append(1).AppendNull().Append(0).Build(); + var array = extensionType.CreateArray(new DictionaryArray(dictionaryType, indices, dictionary)); + var batch = new RecordBatch(new Schema(new[] { field }, null), new[] { array }, array.Length); + + // An extension type whose storage is a dictionary keeps its metadata on the field + Flatbuf.Field written = WriteSchemaField(field); + Assert.Equal( + new Dictionary { ["ARROW:extension:name"] = DictionaryStorageType.ExtensionName, ["ARROW:extension:metadata"] = "" }, + ReadFieldMetadata(written)); + Assert.Empty(ReadDictionaryMetadata(written)); + + var registry = new ExtensionTypeRegistry(); + registry.Register(DictionaryStorageDefinition.Instance); + RecordBatch readBatch = IpcRoundTrip(batch, registry); + + DictionaryStorageType readType = Assert.IsType(readBatch.Schema.GetFieldByIndex(0).DataType); + Assert.IsType(Assert.IsType(readType.StorageType).ValueType); + ExtensionArray readArray = Assert.IsAssignableFrom(readBatch.Column(0)); + DictionaryArray readStorage = Assert.IsType(readArray.Storage); + StringArray readDictionary = Assert.IsType(readStorage.Dictionary); + Assert.Equal(new[] { "a", "b" }, new[] { readDictionary.GetString(0), readDictionary.GetString(1) }); + Assert.Equal(new int?[] { 1, null, 0 }, (Int32Array)readStorage.Indices); + } + + private static RecordBatch BuildGuidDictionaryRecordBatch(Guid?[] values) + { + GuidArray dictionary = (GuidArray)BuildGuidRecordBatch(values).Column(0); + Int32Array indices = new Int32Array.Builder().Append(1).AppendNull().Append(0).Append(1).Build(); + var array = new DictionaryArray(s_guidDictionaryType, indices, dictionary); + var schema = new Schema(new[] { new Field("f", s_guidDictionaryType, true) }, null); + return new RecordBatch(schema, new[] { array }, array.Length); + } + + private static RecordBatch IpcRoundTrip(RecordBatch batch, ExtensionTypeRegistry registry) + { + var stream = new MemoryStream(); + var writer = new ArrowStreamWriter(stream, batch.Schema); + writer.WriteRecordBatch(batch); + writer.WriteEnd(); + stream.Position = 0; + + var reader = new ArrowStreamReader(new ArrowContext(extensionRegistry: registry), stream); + RecordBatch readBatch = reader.ReadNextRecordBatch(); + Assert.NotNull(readBatch); + return readBatch; + } + + private static Flatbuf.Field WriteSchemaField(Field field) + { + var stream = new MemoryStream(); + var writer = new ArrowStreamWriter(stream, new Schema(new[] { field }, null)); + writer.WriteStart(); + + // Skip the continuation marker and the metadata length + var message = Flatbuf.Message.GetRootAsMessage(new Google.FlatBuffers.ByteBuffer(stream.ToArray(), 8)); + return message.Header().Value.Fields(0).Value; + } + + private static Dictionary ReadFieldMetadata(Flatbuf.Field field) + { + var result = new Dictionary(); + for (int i = 0; i < field.CustomMetadataLength; i++) + { + Flatbuf.KeyValue keyValue = field.CustomMetadata(i).Value; + result[keyValue.Key] = keyValue.Value; + } + return result; + } + + private static Dictionary ReadDictionaryMetadata(Flatbuf.Field field) + { + Flatbuf.DictionaryEncoding encoding = field.Dictionary.Value; + var result = new Dictionary(); + for (int i = 0; i < encoding.CustomMetadataLength; i++) + { + Flatbuf.KeyValue keyValue = encoding.CustomMetadata(i).Value; + result[keyValue.Key] = keyValue.Value; + } + return result; + } + + private sealed class DictionaryStorageType : ExtensionType + { + public const string ExtensionName = "test.dict_storage"; + + public DictionaryStorageType(IArrowType storageType) : base(storageType) { } + + public override string Name => ExtensionName; + public override string ExtensionMetadata => ""; + public override ExtensionArray CreateArray(IArrowArray storageArray) => new ExtensionArray(this, storageArray); + } + + private sealed class DictionaryStorageDefinition : ExtensionDefinition + { + public static readonly DictionaryStorageDefinition Instance = new DictionaryStorageDefinition(); + + public override string ExtensionName => DictionaryStorageType.ExtensionName; + + public override bool TryCreateType(IArrowType storageType, string metadata, out ExtensionType type) + { + type = storageType is DictionaryType ? new DictionaryStorageType(storageType) : null; + return type != null; + } + } + private sealed class SingleBatchStream : IArrowArrayStream { private RecordBatch _batch;