diff --git a/src/nanoarrow/common/array.c b/src/nanoarrow/common/array.c index 68c42d0d2..de3def350 100644 --- a/src/nanoarrow/common/array.c +++ b/src/nanoarrow/common/array.c @@ -77,6 +77,455 @@ int ArrowArrayIsInternal(struct ArrowArray* array) { return array->release == &ArrowArrayReleaseInternal; } +static int ArrowTypeIsSignedInteger(enum ArrowType type) { + switch (type) { + case NANOARROW_TYPE_BOOL: + case NANOARROW_TYPE_INT8: + case NANOARROW_TYPE_INT16: + case NANOARROW_TYPE_INT32: + case NANOARROW_TYPE_INT64: + case NANOARROW_TYPE_DATE32: + case NANOARROW_TYPE_DATE64: + case NANOARROW_TYPE_TIMESTAMP: + case NANOARROW_TYPE_TIME32: + case NANOARROW_TYPE_TIME64: + case NANOARROW_TYPE_DURATION: + return 1; + default: + return 0; + } +} + +static int ArrowTypeIsUnsignedInteger(enum ArrowType type) { + switch (type) { + case NANOARROW_TYPE_UINT8: + case NANOARROW_TYPE_UINT16: + case NANOARROW_TYPE_UINT32: + case NANOARROW_TYPE_UINT64: + return 1; + default: + return 0; + } +} + +static int ArrowTypeIsFloatingPoint(enum ArrowType type) { + return type == NANOARROW_TYPE_HALF_FLOAT || type == NANOARROW_TYPE_FLOAT || + type == NANOARROW_TYPE_DOUBLE; +} + +static int ArrowTypeIsBinaryLike(enum ArrowType type) { + switch (type) { + case NANOARROW_TYPE_STRING: + case NANOARROW_TYPE_BINARY: + case NANOARROW_TYPE_FIXED_SIZE_BINARY: + case NANOARROW_TYPE_LARGE_STRING: + case NANOARROW_TYPE_LARGE_BINARY: + case NANOARROW_TYPE_BINARY_VIEW: + case NANOARROW_TYPE_STRING_VIEW: + return 1; + default: + return 0; + } +} + +static int ArrowTypeIsListLike(enum ArrowType type) { + switch (type) { + case NANOARROW_TYPE_LIST: + case NANOARROW_TYPE_LARGE_LIST: + case NANOARROW_TYPE_MAP: + case NANOARROW_TYPE_LIST_VIEW: + case NANOARROW_TYPE_LARGE_LIST_VIEW: + case NANOARROW_TYPE_FIXED_SIZE_LIST: + return 1; + default: + return 0; + } +} + +static int ArrowArrayCanAppendStorageType(enum ArrowType dst_type, + enum ArrowType src_type) { + if (src_type == NANOARROW_TYPE_NA) { + return 1; + } + + if (ArrowTypeIsSignedInteger(src_type) || ArrowTypeIsUnsignedInteger(src_type)) { + return ArrowTypeIsSignedInteger(dst_type) || ArrowTypeIsUnsignedInteger(dst_type) || + ArrowTypeIsFloatingPoint(dst_type); + } + + if (ArrowTypeIsFloatingPoint(src_type)) { + return ArrowTypeIsFloatingPoint(dst_type); + } + + if (ArrowTypeIsBinaryLike(src_type)) { + return ArrowTypeIsBinaryLike(dst_type); + } + + if (ArrowTypeIsListLike(src_type)) { + return ArrowTypeIsListLike(dst_type); + } + + return dst_type == src_type; +} + +static ArrowErrorCode ArrowArrayCheckCanAppendStorageFromArrayView( + struct ArrowArray* dst, const struct ArrowArrayView* src, struct ArrowError* error) { + if (src->offset < 0 || src->length < 0 || src->offset > INT64_MAX - src->length) { + ArrowErrorSet(error, + "Expected source offset and length to describe a valid int64 " + "range but found %" PRId64 " and %" PRId64, + src->offset, src->length); + return EINVAL; + } + + if (!ArrowArrayIsInternal(dst)) { + ArrowErrorSet(error, "Expected destination to be an internal ArrowArray"); + return EINVAL; + } + + struct ArrowArrayPrivateData* private_data = + (struct ArrowArrayPrivateData*)dst->private_data; + enum ArrowType dst_type = private_data->storage_type; + + if (src->storage_type == NANOARROW_TYPE_DENSE_UNION || + src->storage_type == NANOARROW_TYPE_SPARSE_UNION) { + ArrowErrorSet(error, "Appending array views is not supported for %s", + ArrowTypeString(src->storage_type)); + return ENOTSUP; + } + + if (!ArrowArrayCanAppendStorageType(dst_type, src->storage_type)) { + ArrowErrorSet(error, "Can't append %s storage to an array with %s storage", + ArrowTypeString(src->storage_type), ArrowTypeString(dst_type)); + return EINVAL; + } + + if ((src->dictionary == NULL) != (dst->dictionary == NULL)) { + ArrowErrorSet(error, + "Can't append storage when exactly one of source and destination " + "is dictionary-encoded"); + return EINVAL; + } + + if (src->storage_type == NANOARROW_TYPE_NA) { + return NANOARROW_OK; + } + + if (src->n_children != dst->n_children) { + ArrowErrorSet(error, + "Expected source and destination to have the same number of " + "children but found %" PRId64 " and %" PRId64, + src->n_children, dst->n_children); + return EINVAL; + } + + if (src->storage_type == NANOARROW_TYPE_FIXED_SIZE_LIST && + dst_type == NANOARROW_TYPE_FIXED_SIZE_LIST && + src->layout.child_size_elements != private_data->layout.child_size_elements) { + ArrowErrorSet(error, + "Expected source and destination fixed-size list child sizes to " + "match but found %" PRId64 " and %" PRId64, + src->layout.child_size_elements, + private_data->layout.child_size_elements); + return EINVAL; + } + + if (src->storage_type == NANOARROW_TYPE_FIXED_SIZE_BINARY && + dst_type == NANOARROW_TYPE_FIXED_SIZE_BINARY && + src->layout.element_size_bits[1] != private_data->layout.element_size_bits[1]) { + ArrowErrorSet(error, + "Expected source and destination fixed-size binary widths to " + "match but found %" PRId64 " and %" PRId64 " bits", + src->layout.element_size_bits[1], + private_data->layout.element_size_bits[1]); + return EINVAL; + } + + for (int64_t i = 0; i < src->n_children; i++) { + NANOARROW_RETURN_NOT_OK(ArrowArrayCheckCanAppendStorageFromArrayView( + dst->children[i], src->children[i], error)); + } + + return NANOARROW_OK; +} + +static ArrowErrorCode ArrowArrayAppendStorageFromArrayViewElement( + struct ArrowArray* dst, const struct ArrowArrayView* src, int64_t i, + struct ArrowError* error) { + if (ArrowArrayViewIsNull(src, i)) { + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowArrayAppendNull(dst, 1), error); + return NANOARROW_OK; + } + + switch (src->storage_type) { + case NANOARROW_TYPE_NA: + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowArrayAppendNull(dst, 1), error); + return NANOARROW_OK; + case NANOARROW_TYPE_BOOL: + case NANOARROW_TYPE_INT8: + case NANOARROW_TYPE_INT16: + case NANOARROW_TYPE_INT32: + case NANOARROW_TYPE_INT64: + case NANOARROW_TYPE_DATE32: + case NANOARROW_TYPE_DATE64: + case NANOARROW_TYPE_TIMESTAMP: + case NANOARROW_TYPE_TIME32: + case NANOARROW_TYPE_TIME64: + case NANOARROW_TYPE_DURATION: + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowArrayAppendInt(dst, ArrowArrayViewGetIntUnsafe(src, i)), error); + return NANOARROW_OK; + case NANOARROW_TYPE_UINT8: + case NANOARROW_TYPE_UINT16: + case NANOARROW_TYPE_UINT32: + case NANOARROW_TYPE_UINT64: + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowArrayAppendUInt(dst, ArrowArrayViewGetUIntUnsafe(src, i)), error); + return NANOARROW_OK; + case NANOARROW_TYPE_HALF_FLOAT: + case NANOARROW_TYPE_FLOAT: + case NANOARROW_TYPE_DOUBLE: + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowArrayAppendDouble(dst, ArrowArrayViewGetDoubleUnsafe(src, i)), error); + return NANOARROW_OK; + case NANOARROW_TYPE_STRING: + case NANOARROW_TYPE_BINARY: + case NANOARROW_TYPE_FIXED_SIZE_BINARY: + case NANOARROW_TYPE_LARGE_STRING: + case NANOARROW_TYPE_LARGE_BINARY: + case NANOARROW_TYPE_BINARY_VIEW: + case NANOARROW_TYPE_STRING_VIEW: + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowArrayAppendBytes(dst, ArrowArrayViewGetBytesUnsafe(src, i)), error); + return NANOARROW_OK; + case NANOARROW_TYPE_INTERVAL_MONTHS: + case NANOARROW_TYPE_INTERVAL_DAY_TIME: + case NANOARROW_TYPE_INTERVAL_MONTH_DAY_NANO: { + struct ArrowInterval interval; + ArrowIntervalInit(&interval, src->storage_type); + ArrowArrayViewGetIntervalUnsafe(src, i, &interval); + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowArrayAppendInterval(dst, &interval), error); + return NANOARROW_OK; + } + case NANOARROW_TYPE_STRUCT: + for (int64_t child_i = 0; child_i < src->n_children; child_i++) { + NANOARROW_RETURN_NOT_OK(ArrowArrayAppendStorageFromArrayViewElement( + dst->children[child_i], src->children[child_i], src->offset + i, error)); + } + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowArrayFinishElement(dst), error); + return NANOARROW_OK; + case NANOARROW_TYPE_LIST: + case NANOARROW_TYPE_LARGE_LIST: + case NANOARROW_TYPE_MAP: + case NANOARROW_TYPE_LIST_VIEW: + case NANOARROW_TYPE_LARGE_LIST_VIEW: { + int64_t logical_i = src->offset + i; + int64_t child_offset = ArrowArrayViewListChildOffset(src, logical_i); + int64_t child_length; + if (src->storage_type == NANOARROW_TYPE_LIST_VIEW) { + child_length = src->buffer_views[2].data.as_int32[logical_i]; + } else if (src->storage_type == NANOARROW_TYPE_LARGE_LIST_VIEW) { + child_length = src->buffer_views[2].data.as_int64[logical_i]; + } else { + child_length = ArrowArrayViewListChildOffset(src, logical_i + 1) - child_offset; + } + + for (int64_t child_i = 0; child_i < child_length; child_i++) { + NANOARROW_RETURN_NOT_OK(ArrowArrayAppendStorageFromArrayViewElement( + dst->children[0], src->children[0], child_offset + child_i, error)); + } + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowArrayFinishElement(dst), error); + return NANOARROW_OK; + } + case NANOARROW_TYPE_FIXED_SIZE_LIST: { + int64_t child_offset = (src->offset + i) * src->layout.child_size_elements; + for (int64_t child_i = 0; child_i < src->layout.child_size_elements; child_i++) { + NANOARROW_RETURN_NOT_OK(ArrowArrayAppendStorageFromArrayViewElement( + dst->children[0], src->children[0], child_offset + child_i, error)); + } + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowArrayFinishElement(dst), error); + return NANOARROW_OK; + } + case NANOARROW_TYPE_DECIMAL32: + case NANOARROW_TYPE_DECIMAL64: + case NANOARROW_TYPE_DECIMAL128: + case NANOARROW_TYPE_DECIMAL256: + case NANOARROW_TYPE_DENSE_UNION: + case NANOARROW_TYPE_SPARSE_UNION: + case NANOARROW_TYPE_RUN_END_ENCODED: + case NANOARROW_TYPE_UNINITIALIZED: + default: + ArrowErrorSet(error, "Appending array views is not supported for %s", + ArrowTypeString(src->storage_type)); + return ENOTSUP; + } +} + +static int ArrowArrayCanAppendFixedWidthStorage(struct ArrowArray* dst, + const struct ArrowArrayView* src) { + struct ArrowArrayPrivateData* private_data = + (struct ArrowArrayPrivateData*)dst->private_data; + return private_data->storage_type == src->storage_type && dst->n_buffers == 2 && + src->n_children == 0 && + src->layout.buffer_type[1] == NANOARROW_BUFFER_TYPE_DATA && + src->layout.element_size_bits[1] > 0 && + src->layout.element_size_bits[1] % 8 == 0; +} + +static ArrowErrorCode ArrowArrayAppendValidityFromArrayView( + struct ArrowArray* dst, const struct ArrowArrayView* src, struct ArrowError* error) { + struct ArrowBitmap* dst_validity = ArrowArrayValidityBitmap(dst); + const uint8_t* src_validity = src->buffer_views[0].data.as_uint8; + if (src_validity == NULL && dst_validity->buffer.data == NULL) { + return NANOARROW_OK; + } + + if (dst_validity->buffer.data == NULL) { + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBitmapAppend(dst_validity, 1, dst->length), + error); + } + + if (src_validity == NULL) { + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBitmapAppend(dst_validity, 1, src->length), + error); + return NANOARROW_OK; + } + + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBitmapReserve(dst_validity, src->length), + error); + int8_t validity[1024] = {0}; + for (int64_t offset = 0; offset < src->length;) { + int64_t remaining = src->length - offset; + int64_t chunk_size = + remaining < (int64_t)sizeof(validity) ? remaining : (int64_t)sizeof(validity); + ArrowBitsUnpackInt8(src_validity, src->offset + offset, chunk_size, validity); + ArrowBitmapAppendInt8Unsafe(dst_validity, validity, chunk_size); + offset += chunk_size; + } + + return NANOARROW_OK; +} + +static ArrowErrorCode ArrowArrayAppendFixedWidthStorageFromArrayView( + struct ArrowArray* dst, const struct ArrowArrayView* src, struct ArrowError* error) { + if (src->length == 0) { + return NANOARROW_OK; + } + + int64_t element_size_bytes = src->layout.element_size_bits[1] / 8; + if (src->offset > INT64_MAX / element_size_bytes || + src->length > INT64_MAX / element_size_bytes) { + ArrowErrorSet(error, + "Expected fixed-width append size to fit in int64 but found " + "element size %" PRId64 ", offset %" PRId64 ", and length %" PRId64, + element_size_bytes, src->offset, src->length); + return EOVERFLOW; + } + + int64_t src_offset_bytes = src->offset * element_size_bytes; + int64_t src_size_bytes = src->length * element_size_bytes; + if (src->buffer_views[1].data.as_uint8 == NULL || + src_offset_bytes > src->buffer_views[1].size_bytes || + src_size_bytes > src->buffer_views[1].size_bytes - src_offset_bytes) { + ArrowErrorSet(error, + "Expected fixed-width source buffer to contain %" PRId64 + " bytes at offset %" PRId64 " but its size is %" PRId64, + src_size_bytes, src_offset_bytes, src->buffer_views[1].size_bytes); + return EINVAL; + } + + NANOARROW_RETURN_NOT_OK(ArrowArrayAppendValidityFromArrayView(dst, src, error)); + + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowBufferAppend(ArrowArrayBuffer(dst, 1), + src->buffer_views[1].data.as_uint8 + src_offset_bytes, + src_size_bytes), + error); + + if (src->buffer_views[0].data.as_uint8 != NULL) { + dst->null_count += src->length - ArrowBitCountSet(src->buffer_views[0].data.as_uint8, + src->offset, src->length); + } + dst->length += src->length; + return NANOARROW_OK; +} + +static int64_t ArrowArrayViewResolveRun(const struct ArrowArrayView* run_ends, + int64_t logical_offset) { + if (run_ends->length <= 1) { + return 0; + } + + switch (run_ends->storage_type) { + case NANOARROW_TYPE_INT32: + return ArrowResolveChunk32( + (int32_t)(logical_offset + 1), + run_ends->buffer_views[1].data.as_int32 + run_ends->offset, 0, + (int32_t)run_ends->length); + case NANOARROW_TYPE_INT64: + return ArrowResolveChunk64( + logical_offset + 1, run_ends->buffer_views[1].data.as_int64 + run_ends->offset, + 0, run_ends->length); + case NANOARROW_TYPE_INT16: + default: { + int64_t run = 0; + while (run + 1 < run_ends->length && + ArrowArrayViewGetIntUnsafe(run_ends, run) <= logical_offset) { + run++; + } + return run; + } + } +} + +ArrowErrorCode ArrowArrayAppendStorageFromArrayView(struct ArrowArray* dst, + const struct ArrowArrayView* src, + struct ArrowError* error) { + NANOARROW_RETURN_NOT_OK(ArrowArrayCheckCanAppendStorageFromArrayView(dst, src, error)); + + if (ArrowArrayCanAppendFixedWidthStorage(dst, src)) { + return ArrowArrayAppendFixedWidthStorageFromArrayView(dst, src, error); + } + + if (src->storage_type == NANOARROW_TYPE_RUN_END_ENCODED) { + if (src->length == 0) { + return NANOARROW_OK; + } + + const struct ArrowArrayView* run_ends = src->children[0]; + if (run_ends->length == 0) { + ArrowErrorSet(error, + "Expected a non-empty run ends array for a non-empty " + "run-end encoded array"); + return EINVAL; + } + + int64_t first_run = ArrowArrayViewResolveRun(run_ends, src->offset); + int64_t slice_end = src->offset + src->length; + int64_t run_end_offset = dst->length; + for (int64_t i = first_run; i < run_ends->length; i++) { + int64_t src_run_end = ArrowArrayViewGetIntUnsafe(run_ends, i); + int64_t clipped_run_end = src_run_end < slice_end ? src_run_end : slice_end; + int64_t dst_run_end = run_end_offset + clipped_run_end - src->offset; + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowArrayAppendInt(dst->children[0], dst_run_end), error); + NANOARROW_RETURN_NOT_OK(ArrowArrayAppendStorageFromArrayViewElement( + dst->children[1], src->children[1], i, error)); + if (src_run_end >= slice_end) { + break; + } + } + dst->length += src->length; + return NANOARROW_OK; + } + + for (int64_t i = 0; i < src->length; i++) { + NANOARROW_RETURN_NOT_OK( + ArrowArrayAppendStorageFromArrayViewElement(dst, src, i, error)); + } + return NANOARROW_OK; +} + static ArrowErrorCode ArrowArraySetStorageType(struct ArrowArray* array, enum ArrowType storage_type) { switch (storage_type) { diff --git a/src/nanoarrow/common/array_test.cc b/src/nanoarrow/common/array_test.cc index bc4686c2a..e8c9eea9c 100644 --- a/src/nanoarrow/common/array_test.cc +++ b/src/nanoarrow/common/array_test.cc @@ -33,6 +33,7 @@ #include #include #include +#include #include #if defined(NANOARROW_BUILD_TESTS_WITH_ARROW) && defined(ARROW_VERSION_MAJOR) && \ @@ -5087,3 +5088,716 @@ TEST(ArrayMoveSharedTest, ArrayWithNullBuffers) { ArrowArrayRelease(&shared); } + +static ArrowErrorCode AppendStorageFromArrayViewForTest(const struct ArrowArrayView* src, + struct ArrowArray* dst, + struct ArrowError* error) { + NANOARROW_RETURN_NOT_OK(ArrowArrayInitFromArrayView(dst, src, error)); + NANOARROW_RETURN_NOT_OK(ArrowArrayStartAppending(dst)); + NANOARROW_RETURN_NOT_OK(ArrowArrayAppendStorageFromArrayView(dst, src, error)); + return ArrowArrayFinishBuildingDefault(dst, error); +} + +static void ExpectPrimitiveStorageAppendIdentical(struct ArrowArray* src, + enum ArrowType type) { + struct ArrowError error; + struct ArrowArrayView src_view; + ArrowArrayViewInitFromType(&src_view, type); + ASSERT_EQ(ArrowArrayViewSetArray(&src_view, src, &error), NANOARROW_OK) + << error.message; + + struct ArrowArray dst; + ASSERT_EQ(AppendStorageFromArrayViewForTest(&src_view, &dst, &error), NANOARROW_OK) + << error.message; + struct ArrowArrayView dst_view; + ArrowArrayViewInitFromType(&dst_view, type); + ASSERT_EQ(ArrowArrayViewSetArray(&dst_view, &dst, &error), NANOARROW_OK) + << error.message; + + int identical = 0; + ASSERT_EQ(ArrowArrayViewCompare(&src_view, &dst_view, NANOARROW_COMPARE_IDENTICAL, + &identical, &error), + NANOARROW_OK); + EXPECT_EQ(identical, 1) << error.message; + + ArrowArrayViewReset(&dst_view); + ArrowArrayRelease(&dst); + ArrowArrayViewReset(&src_view); +} + +TEST(ArrayTest, ArrayAppendStorageFromArrayViewPrimitiveTypes) { + struct ArrowArray array; + + ASSERT_EQ(ArrowArrayInitFromType(&array, NANOARROW_TYPE_INT64), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&array), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(&array, -42), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendNull(&array, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&array, nullptr), NANOARROW_OK); + ExpectPrimitiveStorageAppendIdentical(&array, NANOARROW_TYPE_INT64); + ArrowArrayRelease(&array); + + ASSERT_EQ(ArrowArrayInitFromType(&array, NANOARROW_TYPE_UINT64), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&array), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendUInt(&array, UINT64_MAX), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&array, nullptr), NANOARROW_OK); + ExpectPrimitiveStorageAppendIdentical(&array, NANOARROW_TYPE_UINT64); + ArrowArrayRelease(&array); + + ASSERT_EQ(ArrowArrayInitFromType(&array, NANOARROW_TYPE_DOUBLE), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&array), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendDouble(&array, 1.25), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&array, nullptr), NANOARROW_OK); + ExpectPrimitiveStorageAppendIdentical(&array, NANOARROW_TYPE_DOUBLE); + ArrowArrayRelease(&array); + + ASSERT_EQ(ArrowArrayInitFromType(&array, NANOARROW_TYPE_STRING), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&array), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendString(&array, "some value"_asv), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&array, nullptr), NANOARROW_OK); + ExpectPrimitiveStorageAppendIdentical(&array, NANOARROW_TYPE_STRING); + ArrowArrayRelease(&array); + + ASSERT_EQ(ArrowArrayInitFromType(&array, NANOARROW_TYPE_INTERVAL_MONTHS), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&array), NANOARROW_OK); + struct ArrowInterval interval; + ArrowIntervalInit(&interval, NANOARROW_TYPE_INTERVAL_MONTHS); + interval.months = 42; + ASSERT_EQ(ArrowArrayAppendInterval(&array, &interval), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&array, nullptr), NANOARROW_OK); + ExpectPrimitiveStorageAppendIdentical(&array, NANOARROW_TYPE_INTERVAL_MONTHS); + ArrowArrayRelease(&array); + + ASSERT_EQ(ArrowArrayInitFromType(&array, NANOARROW_TYPE_DECIMAL128), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&array), NANOARROW_OK); + struct ArrowDecimal decimal; + ArrowDecimalInit(&decimal, 128, 10, 2); + ArrowDecimalSetInt(&decimal, 1234); + ASSERT_EQ(ArrowArrayAppendDecimal(&array, &decimal), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&array, nullptr), NANOARROW_OK); + ExpectPrimitiveStorageAppendIdentical(&array, NANOARROW_TYPE_DECIMAL128); + ArrowArrayRelease(&array); +} + +TEST(ArrayTest, ArrayAppendStorageFromArrayViewNestedAndSliced) { + struct ArrowError error; + struct ArrowSchema schema; + ArrowSchemaInit(&schema); + ASSERT_EQ(ArrowSchemaSetTypeStruct(&schema, 2), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetType(schema.children[0], NANOARROW_TYPE_LIST), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetType(schema.children[0]->children[0], NANOARROW_TYPE_INT32), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetType(schema.children[1], NANOARROW_TYPE_STRING), NANOARROW_OK); + + struct ArrowArray src; + ASSERT_EQ(ArrowArrayInitFromSchema(&src, &schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&src), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(src.children[0]->children[0], 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(src.children[0]->children[0], 2), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishElement(src.children[0]), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendString(src.children[1], "first"_asv), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishElement(&src), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendNull(&src, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishElement(src.children[0]), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendString(src.children[1], "third"_asv), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishElement(&src), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&src, &error), NANOARROW_OK) << error.message; + + struct ArrowArrayView src_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(&src_view, &schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(&src_view, &src, &error), NANOARROW_OK); + src_view.offset = 1; + src_view.length = 2; + src_view.null_count = 1; + + struct ArrowArray dst; + ASSERT_EQ(AppendStorageFromArrayViewForTest(&src_view, &dst, &error), NANOARROW_OK) + << error.message; + EXPECT_EQ(dst.length, 2); + EXPECT_EQ(dst.children[0]->length, 2); + EXPECT_EQ(dst.children[0]->children[0]->length, 0); + EXPECT_EQ(dst.children[1]->length, 2); + + struct ArrowArrayView dst_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(&dst_view, &schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(&dst_view, &dst, &error), NANOARROW_OK); + EXPECT_TRUE(ArrowArrayViewIsNull(&dst_view, 0)); + EXPECT_EQ(ArrowArrayViewGetStringUnsafe(dst_view.children[1], 1), "third"_asv); + + ArrowArrayViewReset(&dst_view); + ArrowArrayRelease(&dst); + ArrowArrayViewReset(&src_view); + ArrowArrayRelease(&src); + ArrowSchemaRelease(&schema); +} + +TEST(ArrayTest, ArrayAppendStorageFromArrayViewRejectsUnions) { + for (enum ArrowType union_type : + {NANOARROW_TYPE_DENSE_UNION, NANOARROW_TYPE_SPARSE_UNION}) { + struct ArrowError error; + struct ArrowSchema schema; + ArrowSchemaInit(&schema); + ASSERT_EQ(ArrowSchemaSetTypeUnion(&schema, union_type, 2), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetType(schema.children[0], NANOARROW_TYPE_INT32), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetType(schema.children[1], NANOARROW_TYPE_STRING), + NANOARROW_OK); + + struct ArrowArray src; + ASSERT_EQ(ArrowArrayInitFromSchema(&src, &schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&src), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(src.children[0], 42), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishUnionElement(&src, 0), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendString(src.children[1], "value"_asv), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishUnionElement(&src, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&src, &error), NANOARROW_OK); + + struct ArrowArrayView src_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(&src_view, &schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(&src_view, &src, &error), NANOARROW_OK); + struct ArrowArray dst; + ASSERT_EQ(AppendStorageFromArrayViewForTest(&src_view, &dst, &error), ENOTSUP); + EXPECT_EQ(dst.length, 0); + EXPECT_EQ(std::string(error.message), "Appending array views is not supported for " + + std::string(ArrowTypeString(union_type))); + + ArrowArrayRelease(&dst); + ArrowArrayViewReset(&src_view); + ArrowArrayRelease(&src); + ArrowSchemaRelease(&schema); + } +} + +TEST(ArrayTest, ArrayAppendStorageFromArrayViewRunEndEncoded) { + struct ArrowError error; + struct ArrowSchema schema; + ArrowSchemaInit(&schema); + ASSERT_EQ(ArrowSchemaSetTypeRunEndEncoded(&schema, NANOARROW_TYPE_INT32), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetType(schema.children[1], NANOARROW_TYPE_STRING), NANOARROW_OK); + + struct ArrowArray src; + ASSERT_EQ(ArrowArrayInitFromSchema(&src, &schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&src), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(src.children[0], 2), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(src.children[0], 3), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendString(src.children[1], "a"_asv), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendString(src.children[1], "b"_asv), NANOARROW_OK); + src.length = 3; + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&src, &error), NANOARROW_OK); + + struct ArrowArrayView src_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(&src_view, &schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(&src_view, &src, &error), NANOARROW_OK); + struct ArrowArray dst; + ASSERT_EQ(ArrowArrayInitFromArrayView(&dst, &src_view, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&dst), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendStorageFromArrayView(&dst, &src_view, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendStorageFromArrayView(&dst, &src_view, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&dst, &error), NANOARROW_OK) << error.message; + EXPECT_EQ(dst.length, 6); + + struct ArrowArrayView dst_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(&dst_view, &schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(&dst_view, &dst, &error), NANOARROW_OK); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(dst_view.children[0], 2), 5); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(dst_view.children[0], 3), 6); + ArrowArrayViewReset(&dst_view); + + src_view.offset = 1; + src_view.length = 1; + ASSERT_EQ(ArrowArrayAppendStorageFromArrayView(&dst, &src_view, &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&dst, &error), NANOARROW_OK) << error.message; + EXPECT_EQ(dst.length, 7); + EXPECT_EQ(dst.children[0]->length, 5); + EXPECT_EQ(dst.children[1]->length, 5); + ASSERT_EQ(ArrowArrayViewInitFromSchema(&dst_view, &schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(&dst_view, &dst, &error), NANOARROW_OK); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(dst_view.children[0], 4), 7); + EXPECT_EQ(ArrowArrayViewGetStringUnsafe(dst_view.children[1], 4), "a"_asv); + + ArrowArrayViewReset(&dst_view); + ArrowArrayRelease(&dst); + ArrowArrayViewReset(&src_view); + ArrowArrayRelease(&src); + ArrowSchemaRelease(&schema); +} + +TEST(ArrayTest, ArrayAppendStorageFromArrayViewFixedWidthSlicedValidity) { + struct ArrowError error; + struct ArrowArray src; + ASSERT_EQ(ArrowArrayInitFromType(&src, NANOARROW_TYPE_INT32), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&src), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendNull(&src, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(&src, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendNull(&src, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(&src, 3), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendNull(&src, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&src, &error), NANOARROW_OK); + + struct ArrowArrayView src_view; + ArrowArrayViewInitFromType(&src_view, NANOARROW_TYPE_INT32); + ASSERT_EQ(ArrowArrayViewSetArray(&src_view, &src, &error), NANOARROW_OK); + src_view.offset = 1; + src_view.length = 3; + src_view.null_count = 1; + + struct ArrowArray dst; + ASSERT_EQ(ArrowArrayInitFromType(&dst, NANOARROW_TYPE_INT32), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&dst), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(&dst, 10), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendStorageFromArrayView(&dst, &src_view, &error), NANOARROW_OK) + << error.message; + + struct ArrowArray all_valid; + ASSERT_EQ(ArrowArrayInitFromType(&all_valid, NANOARROW_TYPE_INT32), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&all_valid), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(&all_valid, 4), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(&all_valid, 5), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&all_valid, &error), NANOARROW_OK); + struct ArrowArrayView all_valid_view; + ArrowArrayViewInitFromType(&all_valid_view, NANOARROW_TYPE_INT32); + ASSERT_EQ(ArrowArrayViewSetArray(&all_valid_view, &all_valid, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendStorageFromArrayView(&dst, &all_valid_view, &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayAppendStorageFromArrayView(&dst, &src_view, &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&dst, &error), NANOARROW_OK) << error.message; + + struct ArrowArrayView dst_view; + ArrowArrayViewInitFromType(&dst_view, NANOARROW_TYPE_INT32); + ASSERT_EQ(ArrowArrayViewSetArray(&dst_view, &dst, &error), NANOARROW_OK); + EXPECT_EQ(dst.length, 9); + EXPECT_EQ(dst.null_count, 2); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(&dst_view, 0), 10); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(&dst_view, 1), 1); + EXPECT_TRUE(ArrowArrayViewIsNull(&dst_view, 2)); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(&dst_view, 3), 3); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(&dst_view, 4), 4); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(&dst_view, 5), 5); + EXPECT_TRUE(ArrowArrayViewIsNull(&dst_view, 7)); + + struct ArrowArray dst_all_valid; + ASSERT_EQ(ArrowArrayInitFromType(&dst_all_valid, NANOARROW_TYPE_INT32), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&dst_all_valid), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendStorageFromArrayView(&dst_all_valid, &all_valid_view, &error), + NANOARROW_OK); + EXPECT_EQ(dst_all_valid.buffers[0], nullptr); + + ArrowArrayRelease(&dst_all_valid); + ArrowArrayViewReset(&dst_view); + ArrowArrayViewReset(&all_valid_view); + ArrowArrayRelease(&all_valid); + ArrowArrayRelease(&dst); + ArrowArrayViewReset(&src_view); + ArrowArrayRelease(&src); +} + +TEST(ArrayTest, ArrayAppendStorageFromArrayViewFixedWidthChunkedValidity) { + constexpr int64_t kLength = 2050; + struct ArrowError error; + struct ArrowArray src; + ASSERT_EQ(ArrowArrayInitFromType(&src, NANOARROW_TYPE_INT32), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&src), NANOARROW_OK); + for (int64_t i = 0; i < kLength; i++) { + if (i % 3 == 0) { + ASSERT_EQ(ArrowArrayAppendNull(&src, 1), NANOARROW_OK); + } else { + ASSERT_EQ(ArrowArrayAppendInt(&src, i), NANOARROW_OK); + } + } + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&src, &error), NANOARROW_OK); + + struct ArrowArrayView src_view; + ArrowArrayViewInitFromType(&src_view, NANOARROW_TYPE_INT32); + ASSERT_EQ(ArrowArrayViewSetArray(&src_view, &src, &error), NANOARROW_OK); + + struct ArrowArray dst; + ASSERT_EQ(AppendStorageFromArrayViewForTest(&src_view, &dst, &error), NANOARROW_OK) + << error.message; + EXPECT_EQ(dst.length, kLength); + EXPECT_EQ(dst.null_count, 684); + + struct ArrowArrayView dst_view; + ArrowArrayViewInitFromType(&dst_view, NANOARROW_TYPE_INT32); + ASSERT_EQ(ArrowArrayViewSetArray(&dst_view, &dst, &error), NANOARROW_OK); + EXPECT_TRUE(ArrowArrayViewIsNull(&dst_view, 1023)); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(&dst_view, 1024), 1024); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(&dst_view, 2048), 2048); + EXPECT_TRUE(ArrowArrayViewIsNull(&dst_view, 2049)); + + ArrowArrayViewReset(&dst_view); + ArrowArrayRelease(&dst); + ArrowArrayViewReset(&src_view); + ArrowArrayRelease(&src); +} + +TEST(ArrayTest, ArrayAppendStorageFromArrayViewFixedToVariableWidth) { + struct ArrowError error; + + struct ArrowSchema fixed_binary_schema; + ArrowSchemaInit(&fixed_binary_schema); + ASSERT_EQ(ArrowSchemaSetTypeFixedSize(&fixed_binary_schema, + NANOARROW_TYPE_FIXED_SIZE_BINARY, 3), + NANOARROW_OK); + struct ArrowArray fixed_binary; + ASSERT_EQ(ArrowArrayInitFromSchema(&fixed_binary, &fixed_binary_schema, &error), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&fixed_binary), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendBytes(&fixed_binary, {{"abc"}, 3}), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&fixed_binary, &error), NANOARROW_OK); + struct ArrowArrayView fixed_binary_view; + ASSERT_EQ( + ArrowArrayViewInitFromSchema(&fixed_binary_view, &fixed_binary_schema, &error), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(&fixed_binary_view, &fixed_binary, &error), + NANOARROW_OK); + + struct ArrowArray binary; + ASSERT_EQ(ArrowArrayInitFromType(&binary, NANOARROW_TYPE_BINARY), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&binary), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendStorageFromArrayView(&binary, &fixed_binary_view, &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&binary, &error), NANOARROW_OK); + struct ArrowArrayView binary_view; + ArrowArrayViewInitFromType(&binary_view, NANOARROW_TYPE_BINARY); + ASSERT_EQ(ArrowArrayViewSetArray(&binary_view, &binary, &error), NANOARROW_OK); + EXPECT_EQ(ArrowArrayViewGetStringUnsafe(&binary_view, 0), "abc"_asv); + + ArrowArrayViewReset(&binary_view); + ArrowArrayRelease(&binary); + ArrowArrayViewReset(&fixed_binary_view); + ArrowArrayRelease(&fixed_binary); + ArrowSchemaRelease(&fixed_binary_schema); + + struct ArrowSchema fixed_list_schema; + ArrowSchemaInit(&fixed_list_schema); + ASSERT_EQ( + ArrowSchemaSetTypeFixedSize(&fixed_list_schema, NANOARROW_TYPE_FIXED_SIZE_LIST, 2), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetType(fixed_list_schema.children[0], NANOARROW_TYPE_INT32), + NANOARROW_OK); + struct ArrowArray fixed_list; + ASSERT_EQ(ArrowArrayInitFromSchema(&fixed_list, &fixed_list_schema, &error), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&fixed_list), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(fixed_list.children[0], 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(fixed_list.children[0], 2), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishElement(&fixed_list), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&fixed_list, &error), NANOARROW_OK); + struct ArrowArrayView fixed_list_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(&fixed_list_view, &fixed_list_schema, &error), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(&fixed_list_view, &fixed_list, &error), NANOARROW_OK); + + struct ArrowSchema list_schema; + ASSERT_EQ(ArrowSchemaInitFromType(&list_schema, NANOARROW_TYPE_LIST), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetType(list_schema.children[0], NANOARROW_TYPE_INT32), + NANOARROW_OK); + struct ArrowArray list; + ASSERT_EQ(ArrowArrayInitFromSchema(&list, &list_schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&list), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendStorageFromArrayView(&list, &fixed_list_view, &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&list, &error), NANOARROW_OK); + struct ArrowArrayView list_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(&list_view, &list_schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(&list_view, &list, &error), NANOARROW_OK); + EXPECT_EQ(ArrowArrayViewListChildOffset(&list_view, 0), 0); + EXPECT_EQ(ArrowArrayViewListChildOffset(&list_view, 1), 2); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(list_view.children[0], 0), 1); + EXPECT_EQ(ArrowArrayViewGetIntUnsafe(list_view.children[0], 1), 2); + + ArrowArrayViewReset(&list_view); + ArrowArrayRelease(&list); + ArrowSchemaRelease(&list_schema); + ArrowArrayViewReset(&fixed_list_view); + ArrowArrayRelease(&fixed_list); + ArrowSchemaRelease(&fixed_list_schema); +} + +TEST(ArrayTest, ArrayAppendStorageFromArrayViewErrors) { + struct ArrowError error; + struct ArrowArray src; + ASSERT_EQ(ArrowArrayInitFromType(&src, NANOARROW_TYPE_INT64), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&src), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(&src, INT8_MAX + 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&src, &error), NANOARROW_OK); + struct ArrowArrayView src_view; + ArrowArrayViewInitFromType(&src_view, NANOARROW_TYPE_INT64); + ASSERT_EQ(ArrowArrayViewSetArray(&src_view, &src, &error), NANOARROW_OK); + + struct ArrowArray dst; + ASSERT_EQ(ArrowArrayInitFromType(&dst, NANOARROW_TYPE_INT8), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&dst), NANOARROW_OK); + EXPECT_EQ(ArrowArrayAppendStorageFromArrayView(&dst, &src_view, &error), EINVAL); + EXPECT_NE(error.message[0], '\0'); + EXPECT_EQ(dst.length, 0); + + ArrowArrayRelease(&dst); + ArrowArrayViewReset(&src_view); + ArrowArrayRelease(&src); + + struct ArrowSchema list_schema; + ASSERT_EQ(ArrowSchemaInitFromType(&list_schema, NANOARROW_TYPE_LIST), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetType(list_schema.children[0], NANOARROW_TYPE_NA), NANOARROW_OK); + ASSERT_EQ(ArrowArrayInitFromSchema(&src, &list_schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&src), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendNull(src.children[0], 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishElement(&src), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&src, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewInitFromSchema(&src_view, &list_schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(&src_view, &src, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayInitFromSchema(&dst, &list_schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&dst), NANOARROW_OK); + dst.children[0]->length = INT32_MAX; + EXPECT_EQ(ArrowArrayAppendStorageFromArrayView(&dst, &src_view, &error), EOVERFLOW); + EXPECT_NE(error.message[0], '\0'); + + ArrowArrayRelease(&dst); + ArrowArrayViewReset(&src_view); + ArrowArrayRelease(&src); + ArrowSchemaRelease(&list_schema); +} + +TEST(ArrayTest, ArrayAppendStorageFromArrayViewDictionary) { + struct ArrowError error; + struct ArrowSchema schema; + ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_INT8), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateDictionary(&schema), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaInitFromType(schema.dictionary, NANOARROW_TYPE_STRING), + NANOARROW_OK); + + struct ArrowArray src; + ASSERT_EQ(ArrowArrayInitFromSchema(&src, &schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&src), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendString(src.dictionary, "a"_asv), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendString(src.dictionary, "b"_asv), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(&src, 0), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(&src, 1), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendInt(&src, 0), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuildingDefault(&src, &error), NANOARROW_OK); + struct ArrowArrayView src_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(&src_view, &schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(&src_view, &src, &error), NANOARROW_OK); + + struct ArrowArray dst; + ASSERT_EQ(ArrowArrayInitFromSchema(&dst, &schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&dst), NANOARROW_OK); + ASSERT_EQ( + ArrowArrayAppendStorageFromArrayView(dst.dictionary, src_view.dictionary, &error), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendStorageFromArrayView(&dst, &src_view, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishBuilding(&dst, NANOARROW_VALIDATION_LEVEL_FULL, &error), + NANOARROW_OK) + << error.message; + + struct ArrowArrayView dst_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(&dst_view, &schema, &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(&dst_view, &dst, &error), NANOARROW_OK); + int identical = 0; + ASSERT_EQ(ArrowArrayViewCompare(&src_view, &dst_view, NANOARROW_COMPARE_IDENTICAL, + &identical, &error), + NANOARROW_OK); + EXPECT_EQ(identical, 1) << error.message; + + struct ArrowArray plain_dst; + ASSERT_EQ(ArrowArrayInitFromType(&plain_dst, NANOARROW_TYPE_INT8), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(&plain_dst), NANOARROW_OK); + EXPECT_EQ(ArrowArrayAppendStorageFromArrayView(&plain_dst, &src_view, &error), EINVAL); + EXPECT_STREQ(error.message, + "Can't append storage when exactly one of source and destination is " + "dictionary-encoded"); + + struct ArrowArrayView plain_view; + ArrowArrayViewInitFromType(&plain_view, NANOARROW_TYPE_INT8); + ASSERT_EQ(ArrowArrayViewSetArray(&plain_view, &plain_dst, &error), NANOARROW_OK); + EXPECT_EQ(ArrowArrayAppendStorageFromArrayView(&dst, &plain_view, &error), EINVAL); + EXPECT_STREQ(error.message, + "Can't append storage when exactly one of source and destination is " + "dictionary-encoded"); + + ArrowArrayViewReset(&plain_view); + ArrowArrayRelease(&plain_dst); + ArrowArrayViewReset(&dst_view); + ArrowArrayRelease(&dst); + ArrowArrayViewReset(&src_view); + ArrowArrayRelease(&src); + ArrowSchemaRelease(&schema); +} + +#if defined(NANOARROW_BUILD_TESTS_WITH_ARROW) + +enum class ArrowCppAppendValueKind { kSigned, kUnsigned, kString }; + +static std::shared_ptr BuildArrowCppAppendArray( + const std::shared_ptr& type, ArrowCppAppendValueKind value_kind) { + auto maybe_builder = MakeBuilder(type); + ARROW_EXPECT_OK(maybe_builder.status()); + auto builder = std::move(maybe_builder).ValueUnsafe(); + + for (int i = 0; i < 3; i++) { + if (i == 1) { + ARROW_EXPECT_OK(builder->AppendNull()); + continue; + } + + std::shared_ptr scalar; + switch (value_kind) { + case ArrowCppAppendValueKind::kSigned: { + auto maybe_scalar = MakeScalar(type, (int64_t)(i == 0 ? -1 : 42)); + ARROW_EXPECT_OK(maybe_scalar.status()); + scalar = maybe_scalar.ValueUnsafe(); + break; + } + case ArrowCppAppendValueKind::kUnsigned: { + auto maybe_scalar = MakeScalar(type, (uint64_t)(i == 0 ? 1 : 42)); + ARROW_EXPECT_OK(maybe_scalar.status()); + scalar = maybe_scalar.ValueUnsafe(); + break; + } + case ArrowCppAppendValueKind::kString: { + auto maybe_scalar = MakeScalar(type, std::string(i == 0 ? "abc" : "def")); + ARROW_EXPECT_OK(maybe_scalar.status()); + scalar = maybe_scalar.ValueUnsafe(); + break; + } + } + + ARROW_EXPECT_OK(builder->AppendScalar(*scalar)); + } + + auto maybe_array = builder->Finish(); + ARROW_EXPECT_OK(maybe_array.status()); + return maybe_array.ValueUnsafe(); +} + +static std::shared_ptr BuildArrowCppAppendListArray( + const std::shared_ptr& type) { + auto child_builder = std::make_shared(); + if (type->id() == Type::LIST) { + ListBuilder builder(default_memory_pool(), child_builder, type); + ARROW_EXPECT_OK(builder.Append()); + ARROW_EXPECT_OK(child_builder->Append(1)); + ARROW_EXPECT_OK(child_builder->Append(2)); + ARROW_EXPECT_OK(builder.AppendNull()); + ARROW_EXPECT_OK(builder.Append()); + ARROW_EXPECT_OK(child_builder->Append(3)); + auto maybe_array = builder.Finish(); + ARROW_EXPECT_OK(maybe_array.status()); + return maybe_array.ValueUnsafe(); + } else { + LargeListBuilder builder(default_memory_pool(), child_builder, type); + ARROW_EXPECT_OK(builder.Append()); + ARROW_EXPECT_OK(child_builder->Append(1)); + ARROW_EXPECT_OK(child_builder->Append(2)); + ARROW_EXPECT_OK(builder.AppendNull()); + ARROW_EXPECT_OK(builder.Append()); + ARROW_EXPECT_OK(child_builder->Append(3)); + auto maybe_array = builder.Finish(); + ARROW_EXPECT_OK(maybe_array.status()); + return maybe_array.ValueUnsafe(); + } +} + +static void ExpectArrowCppStorageAppendEquals(const std::shared_ptr& src, + const std::shared_ptr& expected) { + struct ArrowError error; + struct ArrowArray src_array; + struct ArrowSchema src_schema; + ARROW_EXPECT_OK(ExportArray(*src, &src_array, &src_schema)); + struct ArrowArrayView src_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(&src_view, &src_schema, &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayViewSetArray(&src_view, &src_array, &error), NANOARROW_OK) + << error.message; + + struct ArrowSchema dst_schema; + ARROW_EXPECT_OK(ExportType(*expected->type(), &dst_schema)); + struct ArrowArray dst; + ASSERT_EQ(ArrowArrayInitFromSchema(&dst, &dst_schema, &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayStartAppending(&dst), NANOARROW_OK); + ASSERT_EQ(ArrowArrayAppendStorageFromArrayView(&dst, &src_view, &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayFinishBuilding(&dst, NANOARROW_VALIDATION_LEVEL_FULL, &error), + NANOARROW_OK) + << error.message; + + auto maybe_actual = ImportArray(&dst, &dst_schema); + ARROW_EXPECT_OK(maybe_actual.status()); + EXPECT_TRUE(maybe_actual.ValueUnsafe()->Equals(expected)) + << maybe_actual.ValueUnsafe()->ToString() << "\n!=\n" + << expected->ToString(); + + ArrowArrayViewReset(&src_view); + ArrowArrayRelease(&src_array); + ArrowSchemaRelease(&src_schema); +} + +TEST(ArrayTest, ArrayAppendStorageFromArrayViewArrowCppSignedIntegers) { + std::vector> types = {int8(), int16(), int32(), int64()}; + for (const auto& src_type : types) { + auto src = BuildArrowCppAppendArray(src_type, ArrowCppAppendValueKind::kSigned); + for (const auto& dst_type : types) { + SCOPED_TRACE(src_type->ToString() + " -> " + dst_type->ToString()); + auto expected = + BuildArrowCppAppendArray(dst_type, ArrowCppAppendValueKind::kSigned); + ExpectArrowCppStorageAppendEquals(src, expected); + } + } +} + +TEST(ArrayTest, ArrayAppendStorageFromArrayViewArrowCppUnsignedIntegers) { + std::vector> types = {uint8(), uint16(), uint32(), uint64()}; + for (const auto& src_type : types) { + auto src = BuildArrowCppAppendArray(src_type, ArrowCppAppendValueKind::kUnsigned); + for (const auto& dst_type : types) { + SCOPED_TRACE(src_type->ToString() + " -> " + dst_type->ToString()); + auto expected = + BuildArrowCppAppendArray(dst_type, ArrowCppAppendValueKind::kUnsigned); + ExpectArrowCppStorageAppendEquals(src, expected); + } + } +} + +TEST(ArrayTest, ArrayAppendStorageFromArrayViewArrowCppStrings) { + std::vector> types = {utf8(), large_utf8()}; + for (const auto& src_type : types) { + auto src = BuildArrowCppAppendArray(src_type, ArrowCppAppendValueKind::kString); + for (const auto& dst_type : types) { + SCOPED_TRACE(src_type->ToString() + " -> " + dst_type->ToString()); + auto expected = + BuildArrowCppAppendArray(dst_type, ArrowCppAppendValueKind::kString); + ExpectArrowCppStorageAppendEquals(src, expected); + } + } +} + +TEST(ArrayTest, ArrayAppendStorageFromArrayViewArrowCppBinary) { + std::vector> types = {binary(), large_binary()}; + for (const auto& src_type : types) { + auto src = BuildArrowCppAppendArray(src_type, ArrowCppAppendValueKind::kString); + for (const auto& dst_type : types) { + SCOPED_TRACE(src_type->ToString() + " -> " + dst_type->ToString()); + auto expected = + BuildArrowCppAppendArray(dst_type, ArrowCppAppendValueKind::kString); + ExpectArrowCppStorageAppendEquals(src, expected); + } + } +} + +TEST(ArrayTest, ArrayAppendStorageFromArrayViewArrowCppLists) { + std::vector> types = {list(int32()), large_list(int32())}; + for (const auto& src_type : types) { + auto src = BuildArrowCppAppendListArray(src_type); + for (const auto& dst_type : types) { + SCOPED_TRACE(src_type->ToString() + " -> " + dst_type->ToString()); + auto expected = BuildArrowCppAppendListArray(dst_type); + ExpectArrowCppStorageAppendEquals(src, expected); + } + } +} + +#endif diff --git a/src/nanoarrow/common/buffer_test.cc b/src/nanoarrow/common/buffer_test.cc index 3b3cc8af6..51c0964e5 100644 --- a/src/nanoarrow/common/buffer_test.cc +++ b/src/nanoarrow/common/buffer_test.cc @@ -647,6 +647,16 @@ TEST(BitmapTest, BitmapTestAppendInt8Unsafe) { EXPECT_EQ(ArrowBitGet(bitmap.buffer.data, i), test_values[i - 136]); } + // Append fewer values than are needed to complete the current byte + int8_t short_values[] = {1, 0, 1}; + ASSERT_EQ(ArrowBitmapReserve(&bitmap, 3), NANOARROW_OK); + ArrowBitmapAppendInt8Unsafe(&bitmap, short_values, 3); + EXPECT_EQ(bitmap.size_bits, 207); + EXPECT_EQ(bitmap.buffer.size_bytes, 26); + for (int i = 0; i < 3; i++) { + EXPECT_EQ(ArrowBitGet(bitmap.buffer.data, 204 + i), short_values[i]); + } + ArrowBitmapReset(&bitmap); } @@ -702,6 +712,16 @@ TEST(BitmapTest, BitmapTestAppendInt32Unsafe) { EXPECT_EQ(ArrowBitGet(bitmap.buffer.data, i), test_values[i - 136]); } + // Append fewer values than are needed to complete the current byte + int32_t short_values[] = {1, 0, 1}; + ASSERT_EQ(ArrowBitmapReserve(&bitmap, 3), NANOARROW_OK); + ArrowBitmapAppendInt32Unsafe(&bitmap, short_values, 3); + EXPECT_EQ(bitmap.size_bits, 207); + EXPECT_EQ(bitmap.buffer.size_bytes, 26); + for (int i = 0; i < 3; i++) { + EXPECT_EQ(ArrowBitGet(bitmap.buffer.data, 204 + i), short_values[i]); + } + ArrowBitmapReset(&bitmap); } diff --git a/src/nanoarrow/common/inline_buffer.h b/src/nanoarrow/common/inline_buffer.h index c9a78602b..1d4997380 100644 --- a/src/nanoarrow/common/inline_buffer.h +++ b/src/nanoarrow/common/inline_buffer.h @@ -635,6 +635,10 @@ static inline void ArrowBitmapAppendInt8Unsafe(struct ArrowBitmap* bitmap, // First byte if ((out_i_cursor % 8) != 0) { int64_t n_partial_bits = _ArrowRoundUpToMultipleOf8(out_i_cursor) - out_i_cursor; + if (n_partial_bits > n_remaining) { + n_partial_bits = n_remaining; + } + for (int i = 0; i < n_partial_bits; i++) { ArrowBitSetTo(bitmap->buffer.data, out_i_cursor++, values[i]); } @@ -685,6 +689,10 @@ static inline void ArrowBitmapAppendInt32Unsafe(struct ArrowBitmap* bitmap, // First byte if ((out_i_cursor % 8) != 0) { int64_t n_partial_bits = _ArrowRoundUpToMultipleOf8(out_i_cursor) - out_i_cursor; + if (n_partial_bits > n_remaining) { + n_partial_bits = n_remaining; + } + for (int i = 0; i < n_partial_bits; i++) { ArrowBitSetTo(bitmap->buffer.data, out_i_cursor++, (uint8_t)values[i]); } diff --git a/src/nanoarrow/nanoarrow.h b/src/nanoarrow/nanoarrow.h index 564db40a4..c165317b3 100644 --- a/src/nanoarrow/nanoarrow.h +++ b/src/nanoarrow/nanoarrow.h @@ -122,6 +122,8 @@ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowArraySetValidityBitmap) #define ArrowArraySetBuffer NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowArraySetBuffer) #define ArrowArrayReserve NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowArrayReserve) +#define ArrowArrayAppendStorageFromArrayView \ + NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowArrayAppendStorageFromArrayView) #define ArrowArrayFinishBuilding \ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowArrayFinishBuilding) #define ArrowArrayFinishBuildingDefault \ @@ -1032,6 +1034,17 @@ static inline ArrowErrorCode ArrowArrayStartAppending(struct ArrowArray* array); NANOARROW_DLL ArrowErrorCode ArrowArrayReserve(struct ArrowArray* array, int64_t additional_size_elements); +/// \brief Append storage from an ArrowArrayView to an ArrowArray +/// +/// Appends each logical storage element of array_view to array. array must have +/// been initialized with compatible storage and prepared using +/// ArrowArrayStartAppending(). Dictionary values referenced by array_view are +/// not copied; dictionary-encoded inputs require a dictionary-encoded output. +/// Returns EINVAL for incompatible storage and ENOTSUP for unsupported storage. +NANOARROW_DLL ArrowErrorCode ArrowArrayAppendStorageFromArrayView( + struct ArrowArray* array, const struct ArrowArrayView* array_view, + struct ArrowError* error); + /// \brief Append a null value to an array static inline ArrowErrorCode ArrowArrayAppendNull(struct ArrowArray* array, int64_t n);