Examples with RecordSet io.prestosql.spi.connector.RecordSet used on opensource projects

Example 1 with RecordSet

use of io.prestosql.spi.connector.RecordSet in project hetu-core by openlookeng.

the class TpchIndexedData method indexTable.

private static IndexedTable indexTable(RecordSet recordSet, final List<String> outputColumns, List<String> keyColumns) {
    List<Integer> keyPositions = keyColumns.stream().map(columnName -> {
        int position = outputColumns.indexOf(columnName);
        checkState(position != -1);
        return position;
    }).collect(toImmutableList());
    ImmutableListMultimap.Builder<MaterializedTuple, MaterializedTuple> indexedValuesBuilder = ImmutableListMultimap.builder();
    List<Type> outputTypes = recordSet.getColumnTypes();
    List<Type> keyTypes = extractPositionValues(outputTypes, keyPositions);
    RecordCursor cursor = recordSet.cursor();
    while (cursor.advanceNextPosition()) {
        List<Object> values = extractValues(cursor, outputTypes);
        List<Object> keyValues = extractPositionValues(values, keyPositions);
        indexedValuesBuilder.put(new MaterializedTuple(keyValues), new MaterializedTuple(values));
    }
    return new IndexedTable(keyColumns, keyTypes, outputColumns, outputTypes, indexedValuesBuilder.build());
}

Example 2 with RecordSet

use of io.prestosql.spi.connector.RecordSet in project hetu-core by openlookeng.

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 3 with RecordSet

use of io.prestosql.spi.connector.RecordSet in project hetu-core by openlookeng.

the class TestHBase method testGetSRecordSet.

/**
 * testGetSRecordSet
 */
@Test
public void testGetSRecordSet() {
    List<HostAddress> hostAddressList = new ArrayList<>(1);
    Map<Integer, List<Range>> ranges = new HashMap<>();
    HBaseSplit split = new HBaseSplit("rowkey", TestUtils.createHBaseTableHandle(), hostAddressList, null, null, ranges, 0, false, null);
    HBaseRecordSetProvider hrsp = new HBaseRecordSetProvider(hconn);
    RecordSet rs = hrsp.getRecordSet(new HBaseTransactionHandle(), session, split, TestUtils.createHBaseTableHandle(), hconn.getTable("hbase.test_table").getColumns());
    assertEquals(5, rs.getColumnTypes().size());
}

Example 4 with RecordSet

use of io.prestosql.spi.connector.RecordSet in project hetu-core by openlookeng.

the class HBasePageSourceProvider method createPageSource.

@Override
public ConnectorPageSource createPageSource(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<ColumnHandle> columns) {
    // if delete rows, we should replace $rowId -> real rowkey name
    List<ColumnHandle> columnsReplaceRowKey = new ArrayList<>();
    columns.forEach(ch -> {
        if (Constants.HBASE_ROWID_NAME.equals(ch.getColumnName())) {
            if (ch instanceof HBaseColumnHandle && table instanceof HBaseTableHandle) {
                HBaseColumnHandle rowColumnHandle = (HBaseColumnHandle) ch;
                columnsReplaceRowKey.add(new HBaseColumnHandle(((HBaseTableHandle) table).getRowId(), rowColumnHandle.getFamily(), rowColumnHandle.getQualifier(), rowColumnHandle.getType(), rowColumnHandle.getOrdinal(), rowColumnHandle.getComment(), rowColumnHandle.isIndexed()));
            }
        } else {
            columnsReplaceRowKey.add(ch);
        }
    });
    RecordSet recordSet = recordSetProvider.getRecordSet(transactionHandle, session, split, table, columnsReplaceRowKey);
    HBaseRecordSet hbaseRecordSet = null;
    if (recordSet instanceof HBaseRecordSet) {
        hbaseRecordSet = (HBaseRecordSet) recordSet;
    }
    if (columnsReplaceRowKey.stream().anyMatch(ch -> (ch instanceof HBaseColumnHandle) && (table instanceof HBaseTableHandle) && ((HBaseColumnHandle) ch).getOrdinal() == ((HBaseTableHandle) table).getRowIdOrdinal())) {
        return new HBaseUpdatablePageSource(hbaseRecordSet, hbaseConnection);
    } else {
        return new RecordPageSource(recordSet);
    }
}

Example 5 with RecordSet

use of io.prestosql.spi.connector.RecordSet in project hetu-core by openlookeng.

the class TestJdbcRecordSet method testCursorSimple.

@Test
public void testCursorSimple() {
    RecordSet recordSet = new JdbcRecordSet(jdbcClient, SESSION, split, table, ImmutableList.of(columnHandles.get("text"), columnHandles.get("text_short"), columnHandles.get("value")));
    try (RecordCursor cursor = recordSet.cursor()) {
        assertEquals(cursor.getType(0), VARCHAR);
        assertEquals(cursor.getType(1), createVarcharType(32));
        assertEquals(cursor.getType(2), BIGINT);
        Map<String, Long> data = new LinkedHashMap<>();
        while (cursor.advanceNextPosition()) {
            data.put(cursor.getSlice(0).toStringUtf8(), cursor.getLong(2));
            assertEquals(cursor.getSlice(0), cursor.getSlice(1));
            assertFalse(cursor.isNull(0));
            assertFalse(cursor.isNull(1));
            assertFalse(cursor.isNull(2));
        }
        assertEquals(data, ImmutableMap.<String, Long>builder().put("one", 1L).put("two", 2L).put("three", 3L).put("ten", 10L).put("eleven", 11L).put("twelve", 12L).build());
    }
}