From bd73528b2f9fb7bd8080a7b41ad395f104d546e6 Mon Sep 17 00:00:00 2001
From: rui-mo <rui.mo@intel.com>
Date: Wed, 20 Mar 2024 15:51:25 +0800
Subject: [PATCH] Fix gatherFromTimestampBuffer

---
 velox/vector/arrow/Bridge.cpp                 | 21 ++++----
 .../arrow/tests/ArrowBridgeArrayTest.cpp      | 48 +++++++++++++++++++
 2 files changed, 61 insertions(+), 8 deletions(-)

diff --git a/velox/vector/arrow/Bridge.cpp b/velox/vector/arrow/Bridge.cpp
index 04b18dc976d0..dc4a400134c9 100644
--- 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
index 6fb8af3fb339..e93becebc419 100644
--- 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,