paleolimbot commented on code in PR #935:
URL: https://github.com/apache/arrow-nanoarrow/pull/935#discussion_r3969721216
##########
src/nanoarrow/ipc/encoder.c:
##########
@@ -695,40 +800,207 @@ static ArrowErrorCode
ArrowIpcEncoderBuildContiguousBodyBufferCallback(
return NANOARROW_OK;
}
-static ArrowErrorCode ArrowIpcEncoderEncodeRecordBatchImpl(
+// Make sure there is a scratch ArrowBuffer for each of n compressed buffers
+static ArrowErrorCode ArrowIpcEncoderReserveCompressedBuffers(
+ struct ArrowIpcEncoderPrivate* private, int64_t n) {
+ if (n <= private->n_compressed_buffers) {
+ return NANOARROW_OK;
+ }
+
+ struct ArrowBuffer* buffers =
+ (struct ArrowBuffer*)ArrowMalloc(n * sizeof(struct ArrowBuffer));
+ if (buffers == NULL) {
+ return ENOMEM;
+ }
+
+ if (private->n_compressed_buffers > 0) {
+ memcpy(buffers, private->compressed_buffers,
+ private->n_compressed_buffers * sizeof(struct ArrowBuffer));
+ ArrowFree(private->compressed_buffers);
+ }
+ for (int64_t i = private->n_compressed_buffers; i < n; i++) {
+ ArrowBufferInit(&buffers[i]);
+ }
+
+ private->compressed_buffers = buffers;
+ private->n_compressed_buffers = n;
+ return NANOARROW_OK;
+}
+
+// Replace the collected buffer views of the message being encoded with views
of their
+// compressed form: the uncompressed length as a little-endian int64 followed
by the
+// compressed bytes. Buffers that do not shrink are stored uncompressed with a
prefix of
+// -1 and empty buffers are left as they are, matching Arrow C++. All buffers
are queued
+// with the compressor before waiting, so that a compressor may compress them
in
+// parallel.
+static ArrowErrorCode ArrowIpcEncoderCompressBuffers(
+ struct ArrowIpcEncoderPrivate* private, struct ArrowError* error) {
+ struct ArrowBufferView* views = (struct
ArrowBufferView*)private->buffer_views.data;
+ int64_t n_views = private->buffer_views.size_bytes / (int64_t)sizeof(*views);
+ NANOARROW_RETURN_NOT_OK_WITH_ERROR(
+ ArrowIpcEncoderReserveCompressedBuffers(private, n_views), error);
+
+ // Allocate every prefix before queueing work so allocation failures cannot
leave
+ // jobs referring to the caller's source buffers or our scratch buffers.
+ for (int64_t i = 0; i < n_views; i++) {
+ if (views[i].size_bytes == 0) {
+ continue;
+ }
+
+ // placeholder for the prefix, then the compressed bytes
+ struct ArrowBuffer* dst = &private->compressed_buffers[i];
+ NANOARROW_ASSERT_OK(ArrowBufferResize(dst, 0, 0));
+ NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBufferAppendInt64(dst, 0), error);
+ }
+
+ for (int64_t i = 0; i < n_views; i++) {
+ if (views[i].size_bytes == 0) {
+ continue;
+ }
+
+ int result = private->compressor.compress_add(&private->compressor,
views[i],
+
&private->compressed_buffers[i], error);
+ if (result != NANOARROW_OK) {
+ // don't leave queued work referring to our buffers behind
+ struct ArrowError ignored;
+ NANOARROW_UNUSED(
+ private->compressor.compress_wait(&private->compressor, -1,
&ignored));
Review Comment:
tiny nit: You should be able to pass `NULL` for the error here.
--
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]
