diff --git a/velox/vector/arrow/Bridge.cpp b/velox/vector/arrow/Bridge.cpp index dd563315131af..cf5b07a815c2f 100644 --- a/velox/vector/arrow/Bridge.cpp +++ b/velox/vector/arrow/Bridge.cpp @@ -1396,9 +1396,12 @@ VectorPtr createTimestampVector( int64_t nullCount) { BufferPtr timestamps = AlignedBuffer::allocate(length, pool); auto* rawTimestamps = timestamps->asMutable(); - const auto* rawNulls = nulls->as(); - - if (length > nullCount) { + if (nulls == nullptr) { + for (size_t i = 0; i < length; ++i) { + rawTimestamps[i] = Timestamp::fromNanos(input[i]); + } + } else if (length > nullCount) { + const auto* rawNulls = nulls->as(); for (size_t i = 0; i < length; ++i) { if (!bits::isBitNull(rawNulls, i)) { rawTimestamps[i] = Timestamp::fromNanos(input[i]); diff --git a/velox/vector/arrow/tests/ArrowBridgeArrayTest.cpp b/velox/vector/arrow/tests/ArrowBridgeArrayTest.cpp index afde498c82c2d..74c996533bccf 100644 --- a/velox/vector/arrow/tests/ArrowBridgeArrayTest.cpp +++ b/velox/vector/arrow/tests/ArrowBridgeArrayTest.cpp @@ -1135,14 +1135,16 @@ class ArrowBridgeArrayImportTest : public ArrowBridgeArrayExportTest { "ttn", {Timestamp(0, 0), std::nullopt, Timestamp(1699308257, 1234)}); } - void testImportWithoutNullsBuffer() { - std::vector> inputValues = {1, 2, 3, 4, 5}; + template + void testImportWithoutNullsBuffer( + std::vector> inputValues, + const char* format) { auto length = inputValues.size(); // Construct Arrow array without nulls value and nulls buffer ArrowContextHolder holder1; - holder1.values = AlignedBuffer::allocate(length, pool_.get()); - auto rawValues1 = holder1.values->asMutable(); + holder1.values = AlignedBuffer::allocate(length, pool_.get()); + auto rawValues1 = holder1.values->asMutable(); for (size_t i = 0; i < length; ++i) { rawValues1[i] = *inputValues[i]; } @@ -1150,10 +1152,14 @@ class ArrowBridgeArrayImportTest : public ArrowBridgeArrayExportTest { holder1.buffers[0] = nullptr; holder1.buffers[1] = (const void*)rawValues1; auto arrowArray1 = makeArrowArray(holder1.buffers, 2, length, 0); - auto arrowSchema1 = makeArrowSchema("l"); + auto arrowSchema1 = makeArrowSchema(format); auto output = importFromArrow(arrowSchema1, arrowArray1, pool_.get()); - assertVectorContent(inputValues, output, arrowArray1.null_count); + if constexpr (std::is_same_v) { + assertTimestampVectorContent(inputValues, output, arrowArray1.null_count); + } else { + assertVectorContent(inputValues, output, arrowArray1.null_count); + } // However, convert from an Arrow array without nulls buffer but non-zero // null count should fail @@ -1166,7 +1172,7 @@ class ArrowBridgeArrayImportTest : public ArrowBridgeArrayExportTest { holder2.buffers[0] = nullptr; holder2.buffers[1] = (const void*)rawValues2; - auto arrowSchema2 = makeArrowSchema("l"); + auto arrowSchema2 = makeArrowSchema(format); auto arrowArray2 = makeArrowArray(holder2.buffers, 2, length, 1); EXPECT_THROW( importFromArrow(arrowSchema2, arrowArray2, pool_.get()), @@ -1205,6 +1211,24 @@ class ArrowBridgeArrayImportTest : public ArrowBridgeArrayExportTest { } private: + // Creates timestamp from bigint and asserts the content of actual vector with + // the expected timestamp values. + void assertTimestampVectorContent( + const std::vector>& expectedValues, + const VectorPtr& actual, + size_t nullCount) { + std::vector> tsValues; + tsValues.reserve(expectedValues.size()); + for (const auto& value : expectedValues) { + if (value.has_value()) { + tsValues.emplace_back(Timestamp::fromNanos(value.value())); + } else { + tsValues.emplace_back(std::nullopt); + } + } + assertVectorContent(tsValues, actual, nullCount); + } + void testImportRowFull() { // Manually create a ROW type. ArrowSchema arrowSchema; @@ -1580,7 +1604,9 @@ TEST_F(ArrowBridgeArrayImportAsViewerTest, scalar) { } TEST_F(ArrowBridgeArrayImportAsViewerTest, without_nulls_buffer) { - testImportWithoutNullsBuffer(); + std::vector> inputValues = {1, 2, 3, 4, 5}; + testImportWithoutNullsBuffer(inputValues, "l"); + testImportWithoutNullsBuffer(inputValues, "ttn"); } TEST_F(ArrowBridgeArrayImportAsViewerTest, string) { @@ -1631,7 +1657,9 @@ TEST_F(ArrowBridgeArrayImportAsOwnerTest, scalar) { } TEST_F(ArrowBridgeArrayImportAsOwnerTest, without_nulls_buffer) { - testImportWithoutNullsBuffer(); + std::vector> inputValues = {1, 2, 3, 4, 5}; + testImportWithoutNullsBuffer(inputValues, "l"); + testImportWithoutNullsBuffer(inputValues, "ttn"); } TEST_F(ArrowBridgeArrayImportAsOwnerTest, string) {