Skip to content

Commit aef0a7f

Browse files
predator4annpredator4ann
committed
[Feature] Add Lineage Information to EXPLAIN (TYPE IO)
Signed-off-by: predator4ann <[email protected]>
1 parent 4c2dda8 commit aef0a7f

File tree

4 files changed

+115
-22
lines changed

4 files changed

+115
-22
lines changed

Diff for: ‎core/trino-main/src/main/java/io/trino/sql/analyzer/QueryExplainer.java

+13-3
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ public String getPlan(Session session, Statement statement, Type planType, List<
114114
session,
115115
false,
116116
version);
117-
case IO -> textIoPlan(getLogicalPlan(session, statement, parameters, warningCollector, planOptimizersStatsCollector), plannerContext, session);
117+
case IO -> getIoPlan(session, statement, parameters, warningCollector, planOptimizersStatsCollector);
118118
default -> throw new IllegalArgumentException("Unhandled plan type: " + planType);
119119
};
120120
}
@@ -146,7 +146,7 @@ public String getJsonPlan(Session session, Statement statement, Type planType, L
146146
}
147147

148148
return switch (planType) {
149-
case IO -> textIoPlan(getLogicalPlan(session, statement, parameters, warningCollector, planOptimizersStatsCollector), plannerContext, session);
149+
case IO -> getIoPlan(session, statement, parameters, warningCollector, planOptimizersStatsCollector);
150150
case LOGICAL -> {
151151
Plan plan = getLogicalPlan(session, statement, parameters, warningCollector, planOptimizersStatsCollector);
152152
yield jsonLogicalPlan(plan.getRoot(), session, plannerContext.getMetadata(), plannerContext.getFunctionManager(), plan.getStatsAndCosts());
@@ -160,11 +160,21 @@ public String getJsonPlan(Session session, Statement statement, Type planType, L
160160
};
161161
}
162162

163+
private String getIoPlan(Session session, Statement statement, List<Expression> parameters, WarningCollector warningCollector, PlanOptimizersStatsCollector planOptimizersStatsCollector)
164+
{
165+
Analysis analysis = analyze(session, statement, parameters, warningCollector, planOptimizersStatsCollector);
166+
Plan plan = getLogicalPlan(session, statement, parameters, warningCollector, planOptimizersStatsCollector, analysis);
167+
return textIoPlan(plan, plannerContext, session, analysis);
168+
}
169+
163170
public Plan getLogicalPlan(Session session, Statement statement, List<Expression> parameters, WarningCollector warningCollector, PlanOptimizersStatsCollector planOptimizersStatsCollector)
164171
{
165-
// analyze statement
166172
Analysis analysis = analyze(session, statement, parameters, warningCollector, planOptimizersStatsCollector);
173+
return getLogicalPlan(session, statement, parameters, warningCollector, planOptimizersStatsCollector, analysis);
174+
}
167175

176+
public Plan getLogicalPlan(Session session, Statement statement, List<Expression> parameters, WarningCollector warningCollector, PlanOptimizersStatsCollector planOptimizersStatsCollector, Analysis analysis)
177+
{
168178
PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator();
169179

170180
// plan statement

Diff for: ‎core/trino-main/src/main/java/io/trino/sql/planner/planprinter/IoPlanPrinter.java

+58-8
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515

1616
import com.fasterxml.jackson.annotation.JsonCreator;
1717
import com.fasterxml.jackson.annotation.JsonProperty;
18+
import com.google.common.collect.ImmutableList;
1819
import com.google.common.collect.ImmutableSet;
1920
import io.trino.Session;
2021
import io.trino.cost.PlanCostEstimate;
@@ -24,11 +25,14 @@
2425
import io.trino.spi.connector.CatalogSchemaTableName;
2526
import io.trino.spi.connector.ColumnHandle;
2627
import io.trino.spi.connector.ColumnMetadata;
28+
import io.trino.spi.eventlistener.OutputColumnMetadata;
2729
import io.trino.spi.predicate.Domain;
2830
import io.trino.spi.predicate.Range;
2931
import io.trino.spi.predicate.TupleDomain;
3032
import io.trino.spi.type.Type;
3133
import io.trino.sql.PlannerContext;
34+
import io.trino.sql.analyzer.Analysis;
35+
import io.trino.sql.analyzer.OutputColumn;
3236
import io.trino.sql.planner.DomainTranslator;
3337
import io.trino.sql.planner.Plan;
3438
import io.trino.sql.planner.plan.FilterNode;
@@ -47,12 +51,14 @@
4751
import io.trino.sql.planner.planprinter.IoPlanPrinter.IoPlan.IoPlanBuilder;
4852

4953
import java.util.HashSet;
54+
import java.util.List;
5055
import java.util.Map;
5156
import java.util.Objects;
5257
import java.util.Optional;
5358
import java.util.Set;
5459

5560
import static com.google.common.base.MoreObjects.toStringHelper;
61+
import static com.google.common.collect.ImmutableList.toImmutableList;
5662
import static com.google.common.collect.ImmutableSet.toImmutableSet;
5763
import static io.airlift.json.JsonCodec.jsonCodec;
5864
import static java.lang.String.format;
@@ -64,26 +70,28 @@ public class IoPlanPrinter
6470
private final PlannerContext plannerContext;
6571
private final Session session;
6672
private final ValuePrinter valuePrinter;
73+
private final Optional<Analysis> analysis;
6774

68-
private IoPlanPrinter(Plan plan, PlannerContext plannerContext, Session session)
75+
private IoPlanPrinter(Plan plan, PlannerContext plannerContext, Session session, Optional<Analysis> analysis)
6976
{
7077
this.plan = requireNonNull(plan, "plan is null");
7178
this.plannerContext = requireNonNull(plannerContext, "plannerContext is null");
7279
this.session = requireNonNull(session, "session is null");
80+
this.analysis = requireNonNull(analysis, "analysis is null");
7381
this.valuePrinter = new ValuePrinter(plannerContext.getMetadata(), plannerContext.getFunctionManager(), session);
7482
}
7583

7684
/**
7785
* @throws io.trino.NotInTransactionException if called without an active transaction
7886
*/
79-
public static String textIoPlan(Plan plan, PlannerContext plannerContext, Session session)
87+
public static String textIoPlan(Plan plan, PlannerContext plannerContext, Session session, Analysis analysis)
8088
{
81-
return new IoPlanPrinter(plan, plannerContext, session).print();
89+
return new IoPlanPrinter(plan, plannerContext, session, Optional.of(analysis)).print();
8290
}
8391

8492
private String print()
8593
{
86-
IoPlanBuilder ioPlanBuilder = new IoPlanBuilder(plan);
94+
IoPlanBuilder ioPlanBuilder = new IoPlanBuilder(plan, analysis);
8795
plan.getRoot().accept(new IoPlanVisitor(), ioPlanBuilder);
8896
return jsonCodec(IoPlan.class).toJson(ioPlanBuilder.build());
8997
}
@@ -93,16 +101,19 @@ public static class IoPlan
93101
private final Set<TableColumnInfo> inputTableColumnInfos;
94102
private final Optional<CatalogSchemaTableName> outputTable;
95103
private final EstimatedStatsAndCost estimate;
104+
private final Optional<List<OutputColumnMetadata>> columns;
96105

97106
@JsonCreator
98107
public IoPlan(
99108
@JsonProperty("inputTableColumnInfos") Set<TableColumnInfo> inputTableColumnInfos,
100109
@JsonProperty("outputTable") Optional<CatalogSchemaTableName> outputTable,
101-
@JsonProperty("estimate") EstimatedStatsAndCost estimate)
110+
@JsonProperty("estimate") EstimatedStatsAndCost estimate,
111+
@JsonProperty("outputColumns") Optional<List<OutputColumnMetadata>> columns)
102112
{
103113
this.inputTableColumnInfos = ImmutableSet.copyOf(requireNonNull(inputTableColumnInfos, "inputTableColumnInfos is null"));
104114
this.outputTable = requireNonNull(outputTable, "outputTable is null");
105115
this.estimate = requireNonNull(estimate, "estimate is null");
116+
this.columns = requireNonNull(columns, "columns is null").map(ImmutableList::copyOf);
106117
}
107118

108119
@JsonProperty
@@ -123,6 +134,12 @@ public EstimatedStatsAndCost getEstimate()
123134
return estimate;
124135
}
125136

