Skip to content

Commit 3785079

Browse files
kevintang2022kevintang2022
committed
Add function metadata ability to push down struct argument in optimizer
1 parent 274e05a commit 3785079

File tree

14 files changed

+277
-9
lines changed

14 files changed

+277
-9
lines changed

‎presto-hive/src/test/java/com/facebook/presto/hive/TestHiveLogicalPlanner.java

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,13 @@ protected QueryRunner createQueryRunner()
166166
Optional.empty());
167167
}
168168

169+
@Override
170+
protected QueryRunner createExpectedQueryRunner()
171+
throws Exception
172+
{
173+
return getQueryRunner();
174+
}
175+
169176
@Test
170177
public void testMetadataQueryOptimizationWithLimit()
171178
{
@@ -1366,6 +1373,18 @@ public void testPushdownSubfields()
13661373
assertPushdownSubfields("SELECT x.a FROM test_pushdown_struct_subfields WHERE x.a > 10 AND x.b LIKE 'abc%'", "test_pushdown_struct_subfields",
13671374
ImmutableMap.of("x", toSubfields("x.a", "x.b")));
13681375

1376+
assertQuery("SELECT struct.b FROM (SELECT CUSTOM_STRUCT_WITH_PASSTHROUGH(x) AS struct FROM test_pushdown_struct_subfields)");
1377+
assertQuery("SELECT struct.b FROM (SELECT CUSTOM_STRUCT_WITHOUT_PASSTHROUGH(x) AS struct FROM test_pushdown_struct_subfields)");
1378+
1379+
assertPushdownSubfields("SELECT struct.b FROM (SELECT CUSTOM_STRUCT_WITH_PASSTHROUGH(x) AS struct FROM test_pushdown_struct_subfields)", "test_pushdown_struct_subfields",
1380+
ImmutableMap.of("x", toSubfields("x.b")));
1381+
1382+
assertPushdownSubfields("SELECT struct.b FROM (SELECT CUSTOM_STRUCT_WITHOUT_PASSTHROUGH(x) AS struct FROM test_pushdown_struct_subfields)", "test_pushdown_struct_subfields",
1383+
ImmutableMap.of());
1384+
1385+
assertPushdownSubfields("SELECT struct.b FROM (SELECT x AS struct FROM test_pushdown_struct_subfields)", "test_pushdown_struct_subfields",
1386+
ImmutableMap.of("x", toSubfields("x.b")));
1387+
13691388
// Join
13701389
assertPlan("SELECT l.orderkey, x.a, mod(x.d.d1, 2) FROM lineitem l, test_pushdown_struct_subfields a WHERE l.linenumber = a.id",
13711390
anyTree(

‎presto-main-base/src/main/java/com/facebook/presto/operator/annotations/FunctionsParserHelper.java

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,9 @@
1515

1616
import com.facebook.presto.common.function.OperatorType;
1717
import com.facebook.presto.common.type.TypeSignature;
18+
import com.facebook.presto.spi.function.ComplexTypeFunctionDescriptor;
1819
import com.facebook.presto.spi.function.Description;
20+
import com.facebook.presto.spi.function.FunctionDescriptor;
1921
import com.facebook.presto.spi.function.IsNull;
2022
import com.facebook.presto.spi.function.LiteralParameters;
2123
import com.facebook.presto.spi.function.LongVariableConstraint;
@@ -58,12 +60,16 @@
5860
import static com.facebook.presto.common.function.OperatorType.NOT_EQUAL;
5961
import static com.facebook.presto.common.type.StandardTypes.PARAMETRIC_TYPES;
6062
import static com.facebook.presto.operator.annotations.ImplementationDependency.isImplementationDependencyAnnotation;
63+
import static com.facebook.presto.spi.StandardErrorCode.FUNCTION_IMPLEMENTATION_ERROR;
64+
import static com.facebook.presto.util.Failures.checkCondition;
6165
import static com.google.common.base.Preconditions.checkArgument;
6266
import static com.google.common.collect.ImmutableList.toImmutableList;
6367
import static com.google.common.collect.ImmutableSet.toImmutableSet;
6468
import static java.lang.reflect.Modifier.isPublic;
6569
import static java.lang.reflect.Modifier.isStatic;
6670
import static java.util.Arrays.asList;
71+
import static java.util.Collections.emptyList;
72+
import static java.util.Collections.emptySet;
6773

6874
public class FunctionsParserHelper
6975
{
@@ -254,6 +260,30 @@ public static Optional<String> parseDescription(AnnotatedElement base)
254260
return (description == null) ? Optional.empty() : Optional.of(description.value());
255261
}
256262

263+
public static ComplexTypeFunctionDescriptor parseFunctionDescriptor(AnnotatedElement base)
264+
{
265+
FunctionDescriptor descriptor = base.getAnnotation(FunctionDescriptor.class);
266+
if (descriptor == null) {
267+
return ComplexTypeFunctionDescriptor.DEFAULT;
268+
}
269+
270+
int pushdownSubfieldArgIndex = descriptor.pushdownSubfieldArgIndex();
271+
Optional<Integer> descriptorPushdownIndex;
272+
if (pushdownSubfieldArgIndex < 0) {
273+
descriptorPushdownIndex = Optional.empty();
274+
}
275+
else {
276+
descriptorPushdownIndex = Optional.of(pushdownSubfieldArgIndex);
277+
}
278+
279+
return new ComplexTypeFunctionDescriptor(
280+
true,
281+
emptyList(),
282+
Optional.of(emptySet()),
283+
Optional.of(ComplexTypeFunctionDescriptor::allSubfieldsRequired),
284+
descriptorPushdownIndex);
285+
}
286+
257287
public static List<LongVariableConstraint> parseLongVariableConstraints(Method inputFunction)
258288
{
259289
return Stream.of(inputFunction.getAnnotationsByType(Constraint.class))
@@ -277,4 +307,25 @@ public static Map<String, Class<?>> getDeclaredSpecializedTypeParameters(Method
277307
}
278308
return specializedTypeParameters;
279309
}
310+
311+
public static void checkPushdownSubfieldArgIndex(Method method, Signature signature, Optional<Integer> pushdownSubfieldArgIndex)
312+
{
313+
if (pushdownSubfieldArgIndex.isPresent()) {
314+
Map<String, TypeVariableConstraint> typeConstraintMapping = new HashMap<>();
315+
for (TypeVariableConstraint constraint : signature.getTypeVariableConstraints()) {
316+
typeConstraintMapping.put(constraint.getName(), constraint);
317+
}
318+
checkCondition(signature.getArgumentTypes().size() > pushdownSubfieldArgIndex.get(), FUNCTION_IMPLEMENTATION_ERROR, "Method [%s] has out of range pushdown subfield arg index", method);
319+
String typeVariableName = signature.getArgumentTypes().get(pushdownSubfieldArgIndex.get()).toString();
320+
321+
// The type variable must be directly a ROW type
322+
// or (it is a type alias that is not bounded by a type)
323+
// or (it is a type alias that maps to a row type)
324+
boolean meetsTypeConstraint = (!typeConstraintMapping.containsKey(typeVariableName) && typeVariableName.equals(com.facebook.presto.common.type.StandardTypes.ROW)) ||
325+
(typeConstraintMapping.containsKey(typeVariableName) && typeConstraintMapping.get(typeVariableName).getVariadicBound() == null && !typeConstraintMapping.get(typeVariableName).isNonDecimalNumericRequired()) ||
326+
(typeConstraintMapping.containsKey(typeVariableName) && typeConstraintMapping.get(typeVariableName).getVariadicBound().equals(com.facebook.presto.common.type.StandardTypes.ROW));
327+
328+
checkCondition(meetsTypeConstraint, FUNCTION_IMPLEMENTATION_ERROR, "Method [%s] does not have a struct or row type as pushdown subfield arg", method);
329+
}
330+
}
280331
}

‎presto-main-base/src/main/java/com/facebook/presto/operator/scalar/ParametricScalar.java

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import com.facebook.presto.operator.ParametricImplementationsGroup;
2020
import com.facebook.presto.operator.scalar.annotations.ParametricScalarImplementation;
2121
import com.facebook.presto.spi.PrestoException;
22+
import com.facebook.presto.spi.function.ComplexTypeFunctionDescriptor;
2223
import com.facebook.presto.spi.function.Signature;
2324
import com.facebook.presto.spi.function.SqlFunctionVisibility;
2425
import com.google.common.annotations.VisibleForTesting;
@@ -67,6 +68,12 @@ public boolean isCalledOnNullInput()
6768
return details.isCalledOnNullInput();
6869
}
6970

71+
@Override
72+
public ComplexTypeFunctionDescriptor getComplexTypeFunctionDescriptor()
73+
{
74+
return details.getComplexTypeFunctionDescriptor();
75+
}
76+
7077
@Override
7178
public String getDescription()
7279
{

‎presto-main-base/src/main/java/com/facebook/presto/operator/scalar/ScalarHeader.java

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
*/
1414
package com.facebook.presto.operator.scalar;
1515

16+
import com.facebook.presto.spi.function.ComplexTypeFunctionDescriptor;
1617
import com.facebook.presto.spi.function.SqlFunctionVisibility;
1718

1819
import java.util.Optional;
@@ -23,13 +24,15 @@ public class ScalarHeader
2324
private final SqlFunctionVisibility visibility;
2425
private final boolean deterministic;
2526
private final boolean calledOnNullInput;
27+
private final ComplexTypeFunctionDescriptor complexTypeFunctionDescriptor;
2628

27-
public ScalarHeader(Optional<String> description, SqlFunctionVisibility visibility, boolean deterministic, boolean calledOnNullInput)
29+
public ScalarHeader(Optional<String> description, SqlFunctionVisibility visibility, boolean deterministic, boolean calledOnNullInput, ComplexTypeFunctionDescriptor complexTypeFunctionDescriptor)
2830
{
2931
this.description = description;
3032
this.visibility = visibility;
3133
this.deterministic = deterministic;
3234
this.calledOnNullInput = calledOnNullInput;
35+
this.complexTypeFunctionDescriptor = complexTypeFunctionDescriptor;
3336
}
3437

3538
public Optional<String> getDescription()
@@ -51,4 +54,9 @@ public boolean isCalledOnNullInput()
5154
{
5255
return calledOnNullInput;
5356
}
57+
58+
public ComplexTypeFunctionDescriptor getComplexTypeFunctionDescriptor()
59+
{
60+
return complexTypeFunctionDescriptor;
61+
}
5462
}

‎presto-main-base/src/main/java/com/facebook/presto/operator/scalar/annotations/CodegenScalarFromAnnotationsParser.java

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import com.facebook.presto.spi.PrestoException;
2626
import com.facebook.presto.spi.function.BlockPosition;
2727
import com.facebook.presto.spi.function.CodegenScalarFunction;
28+
import com.facebook.presto.spi.function.ComplexTypeFunctionDescriptor;
2829
import com.facebook.presto.spi.function.Description;
2930
import com.facebook.presto.spi.function.FunctionKind;
3031
import com.facebook.presto.spi.function.IsNull;
@@ -51,7 +52,9 @@
5152
import static com.facebook.presto.common.type.TypeSignature.parseTypeSignature;
5253
import static com.facebook.presto.metadata.BuiltInTypeAndFunctionNamespaceManager.JAVA_BUILTIN_NAMESPACE;
5354
import static com.facebook.presto.metadata.SignatureBinder.applyBoundVariables;
55+
import static com.facebook.presto.operator.annotations.FunctionsParserHelper.checkPushdownSubfieldArgIndex;
5456
import static com.facebook.presto.operator.annotations.FunctionsParserHelper.findPublicStaticMethods;
57+
import static com.facebook.presto.operator.annotations.FunctionsParserHelper.parseFunctionDescriptor;
5558
import static com.facebook.presto.operator.scalar.ScalarFunctionImplementationChoice.ArgumentProperty.functionTypeArgumentProperty;
5659
import static com.facebook.presto.operator.scalar.ScalarFunctionImplementationChoice.ArgumentProperty.valueTypeArgumentProperty;
5760
import static com.facebook.presto.operator.scalar.ScalarFunctionImplementationChoice.NullConvention.RETURN_NULL_ON_NULL;
@@ -122,6 +125,8 @@ private static SqlScalarFunction createSqlScalarFunction(Method method)
122125
Arrays.stream(method.getParameters()).map(p -> parseTypeSignature(p.getAnnotation(SqlType.class).value())).collect(toImmutableList()),
123126
false);
124127

