diff --git a/velox/expression/ComplexWriterTypes.h b/velox/expression/ComplexWriterTypes.h index 328731919a9e..a39cfe892452 100644 --- a/velox/expression/ComplexWriterTypes.h +++ b/velox/expression/ComplexWriterTypes.h @@ -137,6 +137,42 @@ class ArrayWriter { } } + void push_back(const OptionalAccessor& item) { + if (!item.has_value()) { + add_null(); + } else { + if constexpr (provide_std_interface) { + push_back(item.value()); + } else { + add_item().copy_from(item.value()); + } + } + } + + void push_back(const std::optional& inputView) { + if (inputView.has_value()) { + push_back(*inputView); + } else { + add_null(); + } + } + + void push_back(const GenericView& item) { + add_item().copy_from(item); + } + + void push_back(const StringView& item) { + add_item().copy_from(item); + } + + void push_back(const std::optional& item) { + if (item.has_value()) { + push_back(*item); + } else { + add_null(); + } + } + // Add a new null item to the array. void add_null() { resize(length_ + 1); diff --git a/velox/expression/tests/ArrayWriterTest.cpp b/velox/expression/tests/ArrayWriterTest.cpp index d1fc91918afe..05115c855c23 100644 --- a/velox/expression/tests/ArrayWriterTest.cpp +++ b/velox/expression/tests/ArrayWriterTest.cpp @@ -316,12 +316,16 @@ TEST_F(ArrayWriterTest, testVarChar) { stringWriter, "test a long string, a bit longer than that, longer, and longer"); } + + arrayWriter.push_back("new_feature"_sv); + vectorWriter.commit(); auto expected = std::vector>>{ {"hi"_sv, std::nullopt, "welcome"_sv, - "test a long string, a bit longer than that, longer, and longer"_sv}}; + "test a long string, a bit longer than that, longer, and longer"_sv, + "new_feature"_sv}}; assertEqualVectors(result, makeNullableArrayVector(expected)); } @@ -853,7 +857,6 @@ TEST_F(ArrayWriterTest, nestedArrayWriteThenCommitNull) { auto& nestedArray = current.add_item(); nestedArray.push_back(1); nestedArray.push_back(2); - writer.commitNull(); } diff --git a/velox/expression/tests/GenericWriterTest.cpp b/velox/expression/tests/GenericWriterTest.cpp index 5366044ed8a2..9f016d0be875 100644 --- a/velox/expression/tests/GenericWriterTest.cpp +++ b/velox/expression/tests/GenericWriterTest.cpp @@ -116,12 +116,7 @@ struct ArrayTrimFunction { const int64_t& n) { int64_t end = inputArray.size() - n; for (int i = 0; i < end; ++i) { - if (inputArray[i].has_value()) { - auto& newItem = out.add_item(); - newItem.copy_from(inputArray[i].value()); - } else { - out.add_null(); - } + out.push_back(inputArray[i]); } } }; diff --git a/velox/functions/prestosql/ArrayFunctions.h b/velox/functions/prestosql/ArrayFunctions.h index d06a9dd92676..cc0f3ab18d2b 100644 --- a/velox/functions/prestosql/ArrayFunctions.h +++ b/velox/functions/prestosql/ArrayFunctions.h @@ -276,10 +276,11 @@ struct CombinationsFunction { innerArray.add_null(); continue; } + if constexpr (std::is_same_v) { innerArray.add_item().setNoCopy(array[idx].value()); } else { - innerArray.add_item() = array[idx].value(); + innerArray.push_back(array[idx].value()); } } } @@ -573,24 +574,13 @@ struct ArrayConcatFunction { } } - template - void assignElement(B& out, const U& input) { - if constexpr (SimpleTypeTrait::isPrimitiveType) { - // Primitives we just assign. copy_from not defined for primitives. - out = input; - } else { - out.copy_from(input); - } - } - void call( out_type>& out, const arg_type>& array, const arg_type& element) { out.reserve(array.size() + 1); out.add_items(array); - auto& newItem = out.add_item(); - assignElement(newItem, element); + out.push_back(element); } void call( @@ -598,8 +588,7 @@ struct ArrayConcatFunction { const arg_type& element, const arg_type>& array) { out.reserve(array.size() + 1); - auto& newItem = out.add_item(); - assignElement(newItem, element); + out.push_back(element); out.add_items(array); } }; @@ -628,12 +617,7 @@ struct ArrayTrimFunction { int64_t end = inputArray.size() - size; for (int i = 0; i < end; ++i) { - if (inputArray[i].has_value()) { - auto& newItem = out.add_item(); - newItem = inputArray[i].value(); - } else { - out.add_null(); - } + out.push_back(inputArray[i]); } } @@ -646,12 +630,7 @@ struct ArrayTrimFunction { int64_t end = inputArray.size() - size; for (int i = 0; i < end; ++i) { - if (inputArray[i].has_value()) { - auto& newItem = out.add_item(); - newItem.copy_from(inputArray[i].value()); - } else { - out.add_null(); - } + out.push_back(inputArray[i]); } } }; @@ -690,8 +669,7 @@ struct ArrayRemoveNullFunction { FOLLY_ALWAYS_INLINE void call(Out& out, const In& inputArray) { for (int i = 0; i < inputArray.size(); ++i) { if (inputArray[i].has_value()) { - auto& newItem = out.add_item(); - newItem = inputArray[i].value(); + out.push_back(inputArray[i].value()); } } } @@ -702,8 +680,7 @@ struct ArrayRemoveNullFunction { const arg_type>>& inputArray) { for (int i = 0; i < inputArray.size(); ++i) { if (inputArray[i].has_value()) { - auto& newItem = out.add_item(); - newItem.copy_from(inputArray[i].value()); + out.push_back(inputArray[i].value()); } } } @@ -766,7 +743,6 @@ template struct ArrayUnionFunction { VELOX_DEFINE_FUNCTION_TYPES(T) - // Fast path for primitives. template void call(Out& out, const In& inputArray1, const In& inputArray2) { folly::F14FastSet elementSet; @@ -775,31 +751,7 @@ struct ArrayUnionFunction { for (const auto& item : inputArray) { if (item.has_value()) { if (elementSet.insert(item.value()).second) { - auto& newItem = out.add_item(); - newItem = item.value(); - } - } else if (!nullAdded) { - nullAdded = true; - out.add_null(); - } - } - }; - addItems(inputArray1); - addItems(inputArray2); - } - - void call( - out_type>>& out, - const arg_type>>& inputArray1, - const arg_type>>& inputArray2) { - folly::F14FastSet elementSet; - bool nullAdded = false; - auto addItems = [&](auto& inputArray) { - for (const auto& item : inputArray) { - if (item.has_value()) { - if (elementSet.insert(item.value()).second) { - auto& newItem = out.add_item(); - newItem.copy_from(item.value()); + out.push_back(item.value()); } } else if (!nullAdded) { nullAdded = true; @@ -827,8 +779,7 @@ struct ArrayRemoveFunction { for (const auto& item : inputArray) { if (item.has_value()) { if (element != item.value()) { - auto& newItem = out.add_item(); - newItem = item.value(); + out.push_back(item.value()); } } else { out.add_null(); @@ -845,26 +796,22 @@ struct ArrayRemoveFunction { CompareFlags::NullHandlingMode::kNullAsIndeterminate); std::vector> toCopyItems; for (const auto& item : array) { - if (item.has_value()) { - auto result = element.compare(item.value(), kFlags); - VELOX_USER_CHECK( - result.has_value(), - "array_remove does not support arrays with elements that are null or contain null") - if (result.value()) { - toCopyItems.push_back(item.value()); - } - } else { + if (!item.has_value()) { toCopyItems.push_back(std::nullopt); + continue; + } + + auto result = element.compare(item.value(), kFlags); + VELOX_USER_CHECK( + result.has_value(), + "array_remove does not support arrays with elements that are null or contain null") + if (result.value()) { + toCopyItems.push_back(item.value()); } } for (const auto& item : toCopyItems) { - if (item.has_value()) { - auto& newItem = out.add_item(); - newItem.copy_from(item.value()); - } else { - out.add_null(); - } + out.push_back(item); } } }; @@ -884,8 +831,7 @@ struct ArrayRemoveFunctionString { if (item.has_value()) { auto result = element.compare(item.value()); if (result) { - auto& newItem = out.add_item(); - newItem.setNoCopy(item.value()); + out.push_back(item.value()); } } else { out.add_null(); diff --git a/velox/functions/prestosql/MultimapFromEntries.h b/velox/functions/prestosql/MultimapFromEntries.h index 17355f6202b1..a686a5083ca8 100644 --- a/velox/functions/prestosql/MultimapFromEntries.h +++ b/velox/functions/prestosql/MultimapFromEntries.h @@ -45,12 +45,9 @@ struct MultimapFromEntriesFunction { for (const auto& [key, values] : keyValuesMap_) { auto [keyWriter, valueWriter] = out.add_item(); keyWriter.copy_from(key); + for (const auto& value : values) { - if (value.has_value()) { - valueWriter.add_item().copy_from(value.value()); - } else { - valueWriter.add_null(); - } + valueWriter.push_back(value); } } }