137+
@JsonProperty
138+
public Optional<List<OutputColumnMetadata>> getColumns()
139+
{
140+
return columns;
141+
}
142+
126143
@Override
127144
public boolean equals(Object obj)
128145
{
@@ -134,7 +151,8 @@ public boolean equals(Object obj)
134151
}
135152
IoPlan o = (IoPlan) obj;
136153
return Objects.equals(inputTableColumnInfos, o.inputTableColumnInfos) &&
137-
Objects.equals(outputTable, o.outputTable);
154+
Objects.equals(outputTable, o.outputTable) &&
155+
Objects.equals(columns, o.columns);
138156
}
139157

140158
@Override
@@ -150,20 +168,25 @@ public String toString()
150168
.add("inputTableColumnInfos", inputTableColumnInfos)
151169
.add("outputTable", outputTable)
152170
.add("estimate", estimate)
171+
.add("outputColumns", columns)
153172
.toString();
154173
}
155174

156175
protected static class IoPlanBuilder
157176
{
158177
private final Plan plan;
178+
private final Optional<Analysis> analysis;
159179
private final Set<TableColumnInfo> inputTableColumnInfos;
160180
private Optional<CatalogSchemaTableName> outputTable;
181+
private Optional<List<OutputColumnMetadata>> columns;
161182

162-
private IoPlanBuilder(Plan plan)
183+
private IoPlanBuilder(Plan plan, Optional<Analysis> analysis)
163184
{
164185
this.plan = plan;
186+
this.analysis = analysis;
165187
this.inputTableColumnInfos = new HashSet<>();
166188
this.outputTable = Optional.empty();
189+
this.columns = Optional.empty();
167190
}
168191

169192
private IoPlanBuilder addInputTableColumnInfo(TableColumnInfo tableColumnInfo)
@@ -172,6 +195,32 @@ private IoPlanBuilder addInputTableColumnInfo(TableColumnInfo tableColumnInfo)
172195
return this;
173196
}
174197

