Examples with RecordSet com.facebook.presto.spi.RecordSet used on opensource projects

Example 21 with RecordSet

use of com.facebook.presto.spi.RecordSet in project presto by prestodb.

the class TestExampleRecordSet method testCursorMixedOrder.

@Test
public void testCursorMixedOrder() {
    RecordSet recordSet = new ExampleRecordSet(new ExampleSplit("test", "schema", "table", dataUri), ImmutableList.of(new ExampleColumnHandle("test", "value", BIGINT, 1), new ExampleColumnHandle("test", "value", BIGINT, 1), new ExampleColumnHandle("test", "text", createUnboundedVarcharType(), 0)));
    RecordCursor cursor = recordSet.cursor();
    Map<String, Long> data = new LinkedHashMap<>();
    while (cursor.advanceNextPosition()) {
        assertEquals(cursor.getLong(0), cursor.getLong(1));
        data.put(cursor.getSlice(2).toStringUtf8(), cursor.getLong(0));
    }
    assertEquals(data, ImmutableMap.<String, Long>builder().put("ten", 10L).put("eleven", 11L).put("twelve", 12L).build());
}

Example 22 with RecordSet

use of com.facebook.presto.spi.RecordSet in project presto by prestodb.

the class TestJdbcRecordSetProvider method getCursor.

private RecordCursor getCursor(JdbcTableHandle jdbcTableHandle, List<JdbcColumnHandle> columns, TupleDomain<ColumnHandle> domain) {
    JdbcTableLayoutHandle layoutHandle = new JdbcTableLayoutHandle(SESSION.getSqlFunctionProperties(), jdbcTableHandle, domain, Optional.empty());
    ConnectorSplitSource splits = jdbcClient.getSplits(IDENTITY, layoutHandle);
    JdbcSplit split = (JdbcSplit) getOnlyElement(getFutureValue(splits.getNextBatch(NOT_PARTITIONED, 1000)).getSplits());
    ConnectorTransactionHandle transaction = new JdbcTransactionHandle();
    JdbcRecordSetProvider recordSetProvider = new JdbcRecordSetProvider(jdbcClient);
    RecordSet recordSet = recordSetProvider.getRecordSet(transaction, SESSION, split, columns);
    return recordSet.cursor();
}

Example 23 with RecordSet

use of com.facebook.presto.spi.RecordSet in project presto by prestodb.

the class TestJdbcRecordSet method testGetColumnTypes.

@Test
public void testGetColumnTypes() {
    RecordSet recordSet = new JdbcRecordSet(jdbcClient, session, split, ImmutableList.of(new JdbcColumnHandle("test", "text", JDBC_VARCHAR, VARCHAR, true, Optional.empty()), new JdbcColumnHandle("test", "text_short", JDBC_VARCHAR, createVarcharType(32), true, Optional.empty()), new JdbcColumnHandle("test", "value", JDBC_BIGINT, BIGINT, true, Optional.empty())));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(VARCHAR, createVarcharType(32), BIGINT));
    recordSet = new JdbcRecordSet(jdbcClient, session, split, ImmutableList.of(new JdbcColumnHandle("test", "value", JDBC_BIGINT, BIGINT, true, Optional.empty()), new JdbcColumnHandle("test", "text", JDBC_VARCHAR, VARCHAR, true, Optional.empty())));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, VARCHAR));
    recordSet = new JdbcRecordSet(jdbcClient, session, split, ImmutableList.of(new JdbcColumnHandle("test", "value", JDBC_BIGINT, BIGINT, true, Optional.empty()), new JdbcColumnHandle("test", "value", JDBC_BIGINT, BIGINT, true, Optional.empty()), new JdbcColumnHandle("test", "text", JDBC_VARCHAR, VARCHAR, true, Optional.empty())));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, BIGINT, VARCHAR));
    recordSet = new JdbcRecordSet(jdbcClient, session, split, ImmutableList.of());
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of());
}

Example 24 with RecordSet

use of com.facebook.presto.spi.RecordSet in project presto by prestodb.

the class H2QueryRunner method insertRows.

private static void insertRows(ConnectorTableMetadata tableMetadata, Handle handle, RecordSet data) {
    List<ColumnMetadata> columns = tableMetadata.getColumns().stream().filter(columnMetadata -> !columnMetadata.isHidden()).collect(toImmutableList());
    String vars = Joiner.on(',').join(nCopies(columns.size(), "?"));
    String sql = format("INSERT INTO %s VALUES (%s)", tableMetadata.getTable().getTableName(), vars);
    RecordCursor cursor = data.cursor();
    while (true) {
        // insert 1000 rows at a time
        PreparedBatch batch = handle.prepareBatch(sql);
        for (int row = 0; row < 1000; row++) {
            if (!cursor.advanceNextPosition()) {
                if (batch.size() > 0) {
                    batch.execute();
                }
                return;
            }
            for (int column = 0; column < columns.size(); column++) {
                Type type = columns.get(column).getType();
                if (BOOLEAN.equals(type)) {
                    batch.bind(column, cursor.getBoolean(column));
                } else if (BIGINT.equals(type)) {
                    batch.bind(column, cursor.getLong(column));
                } else if (INTEGER.equals(type)) {
                    batch.bind(column, (int) cursor.getLong(column));
                } else if (DOUBLE.equals(type)) {
                    batch.bind(column, cursor.getDouble(column));
                } else if (type instanceof VarcharType) {
                    batch.bind(column, cursor.getSlice(column).toStringUtf8());
                } else if (DATE.equals(type)) {
                    long millisUtc = TimeUnit.DAYS.toMillis(cursor.getLong(column));
                    // H2 expects dates in to be millis at midnight in the JVM timezone
                    long localMillis = DateTimeZone.UTC.getMillisKeepLocal(DateTimeZone.getDefault(), millisUtc);
                    batch.bind(column, new Date(localMillis));
                } else {
                    throw new IllegalArgumentException("Unsupported type " + type);
                }
            }
            batch.add();
        }
        batch.execute();
    }
}

Example 25 with RecordSet

use of com.facebook.presto.spi.RecordSet in project presto by prestodb.

the class TpchConnectorIndex method lookup.

@Override
public ConnectorPageSource lookup(RecordSet rawInputRecordSet) {
    // convert the input record set from the column ordering in the query to
    // match the column ordering of the index
    RecordSet inputRecordSet = keyFormatter.apply(rawInputRecordSet);
    // lookup the values in the index
    RecordSet rawOutputRecordSet = indexedTable.lookupKeys(inputRecordSet);
    // expect by the query
    return new RecordPageSource(outputFormatter.apply(rawOutputRecordSet));
}