Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
101 changes: 74 additions & 27 deletions src/Apache.Arrow/C/CArrowSchemaExporter.cs
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,9 @@ namespace Apache.Arrow.C
{
public static class CArrowSchemaExporter
{
private const string ExtensionNameKey = "ARROW:extension:name";
private const string ExtensionMetadataKey = "ARROW:extension:metadata";

#if NET5_0_OR_GREATER
private static unsafe delegate* unmanaged<CArrowSchema*, void> ReleaseSchemaPtr => &ReleaseCArrowSchema;
#else
Expand All @@ -49,6 +52,11 @@ public static class CArrowSchemaExporter
/// </code>
/// </example>
public static unsafe void ExportType(IArrowType datatype, CArrowSchema* schema)
{
ExportType(datatype, schema, null);
}

private static unsafe void ExportType(IArrowType datatype, CArrowSchema* schema, IReadOnlyDictionary<string, string> metadata)
{
if (datatype == null)
{
Expand All @@ -59,17 +67,31 @@ public static unsafe void ExportType(IArrowType datatype, CArrowSchema* schema)
throw new ArgumentNullException(nameof(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.
IReadOnlyDictionary<string, string> dictionaryMetadata = null;
if (datatype is DictionaryType dictionaryType &&
dictionaryType.ValueType is ExtensionType valueType &&
metadata != null &&
metadata.TryGetValue(ExtensionNameKey, out string extName) &&
extName == valueType.Name)
{
metadata = SplitExtensionMetadata(metadata, out dictionaryMetadata);
}

schema->format = StringUtil.ToCStringUtf8(GetFormat(datatype));
schema->name = null;
schema->metadata = null;
schema->metadata = ConstructMetadata(WithExtensionMetadata(datatype, metadata));

IArrowType physicalType = datatype is ExtensionType ext ? ext.StorageType : datatype;
schema->flags = GetFlags(physicalType);

schema->children = ConstructChildren(physicalType, out var numChildren);
schema->n_children = numChildren;

schema->dictionary = ConstructDictionary(physicalType);
schema->dictionary = ConstructDictionary(physicalType, dictionaryMetadata);

schema->release = ReleaseSchemaPtr;

Expand All @@ -91,29 +113,9 @@ public static unsafe void ExportType(IArrowType datatype, CArrowSchema* schema)
/// </example>
public static unsafe void ExportField(Field field, CArrowSchema* schema)
{
ExportType(field.DataType, schema);
ExportType(field.DataType, schema, field.Metadata);
schema->name = StringUtil.ToCStringUtf8(field.Name);

// Inject extension metadata if the field type is an ExtensionType
IReadOnlyDictionary<string, string> metadata = field.Metadata;
ExtensionType extType = field.DataType as ExtensionType;
if (extType == null && field.DataType is DictionaryType dt && dt.ValueType is ExtensionType dext)
extType = dext;

if (extType != null)
{
var merged = metadata != null
? new Dictionary<string, string>((IDictionary<string, string>)metadata)
: new Dictionary<string, string>();
if (!merged.ContainsKey("ARROW:extension:name"))
merged["ARROW:extension:name"] = extType.Name;
if (!merged.ContainsKey("ARROW:extension:metadata"))
merged["ARROW:extension:metadata"] = extType.ExtensionMetadata ?? "";
metadata = merged;
}

schema->metadata = ConstructMetadata(metadata);

IArrowType physicalType = field.DataType is ExtensionType ext ? ext.StorageType : field.DataType;
schema->flags = GetFlags(physicalType, field.IsNullable);
}
Expand All @@ -134,8 +136,7 @@ public static unsafe void ExportField(Field field, CArrowSchema* schema)
public static unsafe void ExportSchema(Schema schema, CArrowSchema* out_schema)
{
var structType = new StructType(schema.FieldsList);
ExportType(structType, out_schema);
out_schema->metadata = ConstructMetadata(schema.Metadata);
ExportType(structType, out_schema, schema.Metadata);
}

private static char FormatTimeUnit(TimeUnit unit) => unit switch
Expand Down Expand Up @@ -301,12 +302,12 @@ private static long GetFlags(IArrowType datatype, bool nullable = true)
}
}

private static unsafe CArrowSchema* ConstructDictionary(IArrowType datatype)
private static unsafe CArrowSchema* ConstructDictionary(IArrowType datatype, IReadOnlyDictionary<string, string> metadata)
{
if (datatype is DictionaryType dictType)
{
CArrowSchema* cSchema = CArrowSchema.Create();
ExportType(dictType.ValueType, cSchema);
ExportType(dictType.ValueType, cSchema, metadata);
return cSchema;
}
else
Expand All @@ -315,6 +316,52 @@ private static long GetFlags(IArrowType datatype, bool nullable = true)
}
}

/// <summary>
/// Adds the extension name and metadata of <paramref name="datatype"/>, if it is an
/// <see cref="ExtensionType"/>, to <paramref name="metadata"/>. Existing keys are kept.
/// </summary>
private static IReadOnlyDictionary<string, string> WithExtensionMetadata(IArrowType datatype, IReadOnlyDictionary<string, string> metadata)
{
if (!(datatype is ExtensionType extType))
{
return metadata;
}

var merged = metadata != null
? new Dictionary<string, string>((IDictionary<string, string>)metadata)
: new Dictionary<string, string>();
if (!merged.ContainsKey(ExtensionNameKey))
merged[ExtensionNameKey] = extType.Name;
if (!merged.ContainsKey(ExtensionMetadataKey))
merged[ExtensionMetadataKey] = extType.ExtensionMetadata ?? "";
return merged;
}

/// <summary>
/// Splits the extension name and metadata keys out of <paramref name="metadata"/>.
/// </summary>
/// <returns>The remaining metadata, or null if none remains.</returns>
private static IReadOnlyDictionary<string, string> SplitExtensionMetadata(
IReadOnlyDictionary<string, string> metadata, out IReadOnlyDictionary<string, string> extensionMetadata)
{
extensionMetadata = null;
if (metadata == null ||
!(metadata.ContainsKey(ExtensionNameKey) || metadata.ContainsKey(ExtensionMetadataKey)))
{
return metadata;
}

var remaining = new Dictionary<string, string>();
var extension = new Dictionary<string, string>();
foreach (KeyValuePair<string, string> pair in metadata)
{
bool isExtensionKey = pair.Key == ExtensionNameKey || pair.Key == ExtensionMetadataKey;
(isExtensionKey ? extension : remaining)[pair.Key] = pair.Value;
}
extensionMetadata = extension;
return remaining.Count > 0 ? remaining : null;
}

private unsafe static byte* ConstructMetadata(IReadOnlyDictionary<string, string> metadata)
{
if (metadata == null || metadata.Count == 0)
Expand Down
46 changes: 33 additions & 13 deletions src/Apache.Arrow/C/CArrowSchemaImporter.cs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@ namespace Apache.Arrow.C
{
public static class CArrowSchemaImporter
{
private const string ExtensionNameKey = "ARROW:extension:name";
private const string ExtensionMetadataKey = "ARROW:extension:metadata";

/// <summary>
/// Import C pointer as an <see cref="ArrowType"/>.
/// </summary>
Expand Down Expand Up @@ -186,7 +189,10 @@ public ArrowType GetAsType(ExtensionTypeRegistry extensionRegistry)
};

var dictionarySchema = new ImportedArrowSchema(_cSchema->dictionary, isRoot: false);
ArrowType dictionaryType = dictionarySchema.GetAsType(extensionRegistry);
IArrowType dictionaryType = ResolveExtensionType(
dictionarySchema.GetAsType(extensionRegistry),
GetMetadata(_cSchema->dictionary->metadata),
extensionRegistry);

bool ordered = _cSchema->GetFlag(CArrowSchema.ArrowFlagDictionaryOrdered);

Expand Down Expand Up @@ -384,18 +390,10 @@ public Field GetAsField(ExtensionTypeRegistry extensionRegistry)
IArrowType type = GetAsType(extensionRegistry);
IReadOnlyDictionary<string, string> metadata = GetMetadata(_cSchema->metadata);

// Resolve extension type from metadata
if (metadata != null &&
extensionRegistry != null &&
metadata.TryGetValue("ARROW:extension:name", out string extName) &&
extensionRegistry.TryGetDefinition(extName, out ExtensionDefinition extDef))
{
metadata.TryGetValue("ARROW:extension:metadata", out string extMeta);
if (extDef.TryCreateType(type, extMeta, out ExtensionType extType))
{
type = extType;
}
}
// On a dictionary field, extension metadata means the extension's storage is the
// whole dictionary. Extension metadata on the dictionary child isn't copied to the
// field for that reason; like Arrow C++, an unresolved value type's is dropped.
type = ResolveExtensionType(type, metadata, extensionRegistry);

return new Field(fieldName, type, nullable, metadata);
}
Expand All @@ -418,6 +416,28 @@ public Schema GetAsSchema(ExtensionTypeRegistry extensionRegistry)
}
}

/// <summary>
/// Returns the extension type named in <paramref name="metadata"/> if
/// <paramref name="extensionRegistry"/> defines it for <paramref name="storageType"/>;
/// otherwise returns <paramref name="storageType"/>.
/// </summary>
private static IArrowType ResolveExtensionType(
IArrowType storageType, IReadOnlyDictionary<string, string> metadata, ExtensionTypeRegistry extensionRegistry)
{
if (metadata != null &&
extensionRegistry != null &&
metadata.TryGetValue(ExtensionNameKey, out string extName) &&
extensionRegistry.TryGetDefinition(extName, out ExtensionDefinition extDef))
{
metadata.TryGetValue(ExtensionMetadataKey, out string extMeta);
if (extDef.TryCreateType(storageType, extMeta, out ExtensionType extType))
{
return extType;
}
}
return storageType;
}

private List<Field> ParseChildren(string typeName, ExtensionTypeRegistry extensionRegistry)
{
var child_schemas = new ImportedArrowSchema[_cSchema->n_children];
Expand Down
95 changes: 95 additions & 0 deletions test/Apache.Arrow.Tests/CDataInterfacePythonTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -1225,6 +1225,101 @@ public unsafe void ImportArrayStreamWithExtensionTypes()
Assert.Equal(new Guid?[] { guid2 }, Assert.IsType<GuidArray>(uuidLists.GetSlicedValues(1)));
}

[SkippableFact]
public unsafe void ExportDictionaryOfGuidArray()
{
// Export a dictionary-encoded GuidArray and verify that pyarrow sees a
// dictionary whose value type is arrow.uuid.

var guid1 = Guid.Parse("01234567-89ab-cdef-0123-456789abcdef");
var guid2 = Guid.Parse("fedcba98-7654-3210-fedc-ba9876543210");
GuidArray dictionary = new GuidArray.Builder().Append(guid1).Append(guid2).Build();
Int32Array indices = new Int32Array.Builder().Append(1).AppendNull().Append(0).Build();
var dictionaryType = new DictionaryType(Int32Type.Default, GuidType.Default, false);
var array = new DictionaryArray(dictionaryType, indices, dictionary);
var schema = new Schema(new[] { new Field("uuids", dictionaryType, true) }, null);
var batch = new RecordBatch(schema, new[] { array }, array.Length);

CArrowArray* cArray = CArrowArray.Create();
CArrowArrayExporter.ExportRecordBatch(batch, cArray);
CArrowSchema* cSchema = CArrowSchema.Create();
CArrowSchemaExporter.ExportSchema(batch.Schema, cSchema);

long arrayPtr = ((IntPtr)cArray).ToInt64();
long schemaPtr = ((IntPtr)cSchema).ToInt64();

using (Py.GIL())
{
dynamic pa = Py.Import("pyarrow");
dynamic uuid_mod = Py.Import("uuid");

dynamic pyBatch = pa.RecordBatch._import_from_c(arrayPtr, schemaPtr);
dynamic pyArray = pyBatch.column(0);

Assert.True((bool)pyArray.type.equals(pa.dictionary(pa.int32(), pa.uuid())));

dynamic expectedArray = pa.DictionaryArray.from_arrays(
pa.array(new PyList(new PyObject[] { new PyInt(1), PyObject.None, new PyInt(0) }), type: pa.int32()),
pa.array(
new PyList(new PyObject[]
{
(PyObject)uuid_mod.UUID(guid1.ToString()).bytes,
(PyObject)uuid_mod.UUID(guid2.ToString()).bytes,
}),
type: pa.uuid()));

Assert.True((bool)pyArray.equals(expectedArray));
}

CArrowArray.Free(cArray);
CArrowSchema.Free(cSchema);
}

[SkippableFact]
public unsafe void ImportDictionaryOfGuidArray()
{
// Import a pyarrow dictionary<int32, arrow.uuid> array and verify that the
// dictionary values resolve to a GuidArray.

var guid1 = Guid.Parse("01234567-89ab-cdef-0123-456789abcdef");
var guid2 = Guid.Parse("fedcba98-7654-3210-fedc-ba9876543210");

CArrowArray* cArray = CArrowArray.Create();
CArrowSchema* cSchema = CArrowSchema.Create();

using (Py.GIL())
{
dynamic pa = Py.Import("pyarrow");
dynamic uuid_mod = Py.Import("uuid");

dynamic pyArray = pa.DictionaryArray.from_arrays(
pa.array(new[] { 1, 0, 1 }, type: pa.int32()),
pa.array(
new PyList(new PyObject[]
{
(PyObject)uuid_mod.UUID(guid1.ToString()).bytes,
(PyObject)uuid_mod.UUID(guid2.ToString()).bytes,
}),
type: pa.uuid()));

pyArray._export_to_c(((IntPtr)cArray).ToInt64(), ((IntPtr)cSchema).ToInt64());
}

var registry = new ExtensionTypeRegistry();
registry.Register(GuidExtensionDefinition.Instance);

Field field = CArrowSchemaImporter.ImportField(cSchema, registry);
CArrowSchema.Free(cSchema);
DictionaryType dictionaryType = Assert.IsType<DictionaryType>(field.DataType);
Assert.IsType<GuidType>(dictionaryType.ValueType);

var importedArray = (DictionaryArray)CArrowArrayImporter.ImportArray(cArray, field.DataType);
Assert.Equal(new Guid?[] { guid1, guid2 }, Assert.IsType<GuidArray>(importedArray.Dictionary));
Assert.Equal(new int?[] { 1, 0, 1 }, (Int32Array)importedArray.Indices);

CArrowArray.Free(cArray);
}

[SkippableFact]
public unsafe void ExportBool8Array()
{
Expand Down
Loading
Loading