From 5a34db303e27d89a1d6a957f458082b5230ff7bc Mon Sep 17 00:00:00 2001
From: Wei He <weihe@meta.com>
Date: Fri, 2 Dec 2022 13:53:51 -0800
Subject: [PATCH] Fix wrapping of 0-length vector in element_at and subscript
 functions (#3284)

Summary:
Pull Request resolved: https://github.com/facebookincubator/velox/pull/3284

Expression fuzzer found a bug in element_at() and subscript() functions when the input vector has a 0-length element vector and there are unselected rows of this evaluation: https://github.com/facebookincubator/velox/issues/3216. In this situation, SubscriptImpl builds up an `indices` where unselected rows point to 0 (i.e., default value) and use it to wrap the 0-length element vector via BaseVector::wrapInDictionary. DictionaryVector<T>::setInternalState() checks that the index pointed to by `indices` is less than the length of the vector being wrapped, which are both 0 in this case. Hence the error happens.

This diff fixes this bug by directly creating a null vector when the input element vector has 0 length, instead of wrapping it in a dictionary.

Reviewed By: bikramSingh91

Differential Revision: D41357753

fbshipit-source-id: 71d217f73733cb12914879b27b0b61ead7c8bf34
---
 velox/functions/lib/SubscriptUtil.h           | 12 +++++
 .../prestosql/tests/ElementAtTest.cpp         | 50 +++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/velox/functions/lib/SubscriptUtil.h b/velox/functions/lib/SubscriptUtil.h
index af6f465a930a..020beedd37bf 100644
--- a/velox/functions/lib/SubscriptUtil.h
+++ b/velox/functions/lib/SubscriptUtil.h
@@ -167,6 +167,12 @@ class SubscriptImpl : public exec::VectorFunction {
     auto baseArray = decodedArray->base()->as<ArrayVector>();
     auto arrayIndices = decodedArray->indices();
 
+    // Subscript into empty arrays always returns NULLs.
+    if (baseArray->elements()->size() == 0) {
+      return BaseVector::createNullConstant(
+          baseArray->elements()->type(), rows.size(), context.pool());
+    }
+
     exec::LocalDecodedVector indexHolder(context, *indexArg, rows);
     auto decodedIndices = indexHolder.get();
 
@@ -295,6 +301,12 @@ class SubscriptImpl : public exec::VectorFunction {
     auto baseMap = decodedMap->base()->as<MapVector>();
     auto mapIndices = decodedMap->indices();
 
+    // Subscript into empty maps always returns NULLs.
+    if (baseMap->mapValues()->size() == 0) {
+      return BaseVector::createNullConstant(
+          baseMap->mapValues()->type(), rows.size(), context.pool());
+    }
+
     // Get map keys.
     auto mapKeys = baseMap->mapKeys();
     exec::LocalSelectivityVector allElementRows(context, mapKeys->size());
diff --git a/velox/functions/prestosql/tests/ElementAtTest.cpp b/velox/functions/prestosql/tests/ElementAtTest.cpp
index bb9342966748..42d7180cc699 100644
--- a/velox/functions/prestosql/tests/ElementAtTest.cpp
+++ b/velox/functions/prestosql/tests/ElementAtTest.cpp
@@ -15,6 +15,7 @@
  */
 
 #include <optional>
+#include "velox/expression/Expr.h"
 #include "velox/functions/lib/SubscriptUtil.h"
 #include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h"
 
@@ -543,3 +544,52 @@ TEST_F(ElementAtTest, errorStatesArray) {
       expectedValueAt,
       [](auto row) { return row == 40; });
 }
+
+TEST_F(ElementAtTest, emptyElementVector) {
+  auto keys = makeFlatVector<int64_t>({});
+  auto values = makeFlatVector<int64_t>({});
+
+  auto offsets = allocateOffsets(2, pool());
+  auto rawOffsets = offsets->asMutable<vector_size_t>();
+
+  auto sizes = allocateSizes(2, pool());
+  auto rawSizes = sizes->asMutable<vector_size_t>();
+
+  SelectivityVector rows(2);
+  rows.setValid(0, false);
+  rows.updateBounds();
+
+  VectorPtr result;
+
+  // Test map vector.
+  {
+    auto map = std::make_shared<MapVector>(
+        pool(),
+        MAP(BIGINT(), BIGINT()),
+        nullptr,
+        2,
+        offsets,
+        sizes,
+        keys,
+        values);
+    auto expected = makeNullConstant(TypeKind::BIGINT, 2);
+
+    evaluate<SimpleVector<int64_t>>(
+        "element_at(c0, 1)", makeRowVector({map}), rows, result);
+    test::assertEqualVectors(expected, result);
+    evaluate<SimpleVector<int64_t>>(
+        "c0[1]", makeRowVector({map}), rows, result);
+    test::assertEqualVectors(expected, result);
+  }
+
+  // Test array vector.
+  {
+    auto array = std::make_shared<ArrayVector>(
+        pool(), ARRAY(BIGINT()), nullptr, 2, offsets, sizes, values);
+    auto expected = makeNullConstant(TypeKind::BIGINT, 2);
+
+    evaluate<SimpleVector<int64_t>>(
+        "element_at(c0, 1)", makeRowVector({array}), rows, result);
+    test::assertEqualVectors(expected, result);
+  }
+}