198+
private IoPlanBuilder setOutputColumns(Optional<List<OutputColumnMetadata>> columns)
199+
{
200+
this.columns = columns;
201+
return this;
202+
}
203+
204+
private void setOutputColumns()
205+
{
206+
analysis.ifPresent(analysis ->
207+
setOutputColumns(analysis.getTarget()
208+
.flatMap(target -> target.getColumns()
209+
.map(columns -> columns.stream()
210+
.map(this::createOutputColumnMetadata)
211+
.collect(toImmutableList())))));
212+
}
213+
214+
private OutputColumnMetadata createOutputColumnMetadata(OutputColumn column)
215+
{
216+
return new OutputColumnMetadata(
217+
column.getColumn().getName(),
218+
column.getColumn().getType(),
219+
column.getSourceColumns().stream()
220+
.map(Analysis.SourceColumn::getColumnDetail)
221+
.collect(toImmutableSet()));
222+
}
223+
175224
private IoPlanBuilder setOutputTable(CatalogSchemaTableName outputTable)
176225
{
177226
this.outputTable = Optional.of(outputTable);
@@ -180,7 +229,7 @@ private IoPlanBuilder setOutputTable(CatalogSchemaTableName outputTable)
180229

181230
private IoPlan build()
182231
{
183-
return new IoPlan(inputTableColumnInfos, outputTable, getEstimatedStatsAndCost());
232+
return new IoPlan(inputTableColumnInfos, outputTable, getEstimatedStatsAndCost(), columns);
184233
}
185234

186235
private EstimatedStatsAndCost getEstimatedStatsAndCost()
@@ -721,6 +770,7 @@ else if (writerTarget instanceof CreateReference || writerTarget instanceof Inse
721770
else {
722771
throw new IllegalStateException(format("Unknown WriterTarget subclass %s", writerTarget.getClass().getSimpleName()));
723772
}
773+
context.setOutputColumns();
724774
return processChildren(node, context);
725775
}
726776

Diff for: ‎plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseHiveConnectorTest.java

+43-10
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@
4040
import io.trino.spi.connector.CatalogSchemaTableName;
4141
import io.trino.spi.connector.ColumnMetadata;
4242
import io.trino.spi.connector.Constraint;
43+
import io.trino.spi.eventlistener.ColumnDetail;
44+
import io.trino.spi.eventlistener.OutputColumnMetadata;
4345
import io.trino.spi.security.ConnectorIdentity;
4446
import io.trino.spi.security.Identity;
4547
import io.trino.spi.security.SelectedRole;
@@ -1253,7 +1255,17 @@ public void testIoExplain()
12531255
new FormattedMarker(Optional.of("10"), EXACTLY))))))),
12541256
estimate)),
12551257
Optional.of(new CatalogSchemaTableName(catalog, "tpch", "test_io_explain")),
1256-
estimate));
1258+
estimate,
1259+
Optional.of(ImmutableList.of(
1260+
new OutputColumnMetadata("custkey", "bigint",
1261+
ImmutableSet.of(
1262+
new ColumnDetail("hive", "tpch", "test_io_explain", "custkey"))),
1263+
new OutputColumnMetadata("orderkey", "bigint",
1264+
ImmutableSet.of(
1265+
new ColumnDetail("hive", "tpch", "test_io_explain", "orderkey"))),
1266+
new OutputColumnMetadata("processing", "boolean",
1267+
ImmutableSet.of(
1268+
new ColumnDetail("hive", "tpch", "test_io_explain", "processing")))))));
12571269

