Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Simplify array functions #18094

Merged
merged 4 commits into from
Jun 30, 2023
Merged

Simplify array functions #18094

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d6aa3e3
Add BLOCK_BUILDER return convention
dain May 19, 2023
10bff93
Add READ_VALUE operator to types
dain Oct 12, 2022
b2af78d
Convert array functions to use READ_VALUE operator
dain Oct 13, 2022
7b6a8c3
Migrate functions to BufferedArrayValueBuilder
dain Jun 28, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions core/trino-main/src/main/java/io/trino/metadata/FunctionManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import io.trino.connector.system.GlobalSystemConnector;
import io.trino.spi.TrinoException;
import io.trino.spi.block.Block;
import io.trino.spi.block.BlockBuilder;
import io.trino.spi.connector.CatalogHandle;
import io.trino.spi.connector.ConnectorSession;
import io.trino.spi.function.AggregationImplementation;
Expand Down Expand Up @@ -192,6 +193,7 @@ private static void verifyMethodHandleSignature(BoundSignature boundSignature, S
.mapToInt(InvocationArgumentConvention::getParameterCount)
.sum();
expectedParameterCount += methodType.parameterList().stream().filter(ConnectorSession.class::equals).count();
expectedParameterCount += convention.getReturnConvention().getParameterCount();
if (scalarFunctionImplementation.getInstanceFactory().isPresent()) {
expectedParameterCount++;
}
Expand Down Expand Up @@ -262,6 +264,12 @@ private static void verifyMethodHandleSignature(BoundSignature boundSignature, S
verifyFunctionSignature(methodType.returnType().isAssignableFrom(wrap(returnType.getJavaType())),
"Expected return type to be %s, but is %s", returnType.getJavaType(), wrap(methodType.returnType()));
break;
case BLOCK_BUILDER:
verifyFunctionSignature(methodType.lastParameterType().equals(BlockBuilder.class),
"Expected last argument type to be BlockBuilder, but is %s", methodType.lastParameterType());
verifyFunctionSignature(methodType.returnType().equals(void.class),
"Expected return type to be void, but is %s", methodType.returnType());
break;
default:
throw new UnsupportedOperationException("Unknown return convention: " + convention.getReturnConvention());
}
Expand Down
4 changes: 4 additions & 0 deletions core/trino-main/src/main/java/io/trino/metadata/PolymorphicScalarFunction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,9 @@ private static boolean matchesParameterAndReturnTypes(
}
methodParameterIndex += argumentConvention.getParameterCount();
}
if (returnConvention == InvocationReturnConvention.BLOCK_BUILDER) {
throw new UnsupportedOperationException("BLOCK_BUILDER return convention is not yet supported");
}
return method.getReturnType().equals(getNullAwareContainerType(boundSignature.getReturnType().getJavaType(), returnConvention));
}

Expand Down Expand Up @@ -174,6 +177,7 @@ private static Class<?> getNullAwareContainerType(Class<?> clazz, InvocationRetu
return switch (returnConvention) {
case NULLABLE_RETURN -> Primitives.wrap(clazz);
case DEFAULT_ON_NULL, FAIL_ON_NULL -> clazz;
case BLOCK_BUILDER -> void.class;
};
}

Expand Down
2 changes: 2 additions & 0 deletions core/trino-main/src/main/java/io/trino/metadata/SystemFunctionBundle.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,7 @@
import io.trino.operator.scalar.GenericIndeterminateOperator;
import io.trino.operator.scalar.GenericLessThanOperator;
import io.trino.operator.scalar.GenericLessThanOrEqualOperator;
import io.trino.operator.scalar.GenericReadValueOperator;
import io.trino.operator.scalar.GenericXxHash64Operator;
import io.trino.operator.scalar.HmacFunctions;
import io.trino.operator.scalar.HyperLogLogFunctions;
Expand Down Expand Up @@ -569,6 +570,7 @@ public static FunctionBundle create(FeaturesConfig featuresConfig, TypeOperators
.functions(MAP_FILTER_FUNCTION, new MapTransformKeysFunction(blockTypeOperators), MAP_TRANSFORM_VALUES_FUNCTION)
.function(FORMAT_FUNCTION)
.function(TRY_CAST)
.function(new GenericReadValueOperator(typeOperators))
.function(new GenericEqualOperator(typeOperators))
.function(new GenericHashCodeOperator(typeOperators))
.function(new GenericXxHash64Operator(typeOperators))
Expand Down
4 changes: 4 additions & 0 deletions core/trino-main/src/main/java/io/trino/operator/annotations/FunctionsParserHelper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@
import static io.trino.spi.function.OperatorType.IS_DISTINCT_FROM;
import static io.trino.spi.function.OperatorType.LESS_THAN;
import static io.trino.spi.function.OperatorType.LESS_THAN_OR_EQUAL;
import static io.trino.spi.function.OperatorType.READ_VALUE;
import static io.trino.spi.function.OperatorType.XX_HASH_64;
import static io.trino.sql.analyzer.TypeSignatureTranslator.parseTypeSignature;
import static java.lang.String.CASE_INSENSITIVE_ORDER;
Expand Down Expand Up @@ -119,6 +120,9 @@ else if (ORDERABLE_TYPE_OPERATORS.contains(operator)) {
verifyTypeSignatureDoesNotContainAnyTypeParameters(typeSignature, typeSignature, typeParameterNames);
}
}
else if (operator == READ_VALUE) {
verifyOperatorSignature(operator, argumentTypes);
}
else {
throw new IllegalArgumentException("Operator dependency on " + operator + " is not allowed");
}
Expand Down
107 changes: 11 additions & 96 deletions core/trino-main/src/main/java/io/trino/operator/scalar/ArrayAllMatchFunction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,20 @@
package io.trino.operator.scalar;

