Skip to content

Commit 2c2555f

Browse files
kevintang2022kevintang2022
committed
fb_reshape_row always push down subfields
Summary: T224244100 For fb_reshape_row udf, Presto should pushdown the utilized subfields of a struct during the optimizer phase. Presto version 0.293-20250521.210824-350 Differential Revision: D74738214
1 parent 274e05a commit 2c2555f

File tree

18 files changed

+219
-12
lines changed

18 files changed

+219
-12
lines changed

‎presto-expressions/src/main/java/com/facebook/presto/expressions/translator/TranslatorAnnotationParser.java

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,7 @@ private static class ScalarTranslationHeader
208208
private final Optional<OperatorType> operatorType;
209209
private final boolean deterministic;
210210
private final boolean calledOnNullInput;
211+
private final int pushdownSubfieldArgIndex;
211212

212213
public static List<ScalarTranslationHeader> fromAnnotatedElement(AnnotatedElement annotated)
213214
{
@@ -218,37 +219,39 @@ public static List<ScalarTranslationHeader> fromAnnotatedElement(AnnotatedElemen
218219

219220
if (scalarFunction != null) {
220221
String baseName = scalarFunction.value().isEmpty() ? camelToSnake(annotatedName(annotated)) : scalarFunction.value();
221-
builder.add(new ScalarTranslationHeader(baseName, scalarFunction.deterministic(), scalarFunction.calledOnNullInput()));
222+
builder.add(new ScalarTranslationHeader(baseName, scalarFunction.deterministic(), scalarFunction.calledOnNullInput(), scalarFunction.pushdownSubfieldArgIndex()));
222223

223224
for (String alias : scalarFunction.alias()) {
224-
builder.add(new ScalarTranslationHeader(alias, scalarFunction.deterministic(), scalarFunction.calledOnNullInput()));
225+
builder.add(new ScalarTranslationHeader(alias, scalarFunction.deterministic(), scalarFunction.calledOnNullInput(), scalarFunction.pushdownSubfieldArgIndex()));
225226
}
226227
}
227228

228229
if (scalarOperator != null) {
229-
builder.add(new ScalarTranslationHeader(scalarOperator.value(), true, scalarOperator.value().isCalledOnNullInput()));
230+
builder.add(new ScalarTranslationHeader(scalarOperator.value(), true, scalarOperator.value().isCalledOnNullInput(), -1));
230231
}
231232

232233
List<ScalarTranslationHeader> result = builder.build();
233234
checkArgument(!result.isEmpty());
234235
return result;
235236
}
236237

237-
private ScalarTranslationHeader(String name, boolean deterministic, boolean calledOnNullInput)
238+
private ScalarTranslationHeader(String name, boolean deterministic, boolean calledOnNullInput, int pushdownSubfieldArgIndex)
238239
{
239240
// TODO This is a hack. Engine should provide an API for connectors to overwrite functions. Connector should not hard code the builtin function namespace.
240241
this.name = requireNonNull(QualifiedObjectName.valueOf("presto", "default", name));
241242
this.operatorType = Optional.empty();
242243
this.deterministic = deterministic;
243244
this.calledOnNullInput = calledOnNullInput;
245+
this.pushdownSubfieldArgIndex = pushdownSubfieldArgIndex;
244246
}
245247

246-
private ScalarTranslationHeader(OperatorType operatorType, boolean deterministic, boolean calledOnNullInput)
248+
private ScalarTranslationHeader(OperatorType operatorType, boolean deterministic, boolean calledOnNullInput, int pushdownSubfieldArgIndex)
247249
{
248250
this.name = operatorType.getFunctionName();
249251
this.operatorType = Optional.of(operatorType);
250252
this.deterministic = deterministic;
251253
this.calledOnNullInput = calledOnNullInput;
254+
this.pushdownSubfieldArgIndex = pushdownSubfieldArgIndex;
252255
}
253256

254257
private static String annotatedName(AnnotatedElement annotatedElement)

‎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-hive/src/test/java/com/facebook/presto/hive/TestLambdaSubfieldPruning.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -161,4 +161,10 @@ private void testPushDownSubfieldsFromLambdas(String tableName)
161161
assertQuery("SELECT TRANSFORM_VALUES(MAP_REMOVE_NULL_VALUES(row_with_map_varchar_key_row_value.map_varchar_key_row_value), (k,v) -> v.orderkey) FROM " + tableName);
162162
assertQuery("SELECT TRANSFORM_VALUES(MAP_SUBSET(row_with_map_varchar_key_row_value.map_varchar_key_row_value, ARRAY['orderdata_ex']), (k,v) -> v.orderkey) FROM " + tableName);
163163
}
164+
165+
public void testPushdownSubfieldArgIndexForScalar()
166+
{
167+
assertQuery("SELECT struct.a FROM (SELECT CUSTOM_STRUCT_WITH_PASSTHROUGH(CAST(ROW(orderkey, comment) AS ROW(a BIGINT, b VARCHAR))) AS struct FROM lineitem)");
168+
assertQuery("SELECT struct.a FROM (SELECT CUSTOM_STRUCT_WITHOUT_PASSTHROUGH(CAST(ROW(orderkey, comment) AS ROW(a BIGINT, b VARCHAR))) AS struct FROM lineitem)");
169+
}
164170
}