128+
ComplexTypeFunctionDescriptor descriptor = parseAndCheckFunctionDescriptor(method, signature);
129+
125130
return new SqlScalarFunction(signature)
126131
{
127132
@Override
@@ -166,6 +171,19 @@ public boolean isCalledOnNullInput()
166171
{
167172
return codegenScalarFunction.calledOnNullInput();
168173
}
174+
175+
@Override
176+
public ComplexTypeFunctionDescriptor getComplexTypeFunctionDescriptor()
177+
{
178+
return descriptor;
179+
}
169180
};
170181
}
182+
183+
private static ComplexTypeFunctionDescriptor parseAndCheckFunctionDescriptor(Method method, Signature signature)
184+
{
185+
ComplexTypeFunctionDescriptor descriptor = parseFunctionDescriptor(method);
186+
checkPushdownSubfieldArgIndex(method, signature, descriptor.getPushdownSubfieldArgIndex());
187+
return descriptor;
188+
}
171189
}

‎presto-main-base/src/main/java/com/facebook/presto/operator/scalar/annotations/ScalarFromAnnotationsParser.java

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import com.facebook.presto.operator.scalar.ParametricScalar;
2020
import com.facebook.presto.operator.scalar.annotations.ParametricScalarImplementation.SpecializedSignature;
2121
import com.facebook.presto.spi.function.CodegenScalarFunction;
22+
import com.facebook.presto.spi.function.FunctionDescriptor;
2223
import com.facebook.presto.spi.function.ScalarFunction;
2324
import com.facebook.presto.spi.function.ScalarOperator;
2425
import com.facebook.presto.spi.function.Signature;
@@ -35,6 +36,7 @@
3536
import java.util.Optional;
3637
import java.util.Set;
3738

