Skip to content

Add Lineage Information to EXPLAIN (TYPE IO) #24952

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

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
+115 −22
Open

Add Lineage Information to EXPLAIN (TYPE IO) #24952

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
16 changes: 13 additions & 3 deletions core/trino-main/src/main/java/io/trino/sql/analyzer/QueryExplainer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ public String getPlan(Session session, Statement statement, Type planType, List<
session,
false,
version);
case IO -> textIoPlan(getLogicalPlan(session, statement, parameters, warningCollector, planOptimizersStatsCollector), plannerContext, session);
case IO -> getIoPlan(session, statement, parameters, warningCollector, planOptimizersStatsCollector);
default -> throw new IllegalArgumentException("Unhandled plan type: " + planType);
};
}
Expand Down Expand Up @@ -141,7 +141,7 @@ public String getJsonPlan(Session session, Statement statement, Type planType, L
}

return switch (planType) {
case IO -> textIoPlan(getLogicalPlan(session, statement, parameters, warningCollector, planOptimizersStatsCollector), plannerContext, session);
case IO -> getIoPlan(session, statement, parameters, warningCollector, planOptimizersStatsCollector);
case LOGICAL -> {
Plan plan = getLogicalPlan(session, statement, parameters, warningCollector, planOptimizersStatsCollector);
yield jsonLogicalPlan(plan.getRoot(), session, plannerContext.getMetadata(), plannerContext.getFunctionManager(), plan.getStatsAndCosts());
Expand All @@ -155,11 +155,21 @@ public String getJsonPlan(Session session, Statement statement, Type planType, L
};
}

private String getIoPlan(Session session, Statement statement, List<Expression> parameters, WarningCollector warningCollector, PlanOptimizersStatsCollector planOptimizersStatsCollector)
{
Analysis analysis = analyze(session, statement, parameters, warningCollector, planOptimizersStatsCollector);
Plan plan = getLogicalPlan(session, statement, parameters, warningCollector, planOptimizersStatsCollector, analysis);
return textIoPlan(plan, plannerContext, session, analysis);
}

public Plan getLogicalPlan(Session session, Statement statement, List<Expression> parameters, WarningCollector warningCollector, PlanOptimizersStatsCollector planOptimizersStatsCollector)
{
// analyze statement
Analysis analysis = analyze(session, statement, parameters, warningCollector, planOptimizersStatsCollector);
return getLogicalPlan(session, statement, parameters, warningCollector, planOptimizersStatsCollector, analysis);
}

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

// plan statement
Expand Down
66 changes: 58 additions & 8 deletions core/trino-main/src/main/java/io/trino/sql/planner/planprinter/IoPlanPrinter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import io.trino.Session;
import io.trino.cost.PlanCostEstimate;
Expand All @@ -24,11 +25,14 @@
import io.trino.spi.connector.CatalogSchemaTableName;
import io.trino.spi.connector.ColumnHandle;
import io.trino.spi.connector.ColumnMetadata;
import io.trino.spi.eventlistener.OutputColumnMetadata;
import io.trino.spi.predicate.Domain;
import io.trino.spi.predicate.Range;
import io.trino.spi.predicate.TupleDomain;
import io.trino.spi.type.Type;
import io.trino.sql.PlannerContext;
import io.trino.sql.analyzer.Analysis;
import io.trino.sql.analyzer.OutputColumn;
import io.trino.sql.planner.DomainTranslator;
import io.trino.sql.planner.Plan;
import io.trino.sql.planner.plan.FilterNode;
Expand All @@ -47,12 +51,14 @@
import io.trino.sql.planner.planprinter.IoPlanPrinter.IoPlan.IoPlanBuilder;

import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;

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

private IoPlanPrinter(Plan plan, PlannerContext plannerContext, Session session)
private IoPlanPrinter(Plan plan, PlannerContext plannerContext, Session session, Optional<Analysis> analysis)
{
this.plan = requireNonNull(plan, "plan is null");
this.plannerContext = requireNonNull(plannerContext, "plannerContext is null");
this.session = requireNonNull(session, "session is null");
this.analysis = requireNonNull(analysis, "analysis is null");
this.valuePrinter = new ValuePrinter(plannerContext.getMetadata(), plannerContext.getFunctionManager(), session);
}

/**
* @throws io.trino.NotInTransactionException if called without an active transaction
*/
public static String textIoPlan(Plan plan, PlannerContext plannerContext, Session session)
public static String textIoPlan(Plan plan, PlannerContext plannerContext, Session session, Analysis analysis)
{
return new IoPlanPrinter(plan, plannerContext, session).print();
return new IoPlanPrinter(plan, plannerContext, session, Optional.of(analysis)).print();
}

private String print()
{
IoPlanBuilder ioPlanBuilder = new IoPlanBuilder(plan);
IoPlanBuilder ioPlanBuilder = new IoPlanBuilder(plan, analysis);
plan.getRoot().accept(new IoPlanVisitor(), ioPlanBuilder);
return jsonCodec(IoPlan.class).toJson(ioPlanBuilder.build());
}
Expand All @@ -93,16 +101,19 @@ public static class IoPlan
private final Set<TableColumnInfo> inputTableColumnInfos;
private final Optional<CatalogSchemaTableName> outputTable;
private final EstimatedStatsAndCost estimate;
private final Optional<List<OutputColumnMetadata>> columns;

@JsonCreator
public IoPlan(
@JsonProperty("inputTableColumnInfos") Set<TableColumnInfo> inputTableColumnInfos,
@JsonProperty("outputTable") Optional<CatalogSchemaTableName> outputTable,
@JsonProperty("estimate") EstimatedStatsAndCost estimate)
@JsonProperty("estimate") EstimatedStatsAndCost estimate,
@JsonProperty("outputColumns") Optional<List<OutputColumnMetadata>> columns)
{
this.inputTableColumnInfos = ImmutableSet.copyOf(requireNonNull(inputTableColumnInfos, "inputTableColumnInfos is null"));
this.outputTable = requireNonNull(outputTable, "outputTable is null");
this.estimate = requireNonNull(estimate, "estimate is null");
this.columns = requireNonNull(columns, "columns is null").map(ImmutableList::copyOf);
}

