From 556851897a42a04ab2c1387a51d47a0310d94ba4 Mon Sep 17 00:00:00 2001
From: Masha Basmanova <mbasmanova@meta.com>
Date: Tue, 11 Feb 2025 11:01:21 -0800
Subject: [PATCH] feat: Add convenience API to evalute constant expression
 (#12306)

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

Useful for constant folding.

Reviewed By: Yuhta

Differential Revision: D69463896

fbshipit-source-id: 6ca71daea190c83cdaf816c3945c6624f1639325
---
 velox/expression/Expr.cpp           | 16 ++++++++++++++++
 velox/expression/Expr.h             |  6 ++++++
 velox/expression/tests/ExprTest.cpp | 13 +++++++++++++
 3 files changed, 35 insertions(+)

diff --git a/velox/expression/Expr.cpp b/velox/expression/Expr.cpp
index 4edf2810db3e..779da36937d9 100644
--- a/velox/expression/Expr.cpp
+++ b/velox/expression/Expr.cpp
@@ -2014,4 +2014,20 @@ core::ExecCtx* SimpleExpressionEvaluator::ensureExecCtx() {
   return execCtx_.get();
 }
 
+VectorPtr evaluateConstantExpression(
+    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);
+}
+
 } // namespace facebook::velox::exec
diff --git a/velox/expression/Expr.h b/velox/expression/Expr.h
index c0cdf0fb85e4..a2655ed6d442 100644
--- a/velox/expression/Expr.h
+++ b/velox/expression/Expr.h
@@ -830,6 +830,12 @@ 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.
+VectorPtr evaluateConstantExpression(
+    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
diff --git a/velox/expression/tests/ExprTest.cpp b/velox/expression/tests/ExprTest.cpp
index 2d0ff0ed5886..3784d7adb156 100644
--- a/velox/expression/tests/ExprTest.cpp
+++ b/velox/expression/tests/ExprTest.cpp
@@ -4970,5 +4970,18 @@ TEST_F(ExprTest, disabledeferredLazyLoading) {
       expressions, makeRowVector({c0, c1}), {}, execCtx.get());
 }
 
+TEST_F(ExprTest, evaluateConstantExpression) {
+  auto eval = [&](const std::string& sql) {
+    auto expr = parseExpression(sql, ROW({}));
+    return exec::evaluateConstantExpression(expr, pool());
+  };
+
+  assertEqualVectors(eval("1 + 2"), makeConstant<int64_t>(3, 1));
+
+  assertEqualVectors(
+      eval("transform(array[1, 2, 3], x -> (x * 2))"),
+      makeArrayVectorFromJson<int64_t>({"[2, 4, 6]"}));
+}
+
 } // namespace
 } // namespace facebook::velox::test