feat: Handle errors during constant folding in 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/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({}, {});
@@ -329,13 +332,26 @@ ExprPtr tryFoldIfConstant(const ExprPtr& expr, Scope* scope) {
     // 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&) {
+    // If folding an expression throws an exception, the user can choose whether
+    // to ignore it or replace the expression with a FailFunction by setting the
+    // query config constant_fold_error_with_fail_function to either false or
+    // true respectively. By default, such exceptions are ignored. 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& 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 +538,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()});