Examples with Plan io.trino.sql.planner.Plan used on opensource projects

Example 11 with Plan

use of io.trino.sql.planner.Plan in project trino by trinodb.

the class TestUnion method testUnionOverSingleNodeAggregationAndUnion.

@Test
public void testUnionOverSingleNodeAggregationAndUnion() {
    Plan plan = plan("SELECT count(*) FROM (" + "SELECT 1 FROM nation GROUP BY regionkey " + "UNION ALL (" + "   SELECT 1 FROM nation " + "   UNION ALL " + "   SELECT 1 FROM nation))", OPTIMIZED_AND_VALIDATED, false);
    List<PlanNode> remotes = searchFrom(plan.getRoot()).where(TestUnion::isRemoteExchange).findAll();
    assertEquals(remotes.size(), 2, "There should be exactly two RemoteExchanges");
    assertEquals(((ExchangeNode) remotes.get(0)).getType(), GATHER);
    assertEquals(((ExchangeNode) remotes.get(1)).getType(), REPARTITION);
}

Example 12 with Plan

use of io.trino.sql.planner.Plan in project trino by trinodb.

the class QueryAssertions method assertQuery.

private static void assertQuery(QueryRunner actualQueryRunner, Session session, @Language("SQL") String actual, H2QueryRunner h2QueryRunner, @Language("SQL") String expected, boolean ensureOrdering, boolean compareUpdate, Optional<Consumer<Plan>> planAssertion) {
    long start = System.nanoTime();
    MaterializedResult actualResults = null;
    Plan queryPlan = null;
    if (planAssertion.isPresent()) {
        try {
            MaterializedResultWithPlan resultWithPlan = actualQueryRunner.executeWithPlan(session, actual, WarningCollector.NOOP);
            queryPlan = resultWithPlan.getQueryPlan();
            actualResults = resultWithPlan.getMaterializedResult().toTestTypes();
        } catch (RuntimeException ex) {
            fail("Execution of 'actual' query failed: " + actual, ex);
        }
    } else {
        try {
            actualResults = actualQueryRunner.execute(session, actual).toTestTypes();
        } catch (RuntimeException ex) {
            fail("Execution of 'actual' query failed: " + actual, ex);
        }
    }
    if (planAssertion.isPresent()) {
        planAssertion.get().accept(queryPlan);
    }
    Duration actualTime = nanosSince(start);
    long expectedStart = System.nanoTime();
    MaterializedResult expectedResults = null;
    try {
        expectedResults = h2QueryRunner.execute(session, expected, actualResults.getTypes());
    } catch (RuntimeException ex) {
        fail("Execution of 'expected' query failed: " + expected, ex);
    }
    Duration totalTime = nanosSince(start);
    if (totalTime.compareTo(Duration.succinctDuration(1, SECONDS)) > 0) {
        log.info("FINISHED in Trino: %s, H2: %s, total: %s", actualTime, nanosSince(expectedStart), totalTime);
    }
    if (actualResults.getUpdateType().isPresent() || actualResults.getUpdateCount().isPresent()) {
        if (actualResults.getUpdateType().isEmpty()) {
            fail("update count present without update type for query: \n" + actual);
        }
        if (!compareUpdate) {
            fail("update type should not be present (use assertUpdate) for query: \n" + actual);
        }
    }
    List<MaterializedRow> actualRows = actualResults.getMaterializedRows();
    List<MaterializedRow> expectedRows = expectedResults.getMaterializedRows();
    if (compareUpdate) {
        if (actualResults.getUpdateType().isEmpty()) {
            fail("update type not present for query: \n" + actual);
        }
        if (actualResults.getUpdateCount().isEmpty()) {
            fail("update count not present for query: \n" + actual);
        }
        assertEquals(actualRows.size(), 1, "For query: \n " + actual + "\n:");
        assertEquals(expectedRows.size(), 1, "For query: \n " + actual + "\n:");
        MaterializedRow row = expectedRows.get(0);
        assertEquals(row.getFieldCount(), 1, "For query: \n " + actual + "\n:");
        assertEquals(row.getField(0), actualResults.getUpdateCount().getAsLong(), "For query: \n " + actual + "\n:");
    }
    if (ensureOrdering) {
        if (!actualRows.equals(expectedRows)) {
            assertEquals(actualRows, expectedRows, "For query: \n " + actual + "\n:");
        }
    } else {
        assertEqualsIgnoreOrder(actualRows, expectedRows, "For query: \n " + actual);
    }
}

