feat: Capture constant folding errors during expression compilation

velox/core/QueryConfig.h

@@ -505,6 +505,11 @@ class QueryConfig {
   static constexpr const char* kShuffleCompressionKind =
       "shuffle_compression_codec";
 
+  /// Specifies whether the expression compiler should ignore errors encountered
+  /// during constant folding or return a FailFunction instead.
+  static constexpr const char* kConstantFoldErrorWithFailFunction =
+      "constant_fold_error_with_fail_function";
+
   bool selectiveNimbleReaderEnabled() const {
     return get<bool>(kSelectiveNimbleReaderEnabled, false);
   }
@@ -930,6 +935,10 @@ class QueryConfig {
     return get<std::string>(kShuffleCompressionKind, "none");
   }
 
+  bool constantFoldErrorWithFailFunction() const {
+    return get<bool>(kConstantFoldErrorWithFailFunction, false);
+  }
+
   template <typename T>
   T get(const std::string& key, const T& defaultValue) const {
     return config_->get<T>(key, defaultValue);

velox/expression/Expr.h

@@ -702,6 +702,23 @@ using ExprPtr = std::shared_ptr<Expr>;
 // this ExprSet. If however such an association cannot be made with certainty,
 // then its advisable to pre-load all lazy vectors to avoid issues associated
 // with partial loading.
+//
+// The input typed expressions are compiled in the constructor to exec::Expr
+// using the helper function compileExpressions. Constant folding is enabled by
+// default during expression compilation and it can be disabled by setting the
+// argument enableConstantFolding to false. Constant folding has a subtle
+// gotcha: if folding a constant expression deterministically throws, we can't
+// throw at expression compilation time yet because we can't guarantee that this
+// expression would actually need to be evaluated.
+//
+// If folding an expression throws an exception, the user can choose whether to
+// ignore it or receive a FailFunction in place of the expression by setting the
+// query config constant_fold_error_with_fail_function to either false or true
+// respectively. By default, we just ignore such exceptions and leave the
+// expression as-is. If this expression is hit at execution time and needs to be
+// evaluated, it will throw and fail the query anyway. If not, in case this
+// expression is never hit at execution time (for instance, if other arguments
+// are all null in a function with default null behavior), the query won't fail.
 class ExprSet {
  public:
   explicit ExprSet(

velox/expression/ExprCompiler.cpp

@@ -299,7 +299,10 @@ std::shared_ptr<Expr> compileLambda(
       config.exprTrackCpuUsage());
 }
 
-ExprPtr tryFoldIfConstant(const ExprPtr& expr, Scope* scope) {
+ExprPtr tryFoldIfConstant(
+    const ExprPtr& expr,
+    Scope* scope,
+    bool constantFoldErrorWithFailFunction) {
   if (expr->isConstant() && scope->exprSet->execCtx()) {
     try {
       auto rowType = ROW({}, {});
@@ -323,19 +326,18 @@ ExprPtr tryFoldIfConstant(const ExprPtr& expr, Scope* scope) {
         resultExpr->setDefaultNullRowsSkipped(true);
       }
       return resultExpr;
-    }
-    // Constant folding has a subtle gotcha: if folding a constant expression
-    // deterministically throws, we can't throw at expression compilation time
-    // yet because we can't guarantee that this expression would actually need
-    // to be evaluated.
-    //
-    // So, here, if folding an expression throws an exception, we just ignore it
-    // and leave the expression as-is. If this expression is hit at execution
-    // time and needs to be evaluated, it will throw and fail the query anyway.
-    // If not, in case this expression is never hit at execution time (for
-    // instance, if other arguments are all null in a function with default null
-    // behavior), the query won't fail.
-    catch (const VeloxUserError&) {
+    } catch (const VeloxUserError& error) {
+      if (constantFoldErrorWithFailFunction) {
+        auto failTypedExpr = std::make_shared<core::CallTypedExpr>(
+            UNKNOWN(),
+            std::vector<core::TypedExprPtr>{
+                std::make_shared<core::ConstantTypedExpr>(
+                    VARCHAR(), error.message())},
+            "fail");
+        auto failExprSet =
+            ExprSet({failTypedExpr}, scope->exprSet->execCtx(), false);
+        return failExprSet.exprs().front();
+      }
     }
   }
   return expr;
@@ -522,7 +524,8 @@ ExprPtr compileRewrittenExpression(
 
   // If the expression is constant folding it is redundant.
   auto folded = enableConstantFolding && !isConstantExpr
-      ? tryFoldIfConstant(result, scope)
+      ? tryFoldIfConstant(
+            result, scope, config.constantFoldErrorWithFailFunction())
       : result;
   scope->visited[expr.get()] = folded;
   return folded;

velox/expression/tests/ExprCompilerTest.cpp

@@ -123,6 +123,28 @@ TEST_F(ExprCompilerTest, constantFolding) {
       "plus(plus(a, 1:BIGINT), 5:BIGINT)", compile(expression)->toString());
 }
 
+TEST_F(ExprCompilerTest, constantFoldingError) {
+  // When query config kConstantFoldErrorWithFailFunction is false, the compiled
+  // expression should be 'divide(0, 0)', since the 'divide by zero' exception
+  // seen during constant folding is ignored by the expression compiler.
+  queryCtx_->testingOverrideConfigUnsafe({
+      {core::QueryConfig::kConstantFoldErrorWithFailFunction, "false"},
+  });
+  auto exprSet = compile(makeTypedExpr("0 / 0", ROW({}, {})));
+  auto compiled = exprSet->exprs().front();
+  ASSERT_EQ(compiled->name(), "divide");
+
+  // When query config kConstantFoldErrorWithFailFunction is true, the compiled
+  // expression should be 'fail(errorMessage)', with the 'divide by zero' error
+  // seen during constant folding as errorMessage.
+  queryCtx_->testingOverrideConfigUnsafe({
+      {core::QueryConfig::kConstantFoldErrorWithFailFunction, "true"},
+  });
+  exprSet = compile(makeTypedExpr("0 / 0", ROW({}, {})));
+  compiled = exprSet->exprs().front();
+  ASSERT_EQ(compiled->name(), "fail");
+}
+
 TEST_F(ExprCompilerTest, andFlattening) {
   auto rowType =
       ROW({"a", "b", "c", "d"}, {BOOLEAN(), BOOLEAN(), BOOLEAN(), BOOLEAN()});