paleolimbot commented on code in PR #926:
URL: https://github.com/apache/arrow-nanoarrow/pull/926#discussion_r3939247565
##########
src/nanoarrow/ipc/writer_test.cc:
##########
@@ -204,3 +204,180 @@ TEST(NanoarrowIpcWriter, FileWriting) {
auto after_footer = p->bytes_written;
EXPECT_GT(after_footer, after_eos);
}
+
+TEST(NanoarrowIpcWriter, WriteDictionaryBatch) {
+ struct ArrowError error;
+
+ nanoarrow::UniqueBuffer output;
+ nanoarrow::ipc::UniqueOutputStream stream;
+ ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), output.get()),
NANOARROW_OK);
+
+ nanoarrow::ipc::UniqueWriter writer;
+ ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK);
+
+ auto* p = static_cast<struct ArrowIpcWriterPrivate*>(writer->private_data);
+
+ // Build a simple Utf8 values array
+ nanoarrow::UniqueSchema values_schema;
+ ASSERT_EQ(ArrowSchemaInitFromType(values_schema.get(),
NANOARROW_TYPE_STRING),
+ NANOARROW_OK);
+
+ nanoarrow::UniqueArray values_array;
+ ASSERT_EQ(ArrowArrayInitFromSchema(values_array.get(), values_schema.get(),
nullptr),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayStartAppending(values_array.get()), NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendString(values_array.get(), ArrowCharView("foo")),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendString(values_array.get(), ArrowCharView("bar")),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayFinishBuildingDefault(values_array.get(), &error),
NANOARROW_OK)
+ << error.message;
+
+ nanoarrow::UniqueArrayView values_view;
+ ASSERT_EQ(ArrowArrayViewInitFromSchema(values_view.get(),
values_schema.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowArrayViewSetArray(values_view.get(), values_array.get(),
&error),
+ NANOARROW_OK)
+ << error.message;
+
+ // stream mode: write a DictionaryBatch — bytes are emitted but no block is
tracked
+ EXPECT_EQ(p->bytes_written, 0);
+ EXPECT_EQ(p->footer.dictionary_blocks.size_bytes, 0);
+
+ EXPECT_EQ(ArrowIpcWriterWriteDictionaryBatch(writer.get(),
/*dictionary_id=*/0,
+ /*is_delta=*/0,
values_view.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+
+ auto after_dict_stream = p->bytes_written;
+ EXPECT_GT(after_dict_stream, 0);
+ // no block tracked in stream mode
+ EXPECT_EQ(p->footer.dictionary_blocks.size_bytes, 0);
+
+ // file mode: the block is tracked in the footer
+ nanoarrow::ipc::UniqueOutputStream stream2;
+ nanoarrow::UniqueBuffer output2;
+ ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream2.get(), output2.get()),
NANOARROW_OK);
+
+ nanoarrow::ipc::UniqueWriter writer2;
+ ASSERT_EQ(ArrowIpcWriterInit(writer2.get(), stream2.get()), NANOARROW_OK);
+
+ auto* p2 = static_cast<struct ArrowIpcWriterPrivate*>(writer2->private_data);
+
+ ASSERT_EQ(ArrowIpcWriterStartFile(writer2.get(), &error), NANOARROW_OK)
+ << error.message;
+ EXPECT_EQ(p2->footer.dictionary_blocks.size_bytes, 0);
+
+ EXPECT_EQ(ArrowIpcWriterWriteDictionaryBatch(writer2.get(),
/*dictionary_id=*/0,
+ /*is_delta=*/0,
values_view.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+
+ // one block tracked in file mode
+ EXPECT_EQ(p2->footer.dictionary_blocks.size_bytes, sizeof(struct
ArrowIpcFileBlock));
+}
+
+// Build a struct array with a single dictionary-encoded (int32 -> utf8) child.
+static void MakeDictionaryStructArray(struct ArrowArray* array,
+ struct ArrowSchema* schema) {
+ ASSERT_EQ(ArrowSchemaInitFromType(schema, NANOARROW_TYPE_STRUCT),
NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaAllocateChildren(schema, 1), NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaSetName(schema->children[0], "dict_col"), NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaAllocateDictionary(schema->children[0]), NANOARROW_OK);
+ ASSERT_EQ(
+ ArrowSchemaInitFromType(schema->children[0]->dictionary,
NANOARROW_TYPE_STRING),
+ NANOARROW_OK);
+
+ ASSERT_EQ(ArrowArrayInitFromSchema(array, schema, nullptr), NANOARROW_OK);
+ struct ArrowArray* indices = array->children[0];
+ struct ArrowArray* values = indices->dictionary;
+
+ ASSERT_EQ(ArrowArrayStartAppending(array), NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendString(values, ArrowCharView("foo")),
NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendString(values, ArrowCharView("bar")),
NANOARROW_OK);
+
+ ASSERT_EQ(ArrowArrayAppendInt(indices, 0), NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendInt(indices, 1), NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendInt(indices, 0), NANOARROW_OK);
+ array->length = 3;
+
+ ASSERT_EQ(ArrowArrayFinishBuildingDefault(array, nullptr), NANOARROW_OK);
+}
+
+// Write a dictionary-encoded stream through the high-level WriteArrayStream
path
+// and read it back through the IPC reader, confirming the DictionaryBatch is
+// emitted automatically and the decoded values match.
+TEST(NanoarrowIpcWriter, RoundtripDictionaryStream) {
+ struct ArrowError error;
+
+ nanoarrow::UniqueSchema schema;
+ nanoarrow::UniqueArray array;
+ MakeDictionaryStructArray(array.get(), schema.get());
+
+ nanoarrow::UniqueArrayStream array_stream;
+ ASSERT_EQ(ArrowBasicArrayStreamInit(array_stream.get(), schema.get(), 1),
NANOARROW_OK);
+ ArrowBasicArrayStreamSetArray(array_stream.get(), 0, array.get());
+
+ nanoarrow::UniqueBuffer output;
+ nanoarrow::ipc::UniqueOutputStream out_stream;
+ ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(out_stream.get(), output.get()),
NANOARROW_OK);
+
+ nanoarrow::ipc::UniqueWriter writer;
+ ASSERT_EQ(ArrowIpcWriterInit(writer.get(), out_stream.get()), NANOARROW_OK);
+ ASSERT_EQ(ArrowIpcWriterWriteArrayStream(writer.get(), array_stream.get(),
&error),
+ NANOARROW_OK)
+ << error.message;
+
+ // Read the encoded bytes back
+ struct ArrowIpcInputStream input;
+ ASSERT_EQ(ArrowIpcInputStreamInitBuffer(&input, output.get()), NANOARROW_OK);
+
+ nanoarrow::UniqueArrayStream reader;
+ ASSERT_EQ(ArrowIpcArrayStreamReaderInit(reader.get(), &input, nullptr),
NANOARROW_OK);
+
+ nanoarrow::UniqueSchema roundtrip_schema;
+ ASSERT_EQ(ArrowArrayStreamGetSchema(reader.get(), roundtrip_schema.get(),
&error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(roundtrip_schema->n_children, 1);
+ ASSERT_NE(roundtrip_schema->children[0]->dictionary, nullptr);
+ EXPECT_STREQ(roundtrip_schema->children[0]->dictionary->format, "u");
+
+ nanoarrow::UniqueArray roundtrip_array;
+ ASSERT_EQ(ArrowArrayStreamGetNext(reader.get(), roundtrip_array.get(),
&error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(roundtrip_array->length, 3);
+ ASSERT_EQ(roundtrip_array->n_children, 1);
+ ASSERT_NE(roundtrip_array->children[0]->dictionary, nullptr);
+ EXPECT_EQ(roundtrip_array->children[0]->dictionary->length, 2);
+
+ // Validate the decoded indices resolve to the original values
+ nanoarrow::UniqueArrayView view;
+ ASSERT_EQ(ArrowArrayViewInitFromSchema(view.get(), roundtrip_schema.get(),
&error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowArrayViewSetArray(view.get(), roundtrip_array.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+
+ struct ArrowArrayView* indices_view = view->children[0];
+ struct ArrowArrayView* values_view = indices_view->dictionary;
+ ASSERT_NE(values_view, nullptr);
+ EXPECT_EQ(ArrowArrayViewGetIntUnsafe(indices_view, 0), 0);
+ EXPECT_EQ(ArrowArrayViewGetIntUnsafe(indices_view, 1), 1);
+ EXPECT_EQ(ArrowArrayViewGetIntUnsafe(indices_view, 2), 0);
+
+ struct ArrowStringView v0 = ArrowArrayViewGetStringUnsafe(values_view, 0);
+ struct ArrowStringView v1 = ArrowArrayViewGetStringUnsafe(values_view, 1);
+ EXPECT_EQ(std::string(v0.data, v0.size_bytes), "foo");
+ EXPECT_EQ(std::string(v1.data, v1.size_bytes), "bar");
+
+ ASSERT_EQ(ArrowArrayStreamGetNext(reader.get(), roundtrip_array.get(),
&error),
Review Comment:
I believe the CI error is is a memory leak coming from here (just need to
reset the array before reading into it)
```suggestion
roundtrip_array.reset();
ASSERT_EQ(ArrowArrayStreamGetNext(reader.get(), roundtrip_array.get(),
&error),
```
##########
src/nanoarrow/ipc/writer.c:
##########
@@ -317,6 +317,73 @@ ArrowErrorCode ArrowIpcWriterWriteArrayView(struct
ArrowIpcWriter* writer,
return NANOARROW_OK;
}
+ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch(
+ struct ArrowIpcWriter* writer, int64_t dictionary_id, char is_delta,
+ const struct ArrowArrayView* values_view, struct ArrowError* error) {
+ NANOARROW_DCHECK(writer != NULL && writer->private_data != NULL &&
values_view != NULL);
+ struct ArrowIpcWriterPrivate* private =
+ (struct ArrowIpcWriterPrivate*)writer->private_data;
+
+ NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffer, 0, 0));
+ NANOARROW_ASSERT_OK(ArrowBufferResize(&private->body_buffer, 0, 0));
+
+ NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeSimpleDictionaryBatch(
+ &private->encoder, dictionary_id, is_delta, values_view,
&private->body_buffer,
+ error));
+ NANOARROW_RETURN_NOT_OK_WITH_ERROR(
+ ArrowIpcEncoderFinalizeBuffer(&private->encoder, /*encapsulate=*/1,
+ &private->buffer),
+ error);
+
+ if (private->writing_file) {
+ _NANOARROW_CHECK_RANGE(private->buffer.size_bytes, 0, INT32_MAX);
+ struct ArrowIpcFileBlock block = {
+ .offset = private->bytes_written,
+ .metadata_length = (int32_t) private->buffer.size_bytes,
+ .body_length = private->body_buffer.size_bytes,
+ };
+ NANOARROW_RETURN_NOT_OK_WITH_ERROR(
+ ArrowBufferAppend(&private->footer.dictionary_blocks, &block,
sizeof(block)),
+ error);
+ }
+ private->bytes_written += private->buffer.size_bytes;
+ private->bytes_written += private->body_buffer.size_bytes;
+
+ NANOARROW_RETURN_NOT_OK(ArrowIpcOutputStreamWrite(
+ &private->output_stream, ArrowBufferToBufferView(&private->buffer),
error));
+ NANOARROW_RETURN_NOT_OK(ArrowIpcOutputStreamWrite(
+ &private->output_stream, ArrowBufferToBufferView(&private->body_buffer),
error));
+ return NANOARROW_OK;
+}
+
+// Walk the array in the same depth-first order the schema encoder uses to
assign
+// dictionary ids (see ArrowIpcDictionaryEncodingsAppendSchema): a
dictionary-encoded
+// node claims the next id before descending into its children and then its
values.
+// Emitting a full (non-delta) DictionaryBatch for each dictionary before every
+// RecordBatch keeps each batch's indices valid against the dictionary that
precedes
+// it, which is required because each array in the stream carries its own
dictionary.
Review Comment:
I think this is a good scope for this PR, but we will want to ensure we have
some way to avoid emitting duplicate dictionaries as tracked follow-up.
##########
src/nanoarrow/ipc/writer.c:
##########
@@ -317,6 +317,73 @@ ArrowErrorCode ArrowIpcWriterWriteArrayView(struct
ArrowIpcWriter* writer,
return NANOARROW_OK;
}
+ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch(
+ struct ArrowIpcWriter* writer, int64_t dictionary_id, char is_delta,
+ const struct ArrowArrayView* values_view, struct ArrowError* error) {
+ NANOARROW_DCHECK(writer != NULL && writer->private_data != NULL &&
values_view != NULL);
+ struct ArrowIpcWriterPrivate* private =
+ (struct ArrowIpcWriterPrivate*)writer->private_data;
+
+ NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffer, 0, 0));
+ NANOARROW_ASSERT_OK(ArrowBufferResize(&private->body_buffer, 0, 0));
+
+ NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeSimpleDictionaryBatch(
+ &private->encoder, dictionary_id, is_delta, values_view,
&private->body_buffer,
+ error));
+ NANOARROW_RETURN_NOT_OK_WITH_ERROR(
+ ArrowIpcEncoderFinalizeBuffer(&private->encoder, /*encapsulate=*/1,
+ &private->buffer),
+ error);
+
+ if (private->writing_file) {
+ _NANOARROW_CHECK_RANGE(private->buffer.size_bytes, 0, INT32_MAX);
+ struct ArrowIpcFileBlock block = {
+ .offset = private->bytes_written,
+ .metadata_length = (int32_t) private->buffer.size_bytes,
+ .body_length = private->body_buffer.size_bytes,
+ };
+ NANOARROW_RETURN_NOT_OK_WITH_ERROR(
+ ArrowBufferAppend(&private->footer.dictionary_blocks, &block,
sizeof(block)),
+ error);
+ }
+ private->bytes_written += private->buffer.size_bytes;
+ private->bytes_written += private->body_buffer.size_bytes;
+
+ NANOARROW_RETURN_NOT_OK(ArrowIpcOutputStreamWrite(
+ &private->output_stream, ArrowBufferToBufferView(&private->buffer),
error));
+ NANOARROW_RETURN_NOT_OK(ArrowIpcOutputStreamWrite(
+ &private->output_stream, ArrowBufferToBufferView(&private->body_buffer),
error));
+ return NANOARROW_OK;
+}
+
+// Walk the array in the same depth-first order the schema encoder uses to
assign
+// dictionary ids (see ArrowIpcDictionaryEncodingsAppendSchema): a
dictionary-encoded
+// node claims the next id before descending into its children and then its
values.
+// Emitting a full (non-delta) DictionaryBatch for each dictionary before every
+// RecordBatch keeps each batch's indices valid against the dictionary that
precedes
+// it, which is required because each array in the stream carries its own
dictionary.
+static ArrowErrorCode ArrowIpcWriterWriteDictionariesForArrayView(
+ struct ArrowIpcWriter* writer, const struct ArrowArrayView* array_view,
+ int64_t* next_id, struct ArrowError* error) {
+ if (array_view->dictionary != NULL) {
+ int64_t dictionary_id = (*next_id)++;
+ NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionaryBatch(
+ writer, dictionary_id, /*is_delta=*/0, array_view->dictionary, error));
+ }
+
+ for (int64_t i = 0; i < array_view->n_children; i++) {
+ NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArrayView(
+ writer, array_view->children[i], next_id, error));
+ }
+
+ if (array_view->dictionary != NULL) {
+ NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArrayView(
+ writer, array_view->dictionary, next_id, error));
+ }
Review Comment:
The encoder rejects dictionary values that have a dictionary. Should this
error instead? (Or if it works, can we add a test?)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
