Skip to content

Support DictionaryEncoding.custom_metadata (format 1.6) for dictionaries of extension types #476

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
CurtHagenlocher wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Support DictionaryEncoding.custom_metadata (format 1.6) for dictionaries of extension types #476

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions format/Schema.fbs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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 ];
}

/// ----------------------------------------------------------------------
Expand Down
5 changes: 3 additions & 2 deletions src/Apache.Arrow/C/CArrowSchemaExporter.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<string, string> dictionaryMetadata = null;
if (datatype is DictionaryType dictionaryType &&
dictionaryType.ValueType is ExtensionType valueType &&
Expand Down
18 changes: 15 additions & 3 deletions src/Apache.Arrow/Flatbuf/DictionaryEncoding.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<DictionaryEncoding> CreateDictionaryEncoding(FlatBufferBuilder builder,
long id = 0,
Offset<Int> indexTypeOffset = default(Offset<Int>),
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<Int> 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<KeyValue>[] 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<KeyValue>[] data) { builder.StartVector(4, data.Length, 4); builder.Add(data); return builder.EndVector(); }
public static VectorOffset CreateCustomMetadataVectorBlock(FlatBufferBuilder builder, ArraySegment<Offset<KeyValue>> 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<Offset<KeyValue>>(dataPtr, sizeInBytes); return builder.EndVector(); }
public static void StartCustomMetadataVector(FlatBufferBuilder builder, int numElems) { builder.StartVector(4, numElems, 4); }
public static Offset<DictionaryEncoding> EndDictionaryEncoding(FlatBufferBuilder builder) {
int o = builder.EndTable();
return new Offset<DictionaryEncoding>(o);
Expand All @@ -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);
}
}
Expand Down
40 changes: 34 additions & 6 deletions src/Apache.Arrow/Ipc/ArrowStreamWriter.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
{
Expand Down Expand Up @@ -1279,12 +1281,24 @@ private VectorOffset GetFieldMetadataOffset(Field field)
{
IReadOnlyDictionary<string, string> 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<string, string>((IDictionary<string, string>)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
Expand Down Expand Up @@ -1319,7 +1333,21 @@ private VectorOffset GetFieldMetadataOffset(Field field)
var indexType = dicType.IndexType as NumberType;

Offset<Flatbuf.Int> 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<Flatbuf.KeyValue>[] valueMetadataOffsets = GetMetadataOffsets(new Dictionary<string, string>
{
["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);
}

/// <summary>
Expand Down
5 changes: 3 additions & 2 deletions src/Apache.Arrow/Ipc/DictionaryMemo.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
{
Expand Down
57 changes: 38 additions & 19 deletions src/Apache.Arrow/Ipc/MessageSerializer.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<string, string> metadata = flatbufField.CustomMetadataLength > 0 ? new Dictionary<string, string>() : 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<string, string> valueMetadata = null;
for (int i = 0; i < dictionaryEncoding.Value.CustomMetadataLength; i++)
{
type = extType;
Flatbuf.KeyValue keyValue = dictionaryEncoding.Value.CustomMetadata(i).GetValueOrDefault();
valueMetadata ??= new Dictionary<string, string>();
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) :
Expand All @@ -121,6 +112,17 @@ private static Field FieldFromFlatbuffer(Flatbuf.Field flatbufField, ref Diction
type = new DictionaryType(indexType, type, dictionaryEncoding.Value.IsOrdered);
}

Dictionary<string, string> metadata = flatbufField.CustomMetadataLength > 0 ? new Dictionary<string, string>() : 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)
Expand All @@ -132,6 +134,23 @@ private static Field FieldFromFlatbuffer(Flatbuf.Field flatbufField, ref Diction
return arrowField;
}

private static IArrowType ResolveExtensionType(IArrowType type, Dictionary<string, string> 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)
Expand Down
Loading
Loading