Example 13 with Plan

use of io.trino.sql.planner.Plan in project trino by trinodb.

the class QueryAssertions method assertUpdate.

public static void assertUpdate(QueryRunner queryRunner, Session session, @Language("SQL") String sql, OptionalLong count, Optional<Consumer<Plan>> planAssertion) {
    long start = System.nanoTime();
    MaterializedResult results;
    Plan queryPlan;
    if (planAssertion.isPresent()) {
        MaterializedResultWithPlan resultWithPlan = queryRunner.executeWithPlan(session, sql, WarningCollector.NOOP);
        queryPlan = resultWithPlan.getQueryPlan();
        results = resultWithPlan.getMaterializedResult().toTestTypes();
    } else {
        queryPlan = null;
        results = queryRunner.execute(session, sql);
    }
    Duration queryTime = nanosSince(start);
    if (queryTime.compareTo(Duration.succinctDuration(1, SECONDS)) > 0) {
        log.info("FINISHED in Trino: %s", queryTime);
    }
    if (planAssertion.isPresent()) {
        planAssertion.get().accept(queryPlan);
    }
    if (results.getUpdateType().isEmpty()) {
        fail("update type is not set");
    }
    if (results.getUpdateCount().isPresent()) {
        if (count.isEmpty()) {
            fail("expected no update count, but got " + results.getUpdateCount().getAsLong());
        }
        assertEquals(results.getUpdateCount().getAsLong(), count.getAsLong(), "update count");
    } else if (count.isPresent()) {
        fail("update count is not present");
    }
}

Example 14 with Plan

use of io.trino.sql.planner.Plan in project trino by trinodb.

the class TestUnaliasSymbolReferences method assertOptimizedPlan.

private void assertOptimizedPlan(PlanOptimizer optimizer, PlanCreator planCreator, PlanMatchPattern pattern) {
    LocalQueryRunner queryRunner = getQueryRunner();
    queryRunner.inTransaction(session -> {
        Metadata metadata = queryRunner.getMetadata();
        session.getCatalog().ifPresent(catalog -> metadata.getCatalogHandle(session, catalog));
        PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator();
        PlanBuilder planBuilder = new PlanBuilder(idAllocator, metadata, session);
        SymbolAllocator symbolAllocator = new SymbolAllocator();
        PlanNode plan = planCreator.create(planBuilder, session, metadata);
        PlanNode optimized = optimizer.optimize(plan, session, planBuilder.getTypes(), symbolAllocator, idAllocator, WarningCollector.NOOP);
        Plan actual = new Plan(optimized, planBuilder.getTypes(), StatsAndCosts.empty());
        PlanAssert.assertPlan(session, queryRunner.getMetadata(), queryRunner.getFunctionManager(), queryRunner.getStatsCalculator(), actual, pattern);
        return null;
    });
}

Example 15 with Plan

use of io.trino.sql.planner.Plan in project trino by trinodb.

the class TestUnion method testAggregationWithUnionAndValues.

@Test
public void testAggregationWithUnionAndValues() {
    Plan plan = plan("SELECT regionkey, count(*) FROM (SELECT regionkey FROM nation UNION ALL SELECT * FROM (VALUES 2, 100) t(regionkey)) GROUP BY regionkey", OPTIMIZED_AND_VALIDATED, false);
    assertAtMostOneAggregationBetweenRemoteExchanges(plan);
// TODO: Enable this check once distributed UNION can handle both partitioned and single node sources at the same time
// assertPlanIsFullyDistributed(plan);
}