unlock unlimited code runs for $79, limited time →
NOSSLETTER.techLIVE
back to OSSpedia
apache/arrow
mid8/10
testing-strategyapi-designbackwards-compatibility

GH-50993: [CI][Integration] Add extension-wrapped union to integration data

Project: apache/arrow | Language: C++ | Concepts: testing-strategy, api-design, backwards-compatibility | Difficulty: mid
Link: https://github.com/apache/arrow/pull/51027

The problem

The Arrow integration test suite did not include a case where an extension type wraps a union type (dense or sparse). This scenario had previously exposed bugs in C++ (GH-50623), but was not exercised across all implementations in the cross-language integration tests. Without this, implementations could miss handling extension types whose storage is a union.

Why it's non-trivial

A naive approach might simply add test data with extension-wrapped unions, but:

  • Extension types must be registered in each test harness, or the tests are vacuously passing (i.e., not actually exercising the extension logic).
  • The logic for buffer layout (such as validity bitmaps) must use the storage type, not the logical extension type, or serialization/deserialization will break for unions.
  • The integration data generator and consumers must handle both dense and sparse unions correctly, including for zero-length batches.

The approach

1. Add integration data generation for extension-wrapped unions

A new test case is added to the integration data generator:

def generate_extension_wrapped_union_case():
    # Unions wrapped in an extension type, exercising the extension/union
    # interaction across implementations (see the C++ fix in GH-50623).
    sparse_union_type = ExtensionType(
        'sparse-union-extension', 'sparse-union-extension',
        SparseUnionField('', [get_field('floats', 'float64'),
                              get_field('strings', 'largeutf8')],
                         type_ids=[0, 1]))
    dense_union_type = ExtensionType(
        'dense-union-extension', 'dense-union-extension',
        DenseUnionField('', [get_field('floats', 'float64'),
                             get_field('strings', 'largeutf8')],
                        type_ids=[0, 1]))

    fields = [
        ExtensionField('sparse_union_ext', sparse_union_type),
        ExtensionField('dense_union_ext', dense_union_type),
    ]

    batch_sizes = [0, 7]
    return _generate_file("extension_union", fields, batch_sizes)

This defines two columns: one with a sparse union wrapped in an extension, one with a dense union wrapped in an extension, and generates batches of size 0 and 7.

2. Register extension types in all C++ integration test harnesses

Extension types must be registered so the test is meaningful:

// In multiple places, e.g. json_integration_test.cc, c_data_integration_internal.cc, flight/integration_tests/test_integration_client.cc
-  ExtensionTypeGuard ext_guard({uuid(), dict_extension_type()});
+  ExtensionTypeGuard ext_guard({uuid(), dict_extension_type(),
+                                dense_union_extension_type(),
+                                sparse_union_extension_type()});

This ensures the new extension types are registered in all relevant test harnesses, so the extension logic is exercised.

3. Fix buffer layout logic to use storage type, not logical type

When exporting arrays via the Arrow C Data interface, buffer layout decisions must use the storage type:

-    if (n_buffers > 0 && !internal::may_have_validity_bitmap(data->type->id())) {
+    if (n_buffers > 0 && !internal::may_have_validity_bitmap(physical_type_id)) {

Here, physical_type_id is obtained via data->type->storage_id(), ensuring the buffer layout matches the storage type, not the extension wrapper.

4. Add test helpers for extension-wrapped unions

Helpers are added to create example arrays for dense and sparse union extensions:

std::shared_ptr<Array> ExampleDenseUnionExtension() {
  auto type = dense_union_extension_type();
  auto storage_type = checked_cast<const ExtensionType&>(*type).storage_type();
  return ExtensionType::WrapArray(
      type, ArrayFromJSON(storage_type, R"([[0, 1.5], [1, "abc"]])"));
}

This makes it easy to construct test arrays of the new types for use in unit and integration tests.

5. Update tests to cover the new extension-wrapped unions

The extension array export test now covers the new cases:

TEST_F(TestArrayExport, Extension) {
  TestPrimitive(ExampleUuid);
  TestPrimitive(ExampleSmallint);
  TestPrimitive(ExampleComplex128);
+  TestPrimitive(ExampleDenseUnionExtension);
+  TestPrimitive(ExampleSparseUnionExtension);
}

This ensures the C Data export/import logic is exercised for the new types.

6. Correct validity bitmap handling in JSON integration reader

When reading JSON integration data, validity bitmaps must be determined from the storage type:

-    if (::arrow::internal::may_have_validity_bitmap(type_->id())) {
+    if (::arrow::internal::may_have_validity_bitmap(type_->storage_id())) {
       // Null and union types don't have a validity bitmap
       RETURN_NOT_OK(ParseValidityBitmap());
     }

This prevents requiring a validity bitmap for union-backed extensions, matching the IPC writer.

From the review thread

pitrou: I've tested this without the fix from https://github.com/apache/arrow/commit/3e6f730bd0d04d972d60d4efee12b6be9c53cf25, and it still succeed.

I think this is due to the fact that the extension type is not registered, therefore there is no wrapping within ExtensionType that would exhibit the issue.

@Alb3e3 Can we remove ARROW:integration:allow_unregistered_extension?

Alb3e3: Fixed in de746b3cb.

I removed ARROW:integration:allow_unregistered_extension and registered the existing dense_union_extension_type() / sparse_union_extension_type() test types in the C++ integration executable. The generated storage schemas and serialized metadata now match those registered types.

Removing the bypass also exposed a related integration JSON reader issue: validity parsing was selected from the logical EXTENSION type, so it incorrectly required VALIDITY for union storage. It now selects validity layout from the physical storage type, matching the IPC writer's behavior.

pitrou: I'm just noticing that we can simply call DataType::storage_id here:

  const auto physical_type_id = data->type->storage_id();

and then:

  if (n_buffers > 0 && !internal::may_have_validity_bitmap(physical_type_id)) {

Takeaways

  • Integration tests must register all custom extension types to ensure the test actually exercises extension handling logic. Otherwise, bugs can be masked by fallback or bypass behavior.
  • Buffer layout and validity bitmap logic should be determined by the physical storage type, not the logical extension type. This is especially important for extension types wrapping unions or other complex types.
  • Adding focused helpers and test cases for new type combinations (like extension-wrapped unions) ensures that all code paths are exercised and regressions are caught.