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(); auto dst = out.asMutable(); + 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(5, [](vector_size_t i) { return i; }); + auto values = vectorMaker_.flatVector( + 5, [](vector_size_t row) { return Timestamp(row, row); }); + const auto offsets = makeBuffer({1, 4, 2}); + const auto sizes = makeBuffer({1, 1, 1}); + const auto type = MAP(INTEGER(), TIMESTAMP()); + auto vec = std::make_shared( + 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(*array); + auto& values = static_cast(*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( + 5, + [](vector_size_t row) { return Timestamp(row, row); }, + [](vector_size_t row) { return row % 2 == 1; }); + vec = std::make_shared( + 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(*array); + auto& values = static_cast(*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,