From ebd15abb17b333b586236d9a71abe87688b974df Mon Sep 17 00:00:00 2001 From: Bikramjeet Vig Date: Tue, 10 Dec 2024 16:32:59 -0800 Subject: [PATCH] fix: Ensure flat input for functions that consume a single field Ensure that Velox's guarantee that functions consuming a single field reference as input always receive a flat input holds, even when the debug query config disables peeling is turned on. Test Plan: Added a unit test --- velox/expression/Expr.cpp | 10 +++++++--- velox/expression/tests/ExprTest.cpp | 9 +++++++++ 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/velox/expression/Expr.cpp b/velox/expression/Expr.cpp index d4d0404cea1b..cf8da93b2c1c 100644 --- a/velox/expression/Expr.cpp +++ b/velox/expression/Expr.cpp @@ -1470,10 +1470,14 @@ bool Expr::applyFunctionWithPeeling( LocalDecodedVector localDecoded(context); LocalSelectivityVector newRowsHolder(context); if (!context.peelingEnabled()) { - if (inputValues_.size() == 1) { + if (distinctFields_.size() < 2) { // If we have a single input, velox needs to ensure that the - // vectorFunction would receive a flat input. - BaseVector::flattenVector(inputValues_[0]); + // vectorFunction would receive a flat or constant input. + for (int i = 0; i < inputValues_.size(); ++i) { + if (inputValues_[i]->encoding() == VectorEncoding::Simple::DICTIONARY) { + BaseVector::flattenVector(inputValues_[i]); + } + } applyFunction(applyRows, context, result); return true; } diff --git a/velox/expression/tests/ExprTest.cpp b/velox/expression/tests/ExprTest.cpp index 507c75853f5d..2636c95cf7f3 100644 --- a/velox/expression/tests/ExprTest.cpp +++ b/velox/expression/tests/ExprTest.cpp @@ -4861,6 +4861,15 @@ TEST_F(ExprTest, disablePeeling) { makeRowVector({flatInput}), {}, execCtx.get())); + + // Ensure functions that take a single column as input but can have more + // constant inputs also receive a flat vector. We use the in-predicate in this + // case which has a check for ensuring flat input. + ASSERT_NO_THROW(evaluateMultiple( + {"dict_wrap(c0) in (40, 42)"}, + makeRowVector({flatInput}), + {}, + execCtx.get())); } TEST_F(ExprTest, disableSharedSubExpressionReuse) {