Examples with MaterializedRow io.prestosql.testing.MaterializedRow used on opensource projects

Example 16 with MaterializedRow

use of io.prestosql.testing.MaterializedRow in project hetu-core by openlookeng.

the class H2QueryRunner method rowMapper.

private static RowMapper<MaterializedRow> rowMapper(List<? extends Type> types) {
    return new RowMapper<MaterializedRow>() {

        @Override
        public MaterializedRow map(ResultSet resultSet, StatementContext context) throws SQLException {
            int count = resultSet.getMetaData().getColumnCount();
            checkArgument(types.size() == count, "expected types count (%s) does not match actual column count (%s)", types.size(), count);
            List<Object> row = new ArrayList<>(count);
            for (int i = 1; i <= count; i++) {
                Type type = types.get(i - 1);
                if (BOOLEAN.equals(type)) {
                    boolean booleanValue = resultSet.getBoolean(i);
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(booleanValue);
                    }
                } else if (TINYINT.equals(type)) {
                    byte byteValue = resultSet.getByte(i);
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(byteValue);
                    }
                } else if (SMALLINT.equals(type)) {
                    short shortValue = resultSet.getShort(i);
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(shortValue);
                    }
                } else if (INTEGER.equals(type)) {
                    int intValue = resultSet.getInt(i);
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(intValue);
                    }
                } else if (BIGINT.equals(type)) {
                    long longValue = resultSet.getLong(i);
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(longValue);
                    }
                } else if (REAL.equals(type)) {
                    float floatValue = resultSet.getFloat(i);
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(floatValue);
                    }
                } else if (DOUBLE.equals(type)) {
                    double doubleValue = resultSet.getDouble(i);
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(doubleValue);
                    }
                } else if (JSON.equals(type)) {
                    String stringValue = resultSet.getString(i);
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(jsonParse(utf8Slice(stringValue)).toStringUtf8());
                    }
                } else if (isVarcharType(type)) {
                    String stringValue = resultSet.getString(i);
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(stringValue);
                    }
                } else if (isCharType(type)) {
                    String stringValue = resultSet.getString(i);
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(padEnd(stringValue, ((CharType) type).getLength(), ' '));
                    }
                } else if (DATE.equals(type)) {
                    // resultSet.getDate(i) doesn't work if JVM's zone skipped day being retrieved (e.g. 2011-12-30 and Pacific/Apia zone)
                    LocalDate dateValue = resultSet.getObject(i, LocalDate.class);
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(dateValue);
                    }
                } else if (TIME.equals(type)) {
                    // resultSet.getTime(i) doesn't work if JVM's zone had forward offset change during 1970-01-01 (e.g. America/Hermosillo zone)
                    LocalTime timeValue = resultSet.getObject(i, LocalTime.class);
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(timeValue);
                    }
                } else if (TIME_WITH_TIME_ZONE.equals(type)) {
                    throw new UnsupportedOperationException("H2 does not support TIME WITH TIME ZONE");
                } else if (TIMESTAMP.equals(type)) {
                    // resultSet.getTimestamp(i) doesn't work if JVM's zone had forward offset at the date/time being retrieved
                    LocalDateTime timestampValue;
                    try {
                        timestampValue = resultSet.getObject(i, LocalDateTime.class);
                    } catch (SQLException first) {
                        // H2 cannot convert DATE to LocalDateTime in their JDBC driver (even though it can convert to java.sql.Timestamp), we need to do this manually
                        try {
                            timestampValue = Optional.ofNullable(resultSet.getObject(i, LocalDate.class)).map(LocalDate::atStartOfDay).orElse(null);
                        } catch (RuntimeException e) {
                            first.addSuppressed(e);
                            throw first;
                        }
                    }
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(timestampValue);
                    }
                } else if (TIMESTAMP_WITH_TIME_ZONE.equals(type)) {
                    // This means H2 is unsuitable for testing TIMESTAMP WITH TIME ZONE-bearing queries. Those need to be tested manually.
                    throw new UnsupportedOperationException();
                } else if (UNKNOWN.equals(type)) {
                    Object objectValue = resultSet.getObject(i);
                    checkState(resultSet.wasNull(), "Expected a null value, but got %s", objectValue);
                    row.add(null);
                } else if (type instanceof DecimalType) {
                    DecimalType decimalType = (DecimalType) type;
                    BigDecimal decimalValue = resultSet.getBigDecimal(i);
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(decimalValue.setScale(decimalType.getScale(), BigDecimal.ROUND_HALF_UP).round(new MathContext(decimalType.getPrecision())));
                    }
                } else if (type instanceof ArrayType) {
                    Array array = resultSet.getArray(i);
                    if (resultSet.wasNull()) {
                        row.add(null);
                    } else {
                        row.add(newArrayList((Object[]) array.getArray()));
                    }
                } else {
                    throw new AssertionError("unhandled type: " + type);
                }
            }
            return new MaterializedRow(MaterializedResult.DEFAULT_PRECISION, row);
        }
    };
}