@JsonProperty
Expand All @@ -123,6 +134,12 @@ public EstimatedStatsAndCost getEstimate()
return estimate;
}

@JsonProperty
public Optional<List<OutputColumnMetadata>> getColumns()
{
return columns;
}

@Override
public boolean equals(Object obj)
{
Expand All @@ -134,7 +151,8 @@ public boolean equals(Object obj)
}
IoPlan o = (IoPlan) obj;
return Objects.equals(inputTableColumnInfos, o.inputTableColumnInfos) &&
Objects.equals(outputTable, o.outputTable);
Objects.equals(outputTable, o.outputTable) &&
Objects.equals(columns, o.columns);
}

@Override
Expand All @@ -150,20 +168,25 @@ public String toString()
.add("inputTableColumnInfos", inputTableColumnInfos)
.add("outputTable", outputTable)
.add("estimate", estimate)
.add("outputColumns", columns)
.toString();
}

protected static class IoPlanBuilder
{
private final Plan plan;
private final Optional<Analysis> analysis;
private final Set<TableColumnInfo> inputTableColumnInfos;
private Optional<CatalogSchemaTableName> outputTable;
private Optional<List<OutputColumnMetadata>> columns;

private IoPlanBuilder(Plan plan)
private IoPlanBuilder(Plan plan, Optional<Analysis> analysis)
{
this.plan = plan;
this.analysis = analysis;
this.inputTableColumnInfos = new HashSet<>();
this.outputTable = Optional.empty();
this.columns = Optional.empty();
}

private IoPlanBuilder addInputTableColumnInfo(TableColumnInfo tableColumnInfo)
Expand All @@ -172,6 +195,32 @@ private IoPlanBuilder addInputTableColumnInfo(TableColumnInfo tableColumnInfo)
return this;
}