39+
import static com.facebook.presto.operator.annotations.FunctionsParserHelper.checkPushdownSubfieldArgIndex;
3840
import static com.facebook.presto.operator.scalar.annotations.OperatorValidator.validateOperator;
3941
import static com.facebook.presto.spi.StandardErrorCode.FUNCTION_IMPLEMENTATION_ERROR;
4042
import static com.facebook.presto.util.Failures.checkCondition;
@@ -88,7 +90,7 @@ private static List<ScalarHeaderAndMethods> findScalarsInFunctionSetClass(Class<
8890
ImmutableList.Builder<ScalarHeaderAndMethods> builder = ImmutableList.builder();
8991
for (Method method : FunctionsParserHelper.findPublicMethods(
9092
annotated,
91-
ImmutableSet.of(SqlType.class, ScalarFunction.class, ScalarOperator.class),
93+
ImmutableSet.of(SqlType.class, ScalarFunction.class, ScalarOperator.class, FunctionDescriptor.class),
9294
ImmutableSet.of(SqlInvokedScalarFunction.class, CodegenScalarFunction.class))) {
9395
checkCondition((method.getAnnotation(ScalarFunction.class) != null) || (method.getAnnotation(ScalarOperator.class) != null),
9496
FUNCTION_IMPLEMENTATION_ERROR, "Method [%s] annotated with @SqlType is missing @ScalarFunction or @ScalarOperator", method);
@@ -106,6 +108,7 @@ private static SqlScalarFunction parseParametricScalar(ScalarHeaderAndMethods sc
106108
Map<SpecializedSignature, ParametricScalarImplementation.Builder> signatures = new HashMap<>();
107109
for (Method method : scalar.getMethods()) {
108110
ParametricScalarImplementation implementation = ParametricScalarImplementation.Parser.parseImplementation(header, method, constructor);
111+
checkPushdownSubfieldArgIndex(method, implementation.getSignature(), header.getHeader().getComplexTypeFunctionDescriptor().getPushdownSubfieldArgIndex());
109112
if (!signatures.containsKey(implementation.getSpecializedSignature())) {
110113
ParametricScalarImplementation.Builder builder = new ParametricScalarImplementation.Builder(
111114
implementation.getSignature(),

‎presto-main-base/src/main/java/com/facebook/presto/operator/scalar/annotations/ScalarImplementationHeader.java

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import com.facebook.presto.common.QualifiedObjectName;
1717
import com.facebook.presto.common.function.OperatorType;
1818
import com.facebook.presto.operator.scalar.ScalarHeader;
19+
import com.facebook.presto.spi.function.ComplexTypeFunctionDescriptor;
1920
import com.facebook.presto.spi.function.ScalarFunction;
2021
import com.facebook.presto.spi.function.ScalarOperator;
2122
import com.facebook.presto.spi.function.SqlFunctionVisibility;
@@ -28,6 +29,7 @@
2829

2930
import static com.facebook.presto.metadata.BuiltInTypeAndFunctionNamespaceManager.JAVA_BUILTIN_NAMESPACE;
3031
import static com.facebook.presto.operator.annotations.FunctionsParserHelper.parseDescription;
32+
import static com.facebook.presto.operator.annotations.FunctionsParserHelper.parseFunctionDescriptor;
3133
import static com.facebook.presto.spi.function.SqlFunctionVisibility.HIDDEN;
3234
import static com.google.common.base.CaseFormat.LOWER_CAMEL;
3335
import static com.google.common.base.CaseFormat.LOWER_UNDERSCORE;
@@ -76,20 +78,21 @@ public static List<ScalarImplementationHeader> fromAnnotatedElement(AnnotatedEle
7678
ScalarFunction scalarFunction = annotated.getAnnotation(ScalarFunction.class);
7779
ScalarOperator scalarOperator = annotated.getAnnotation(ScalarOperator.class);
7880
Optional<String> description = parseDescription(annotated);
81+
ComplexTypeFunctionDescriptor descriptor = parseFunctionDescriptor(annotated);
7982

8083
ImmutableList.Builder<ScalarImplementationHeader> builder = ImmutableList.builder();
8184

8285
if (scalarFunction != null) {
8386
String baseName = scalarFunction.value().isEmpty() ? camelToSnake(annotatedName(annotated)) : scalarFunction.value();
84-
builder.add(new ScalarImplementationHeader(baseName, new ScalarHeader(description, scalarFunction.visibility(), scalarFunction.deterministic(), scalarFunction.calledOnNullInput())));
87+
builder.add(new ScalarImplementationHeader(baseName, new ScalarHeader(description, scalarFunction.visibility(), scalarFunction.deterministic(), scalarFunction.calledOnNullInput(), descriptor)));
8588

8689
for (String alias : scalarFunction.alias()) {
87-
builder.add(new ScalarImplementationHeader(alias, new ScalarHeader(description, scalarFunction.visibility(), scalarFunction.deterministic(), scalarFunction.calledOnNullInput())));
90+
builder.add(new ScalarImplementationHeader(alias, new ScalarHeader(description, scalarFunction.visibility(), scalarFunction.deterministic(), scalarFunction.calledOnNullInput(), descriptor)));
8891
}
8992
}
9093

9194
if (scalarOperator != null) {
92-
builder.add(new ScalarImplementationHeader(scalarOperator.value(), new ScalarHeader(description, HIDDEN, true, scalarOperator.value().isCalledOnNullInput())));
95+
builder.add(new ScalarImplementationHeader(scalarOperator.value(), new ScalarHeader(description, HIDDEN, true, scalarOperator.value().isCalledOnNullInput(), descriptor)));
9396
}
9497

9598
List<ScalarImplementationHeader> result = builder.build();

‎presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PushdownSubfields.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -582,6 +582,16 @@ private static Optional<Subfield> toSubfield(
582582
if (expression instanceof VariableReferenceExpression) {
583583
return Optional.of(new Subfield(((VariableReferenceExpression) expression).getName(), elements.build().reverse()));
584584
}
585+
if (expression instanceof CallExpression) {
586+
ComplexTypeFunctionDescriptor functionDescriptor = functionAndTypeManager.getFunctionMetadata(((CallExpression) expression).getFunctionHandle()).getDescriptor();
587+
Optional<Integer> pushdownSubfieldArgIndex = functionDescriptor.getPushdownSubfieldArgIndex();
588+
if (pushdownSubfieldArgIndex.isPresent() &&
589+
((CallExpression) expression).getArguments().size() > pushdownSubfieldArgIndex.get() &&
590+
((CallExpression) expression).getArguments().get(pushdownSubfieldArgIndex.get()).getType() instanceof RowType) {
591+
expression = ((CallExpression) expression).getArguments().get(pushdownSubfieldArgIndex.get());
592+
continue;
593+
}
594+
}
585595

586596
if (expression instanceof SpecialFormExpression && ((SpecialFormExpression) expression).getForm() == DEREFERENCE) {
587597
SpecialFormExpression dereference = (SpecialFormExpression) expression;

0 commit comments

Comments
 (0)