This is an automated email from the ASF dual-hosted git repository.
Jackie-Jiang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 881a1f74dc6 Honor nullableParameters for scalar functions in the
single-stage engine (#19983)
881a1f74dc6 is described below
commit 881a1f74dc660983a82080e4f78eca207c37cb17
Author: Xiaotian (Jackie) Jiang <[email protected]>
AuthorDate: Tue Oct 20 21:43:58 2026 -0700
Honor nullableParameters for scalar functions in the single-stage engine
(#19983)
---
.../common/function/scalar/StringFunctions.java | 18 +-
.../function/scalar/StringFunctionsTest.java | 6 +
.../function/ScalarTransformFunctionWrapper.java | 250 +++++++++++++++++++--
.../ScalarTransformFunctionWrapperTest.java | 65 ++++++
.../queries/NullHandlingEnabledQueriesTest.java | 24 ++
5 files changed, 341 insertions(+), 22 deletions(-)
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/StringFunctions.java
b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/StringFunctions.java
index 8482797325a..fcbac14c673 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/StringFunctions.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/StringFunctions.java
@@ -449,9 +449,14 @@ public class StringFunctions {
/// @param input an input string for prefix strings generations.
/// @param maxlength the max length of the prefix strings for the string.
/// @param prefix the prefix to be prepended to prefix strings generated.
e.g. '^' for regex matching
- /// @return generate an array of prefix matchers of the string that are
shorter than the specified length.
+ /// @return generate an array of prefix matchers of the string that are
shorter than the specified length, or `null`
+ /// if the input is `null`.
+ @Nullable
@ScalarFunction(nullableParameters = true)
- public static String[] prefixesWithPrefix(String input, int maxlength,
@Nullable String prefix) {
+ public static String[] prefixesWithPrefix(@Nullable String input, int
maxlength, @Nullable String prefix) {
+ if (input == null) {
+ return null;
+ }
if (prefix == null) {
return prefixes(input, maxlength);
}
@@ -479,9 +484,14 @@ public class StringFunctions {
/// @param input an input string for suffix strings generations.
/// @param maxlength the max length of the suffix strings for the string.
/// @param suffix the suffix string to be appended for suffix strings
generated. e.g. '$' for regex matching.
- /// @return generate an array of suffix matchers of the string that are
shorter than the specified length.
+ /// @return generate an array of suffix matchers of the string that are
shorter than the specified length, or `null`
+ /// if the input is `null`.
+ @Nullable
@ScalarFunction(nullableParameters = true)
- public static String[] suffixesWithSuffix(String input, int maxlength,
@Nullable String suffix) {
+ public static String[] suffixesWithSuffix(@Nullable String input, int
maxlength, @Nullable String suffix) {
+ if (input == null) {
+ return null;
+ }
if (suffix == null) {
return suffixes(input, maxlength);
}
diff --git
a/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/StringFunctionsTest.java
b/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/StringFunctionsTest.java
index ed6f3359df1..d705c65559e 100644
---
a/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/StringFunctionsTest.java
+++
b/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/StringFunctionsTest.java
@@ -381,6 +381,12 @@ public class StringFunctionsTest {
assertEquals(StringFunctions.suffixesWithSuffix(input, length, "$"),
expectedSuffixWithRegexChar);
}
+ @Test
+ public void testPrefixAndSuffixWithNullInput() {
+ assertNull(StringFunctions.prefixesWithPrefix(null, 3, "^"));
+ assertNull(StringFunctions.suffixesWithSuffix(null, 3, "$"));
+ }
+
@Test(dataProvider = "initcapTestCases")
public void testInitcap(String input, String expected) {
assertEquals(StringFunctions.initcap(input), expected);
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapper.java
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapper.java
index f55123a3295..e6636906663 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapper.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapper.java
@@ -24,6 +24,7 @@ import java.sql.Timestamp;
import java.util.List;
import java.util.Map;
import java.util.UUID;
+import javax.annotation.Nullable;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.pinot.common.function.FunctionInfo;
import org.apache.pinot.common.function.FunctionUtils;
@@ -37,12 +38,20 @@ import org.apache.pinot.spi.utils.ByteArray;
import org.apache.pinot.spi.utils.CommonConstants.NullValuePlaceHolder;
import org.apache.pinot.spi.utils.PinotDataType;
import org.apache.pinot.spi.utils.UuidUtils;
+import org.roaringbitmap.PeekableIntIterator;
+import org.roaringbitmap.RoaringBitmap;
/// Wrapper transform function on the annotated scalar function.
+///
+/// When null handling is enabled and the function accepts `null` arguments
(see `ScalarFunction#nullableParameters`),
+/// the function is invoked with `null` for the null rows of the arguments,
and the rows where it returns `null` are the
+/// null rows of the result. Otherwise, the function is invoked with the
values stored in the null rows (e.g. the
+/// default null value of a column), and a row of the result is null when any
argument is null in it.
public class ScalarTransformFunctionWrapper extends BaseTransformFunction {
private final String _name;
private final QueryFunctionInvoker _functionInvoker;
+ private final boolean _nullableParameters;
private final ColumnDataType _resultType;
private final TransformResultMetadata _resultMetadata;
@@ -53,9 +62,19 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
private PinotDataType[] _nonLiteralTypes;
private Object[][] _nonLiteralValues;
+ // Whether the function is invoked with null for null argument rows, and
decides the null rows of the result
+ private boolean _invokeWithNulls;
+ // When invoking with nulls, the value block evaluated last, whose values
are in the result buffer, and the rows of it
+ // where the function returns null. Value blocks are not reused across
batches, so these stay valid for the block.
+ @Nullable
+ private ValueBlock _evaluatedValueBlock;
+ @Nullable
+ private RoaringBitmap _resultNullBitmap;
+
public ScalarTransformFunctionWrapper(FunctionInfo functionInfo) {
_name = functionInfo.getMethod().getName();
_functionInvoker = new QueryFunctionInvoker(functionInfo);
+ _nullableParameters = functionInfo.hasNullableParameters();
Class<?>[] parameterClasses = _functionInvoker.getParameterClasses();
PinotDataType[] parameterTypes = _functionInvoker.getParameterTypes();
int numParameters = parameterClasses.length;
@@ -157,6 +176,13 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
_nonLiteralValues = new Object[_numNonLiteralArguments][];
}
+ @Override
+ public void init(List<TransformFunction> arguments, Map<String,
ColumnContext> columnContextMap,
+ boolean nullHandlingEnabled) {
+ super.init(arguments, columnContextMap, nullHandlingEnabled);
+ _invokeWithNulls = nullHandlingEnabled && _nullableParameters;
+ }
+
@Override
public TransformResultMetadata getResultMetadata() {
return _resultMetadata;
@@ -167,6 +193,9 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
if (_resultMetadata.getDataType().getStoredType() != DataType.INT) {
return super.transformToIntValuesSV(valueBlock);
}
+ if (valueBlock == _evaluatedValueBlock) {
+ return _intValuesSV;
+ }
int length = valueBlock.getNumDocs();
initIntValuesSV(length);
getNonLiteralValues(valueBlock);
@@ -175,7 +204,12 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
_scalarArguments[_nonLiteralIndices[j]] = _nonLiteralValues[j][i];
}
Object value = _functionInvoker.invoke(_scalarArguments);
- _intValuesSV[i] = value != null ? (int) _resultType.toInternal(value) :
NullValuePlaceHolder.INT;
+ if (value != null) {
+ _intValuesSV[i] = (int) _resultType.toInternal(value);
+ } else {
+ _intValuesSV[i] = NullValuePlaceHolder.INT;
+ addNullResult(i);
+ }
}
return _intValuesSV;
}
@@ -185,6 +219,9 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
if (_resultMetadata.getDataType().getStoredType() != DataType.LONG) {
return super.transformToLongValuesSV(valueBlock);
}
+ if (valueBlock == _evaluatedValueBlock) {
+ return _longValuesSV;
+ }
int length = valueBlock.getNumDocs();
initLongValuesSV(length);
getNonLiteralValues(valueBlock);
@@ -193,7 +230,12 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
_scalarArguments[_nonLiteralIndices[j]] = _nonLiteralValues[j][i];
}
Object value = _functionInvoker.invoke(_scalarArguments);
- _longValuesSV[i] = value != null ? (long) _resultType.toInternal(value)
: NullValuePlaceHolder.LONG;
+ if (value != null) {
+ _longValuesSV[i] = (long) _resultType.toInternal(value);
+ } else {
+ _longValuesSV[i] = NullValuePlaceHolder.LONG;
+ addNullResult(i);
+ }
}
return _longValuesSV;
}
@@ -203,6 +245,9 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
if (_resultMetadata.getDataType().getStoredType() != DataType.FLOAT) {
return super.transformToFloatValuesSV(valueBlock);
}
+ if (valueBlock == _evaluatedValueBlock) {
+ return _floatValuesSV;
+ }
int length = valueBlock.getNumDocs();
initFloatValuesSV(length);
getNonLiteralValues(valueBlock);
@@ -211,8 +256,12 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
_scalarArguments[_nonLiteralIndices[j]] = _nonLiteralValues[j][i];
}
Object value = _functionInvoker.invoke(_scalarArguments);
- _floatValuesSV[i] =
- value != null ? (float) _resultType.toInternal(value) :
NullValuePlaceHolder.FLOAT;
+ if (value != null) {
+ _floatValuesSV[i] = (float) _resultType.toInternal(value);
+ } else {
+ _floatValuesSV[i] = NullValuePlaceHolder.FLOAT;
+ addNullResult(i);
+ }
}
return _floatValuesSV;
}
@@ -222,6 +271,9 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
if (_resultMetadata.getDataType().getStoredType() != DataType.DOUBLE) {
return super.transformToDoubleValuesSV(valueBlock);
}
+ if (valueBlock == _evaluatedValueBlock) {
+ return _doubleValuesSV;
+ }
int length = valueBlock.getNumDocs();
initDoubleValuesSV(length);
getNonLiteralValues(valueBlock);
@@ -230,8 +282,12 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
_scalarArguments[_nonLiteralIndices[j]] = _nonLiteralValues[j][i];
}
Object value = _functionInvoker.invoke(_scalarArguments);
- _doubleValuesSV[i] =
- value != null ? (double) _resultType.toInternal(value) :
NullValuePlaceHolder.DOUBLE;
+ if (value != null) {
+ _doubleValuesSV[i] = (double) _resultType.toInternal(value);
+ } else {
+ _doubleValuesSV[i] = NullValuePlaceHolder.DOUBLE;
+ addNullResult(i);
+ }
}
return _doubleValuesSV;
}
@@ -241,6 +297,9 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
if (_resultMetadata.getDataType().getStoredType() != DataType.BIG_DECIMAL)
{
return super.transformToBigDecimalValuesSV(valueBlock);
}
+ if (valueBlock == _evaluatedValueBlock) {
+ return _bigDecimalValuesSV;
+ }
int length = valueBlock.getNumDocs();
initBigDecimalValuesSV(length);
getNonLiteralValues(valueBlock);
@@ -249,8 +308,12 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
_scalarArguments[_nonLiteralIndices[j]] = _nonLiteralValues[j][i];
}
Object value = _functionInvoker.invoke(_scalarArguments);
- _bigDecimalValuesSV[i] =
- value != null ? (BigDecimal) _resultType.toInternal(value) :
NullValuePlaceHolder.BIG_DECIMAL;
+ if (value != null) {
+ _bigDecimalValuesSV[i] = (BigDecimal) _resultType.toInternal(value);
+ } else {
+ _bigDecimalValuesSV[i] = NullValuePlaceHolder.BIG_DECIMAL;
+ addNullResult(i);
+ }
}
return _bigDecimalValuesSV;
}
@@ -260,6 +323,9 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
if (_resultMetadata.getDataType().getStoredType() != DataType.STRING) {
return super.transformToStringValuesSV(valueBlock);
}
+ if (valueBlock == _evaluatedValueBlock) {
+ return _stringValuesSV;
+ }
int length = valueBlock.getNumDocs();
initStringValuesSV(length);
getNonLiteralValues(valueBlock);
@@ -268,8 +334,12 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
_scalarArguments[_nonLiteralIndices[j]] = _nonLiteralValues[j][i];
}
Object value = _functionInvoker.invoke(_scalarArguments);
- _stringValuesSV[i] =
- value != null ? (String) _resultType.toInternal(value) :
NullValuePlaceHolder.STRING;
+ if (value != null) {
+ _stringValuesSV[i] = (String) _resultType.toInternal(value);
+ } else {
+ _stringValuesSV[i] = NullValuePlaceHolder.STRING;
+ addNullResult(i);
+ }
}
return _stringValuesSV;
}
@@ -279,6 +349,9 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
if (_resultMetadata.getDataType().getStoredType() != DataType.BYTES) {
return super.transformToBytesValuesSV(valueBlock);
}
+ if (valueBlock == _evaluatedValueBlock) {
+ return _bytesValuesSV;
+ }
int length = valueBlock.getNumDocs();
initBytesValuesSV(length);
getNonLiteralValues(valueBlock);
@@ -287,9 +360,12 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
_scalarArguments[_nonLiteralIndices[j]] = _nonLiteralValues[j][i];
}
Object value = _functionInvoker.invoke(_scalarArguments);
- byte[] bytes =
- value != null ? ((ByteArray)
_resultType.toInternal(value)).getBytes() : NullValuePlaceHolder.BYTES;
- _bytesValuesSV[i] = bytes;
+ if (value != null) {
+ _bytesValuesSV[i] = ((ByteArray)
_resultType.toInternal(value)).getBytes();
+ } else {
+ _bytesValuesSV[i] = NullValuePlaceHolder.BYTES;
+ addNullResult(i);
+ }
}
return _bytesValuesSV;
}
@@ -299,6 +375,9 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
if (_resultMetadata.getDataType().getStoredType() != DataType.INT) {
return super.transformToIntValuesMV(valueBlock);
}
+ if (valueBlock == _evaluatedValueBlock) {
+ return _intValuesMV;
+ }
int length = valueBlock.getNumDocs();
initIntValuesMV(length);
getNonLiteralValues(valueBlock);
@@ -306,7 +385,13 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
for (int j = 0; j < _numNonLiteralArguments; j++) {
_scalarArguments[_nonLiteralIndices[j]] = _nonLiteralValues[j][i];
}
- _intValuesMV[i] = (int[])
_resultType.toInternal(_functionInvoker.invoke(_scalarArguments));
+ Object value = _functionInvoker.invoke(_scalarArguments);
+ if (value != null) {
+ _intValuesMV[i] = (int[]) _resultType.toInternal(value);
+ } else {
+ _intValuesMV[i] = NullValuePlaceHolder.INT_ARRAY;
+ addNullResult(i);
+ }
}
return _intValuesMV;
}
@@ -316,6 +401,9 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
if (_resultMetadata.getDataType().getStoredType() != DataType.LONG) {
return super.transformToLongValuesMV(valueBlock);
}
+ if (valueBlock == _evaluatedValueBlock) {
+ return _longValuesMV;
+ }
int length = valueBlock.getNumDocs();
initLongValuesMV(length);
getNonLiteralValues(valueBlock);
@@ -323,7 +411,13 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
for (int j = 0; j < _numNonLiteralArguments; j++) {
_scalarArguments[_nonLiteralIndices[j]] = _nonLiteralValues[j][i];
}
- _longValuesMV[i] = (long[])
_resultType.toInternal(_functionInvoker.invoke(_scalarArguments));
+ Object value = _functionInvoker.invoke(_scalarArguments);
+ if (value != null) {
+ _longValuesMV[i] = (long[]) _resultType.toInternal(value);
+ } else {
+ _longValuesMV[i] = NullValuePlaceHolder.LONG_ARRAY;
+ addNullResult(i);
+ }
}
return _longValuesMV;
}
@@ -333,6 +427,9 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
if (_resultMetadata.getDataType().getStoredType() != DataType.FLOAT) {
return super.transformToFloatValuesMV(valueBlock);
}
+ if (valueBlock == _evaluatedValueBlock) {
+ return _floatValuesMV;
+ }
int length = valueBlock.getNumDocs();
initFloatValuesMV(length);
getNonLiteralValues(valueBlock);
@@ -340,7 +437,13 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
for (int j = 0; j < _numNonLiteralArguments; j++) {
_scalarArguments[_nonLiteralIndices[j]] = _nonLiteralValues[j][i];
}
- _floatValuesMV[i] = (float[])
_resultType.toInternal(_functionInvoker.invoke(_scalarArguments));
+ Object value = _functionInvoker.invoke(_scalarArguments);
+ if (value != null) {
+ _floatValuesMV[i] = (float[]) _resultType.toInternal(value);
+ } else {
+ _floatValuesMV[i] = NullValuePlaceHolder.FLOAT_ARRAY;
+ addNullResult(i);
+ }
}
return _floatValuesMV;
}
@@ -350,6 +453,9 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
if (_resultMetadata.getDataType().getStoredType() != DataType.DOUBLE) {
return super.transformToDoubleValuesMV(valueBlock);
}
+ if (valueBlock == _evaluatedValueBlock) {
+ return _doubleValuesMV;
+ }
int length = valueBlock.getNumDocs();
initDoubleValuesMV(length);
getNonLiteralValues(valueBlock);
@@ -357,7 +463,13 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
for (int j = 0; j < _numNonLiteralArguments; j++) {
_scalarArguments[_nonLiteralIndices[j]] = _nonLiteralValues[j][i];
}
- _doubleValuesMV[i] = (double[])
_resultType.toInternal(_functionInvoker.invoke(_scalarArguments));
+ Object value = _functionInvoker.invoke(_scalarArguments);
+ if (value != null) {
+ _doubleValuesMV[i] = (double[]) _resultType.toInternal(value);
+ } else {
+ _doubleValuesMV[i] = NullValuePlaceHolder.DOUBLE_ARRAY;
+ addNullResult(i);
+ }
}
return _doubleValuesMV;
}
@@ -367,6 +479,9 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
if (_resultMetadata.getDataType().getStoredType() != DataType.STRING) {
return super.transformToStringValuesMV(valueBlock);
}
+ if (valueBlock == _evaluatedValueBlock) {
+ return _stringValuesMV;
+ }
int length = valueBlock.getNumDocs();
initStringValuesMV(length);
getNonLiteralValues(valueBlock);
@@ -374,11 +489,91 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
for (int j = 0; j < _numNonLiteralArguments; j++) {
_scalarArguments[_nonLiteralIndices[j]] = _nonLiteralValues[j][i];
}
- _stringValuesMV[i] = (String[])
_resultType.toInternal(_functionInvoker.invoke(_scalarArguments));
+ Object value = _functionInvoker.invoke(_scalarArguments);
+ if (value != null) {
+ _stringValuesMV[i] = (String[]) _resultType.toInternal(value);
+ } else {
+ _stringValuesMV[i] = NullValuePlaceHolder.STRING_ARRAY;
+ addNullResult(i);
+ }
}
return _stringValuesMV;
}
+ @Nullable
+ @Override
+ public RoaringBitmap getNullBitmap(ValueBlock valueBlock) {
+ if (!_invokeWithNulls) {
+ return super.getNullBitmap(valueBlock);
+ }
+ if (valueBlock != _evaluatedValueBlock) {
+ evaluate(valueBlock);
+ }
+ return _resultNullBitmap;
+ }
+
+ /// Evaluates the function on the given value block with the transform
method of the result type, which also records
+ /// the null rows of the result.
+ private void evaluate(ValueBlock valueBlock) {
+ DataType storedType = _resultMetadata.getDataType().getStoredType();
+ if (_resultMetadata.isSingleValue()) {
+ switch (storedType) {
+ case INT:
+ transformToIntValuesSV(valueBlock);
+ break;
+ case LONG:
+ transformToLongValuesSV(valueBlock);
+ break;
+ case FLOAT:
+ transformToFloatValuesSV(valueBlock);
+ break;
+ case DOUBLE:
+ transformToDoubleValuesSV(valueBlock);
+ break;
+ case BIG_DECIMAL:
+ transformToBigDecimalValuesSV(valueBlock);
+ break;
+ case STRING:
+ transformToStringValuesSV(valueBlock);
+ break;
+ case BYTES:
+ transformToBytesValuesSV(valueBlock);
+ break;
+ default:
+ throw new IllegalStateException("Unsupported single-value result
type: " + storedType);
+ }
+ } else {
+ switch (storedType) {
+ case INT:
+ transformToIntValuesMV(valueBlock);
+ break;
+ case LONG:
+ transformToLongValuesMV(valueBlock);
+ break;
+ case FLOAT:
+ transformToFloatValuesMV(valueBlock);
+ break;
+ case DOUBLE:
+ transformToDoubleValuesMV(valueBlock);
+ break;
+ case STRING:
+ transformToStringValuesMV(valueBlock);
+ break;
+ default:
+ throw new IllegalStateException("Unsupported multi-value result
type: " + storedType);
+ }
+ }
+ }
+
+ private void addNullResult(int row) {
+ if (_invokeWithNulls) {
+ if (_resultNullBitmap == null) {
+ _resultNullBitmap = new RoaringBitmap();
+ }
+ _resultNullBitmap.add(row);
+ }
+ }
+
/// Returns the type to read an argument for an `Object` parameter, so that
the function gets the external Java value
/// of the argument result type, the same as `ColumnDataType.toExternal()`
in the multi-stage engine `FunctionOperand`
/// for that type. Complex types such as MAP are not supported.
@@ -416,6 +611,9 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
/// Helper method to fetch values for the non-literal transform functions
based on the parameter types. For an
/// `Object` parameter, the values are read based on the argument result
type (see [#getObjectArgumentType]).
+ ///
+ /// When invoking with nulls, also replaces the values of the null rows with
`null`, and starts recording the null
+ /// rows of the result for the value block.
private void getNonLiteralValues(ValueBlock valueBlock) {
for (int i = 0; i < _numNonLiteralArguments; i++) {
PinotDataType parameterType = _nonLiteralTypes[i];
@@ -545,5 +743,21 @@ public class ScalarTransformFunctionWrapper extends
BaseTransformFunction {
throw new IllegalStateException("Unsupported parameter type: " +
parameterType);
}
}
+ if (_invokeWithNulls) {
+ for (int i = 0; i < _numNonLiteralArguments; i++) {
+ RoaringBitmap nullBitmap =
_nonLiteralFunctions[i].getNullBitmap(valueBlock);
+ if (nullBitmap != null && !nullBitmap.isEmpty()) {
+ // The values can be the buffer of the argument, so replace the null
rows in a copy
+ Object[] values = _nonLiteralValues[i].clone();
+ PeekableIntIterator iterator = nullBitmap.getIntIterator();
+ while (iterator.hasNext()) {
+ values[iterator.next()] = null;
+ }
+ _nonLiteralValues[i] = values;
+ }
+ }
+ _evaluatedValueBlock = valueBlock;
+ _resultNullBitmap = null;
+ }
}
}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapperTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapperTest.java
index 90ab19eea37..24e07be6525 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapperTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapperTest.java
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.core.operator.transform.function;
+import com.google.common.collect.Maps;
import it.unimi.dsi.fastutil.ints.IntLinkedOpenHashSet;
import it.unimi.dsi.fastutil.objects.ObjectLinkedOpenHashSet;
import java.nio.charset.StandardCharsets;
@@ -40,8 +41,11 @@ import
org.apache.pinot.common.function.scalar.StringFunctions;
import org.apache.pinot.common.function.scalar.uuid.UuidToStringScalarFunction;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.RequestContextUtils;
+import org.apache.pinot.core.operator.ColumnContext;
import org.apache.pinot.core.operator.blocks.ValueBlock;
import org.apache.pinot.core.operator.transform.TransformResultMetadata;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import
org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
import org.apache.pinot.spi.annotations.ScalarFunction;
import org.apache.pinot.spi.data.FieldSpec.DataType;
import org.apache.pinot.spi.utils.ArrayCopyUtils;
@@ -54,6 +58,7 @@ import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertTrue;
@@ -1305,6 +1310,66 @@ public class ScalarTransformFunctionWrapperTest extends
BaseTransformFunctionTes
testTransformFunctionWithNull(transformFunction, expectedValues, bitmap);
}
+ @Test
+ public void testNullArgumentsWithNullHandling() {
+ String expression = String.format("prefixesWithPrefix(%s, 3, %s)",
STRING_SV_COLUMN, STRING_SV_NULL_COLUMN);
+ TransformFunction transformFunction =
getTransformFunctionWithNullHandling(expression);
+ assertTrue(transformFunction instanceof ScalarTransformFunctionWrapper);
+ assertNull(transformFunction.getNullBitmap(_projectionBlock));
+ String[][] values =
transformFunction.transformToStringValuesMV(_projectionBlock);
+ for (int i = 0; i < NUM_ROWS; i++) {
+ String prefix = isNullRow(i) ? null : _stringSVValues[i];
+ assertEquals(values[i],
StringFunctions.prefixesWithPrefix(_stringSVValues[i], 3, prefix));
+ }
+
+ // Without null handling, a null argument makes the row null
+ transformFunction =
TransformFunctionFactory.get(RequestContextUtils.getExpression(expression),
_dataSourceMap);
+ assertEquals(transformFunction.getNullBitmap(_projectionBlock),
getNullRows());
+ }
+
+ @Test
+ public void testNullLiteralArgumentWithNullHandling() {
+ TransformFunction transformFunction =
+
getTransformFunctionWithNullHandling(String.format("prefixesWithPrefix(%s, 3,
null)", STRING_SV_COLUMN));
+ assertNull(transformFunction.getNullBitmap(_projectionBlock));
+ String[][] values =
transformFunction.transformToStringValuesMV(_projectionBlock);
+ for (int i = 0; i < NUM_ROWS; i++) {
+ assertEquals(values[i], StringFunctions.prefixes(_stringSVValues[i], 3));
+ }
+ }
+
+ @Test
+ public void testNullResultsWithNullHandling() {
+ TransformFunction transformFunction =
+
getTransformFunctionWithNullHandling(String.format("prefixesWithPrefix(%s, 3,
'^')", STRING_SV_NULL_COLUMN));
+ String[][] values =
transformFunction.transformToStringValuesMV(_projectionBlock);
+ for (int i = 0; i < NUM_ROWS; i++) {
+ String[] expectedValues = isNullRow(i) ?
NullValuePlaceHolder.STRING_ARRAY
+ : StringFunctions.prefixesWithPrefix(_stringSVValues[i], 3, "^");
+ assertEquals(values[i], expectedValues);
+ }
+ assertEquals(transformFunction.getNullBitmap(_projectionBlock),
getNullRows());
+ }
+
+ private TransformFunction getTransformFunctionWithNullHandling(String
expression) {
+ Map<String, ColumnContext> columnContextMap =
Maps.newHashMapWithExpectedSize(_dataSourceMap.size());
+ _dataSourceMap.forEach(
+ (column, dataSource) -> columnContextMap.put(column,
ColumnContext.fromDataSource(dataSource)));
+ QueryContext queryContext =
+ QueryContextConverterUtils.getQueryContext("SET
enableNullHandling=true; SELECT * FROM testTable");
+ return
TransformFunctionFactory.get(RequestContextUtils.getExpression(expression),
columnContextMap, queryContext);
+ }
+
+ private RoaringBitmap getNullRows() {
+ RoaringBitmap nullRows = new RoaringBitmap();
+ for (int i = 0; i < NUM_ROWS; i++) {
+ if (isNullRow(i)) {
+ nullRows.add(i);
+ }
+ }
+ return nullRows;
+ }
+
@Test
public void testCbrtTransformFunction() {
ExpressionContext expression =
diff --git
a/pinot-core/src/test/java/org/apache/pinot/queries/NullHandlingEnabledQueriesTest.java
b/pinot-core/src/test/java/org/apache/pinot/queries/NullHandlingEnabledQueriesTest.java
index cb48a2e910a..b9e1cff051e 100644
---
a/pinot-core/src/test/java/org/apache/pinot/queries/NullHandlingEnabledQueriesTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/queries/NullHandlingEnabledQueriesTest.java
@@ -333,6 +333,30 @@ public class NullHandlingEnabledQueriesTest extends
BaseQueriesTest {
COLUMN1, COLUMN2)).get(0)[0], (long) NUM_OF_SEGMENT_COPIES);
}
+ /// A scalar function that accepts null arguments gets null for the null
row, and returns its default value there,
+ /// also when nested inside another function.
+ @Test
+ public void testScalarFunctionWithNullableParameters()
+ throws Exception {
+ initializeRows();
+ insertRow("{\"a\":\"x\",\"b\":1}");
+ insertRow(null);
+ Schema schema = schemaBuilder().addSingleValueDimension(COLUMN1,
DataType.STRING).build();
+ setUpSegments(TABLE_CONFIG, schema);
+
+ List<Object[]> rows = getRows(String.format(
+ "SELECT jsonPathString(%s, '$.a', 'none'), COUNT(*) FROM testTable
GROUP BY 1 ORDER BY 1", COLUMN1));
+ assertEquals(rows.size(), 2);
+ assertEquals(rows.get(0), new Object[]{"none", (long)
NUM_OF_SEGMENT_COPIES});
+ assertEquals(rows.get(1), new Object[]{"x", (long) NUM_OF_SEGMENT_COPIES});
+
+ rows = getRows(String.format(
+ "SELECT jsonPathLong(%s, '$.b', 5) + 0, COUNT(*) FROM testTable GROUP
BY 1 ORDER BY 1", COLUMN1));
+ assertEquals(rows.size(), 2);
+ assertEquals(rows.get(0), new Object[]{1.0, (long) NUM_OF_SEGMENT_COPIES});
+ assertEquals(rows.get(1), new Object[]{5.0, (long) NUM_OF_SEGMENT_COPIES});
+ }
+
/// A multi-value column's null row is stored as a single default value, and
its posting lists overlap across values,
/// so the count is taken over their union.
@Test
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
