Java Code Examples for io.prestosql.spi.connector.RecordSet#cursor()

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());
}
 

@Test
public void testGetRecordSet()
{
    ConnectorTableHandle tableHandle = new ExampleTableHandle("schema", "table");
    ExampleRecordSetProvider recordSetProvider = new ExampleRecordSetProvider();
    RecordSet recordSet = recordSetProvider.getRecordSet(ExampleTransactionHandle.INSTANCE, SESSION, new ExampleSplit(dataUri), tableHandle, ImmutableList.of(
            new ExampleColumnHandle("text", createUnboundedVarcharType(), 0),
            new ExampleColumnHandle("value", BIGINT, 1)));
    assertNotNull(recordSet, "recordSet is null");

    RecordCursor cursor = recordSet.cursor();
    assertNotNull(cursor, "cursor is null");

    Map<String, Long> data = new LinkedHashMap<>();
    while (cursor.advanceNextPosition()) {
        data.put(cursor.getSlice(0).toStringUtf8(), cursor.getLong(1));
    }
    assertEquals(data, ImmutableMap.<String, Long>builder()
            .put("ten", 10L)
            .put("eleven", 11L)
            .put("twelve", 12L)
            .build());
}
 

@Test
public void testCursorSimple()
{
    RecordSet recordSet = new ExampleRecordSet(new ExampleSplit(dataUri), ImmutableList.of(
            new ExampleColumnHandle("text", createUnboundedVarcharType(), 0),
            new ExampleColumnHandle("value", BIGINT, 1)));
    RecordCursor cursor = recordSet.cursor();

    assertEquals(cursor.getType(0), createUnboundedVarcharType());
    assertEquals(cursor.getType(1), BIGINT);

    Map<String, Long> data = new LinkedHashMap<>();
    while (cursor.advanceNextPosition()) {
        data.put(cursor.getSlice(0).toStringUtf8(), cursor.getLong(1));
        assertFalse(cursor.isNull(0));
        assertFalse(cursor.isNull(1));
    }
    assertEquals(data, ImmutableMap.<String, Long>builder()
            .put("ten", 10L)
            .put("eleven", 11L)
            .put("twelve", 12L)
            .build());
}
 

@Test
public void testCursorMixedOrder()
{
    RecordSet recordSet = new ExampleRecordSet(new ExampleSplit(dataUri), ImmutableList.of(
            new ExampleColumnHandle("value", BIGINT, 1),
            new ExampleColumnHandle("value", BIGINT, 1),
            new ExampleColumnHandle("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());
}
 

private List<Long> readTimeStampsFrom(RecordSet recordSet)
{
    ImmutableList.Builder<Long> result = ImmutableList.builder();
    try (RecordCursor cursor = recordSet.cursor()) {
        while (cursor.advanceNextPosition()) {
            for (int i = 0; i < recordSet.getColumnTypes().size(); i++) {
                cursor.isNull(i);
            }
            if (cursor.isNull(0)) {
                return result.build();
            }
            assertTrue(recordSet.getColumnTypes().get(0) instanceof TimestampType);
            result.add(cursor.getLong(0));
        }
    }
    return result.build();
}
 

public static PrestoThriftBlock fromIntBasedColumn(RecordSet recordSet, int columnIndex, int positions, BiFunction<boolean[], int[], PrestoThriftBlock> result)
{
    if (positions == 0) {
        return result.apply(null, null);
    }
    boolean[] nulls = null;
    int[] ints = null;
    RecordCursor cursor = recordSet.cursor();
    for (int position = 0; position < positions; position++) {
        checkState(cursor.advanceNextPosition(), "cursor has less values than expected");
        if (cursor.isNull(columnIndex)) {
            if (nulls == null) {
                nulls = new boolean[positions];
            }
            nulls[position] = true;
        }
        else {
            if (ints == null) {
                ints = new int[positions];
            }
            ints[position] = (int) cursor.getLong(columnIndex);
        }
    }
    checkState(!cursor.advanceNextPosition(), "cursor has more values than expected");
    return result.apply(nulls, ints);
}
 

public static PrestoThriftBlock fromLongBasedColumn(RecordSet recordSet, int columnIndex, int positions, BiFunction<boolean[], long[], PrestoThriftBlock> result)
{
    if (positions == 0) {
        return result.apply(null, null);
    }
    boolean[] nulls = null;
    long[] longs = null;
    RecordCursor cursor = recordSet.cursor();
    for (int position = 0; position < positions; position++) {
        checkState(cursor.advanceNextPosition(), "cursor has less values than expected");
        if (cursor.isNull(columnIndex)) {
            if (nulls == null) {
                nulls = new boolean[positions];
            }
            nulls[position] = true;
        }
        else {
            if (longs == null) {
                longs = new long[positions];
            }
            longs[position] = cursor.getLong(columnIndex);
        }
    }
    checkState(!cursor.advanceNextPosition(), "cursor has more values than expected");
    return result.apply(nulls, longs);
}
 

@Test
public void testGetRecordSet()
{
    ConnectorTransactionHandle transaction = new JdbcTransactionHandle();
    JdbcRecordSetProvider recordSetProvider = new JdbcRecordSetProvider(jdbcClient);
    RecordSet recordSet = recordSetProvider.getRecordSet(transaction, SESSION, split, table, ImmutableList.of(textColumn, textShortColumn, valueColumn));
    assertNotNull(recordSet, "recordSet is null");

    RecordCursor cursor = recordSet.cursor();
    assertNotNull(cursor, "cursor is null");

    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));
    }
    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());
}
 