‎presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1187,6 +1187,7 @@ else if (function instanceof SqlInvokedFunction) {
11871187
JAVA,
11881188
function.isDeterministic(),
11891189
function.isCalledOnNullInput(),
1190+
function.getPushdownSubfieldArgIndex(),
11901191
function.getComplexTypeFunctionDescriptor());
11911192
}
11921193
}

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,12 @@ public boolean isCalledOnNullInput()
6767
return details.isCalledOnNullInput();
6868
}
6969

70+
@Override
71+
public Optional<Integer> getPushdownSubfieldArgIndex()
72+
{
73+
return details.getPushdownSubfieldArgIndex();
74+
}
75+
7076
@Override
7177
public String getDescription()
7278
{

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

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,13 +23,15 @@ public class ScalarHeader
2323
private final SqlFunctionVisibility visibility;
2424
private final boolean deterministic;
2525
private final boolean calledOnNullInput;
26+
private final int pushdownSubfieldArgIndex;
2627

27-
public ScalarHeader(Optional<String> description, SqlFunctionVisibility visibility, boolean deterministic, boolean calledOnNullInput)
28+
public ScalarHeader(Optional<String> description, SqlFunctionVisibility visibility, boolean deterministic, boolean calledOnNullInput, int pushdownSubfieldArgIndex)
2829
{
2930
this.description = description;
3031
this.visibility = visibility;
3132
this.deterministic = deterministic;
3233
this.calledOnNullInput = calledOnNullInput;
34+
this.pushdownSubfieldArgIndex = pushdownSubfieldArgIndex;
3335
}
3436

3537
public Optional<String> getDescription()
@@ -51,4 +53,12 @@ public boolean isCalledOnNullInput()
5153
{
5254
return calledOnNullInput;
5355
}
56+
57+
public Optional<Integer> getPushdownSubfieldArgIndex()
58+
{
59+
if (pushdownSubfieldArgIndex < 0) {
60+
return Optional.empty();
61+
}
62+
return Optional.of(pushdownSubfieldArgIndex);
63+
}
5464
}

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import com.facebook.presto.spi.function.SqlNullable;
3737
import com.facebook.presto.spi.function.SqlType;
3838
import com.facebook.presto.spi.function.TypeParameter;
39+
import com.facebook.presto.spi.function.TypeVariableConstraint;
3940
import com.facebook.presto.sql.gen.lambda.BinaryFunctionInterface;
4041
import com.facebook.presto.sql.gen.lambda.UnaryFunctionInterface;
4142
import com.google.common.collect.ImmutableList;
@@ -44,7 +45,9 @@
4445
import java.lang.invoke.MethodHandle;
4546
import java.lang.reflect.Method;
4647
import java.util.Arrays;
48+
import java.util.HashMap;
4749
import java.util.List;
50+
import java.util.Map;
4851
import java.util.Optional;
4952
import java.util.Set;
5053