12581270
assertUpdate("DROP TABLE test_io_explain");
12591271

@@ -1289,7 +1301,14 @@ public void testIoExplain()
12891301
new FormattedMarker(Optional.of("10"), EXACTLY))))))),
12901302
estimate)),
12911303
Optional.of(new CatalogSchemaTableName(catalog, "tpch", "test_io_explain")),
1292-
estimate));
1304+
estimate,
1305+
Optional.of(ImmutableList.of(
1306+
new OutputColumnMetadata("custkey", "bigint",
1307+
ImmutableSet.of(
1308+
new ColumnDetail("hive", "tpch", "test_io_explain", "custkey"))),
1309+
new OutputColumnMetadata("orderkey", "bigint",
1310+
ImmutableSet.of(
1311+
new ColumnDetail("hive", "tpch", "test_io_explain", "orderkey")))))));
12931312

12941313
EstimatedStatsAndCost finalEstimate = new EstimatedStatsAndCost(Double.NaN, Double.NaN, Double.NaN, Double.NaN, Double.NaN);
12951314
estimate = new EstimatedStatsAndCost(1.0, 18.0, 18, 0.0, 0.0);
@@ -1312,7 +1331,14 @@ public void testIoExplain()
13121331
new FormattedMarker(Optional.of("100"), EXACTLY))))))),
13131332
estimate)),
13141333
Optional.of(new CatalogSchemaTableName(catalog, "tpch", "test_io_explain")),
1315-
finalEstimate));
1334+
finalEstimate,
1335+
Optional.of(ImmutableList.of(
1336+
new OutputColumnMetadata("custkey", "bigint",
1337+
ImmutableSet.of(
1338+
new ColumnDetail("hive", "tpch", "test_io_explain", "custkey"))),
1339+
new OutputColumnMetadata("orderkey", "bigint",
1340+
ImmutableSet.of(
1341+
new ColumnDetail("hive", "tpch", "test_io_explain", "orderkey")))))));
13161342