private RecordCursor getCursor(JdbcTableHandle jdbcTableHandle, List<JdbcColumnHandle> columns, TupleDomain<ColumnHandle> domain)
{
    jdbcTableHandle = new JdbcTableHandle(
            jdbcTableHandle.getSchemaTableName(),
            jdbcTableHandle.getCatalogName(),
            jdbcTableHandle.getSchemaName(),
            jdbcTableHandle.getTableName(),
            domain,
            Optional.empty(),
            OptionalLong.empty(),
            Optional.empty());

    ConnectorSplitSource splits = jdbcClient.getSplits(SESSION, jdbcTableHandle);
    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, jdbcTableHandle, columns);

    return recordSet.cursor();
}
 

LazyRecordPageSource(int maxRowsPerPage, RecordSet recordSet)
{
    requireNonNull(recordSet, "recordSet is null");

    this.maxRowsPerPage = maxRowsPerPage;
    this.cursor = recordSet.cursor();
    this.types = ImmutableList.copyOf(recordSet.getColumnTypes());
    this.pageBuilder = new PageBuilder(this.types);
}
 

private static int totalRecords(RecordSet recordSet)
{
    RecordCursor cursor = recordSet.cursor();
    int result = 0;
    while (cursor.advanceNextPosition()) {
        result++;
    }
    return result;
}
 

private static Block convertColumnToBlock(RecordSet recordSet, int columnIndex, int positions)
{
    Type type = recordSet.getColumnTypes().get(columnIndex);
    BlockBuilder output = type.createBlockBuilder(null, positions);
    Class<?> javaType = type.getJavaType();
    RecordCursor cursor = recordSet.cursor();
    for (int position = 0; position < positions; position++) {
        checkState(cursor.advanceNextPosition(), "cursor has less values than expected");
        if (cursor.isNull(columnIndex)) {
            output.appendNull();
        }
        else {
            if (javaType == boolean.class) {
                type.writeBoolean(output, cursor.getBoolean(columnIndex));
            }
            else if (javaType == long.class) {
                type.writeLong(output, cursor.getLong(columnIndex));
            }
            else if (javaType == double.class) {
                type.writeDouble(output, cursor.getDouble(columnIndex));
            }
            else if (javaType == Slice.class) {
                Slice slice = cursor.getSlice(columnIndex);
                type.writeSlice(output, slice, 0, slice.length());
            }
            else {
                type.writeObject(output, cursor.getObject(columnIndex));
            }
        }
    }
    checkState(!cursor.advanceNextPosition(), "cursor has more values than expected");
    return output.build();
}
 

@Test
public void testIdempotentClose()
{
    RecordSet recordSet = createRecordSet(ImmutableList.of(
            columnHandles.get("value"),
            columnHandles.get("value"),
            columnHandles.get("text")));

    RecordCursor cursor = recordSet.cursor();
    cursor.close();
    cursor.close();
}
 

