Fix gatherFromTimestampBuffer

rui-mo · Mar 20, 2024 · bd73528 · bd73528
1 parent 5e07790
commit bd73528
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 8 deletions.
diff --git a/velox/vector/arrow/Bridge.cpp b/velox/vector/arrow/Bridge.cpp
@@ -353,19 +353,20 @@ void gatherFromTimestampBuffer(
     Buffer& out) {
   auto src = (*vec.values()).as<Timestamp>();
   auto dst = out.asMutable<int64_t>();
+  vector_size_t j = 0; // index into dst
   if (!vec.mayHaveNulls() || vec.getNullCount() == 0) {
     switch (unit) {
       case TimestampUnit::kSecond:
-        rows.apply([&](vector_size_t i) { dst[i] = src[i].getSeconds(); });
+        rows.apply([&](vector_size_t i) { dst[j++] = src[i].getSeconds(); });
         break;
       case TimestampUnit::kMilli:
-        rows.apply([&](vector_size_t i) { dst[i] = src[i].toMillis(); });
+        rows.apply([&](vector_size_t i) { dst[j++] = src[i].toMillis(); });
         break;
       case TimestampUnit::kMicro:
-        rows.apply([&](vector_size_t i) { dst[i] = src[i].toMicros(); });
+        rows.apply([&](vector_size_t i) { dst[j++] = src[i].toMicros(); });
         break;
       case TimestampUnit::kNano:
-        rows.apply([&](vector_size_t i) { dst[i] = src[i].toNanos(); });
+        rows.apply([&](vector_size_t i) { dst[j++] = src[i].toNanos(); });
         break;
       default:
         VELOX_UNREACHABLE();
@@ -376,29 +377,33 @@ void gatherFromTimestampBuffer(
     case TimestampUnit::kSecond:
       rows.apply([&](vector_size_t i) {
         if (!vec.isNullAt(i)) {
-          dst[i] = src[i].getSeconds();
+          dst[j] = src[i].getSeconds();
         }
+        j++;
       });
       break;
     case TimestampUnit::kMilli:
       rows.apply([&](vector_size_t i) {
         if (!vec.isNullAt(i)) {
-          dst[i] = src[i].toMillis();
+          dst[j] = src[i].toMillis();
         }
+        j++;
       });
       break;
     case TimestampUnit::kMicro:
       rows.apply([&](vector_size_t i) {
         if (!vec.isNullAt(i)) {
-          dst[i] = src[i].toMicros();
+          dst[j] = src[i].toMicros();
         };
+        j++;
       });
       break;
     case TimestampUnit::kNano:
       rows.apply([&](vector_size_t i) {
         if (!vec.isNullAt(i)) {
-          dst[i] = src[i].toNanos();
+          dst[j] = src[i].toNanos();
         }
+        j++;
       });
       break;
     default:

diff --git a/velox/vector/arrow/tests/ArrowBridgeArrayTest.cpp b/velox/vector/arrow/tests/ArrowBridgeArrayTest.cpp
@@ -869,6 +869,54 @@ TEST_F(ArrowBridgeArrayExportTest, mapNested) {
   EXPECT_EQ(values.Value(1), 0);
 }
 
+TEST_F(ArrowBridgeArrayExportTest, mapTimestamp) {
+  const auto keys =
+      vectorMaker_.flatVector<int32_t>(5, [](vector_size_t i) { return i; });
+  auto values = vectorMaker_.flatVector<Timestamp>(
+      5, [](vector_size_t row) { return Timestamp(row, row); });
+  const auto offsets = makeBuffer<vector_size_t>({1, 4, 2});
+  const auto sizes = makeBuffer<vector_size_t>({1, 1, 1});
+  const auto type = MAP(INTEGER(), TIMESTAMP());
+  auto vec = std::make_shared<MapVector>(
+      pool_.get(), type, nullptr, 3, offsets, sizes, keys, values);
+  auto array = toArrow(vec, options_, pool_.get());
+  ASSERT_OK(array->ValidateFull());
+  EXPECT_EQ(array->null_count(), 0);
+  ASSERT_EQ(
+      *array->type(),
+      *arrow::map(arrow::int32(), arrow::timestamp(arrow::TimeUnit::NANO)));
+  {
+    auto& mapArray = static_cast<const arrow::MapArray&>(*array);
+    auto& values = static_cast<const arrow::TimestampArray&>(*mapArray.items());
+    ASSERT_EQ(values.length(), 3);
+    EXPECT_EQ(values.Value(0), 1'000'000'001L);
+    EXPECT_EQ(values.Value(1), 4'000'000'004L);
+    EXPECT_EQ(values.Value(2), 2'000'000'002L);
+  }
+
+  // Nullable timestamp vector.
+  values = vectorMaker_.flatVector<Timestamp>(
+      5,
+      [](vector_size_t row) { return Timestamp(row, row); },
+      [](vector_size_t row) { return row % 2 == 1; });
+  vec = std::make_shared<MapVector>(
+      pool_.get(), type, nullptr, 3, offsets, sizes, keys, values);
+  array = toArrow(vec, options_, pool_.get());
+  ASSERT_OK(array->ValidateFull());
+  EXPECT_EQ(array->null_count(), 0);
+  ASSERT_EQ(
+      *array->type(),
+      *arrow::map(arrow::int32(), arrow::timestamp(arrow::TimeUnit::NANO)));
+  {
+    auto& mapArray = static_cast<const arrow::MapArray&>(*array);
+    auto& values = static_cast<const arrow::TimestampArray&>(*mapArray.items());
+    ASSERT_EQ(values.length(), 3);
+    EXPECT_TRUE(values.IsNull(0));
+    EXPECT_EQ(values.Value(1), 4'000'000'004L);
+    EXPECT_EQ(values.Value(2), 2'000'000'002L);
+  }
+}
+
 TEST_F(ArrowBridgeArrayExportTest, dictionarySimple) {
   auto vec = BaseVector::wrapInDictionary(
       nullptr,