Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 6 additions & 61 deletions src/nanoarrow/ipc/decoder.c
Original file line number Diff line number Diff line change
Expand Up @@ -644,19 +644,14 @@ static inline int32_t ArrowIpcReadInt32LE(struct ArrowBufferView* data, int swap
return value;
}

// Returned by an internal ArrowIpcMetadataVisitFunction to stop iterating early.
// This is never returned to a caller of the public API.
#define _NANOARROW_IPC_VISIT_STOP (-1)

// Visits each key/value pair in a flatbuffers vector of KeyValue.
//
// Keys and values point into the message and are passed with an explicit size because
// both are optional fields whose content may contain embedded nulls; a KeyValue with no
// key or no value is visited with an empty string view.
static ArrowErrorCode ArrowIpcDecoderVisitMetadata(ns(KeyValue_vec_t) kv_vec,
ArrowIpcMetadataVisitFunction visit,
void* private_data,
struct ArrowError* error) {
static ArrowErrorCode ArrowIpcDecoderVisitMetadata(
ns(KeyValue_vec_t) kv_vec, ArrowIpcMessageMetadataVisitFunction visit,
void* private_data, struct ArrowError* error) {
int64_t n_pairs = ns(KeyValue_vec_len(kv_vec));

for (int64_t i = 0; i < n_pairs; i++) {
Expand Down Expand Up @@ -1767,59 +1762,9 @@ ArrowErrorCode ArrowIpcDecoderGetMessageMetadata(struct ArrowIpcDecoder* decoder
return ArrowIpcDecoderBuildMetadata(private_data->last_message_metadata, out, error);
}

struct ArrowIpcMetadataValueLookup {
struct ArrowStringView key;
struct ArrowStringView* value_out;
};

static ArrowErrorCode ArrowIpcDecoderMatchMetadataKey(struct ArrowStringView key,
struct ArrowStringView value,
void* private_data,
struct ArrowError* error) {
NANOARROW_UNUSED(error);
struct ArrowIpcMetadataValueLookup* lookup =
(struct ArrowIpcMetadataValueLookup*)private_data;

if (key.size_bytes != lookup->key.size_bytes) {
return NANOARROW_OK;
}

if (key.size_bytes > 0 &&
memcmp(key.data, lookup->key.data, (size_t)key.size_bytes) != 0) {
return NANOARROW_OK;
}

*lookup->value_out = value;
return _NANOARROW_IPC_VISIT_STOP;
}

ArrowErrorCode ArrowIpcDecoderGetMessageMetadataValue(struct ArrowIpcDecoder* decoder,
struct ArrowStringView key,
struct ArrowStringView* value_out,
struct ArrowError* error) {
NANOARROW_DCHECK(decoder != NULL && decoder->private_data != NULL && value_out != NULL);
struct ArrowIpcDecoderPrivate* private_data =
(struct ArrowIpcDecoderPrivate*)decoder->private_data;

struct ArrowIpcMetadataValueLookup lookup;
lookup.key = key;
lookup.value_out = value_out;

int result =
ArrowIpcDecoderVisitMetadata(private_data->last_message_metadata,
&ArrowIpcDecoderMatchMetadataKey, &lookup, error);
if (result == _NANOARROW_IPC_VISIT_STOP) {
// key was found and value_out was set
return NANOARROW_OK;
}

return result;
}

ArrowErrorCode ArrowIpcDecoderVisitMessageMetadata(struct ArrowIpcDecoder* decoder,
ArrowIpcMetadataVisitFunction visit,
void* private_data,
struct ArrowError* error) {
ArrowErrorCode ArrowIpcDecoderVisitMessageMetadata(
struct ArrowIpcDecoder* decoder, ArrowIpcMessageMetadataVisitFunction visit,
void* private_data, struct ArrowError* error) {
NANOARROW_DCHECK(decoder != NULL && decoder->private_data != NULL && visit != NULL);
struct ArrowIpcDecoderPrivate* decoder_private =
(struct ArrowIpcDecoderPrivate*)decoder->private_data;
Expand Down
106 changes: 72 additions & 34 deletions src/nanoarrow/ipc/decoder_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,13 @@
// For bswap32()
#include "flatcc/portable/pendian.h"

#include "flatcc/flatcc_builder.h"
#include "flatcc_generated.h"
#include "nanoarrow/nanoarrow_gtest_util.hpp"
#include "nanoarrow/nanoarrow_ipc.hpp"

#define ns(x) FLATBUFFERS_WRAP_NAMESPACE(org_apache_arrow_flatbuf, x)

#if defined(NANOARROW_BUILD_TESTS_WITH_ARROW)
using namespace arrow;
#endif
Expand Down Expand Up @@ -1382,28 +1386,69 @@ void AssertArrayViewIdentical(const struct ArrowArrayView* actual,
}
}

// A Message whose Message.custom_metadata and Schema.custom_metadata each contain a
// KeyValue with no key. Both `key` and `value` are optional fields of KeyValue in the
// IPC format, so this passes flatbuffer verification and must not crash the decoder.
// Generated with flatcc: a Schema message with no fields, whose two metadata vectors
// each contain KeyValue{value: "message_value" / "schema_value"} and no key.
alignas(8) static uint8_t kKeylessMetadataSchema[] = {
0xff, 0xff, 0xff, 0xff, 0x90, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x86, 0xff,
0xff, 0xff, 0x04, 0x00, 0x01, 0x00, 0x2c, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00,
0x01, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0xb0, 0xff, 0xff, 0xff, 0x04, 0x00,
0x00, 0x00, 0x0d, 0x00, 0x00, 0x00, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x5f,
0x76, 0x61, 0x6c, 0x75, 0x65, 0x00, 0x00, 0x00, 0xc4, 0xff, 0xff, 0xff, 0x2c, 0x00,
0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00,
0xe0, 0xff, 0xff, 0xff, 0x04, 0x00, 0x00, 0x00, 0x0c, 0x00, 0x00, 0x00, 0x73, 0x63,
0x68, 0x65, 0x6d, 0x61, 0x5f, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x08, 0x00, 0x08, 0x00, 0x00, 0x00, 0x04, 0x00, 0x0a, 0x00,
0x0c, 0x00, 0x00, 0x00, 0x04, 0x00, 0x08, 0x00, 0x0e, 0x00, 0x10, 0x00, 0x04, 0x00,
0x06, 0x00, 0x08, 0x00, 0x00, 0x00, 0x0c, 0x00, 0x00, 0x00, 0x00, 0x00};
// Builds an encapsulated Schema message whose Message.custom_metadata and
// Schema.custom_metadata each contain a KeyValue with no key. Both `key` and `value`
// are optional fields of KeyValue in the IPC format, so this passes flatbuffer
// verification and must not crash the decoder. The nanoarrow encoder always writes
// both, so this message has to be built with flatcc directly.
static void MakeKeylessMetadataSchemaMessage(struct ArrowBuffer* out) {
flatcc_builder_t builder;
flatcc_builder_init(&builder);

ASSERT_EQ(ns(Message_start_as_root(&builder)), 0);
ASSERT_EQ(ns(Message_version_add(&builder, ns(MetadataVersion_V5))), 0);

ASSERT_EQ(ns(Message_header_Schema_start(&builder)), 0);
ASSERT_EQ(ns(Schema_endianness_add(&builder, ns(Endianness_Little))), 0);
ASSERT_EQ(ns(Schema_fields_start(&builder)), 0);
ASSERT_EQ(ns(Schema_fields_end(&builder)), 0);
ASSERT_EQ(ns(Schema_custom_metadata_start(&builder)), 0);
ASSERT_EQ(ns(Schema_custom_metadata_push_start(&builder)), 0);
ASSERT_EQ(ns(KeyValue_value_create_str(&builder, "schema_value")), 0);
ASSERT_NE(ns(Schema_custom_metadata_push_end(&builder)), nullptr);
ASSERT_EQ(ns(Schema_custom_metadata_end(&builder)), 0);
ASSERT_EQ(ns(Message_header_Schema_end(&builder)), 0);

ASSERT_EQ(ns(Message_custom_metadata_start(&builder)), 0);
ASSERT_EQ(ns(Message_custom_metadata_push_start(&builder)), 0);
ASSERT_EQ(ns(KeyValue_value_create_str(&builder, "message_value")), 0);
ASSERT_NE(ns(Message_custom_metadata_push_end(&builder)), nullptr);
ASSERT_EQ(ns(Message_custom_metadata_end(&builder)), 0);

ASSERT_EQ(ns(Message_bodyLength_add(&builder, 0)), 0);
ASSERT_NE(ns(Message_end_as_root(&builder)), 0);

// Encapsulate: continuation, little endian header size, header, padding to 8 bytes
size_t size = flatcc_builder_get_buffer_size(&builder);
int64_t padded_size = _ArrowRoundUpToMultipleOf8(static_cast<int64_t>(size));
ASSERT_EQ(
ArrowBufferReserve(out, 2 * static_cast<int64_t>(sizeof(int32_t)) + padded_size),
NANOARROW_OK);
ASSERT_EQ(ArrowBufferAppendInt32(out, -1), NANOARROW_OK);
#if defined(__BIG_ENDIAN__)
ASSERT_EQ(ArrowBufferAppendInt32(
out, static_cast<int32_t>(bswap32(static_cast<uint32_t>(padded_size)))),
NANOARROW_OK);
#else
ASSERT_EQ(ArrowBufferAppendInt32(out, static_cast<int32_t>(padded_size)), NANOARROW_OK);
#endif
ASSERT_NE(flatcc_builder_copy_buffer(&builder, out->data + out->size_bytes, size),
nullptr);
out->size_bytes += size;
while (out->size_bytes % 8 != 0) {
out->data[out->size_bytes++] = 0;
}

flatcc_builder_clear(&builder);
}

TEST(NanoarrowIpcTest, NanoarrowIpcDecodeMetadataWithoutKey) {
nanoarrow::UniqueBuffer message;
ASSERT_NO_FATAL_FAILURE(MakeKeylessMetadataSchemaMessage(message.get()));

struct ArrowBufferView data;
data.data.as_uint8 = kKeylessMetadataSchema;
data.size_bytes = sizeof(kKeylessMetadataSchema);
data.data.as_uint8 = message->data;
data.size_bytes = message->size_bytes;

struct ArrowError error;
nanoarrow::ipc::UniqueDecoder decoder;
Expand All @@ -1419,15 +1464,11 @@ TEST(NanoarrowIpcTest, NanoarrowIpcDecodeMetadataWithoutKey) {
ArrowIpcDecoderGetMessageMetadata(decoder.get(), message_metadata.get(), &error),
NANOARROW_OK)
<< error.message;
EXPECT_EQ(
ArrowSchemaMetadataToString(reinterpret_cast<const char*>(message_metadata->data)),
"=message_value");
const char* packed = reinterpret_cast<const char*>(message_metadata->data);
EXPECT_EQ(ArrowSchemaMetadataToString(packed), "=message_value");

struct ArrowStringView value = ArrowCharView(nullptr);
ASSERT_EQ(ArrowIpcDecoderGetMessageMetadataValue(decoder.get(), ArrowCharView(""),
&value, &error),
NANOARROW_OK)
<< error.message;
ASSERT_EQ(ArrowMetadataGetValue(packed, ArrowCharView(""), &value), NANOARROW_OK);
EXPECT_EQ(std::string(value.data, value.size_bytes), "message_value");

// ...including the metadata of the schema the message contains
Expand Down Expand Up @@ -1885,17 +1926,14 @@ TEST(NanoarrowIpcTest, NanoarrowIpcMessageMetadataArrowInterop) {
&error),
NANOARROW_OK)
<< error.message;
EXPECT_EQ(ArrowSchemaMetadataToString(
reinterpret_cast<const char*>(batch_message_metadata->data)),
"key1=value1, key2=value2");
const char* packed = reinterpret_cast<const char*>(batch_message_metadata->data);
EXPECT_EQ(ArrowSchemaMetadataToString(packed), "key1=value1, key2=value2");

for (int64_t i = 0; i < custom_metadata->size(); i++) {
struct ArrowStringView value = ArrowCharView(nullptr);
ASSERT_EQ(ArrowIpcDecoderGetMessageMetadataValue(
decoder.get(), ArrowCharView(custom_metadata->key(i).c_str()), &value,
&error),
NANOARROW_OK)
<< error.message;
ASSERT_EQ(ArrowMetadataGetValue(
packed, ArrowCharView(custom_metadata->key(i).c_str()), &value),
NANOARROW_OK);
EXPECT_EQ(std::string(value.data, value.size_bytes), custom_metadata->value(i));
}
}
Expand Down
41 changes: 25 additions & 16 deletions src/nanoarrow/ipc/encoder_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -248,28 +248,29 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderMessageMetadataRoundtrip) {

EXPECT_EQ(DecodeMessageMetadata(message.get(), decoder.get()), key_values);

// Values can also be read in place, without copying
struct ArrowStringView value = ArrowCharView(nullptr);
ASSERT_EQ(ArrowIpcDecoderGetMessageMetadataValue(
decoder.get(), ArrowCharView("cache-control"), &value, &error),
// Single keys are read from the decoded metadata with ArrowMetadataGetValue()
nanoarrow::UniqueBuffer packed;
ASSERT_EQ(ArrowIpcDecoderGetMessageMetadata(decoder.get(), packed.get(), &error),
NANOARROW_OK)
<< error.message;
const char* packed_metadata = reinterpret_cast<const char*>(packed->data);

struct ArrowStringView value = ArrowCharView(nullptr);
ASSERT_EQ(
ArrowMetadataGetValue(packed_metadata, ArrowCharView("cache-control"), &value),
NANOARROW_OK);
EXPECT_EQ(std::string(value.data, value.size_bytes), "no-store");

// A key that isn't present leaves value_out untouched
value = ArrowCharView(nullptr);
ASSERT_EQ(ArrowIpcDecoderGetMessageMetadataValue(
decoder.get(), ArrowCharView("not-a-key"), &value, &error),
NANOARROW_OK)
<< error.message;
ASSERT_EQ(ArrowMetadataGetValue(packed_metadata, ArrowCharView("not-a-key"), &value),
NANOARROW_OK);
EXPECT_EQ(value.data, nullptr);

// A key which is a prefix of a present key is not a match
value = ArrowCharView(nullptr);
ASSERT_EQ(ArrowIpcDecoderGetMessageMetadataValue(decoder.get(), ArrowCharView("cache"),
&value, &error),
NANOARROW_OK)
<< error.message;
ASSERT_EQ(ArrowMetadataGetValue(packed_metadata, ArrowCharView("cache"), &value),
NANOARROW_OK);
EXPECT_EQ(value.data, nullptr);

// The metadata applied to exactly one message: the next one has none
Expand Down Expand Up @@ -395,12 +396,20 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderMessageMetadataEmpty) {
ASSERT_EQ(message->size_bytes, baseline->size_bytes);
EXPECT_EQ(memcmp(message->data, baseline->data, message->size_bytes), 0);

// Reading in place from a message without metadata finds nothing
struct ArrowStringView value = ArrowCharView(nullptr);
ASSERT_EQ(ArrowIpcDecoderGetMessageMetadataValue(decoder.get(), ArrowCharView("key"),
&value, &error),
// The decoded metadata of a message without any is empty, and is still safe to
// hand to ArrowMetadataGetValue()
EXPECT_EQ(DecodeMessageMetadata(message.get(), decoder.get()), KeyValues{});

nanoarrow::UniqueBuffer packed;
ASSERT_EQ(ArrowIpcDecoderGetMessageMetadata(decoder.get(), packed.get(), &error),
NANOARROW_OK)
<< error.message;
ASSERT_EQ(packed->data, nullptr);

struct ArrowStringView value = ArrowCharView(nullptr);
ASSERT_EQ(ArrowMetadataGetValue(reinterpret_cast<const char*>(packed->data),
ArrowCharView("key"), &value),
NANOARROW_OK);
EXPECT_EQ(value.data, nullptr);
}

Expand Down
34 changes: 10 additions & 24 deletions src/nanoarrow/nanoarrow_ipc.h
Original file line number Diff line number Diff line change
Expand Up @@ -82,8 +82,6 @@
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcDecoderSetEndianness)
#define ArrowIpcDecoderGetMessageMetadata \
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcDecoderGetMessageMetadata)
#define ArrowIpcDecoderGetMessageMetadataValue \
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcDecoderGetMessageMetadataValue)
#define ArrowIpcDecoderVisitMessageMetadata \
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcDecoderVisitMessageMetadata)
#define ArrowIpcDecoderPeekFooter \
Expand Down Expand Up @@ -623,10 +621,9 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcDecoderSetDecompressor(
///
/// Returning any value other than NANOARROW_OK will stop the visit and cause that
/// value to be returned by ArrowIpcDecoderVisitMessageMetadata().
typedef ArrowErrorCode (*ArrowIpcMetadataVisitFunction)(struct ArrowStringView key,
struct ArrowStringView value,
void* private_data,
struct ArrowError* error);
typedef ArrowErrorCode (*ArrowIpcMessageMetadataVisitFunction)(
struct ArrowStringView key, struct ArrowStringView value, void* private_data,
struct ArrowError* error);

/// \brief Peek at a message header
///
Expand Down Expand Up @@ -691,36 +688,25 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcDecoderDecodeHeader(struct ArrowIpcDecoder*
/// (const char*)out->data can be passed to ArrowSchemaSetMetadata() or
/// ArrowMetadataReaderInit().
///
/// To read a single key, pass (const char*)out->data to ArrowMetadataGetValue(); to
/// read every pair without copying, use ArrowIpcDecoderVisitMessageMetadata().
///
/// Returns ENOMEM if allocation fails, EINVAL if the metadata cannot be decoded, or
/// NANOARROW_OK otherwise.
NANOARROW_DLL ArrowErrorCode ArrowIpcDecoderGetMessageMetadata(
struct ArrowIpcDecoder* decoder, struct ArrowBuffer* out, struct ArrowError* error);

/// \brief Get one value from the custom metadata of the most recently decoded message
///
/// Unlike ArrowIpcDecoderGetMessageMetadata(), this does not copy: the value returned
/// points into the message header passed to ArrowIpcDecoderVerifyHeader() or
/// ArrowIpcDecoderDecodeHeader() and is only valid until that data is invalidated or
/// another message header is decoded.
///
/// If key occurs more than once, the first value is returned. If key does not occur,
/// value_out is left unmodified: initialize it with ArrowCharView(NULL) and check
/// value_out->data for NULL to detect a missing key.
NANOARROW_DLL ArrowErrorCode ArrowIpcDecoderGetMessageMetadataValue(
struct ArrowIpcDecoder* decoder, struct ArrowStringView key,
struct ArrowStringView* value_out, struct ArrowError* error);

/// \brief Visit each key/value pair in the most recently decoded message's metadata
///
/// Like ArrowIpcDecoderGetMessageMetadataValue(), the keys and values passed to visit
/// point into the message header passed to ArrowIpcDecoderVerifyHeader() or
/// ArrowIpcDecoderDecodeHeader() and must not be retained beyond the lifetime of that
/// Unlike ArrowIpcDecoderGetMessageMetadata(), this does not copy: the keys and values
/// passed to visit point into the message header passed to ArrowIpcDecoderVerifyHeader()
/// or ArrowIpcDecoderDecodeHeader() and must not be retained beyond the lifetime of that
/// data. private_data and error are passed to each invocation of visit.
///
/// Returns the first non-NANOARROW_OK value returned by visit, or NANOARROW_OK if all
/// pairs were visited.
NANOARROW_DLL ArrowErrorCode ArrowIpcDecoderVisitMessageMetadata(
struct ArrowIpcDecoder* decoder, ArrowIpcMetadataVisitFunction visit,
struct ArrowIpcDecoder* decoder, ArrowIpcMessageMetadataVisitFunction visit,
void* private_data, struct ArrowError* error);

/// \brief Decode an ArrowSchema
Expand Down
Loading