diff --git a/src/Apache.Arrow/C/CArrowSchemaExporter.cs b/src/Apache.Arrow/C/CArrowSchemaExporter.cs index acc44e5..da875c0 100644 --- a/src/Apache.Arrow/C/CArrowSchemaExporter.cs +++ b/src/Apache.Arrow/C/CArrowSchemaExporter.cs @@ -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 ReleaseSchemaPtr => &ReleaseCArrowSchema; #else @@ -49,6 +52,11 @@ public static class CArrowSchemaExporter /// /// public static unsafe void ExportType(IArrowType datatype, CArrowSchema* schema) + { + ExportType(datatype, schema, null); + } + + private static unsafe void ExportType(IArrowType datatype, CArrowSchema* schema, IReadOnlyDictionary metadata) { if (datatype == null) { @@ -59,9 +67,23 @@ 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 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); @@ -69,7 +91,7 @@ public static unsafe void ExportType(IArrowType datatype, CArrowSchema* schema) schema->children = ConstructChildren(physicalType, out var numChildren); schema->n_children = numChildren; - schema->dictionary = ConstructDictionary(physicalType); + schema->dictionary = ConstructDictionary(physicalType, dictionaryMetadata); schema->release = ReleaseSchemaPtr; @@ -91,29 +113,9 @@ public static unsafe void ExportType(IArrowType datatype, CArrowSchema* schema) /// 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 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((IDictionary)metadata) - : new Dictionary(); - 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); } @@ -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 @@ -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 metadata) { if (datatype is DictionaryType dictType) { CArrowSchema* cSchema = CArrowSchema.Create(); - ExportType(dictType.ValueType, cSchema); + ExportType(dictType.ValueType, cSchema, metadata); return cSchema; } else @@ -315,6 +316,52 @@ private static long GetFlags(IArrowType datatype, bool nullable = true) } } + /// + /// Adds the extension name and metadata of , if it is an + /// , to . Existing keys are kept. + /// + private static IReadOnlyDictionary WithExtensionMetadata(IArrowType datatype, IReadOnlyDictionary metadata) + { + if (!(datatype is ExtensionType extType)) + { + return metadata; + } + + var merged = metadata != null + ? new Dictionary((IDictionary)metadata) + : new Dictionary(); + if (!merged.ContainsKey(ExtensionNameKey)) + merged[ExtensionNameKey] = extType.Name; + if (!merged.ContainsKey(ExtensionMetadataKey)) + merged[ExtensionMetadataKey] = extType.ExtensionMetadata ?? ""; + return merged; + } + + /// + /// Splits the extension name and metadata keys out of . + /// + /// The remaining metadata, or null if none remains. + private static IReadOnlyDictionary SplitExtensionMetadata( + IReadOnlyDictionary metadata, out IReadOnlyDictionary extensionMetadata) + { + extensionMetadata = null; + if (metadata == null || + !(metadata.ContainsKey(ExtensionNameKey) || metadata.ContainsKey(ExtensionMetadataKey))) + { + return metadata; + } + + var remaining = new Dictionary(); + var extension = new Dictionary(); + foreach (KeyValuePair 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 metadata) { if (metadata == null || metadata.Count == 0) diff --git a/src/Apache.Arrow/C/CArrowSchemaImporter.cs b/src/Apache.Arrow/C/CArrowSchemaImporter.cs index 4bf7f44..c92c0fa 100644 --- a/src/Apache.Arrow/C/CArrowSchemaImporter.cs +++ b/src/Apache.Arrow/C/CArrowSchemaImporter.cs @@ -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"; + /// /// Import C pointer as an . /// @@ -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); @@ -384,18 +390,10 @@ public Field GetAsField(ExtensionTypeRegistry extensionRegistry) IArrowType type = GetAsType(extensionRegistry); IReadOnlyDictionary 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); } @@ -418,6 +416,28 @@ public Schema GetAsSchema(ExtensionTypeRegistry extensionRegistry) } } + /// + /// Returns the extension type named in if + /// defines it for ; + /// otherwise returns . + /// + private static IArrowType ResolveExtensionType( + IArrowType storageType, IReadOnlyDictionary 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 ParseChildren(string typeName, ExtensionTypeRegistry extensionRegistry) { var child_schemas = new ImportedArrowSchema[_cSchema->n_children]; diff --git a/test/Apache.Arrow.Tests/CDataInterfacePythonTests.cs b/test/Apache.Arrow.Tests/CDataInterfacePythonTests.cs index ee5134a..48db2b0 100644 --- a/test/Apache.Arrow.Tests/CDataInterfacePythonTests.cs +++ b/test/Apache.Arrow.Tests/CDataInterfacePythonTests.cs @@ -1225,6 +1225,101 @@ public unsafe void ImportArrayStreamWithExtensionTypes() Assert.Equal(new Guid?[] { guid2 }, Assert.IsType(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 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(field.DataType); + Assert.IsType(dictionaryType.ValueType); + + var importedArray = (DictionaryArray)CArrowArrayImporter.ImportArray(cArray, field.DataType); + Assert.Equal(new Guid?[] { guid1, guid2 }, Assert.IsType(importedArray.Dictionary)); + Assert.Equal(new int?[] { 1, 0, 1 }, (Int32Array)importedArray.Indices); + + CArrowArray.Free(cArray); + } + [SkippableFact] public unsafe void ExportBool8Array() { diff --git a/test/Apache.Arrow.Tests/ExtensionTypeTests.cs b/test/Apache.Arrow.Tests/ExtensionTypeTests.cs index 308ca26..d87823f 100644 --- a/test/Apache.Arrow.Tests/ExtensionTypeTests.cs +++ b/test/Apache.Arrow.Tests/ExtensionTypeTests.cs @@ -492,6 +492,259 @@ public async Task CDataArrayStreamResolvesListOfExtensionType() } } + private static readonly DictionaryType s_guidDictionaryType = new DictionaryType(Int32Type.Default, GuidType.Default, false); + + private static ExtensionTypeRegistry GuidRegistry() + { + var registry = new ExtensionTypeRegistry(); + registry.Register(GuidExtensionDefinition.Instance); + return registry; + } + + private static unsafe Dictionary ReadMetadata(byte* metadata) + { + var result = new Dictionary(); + if (metadata == null) + { + return result; + } + + int count = *(int*)metadata; + byte* ptr = metadata + 4; + for (int i = 0; i < count; i++) + { + int keyLength = *(int*)ptr; + string key = System.Text.Encoding.UTF8.GetString(ptr + 4, keyLength); + ptr += 4 + keyLength; + int valueLength = *(int*)ptr; + string value = System.Text.Encoding.UTF8.GetString(ptr + 4, valueLength); + ptr += 4 + valueLength; + result[key] = value; + } + return result; + } + + [Fact] + public unsafe void CDataExportTypeIncludesExtensionMetadata() + { + CArrowSchema* cSchema = CArrowSchema.Create(); + try + { + CArrowSchemaExporter.ExportType(GuidType.Default, cSchema); + + Dictionary metadata = ReadMetadata(cSchema->metadata); + Assert.Equal("arrow.uuid", metadata["ARROW:extension:name"]); + Assert.Equal("", metadata["ARROW:extension:metadata"]); + } + finally + { + CArrowSchema.Free(cSchema); + } + } + + [Fact] + public unsafe void CDataExportDictionaryOfExtensionPutsMetadataOnDictionary() + { + var field = new Field("f", s_guidDictionaryType, true, + new Dictionary { ["other"] = "value" }); + + CArrowSchema* cSchema = CArrowSchema.Create(); + try + { + CArrowSchemaExporter.ExportField(field, cSchema); + + Assert.Equal(new Dictionary { ["other"] = "value" }, ReadMetadata(cSchema->metadata)); + Dictionary dictionaryMetadata = ReadMetadata(cSchema->dictionary->metadata); + Assert.Equal("arrow.uuid", dictionaryMetadata["ARROW:extension:name"]); + } + finally + { + CArrowSchema.Free(cSchema); + } + } + + [Fact] + public unsafe void CDataExportMovesValueExtensionMetadataFromDictionaryField() + { + // The IPC reader leaves 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", + }); + + CArrowSchema* cSchema = CArrowSchema.Create(); + try + { + CArrowSchemaExporter.ExportField(field, cSchema); + + Assert.Equal(new Dictionary { ["other"] = "value" }, ReadMetadata(cSchema->metadata)); + Assert.Equal( + new Dictionary { ["ARROW:extension:name"] = "arrow.uuid", ["ARROW:extension:metadata"] = "" }, + ReadMetadata(cSchema->dictionary->metadata)); + } + finally + { + CArrowSchema.Free(cSchema); + } + } + + [Fact] + public unsafe void CDataExportKeepsUnresolvedDictionaryFieldExtensionMetadata() + { + // Without a resolved value type, extension metadata on a dictionary field means + // the extension's storage is the whole dictionary, so it must stay on the field. + var metadata = new Dictionary + { + ["ARROW:extension:name"] = "test.dict_storage", + ["ARROW:extension:metadata"] = "", + }; + var field = new Field("f", new DictionaryType(Int32Type.Default, StringType.Default, false), true, metadata); + + CArrowSchema* cSchema = CArrowSchema.Create(); + try + { + CArrowSchemaExporter.ExportField(field, cSchema); + + Assert.Equal(metadata, ReadMetadata(cSchema->metadata)); + Assert.Empty(ReadMetadata(cSchema->dictionary->metadata)); + } + finally + { + CArrowSchema.Free(cSchema); + } + } + + [Fact] + public unsafe void CDataDictionaryStorageExtensionPassThroughWithoutRegistry() + { + // Arrow C++ exports an extension type whose storage is a dictionary with its + // metadata on the field. Importing it without the definition and exporting it + // again must not move that metadata to the dictionary values. + var metadata = new Dictionary + { + ["ARROW:extension:name"] = "test.dict_storage", + ["ARROW:extension:metadata"] = "", + }; + CArrowSchema* cSchema = CArrowSchema.Create(); + CArrowSchema* cMetadataSchema = CArrowSchema.Create(); + CArrowSchemaExporter.ExportField(new Field("f", new DictionaryType(Int32Type.Default, StringType.Default, false), true), cSchema); + CArrowSchemaExporter.ExportField(new Field("m", Int32Type.Default, true, metadata), cMetadataSchema); + byte* swapped = cSchema->metadata; + cSchema->metadata = cMetadataSchema->metadata; + cMetadataSchema->metadata = swapped; + CArrowSchema.Free(cMetadataSchema); + + Field imported = CArrowSchemaImporter.ImportField(cSchema, GuidRegistry()); + CArrowSchema.Free(cSchema); + Assert.Equal(metadata, imported.Metadata); + + CArrowSchema* cReexported = CArrowSchema.Create(); + try + { + CArrowSchemaExporter.ExportField(imported, cReexported); + + Assert.Equal(metadata, ReadMetadata(cReexported->metadata)); + Assert.Empty(ReadMetadata(cReexported->dictionary->metadata)); + } + finally + { + CArrowSchema.Free(cReexported); + } + } + + [Fact] + public unsafe void CDataDictionaryOfExtensionRoundTrip() + { + var field = new Field("f", s_guidDictionaryType, true); + + CArrowSchema* cSchema = CArrowSchema.Create(); + CArrowSchemaExporter.ExportField(field, cSchema); + Field imported = CArrowSchemaImporter.ImportField(cSchema, GuidRegistry()); + CArrowSchema.Free(cSchema); + + DictionaryType dictionaryType = Assert.IsType(imported.DataType); + Assert.IsType(dictionaryType.IndexType); + Assert.IsType(dictionaryType.ValueType); + Assert.Null(imported.Metadata); + } + + [Fact] + public unsafe void CDataDictionaryOfExtensionRoundTripWithoutRegistry() + { + var field = new Field("f", s_guidDictionaryType, true); + + CArrowSchema* cSchema = CArrowSchema.Create(); + CArrowSchemaExporter.ExportField(field, cSchema); + Field imported = CArrowSchemaImporter.ImportField(cSchema); + CArrowSchema.Free(cSchema); + + // Like Arrow C++, the unresolved value type's extension metadata is dropped rather + // than moved to the field, where it would describe the whole dictionary. + DictionaryType dictionaryType = Assert.IsType(imported.DataType); + Assert.IsType(dictionaryType.ValueType); + Assert.Null(imported.Metadata); + } + + [Fact] + public unsafe void CDataImportDoesNotApplyDictionaryFieldExtensionMetadataToValues() + { + // On a dictionary field, extension metadata describes the whole dictionary, so as in + // Arrow C++ a field-level arrow.uuid doesn't resolve. (Older versions of this library + // exported a value type's extension metadata this way.) Build such a schema by + // swapping in the metadata of another export. + var storageField = new Field("f", new DictionaryType(Int32Type.Default, new FixedSizeBinaryType(16), false), true); + var metadataField = new Field("m", Int32Type.Default, true, + new Dictionary + { + ["ARROW:extension:name"] = "arrow.uuid", + ["ARROW:extension:metadata"] = "", + }); + + CArrowSchema* cSchema = CArrowSchema.Create(); + CArrowSchema* cMetadataSchema = CArrowSchema.Create(); + CArrowSchemaExporter.ExportField(storageField, cSchema); + CArrowSchemaExporter.ExportField(metadataField, cMetadataSchema); + byte* metadata = cSchema->metadata; + cSchema->metadata = cMetadataSchema->metadata; + cMetadataSchema->metadata = metadata; + CArrowSchema.Free(cMetadataSchema); + + Field imported = CArrowSchemaImporter.ImportField(cSchema, GuidRegistry()); + CArrowSchema.Free(cSchema); + + DictionaryType dictionaryType = Assert.IsType(imported.DataType); + Assert.IsType(dictionaryType.ValueType); + Assert.Equal("arrow.uuid", imported.Metadata["ARROW:extension:name"]); + } + + [Fact] + public unsafe void CDataDictionaryOfExtensionArrayRoundTrip() + { + var values = new Guid?[] { Guid.NewGuid(), Guid.NewGuid() }; + 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); + var batch = new RecordBatch(schema, new[] { array }, array.Length); + + CArrowArray* cArray = CArrowArray.Create(); + CArrowSchema* cSchema = CArrowSchema.Create(); + CArrowArrayExporter.ExportRecordBatch(batch, cArray); + CArrowSchemaExporter.ExportSchema(schema, cSchema); + + Schema importedSchema = CArrowSchemaImporter.ImportSchema(cSchema, GuidRegistry()); + using RecordBatch importedBatch = CArrowArrayImporter.ImportRecordBatch(cArray, importedSchema); + CArrowSchema.Free(cSchema); + CArrowArray.Free(cArray); + + DictionaryArray importedArray = Assert.IsType(importedBatch.Column(0)); + Assert.Equal(values, Assert.IsType(importedArray.Dictionary)); + Assert.Equal(new int?[] { 1, null, 0, 1 }, (Int32Array)importedArray.Indices); + } + private sealed class SingleBatchStream : IArrowArrayStream { private RecordBatch _batch;