13171343
assertUpdate("DROP TABLE test_io_explain");
13181344
}
@@ -1365,7 +1391,8 @@ public void testIoExplainColumnFilters()
13651391
new FormattedMarker(Optional.of("P"), EXACTLY))))))),
13661392
estimate)),
13671393
Optional.empty(),
1368-
finalEstimate));
1394+
finalEstimate,
1395+
Optional.empty()));
13691396
result = computeActual("EXPLAIN (TYPE IO, FORMAT JSON) SELECT custkey, orderkey, orderstatus FROM test_io_explain_column_filters WHERE custkey <= 10 and (orderstatus='P' or orderstatus='S')");
13701397
assertThat(getIoPlanCodec().fromJson((String) getOnlyElement(result.getOnlyColumnAsSet()))).isEqualTo(new IoPlan(
13711398
ImmutableSet.of(
@@ -1409,7 +1436,8 @@ public void testIoExplainColumnFilters()
14091436
new FormattedMarker(Optional.of("10"), EXACTLY))))))),
14101437
estimate)),
14111438
Optional.empty(),
1412-
finalEstimate));
1439+
finalEstimate,
1440+
Optional.empty()));
14131441
result = computeActual("EXPLAIN (TYPE IO, FORMAT JSON) SELECT custkey, orderkey, orderstatus FROM test_io_explain_column_filters WHERE custkey <= 10 and cast(orderstatus as integer) = 5");
14141442
assertThat(getIoPlanCodec().fromJson((String) getOnlyElement(result.getOnlyColumnAsSet()))).isEqualTo(new IoPlan(
14151443
ImmutableSet.of(
@@ -1441,7 +1469,8 @@ public void testIoExplainColumnFilters()
14411469
new FormattedMarker(Optional.of("10"), EXACTLY))))))),
14421470
estimate)),
14431471
Optional.empty(),
1444-
finalEstimate));
1472+
finalEstimate,
1473+
Optional.empty()));
14451474

14461475
assertUpdate("DROP TABLE test_io_explain_column_filters");
14471476
}
@@ -1461,7 +1490,8 @@ public void testIoExplainWithEmptyPartitionedTable()
14611490
new IoPlanPrinter.Constraint(true, ImmutableSet.of()),
14621491
estimate)),
14631492
Optional.empty(),
1464-
estimate));
1493+
estimate,
1494+
Optional.empty()));
14651495

14661496
assertUpdate("DROP TABLE test_io_explain_with_empty_partitioned_table");
14671497
}
@@ -1507,7 +1537,8 @@ public void testIoExplainNoFilter()
15071537
new FormattedMarker(Optional.of("a"), EXACTLY))))))),
15081538
estimate)),
15091539
Optional.empty(),
1510-
finalEstimate));
1540+
finalEstimate,
1541+
Optional.empty()));
15111542
assertUpdate("DROP TABLE io_explain_test_no_filter");
15121543
}
15131544

@@ -1561,7 +1592,8 @@ public void testIoExplainFilterOnAgg()
15611592
new FormattedMarker(Optional.of("b"), EXACTLY))))))),
15621593
estimate)),
15631594
Optional.empty(),
1564-
finalEstimate));
1595+
finalEstimate,
1596+
Optional.empty()));
15651597
assertUpdate("DROP TABLE io_explain_test_filter_on_agg");
15661598
}
15671599

@@ -1616,7 +1648,8 @@ public void testIoExplainWithPrimitiveTypes()
16161648
new FormattedMarker(Optional.of(entry.getKey().toString()), EXACTLY))))))),
16171649
estimate)),
16181650
Optional.empty(),
1619-
estimate));
1651+
estimate,
1652+
Optional.empty()));
16201653

16211654
assertUpdate("DROP TABLE " + tableName);
16221655
}

Diff for: ‎testing/trino-tests/src/test/java/io/trino/tests/tpch/TestTpchConnectorTest.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ public void testIoExplain()
108108
objectMapperProvider.setJsonDeserializers(ImmutableMap.of(Type.class, new TypeDeserializer(getQueryRunner().getPlannerContext().getTypeManager())));
109109
JsonCodec<IoPlanPrinter.IoPlan> codec = new JsonCodecFactory(objectMapperProvider).jsonCodec(IoPlanPrinter.IoPlan.class);
110110

111-
assertThat(codec.fromJson((String) getOnlyElement(result.getOnlyColumnAsSet()))).isEqualTo(new IoPlanPrinter.IoPlan(ImmutableSet.of(input), Optional.empty(), totalEstimate));
111+
assertThat(codec.fromJson((String) getOnlyElement(result.getOnlyColumnAsSet()))).isEqualTo(new IoPlanPrinter.IoPlan(ImmutableSet.of(input), Optional.empty(), totalEstimate, Optional.empty()));
112112
}
113113

114114
@Test

0 commit comments

Comments
 (0)