Allow appending ArrayView to ArrayWriter through add_items() (#1494)

Summary:
Pull Request resolved: #1494

title

Reviewed By: kevinwilfong

Differential Revision: D35981026

fbshipit-source-id: dac1cdb6d24d07f977a8d148f5014be99879e5a0

velox/expression/ComplexWriterTypes.h

@@ -194,17 +194,24 @@ class ArrayWriter {
     return PrimitiveWriter<V>{elementsVector_, valuesOffset_ + length_ - 1};
   }
 
+  // Any vector type with std-like optional-free interface.
+  template <typename T>
+  void copy_from(const T& data) {
+    length_ = 0;
+    add_items(data);
+  }
+
   // Any vector type with std-like optional-free interface.
   template <typename VectorType>
-  void copy_from(const VectorType& data) {
+  void add_items(const VectorType& data) {
     if constexpr (provide_std_interface<V>) {
       // TODO: accelerate this with memcpy.
-      resize(data.size());
+      auto start = length_;
+      resize(length_ + data.size());
       for (auto i = 0; i < data.size(); i++) {
-        this->operator[](i) = data[i];
+        this->operator[](i + start) = data[i];
       }
     } else {
-      length_ = 0;
       for (const auto& item : data) {
         auto& writer = add_item();
         writer.copy_from(item);
@@ -213,31 +220,33 @@ class ArrayWriter {
   }
 
   // Copy from null-free ArrayView.
-  void copy_from(
+  void add_items(
       const typename VectorExec::template resolver<Array<V>>::null_free_in_type&
           arrayView) {
     // If the null buffer is allocated this will read every null bit.
     // TODO: create a copy version that avoids null checks (assumes not null)
     // even when null buffer is allocated.
 
-    // The copy_from above works for null-free ArrayView, but calling copy on
+    // The add_items above works for null-free ArrayView, but calling copy on
     // the vector directly uses memcpy which is more efficient.
-    resize(arrayView.size());
+    auto start = length_;
+    resize(length_ + arrayView.size());
     elementsVector_->copy(
         arrayView.elementsVector(),
-        valuesOffset_,
+        valuesOffset_ + start,
         arrayView.offset(),
         arrayView.size());
   }
 
   // Copy from nullable ArrayView.
-  void copy_from(
+  void add_items(
       const typename VectorExec::template resolver<Array<V>>::in_type&
           arrayView) {
-    resize(arrayView.size());
+    auto start = length_;
+    resize(length_ + arrayView.size());
     elementsVector_->copy(
         arrayView.elementsVector(),
-        valuesOffset_,
+        valuesOffset_ + start,
         arrayView.offset(),
         arrayView.size());
   }

velox/expression/tests/ArrayWriterTest.cpp

@@ -710,5 +710,59 @@ TEST_F(ArrayWriterTest, copyFromNestedNullableArrayView) {
           {{1}}, {{3, std::nullopt, 5}}}));
 }
 
+template <typename T>
+struct AddItemsTestFunc {
+  template <typename TOut, typename TIn>
+  void call(TOut& out, const TIn& input) {
+    out.add_items(input);
+    out.add_items(input);
+    out.add_items(std::vector<int64_t>{1, 2, 3});
+  }
+
+  // Will be called when there is no nulls in the input.
+  template <typename TOut, typename TIn>
+  void callNullFree(TOut& out, const TIn& input) {
+    out.add_items(std::vector<int64_t>{1, 2, 3});
+    out.add_items(input);
+    out.add_items(input);
+  }
+};
+
+TEST_F(ArrayWriterTest, addItems) {
+  registerFunction<AddItemsTestFunc, Array<int64_t>, Array<int64_t>>(
+      {"add_items_test"});
+  DecodedVector decoded;
+  SelectivityVector rows(1);
+
+  {
+    // callNullFree path.
+    auto result = evaluate(
+        "add_items_test(array_constructor(10, 20))",
+        makeRowVector({makeFlatVector<int64_t>(1)}));
+
+    // Test results.
+    decoded.decode(*result, rows);
+    exec::VectorReader<Array<int64_t>> reader(&decoded);
+    ASSERT_EQ(
+        reader.readNullFree(0).materialize(),
+        (std::vector<int64_t>{1, 2, 3, 10, 20, 10, 20}));
+  }
+
+  {
+    // call path.
+    auto result = evaluate(
+        "add_items_test(array_constructor(10, null))",
+        makeRowVector({makeFlatVector<int64_t>(1)}));
+
+    // Test results.
+    decoded.decode(*result, rows);
+    exec::VectorReader<Array<int64_t>> reader(&decoded);
+    ASSERT_EQ(
+        reader[0].materialize(),
+        (std::vector<std::optional<int64_t>>{
+            10, std::nullopt, 10, std::nullopt, 1, 2, 3}));
+  }
+}
+
 } // namespace
 } // namespace facebook::velox