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 17eaeb003f9 vectorize if function when then and else expressions have
the same output type (#18507)
17eaeb003f9 is described below
commit 17eaeb003f9638afa7cefbaca276da988992f124
Author: Clint Wylie <[email protected]>
AuthorDate: Wed Sep 10 11:38:37 2025 -0700
vectorize if function when then and else expressions have the same output
type (#18507)
* vectorize if function when then and else expressions have the same output
type
* fix style
* fix test
* add comments, more test
* return false
---
.../java/org/apache/druid/math/expr/Function.java | 15 +++
.../math/expr/vector/ExprEvalBindingVector.java | 58 +++++++++---
.../math/expr/vector/ExprEvalDoubleVector.java | 10 ++
.../druid/math/expr/vector/ExprEvalLongVector.java | 10 ++
.../math/expr/vector/ExprEvalObjectVector.java | 9 ++
.../druid/math/expr/vector/ExprEvalVector.java | 2 +
.../math/expr/vector/IfDoubleVectorProcessor.java | 100 ++++++++++++++++++++
.../expr/vector/IfFunctionVectorProcessor.java | 59 ++++++++++++
.../math/expr/vector/IfLongVectorProcessor.java | 101 +++++++++++++++++++++
.../math/expr/vector/IfObjectVectorProcessor.java | 82 +++++++++++++++++
.../expr/vector/VectorConditionalProcessors.java | 40 ++++++++
.../math/expr/VectorExprResultConsistencyTest.java | 25 +++++
.../timeseries/TimeseriesQueryRunnerTest.java | 1 -
13 files changed, 500 insertions(+), 12 deletions(-)
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 3f2a357b674..c5b530192a5 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
@@ -2140,6 +2140,21 @@ public interface Function extends NamedFunction
{
return ExpressionTypeConversion.conditional(inspector, args.subList(1,
3));
}
+
+ @Override
+ public boolean canVectorize(Expr.InputBindingInspector inspector,
List<Expr> args)
+ {
+ // vector engine requires consistent typing, but native if function does
not coerce then and else expressions,
+ // so for now we can only vectorize if both args have the same output
type to not have a behavior change
+ final ExpressionType thenType = args.get(1).getOutputType(inspector);
+ return Objects.equals(thenType, args.get(2).getOutputType(inspector));
+ }
+
+ @Override
+ public <T> ExprVectorProcessor<T>
asVectorProcessor(Expr.VectorInputBindingInspector inspector, List<Expr> args)
+ {
+ return VectorConditionalProcessors.ifFunction(inspector, args.get(0),
args.get(1), args.get(2));
+ }
}
/**
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalBindingVector.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalBindingVector.java
index 2a0a43e4edb..72750c977e0 100644
---
a/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalBindingVector.java
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalBindingVector.java
@@ -43,6 +43,7 @@ public class ExprEvalBindingVector<T> implements
ExprEvalVector<T>
@Nullable
private boolean[] numericNulls;
+ private Object[] objects;
public ExprEvalBindingVector(
ExpressionType expressionType,
@@ -76,9 +77,10 @@ public class ExprEvalBindingVector<T> implements
ExprEvalVector<T>
public long[] getLongVector()
{
if (expressionType.isNumeric()) {
- return bindings.getLongVector(bindingName);
+ longs = bindings.getLongVector(bindingName);
+ } else {
+ computeNumbers();
}
- computeNumbers();
return longs;
}
@@ -86,9 +88,10 @@ public class ExprEvalBindingVector<T> implements
ExprEvalVector<T>
public double[] getDoubleVector()
{
if (expressionType.isNumeric()) {
- return bindings.getDoubleVector(bindingName);
+ doubles = bindings.getDoubleVector(bindingName);
+ } else {
+ computeNumbers();
}
- computeNumbers();
return doubles;
}
@@ -97,9 +100,10 @@ public class ExprEvalBindingVector<T> implements
ExprEvalVector<T>
public boolean[] getNullVector()
{
if (expressionType.isNumeric()) {
- return bindings.getNullVector(bindingName);
+ numericNulls = bindings.getNullVector(bindingName);
+ } else {
+ computeNumbers();
}
- computeNumbers();
return numericNulls;
}
@@ -109,7 +113,7 @@ public class ExprEvalBindingVector<T> implements
ExprEvalVector<T>
if (expressionType.is(ExprType.LONG)) {
final long[] values = bindings.getLongVector(bindingName);
final boolean[] nulls = bindings.getNullVector(bindingName);
- final Long[] objects = new Long[values.length];
+ objects = new Long[values.length];
if (nulls != null) {
for (int i = 0; i < values.length; i++) {
objects[i] = nulls[i] ? null : values[i];
@@ -119,11 +123,10 @@ public class ExprEvalBindingVector<T> implements
ExprEvalVector<T>
objects[i] = values[i];
}
}
- return objects;
} else if (expressionType.is(ExprType.DOUBLE)) {
final double[] values = bindings.getDoubleVector(bindingName);
final boolean[] nulls = bindings.getNullVector(bindingName);
- Double[] objects = new Double[values.length];
+ objects = new Double[values.length];
if (nulls != null) {
for (int i = 0; i < values.length; i++) {
objects[i] = nulls[i] ? null : values[i];
@@ -133,9 +136,42 @@ public class ExprEvalBindingVector<T> implements
ExprEvalVector<T>
objects[i] = values[i];
}
}
- return objects;
+ } else {
+ objects = bindings.getObjectVector(bindingName);
+ }
+ return objects;
+ }
+
+ @Override
+ public boolean elementAsBoolean(int index)
+ {
+ if (expressionType.is(ExprType.LONG)) {
+ if (longs == null) {
+ // populate stuff
+ getLongVector();
+ }
+ if (numericNulls != null && numericNulls[index]) {
+ return Evals.asBoolean(0L);
+ }
+ return Evals.asBoolean(longs[index]);
+ } else if (expressionType.is(ExprType.DOUBLE)) {
+ if (doubles == null) {
+ getDoubleVector();
+ }
+ if (numericNulls != null && numericNulls[index]) {
+ return Evals.asBoolean(0.0);
+ }
+ return Evals.asBoolean(doubles[index]);
+ } else {
+ if (objects == null) {
+ getObjectVector();
+ }
+ if (expressionType.is(ExprType.STRING)) {
+ return Evals.asBoolean((String) objects[index]);
+ } else {
+ return ExprEval.ofType(expressionType, objects[index]).asBoolean();
+ }
}
- return bindings.getObjectVector(bindingName);
}
private void computeNumbers()
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalDoubleVector.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalDoubleVector.java
index 877a3a4f726..5f484a25041 100644
---
a/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalDoubleVector.java
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalDoubleVector.java
@@ -19,6 +19,7 @@
package org.apache.druid.math.expr.vector;
+import org.apache.druid.math.expr.Evals;
import org.apache.druid.math.expr.ExpressionType;
import java.util.Arrays;
@@ -69,4 +70,13 @@ public final class ExprEvalDoubleVector extends
BaseExprEvalVector<double[]>
}
return objects;
}
+
+ @Override
+ public boolean elementAsBoolean(int index)
+ {
+ if (nulls != null && nulls[index]) {
+ return false;
+ }
+ return Evals.asBoolean(values[index]);
+ }
}
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalLongVector.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalLongVector.java
index b33eb968f17..1857694efe6 100644
---
a/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalLongVector.java
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalLongVector.java
@@ -19,6 +19,7 @@
package org.apache.druid.math.expr.vector;
+import org.apache.druid.math.expr.Evals;
import org.apache.druid.math.expr.ExpressionType;
import javax.annotation.Nullable;
@@ -65,4 +66,13 @@ public final class ExprEvalLongVector extends
BaseExprEvalVector<long[]>
return objects;
}
+ @Override
+ public boolean elementAsBoolean(int index)
+ {
+ if (nulls != null && nulls[index]) {
+ return false;
+ }
+ return Evals.asBoolean(values[index]);
+ }
+
}
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalObjectVector.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalObjectVector.java
index 888a4c59f8c..0bbabc12714 100644
---
a/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalObjectVector.java
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalObjectVector.java
@@ -113,4 +113,13 @@ public final class ExprEvalObjectVector extends
BaseExprEvalVector<Object[]>
{
return values;
}
+
+ @Override
+ public boolean elementAsBoolean(int index)
+ {
+ if (type.is(ExprType.STRING)) {
+ return Evals.asBoolean((String) values[index]);
+ }
+ return ExprEval.ofType(type, values[index]).asBoolean();
+ }
}
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalVector.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalVector.java
index 6f65ce31fdf..2121001066b 100644
---
a/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalVector.java
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/ExprEvalVector.java
@@ -40,4 +40,6 @@ public interface ExprEvalVector<T>
long[] getLongVector();
double[] getDoubleVector();
Object[] getObjectVector();
+
+ boolean elementAsBoolean(int index);
}
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/IfDoubleVectorProcessor.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/IfDoubleVectorProcessor.java
new file mode 100644
index 00000000000..3919b180fa3
--- /dev/null
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/IfDoubleVectorProcessor.java
@@ -0,0 +1,100 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.math.expr.vector;
+
+import org.apache.druid.math.expr.Expr;
+import org.apache.druid.math.expr.ExpressionType;
+
+public class IfDoubleVectorProcessor extends
IfFunctionVectorProcessor<double[]>
+{
+ private final double[] output;
+ private final boolean[] outputNulls;
+
+ public IfDoubleVectorProcessor(
+ ExprVectorProcessor<?> conditionProcessor,
+ ExprVectorProcessor<double[]> thenProcessor,
+ ExprVectorProcessor<double[]> elseProcessor
+ )
+ {
+ super(
+ ExpressionType.DOUBLE,
+ conditionProcessor,
+ CastToTypeVectorProcessor.cast(thenProcessor, ExpressionType.DOUBLE),
+ CastToTypeVectorProcessor.cast(elseProcessor, ExpressionType.DOUBLE)
+ );
+ this.output = new double[conditionProcessor.maxVectorSize()];
+ this.outputNulls = new boolean[conditionProcessor.maxVectorSize()];
+ }
+
+ @Override
+ public ExprEvalVector<double[]> evalVector(Expr.VectorInputBinding bindings)
+ {
+ thenBindingFilterer.setBindings(bindings);
+ elseBindingFilterer.setBindings(bindings);
+ final ExprEvalVector<?> conditionVector =
conditionProcessor.evalVector(bindings);
+
+ final int[] thenSelection =
thenBindingFilterer.getVectorMatch().getSelection();
+ final int[] elseSelection =
elseBindingFilterer.getVectorMatch().getSelection();
+ int thens = 0;
+ int elses = 0;
+ for (int i = 0; i < bindings.getCurrentVectorSize(); i++) {
+ if (conditionVector.elementAsBoolean(i)) {
+ thenSelection[thens++] = i;
+ } else {
+ elseSelection[elses++] = i;
+ }
+ }
+ thenBindingFilterer.getVectorMatch().setSelectionSize(thens);
+ elseBindingFilterer.getVectorMatch().setSelectionSize(elses);
+
+ if (elses == 0) {
+ return thenProcessor.evalVector(bindings);
+ } else if (thens == 0) {
+ return elseProcessor.evalVector(bindings);
+ }
+
+ final ExprEvalVector<double[]> thenVector =
thenProcessor.evalVector(thenBindingFilterer);
+ final double[] thenValues = thenVector.getDoubleVector();
+ final boolean[] thenNulls = thenVector.getNullVector();
+ for (int i = 0; i < thens; i++) {
+ final int outIndex = thenSelection[i];
+ if (thenNulls != null && thenNulls[i]) {
+ outputNulls[outIndex] = true;
+ } else {
+ output[outIndex] = thenValues[i];
+ outputNulls[outIndex] = false;
+ }
+ }
+
+ final ExprEvalVector<double[]> elseVector =
elseProcessor.evalVector(elseBindingFilterer);
+ final double[] elseValues = elseVector.getDoubleVector();
+ final boolean[] elseNulls = elseVector.getNullVector();
+ for (int i = 0; i < elses; i++) {
+ final int outIndex = elseSelection[i];
+ if (elseNulls != null && elseNulls[i]) {
+ outputNulls[outIndex] = true;
+ } else {
+ output[outIndex] = elseValues[i];
+ outputNulls[outIndex] = false;
+ }
+ }
+ return new ExprEvalDoubleVector(output, outputNulls);
+ }
+}
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/IfFunctionVectorProcessor.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/IfFunctionVectorProcessor.java
new file mode 100644
index 00000000000..0cd08bf6c58
--- /dev/null
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/IfFunctionVectorProcessor.java
@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.math.expr.vector;
+
+import org.apache.druid.math.expr.ExpressionType;
+
+public abstract class IfFunctionVectorProcessor<T> implements
ExprVectorProcessor<T>
+{
+ final ExpressionType outputType;
+ final ExprVectorProcessor<?> conditionProcessor;
+ final ExprVectorProcessor<T> thenProcessor;
+ final ExprVectorProcessor<T> elseProcessor;
+ final FilteredVectorInputBinding thenBindingFilterer;
+ final FilteredVectorInputBinding elseBindingFilterer;
+
+ public IfFunctionVectorProcessor(
+ ExpressionType outputType,
+ ExprVectorProcessor<?> conditionProcessor,
+ ExprVectorProcessor<T> thenProcessor,
+ ExprVectorProcessor<T> elseProcessor
+ )
+ {
+ this.outputType = outputType;
+ this.conditionProcessor = conditionProcessor;
+ this.thenProcessor = thenProcessor;
+ this.elseProcessor = elseProcessor;
+ this.thenBindingFilterer = new
FilteredVectorInputBinding(conditionProcessor.maxVectorSize());
+ this.elseBindingFilterer = new
FilteredVectorInputBinding(conditionProcessor.maxVectorSize());
+ }
+
+ @Override
+ public ExpressionType getOutputType()
+ {
+ return outputType;
+ }
+
+ @Override
+ public int maxVectorSize()
+ {
+ return conditionProcessor.maxVectorSize();
+ }
+}
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/IfLongVectorProcessor.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/IfLongVectorProcessor.java
new file mode 100644
index 00000000000..a8b6c7f171a
--- /dev/null
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/IfLongVectorProcessor.java
@@ -0,0 +1,101 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.math.expr.vector;
+
+import org.apache.druid.math.expr.Expr;
+import org.apache.druid.math.expr.ExpressionType;
+
+public class IfLongVectorProcessor extends IfFunctionVectorProcessor<long[]>
+{
+ private final long[] output;
+ private final boolean[] outputNulls;
+
+ public IfLongVectorProcessor(
+ ExprVectorProcessor<?> conditionProcessor,
+ ExprVectorProcessor<long[]> thenProcessor,
+ ExprVectorProcessor<long[]> elseProcessor
+ )
+ {
+ super(
+ ExpressionType.LONG,
+ conditionProcessor,
+ CastToTypeVectorProcessor.cast(thenProcessor, ExpressionType.LONG),
+ CastToTypeVectorProcessor.cast(elseProcessor, ExpressionType.LONG)
+ );
+ this.output = new long[conditionProcessor.maxVectorSize()];
+ this.outputNulls = new boolean[conditionProcessor.maxVectorSize()];
+ }
+
+ @Override
+ public ExprEvalVector<long[]> evalVector(Expr.VectorInputBinding bindings)
+ {
+ thenBindingFilterer.setBindings(bindings);
+ elseBindingFilterer.setBindings(bindings);
+ final ExprEvalVector<?> conditionVector =
conditionProcessor.evalVector(bindings);
+
+ final int[] thenSelection =
thenBindingFilterer.getVectorMatch().getSelection();
+ final int[] elseSelection =
elseBindingFilterer.getVectorMatch().getSelection();
+ int thens = 0;
+ int elses = 0;
+ for (int i = 0; i < bindings.getCurrentVectorSize(); i++) {
+ if (conditionVector.elementAsBoolean(i)) {
+ thenSelection[thens++] = i;
+ } else {
+ elseSelection[elses++] = i;
+ }
+ }
+ thenBindingFilterer.getVectorMatch().setSelectionSize(thens);
+ elseBindingFilterer.getVectorMatch().setSelectionSize(elses);
+
+ if (elses == 0) {
+ return thenProcessor.evalVector(bindings);
+ } else if (thens == 0) {
+ return elseProcessor.evalVector(bindings);
+ }
+
+
+ final ExprEvalVector<long[]> thenVector =
thenProcessor.evalVector(thenBindingFilterer);
+ final long[] thenValues = thenVector.getLongVector();
+ final boolean[] thenNulls = thenVector.getNullVector();
+ for (int i = 0; i < thens; i++) {
+ final int outIndex = thenSelection[i];
+ if (thenNulls != null && thenNulls[i]) {
+ outputNulls[outIndex] = true;
+ } else {
+ output[outIndex] = thenValues[i];
+ outputNulls[outIndex] = false;
+ }
+ }
+
+ final ExprEvalVector<long[]> elseVector =
elseProcessor.evalVector(elseBindingFilterer);
+ final long[] elseValues = elseVector.getLongVector();
+ final boolean[] elseNulls = elseVector.getNullVector();
+ for (int i = 0; i < elses; i++) {
+ final int outIndex = elseSelection[i];
+ if (elseNulls != null && elseNulls[i]) {
+ outputNulls[outIndex] = true;
+ } else {
+ output[outIndex] = elseValues[i];
+ outputNulls[outIndex] = false;
+ }
+ }
+ return new ExprEvalLongVector(output, outputNulls);
+ }
+}
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/IfObjectVectorProcessor.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/IfObjectVectorProcessor.java
new file mode 100644
index 00000000000..269589a5030
--- /dev/null
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/IfObjectVectorProcessor.java
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.math.expr.vector;
+
+import org.apache.druid.math.expr.Expr;
+import org.apache.druid.math.expr.ExpressionType;
+
+public class IfObjectVectorProcessor extends
IfFunctionVectorProcessor<Object[]>
+{
+ private final Object[] output;
+
+ public IfObjectVectorProcessor(
+ ExpressionType outputType,
+ ExprVectorProcessor<?> conditionProcessor,
+ ExprVectorProcessor<Object[]> thenProcessor,
+ ExprVectorProcessor<Object[]> elseProcessor
+ )
+ {
+ super(outputType, conditionProcessor, thenProcessor, elseProcessor);
+ this.output = new Object[conditionProcessor.maxVectorSize()];
+ }
+
+ @Override
+ public ExprEvalVector<Object[]> evalVector(Expr.VectorInputBinding bindings)
+ {
+ thenBindingFilterer.setBindings(bindings);
+ elseBindingFilterer.setBindings(bindings);
+ final ExprEvalVector<?> conditionVector =
conditionProcessor.evalVector(bindings);
+
+ final int[] thenSelection =
thenBindingFilterer.getVectorMatch().getSelection();
+ final int[] elseSelection =
elseBindingFilterer.getVectorMatch().getSelection();
+ int thens = 0;
+ int elses = 0;
+ for (int i = 0; i < bindings.getCurrentVectorSize(); i++) {
+ if (conditionVector.elementAsBoolean(i)) {
+ thenSelection[thens++] = i;
+ } else {
+ elseSelection[elses++] = i;
+ }
+ }
+ thenBindingFilterer.getVectorMatch().setSelectionSize(thens);
+ elseBindingFilterer.getVectorMatch().setSelectionSize(elses);
+
+ if (elses == 0) {
+ return thenProcessor.evalVector(bindings);
+ } else if (thens == 0) {
+ return elseProcessor.evalVector(bindings);
+ }
+
+ final ExprEvalVector<Object[]> thenVector =
thenProcessor.evalVector(thenBindingFilterer);
+ final Object[] thenValues = thenVector.getObjectVector();
+ for (int i = 0; i < thens; i++) {
+ final int outIndex = thenSelection[i];
+ output[outIndex] = thenValues[i];
+ }
+
+ final ExprEvalVector<Object[]> elseVector =
elseProcessor.evalVector(elseBindingFilterer);
+ final Object[] elseValues = elseVector.getObjectVector();
+ for (int i = 0; i < elses; i++) {
+ final int outIndex = elseSelection[i];
+ output[outIndex] = elseValues[i];
+ }
+ return new ExprEvalObjectVector(output, outputType);
+ }
+}
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/VectorConditionalProcessors.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/VectorConditionalProcessors.java
index 30b775f0236..3fbe711948e 100644
---
a/processing/src/main/java/org/apache/druid/math/expr/vector/VectorConditionalProcessors.java
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/VectorConditionalProcessors.java
@@ -57,4 +57,44 @@ public class VectorConditionalProcessors
}
return (ExprVectorProcessor<T>) processor;
}
+
+ public static <T> ExprVectorProcessor<T> ifFunction(
+ Expr.VectorInputBindingInspector inspector,
+ Expr conditionExpr,
+ Expr thenExpr,
+ Expr elseExpr
+ )
+ {
+ // right now this function can only vectorize if then and else clause have
same output type, if this changes then
+ // we'll need to switch this to use whatever output type logic that is
using
+ final ExpressionType outputType = thenExpr.getOutputType(inspector);
+
+ final ExprVectorProcessor<?> processor;
+ if (outputType == null) {
+ // if output type is null, it means all the input types were null
(non-existent), and if(null, null, null) is null
+ return VectorProcessors.constant((Long) null,
inspector.getMaxVectorSize());
+ }
+ if (outputType.is(ExprType.LONG)) {
+ // long is most restrictive so both processors are definitely long typed
if output is long
+ processor = new IfLongVectorProcessor(
+ conditionExpr.asVectorProcessor(inspector),
+ thenExpr.asVectorProcessor(inspector),
+ elseExpr.asVectorProcessor(inspector)
+ );
+ } else if (outputType.is(ExprType.DOUBLE)) {
+ processor = new IfDoubleVectorProcessor(
+ conditionExpr.asVectorProcessor(inspector),
+ thenExpr.asVectorProcessor(inspector),
+ elseExpr.asVectorProcessor(inspector)
+ );
+ } else {
+ processor = new IfObjectVectorProcessor(
+ outputType,
+ conditionExpr.asVectorProcessor(inspector),
+ thenExpr.asVectorProcessor(inspector),
+ elseExpr.asVectorProcessor(inspector)
+ );
+ }
+ return (ExprVectorProcessor<T>) processor;
+ }
}
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 11580eb9f5b..64a955f2e6f 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
@@ -38,6 +38,7 @@ import org.apache.druid.testing.InitializedNullHandlingTest;
import org.junit.Assert;
import org.junit.Assume;
import org.junit.Test;
+import org.junit.jupiter.api.Assertions;
import javax.annotation.Nullable;
import java.util.ArrayList;
@@ -348,6 +349,30 @@ public class VectorExprResultConsistencyTest extends
InitializedNullHandlingTest
testFunctions(types, templates, functions);
}
+ @Test
+ public void testIfFunction()
+ {
+ testExpression("if(l1, l1, l2)", types);
+ testExpression("if(l1, s1, s2)", types);
+ testExpression("if(l1, d1, d2)", types);
+ testExpression("if(d1, l1, l2)", types);
+ testExpression("if(d1, s1, s2)", types);
+ testExpression("if(d1, d1, d2)", types);
+ testExpression("if(boolString1, s1, s2)", types);
+ testExpression("if(boolString1, l1, l2)", types);
+ testExpression("if(boolString1, d1, d2)", types);
+ // make sure eval of else is lazy, else this would be divide by zero error
+ testExpression("if(l1 % 2 == 0, -1, l2 / (l1 % 2))", types);
+ // cannot vectorize mixed types
+ Assertions.assertFalse(
+ Parser.parse("if(s1, l1, d2)",
MACRO_TABLE).canVectorize(InputBindings.inspectorFromTypeMap(types))
+ );
+ Assertions.assertFalse(
+ Parser.parse("if(s1, d1, s2)",
MACRO_TABLE).canVectorize(InputBindings.inspectorFromTypeMap(types))
+ );
+ }
+
+
@Test
public void testStringFns()
{
diff --git
a/processing/src/test/java/org/apache/druid/query/timeseries/TimeseriesQueryRunnerTest.java
b/processing/src/test/java/org/apache/druid/query/timeseries/TimeseriesQueryRunnerTest.java
index 16254aac72b..3bbd83c7720 100644
---
a/processing/src/test/java/org/apache/druid/query/timeseries/TimeseriesQueryRunnerTest.java
+++
b/processing/src/test/java/org/apache/druid/query/timeseries/TimeseriesQueryRunnerTest.java
@@ -2316,7 +2316,6 @@ public class TimeseriesQueryRunnerTest extends
InitializedNullHandlingTest
@Test
public void testTimeSeriesWithFilteredAggAndExpressionFilteredAgg()
{
- cannotVectorizeUnlessFallback();
TimeseriesQuery query = Druids
.newTimeseriesQueryBuilder()
.dataSource(QueryRunnerTestHelper.DATA_SOURCE)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
