Skip to content

Commit 0e7e4e0

Browse files
zacw7zacw7
committed
Add unit tests for index join subfields pushdown
1 parent 84ae840 commit 0e7e4e0

File tree

3 files changed

+186
-1
lines changed

3 files changed

+186
-1
lines changed

‎presto-tests/src/main/java/com/facebook/presto/tests/tpch/TpchIndexMetadata.java‎

Lines changed: 81 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,24 +13,35 @@
1313
*/
1414
package com.facebook.presto.tests.tpch;
1515

16+
import com.facebook.presto.common.block.MethodHandleUtil;
1617
import com.facebook.presto.common.predicate.NullableValue;
1718
import com.facebook.presto.common.predicate.TupleDomain;
19+
import com.facebook.presto.common.type.MapType;
1820
import com.facebook.presto.spi.ColumnHandle;
21+
import com.facebook.presto.spi.ColumnMetadata;
1922
import com.facebook.presto.spi.ConnectorResolvedIndex;
2023
import com.facebook.presto.spi.ConnectorSession;
2124
import com.facebook.presto.spi.ConnectorTableHandle;
25+
import com.facebook.presto.spi.ConnectorTableLayoutHandle;
26+
import com.facebook.presto.spi.ConnectorTableMetadata;
27+
import com.facebook.presto.spi.Constraint;
28+
import com.facebook.presto.spi.SchemaTableName;
29+
import com.facebook.presto.spi.statistics.TableStatistics;
2230
import com.facebook.presto.tpch.TpchMetadata;
2331
import com.facebook.presto.tpch.TpchTableHandle;
2432
import com.google.common.collect.ImmutableList;
2533
import com.google.common.collect.ImmutableMap;
2634
import com.google.common.collect.ImmutableSet;
2735
import com.google.common.collect.Maps;
2836

37+
import java.util.List;
2938
import java.util.Map;
3039
import java.util.Optional;
3140
import java.util.Set;
3241
import java.util.stream.Collectors;
3342