import io.trino.spi.block.Block;
import io.trino.spi.function.Convention;
import io.trino.spi.function.Description;
import io.trino.spi.function.OperatorDependency;
import io.trino.spi.function.ScalarFunction;
import io.trino.spi.function.SqlNullable;
import io.trino.spi.function.SqlType;
import io.trino.spi.function.TypeParameter;
import io.trino.spi.function.TypeParameterSpecialization;
import io.trino.spi.type.StandardTypes;
import io.trino.spi.type.Type;

import java.lang.invoke.MethodHandle;

import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.BLOCK_POSITION_NOT_NULL;
import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL;
import static io.trino.spi.function.OperatorType.READ_VALUE;
import static java.lang.Boolean.FALSE;

@Description("Returns true if all elements of the array match the given predicate")
Expand All @@ -32,110 +37,20 @@ public final class ArrayAllMatchFunction
private ArrayAllMatchFunction() {}

@TypeParameter("T")
@TypeParameterSpecialization(name = "T", nativeContainerType = Object.class)
@SqlType(StandardTypes.BOOLEAN)
@SqlNullable
public static Boolean allMatchObject(
@TypeParameter("T") Type elementType,
public static Boolean allMatch(
@OperatorDependency(operator = READ_VALUE, argumentTypes = "T", convention = @Convention(arguments = BLOCK_POSITION_NOT_NULL, result = FAIL_ON_NULL)) MethodHandle readValue,
@SqlType("array(T)") Block arrayBlock,
@SqlType("function(T, boolean)") ObjectToBooleanFunction function)
throws Throwable
{
boolean hasNullResult = false;
int positionCount = arrayBlock.getPositionCount();
for (int i = 0; i < positionCount; i++) {
Object element = null;
if (!arrayBlock.isNull(i)) {
element = elementType.getObject(arrayBlock, i);
}
Boolean match = function.apply(element);
if (FALSE.equals(match)) {
return false;
}
if (match == null) {
hasNullResult = true;
}
}
if (hasNullResult) {
return null;
}
return true;
}

@TypeParameter("T")
@TypeParameterSpecialization(name = "T", nativeContainerType = long.class)
@SqlType(StandardTypes.BOOLEAN)
@SqlNullable
public static Boolean allMatchLong(
@TypeParameter("T") Type elementType,
@SqlType("array(T)") Block arrayBlock,
@SqlType("function(T, boolean)") LongToBooleanFunction function)
{
boolean hasNullResult = false;
int positionCount = arrayBlock.getPositionCount();
for (int i = 0; i < positionCount; i++) {
Long element = null;
if (!arrayBlock.isNull(i)) {
element = elementType.getLong(arrayBlock, i);
}
Boolean match = function.apply(element);
if (FALSE.equals(match)) {
return false;
}
if (match == null) {
hasNullResult = true;
}
}
if (hasNullResult) {
return null;
}
return true;
}

@TypeParameter("T")
@TypeParameterSpecialization(name = "T", nativeContainerType = double.class)
@SqlType(StandardTypes.BOOLEAN)
@SqlNullable
public static Boolean allMatchDouble(
@TypeParameter("T") Type elementType,
@SqlType("array(T)") Block arrayBlock,
@SqlType("function(T, boolean)") DoubleToBooleanFunction function)
{
boolean hasNullResult = false;
int positionCount = arrayBlock.getPositionCount();
for (int i = 0; i < positionCount; i++) {
Double element = null;
if (!arrayBlock.isNull(i)) {
element = elementType.getDouble(arrayBlock, i);
}
Boolean match = function.apply(element);
if (FALSE.equals(match)) {
return false;
}
if (match == null) {
hasNullResult = true;
}
}
if (hasNullResult) {
return null;
}
return true;
}

@TypeParameter("T")
@TypeParameterSpecialization(name = "T", nativeContainerType = boolean.class)
@SqlType(StandardTypes.BOOLEAN)
@SqlNullable
public static Boolean allMatchBoolean(
@TypeParameter("T") Type elementType,
@SqlType("array(T)") Block arrayBlock,
@SqlType("function(T, boolean)") BooleanToBooleanFunction function)
{
boolean hasNullResult = false;
int positionCount = arrayBlock.getPositionCount();
for (int i = 0; i < positionCount; i++) {
Boolean element = null;
if (!arrayBlock.isNull(i)) {
element = elementType.getBoolean(arrayBlock, i);
element = readValue.invoke(arrayBlock, i);
}
Boolean match = function.apply(element);
if (FALSE.equals(match)) {
Expand Down
107 changes: 11 additions & 96 deletions core/trino-main/src/main/java/io/trino/operator/scalar/ArrayAnyMatchFunction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,20 @@
package io.trino.operator.scalar;

import io.trino.spi.block.Block;
import io.trino.spi.function.Convention;
import io.trino.spi.function.Description;
import io.trino.spi.function.OperatorDependency;
import io.trino.spi.function.ScalarFunction;
import io.trino.spi.function.SqlNullable;
import io.trino.spi.function.SqlType;
import io.trino.spi.function.TypeParameter;
import io.trino.spi.function.TypeParameterSpecialization;
import io.trino.spi.type.StandardTypes;
import io.trino.spi.type.Type;

import java.lang.invoke.MethodHandle;

import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.BLOCK_POSITION_NOT_NULL;
import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL;
import static io.trino.spi.function.OperatorType.READ_VALUE;
import static java.lang.Boolean.TRUE;

@Description("Returns true if the array contains one or more elements that match the given predicate")
Expand All @@ -32,110 +37,20 @@ public final class ArrayAnyMatchFunction
private ArrayAnyMatchFunction() {}

@TypeParameter("T")
@TypeParameterSpecialization(name = "T", nativeContainerType = Object.class)
@SqlType(StandardTypes.BOOLEAN)
@SqlNullable
public static Boolean anyMatchObject(
@TypeParameter("T") Type elementType,
public static Boolean anyMatch(
@OperatorDependency(operator = READ_VALUE, argumentTypes = "T", convention = @Convention(arguments = BLOCK_POSITION_NOT_NULL, result = FAIL_ON_NULL)) MethodHandle readValue,
@SqlType("array(T)") Block arrayBlock,
@SqlType("function(T, boolean)") ObjectToBooleanFunction function)
throws Throwable
{
boolean hasNullResult = false;
int positionCount = arrayBlock.getPositionCount();
for (int i = 0; i < positionCount; i++) {
Object element = null;
if (!arrayBlock.isNull(i)) {
element = elementType.getObject(arrayBlock, i);
}
Boolean match = function.apply(element);
if (TRUE.equals(match)) {
return true;
}
if (match == null) {
hasNullResult = true;
}
}
if (hasNullResult) {
return null;
}
return false;
}

@TypeParameter("T")
@TypeParameterSpecialization(name = "T", nativeContainerType = long.class)
@SqlType(StandardTypes.BOOLEAN)
@SqlNullable
public static Boolean anyMatchLong(
@TypeParameter("T") Type elementType,
@SqlType("array(T)") Block arrayBlock,
@SqlType("function(T, boolean)") LongToBooleanFunction function)
{
boolean hasNullResult = false;
int positionCount = arrayBlock.getPositionCount();
for (int i = 0; i < positionCount; i++) {
Long element = null;
if (!arrayBlock.isNull(i)) {
element = elementType.getLong(arrayBlock, i);
}
Boolean match = function.apply(element);
if (TRUE.equals(match)) {
return true;
}
if (match == null) {
hasNullResult = true;
}
}
if (hasNullResult) {
return null;
}
return false;
}

@TypeParameter("T")
@TypeParameterSpecialization(name = "T", nativeContainerType = double.class)
@SqlType(StandardTypes.BOOLEAN)
@SqlNullable
public static Boolean anyMatchDouble(
@TypeParameter("T") Type elementType,
@SqlType("array(T)") Block arrayBlock,
@SqlType("function(T, boolean)") DoubleToBooleanFunction function)
{
boolean hasNullResult = false;
int positionCount = arrayBlock.getPositionCount();
for (int i = 0; i < positionCount; i++) {
Double element = null;
if (!arrayBlock.isNull(i)) {
element = elementType.getDouble(arrayBlock, i);
}
Boolean match = function.apply(element);
if (TRUE.equals(match)) {
return true;
}
if (match == null) {
hasNullResult = true;
}
}
if (hasNullResult) {
return null;
}
return false;
}

@TypeParameter("T")
@TypeParameterSpecialization(name = "T", nativeContainerType = boolean.class)
@SqlType(StandardTypes.BOOLEAN)
@SqlNullable
public static Boolean anyMatchBoolean(
@TypeParameter("T") Type elementType,
@SqlType("array(T)") Block arrayBlock,
@SqlType("function(T, boolean)") BooleanToBooleanFunction function)
{
boolean hasNullResult = false;
int positionCount = arrayBlock.getPositionCount();
for (int i = 0; i < positionCount; i++) {
Boolean element = null;
if (!arrayBlock.isNull(i)) {
element = elementType.getBoolean(arrayBlock, i);
element = readValue.invoke(arrayBlock, i);
}
Boolean match = function.apply(element);
if (TRUE.equals(match)) {
Expand Down
Loading