feat: optimize array_agg to avoid copy

Summary:
existing fast path requires a single source and a single contiguous byte range to avoid a copy.
we can apply zero copy to a more sophisticated condition as long as for each group there is a contiguous byte range.

the reason why we see fragments of ranges in one input batch is that in the query plan we have filter and then native join which includes this streamingAgg operator so we fall into the branch where not all rows are selected

Reviewed By: Yuhta, xiaoxmeng

Differential Revision: D88194739

Author: Junjie Wang

velox/functions/prestosql/aggregates/ArrayAggAggregate.cpp

@@ -128,46 +128,94 @@ class ArrayAggAggregate : public exec::Aggregate {
     if (clusteredInput_) {
       bool singleSource{true};
       VectorPtr* currentSource{nullptr};
-      std::vector<BaseVector::CopyRange> ranges;
+      bool contiguousElementsPerGroup = true;
+      // check if we can get a contiguous range for each group
+      std::vector<BaseVector::CopyRange> rangeMap;
+      rangeMap.resize(numGroups, {0, 0, 0});
       for (int32_t i = 0; i < numGroups; ++i) {
         auto* accumulator = value<ClusteredAccumulator>(groups[i]);
-        resultOffsets[i] = arrayOffset;
-        vector_size_t arraySize = 0;
+        if (accumulator->sources.empty()) {
+          continue;
+        }
         for (auto& source : accumulator->sources) {
-          if (currentSource && currentSource->get() != source.vector.get()) {
-            elements->copyRanges(currentSource->get(), ranges);
-            ranges.clear();
-          }
-          if (!currentSource || currentSource->get() != source.vector.get()) {
-            singleSource = currentSource == nullptr;
+          if (!currentSource) {
             currentSource = &source.vector;
-            ranges.push_back({source.offset, arrayOffset, source.size});
-          } else if (
-              ranges.back().sourceIndex + ranges.back().count ==
-              source.offset) {
-            ranges.back().count += source.size;
           } else {
-            VELOX_DCHECK_LT(
-                ranges.back().sourceIndex + ranges.back().count, source.offset);
-            ranges.push_back({source.offset, arrayOffset, source.size});
+            singleSource =
+                singleSource && (currentSource->get() == source.vector.get());
           }
-          arrayOffset += source.size;
-          arraySize += source.size;
         }
-        resultSizes[i] = arraySize;
-        if (arraySize == 0) {
-          vector->setNull(i, true);
+        if (singleSource) {
+          vector_size_t offset = accumulator->sources[0].offset;
+          vector_size_t size = accumulator->sources[0].size;
+          for (auto i = 1; i < accumulator->sources.size(); ++i) {
+            if (accumulator->sources[i].offset !=
+                accumulator->sources[i - 1].offset +
+                    accumulator->sources[i - 1].size) {
+              contiguousElementsPerGroup = false;
+              break;
+            } else {
+              size += accumulator->sources[i].size;
+            }
+          }
+          if (contiguousElementsPerGroup) {
+            rangeMap[i] = {offset, 0, size};
+          }
         } else {
-          clearNull(rawNulls, i);
+          contiguousElementsPerGroup = false;
+          break;
+        }
+      }
+      // fallback to copy if we can't get a contiguous range for each group
+      std::vector<BaseVector::CopyRange> ranges;
+      if (currentSource != nullptr && !contiguousElementsPerGroup) {
+        currentSource = nullptr;
+        for (int32_t i = 0; i < numGroups; ++i) {
+          auto* accumulator = value<ClusteredAccumulator>(groups[i]);
+          resultOffsets[i] = arrayOffset;
+          vector_size_t arraySize = 0;
+          for (auto& source : accumulator->sources) {
+            if (currentSource && currentSource->get() != source.vector.get()) {
+              elements->copyRanges(currentSource->get(), ranges);
+              ranges.clear();
+            }
+            if (!currentSource || currentSource->get() != source.vector.get()) {
+              currentSource = &source.vector;
+              ranges.push_back({source.offset, arrayOffset, source.size});
+            } else if (
+                ranges.back().sourceIndex + ranges.back().count ==
+                source.offset) {
+              ranges.back().count += source.size;
+            } else {
+              VELOX_DCHECK_LT(
+                  ranges.back().sourceIndex + ranges.back().count,
+                  source.offset);
+              ranges.push_back({source.offset, arrayOffset, source.size});
+            }
+            arrayOffset += source.size;
+            arraySize += source.size;
+          }
+          resultSizes[i] = arraySize;
+          if (arraySize == 0) {
+            vector->setNull(i, true);
+          } else {
+            clearNull(rawNulls, i);
+          }
         }
       }
       if (currentSource != nullptr) {
-        VELOX_DCHECK(!ranges.empty());
-        if (singleSource && ranges.size() == 1) {
-          VELOX_CHECK_GE(currentSource->get()->size(), numElements);
-          VELOX_CHECK_EQ(ranges[0].count, numElements);
-          elements = currentSource->get()->slice(
-              ranges[0].sourceIndex, ranges[0].count);
+        VELOX_DCHECK(contiguousElementsPerGroup || !ranges.empty());
+        if (singleSource && contiguousElementsPerGroup) {
+          elements = BaseVector::loadedVectorShared(*currentSource);
+          for (int32_t i = 0; i < numGroups; ++i) {
+            if (rangeMap[i].count > 0) {
+              resultOffsets[i] = rangeMap[i].sourceIndex;
+              resultSizes[i] = rangeMap[i].count;
+              clearNull(rawNulls, i);
+            } else {
+              vector->setNull(i, true);
+            }
+          }
         } else {
           elements->copyRanges(currentSource->get(), ranges);
         }