@Test
public void testCursorMixedOrder()
{
    RecordSet recordSet = createRecordSet(ImmutableList.of(
            columnHandles.get("value"),
            columnHandles.get("value"),
            columnHandles.get("text")));

    try (RecordCursor cursor = recordSet.cursor()) {
        assertEquals(cursor.getType(0), BIGINT);
        assertEquals(cursor.getType(1), BIGINT);
        assertEquals(cursor.getType(2), VARCHAR);

        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("one", 1L)
                .put("two", 2L)
                .put("three", 3L)
                .put("ten", 10L)
                .put("eleven", 11L)
                .put("twelve", 12L)
                .build());
    }
}
 

@Test
public void testCursorSimple()
{
    RecordSet recordSet = createRecordSet(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());
    }
}
 

@Test
public void testRecordSetProvider()
        throws Exception
{
    for (SchemaTableName schemaTableName : metadata.listTables(SESSION, Optional.of(JMX_SCHEMA_NAME))) {
        RecordSet recordSet = getRecordSet(schemaTableName);
        try (RecordCursor cursor = recordSet.cursor()) {
            while (cursor.advanceNextPosition()) {
                for (int i = 0; i < recordSet.getColumnTypes().size(); i++) {
                    cursor.isNull(i);
                }
            }
        }
    }
}
 

@Test
public void testGetRecordSet()
{
    ConnectorTableHandle tableHandle = new PrometheusTableHandle("schema", "table");
    PrometheusRecordSetProvider recordSetProvider = new PrometheusRecordSetProvider();
    RecordSet recordSet = recordSetProvider.getRecordSet(PrometheusTransactionHandle.INSTANCE, SESSION,
            new PrometheusSplit(dataUri), tableHandle, ImmutableList.of(
                    new PrometheusColumnHandle("labels", varcharMapType, 0),
                    new PrometheusColumnHandle("timestamp", TimestampType.TIMESTAMP, 1),
                    new PrometheusColumnHandle("value", DoubleType.DOUBLE, 2)));
    assertNotNull(recordSet, "recordSet is null");

    RecordCursor cursor = recordSet.cursor();
    assertNotNull(cursor, "cursor is null");

    Map<Timestamp, Map> actual = new LinkedHashMap<>();
    while (cursor.advanceNextPosition()) {
        actual.put((Timestamp) cursor.getObject(1), (Map) getMapFromBlock(varcharMapType, (Block) cursor.getObject(0)));
    }
    Map<Timestamp, Map> expected = ImmutableMap.<Timestamp, Map>builder()
            .put(Timestamp.from(Instant.ofEpochMilli(1565962969044L)), ImmutableMap.of("instance",
                    "localhost:9090", "__name__", "up", "job", "prometheus"))
            .put(Timestamp.from(Instant.ofEpochMilli(1565962984045L)), ImmutableMap.of("instance",
                    "localhost:9090", "__name__", "up", "job", "prometheus"))
            .put(Timestamp.from(Instant.ofEpochMilli(1565962999044L)), ImmutableMap.of("instance",
                    "localhost:9090", "__name__", "up", "job", "prometheus"))
            .put(Timestamp.from(Instant.ofEpochMilli(1565963014044L)), ImmutableMap.of("instance",
                    "localhost:9090", "__name__", "up", "job", "prometheus"))
            .build();
    assertEquals(actual, expected);
}
 