43+
import static com.facebook.presto.common.type.BigintType.BIGINT;
44+
import static com.facebook.presto.common.type.VarcharType.VARCHAR;
3445
import static com.facebook.presto.tests.tpch.TpchIndexProvider.handleToNames;
3546
import static com.google.common.base.Predicates.in;
3647
import static com.google.common.base.Predicates.not;
@@ -40,6 +51,8 @@ public class TpchIndexMetadata
4051
extends TpchMetadata
4152
{
4253
private final TpchIndexedData indexedData;
54+
// For tables in this set, add an extra map column to their metadata.
55+
private final Set<String> tableWithExtraColumn = ImmutableSet.of("orders_extra");
4356

4457
public TpchIndexMetadata(String connectorId, TpchIndexedData indexedData)
4558
{
@@ -56,6 +69,10 @@ public Optional<ConnectorResolvedIndex> resolveIndex(
5669
TupleDomain<ColumnHandle> tupleDomain)
5770
{
5871
TpchTableHandle tpchTableHandle = (TpchTableHandle) tableHandle;
72+
String tableName = tpchTableHandle.getTableName();
73+
if (tableWithExtraColumn.contains(tableName)) {
74+
tpchTableHandle = new TpchTableHandle(getOriginalTpchTableName(tableName), tpchTableHandle.getScaleFactor());
75+
}
5976

6077
// Keep the fixed values that don't overlap with the indexableColumns
6178
// Note: technically we could more efficiently utilize the overlapped columns, but this way is simpler for now
@@ -82,10 +99,73 @@ public Optional<ConnectorResolvedIndex> resolveIndex(
8299
filteredTupleDomain = TupleDomain.withColumnDomains(Maps.filterKeys(tupleDomain.getDomains().get(), not(in(fixedValues.keySet()))));
83100
}
84101
TpchIndexHandle indexHandle = new TpchIndexHandle(
85-
tpchTableHandle.getTableName(),
102+
tableName,
86103
tpchTableHandle.getScaleFactor(),
87104
lookupColumnNames,
88105
TupleDomain.fromFixedValues(fixedValues));
89106
return Optional.of(new ConnectorResolvedIndex(indexHandle, filteredTupleDomain));
90107
}
108+
109+
@Override
110+
public TpchTableHandle getTableHandle(ConnectorSession session, SchemaTableName schemaTableName)
111+
{
112+
String tableName = schemaTableName.getTableName();
113+
if (tableWithExtraColumn.contains(tableName)) {
114+
TpchTableHandle originalTableHandle = super.getTableHandle(session, new SchemaTableName(schemaTableName.getSchemaName(), getOriginalTpchTableName(tableName)));
115+
return new TpchTableHandle(schemaTableName.getTableName(), originalTableHandle.getScaleFactor());
116+
}
117+
return super.getTableHandle(session, schemaTableName);
118+
}
119+
120+
@Override
121+
public TableStatistics getTableStatistics(ConnectorSession session, ConnectorTableHandle tableHandle, Optional<ConnectorTableLayoutHandle> tableLayoutHandle, List<ColumnHandle> columnHandles, Constraint<ColumnHandle> constraint)
122+
{
123+
String tableName = ((TpchTableHandle) tableHandle).getTableName();
124+
if (tableWithExtraColumn.contains(tableName)) {
125+
TpchTableHandle originalTableHandle = new TpchTableHandle(getOriginalTpchTableName(tableName), ((TpchTableHandle) tableHandle).getScaleFactor());
126+
return super.getTableStatistics(session, originalTableHandle, tableLayoutHandle, columnHandles, constraint);
127+
}
128+
return super.getTableStatistics(session, tableHandle, tableLayoutHandle, columnHandles, constraint);
129+
}
130+
131+
@Override
132+
public ConnectorTableMetadata getTableMetadata(ConnectorSession session, ConnectorTableHandle tableHandle)
133+
{
134+
TpchTableHandle tpchTableHandle = (TpchTableHandle) tableHandle;
135+
String tableName = tpchTableHandle.getTableName();
136+
if (tableWithExtraColumn.contains(tableName)) {
137+
ConnectorTableMetadata tableMetadata = super.getTableMetadata(session, new TpchTableHandle(getOriginalTpchTableName(tableName), tpchTableHandle.getScaleFactor()));
138+
ImmutableList.Builder<ColumnMetadata> columns = ImmutableList.builder();
139+
columns.addAll(tableMetadata.getColumns());
140+
columns.add(getExtraMapColumnMetadata());
141+
return new ConnectorTableMetadata(new SchemaTableName(tableMetadata.getTable().getSchemaName(), tableName),
142+
columns.build());
143+
}
144+
return super.getTableMetadata(session, tableHandle);
145+
}
146+
147+
private ColumnMetadata getExtraMapColumnMetadata()
148+
{
149+
return ColumnMetadata.builder()
150+
.setName("data")
151+
.setType(new MapType(BIGINT,
152+
VARCHAR,
153+
MethodHandleUtil.methodHandle(TpchIndexMetadata.class, "throwUnsupportedOperation"),
154+
MethodHandleUtil.methodHandle(TpchIndexMetadata.class, "throwUnsupportedOperation")))
155+
.build();
156+
}
157+
158+
private String getOriginalTpchTableName(String tableName)
159+
{
160+
String suffix = "_extra";
161+
if (tableName != null && tableName.endsWith(suffix)) {
162+
return tableName.substring(0, tableName.length() - suffix.length());
163+
}
164+
return tableName;
165+
}
166+
167+
public static void throwUnsupportedOperation()
168+
{
169+
throw new UnsupportedOperationException();
170+
}
91171
}

‎presto-tests/src/test/java/com/facebook/presto/tests/TestNativeIndexJoinLogicalPlanner.java‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,24 +14,36 @@
1414
package com.facebook.presto.tests;
1515

1616
import com.facebook.presto.Session;
17+
import com.facebook.presto.spi.ColumnHandle;
18+
import com.facebook.presto.spi.plan.IndexSourceNode;
19+
import com.facebook.presto.spi.plan.PlanNode;
20+
import com.facebook.presto.spi.relation.VariableReferenceExpression;
21+
import com.facebook.presto.sql.planner.assertions.PlanMatchPattern;
22+
import com.facebook.presto.sql.planner.plan.IndexJoinNode;
1723
import com.facebook.presto.testing.QueryRunner;
1824
import com.facebook.presto.tests.tpch.IndexedTpchPlugin;
25+
import com.facebook.presto.tpch.TpchColumnHandle;
1926
import com.google.common.collect.ImmutableList;
2027
import org.testng.annotations.Test;
2128

2229
import java.util.List;
30+
import java.util.Map;
2331

2432
import static com.facebook.presto.SystemSessionProperties.NATIVE_EXECUTION_ENABLED;
2533
import static com.facebook.presto.SystemSessionProperties.OPTIMIZE_HASH_GENERATION;
34+
import static com.facebook.presto.SystemSessionProperties.PUSHDOWN_SUBFIELDS_ENABLED;
2635
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.anyTree;
2736
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.filter;
2837
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.indexJoin;
2938
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.indexSource;
39+
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.output;
3040
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.project;
3141
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.tableScan;
3242
import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
3343
import static com.facebook.presto.tests.AbstractTestIndexedQueries.INDEX_SPEC;
3444
import static com.facebook.presto.tpch.TpchMetadata.TINY_SCHEMA_NAME;
45+
import static org.testng.Assert.assertEquals;
46+
import static org.testng.Assert.assertNotNull;
3547

3648
public class TestNativeIndexJoinLogicalPlanner
3749
extends AbstractTestQueryFramework
@@ -146,4 +158,68 @@ public void testNonEqualIndexJoin()
146158
filter(indexSource("orders")))));
147159
}
148160
}
161+
162+
@Test
163+
public void testPushdownSubfields()
164+
{
165+
String query = "SELECT \n" +
166+
" MAP_SUBSET( " +
167+
" o.data, " +
168+
" ARRAY[1, 222, 33] " +
169+
" ) \n" +
170+
"FROM lineitem l\n" +
171+
"JOIN orders_extra o\n" +
172+
" ON l.orderkey = o.orderkey";
173+
PlanMatchPattern expectedQueryPlan = output(
174+
project(
175+
indexJoin(
176+
tableScan("lineitem"),
177+
indexSource("orders_extra"))));
178+
179+
Session defaultSession = Session.builder(getSession())
180+
.setSystemProperty(NATIVE_EXECUTION_ENABLED, "true")
181+
.setSystemProperty(OPTIMIZE_HASH_GENERATION, "false")
182+
.build();
183+
assertPlan(defaultSession, query, expectedQueryPlan);
184+
IndexSourceNode indexSourceNode = getIndexSourceNodeFromPlan(plan(query, defaultSession).getRoot());
185+
TpchColumnHandle columnHandle = getColumnHandle(indexSourceNode.getAssignments(), "data");
186+
assertNotNull(columnHandle);
187+
assertEquals(columnHandle.getRequiredSubfields().size(), 0);
188+
189+
Session sessionWithSubfieldPushdownEnabled = Session.builder(defaultSession)
190+
.setSystemProperty(PUSHDOWN_SUBFIELDS_ENABLED, "true")
191+
.build();
192+
assertPlan(sessionWithSubfieldPushdownEnabled, query, expectedQueryPlan);
193+
IndexSourceNode indexSourceNodeWithSubfields = getIndexSourceNodeFromPlan(plan(query, sessionWithSubfieldPushdownEnabled).getRoot());
194+
TpchColumnHandle columnHandleWithSubfields = getColumnHandle(indexSourceNodeWithSubfields.getAssignments(), "data");
195+
assertNotNull(columnHandle);
196+
assertEquals(columnHandleWithSubfields.getRequiredSubfields().size(), 3);
197+
}
198+
199+
private IndexSourceNode getIndexSourceNodeFromPlan(PlanNode node)
200+
{
201+
if (node == null) {
202+
return null;
203+
}
204+
if (node instanceof IndexSourceNode) {
205+
return (IndexSourceNode) node;
206+
}
207+
if (node instanceof IndexJoinNode) {
208+
return getIndexSourceNodeFromPlan(((IndexJoinNode) node).getIndexSource());
209+
}
210+
if (node.getSources().isEmpty()) {
211+
return null;
212+
}
213+
return getIndexSourceNodeFromPlan(node.getSources().get(0));
214+
}
215+
216+
private TpchColumnHandle getColumnHandle(Map<VariableReferenceExpression, ColumnHandle> assignments, String columnName)
217+
{
218+
for (ColumnHandle columnHandle : assignments.values()) {
219+
if (columnHandle instanceof TpchColumnHandle && ((TpchColumnHandle) columnHandle).getColumnName().equals(columnName)) {
220+
return (TpchColumnHandle) columnHandle;
221+
}
222+
}
223+
return null;
224+
}
149225
}