@@ -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+
checkPushdownSubfieldArgIndex(signature, codegenScalarFunction.pushdownSubfieldArgIndex(), method);
129+
125130
return new SqlScalarFunction(signature)
126131
{
127132
@Override
@@ -166,6 +171,28 @@ public boolean isCalledOnNullInput()
166171
{
167172
return codegenScalarFunction.calledOnNullInput();
168173
}
174+
175+
@Override
176+
public Optional<Integer> getPushdownSubfieldArgIndex()
177+
{
178+
if (codegenScalarFunction.pushdownSubfieldArgIndex() < 0) {
179+
return Optional.empty();
180+
}
181+
return Optional.of(codegenScalarFunction.pushdownSubfieldArgIndex());
182+
}
169183
};
170184
}
185+
186+
private static void checkPushdownSubfieldArgIndex(Signature signature, int pushdownSubfieldArgIndex, Method method)
187+
{
188+
if (pushdownSubfieldArgIndex >= 0) {
189+
Map<String, String> typeConstraintMapping = new HashMap<>();
190+
for (TypeVariableConstraint constraint : signature.getTypeVariableConstraints()) {
191+
typeConstraintMapping.put(constraint.getName(), constraint.getVariadicBound());
192+
}
193+
checkCondition(signature.getArgumentTypes().size() > pushdownSubfieldArgIndex, FUNCTION_IMPLEMENTATION_ERROR, "Method [%s] has out of range pushdown subfield arg index", method);
194+
String typeVariableName = signature.getArgumentTypes().get(pushdownSubfieldArgIndex).toString();
195+
checkCondition(typeVariableName.equals(com.facebook.presto.common.type.StandardTypes.ROW) || typeConstraintMapping.get(typeVariableName).equals(com.facebook.presto.common.type.StandardTypes.ROW), FUNCTION_IMPLEMENTATION_ERROR, "Method [%s] does not have a struct or row type as pushdown subfield arg", method);
196+
}
197+
}
171198
}

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import com.facebook.presto.spi.function.Signature;
2525
import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
2626
import com.facebook.presto.spi.function.SqlType;
27+
import com.facebook.presto.spi.function.TypeVariableConstraint;
2728
import com.google.common.collect.ImmutableList;
2829
import com.google.common.collect.ImmutableSet;
2930