private static void assertData(LocalFileTables localFileTables, LocalFileMetadata metadata)
{
    LocalFileTableHandle tableHandle = new LocalFileTableHandle(getSchemaTableName(), OptionalInt.of(0), OptionalInt.of(-1));
    List<LocalFileColumnHandle> columnHandles = metadata.getColumnHandles(SESSION, tableHandle)
            .values().stream().map(column -> (LocalFileColumnHandle) column)
            .collect(Collectors.toList());

    LocalFileSplit split = new LocalFileSplit(address);
    RecordSet recordSet = new LocalFileRecordSet(localFileTables, split, tableHandle, columnHandles);
    RecordCursor cursor = recordSet.cursor();

    for (int i = 0; i < columnHandles.size(); i++) {
        assertEquals(cursor.getType(i), columnHandles.get(i).getColumnType());
    }

    // test one row
    assertTrue(cursor.advanceNextPosition());
    assertEquals(cursor.getSlice(0).toStringUtf8(), address.toString());
    assertEquals(cursor.getSlice(2).toStringUtf8(), "127.0.0.1");
    assertEquals(cursor.getSlice(3).toStringUtf8(), "POST");
    assertEquals(cursor.getSlice(4).toStringUtf8(), "/v1/memory");
    assertTrue(cursor.isNull(5));
    assertTrue(cursor.isNull(6));
    assertEquals(cursor.getLong(7), 200);
    assertEquals(cursor.getLong(8), 0);
    assertEquals(cursor.getLong(9), 1000);
    assertEquals(cursor.getLong(10), 10);
    assertTrue(cursor.isNull(11));

    assertTrue(cursor.advanceNextPosition());
    assertEquals(cursor.getSlice(0).toStringUtf8(), address.toString());
    assertEquals(cursor.getSlice(2).toStringUtf8(), "127.0.0.1");
    assertEquals(cursor.getSlice(3).toStringUtf8(), "GET");
    assertEquals(cursor.getSlice(4).toStringUtf8(), "/v1/service/presto/general");
    assertEquals(cursor.getSlice(5).toStringUtf8(), "foo");
    assertEquals(cursor.getSlice(6).toStringUtf8(), "ffffffff-ffff-ffff-ffff-ffffffffffff");
    assertEquals(cursor.getLong(7), 200);
    assertEquals(cursor.getLong(8), 0);
    assertEquals(cursor.getLong(9), 37);
    assertEquals(cursor.getLong(10), 1094);
    assertEquals(cursor.getSlice(11).toStringUtf8(), "a7229d56-5cbd-4e23-81ff-312ba6be0f12");
}
 

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, toIntExact(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();
    }
}
 

@Test
public void testCursorSimple()
{
    RecordSet recordSet = new PrometheusRecordSet(new PrometheusSplit(dataUri), ImmutableList.of(
            new PrometheusColumnHandle("labels", varcharMapType, 0),
            new PrometheusColumnHandle("timestamp", TimestampType.TIMESTAMP, 1),
            new PrometheusColumnHandle("value", DoubleType.DOUBLE, 2)));
    RecordCursor cursor = recordSet.cursor();

    assertEquals(cursor.getType(0), varcharMapType);
    assertEquals(cursor.getType(1), TimestampType.TIMESTAMP);
    assertEquals(cursor.getType(2), DoubleType.DOUBLE);

    List<PrometheusStandardizedRow> actual = new ArrayList<>();
    while (cursor.advanceNextPosition()) {
        actual.add(new PrometheusStandardizedRow(
                (Block) cursor.getObject(0),
                (Timestamp) cursor.getObject(1),
                cursor.getDouble(2)));
        assertFalse(cursor.isNull(0));
        assertFalse(cursor.isNull(1));
        assertFalse(cursor.isNull(2));
    }
    List<PrometheusStandardizedRow> expected = ImmutableList.<PrometheusStandardizedRow>builder()
            .add(new PrometheusStandardizedRow(getBlockFromMap(varcharMapType,
                    ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")), Timestamp.from(Instant.ofEpochMilli(1565962969044L)), 1.0))
            .add(new PrometheusStandardizedRow(getBlockFromMap(varcharMapType,
                    ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")), Timestamp.from(Instant.ofEpochMilli(1565962984045L)), 1.0))
            .add(new PrometheusStandardizedRow(getBlockFromMap(varcharMapType,
                    ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")), Timestamp.from(Instant.ofEpochMilli(1565962999044L)), 1.0))
            .add(new PrometheusStandardizedRow(getBlockFromMap(varcharMapType,
                    ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")), Timestamp.from(Instant.ofEpochMilli(1565963014044L)), 1.0))
            .build();
    List<PairLike<PrometheusStandardizedRow, PrometheusStandardizedRow>> pairs = Streams.zip(actual.stream(), expected.stream(), PairLike::new)
            .collect(Collectors.toList());
    pairs.stream().forEach(pair -> {
        assertEquals(getMapFromBlock(varcharMapType, pair.first.labels), getMapFromBlock(varcharMapType, pair.second.labels));
        assertEquals(pair.first.timestamp, pair.second.timestamp);
        assertEquals(pair.first.value, pair.second.value);
    });
}