Add ArrowArrayViewGetIntervalUnsafe (#270)

Follow on to #258. Closes #264.
apache · Aug 1, 2023 · 5c3e149 · 5c3e149
1 parent b0e90f8
commit 5c3e149
Show file tree

Hide file tree

Showing 2 changed files with 138 additions and 0 deletions.
diff --git a/src/nanoarrow/array_inline.h b/src/nanoarrow/array_inline.h
@@ -928,6 +928,33 @@ static inline struct ArrowBufferView ArrowArrayViewGetBytesUnsafe(
   return view;
 }
 
+static inline void ArrowArrayViewGetIntervalUnsafe(struct ArrowArrayView* array_view,
+                                                   int64_t i, struct ArrowInterval* out) {
+  const uint8_t* data_view = array_view->buffer_views[1].data.as_uint8;
+  switch (array_view->storage_type) {
+    case NANOARROW_TYPE_INTERVAL_MONTHS: {
+      const size_t size = sizeof(int32_t);
+      memcpy(&out->months, data_view + i * size, sizeof(int32_t));
+      break;
+    }
+    case NANOARROW_TYPE_INTERVAL_DAY_TIME: {
+      const size_t size = sizeof(int32_t) + sizeof(int32_t);
+      memcpy(&out->days, data_view + i * size, sizeof(int32_t));
+      memcpy(&out->ms, data_view + i * size + 4, sizeof(int32_t));
+      break;
+    }
+    case NANOARROW_TYPE_INTERVAL_MONTH_DAY_NANO: {
+      const size_t size = sizeof(int32_t) + sizeof(int32_t) + sizeof(int64_t);
+      memcpy(&out->months, data_view + i * size, sizeof(int32_t));
+      memcpy(&out->days, data_view + i * size + 4, sizeof(int32_t));
+      memcpy(&out->ns, data_view + i * size + 8, sizeof(int64_t));
+      break;
+    }
+    default:
+      break;
+  }
+}
+
 static inline void ArrowArrayViewGetDecimalUnsafe(struct ArrowArrayView* array_view,
                                                   int64_t i, struct ArrowDecimal* out) {
   i += array_view->offset;

diff --git a/src/nanoarrow/array_test.cc b/src/nanoarrow/array_test.cc
@@ -2453,6 +2453,117 @@ TEST(ArrayViewTest, ArrayViewTestGetString) {
   TestGetFromBinary<FixedSizeBinaryBuilder>(fixed_size_builder);
 }
 
+TEST(ArrayViewTest, ArrayViewTestGetIntervalYearMonth) {
+  struct ArrowArray array;
+  struct ArrowSchema schema;
+  struct ArrowArrayView array_view;
+  struct ArrowError error;
+  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, 2), NANOARROW_OK);
+
+  interval.months = -42;
+  EXPECT_EQ(ArrowArrayAppendInterval(&array, &interval), NANOARROW_OK);
+  EXPECT_EQ(ArrowArrayFinishBuildingDefault(&array, nullptr), NANOARROW_OK);
+
+  ASSERT_EQ(ArrowSchemaInitFromType(&schema, NANOARROW_TYPE_INTERVAL_MONTHS),
+            NANOARROW_OK);
+  ASSERT_EQ(ArrowArrayViewInitFromSchema(&array_view, &schema, &error), NANOARROW_OK);
+  ASSERT_EQ(ArrowArrayViewSetArray(&array_view, &array, &error), NANOARROW_OK);
+  EXPECT_EQ(ArrowArrayViewValidate(&array_view, NANOARROW_VALIDATION_LEVEL_FULL, &error),
+            NANOARROW_OK);
+
+  ArrowArrayViewGetIntervalUnsafe(&array_view, 0, &interval);
+  EXPECT_EQ(interval.months, 42);
+
+  ArrowArrayViewGetIntervalUnsafe(&array_view, 3, &interval);
+  EXPECT_EQ(interval.months, -42);
+
+  ArrowArrayViewReset(&array_view);
+  schema.release(&schema);
+  array.release(&array);
+}
+
+TEST(ArrayViewTest, ArrayViewTestGetIntervalDayTime) {
+  struct ArrowArray array;
+  struct ArrowSchema schema;
+  struct ArrowArrayView array_view;
+  struct ArrowError error;
+
+  auto builder = DayTimeIntervalBuilder();
+  ARROW_EXPECT_OK(builder.Append(DayTimeIntervalType::DayMilliseconds{42, 42}));
+  ARROW_EXPECT_OK(builder.AppendNulls(2));
+  ARROW_EXPECT_OK(builder.Append(DayTimeIntervalType::DayMilliseconds{-42, -42}));
+  auto maybe_arrow_array = builder.Finish();
+  ARROW_EXPECT_OK(maybe_arrow_array);
+  auto arrow_array = maybe_arrow_array.ValueUnsafe();
+
+  ARROW_EXPECT_OK(ExportArray(*arrow_array, &array, &schema));
+  ASSERT_EQ(ArrowArrayViewInitFromSchema(&array_view, &schema, &error), NANOARROW_OK);
+  ASSERT_EQ(ArrowArrayViewSetArray(&array_view, &array, &error), NANOARROW_OK);
+  EXPECT_EQ(ArrowArrayViewValidate(&array_view, NANOARROW_VALIDATION_LEVEL_FULL, &error),
+            NANOARROW_OK);
+
+  ArrowInterval interval;
+  ArrowIntervalInit(&interval, NANOARROW_TYPE_INTERVAL_DAY_TIME);
+
+  ArrowArrayViewGetIntervalUnsafe(&array_view, 0, &interval);
+  EXPECT_EQ(interval.days, 42);
+  EXPECT_EQ(interval.ms, 42);
+
+  ArrowArrayViewGetIntervalUnsafe(&array_view, 3, &interval);
+  EXPECT_EQ(interval.days, -42);
+  EXPECT_EQ(interval.ms, -42);
+
+  ArrowArrayViewReset(&array_view);
+  schema.release(&schema);
+  array.release(&array);
+}
+
+TEST(ArrayViewTest, ArrayViewTestGetIntervalMonthDayNano) {
+  struct ArrowArray array;
+  struct ArrowSchema schema;
+  struct ArrowArrayView array_view;
+  struct ArrowError error;
+
+  auto builder = MonthDayNanoIntervalBuilder();
+  ARROW_EXPECT_OK(builder.Append(MonthDayNanoIntervalType::MonthDayNanos{5, 12, 42}));
+  ARROW_EXPECT_OK(builder.AppendNulls(2));
+  ARROW_EXPECT_OK(builder.Append(MonthDayNanoIntervalType::MonthDayNanos{-5, -12, -42}));
+  auto maybe_arrow_array = builder.Finish();
+  ARROW_EXPECT_OK(maybe_arrow_array);
+  auto arrow_array = maybe_arrow_array.ValueUnsafe();
+
+  ARROW_EXPECT_OK(ExportArray(*arrow_array, &array, &schema));
+  ASSERT_EQ(ArrowArrayViewInitFromSchema(&array_view, &schema, &error), NANOARROW_OK);
+  ASSERT_EQ(ArrowArrayViewSetArray(&array_view, &array, &error), NANOARROW_OK);
+  EXPECT_EQ(ArrowArrayViewValidate(&array_view, NANOARROW_VALIDATION_LEVEL_FULL, &error),
+            NANOARROW_OK);
+
+  ArrowInterval interval;
+  ArrowIntervalInit(&interval, NANOARROW_TYPE_INTERVAL_MONTH_DAY_NANO);
+
+  ArrowArrayViewGetIntervalUnsafe(&array_view, 0, &interval);
+  EXPECT_EQ(interval.months, 5);
+  EXPECT_EQ(interval.days, 12);
+  EXPECT_EQ(interval.ns, 42);
+
+  ArrowArrayViewGetIntervalUnsafe(&array_view, 3, &interval);
+  EXPECT_EQ(interval.months, -5);
+  EXPECT_EQ(interval.days, -12);
+  EXPECT_EQ(interval.ns, -42);
+
+  ArrowArrayViewReset(&array_view);
+  schema.release(&schema);
+  array.release(&array);
+}
+
 TEST(ArrayViewTest, ArrayViewTestGetDecimal128) {
   struct ArrowArray array;
   struct ArrowSchema schema;