Java Code Examples for io.prestosql.spi.connector.RecordCursor#advanceNextPosition()

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

@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 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 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 int totalRecords(RecordSet recordSet)
{
    RecordCursor cursor = recordSet.cursor();
    int result = 0;
    while (cursor.advanceNextPosition()) {
        result++;
    }
    return result;
}
 

public TableStatisticsData recordStatistics(Table table, double scaleFactor)
{
    Session session = Session.getDefaultSession()
            .withScale(scaleFactor)
            .withParallelism(1)
            .withNoSexism(false);

    List<Column> columns = ImmutableList.copyOf(table.getColumns());
    RecordCursor recordCursor = new TpcdsRecordSet(Results.constructResults(table, session), columns)
            .cursor();

    List<ColumnStatisticsRecorder> statisticsRecorders = createStatisticsRecorders(columns);
    long rowCount = 0;

    while (recordCursor.advanceNextPosition()) {
        rowCount++;
        for (int columnId = 0; columnId < columns.size(); columnId++) {
            Comparable<?> value = getPrestoValue(recordCursor, columns, columnId);
            statisticsRecorders.get(columnId).record(value);
        }
    }

    Map<String, ColumnStatisticsData> columnSampleStatistics = IntStream.range(0, columns.size())
            .boxed()
            .collect(toImmutableMap(
                    i -> columns.get(i).getName(),
                    i -> statisticsRecorders.get(i).getRecording()));
    return new TableStatisticsData(rowCount, columnSampleStatistics);
}
 

private static List<MaterializedRow> getMaterializedResults(RecordCursor cursor, List<ColumnMetadata> columns)
{
    List<Type> types = columns.stream().map(ColumnMetadata::getType).collect(toList());

    ImmutableList.Builder<MaterializedRow> rowBuilder = ImmutableList.builder();
    for (int i = 0; i < types.size(); i++) {
        assertEquals(cursor.getType(i), types.get(i));
    }

    while (cursor.advanceNextPosition()) {
        List<Object> values = new ArrayList<>();
        for (int i = 0; i < columns.size(); i++) {
            Type type = columns.get(i).getType();
            Class<?> javaType = type.getJavaType();
            if (cursor.isNull(i)) {
                values.add(null);
            }
            else if (javaType == boolean.class) {
                values.add(cursor.getBoolean(i));
            }
            else if (javaType == long.class) {
                values.add(cursor.getLong(i));
            }
            else if (javaType == double.class) {
                values.add(cursor.getDouble(i));
            }
            else if (javaType == Slice.class) {
                values.add(cursor.getSlice(i));
            }
        }
        rowBuilder.add(new MaterializedRow(DEFAULT_PRECISION, values));
    }
    return rowBuilder.build();
}
 

private static void assertRecordCursor(List<Type> types, Iterator<?>[] valuesByField, RecordCursor cursor)
{
    while (cursor.advanceNextPosition()) {
        for (int field = 0; field < types.size(); field++) {
            assertTrue(valuesByField[field].hasNext());
            Object expected = valuesByField[field].next();
            Object actual = getActualCursorValue(cursor, types.get(field), field);
            assertEquals(actual, expected);
        }
    }
}
 

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

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