gianm commented on code in PR #16311:
URL: https://github.com/apache/druid/pull/16311#discussion_r1581233813
##########
processing/src/main/java/org/apache/druid/math/expr/Function.java:
##########
@@ -3742,23 +3746,95 @@ public ExpressionType
getOutputType(Expr.InputBindingInspector inspector, List<E
@Override
Expr getScalarArgument(List<Expr> args)
{
- return args.get(0);
+ return args.get(SCALAR_ARG);
}
@Override
Expr getArrayArgument(List<Expr> args)
{
- return args.get(1);
+ return args.get(ARRAY_ARG);
}
@Override
- ExprEval doApply(ExprEval arrayExpr, ExprEval scalarExpr)
+ ExprEval doApply(ExprEval arrayEval, ExprEval scalarEval)
{
- final Object[] array =
arrayExpr.castTo(scalarExpr.asArrayType()).asArray();
+ final Object[] array = arrayEval.asArray();
if (array == null) {
return ExprEval.ofLong(null);
}
- return
ExprEval.ofLongBoolean(Arrays.asList(array).contains(scalarExpr.value()));
+
+ if (scalarEval.value() == null) {
+ return Arrays.asList(array).contains(null) ?
ExprEval.ofLongBoolean(true) : ExprEval.ofLong(null);
+ }
+
+ final ExpressionType matchType = arrayEval.elementType();
+ final ExprEval<?> scalarEvalForComparison =
ExprEval.castForEqualityComparison(scalarEval, matchType);
+
+ if (scalarEvalForComparison == null) {
+ return ExprEval.ofLongBoolean(false);
+ } else {
+ return
ExprEval.ofLongBoolean(Arrays.asList(array).contains(scalarEvalForComparison.value()));
+ }
+ }
+
+ @Override
+ public Function asSingleThreaded(List<Expr> args,
Expr.InputBindingInspector inspector)
+ {
+ if (args.get(ARRAY_ARG).isLiteral()) {
+ final ExpressionType lhsType =
args.get(SCALAR_ARG).getOutputType(inspector);
+ if (lhsType == null) {
+ return this;
+ }
+
+ final ExprEval<?> rhsEval =
args.get(ARRAY_ARG).eval(InputBindings.nilBindings());
+ return new WithConstantArray(rhsEval);
+ }
+ return this;
+ }
+
+ /**
+ * Specialization of {@link ScalarInArrayFunction} for constant {@link
#ARRAY_ARG}.
+ */
+ private static final class WithConstantArray extends ScalarInArrayFunction
+ {
+ private final Set<Object> matchValues;
+
+ @Nullable
+ private final ExpressionType matchType;
+
+ public WithConstantArray(final ExprEval<?> arrayEval)
+ {
+ final Object[] arrayValues = arrayEval.asArray();
+
+ if (arrayValues == null) {
+ matchValues = Collections.emptySet();
+ matchType = null;
+ } else {
+ matchValues = new HashSet<>();
+ Collections.addAll(matchValues, arrayValues);
+ matchType = arrayEval.elementType();
+ }
+ }
+
+ @Override
+ ExprEval doApply(final ExprEval arrayExpr, final ExprEval scalarEval)
+ {
+ if (matchType == null) {
+ return ExprEval.ofLong(null);
Review Comment:
I added a specialization for `WithNullArray`, so now the `WithConstantArray`
specialization requires a non-null constant array.
While doing this I also realized I had implemented `WithConstantArray`
inefficiently: by overriding `doApply` rather than `apply`, the constant array
expr would still have `eval()` called on it. I changed it to override `apply`
instead.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