Example 17 with MaterializedRow

use of io.prestosql.testing.MaterializedRow in project hetu-core by openlookeng.

the class QueryAssertions method assertQuery.

private static void assertQuery(QueryRunner actualQueryRunner, Session actualQuerySession, @Language("SQL") String actual, H2QueryRunner h2QueryRunner, Session expectedQuerySession, @Language("SQL") String expected, boolean ensureOrdering, boolean compareUpdate, Optional<Consumer<Plan>> planAssertion) {
    long start = System.nanoTime();
    Optional<MaterializedResult> actualResultsOp = Optional.empty();
    Plan queryPlan = null;
    if (planAssertion.isPresent()) {
        try {
            MaterializedResultWithPlan resultWithPlan = actualQueryRunner.executeWithPlan(actualQuerySession, actual, WarningCollector.NOOP);
            queryPlan = resultWithPlan.getQueryPlan();
            actualResultsOp = Optional.of(resultWithPlan.getMaterializedResult().toTestTypes());
        } catch (RuntimeException ex) {
            fail("Execution of 'actual' query failed: " + actual, ex);
        }
    } else {
        try {
            actualResultsOp = Optional.of(actualQueryRunner.execute(actualQuerySession, 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();
    Optional<MaterializedResult> expectedResultsOp = Optional.empty();
    try {
        expectedResultsOp = Optional.of(h2QueryRunner.execute(expectedQuerySession, expected, actualResultsOp.get().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 presto: %s, h2: %s, total: %s", actualTime, nanosSince(expectedStart), totalTime);
    }
    if (actualResultsOp.get().getUpdateType().isPresent() || actualResultsOp.get().getUpdateCount().isPresent()) {
        if (!actualResultsOp.get().getUpdateType().isPresent()) {
            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 = actualResultsOp.get().getMaterializedRows();
    List<MaterializedRow> expectedRows = expectedResultsOp.get().getMaterializedRows();
    if (compareUpdate) {
        if (!actualResultsOp.get().getUpdateType().isPresent()) {
            fail("update type not present for query: \n" + actual);
        }
        if (!actualResultsOp.get().getUpdateCount().isPresent()) {
            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), actualResultsOp.get().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 18 with MaterializedRow

use of io.prestosql.testing.MaterializedRow in project hetu-core by openlookeng.

the class TestCachedSqlQueryExecution method testQueriesWithCurrentPathFunction.

@Test
public void testQueriesWithCurrentPathFunction() throws Exception {
    setupWithExecutionPlanCacheEnabled(DEFAULT_SESSION);
    TestingMetadata metadata = new TestingMetadata();
    queryRunner.installPlugin(new TestPlugin(metadata));
    queryRunner.createCatalog("test", "test");
    SqlQueryManager manager = (SqlQueryManager) queryRunner.getCoordinator().getQueryManager();
    String query1 = "SELECT CURRENT_PATH";
    Plan plan1 = getPlan(query1, manager);
    setupWithExecutionPlanCacheEnabled(TPCH_SESSION);
    TestingMetadata metadata2 = new TestingMetadata();
    queryRunner.installPlugin(new TestPlugin(metadata2));
    queryRunner.createCatalog("test", "test");
    SqlQueryManager manager2 = (SqlQueryManager) queryRunner.getCoordinator().getQueryManager();
    Plan plan2 = getPlan(query1, manager2);
    assertNotNull(plan1);
    assertNotNull(plan2);
    // Should not have the same plan
    assertNotSame(plan1.getStatsAndCosts(), plan2.getStatsAndCosts());
    MaterializedResult rows1 = queryRunner.execute(DEFAULT_SESSION, query1);
    MaterializedResult rows2 = queryRunner.execute(TPCH_SESSION, query1);
    assertEquals(rows1.getRowCount(), rows2.getRowCount());
    MaterializedRow row1 = Iterables.getOnlyElement(rows1);
    MaterializedRow row2 = Iterables.getOnlyElement(rows2);
    assertEquals(row1.getFieldCount(), row2.getFieldCount());
    assertNotSame(row1.getFields(), row2.getFields());
}

Example 19 with MaterializedRow

use of io.prestosql.testing.MaterializedRow in project hetu-core by openlookeng.

the class TestCachedSqlQueryExecution method testQueriesWithCurrentTimeFunction.

@Test
public void testQueriesWithCurrentTimeFunction() throws Exception {
    setupWithExecutionPlanCacheEnabled(DEFAULT_SESSION);
    TestingMetadata metadata = new TestingMetadata();
    queryRunner.installPlugin(new TestPlugin(metadata));
    queryRunner.createCatalog("test", "test");
    SqlQueryManager manager = (SqlQueryManager) queryRunner.getCoordinator().getQueryManager();
    String query1 = "SELECT CURRENT_TIME";
    Plan plan1 = getPlan(query1, manager);
    Plan plan2 = getPlan(query1, manager);
    assertNotNull(plan1);
    assertNotNull(plan2);
    // Should not have the same plan
    assertNotSame(plan1.getStatsAndCosts(), plan2.getStatsAndCosts());
    MaterializedResult rows1 = queryRunner.execute(DEFAULT_SESSION, query1);
    MaterializedResult rows2 = queryRunner.execute(DEFAULT_SESSION, query1);
    assertEquals(rows1.getRowCount(), rows2.getRowCount());
    MaterializedRow row1 = Iterables.getOnlyElement(rows1);
    MaterializedRow row2 = Iterables.getOnlyElement(rows2);
    assertEquals(row1.getFieldCount(), row2.getFieldCount());
    assertNotSame(row1.getFields(), row2.getFields());
}

Example 20 with MaterializedRow

use of io.prestosql.testing.MaterializedRow in project hetu-core by openlookeng.

the class TestHashJoinOperator method testOuterJoinSnapshot.

@Test(dataProvider = "hashJoinTestValues")
public void testOuterJoinSnapshot(boolean parallelBuild, boolean probeHashEnabled, boolean buildHashEnabled) throws ExecutionException, InterruptedException {
    TaskContext taskContext = createSnapshotTaskContext();
    // build factory
    List<Type> buildTypes = ImmutableList.of(VARCHAR, BIGINT, BIGINT);
    RowPagesBuilder buildPages = rowPagesBuilder(buildHashEnabled, Ints.asList(0), ImmutableList.of(VARCHAR, BIGINT, BIGINT)).addSequencePage(10, 20, 30, 40);
    BuildSideSetup buildSideSetup = setupBuildSide(parallelBuild, taskContext, Ints.asList(0), buildPages, Optional.empty(), false, SINGLE_STREAM_SPILLER_FACTORY, true);
    JoinBridgeManager<PartitionedLookupSourceFactory> lookupSourceFactory = buildSideSetup.getLookupSourceFactoryManager();
    // probe factory
    List<Type> probeTypes = ImmutableList.of(VARCHAR, BIGINT, BIGINT);
    RowPagesBuilder probePages = rowPagesBuilder(probeHashEnabled, Ints.asList(0), probeTypes);
    List<Page> probeInput = probePages.addSequencePage(10, 15, 1020, 2020).build();
    // Add markers between pages
    List<Page> withMarkers = new ArrayList<>();
    long snapshotId = 1;
    for (Page page : probeInput) {
        withMarkers.add(page);
        withMarkers.add(MarkerPage.snapshotPage(snapshotId++));
    }
    probeInput = withMarkers;
    OperatorFactory joinOperatorFactory = new LookupJoinOperators().fullOuterJoin(0, new PlanNodeId("test"), lookupSourceFactory, probePages.getTypes(), Ints.asList(0), getHashChannelAsInt(probePages), Optional.empty(), OptionalInt.of(1), PARTITIONING_SPILLER_FACTORY);
    // build drivers and operators
    instantiateBuildDrivers(buildSideSetup, taskContext);
    buildLookupSource(buildSideSetup);
    // Construct lookup-outer operator
    PipelineContext buildPipeline = taskContext.addPipelineContext(2, false, true, false);
    DriverContext outerDriverContext = buildPipeline.addDriverContext(Lifespan.taskWide(), 0);
    Operator lookupOuter = ((LookupJoinOperatorFactory) joinOperatorFactory).createOuterOperatorFactory().get().getOuterOperatorFactory().createOperator(outerDriverContext);
    assertFalse(lookupOuter.isBlocked().isDone());
    // expected
    MaterializedResult expected = MaterializedResult.resultBuilder(taskContext.getSession(), concat(probeTypes, buildTypes)).row("15", 1020L, 2020L, null, null, null).row("16", 1021L, 2021L, null, null, null).row("17", 1022L, 2022L, null, null, null).row("18", 1023L, 2023L, null, null, null).row("19", 1024L, 2024L, null, null, null).row("20", 1025L, 2025L, "20", 30L, 40L).row("21", 1026L, 2026L, "21", 31L, 41L).row("22", 1027L, 2027L, "22", 32L, 42L).row("23", 1028L, 2028L, "23", 33L, 43L).row("24", 1029L, 2029L, "24", 34L, 44L).build();
    List<Integer> hashChannels = getHashChannels(probePages, buildPages);
    assertOperatorEquals(joinOperatorFactory, taskContext.addPipelineContext(0, true, true, false).addDriverContext(), probeInput, expected, true, hashChannels);
    assertTrue(lookupOuter.isBlocked().isDone());
    assertTrue(lookupOuter.getOutput() instanceof MarkerPage);
    Object state = lookupOuter.capture(null);
    lookupOuter.isBlocked().get();
    Page page = lookupOuter.getOutput();
    page = dropChannel(ImmutableList.of(page), hashChannels).get(0);
    MaterializedResult outerExpected = MaterializedResult.resultBuilder(taskContext.getSession(), concat(probeTypes, buildTypes)).row(null, null, null, "25", 35L, 45L).row(null, null, null, "26", 36L, 46L).row(null, null, null, "27", 37L, 47L).row(null, null, null, "28", 38L, 48L).row(null, null, null, "29", 39L, 49L).build();
    List<MaterializedRow> rows = MaterializedResult.resultBuilder(taskContext.getSession(), concat(probeTypes, buildTypes)).page(page).build().getMaterializedRows();
    rows = new ArrayList<>(rows);
    rows.sort(Comparator.comparing(a -> (String) a.getField(3)));
    MaterializedResult outerResult = MaterializedResult.resultBuilder(taskContext.getSession(), concat(probeTypes, buildTypes)).rows(rows).build();
    assertEquals(outerResult, outerExpected);
    long matched;
    if (state instanceof boolean[]) {
        boolean[] positions = (boolean[]) state;
        matched = Booleans.asList(positions).stream().filter(e -> e).count();
    } else {
        ByteBuffer bb = ByteBuffer.wrap((byte[]) state);
        List<RoaringBitmap> visitedPositions = new ArrayList<>();
        for (int i = 0; i < (parallelBuild ? PARTITION_COUNT : 1); i++) {
            ImmutableRoaringBitmap bm = new ImmutableRoaringBitmap(bb);
            visitedPositions.add(new RoaringBitmap(bm));
            bb.position(bb.position() + visitedPositions.get(i).serializedSizeInBytes());
        }
        matched = visitedPositions.stream().mapToLong(rr -> rr.getCardinality()).sum();
    }
    assertEquals(matched, 5);
}