Skip to content

Commit

Permalink
Build and print the new IR program for each query
Browse files Browse the repository at this point in the history
The new IR building and printing is tested in src/test/java/io/trino/sql/dialect/trino
However, those unit tests have limited serialization capacity.
Full serialization requires injected dependencies.

With this change, each query run on the server will be rewritten
into a new IR program and printed on the console. It will use
dependencies for full serialization.

Note that the queries are captured in SqlQueryExecution after they are
planned and optimized, and right before they are fragmented. This is
the moment where CTE reuse will kick in.

Note also that some queries will fail in ProgramBuilder if they contain
operations that we don't support yet.
  • Loading branch information
kasiafi committed Feb 7, 2025
1 parent 61a2bee commit efa6f93
Showing 1 changed file with 14 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,8 @@
import io.trino.sql.analyzer.Analysis;
import io.trino.sql.analyzer.Analyzer;
import io.trino.sql.analyzer.AnalyzerFactory;
import io.trino.sql.dialect.trino.ProgramBuilder;
import io.trino.sql.newir.FormatOptions;
import io.trino.sql.planner.AdaptivePlanner;
import io.trino.sql.planner.InputExtractor;
import io.trino.sql.planner.LogicalPlanner;
Expand Down Expand Up @@ -148,6 +150,7 @@ public class SqlQueryExecution
private final EventDrivenTaskSourceFactory eventDrivenTaskSourceFactory;
private final TaskDescriptorStorage taskDescriptorStorage;
private final PlanOptimizersStatsCollector planOptimizersStatsCollector;
private final FormatOptions formatOptions;

private SqlQueryExecution(
PreparedQuery preparedQuery,
Expand Down Expand Up @@ -185,7 +188,8 @@ private SqlQueryExecution(
SqlTaskManager coordinatorTaskManager,
ExchangeManagerRegistry exchangeManagerRegistry,
EventDrivenTaskSourceFactory eventDrivenTaskSourceFactory,
TaskDescriptorStorage taskDescriptorStorage)
TaskDescriptorStorage taskDescriptorStorage,
FormatOptions formatOptions)
{
try (SetThreadName _ = new SetThreadName("Query-" + stateMachine.getQueryId())) {
this.slug = requireNonNull(slug, "slug is null");
Expand Down Expand Up @@ -240,6 +244,7 @@ private SqlQueryExecution(
this.eventDrivenTaskSourceFactory = requireNonNull(eventDrivenTaskSourceFactory, "taskSourceFactory is null");
this.taskDescriptorStorage = requireNonNull(taskDescriptorStorage, "taskDescriptorStorage is null");
this.planOptimizersStatsCollector = requireNonNull(planOptimizersStatsCollector, "planOptimizersStatsCollector is null");
this.formatOptions = requireNonNull(formatOptions, "formatOptions is null");
}
}

Expand Down Expand Up @@ -503,6 +508,8 @@ private PlanRoot doPlanQuery(CachingTableStatsProvider tableStatsProvider)
Plan plan = logicalPlanner.plan(analysis);
queryPlan.set(plan);

System.out.println(ProgramBuilder.buildProgram(plan.getRoot()).print(1, formatOptions));

// fragment the plan
SubPlan fragmentedPlan;
try (var _ = scopedSpan(tracer, "fragment-plan")) {
Expand Down Expand Up @@ -809,6 +816,7 @@ public static class SqlQueryExecutionFactory
private final ExchangeManagerRegistry exchangeManagerRegistry;
private final EventDrivenTaskSourceFactory eventDrivenTaskSourceFactory;
private final TaskDescriptorStorage taskDescriptorStorage;
private final FormatOptions formatOptions;

@Inject
SqlQueryExecutionFactory(
Expand Down Expand Up @@ -841,7 +849,8 @@ public static class SqlQueryExecutionFactory
SqlTaskManager coordinatorTaskManager,
ExchangeManagerRegistry exchangeManagerRegistry,
EventDrivenTaskSourceFactory eventDrivenTaskSourceFactory,
TaskDescriptorStorage taskDescriptorStorage)
TaskDescriptorStorage taskDescriptorStorage,
FormatOptions formatOptions)
{
this.tracer = requireNonNull(tracer, "tracer is null");
this.schedulerStats = requireNonNull(schedulerStats, "schedulerStats is null");
Expand Down Expand Up @@ -875,6 +884,7 @@ public static class SqlQueryExecutionFactory
this.exchangeManagerRegistry = requireNonNull(exchangeManagerRegistry, "exchangeManagerRegistry is null");
this.eventDrivenTaskSourceFactory = requireNonNull(eventDrivenTaskSourceFactory, "eventDrivenTaskSourceFactory is null");
this.taskDescriptorStorage = requireNonNull(taskDescriptorStorage, "taskDescriptorStorage is null");
this.formatOptions = requireNonNull(formatOptions, "formatOptions is null");
}

@Override
Expand Down Expand Up @@ -925,7 +935,8 @@ public QueryExecution createQueryExecution(
coordinatorTaskManager,
exchangeManagerRegistry,
eventDrivenTaskSourceFactory,
taskDescriptorStorage);
taskDescriptorStorage,
formatOptions);
}
}
}

0 comments on commit efa6f93

Please sign in to comment.