fix: Enable constant folding for lambda functions (facebookincubator#13642)

Summary:
Pull Request resolved: facebookincubator#13642

ExprCompiler used to not constant fold lambda functions. transform(array[1, 2, 3], x -> x * 2) would not be constant folded.

This was due to Expr::isConstant returning false for such an expression. Expr::isConstant returned true iff expression is deterministic and all input are ConstantExpr.

A fix is to modify Expr::isConstant to return true iff expression is deterministic and has no dependencies (distinctFields_ is empty).

Also, added convenience API tryEvaluateConstantExpression to compliment existing evaluateConstantExpression. The new API can be safely called on any expression without ensuring the expression is constant-foldable.

Reviewed By: bikramSingh91, xiaoxmeng

Differential Revision: D76004362

fbshipit-source-id: f3fc5dfcb8bcf93205fcc60b1709ddd09aa75ac3

Author: mbasmanova

velox/expression/Expr.cpp

@@ -1192,20 +1192,36 @@ void Expr::evalWithMemo(
   VectorPtr base;
   distinctFields_[0]->evalSpecialForm(rows, context, base);
 
+  // evalWithNulls may throw an exception. If this happens during constant
+  // folding, the exception is suppressed and the Expr object may be reused.
+  // Hence, it is important to update state in way that ensure "valid" state in
+  // case of exceptions.
+  //
+  // Also, note that the same expression running on same data may pass or may
+  // fail depending on whether it runs under TRY or not.
+  //
+  // An example expression that triggers these edge cases:
+  //
+  //    try(coalesce(array_min_by(array[1, 2, 3], x -> x / 0), 0::INTEGER))
+
   if (base.get() != baseOfDictionaryRawPtr_ ||
       baseOfDictionaryWeakPtr_.expired()) {
     baseOfDictionaryRepeats_ = 0;
-    baseOfDictionaryWeakPtr_ = base;
-    baseOfDictionaryRawPtr_ = base.get();
+    baseOfDictionaryWeakPtr_.reset();
+    baseOfDictionaryRawPtr_ = nullptr;
     context.releaseVector(baseOfDictionary_);
     context.releaseVector(dictionaryCache_);
+
     evalWithNulls(rows, context, result);
+    baseOfDictionaryWeakPtr_ = base;
+    baseOfDictionaryRawPtr_ = base.get();
     return;
   }
-  ++baseOfDictionaryRepeats_;
 
-  if (baseOfDictionaryRepeats_ == 1) {
+  if (baseOfDictionaryRepeats_ == 0) {
     evalWithNulls(rows, context, result);
+
+    ++baseOfDictionaryRepeats_;
     baseOfDictionary_ = base;
     dictionaryCache_ = result;
     if (!cachedDictionaryIndices_) {
@@ -1217,6 +1233,8 @@ void Expr::evalWithMemo(
     return;
   }
 
+  ++baseOfDictionaryRepeats_;
+
   if (cachedDictionaryIndices_) {
     LocalSelectivityVector cachedHolder(context, rows);
     auto cached = cachedHolder.get();
@@ -1242,31 +1260,34 @@ void Expr::evalWithMemo(
 
     evalWithNulls(*uncached, context, result);
     context.deselectErrors(*uncached);
-    context.exprSet()->addToMemo(this);
-    auto newCacheSize = uncached->end();
-
-    // dictionaryCache_ is valid only for cachedDictionaryIndices_. Hence, a
-    // safe call to BaseVector::ensureWritable must include all the rows not
-    // covered by cachedDictionaryIndices_. If BaseVector::ensureWritable is
-    // called only for a subset of rows not covered by
-    // cachedDictionaryIndices_, it will attempt to copy rows that are not
-    // valid leading to a crash.
-    LocalSelectivityVector allUncached(context, dictionaryCache_->size());
-    allUncached.get()->setAll();
-    allUncached.get()->deselect(*cachedDictionaryIndices_);
-    context.ensureWritable(*allUncached.get(), type(), dictionaryCache_);
-
-    if (cachedDictionaryIndices_->size() < newCacheSize) {
-      cachedDictionaryIndices_->resize(newCacheSize, false);
-    }
 
-    cachedDictionaryIndices_->select(*uncached);
+    if (uncached->hasSelections()) {
+      context.exprSet()->addToMemo(this);
+      auto newCacheSize = uncached->end();
+
+      // dictionaryCache_ is valid only for cachedDictionaryIndices_. Hence, a
+      // safe call to BaseVector::ensureWritable must include all the rows not
+      // covered by cachedDictionaryIndices_. If BaseVector::ensureWritable is
+      // called only for a subset of rows not covered by
+      // cachedDictionaryIndices_, it will attempt to copy rows that are not
+      // valid leading to a crash.
+      LocalSelectivityVector allUncached(context, dictionaryCache_->size());
+      allUncached.get()->setAll();
+      allUncached.get()->deselect(*cachedDictionaryIndices_);
+      context.ensureWritable(*allUncached.get(), type(), dictionaryCache_);
+
+      if (cachedDictionaryIndices_->size() < newCacheSize) {
+        cachedDictionaryIndices_->resize(newCacheSize, false);
+      }
 
-    // Resize the dictionaryCache_ to accommodate all the necessary rows.
-    if (dictionaryCache_->size() < uncached->end()) {
-      dictionaryCache_->resize(uncached->end());
+      cachedDictionaryIndices_->select(*uncached);
+
+      // Resize the dictionaryCache_ to accommodate all the necessary rows.
+      if (dictionaryCache_->size() < uncached->end()) {
+        dictionaryCache_->resize(uncached->end());
+      }
+      dictionaryCache_->copy(result.get(), *uncached, nullptr);
     }
-    dictionaryCache_->copy(result.get(), *uncached, nullptr);
   }
   context.releaseVector(base);
 }
@@ -1660,12 +1681,8 @@ bool Expr::isConstant() const {
   if (!isDeterministic()) {
     return false;
   }
-  for (auto& input : inputs_) {
-    if (!input->is<ConstantExpr>()) {
-      return false;
-    }
-  }
-  return true;
+
+  return distinctFields_.empty();
 }
 
 namespace {
@@ -2022,17 +2039,30 @@ core::ExecCtx* SimpleExpressionEvaluator::ensureExecCtx() {
 VectorPtr evaluateConstantExpression(
     const core::TypedExprPtr& expr,
     memory::MemoryPool* pool) {
+  auto result = tryEvaluateConstantExpression(expr, pool);
+  VELOX_USER_CHECK_NOT_NULL(
+      result, "Expression is not constant-foldable: {}", expr->toString());
+  return result;
+}
+
+VectorPtr tryEvaluateConstantExpression(
+    const core::TypedExprPtr& expr,
+    memory::MemoryPool* pool) {
   auto data = BaseVector::create<RowVector>(ROW({}), 1, pool);
 
   auto queryCtx = velox::core::QueryCtx::create();
   velox::core::ExecCtx execCtx{pool, queryCtx.get()};
   velox::exec::ExprSet exprSet({expr}, &execCtx);
-  velox::exec::EvalCtx evalCtx(&execCtx, &exprSet, data.get());
 
-  velox::SelectivityVector singleRow(1);
-  std::vector<velox::VectorPtr> results(1);
-  exprSet.eval(singleRow, evalCtx, results);
-  return results.at(0);
+  if (exprSet.expr(0)->is<ConstantExpr>()) {
+    velox::exec::EvalCtx evalCtx(&execCtx, &exprSet, data.get());
+    velox::SelectivityVector singleRow(1);
+    std::vector<velox::VectorPtr> results(1);
+    exprSet.eval(singleRow, evalCtx, results);
+    return results.at(0);
+  }
+
+  return nullptr;
 }
 
 } // namespace facebook::velox::exec

velox/expression/Expr.h

@@ -836,12 +836,20 @@ std::unique_ptr<ExprSet> makeExprSetFromFlag(
     std::vector<core::TypedExprPtr>&& source,
     core::ExecCtx* execCtx);
 
-/// Evaluates an expression that doesn't depend on any inputs and returns the
-/// result as single-row vector.
+/// Evaluates a deterministic expression that doesn't depend on any inputs and
+/// returns the result as single-row vector. Throws if expression is
+/// non-deterministic or has dependencies.
 VectorPtr evaluateConstantExpression(
     const core::TypedExprPtr& expr,
     memory::MemoryPool* pool);
 
+/// Evaluates a deterministic expression that doesn't depend on any inputs and
+/// returns the result as single-row vector. Returns nullptr if the expression
+/// is non-deterministic or has dependencies.
+VectorPtr tryEvaluateConstantExpression(
+    const core::TypedExprPtr& expr,
+    memory::MemoryPool* pool);
+
 /// Returns a string representation of the expression trees annotated with
 /// runtime statistics. Expected to be called after calling ExprSet::eval one or
 /// more times. If called before ExprSet::eval runtime statistics will be all

velox/expression/tests/ExprTest.cpp

@@ -4986,7 +4986,7 @@ TEST_F(ExprTest, disabledeferredLazyLoading) {
 
 TEST_F(ExprTest, evaluateConstantExpression) {
   auto eval = [&](const std::string& sql) {
-    auto expr = parseExpression(sql, ROW({}));
+    auto expr = parseExpression(sql, ROW({"a"}, {BIGINT()}));
     return exec::evaluateConstantExpression(expr, pool());
   };
 
@@ -4995,6 +4995,37 @@ TEST_F(ExprTest, evaluateConstantExpression) {
   assertEqualVectors(
       eval("transform(array[1, 2, 3], x -> (x * 2))"),
       makeArrayVectorFromJson<int64_t>({"[2, 4, 6]"}));
+
+  assertEqualVectors(
+      eval("transform(array[1, 2, 3], x -> (x * (3 - 1)))"),
+      makeArrayVectorFromJson<int64_t>({"[2, 4, 6]"}));
+
+  assertEqualVectors(
+      eval("transform(array[1, 2, 3], x -> 2)"),
+      makeArrayVectorFromJson<int64_t>({"[2, 2, 2]"}));
+
+  assertEqualVectors(
+      eval(
+          "try(coalesce(array_min_by(array[1, 2, 3], x -> x / 0), 0::INTEGER))"),
+      makeNullConstant(TypeKind::INTEGER, 1));
+
+  auto tryEval = [&](const std::string& sql) {
+    auto expr = parseExpression(sql, ROW({"a"}, {BIGINT()}));
+    return exec::tryEvaluateConstantExpression(expr, pool());
+  };
+
+  VELOX_ASSERT_THROW(eval("a + 1"), "Expression is not constant-foldable");
+  ASSERT_TRUE(tryEval("a + 1") == nullptr);
+
+  VELOX_ASSERT_THROW(
+      eval("rand() + 1.0"), "Expression is not constant-foldable");
+  ASSERT_TRUE(tryEval("rand() + 1.0") == nullptr);
+
+  VELOX_ASSERT_THROW(
+      eval("transform(array[1, 2, 3], x -> (x * 2) + a)"),
+      "Expression is not constant-foldable");
+  ASSERT_TRUE(
+      tryEval("transform(array[1, 2, 3], x -> (x * 2) + a)") == nullptr);
 }
 
 TEST_F(ExprTest, peelingOnDeterministicFunctionInNonDeterministicExpr) {