add error handling

facebookincubator · Mar 4, 2023 · edae9dd · edae9dd
1 parent cada699
commit edae9dd
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 1 deletion.
diff --git a/velox/functions/prestosql/ArrayAllMatch.cpp b/velox/functions/prestosql/ArrayAllMatch.cpp
@@ -14,7 +14,9 @@
  * limitations under the License.
  */
 
+#include "velox/expression/EvalCtx.h"
 #include "velox/expression/Expr.h"
+#include "velox/expression/ScopedVarSetter.h"
 #include "velox/expression/VectorFunction.h"
 #include "velox/functions/lib/LambdaFunctionUtil.h"
 #include "velox/functions/lib/RowsTranslationUtil.h"
@@ -57,12 +59,17 @@ class AllMatchFunction : public exec::VectorFunction {
     auto flatResult = result->asFlatVector<bool>();
     exec::LocalDecodedVector bitsDecoder(context);
     auto it = args[1]->asUnchecked<FunctionVector>()->iterator(&rows);
+
+    // Turn on the ThrowOnError flag and temporarily reset the errors in context
+    // to nullptr, so we only handle errors here.
+    ScopedVarSetter throwOnError(context.mutableThrowOnError(), false);
+    ScopedVarSetter<exec::EvalCtx::ErrorVectorPtr> errorsSetter(
+        context.errorsPtr(), nullptr);
     while (auto entry = it.next()) {
       auto elementRows =
           toElementRows<ArrayVector>(numElements, *entry.rows, flatArray.get());
       auto wrapCapture = toWrapCapture<ArrayVector>(
           numElements, entry.callable, *entry.rows, flatArray);
-      bool hasVeloxUserError = false;
       entry.callable->apply(
           elementRows,
           finalSelection,
@@ -74,6 +81,12 @@ class AllMatchFunction : public exec::VectorFunction {
 
       bitsDecoder.get()->decode(*matchBits, elementRows);
       entry.rows->applyToSelected([&](vector_size_t row) {
+        auto errors = context.errors();
+        if (errors && row < errors->size() && !errors->isNullAt(row)) {
+          flatResult->set(row, false);
+          return;
+        }
+
         auto size = sizes[row];
         auto offset = offsets[row];
         auto allMatch = true;

diff --git a/velox/functions/prestosql/tests/ArrayAllMatchTest.cpp b/velox/functions/prestosql/tests/ArrayAllMatchTest.cpp
@@ -122,3 +122,14 @@ TEST_F(ArrayAllMatchTest, doubles) {
       makeNullableFlatVector<bool>({true, true, false, std::nullopt});
   assertEqualVectors(expectedResult, result);
 }
+
+TEST_F(ArrayAllMatchTest, errorOutputs) {
+  auto input =
+      makeNullableArrayVector<int8_t>({{2, 1}, {3, 0}, {std::nullopt}});
+  auto result = evaluate<SimpleVector<bool>>(
+      "all_match(c0, x -> ((10 / x) > 2))", makeRowVector({input}));
+
+  auto expectedResult =
+      makeNullableFlatVector<bool>({true, false, std::nullopt});
+  assertEqualVectors(expectedResult, result);
+}