@@ -106,6 +107,7 @@ private static SqlScalarFunction parseParametricScalar(ScalarHeaderAndMethods sc
106107
Map<SpecializedSignature, ParametricScalarImplementation.Builder> signatures = new HashMap<>();
107108
for (Method method : scalar.getMethods()) {
108109
ParametricScalarImplementation implementation = ParametricScalarImplementation.Parser.parseImplementation(header, method, constructor);
110+
checkPushdownSubfieldArgIndex(implementation, header, method);
109111
if (!signatures.containsKey(implementation.getSpecializedSignature())) {
110112
ParametricScalarImplementation.Builder builder = new ParametricScalarImplementation.Builder(
111113
implementation.getSignature(),
@@ -155,4 +157,21 @@ public Set<Method> getMethods()
155157
return methods;
156158
}
157159
}
160+
161+
private static void checkPushdownSubfieldArgIndex(ParametricScalarImplementation implementation, ScalarImplementationHeader header, Method method)
162+
{
163+
Optional<Integer> pushdownSubfieldArgIndex = header.getHeader().getPushdownSubfieldArgIndex();
164+
if (pushdownSubfieldArgIndex.isPresent()) {
165+
Map<String, String> typeConstraintMapping = new HashMap<>();
166+
Signature signature = implementation.getSignature();
167+
for (TypeVariableConstraint constraint : signature.getTypeVariableConstraints()) {
168+
typeConstraintMapping.put(constraint.getName(), constraint.getVariadicBound());
169+
}
170+
171+
checkCondition(pushdownSubfieldArgIndex.get() >= 0, FUNCTION_IMPLEMENTATION_ERROR, "Method [%s] has negative pushdown subfield arg index", method);
172+
checkCondition(signature.getArgumentTypes().size() > pushdownSubfieldArgIndex.get(), FUNCTION_IMPLEMENTATION_ERROR, "Method [%s] has out of range pushdown subfield arg index", method);
173+
String typeVariableName = signature.getArgumentTypes().get(pushdownSubfieldArgIndex.get()).toString();
174+
checkCondition(typeVariableName.equals(com.facebook.presto.common.type.StandardTypes.ROW) || typeConstraintMapping.get(typeVariableName).equals(com.facebook.presto.common.type.StandardTypes.ROW), FUNCTION_IMPLEMENTATION_ERROR, "Method [%s] does not have a struct or row type as pushdown subfield arg", method);
175+
}
176+
}
158177
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -81,15 +81,15 @@ public static List<ScalarImplementationHeader> fromAnnotatedElement(AnnotatedEle
8181

8282
if (scalarFunction != null) {
8383
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())));
84+
builder.add(new ScalarImplementationHeader(baseName, new ScalarHeader(description, scalarFunction.visibility(), scalarFunction.deterministic(), scalarFunction.calledOnNullInput(), scalarFunction.pushdownSubfieldArgIndex())));
8585

8686
for (String alias : scalarFunction.alias()) {
87-
builder.add(new ScalarImplementationHeader(alias, new ScalarHeader(description, scalarFunction.visibility(), scalarFunction.deterministic(), scalarFunction.calledOnNullInput())));
87+
builder.add(new ScalarImplementationHeader(alias, new ScalarHeader(description, scalarFunction.visibility(), scalarFunction.deterministic(), scalarFunction.calledOnNullInput(), scalarFunction.pushdownSubfieldArgIndex())));
8888
}
8989
}
9090

9191
if (scalarOperator != null) {
92-
builder.add(new ScalarImplementationHeader(scalarOperator.value(), new ScalarHeader(description, HIDDEN, true, scalarOperator.value().isCalledOnNullInput())));
92+
builder.add(new ScalarImplementationHeader(scalarOperator.value(), new ScalarHeader(description, HIDDEN, true, scalarOperator.value().isCalledOnNullInput(), -1)));
9393
}
9494

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

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -582,6 +582,15 @@ 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+
Optional<Integer> pushdownSubfieldArgIndex = functionAndTypeManager.getFunctionMetadata(((CallExpression) expression).getFunctionHandle()).getPushdownSubfieldArgIndex();
587+
if (pushdownSubfieldArgIndex.isPresent() &&
588+
((CallExpression) expression).getArguments().size() > pushdownSubfieldArgIndex.get() &&
589+
((CallExpression) expression).getArguments().get(pushdownSubfieldArgIndex.get()).getType() instanceof RowType) {
590+
expression = ((CallExpression) expression).getArguments().get(pushdownSubfieldArgIndex.get());
591+
continue;
592+
}
593+
}
585594

586595
if (expression instanceof SpecialFormExpression && ((SpecialFormExpression) expression).getForm() == DEREFERENCE) {
587596
SpecialFormExpression dereference = (SpecialFormExpression) expression;

‎presto-spi/src/main/java/com/facebook/presto/spi/function/CodegenScalarFunction.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,4 +33,6 @@
3333
boolean deterministic() default true;
3434

3535
boolean calledOnNullInput() default false;
36+
37+
int pushdownSubfieldArgIndex() default -1;
3638
}

0 commit comments

Comments
 (0)