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 9e5a991 feat: Add Support for Intervals (#258)
9e5a991 is described below
commit 9e5a9910112a78e1654be2753b34c35ebd74a5c7
Author: William Ayd <[email protected]>
AuthorDate: Thu Jul 20 09:30:26 2023 -0700
feat: Add Support for Intervals (#258)
---
src/nanoarrow/array_inline.h | 43 +++++++++++++++++
src/nanoarrow/array_test.cc | 102 ++++++++++++++++++++++++++++++++++++++++
src/nanoarrow/nanoarrow.h | 7 +++
src/nanoarrow/nanoarrow_types.h | 23 +++++++++
4 files changed, 175 insertions(+)
diff --git a/src/nanoarrow/array_inline.h b/src/nanoarrow/array_inline.h
index 78a8353..1a54caa 100644
--- a/src/nanoarrow/array_inline.h
+++ b/src/nanoarrow/array_inline.h
@@ -529,6 +529,49 @@ static inline ArrowErrorCode ArrowArrayAppendString(struct
ArrowArray* array,
}
}
+static inline ArrowErrorCode ArrowArrayAppendInterval(struct ArrowArray* array,
+ struct ArrowInterval*
value) {
+ struct ArrowArrayPrivateData* private_data =
+ (struct ArrowArrayPrivateData*)array->private_data;
+
+ struct ArrowBuffer* data_buffer = ArrowArrayBuffer(array, 1);
+
+ switch (private_data->storage_type) {
+ case NANOARROW_TYPE_INTERVAL_MONTHS: {
+ if (value->type != NANOARROW_TYPE_INTERVAL_MONTHS) {
+ return EINVAL;
+ }
+
+ NANOARROW_RETURN_NOT_OK(ArrowBufferAppendInt32(data_buffer,
value->months));
+ break;
+ }
+ case NANOARROW_TYPE_INTERVAL_DAY_TIME: {
+ if (value->type != NANOARROW_TYPE_INTERVAL_DAY_TIME) {
+ return EINVAL;
+ }
+
+ NANOARROW_RETURN_NOT_OK(ArrowBufferAppendInt32(data_buffer,
value->days));
+ NANOARROW_RETURN_NOT_OK(ArrowBufferAppendInt32(data_buffer, value->ms));
+ break;
+ }
+ case NANOARROW_TYPE_INTERVAL_MONTH_DAY_NANO: {
+ if (value->type != NANOARROW_TYPE_INTERVAL_MONTH_DAY_NANO) {
+ return EINVAL;
+ }
+
+ NANOARROW_RETURN_NOT_OK(ArrowBufferAppendInt32(data_buffer,
value->months));
+ NANOARROW_RETURN_NOT_OK(ArrowBufferAppendInt32(data_buffer,
value->days));
+ NANOARROW_RETURN_NOT_OK(ArrowBufferAppendInt64(data_buffer, value->ns));
+ break;
+ }
+ default:
+ return EINVAL;
+ }
+
+ array->length++;
+ return NANOARROW_OK;
+}
+
static inline ArrowErrorCode ArrowArrayAppendDecimal(struct ArrowArray* array,
struct ArrowDecimal*
value) {
struct ArrowArrayPrivateData* private_data =
diff --git a/src/nanoarrow/array_test.cc b/src/nanoarrow/array_test.cc
index d9e0de0..9863c8d 100644
--- a/src/nanoarrow/array_test.cc
+++ b/src/nanoarrow/array_test.cc
@@ -23,6 +23,7 @@
#include <arrow/array/builder_decimal.h>
#include <arrow/array/builder_nested.h>
#include <arrow/array/builder_primitive.h>
+#include <arrow/array/builder_time.h>
#include <arrow/array/builder_union.h>
#include <arrow/c/bridge.h>
#include <arrow/compare.h>
@@ -877,6 +878,107 @@ TEST(ArrayTest, ArrayTestAppendToFixedSizeBinaryArray) {
EXPECT_TRUE(arrow_array.ValueUnsafe()->Equals(expected_array.ValueUnsafe()));
}
+TEST(ArrayTest, ArrayTestAppendToIntervalArrayYearMonth) {
+ struct ArrowArray array;
+
+ const int32_t months = 42;
+ struct ArrowInterval interval;
+ ArrowIntervalInit(&interval, ArrowType::NANOARROW_TYPE_INTERVAL_MONTHS);
+ interval.months = 42;
+
+ ASSERT_EQ(ArrowArrayInitFromType(&array, NANOARROW_TYPE_INTERVAL_MONTHS),
NANOARROW_OK);
+ EXPECT_EQ(ArrowArrayStartAppending(&array), NANOARROW_OK);
+ EXPECT_EQ(ArrowArrayAppendInterval(&array, &interval), NANOARROW_OK);
+ EXPECT_EQ(ArrowArrayAppendNull(&array, 1), NANOARROW_OK);
+
+ EXPECT_EQ(ArrowArrayFinishBuildingDefault(&array, nullptr), NANOARROW_OK);
+ EXPECT_EQ(array.length, 2);
+ EXPECT_EQ(array.null_count, 1);
+
+ auto data_buffer = reinterpret_cast<const int32_t*>(array.buffers[1]);
+ EXPECT_EQ(data_buffer[0], months);
+
+ auto arrow_array = ImportArray(&array, month_interval());
+ ARROW_EXPECT_OK(arrow_array);
+
+ // TODO: arrow does not have a builder for MonthIntervals
+ // so no comparison is done after creating the array
+}
+
+TEST(ArrayTest, ArrayTestAppendToIntervalArrayDayTime) {
+ struct ArrowArray array;
+
+ struct ArrowInterval interval;
+ ArrowIntervalInit(&interval, ArrowType::NANOARROW_TYPE_INTERVAL_DAY_TIME);
+ interval.days = 42;
+ interval.ms = 42;
+
+ ASSERT_EQ(ArrowArrayInitFromType(&array, NANOARROW_TYPE_INTERVAL_DAY_TIME),
+ NANOARROW_OK);
+ EXPECT_EQ(ArrowArrayStartAppending(&array), NANOARROW_OK);
+ EXPECT_EQ(ArrowArrayAppendInterval(&array, &interval), NANOARROW_OK);
+ EXPECT_EQ(ArrowArrayAppendNull(&array, 1), NANOARROW_OK);
+
+ EXPECT_EQ(ArrowArrayFinishBuildingDefault(&array, nullptr), NANOARROW_OK);
+ EXPECT_EQ(array.length, 2);
+ EXPECT_EQ(array.null_count, 1);
+
+ auto data_buffer = reinterpret_cast<const uint8_t*>(array.buffers[1]);
+
+ EXPECT_EQ(memcmp(data_buffer, &interval.days, 4), 0);
+ EXPECT_EQ(memcmp(data_buffer + sizeof(interval.days), &interval.ms, 4), 0);
+
+ auto arrow_array = ImportArray(&array, day_time_interval());
+ ARROW_EXPECT_OK(arrow_array);
+
+ auto builder = DayTimeIntervalBuilder();
+ DayTimeIntervalType::DayMilliseconds dm = {42, 42};
+ ARROW_EXPECT_OK(builder.Append(dm));
+ ARROW_EXPECT_OK(builder.AppendNulls(1));
+ auto expected_array = builder.Finish();
+
+ EXPECT_TRUE(arrow_array.ValueUnsafe()->Equals(expected_array.ValueUnsafe()));
+}
+
+TEST(ArrayTest, ArrayTestAppendToIntervalArrayMonthDayNano) {
+ struct ArrowArray array;
+
+ struct ArrowInterval interval;
+ ArrowIntervalInit(&interval,
ArrowType::NANOARROW_TYPE_INTERVAL_MONTH_DAY_NANO);
+ interval.months = 2;
+ interval.days = 12;
+ interval.ns = 42;
+
+ ASSERT_EQ(ArrowArrayInitFromType(&array,
NANOARROW_TYPE_INTERVAL_MONTH_DAY_NANO),
+ NANOARROW_OK);
+ EXPECT_EQ(ArrowArrayStartAppending(&array), NANOARROW_OK);
+ EXPECT_EQ(ArrowArrayAppendInterval(&array, &interval), NANOARROW_OK);
+ EXPECT_EQ(ArrowArrayAppendNull(&array, 1), NANOARROW_OK);
+
+ EXPECT_EQ(ArrowArrayFinishBuildingDefault(&array, nullptr), NANOARROW_OK);
+ EXPECT_EQ(array.length, 2);
+ EXPECT_EQ(array.null_count, 1);
+
+ auto data_buffer = reinterpret_cast<const uint8_t*>(array.buffers[1]);
+
+ EXPECT_EQ(memcmp(data_buffer, &interval.months, 4), 0);
+ EXPECT_EQ(memcmp(data_buffer + sizeof(interval.months), &interval.days, 4),
0);
+ EXPECT_EQ(memcmp(data_buffer + sizeof(interval.months) +
sizeof(interval.days),
+ &interval.ns, 8),
+ 0);
+
+ auto arrow_array = ImportArray(&array, month_day_nano_interval());
+ ARROW_EXPECT_OK(arrow_array);
+
+ auto builder = MonthDayNanoIntervalBuilder();
+ MonthDayNanoIntervalType::MonthDayNanos mdn = {2, 12, 42};
+ ARROW_EXPECT_OK(builder.Append(mdn));
+ ARROW_EXPECT_OK(builder.AppendNulls(1));
+ auto expected_array = builder.Finish();
+
+ EXPECT_TRUE(arrow_array.ValueUnsafe()->Equals(expected_array.ValueUnsafe()));
+}
+
TEST(ArrayTest, ArrayTestAppendToDecimal128Array) {
struct ArrowArray array;
struct ArrowDecimal decimal;
diff --git a/src/nanoarrow/nanoarrow.h b/src/nanoarrow/nanoarrow.h
index 413a2c0..1328964 100644
--- a/src/nanoarrow/nanoarrow.h
+++ b/src/nanoarrow/nanoarrow.h
@@ -870,6 +870,13 @@ static inline ArrowErrorCode ArrowArrayAppendBytes(struct
ArrowArray* array,
static inline ArrowErrorCode ArrowArrayAppendString(struct ArrowArray* array,
struct ArrowStringView
value);
+/// \brief Append a Interval to an array
+///
+/// Returns NANOARROW_OK if value can be exactly represented by
+/// the underlying storage type or EINVAL otherwise.
+static inline ArrowErrorCode ArrowArrayAppendInterval(struct ArrowArray* array,
+ struct ArrowInterval*
value);
+
/// \brief Append a decimal value to an array
///
/// Returns NANOARROW_OK if array is a decimal array with the appropriate
diff --git a/src/nanoarrow/nanoarrow_types.h b/src/nanoarrow/nanoarrow_types.h
index 2408a52..5a373f1 100644
--- a/src/nanoarrow/nanoarrow_types.h
+++ b/src/nanoarrow/nanoarrow_types.h
@@ -656,6 +656,29 @@ struct ArrowArrayPrivateData {
int8_t union_type_id_is_child_index;
};
+/// \brief A representation of an interval.
+/// \ingroup nanoarrow-utils
+struct ArrowInterval {
+ /// \brief The type of interval being used
+ enum ArrowType type;
+ /// \brief The number of months represented by the interval
+ int32_t months;
+ /// \brief The number of days represented by the interval
+ int32_t days;
+ /// \brief The number of ms represented by the interval
+ int32_t ms;
+ /// \brief The number of ns represented by the interval
+ int64_t ns;
+};
+
+/// \brief Zero initialize an Interval with a given unit
+/// \ingroup nanoarrow-utils
+static inline void ArrowIntervalInit(struct ArrowInterval* interval,
+ enum ArrowType type) {
+ memset(interval, 0, sizeof(struct ArrowInterval));
+ interval->type = type;
+}
+
/// \brief A representation of a fixed-precision decimal number
/// \ingroup nanoarrow-utils
///