This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-nanoarrow.git
The following commit(s) were added to refs/heads/main by this push:
new f1a66ffe feat: Use reference counted arrays when decoding batches that
contain dictionaries (#895)
f1a66ffe is described below
commit f1a66ffe8f60957b064e4796772bad7d49db75e8
Author: Dewey Dunnington <[email protected]>
AuthorDate: Wed Jul 29 21:18:09 2026 -0400
feat: Use reference counted arrays when decoding batches that contain
dictionaries (#895)
This PR updates the dictionary decoder to use the reference counted
array clones in its dictionary decoding process. This is already covered
by existing test code (dictionary decoding already worked, the
dictionary values just weren't shared).
---
src/nanoarrow/common/array.c | 14 ++++++++++--
src/nanoarrow/common/utils.c | 3 ++-
src/nanoarrow/ipc/decoder.c | 48 +++++++++++++++++++++++++++++++++------
src/nanoarrow/ipc/decoder_test.cc | 4 ++--
src/nanoarrow/nanoarrow.h | 5 +++-
5 files changed, 61 insertions(+), 13 deletions(-)
diff --git a/src/nanoarrow/common/array.c b/src/nanoarrow/common/array.c
index 7184c827..68c42d0d 100644
--- a/src/nanoarrow/common/array.c
+++ b/src/nanoarrow/common/array.c
@@ -701,8 +701,18 @@ ArrowErrorCode ArrowArrayMoveShared(struct ArrowArray*
array, struct ArrowArray*
struct ArrowArray tmp;
tmp.release = NULL;
ArrowErrorCode result = ArrowArrayMoveSharedInternal(array, shared);
- if (result != NANOARROW_OK && tmp.release != NULL) {
- ArrowArrayRelease(&tmp);
+ if (result != NANOARROW_OK) {
+ // On failure, release the temporary output
+ if (tmp.release != NULL) {
+ ArrowArrayRelease(&tmp);
+ }
+
+ // Because this operation may have partially moved the input array at this
point
+ // we also have to release it on failure to be predictable. These failures
are
+ // usually failed heap allocations and are difficult to trigger in
practice.
+ if (array->release != NULL) {
+ ArrowArrayRelease(array);
+ }
}
return result;
diff --git a/src/nanoarrow/common/utils.c b/src/nanoarrow/common/utils.c
index 976a509b..c86d00b3 100644
--- a/src/nanoarrow/common/utils.c
+++ b/src/nanoarrow/common/utils.c
@@ -434,7 +434,8 @@ int ArrowIsSharedBuffer(struct ArrowBuffer* buffer) {
ArrowErrorCode ArrowSharedBufferClone(struct ArrowBuffer* shared,
struct ArrowBuffer* shared_out) {
- if (shared->size_bytes == 0) {
+ // If the buffer has no data, initialize an empty buffer
+ if (shared->data == NULL) {
ArrowBufferInit(shared_out);
return NANOARROW_OK;
}
diff --git a/src/nanoarrow/ipc/decoder.c b/src/nanoarrow/ipc/decoder.c
index 42ba2457..c9687f0c 100644
--- a/src/nanoarrow/ipc/decoder.c
+++ b/src/nanoarrow/ipc/decoder.c
@@ -325,7 +325,11 @@ static ArrowErrorCode ArrowIpcDictionaryReplace(struct
ArrowIpcDictionary* dicti
ArrowArrayRelease(&dictionary->current_value);
}
- ArrowArrayMove(value, &dictionary->current_value);
+ // Convert to a shared array so that clones can reference the same data
+ // without copying
+ struct ArrowArray shared;
+ NANOARROW_RETURN_NOT_OK(ArrowArrayMoveShared(value, &shared));
+ ArrowArrayMove(&shared, &dictionary->current_value);
return NANOARROW_OK;
}
@@ -417,15 +421,26 @@ static ArrowErrorCode
ArrowIpcDictionariesInitDictionaries(
// Set the initial array value to a valid array with zero length. This is
// needed because empty and/or all null columns may not have a dictionary
// message emitted before a record batch arrives.
- result = ArrowArrayInitFromSchema(&dictionary->current_value,
- encoding->schema->dictionary, error);
+ struct ArrowArray initial_value;
+ result =
+ ArrowArrayInitFromSchema(&initial_value, encoding->schema->dictionary,
error);
if (result != NANOARROW_OK) {
*num_initialized_decoders_out = i + 1;
return result;
}
- result = ArrowArrayFinishBuildingDefault(&dictionary->current_value,
error);
+ result = ArrowArrayFinishBuildingDefault(&initial_value, error);
if (result != NANOARROW_OK) {
+ ArrowArrayRelease(&initial_value);
+ *num_initialized_decoders_out = i + 1;
+ return result;
+ }
+
+ // Convert to a shared array so that clones can reference the same data
+ // without copying
+ result = ArrowArrayMoveShared(&initial_value, &dictionary->current_value);
+ if (result != NANOARROW_OK) {
+ ArrowArrayRelease(&initial_value);
*num_initialized_decoders_out = i + 1;
return result;
}
@@ -2319,9 +2334,12 @@ static int ArrowIpcDecoderWalkGetArray(struct
ArrowArrayView* array_view,
}
if (array_view->dictionary != NULL) {
- // TODO: this currently copies the array for every output.
- NANOARROW_RETURN_NOT_OK(ArrowIpcDecoderWalkGetArray(
- array_view->dictionary, array->dictionary, out->dictionary, error));
+ // Release the dictionary that was pre-initialized by
ArrowArrayInitFromArrayView
+ if (out->dictionary->release != NULL) {
+ ArrowArrayRelease(out->dictionary);
+ }
+ // Move the pre-cloned shared dictionary to output (avoids copying)
+ ArrowArrayMove(array->dictionary, out->dictionary);
}
return NANOARROW_OK;
@@ -2360,6 +2378,22 @@ static int ArrowIpcDecoderWalkSetArrayView(struct
ArrowIpcDecoder* decoder,
// decode.
NANOARROW_RETURN_NOT_OK(
ArrowArrayViewSetArray(array_view->dictionary, dictionary, error));
+
+ // Clone the shared dictionary for output (avoids copying dictionary data)
+ if (array->dictionary != NULL) {
+ if (array->dictionary->release != NULL) {
+ ArrowArrayRelease(array->dictionary);
+ }
+ int clone_result =
+ ArrowArrayCloneShared((struct ArrowArray*)dictionary,
array->dictionary);
+ if (clone_result != NANOARROW_OK) {
+ ArrowErrorSet(error,
+ "Failed to clone shared dictionary with ID %" PRId64
+ " (dictionary may not be shared)",
+ ipc_field->dictionary_id);
+ return clone_result;
+ }
+ }
}
ns(FieldNode_struct_t) field =
diff --git a/src/nanoarrow/ipc/decoder_test.cc
b/src/nanoarrow/ipc/decoder_test.cc
index 3c366017..900bc32f 100644
--- a/src/nanoarrow/ipc/decoder_test.cc
+++ b/src/nanoarrow/ipc/decoder_test.cc
@@ -1312,7 +1312,7 @@ TEST_P(ArrowTypeParameterizedTestFixture,
NanoarrowIpcArrowArrayRoundtrip) {
NANOARROW_OK);
auto maybe_batch = arrow::ImportRecordBatch(&array, dummy_schema);
- ASSERT_TRUE(maybe_batch.ok());
+ ASSERT_TRUE(maybe_batch.ok()) << maybe_batch.status().ToString();
EXPECT_EQ(maybe_batch.ValueUnsafe()->ToString(), empty->ToString());
// Arrow C++ MakeEmpty() loses the ordered=1 flag and unsigned index types
for
@@ -1348,7 +1348,7 @@ TEST_P(ArrowTypeParameterizedTestFixture,
NanoarrowIpcArrowArrayRoundtrip) {
NANOARROW_OK);
maybe_batch = arrow::ImportRecordBatch(&array, dummy_schema);
- ASSERT_TRUE(maybe_batch.ok());
+ ASSERT_TRUE(maybe_batch.ok()) << maybe_batch.status().ToString();
EXPECT_EQ(maybe_batch.ValueUnsafe()->ToString(), nulls->ToString());
EXPECT_TRUE(maybe_batch.ValueUnsafe()->Equals(*nulls));
diff --git a/src/nanoarrow/nanoarrow.h b/src/nanoarrow/nanoarrow.h
index 1ec13876..564db40a 100644
--- a/src/nanoarrow/nanoarrow.h
+++ b/src/nanoarrow/nanoarrow.h
@@ -1154,7 +1154,10 @@ NANOARROW_DLL ArrowErrorCode ArrowArrayFinishBuilding(
/// reference-counted. On success, shared is a new ArrowArray whose buffers
/// are backed by ArrowSharedArray references, and array is consumed
/// (release set to NULL). The resulting shared array can be safely moved
-/// or have its buffers cloned via ArrowSharedBufferClone().
+/// or have its buffers cloned via ArrowSharedBufferClone(). On error,
+/// (e.g., failure to allocate a copy of a buffer), the input array is
+/// released as it may have been partially moved at the point the error
+/// occurs.
NANOARROW_DLL ArrowErrorCode ArrowArrayMoveShared(struct ArrowArray* array,
struct ArrowArray* shared);