‎presto-tpch/src/main/java/com/facebook/presto/tpch/TpchColumnHandle.java‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,14 @@
1313
*/
1414
package com.facebook.presto.tpch;
1515

16+
import com.facebook.presto.common.Subfield;
1617
import com.facebook.presto.common.type.Type;
1718
import com.facebook.presto.spi.ColumnHandle;
1819
import com.fasterxml.jackson.annotation.JsonCreator;
1920
import com.fasterxml.jackson.annotation.JsonProperty;
21+
import com.google.common.collect.ImmutableList;
2022

23+
import java.util.List;
2124
import java.util.Objects;
2225

2326
import static java.util.Objects.requireNonNull;
@@ -27,14 +30,24 @@ public class TpchColumnHandle
2730
{
2831
private final String columnName;
2932
private final Type type;
33+
private final List<Subfield> requiredSubfields;
3034

3135
@JsonCreator
3236
public TpchColumnHandle(
3337
@JsonProperty("columnName") String columnName,
3438
@JsonProperty("type") Type type)
39+
{
40+
this(columnName, type, ImmutableList.of());
41+
}
42+
43+
public TpchColumnHandle(
44+
String columnName,
45+
Type type,
46+
List<Subfield> requiredSubfields)
3547
{
3648
this.columnName = requireNonNull(columnName, "columnName is null");
3749
this.type = requireNonNull(type, "type is null");
50+
this.requiredSubfields = requireNonNull(requiredSubfields, "requiredSubfields is null");
3851
}
3952

4053
@JsonProperty
@@ -55,6 +68,22 @@ public String toString()
5568
return "tpch:" + columnName;
5669
}
5770

71+
public List<Subfield> getRequiredSubfields()
72+
{
73+
return requiredSubfields;
74+
}
75+
76+
@Override
77+
public ColumnHandle withRequiredSubfields(List<Subfield> subfields)
78+
{
79+
if (!requiredSubfields.isEmpty()) {
80+
// This column is already a pushed down subfield column.
81+
return this;
82+
}
83+
84+
return new TpchColumnHandle(columnName, type, subfields);
85+
}
86+
5887
@Override
5988
public boolean equals(Object o)
6089
{

0 commit comments

Comments
 (0)