private IoPlanBuilder setOutputColumns(Optional<List<OutputColumnMetadata>> columns)
{
this.columns = columns;
return this;
}

private void setOutputColumns()
{
analysis.ifPresent(analysis ->
setOutputColumns(analysis.getTarget()
.flatMap(target -> target.getColumns()
.map(columns -> columns.stream()
.map(this::createOutputColumnMetadata)
.collect(toImmutableList())))));
}

private OutputColumnMetadata createOutputColumnMetadata(OutputColumn column)
{
return new OutputColumnMetadata(
column.getColumn().getName(),
column.getColumn().getType(),
column.getSourceColumns().stream()
.map(Analysis.SourceColumn::getColumnDetail)
.collect(toImmutableSet()));
}

private IoPlanBuilder setOutputTable(CatalogSchemaTableName outputTable)
{
this.outputTable = Optional.of(outputTable);
Expand All @@ -180,7 +229,7 @@ private IoPlanBuilder setOutputTable(CatalogSchemaTableName outputTable)

private IoPlan build()
{
return new IoPlan(inputTableColumnInfos, outputTable, getEstimatedStatsAndCost());
return new IoPlan(inputTableColumnInfos, outputTable, getEstimatedStatsAndCost(), columns);
}

private EstimatedStatsAndCost getEstimatedStatsAndCost()
Expand Down Expand Up @@ -721,6 +770,7 @@ else if (writerTarget instanceof CreateReference || writerTarget instanceof Inse
else {
throw new IllegalStateException(format("Unknown WriterTarget subclass %s", writerTarget.getClass().getSimpleName()));
}
context.setOutputColumns();
return processChildren(node, context);
}

Expand Down
53 changes: 43 additions & 10 deletions plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseHiveConnectorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,8 @@
import io.trino.spi.connector.CatalogSchemaTableName;
import io.trino.spi.connector.ColumnMetadata;
import io.trino.spi.connector.Constraint;
import io.trino.spi.eventlistener.ColumnDetail;
import io.trino.spi.eventlistener.OutputColumnMetadata;
import io.trino.spi.security.ConnectorIdentity;
import io.trino.spi.security.Identity;
import io.trino.spi.security.SelectedRole;
Expand Down Expand Up @@ -1253,7 +1255,17 @@ public void testIoExplain()
new FormattedMarker(Optional.of("10"), EXACTLY))))))),
estimate)),
Optional.of(new CatalogSchemaTableName(catalog, "tpch", "test_io_explain")),
estimate));
estimate,
Optional.of(ImmutableList.of(
new OutputColumnMetadata("custkey", "bigint",
ImmutableSet.of(
new ColumnDetail("hive", "tpch", "test_io_explain", "custkey"))),
new OutputColumnMetadata("orderkey", "bigint",
ImmutableSet.of(
new ColumnDetail("hive", "tpch", "test_io_explain", "orderkey"))),
new OutputColumnMetadata("processing", "boolean",
ImmutableSet.of(
new ColumnDetail("hive", "tpch", "test_io_explain", "processing")))))));

assertUpdate("DROP TABLE test_io_explain");

Expand Down Expand Up @@ -1289,7 +1301,14 @@ public void testIoExplain()
new FormattedMarker(Optional.of("10"), EXACTLY))))))),
estimate)),
Optional.of(new CatalogSchemaTableName(catalog, "tpch", "test_io_explain")),
estimate));
estimate,
Optional.of(ImmutableList.of(
new OutputColumnMetadata("custkey", "bigint",
ImmutableSet.of(
new ColumnDetail("hive", "tpch", "test_io_explain", "custkey"))),
new OutputColumnMetadata("orderkey", "bigint",
ImmutableSet.of(
new ColumnDetail("hive", "tpch", "test_io_explain", "orderkey")))))));

