This is an automated email from the ASF dual-hosted git repository.
gian pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new 3fe79b56b84 Give control of fallback vectorization to each function.
(#18522)
3fe79b56b84 is described below
commit 3fe79b56b84cb6464bd03245d1faadf6fec63235
Author: Gian Merlino <[email protected]>
AuthorDate: Mon Sep 15 16:26:54 2025 -0700
Give control of fallback vectorization to each function. (#18522)
* Give control of fallback vectorization to each function.
If a function overrides canVectorize, let it decide when it does and
doesn't vectorize, rather than always attempting to fall back. In situations
like conditionals (#18507, #18512), we are currently refraining from
vectorizing when output types are not all aligned.
* Additional coverage.
---
.../org/apache/druid/math/expr/ApplyFunction.java | 14 +++++++-----
.../main/java/org/apache/druid/math/expr/Expr.java | 7 ------
.../org/apache/druid/math/expr/ExprMacroTable.java | 8 +++++--
.../java/org/apache/druid/math/expr/Function.java | 14 +++++++-----
.../org/apache/druid/math/expr/FunctionalExpr.java | 25 ++++------------------
.../math/expr/vector/FallbackVectorProcessor.java | 16 ++++++++++++++
.../math/expr/VectorExprResultConsistencyTest.java | 25 ++++++++++++++++++++++
7 files changed, 69 insertions(+), 40 deletions(-)
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/ApplyFunction.java
b/processing/src/main/java/org/apache/druid/math/expr/ApplyFunction.java
index 8302e1c901b..bd39110c76a 100644
--- a/processing/src/main/java/org/apache/druid/math/expr/ApplyFunction.java
+++ b/processing/src/main/java/org/apache/druid/math/expr/ApplyFunction.java
@@ -23,8 +23,8 @@ import com.google.common.collect.ImmutableSet;
import it.unimi.dsi.fastutil.objects.Object2IntArrayMap;
import it.unimi.dsi.fastutil.objects.Object2IntMap;
import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
-import org.apache.druid.java.util.common.UOE;
import org.apache.druid.math.expr.vector.ExprVectorProcessor;
+import org.apache.druid.math.expr.vector.FallbackVectorProcessor;
import javax.annotation.Nullable;
import java.util.ArrayList;
@@ -50,9 +50,9 @@ public interface ApplyFunction extends NamedFunction
* @see Expr#canVectorize(Expr.InputBindingInspector)
* @see Function#canVectorize(Expr.InputBindingInspector, List)
*/
- default boolean canVectorize(Expr.InputBindingInspector inspector, Expr
lambda, List<Expr> args)
+ default boolean canVectorize(Expr.InputBindingInspector inspector,
LambdaExpr lambda, List<Expr> args)
{
- return false;
+ return
FallbackVectorProcessor.canFallbackVectorize(getOutputType(inspector, lambda,
args), inspector, args);
}
/**
@@ -64,11 +64,15 @@ public interface ApplyFunction extends NamedFunction
*/
default <T> ExprVectorProcessor<T> asVectorProcessor(
Expr.VectorInputBindingInspector inspector,
- Expr lambda,
+ LambdaExpr lambda,
List<Expr> args
)
{
- throw new UOE("%s is not vectorized", name());
+ if (ExpressionProcessing.allowVectorizeFallback()) {
+ return FallbackVectorProcessor.create(this, lambda, args, inspector);
+ } else {
+ throw Exprs.cannotVectorize(name());
+ }
}
/**
diff --git a/processing/src/main/java/org/apache/druid/math/expr/Expr.java
b/processing/src/main/java/org/apache/druid/math/expr/Expr.java
index 4347ac7af07..08e05c06579 100644
--- a/processing/src/main/java/org/apache/druid/math/expr/Expr.java
+++ b/processing/src/main/java/org/apache/druid/math/expr/Expr.java
@@ -189,13 +189,6 @@ public interface Expr extends Cacheable
return false;
}
-
- default boolean canFallbackVectorize(InputBindingInspector inspector,
List<Expr> args)
- {
- return ExpressionProcessing.allowVectorizeFallback() &&
- getOutputType(inspector) != null &&
- inspector.canVectorize(args);
- }
/**
* Possibly convert the {@link Expr} into an optimized, possibly not
thread-safe {@link Expr}. Does not convert
* child {@link Expr}. Most callers should use {@link
Expr#singleThreaded(Expr, InputBindingInspector)} to convert
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/ExprMacroTable.java
b/processing/src/main/java/org/apache/druid/math/expr/ExprMacroTable.java
index 40d8ba1112c..5c885892efd 100644
--- a/processing/src/main/java/org/apache/druid/math/expr/ExprMacroTable.java
+++ b/processing/src/main/java/org/apache/druid/math/expr/ExprMacroTable.java
@@ -178,13 +178,17 @@ public class ExprMacroTable
@Override
public boolean canVectorize(InputBindingInspector inspector)
{
- return canFallbackVectorize(inspector, args);
+ return
FallbackVectorProcessor.canFallbackVectorize(getOutputType(inspector),
inspector, args);
}
@Override
public <T> ExprVectorProcessor<T>
asVectorProcessor(VectorInputBindingInspector inspector)
{
- return FallbackVectorProcessor.create(macro, args, inspector);
+ if (ExpressionProcessing.allowVectorizeFallback()) {
+ return FallbackVectorProcessor.create(macro, args, inspector);
+ } else {
+ throw Exprs.cannotVectorize(this);
+ }
}
/**
diff --git a/processing/src/main/java/org/apache/druid/math/expr/Function.java
b/processing/src/main/java/org/apache/druid/math/expr/Function.java
index 568d44171b5..8f51a50fd90 100644
--- a/processing/src/main/java/org/apache/druid/math/expr/Function.java
+++ b/processing/src/main/java/org/apache/druid/math/expr/Function.java
@@ -29,9 +29,9 @@ import org.apache.druid.error.DruidException;
import org.apache.druid.java.util.common.DateTimes;
import org.apache.druid.java.util.common.HumanReadableBytes;
import org.apache.druid.java.util.common.StringUtils;
-import org.apache.druid.java.util.common.UOE;
import org.apache.druid.math.expr.vector.CastToTypeVectorProcessor;
import org.apache.druid.math.expr.vector.ExprVectorProcessor;
+import org.apache.druid.math.expr.vector.FallbackVectorProcessor;
import org.apache.druid.math.expr.vector.VectorConditionalProcessors;
import org.apache.druid.math.expr.vector.VectorMathProcessors;
import org.apache.druid.math.expr.vector.VectorProcessors;
@@ -142,11 +142,11 @@ public interface Function extends NamedFunction
* batches to use with vectorized query engines.
*
* @see Expr#canVectorize(Expr.InputBindingInspector)
- * @see ApplyFunction#canVectorize(Expr.InputBindingInspector, Expr, List)
+ * @see ApplyFunction#canVectorize(Expr.InputBindingInspector, LambdaExpr,
List)
*/
default boolean canVectorize(Expr.InputBindingInspector inspector,
List<Expr> args)
{
- return false;
+ return
FallbackVectorProcessor.canFallbackVectorize(getOutputType(inspector, args),
inspector, args);
}
/**
@@ -154,11 +154,15 @@ public interface Function extends NamedFunction
* using {@link Expr#asVectorProcessor}, for use in vectorized query engines.
*
* @see Expr#asVectorProcessor(Expr.VectorInputBindingInspector)
- * @see ApplyFunction#asVectorProcessor(Expr.VectorInputBindingInspector,
Expr, List)
+ * @see ApplyFunction#asVectorProcessor(Expr.VectorInputBindingInspector,
LambdaExpr, List)
*/
default <T> ExprVectorProcessor<T>
asVectorProcessor(Expr.VectorInputBindingInspector inspector, List<Expr> args)
{
- throw new UOE("Function[%s] is not vectorized", name());
+ if (ExpressionProcessing.allowVectorizeFallback()) {
+ return FallbackVectorProcessor.create(this, args, inspector);
+ } else {
+ throw Exprs.cannotVectorize(this);
+ }
}
/**
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/FunctionalExpr.java
b/processing/src/main/java/org/apache/druid/math/expr/FunctionalExpr.java
index 02457e992a3..2159cc40008 100644
--- a/processing/src/main/java/org/apache/druid/math/expr/FunctionalExpr.java
+++ b/processing/src/main/java/org/apache/druid/math/expr/FunctionalExpr.java
@@ -23,7 +23,6 @@ import com.google.common.collect.ImmutableList;
import org.apache.druid.error.DruidException;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.math.expr.vector.ExprVectorProcessor;
-import org.apache.druid.math.expr.vector.FallbackVectorProcessor;
import org.apache.druid.query.filter.ColumnIndexSelector;
import org.apache.druid.segment.column.ColumnIndexSupplier;
import org.apache.druid.segment.column.ColumnType;
@@ -104,17 +103,13 @@ class FunctionExpr implements Expr
@Override
public boolean canVectorize(InputBindingInspector inspector)
{
- return function.canVectorize(inspector, args) ||
canFallbackVectorize(inspector, args);
+ return function.canVectorize(inspector, args);
}
@Override
public ExprVectorProcessor<?> asVectorProcessor(VectorInputBindingInspector
inspector)
{
- if (function.canVectorize(inspector, args)) {
- return function.asVectorProcessor(inspector, args);
- } else {
- return FallbackVectorProcessor.create(function, args, inspector);
- }
+ return function.asVectorProcessor(inspector, args);
}
@Nullable
@@ -248,25 +243,13 @@ class ApplyFunctionExpr implements Expr
@Override
public boolean canVectorize(InputBindingInspector inspector)
{
- return canVectorizeNative(inspector) ||
- (canFallbackVectorize(inspector, argsExpr) &&
lambdaExpr.canVectorize(inspector));
+ return function.canVectorize(inspector, lambdaExpr, argsExpr);
}
@Override
public <T> ExprVectorProcessor<T>
asVectorProcessor(VectorInputBindingInspector inspector)
{
- if (canVectorizeNative(inspector)) {
- return function.asVectorProcessor(inspector, lambdaExpr, argsExpr);
- } else {
- return FallbackVectorProcessor.create(function, lambdaExpr, argsExpr,
inspector);
- }
- }
-
- private boolean canVectorizeNative(InputBindingInspector inspector)
- {
- return function.canVectorize(inspector, lambdaExpr, argsExpr) &&
- lambdaExpr.canVectorize(inspector) &&
- argsExpr.stream().allMatch(expr -> expr.canVectorize(inspector));
+ return function.asVectorProcessor(inspector, lambdaExpr, argsExpr);
}
@Override
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/FallbackVectorProcessor.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/FallbackVectorProcessor.java
index 115c9eed397..e8cf412463d 100644
---
a/processing/src/main/java/org/apache/druid/math/expr/vector/FallbackVectorProcessor.java
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/FallbackVectorProcessor.java
@@ -25,6 +25,7 @@ import org.apache.druid.math.expr.Expr;
import org.apache.druid.math.expr.ExprEval;
import org.apache.druid.math.expr.ExprMacroTable;
import org.apache.druid.math.expr.ExprType;
+import org.apache.druid.math.expr.ExpressionProcessing;
import org.apache.druid.math.expr.ExpressionType;
import org.apache.druid.math.expr.Function;
import org.apache.druid.math.expr.LambdaExpr;
@@ -113,6 +114,21 @@ public abstract class FallbackVectorProcessor<T>
implements ExprVectorProcessor<
);
}
+ /**
+ * Returns whether {@link #create(Function, List,
Expr.VectorInputBindingInspector)} can be used to make
+ * a fallback vectorized processor.
+ */
+ public static boolean canFallbackVectorize(
+ @Nullable final ExpressionType outputType,
+ final Expr.InputBindingInspector inspector,
+ final List<Expr> args
+ )
+ {
+ return ExpressionProcessing.allowVectorizeFallback() &&
+ outputType != null &&
+ inspector.canVectorize(args);
+ }
+
/**
* Helper for the two {@link #create} methods. Makes {@link AdaptedExpr}
that can replace the original args to
* the {@link Expr} that requires fallback.
diff --git
a/processing/src/test/java/org/apache/druid/math/expr/VectorExprResultConsistencyTest.java
b/processing/src/test/java/org/apache/druid/math/expr/VectorExprResultConsistencyTest.java
index aa878cb47cf..c54f01472bc 100644
---
a/processing/src/test/java/org/apache/druid/math/expr/VectorExprResultConsistencyTest.java
+++
b/processing/src/test/java/org/apache/druid/math/expr/VectorExprResultConsistencyTest.java
@@ -441,6 +441,31 @@ public class VectorExprResultConsistencyTest extends
InitializedNullHandlingTest
testExpression("array(d1, d2)", types);
testExpression("array(l1, d2)", types);
testExpression("array(s1, l2)", types);
+
+ testExpression("map((x) -> x + 1, array(l1, l2))", types);
+ testExpression("map((x) -> x * 2.0, array(d1, d2))", types);
+ testExpression("map((x) -> concat(x, '_mapped'), array(s1, s2))", types);
+
+ testExpression("cartesian_map((x, y) -> x + y, array(l1, l2), array(d1,
d2))", types);
+ testExpression("cartesian_map((x, y) -> concat(x, cast(y, 'STRING')),
array(s1, s2), array(l1, l2))", types);
+
+ testExpression("fold((x, acc) -> x + acc, array(l1, l2), 0)", types);
+ testExpression("fold((x, acc) -> x + acc, array(d1, d2), 0.0)", types);
+ testExpression("fold((x, acc) -> concat(acc, x), array(s1, s2), '')",
types);
+
+ testExpression("cartesian_fold((x, y, acc) -> acc + x + y, array(l1,
l2), array(d1, d2), 0)", types);
+
+ testExpression("filter((x) -> x > 0, array(l1, l2))", types);
+ testExpression("filter((x) -> x > 0.0, array(d1, d2))", types);
+ testExpression("filter((x) -> strlen(x) > 0, array(s1, s2))", types);
+
+ testExpression("any((x) -> x > 0, array(l1, l2))", types);
+ testExpression("any((x) -> x > 0.0, array(d1, d2))", types);
+ testExpression("any((x) -> strlen(x) > 0, array(s1, s2))", types);
+
+ testExpression("all((x) -> x != null, array(l1, l2))", types);
+ testExpression("all((x) -> x != null, array(d1, d2))", types);
+ testExpression("all((x) -> x != null, array(s1, s2))", types);
}
finally {
ExpressionProcessing.initializeForTests();
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