EstimatedStatsAndCost finalEstimate = new EstimatedStatsAndCost(Double.NaN, Double.NaN, Double.NaN, Double.NaN, Double.NaN);
estimate = new EstimatedStatsAndCost(1.0, 18.0, 18, 0.0, 0.0);
Expand All @@ -1312,7 +1331,14 @@ public void testIoExplain()
new FormattedMarker(Optional.of("100"), EXACTLY))))))),
estimate)),
Optional.of(new CatalogSchemaTableName(catalog, "tpch", "test_io_explain")),
finalEstimate));
finalEstimate,
Optional.of(ImmutableList.of(
new OutputColumnMetadata("custkey", "bigint",
ImmutableSet.of(
new ColumnDetail("hive", "tpch", "test_io_explain", "custkey"))),
new OutputColumnMetadata("orderkey", "bigint",
ImmutableSet.of(
new ColumnDetail("hive", "tpch", "test_io_explain", "orderkey")))))));

assertUpdate("DROP TABLE test_io_explain");
}
Expand Down Expand Up @@ -1365,7 +1391,8 @@ public void testIoExplainColumnFilters()
new FormattedMarker(Optional.of("P"), EXACTLY))))))),
estimate)),
Optional.empty(),
finalEstimate));
finalEstimate,
Optional.empty()));
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')");
assertThat(getIoPlanCodec().fromJson((String) getOnlyElement(result.getOnlyColumnAsSet()))).isEqualTo(new IoPlan(
ImmutableSet.of(
Expand Down Expand Up @@ -1409,7 +1436,8 @@ public void testIoExplainColumnFilters()
new FormattedMarker(Optional.of("10"), EXACTLY))))))),
estimate)),
Optional.empty(),
finalEstimate));
finalEstimate,
Optional.empty()));
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");
assertThat(getIoPlanCodec().fromJson((String) getOnlyElement(result.getOnlyColumnAsSet()))).isEqualTo(new IoPlan(
ImmutableSet.of(
Expand Down Expand Up @@ -1441,7 +1469,8 @@ public void testIoExplainColumnFilters()
new FormattedMarker(Optional.of("10"), EXACTLY))))))),
estimate)),
Optional.empty(),
finalEstimate));
finalEstimate,
Optional.empty()));

assertUpdate("DROP TABLE test_io_explain_column_filters");
}
Expand All @@ -1461,7 +1490,8 @@ public void testIoExplainWithEmptyPartitionedTable()
new IoPlanPrinter.Constraint(true, ImmutableSet.of()),
estimate)),
Optional.empty(),
estimate));
estimate,
Optional.empty()));

assertUpdate("DROP TABLE test_io_explain_with_empty_partitioned_table");
}
Expand Down Expand Up @@ -1507,7 +1537,8 @@ public void testIoExplainNoFilter()
new FormattedMarker(Optional.of("a"), EXACTLY))))))),
estimate)),
Optional.empty(),
finalEstimate));
finalEstimate,
Optional.empty()));
assertUpdate("DROP TABLE io_explain_test_no_filter");
}

Expand Down Expand Up @@ -1561,7 +1592,8 @@ public void testIoExplainFilterOnAgg()
new FormattedMarker(Optional.of("b"), EXACTLY))))))),
estimate)),
Optional.empty(),
finalEstimate));
finalEstimate,
Optional.empty()));
assertUpdate("DROP TABLE io_explain_test_filter_on_agg");
}

Expand Down Expand Up @@ -1616,7 +1648,8 @@ public void testIoExplainWithPrimitiveTypes()
new FormattedMarker(Optional.of(entry.getKey().toString()), EXACTLY))))))),
estimate)),
Optional.empty(),
estimate));
estimate,
Optional.empty()));

assertUpdate("DROP TABLE " + tableName);
}
Expand Down
2 changes: 1 addition & 1 deletion testing/trino-tests/src/test/java/io/trino/tests/tpch/TestTpchConnectorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ public void testIoExplain()
objectMapperProvider.setJsonDeserializers(ImmutableMap.of(Type.class, new TypeDeserializer(getQueryRunner().getPlannerContext().getTypeManager())));
JsonCodec<IoPlanPrinter.IoPlan> codec = new JsonCodecFactory(objectMapperProvider).jsonCodec(IoPlanPrinter.IoPlan.class);

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

@Test
Expand Down
Loading