Java Code Examples for io.prestosql.spi.connector.RecordSet

The following examples show how to use io.prestosql.spi.connector.RecordSet. These examples are extracted from open source projects. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source Project: presto   Source File: TestPrometheusIntegrationTests1.java    License: Apache License 2.0 6 votes vote down vote up
@Test(dependsOnMethods = "testRetrieveUpValue")
public void testGetColumnTypes()
        throws Exception
{
    URI dataUri = new URI("http://" + server.getAddress().getHost() + ":" + server.getAddress().getPort() + "/");
    RecordSet recordSet = new PrometheusRecordSet(new PrometheusSplit(dataUri), ImmutableList.of(
            new PrometheusColumnHandle("labels", createUnboundedVarcharType(), 0),
            new PrometheusColumnHandle("value", DoubleType.DOUBLE, 1),
            new PrometheusColumnHandle("timestamp", TimestampType.TIMESTAMP, 2)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(createUnboundedVarcharType(), DoubleType.DOUBLE, TimestampType.TIMESTAMP));

    recordSet = new PrometheusRecordSet(new PrometheusSplit(dataUri), ImmutableList.of(
            new PrometheusColumnHandle("value", BIGINT, 1),
            new PrometheusColumnHandle("text", createUnboundedVarcharType(), 0)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, createUnboundedVarcharType()));

    recordSet = new PrometheusRecordSet(new PrometheusSplit(dataUri), ImmutableList.of(
            new PrometheusColumnHandle("value", BIGINT, 1),
            new PrometheusColumnHandle("value", BIGINT, 1),
            new PrometheusColumnHandle("text", createUnboundedVarcharType(), 0)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, BIGINT, createUnboundedVarcharType()));

    recordSet = new PrometheusRecordSet(new PrometheusSplit(dataUri), ImmutableList.of());
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of());
}
 
Example 2
Source Project: presto   Source File: TestExampleRecordSetProvider.java    License: Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example 3
Source Project: presto   Source File: TestExampleRecordSet.java    License: Apache License 2.0 6 votes vote down vote up
@Test
public void testGetColumnTypes()
{
    RecordSet recordSet = new ExampleRecordSet(new ExampleSplit(dataUri), ImmutableList.of(
            new ExampleColumnHandle("text", createUnboundedVarcharType(), 0),
            new ExampleColumnHandle("value", BIGINT, 1)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(createUnboundedVarcharType(), BIGINT));

    recordSet = new ExampleRecordSet(new ExampleSplit(dataUri), ImmutableList.of(
            new ExampleColumnHandle("value", BIGINT, 1),
            new ExampleColumnHandle("text", createUnboundedVarcharType(), 0)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, createUnboundedVarcharType()));

    recordSet = new ExampleRecordSet(new ExampleSplit(dataUri), ImmutableList.of(
            new ExampleColumnHandle("value", BIGINT, 1),
            new ExampleColumnHandle("value", BIGINT, 1),
            new ExampleColumnHandle("text", createUnboundedVarcharType(), 0)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, BIGINT, createUnboundedVarcharType()));

    recordSet = new ExampleRecordSet(new ExampleSplit(dataUri), ImmutableList.of());
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of());
}
 
Example 4
Source Project: presto   Source File: TestExampleRecordSet.java    License: Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example 5
Source Project: presto   Source File: TestExampleRecordSet.java    License: Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example 6
Source Project: presto   Source File: JmxRecordSetProvider.java    License: Apache License 2.0 6 votes vote down vote up
@Override
public RecordSet getRecordSet(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<? extends ColumnHandle> columns)
{
    JmxTableHandle tableHandle = (JmxTableHandle) table;

    requireNonNull(columns, "columns is null");
    checkArgument(!columns.isEmpty(), "must provide at least one column");

    List<List<Object>> rows;
    try {
        if (tableHandle.isLiveData()) {
            rows = getLiveRows(tableHandle, columns);
        }
        else {
            List<Integer> selectedColumns = calculateSelectedColumns(tableHandle.getColumnHandles(), getColumnNames(columns));
            rows = tableHandle.getObjectNames().stream()
                    .flatMap(objectName -> jmxHistoricalData.getRows(objectName, selectedColumns).stream())
                    .collect(toImmutableList());
        }
    }
    catch (JMException e) {
        rows = ImmutableList.of();
    }

    return new InMemoryRecordSet(getColumnTypes(columns), rows);
}
 
Example 7
Source Project: presto   Source File: TestJmxSplitManager.java    License: Apache License 2.0 6 votes vote down vote up
@Test
public void testHistoryRecordSetProvider()
        throws Exception
{
    for (SchemaTableName schemaTableName : metadata.listTables(SESSION, Optional.of(HISTORY_SCHEMA_NAME))) {
        // wait for at least two samples
        List<Long> timeStamps = ImmutableList.of();
        for (int waited = 0; waited < TIMEOUT_TIME; waited += SLEEP_TIME) {
            RecordSet recordSet = getRecordSet(schemaTableName);
            timeStamps = readTimeStampsFrom(recordSet);
            if (timeStamps.size() >= 2) {
                break;
            }
            Thread.sleep(SLEEP_TIME);
        }
        assertTrue(timeStamps.size() >= 2);

        // we don't have equality check here because JmxHistoryDumper scheduling can lag
        assertTrue(timeStamps.get(1) - timeStamps.get(0) >= JMX_STATS_DUMP.toMillis());
    }
}
 
Example 8
Source Project: presto   Source File: TestJmxSplitManager.java    License: Apache License 2.0 6 votes vote down vote up
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();
}
 
Example 9
Source Project: presto   Source File: JdbcRecordSetProvider.java    License: Apache License 2.0 6 votes vote down vote up
@Override
public RecordSet getRecordSet(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<? extends ColumnHandle> columns)
{
    JdbcSplit jdbcSplit = (JdbcSplit) split;
    JdbcTableHandle jdbcTable = (JdbcTableHandle) table;

    // In the current API, the columns (and order) needed by the engine are provided via an argument to this method. Make sure that
    // any columns that were recorded in the table handle match the requested set.
    // If no columns are recorded, it means that applyProjection never got called (e.g., in the case all columns are being used) and all
    // table columns should be returned. TODO: this is something that should be addressed once the getRecordSet API is revamped
    jdbcTable.getColumns()
            .ifPresent(tableColumns -> verify(columns.equals(tableColumns)));

    ImmutableList.Builder<JdbcColumnHandle> handles = ImmutableList.builder();
    for (ColumnHandle handle : columns) {
        handles.add((JdbcColumnHandle) handle);
    }

    return new JdbcRecordSet(jdbcClient, session, jdbcSplit, jdbcTable, handles.build());
}
 
Example 10
Source Project: presto   Source File: TestJdbcRecordSet.java    License: Apache License 2.0 6 votes vote down vote up
@Test
public void testGetColumnTypes()
{
    RecordSet recordSet = createRecordSet(ImmutableList.of(
            new JdbcColumnHandle("text", JDBC_VARCHAR, VARCHAR),
            new JdbcColumnHandle("text_short", JDBC_VARCHAR, createVarcharType(32)),
            new JdbcColumnHandle("value", JDBC_BIGINT, BIGINT)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(VARCHAR, createVarcharType(32), BIGINT));

    recordSet = createRecordSet(ImmutableList.of(
            new JdbcColumnHandle("value", JDBC_BIGINT, BIGINT),
            new JdbcColumnHandle("text", JDBC_VARCHAR, VARCHAR)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, VARCHAR));

    recordSet = createRecordSet(ImmutableList.of(
            new JdbcColumnHandle("value", JDBC_BIGINT, BIGINT),
            new JdbcColumnHandle("value", JDBC_BIGINT, BIGINT),
            new JdbcColumnHandle("text", JDBC_VARCHAR, VARCHAR)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, BIGINT, VARCHAR));

    recordSet = createRecordSet(ImmutableList.of());
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of());
}
 
Example 11
Source Project: presto   Source File: TestJdbcRecordSetProvider.java    License: Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example 12
Source Project: presto   Source File: TestJdbcRecordSetProvider.java    License: Apache License 2.0 6 votes vote down vote up
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();
}
 
Example 13
Source Project: presto   Source File: PrestoThriftBlock.java    License: Apache License 2.0 6 votes vote down vote up
public static PrestoThriftBlock fromRecordSetColumn(RecordSet recordSet, int columnIndex, int totalRecords)
{
    Type type = recordSet.getColumnTypes().get(columnIndex);
    // use more efficient implementations for numeric types which are likely to be used in index join
    if (type instanceof IntegerType) {
        return PrestoThriftInteger.fromRecordSetColumn(recordSet, columnIndex, totalRecords);
    }
    if (type instanceof BigintType) {
        return PrestoThriftBigint.fromRecordSetColumn(recordSet, columnIndex, totalRecords);
    }
    if (type instanceof DateType) {
        return PrestoThriftDate.fromRecordSetColumn(recordSet, columnIndex, totalRecords);
    }
    if (type instanceof TimestampType) {
        return PrestoThriftTimestamp.fromRecordSetColumn(recordSet, columnIndex, totalRecords);
    }
    // less efficient implementation which converts to a block first
    return fromBlock(convertColumnToBlock(recordSet, columnIndex, totalRecords), type);
}
 
Example 14
Source Project: presto   Source File: PrestoThriftTypeUtils.java    License: Apache License 2.0 6 votes vote down vote up
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);
}
 
Example 15
Source Project: presto   Source File: PrestoThriftTypeUtils.java    License: Apache License 2.0 6 votes vote down vote up
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);
}
 
Example 16
Source Project: presto   Source File: TpchIndexedData.java    License: Apache License 2.0 6 votes vote down vote up
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 17
Source Project: presto   Source File: TpchIndexedData.java    License: Apache License 2.0 6 votes vote down vote up
private static Iterable<MaterializedTuple> tupleIterable(final RecordSet recordSet)
{
    return () -> new AbstractIterator<MaterializedTuple>()
    {
        private final RecordCursor cursor = recordSet.cursor();

        @Override
        protected MaterializedTuple computeNext()
        {
            if (!cursor.advanceNextPosition()) {
                return endOfData();
            }
            return new MaterializedTuple(extractValues(cursor, recordSet.getColumnTypes()));
        }
    };
}
 
Example 18
Source Project: presto   Source File: SystemPageSourceProvider.java    License: Apache License 2.0 6 votes vote down vote up
private static RecordSet toRecordSet(ConnectorTransactionHandle sourceTransaction, SystemTable table, ConnectorSession session, TupleDomain<Integer> constraint)
{
    return new RecordSet()
    {
        private final List<Type> types = table.getTableMetadata().getColumns().stream()
                .map(ColumnMetadata::getType)
                .collect(toImmutableList());

        @Override
        public List<Type> getColumnTypes()
        {
            return types;
        }

        @Override
        public RecordCursor cursor()
        {
            return table.cursor(sourceTransaction, session, constraint);
        }
    };
}
 
Example 19
Source Project: presto   Source File: FieldSetFilteringRecordSet.java    License: Apache License 2.0 6 votes vote down vote up
public FieldSetFilteringRecordSet(Metadata metadata, RecordSet delegate, List<Set<Integer>> fieldSets)
{
    requireNonNull(metadata, "metadata is null");
    this.delegate = requireNonNull(delegate, "delegate is null");

    ImmutableList.Builder<Set<Field>> fieldSetsBuilder = ImmutableList.builder();
    List<Type> columnTypes = delegate.getColumnTypes();
    for (Set<Integer> fieldSet : requireNonNull(fieldSets, "fieldSets is null")) {
        ImmutableSet.Builder<Field> fieldSetBuilder = ImmutableSet.builder();
        for (int field : fieldSet) {
            ResolvedFunction resolvedFunction = metadata.resolveOperator(EQUAL, ImmutableList.of(columnTypes.get(field), columnTypes.get(field)));
            MethodHandle methodHandle = metadata.getScalarFunctionInvoker(resolvedFunction, Optional.empty()).getMethodHandle();
            fieldSetBuilder.add(new Field(field, methodHandle));
        }
        fieldSetsBuilder.add(fieldSetBuilder.build());
    }
    this.fieldSets = fieldSetsBuilder.build();
}
 
Example 20
Source Project: presto   Source File: IndexSourceOperator.java    License: Apache License 2.0 6 votes vote down vote up
@Override
public Supplier<Optional<UpdatablePageSource>> addSplit(Split split)
{
    requireNonNull(split, "split is null");
    checkState(source == null, "Index source split already set");

    IndexSplit indexSplit = (IndexSplit) split.getConnectorSplit();

    // Normalize the incoming RecordSet to something that can be consumed by the index
    RecordSet normalizedRecordSet = probeKeyNormalizer.apply(indexSplit.getKeyRecordSet());
    ConnectorPageSource result = index.lookup(normalizedRecordSet);
    source = new PageSourceOperator(result, operatorContext);

    Object splitInfo = split.getInfo();
    if (splitInfo != null) {
        operatorContext.setInfoSupplier(() -> new SplitOperatorInfo(splitInfo));
    }

    return Optional::empty;
}
 
Example 21
Source Project: presto   Source File: TestClassLoaderSafeWrappers.java    License: Apache License 2.0 6 votes vote down vote up
@Test
public void testAllMethodsOverridden()
{
    assertAllMethodsOverridden(ConnectorAccessControl.class, ClassLoaderSafeConnectorAccessControl.class);
    assertAllMethodsOverridden(ConnectorMetadata.class, ClassLoaderSafeConnectorMetadata.class);
    assertAllMethodsOverridden(ConnectorPageSink.class, ClassLoaderSafeConnectorPageSink.class);
    assertAllMethodsOverridden(ConnectorPageSinkProvider.class, ClassLoaderSafeConnectorPageSinkProvider.class);
    assertAllMethodsOverridden(ConnectorPageSourceProvider.class, ClassLoaderSafeConnectorPageSourceProvider.class);
    assertAllMethodsOverridden(ConnectorSplitManager.class, ClassLoaderSafeConnectorSplitManager.class);
    assertAllMethodsOverridden(ConnectorNodePartitioningProvider.class, ClassLoaderSafeNodePartitioningProvider.class);
    assertAllMethodsOverridden(ConnectorSplitSource.class, ClassLoaderSafeConnectorSplitSource.class);
    assertAllMethodsOverridden(SystemTable.class, ClassLoaderSafeSystemTable.class);
    assertAllMethodsOverridden(ConnectorRecordSetProvider.class, ClassLoaderSafeConnectorRecordSetProvider.class);
    assertAllMethodsOverridden(RecordSet.class, ClassLoaderSafeRecordSet.class);
    assertAllMethodsOverridden(EventListener.class, ClassLoaderSafeEventListener.class);
}
 
Example 22
Source Project: presto   Source File: TpchRecordSetProvider.java    License: Apache License 2.0 6 votes vote down vote up
public <E extends TpchEntity> RecordSet getRecordSet(
        TpchTable<E> table,
        List<? extends ColumnHandle> columns,
        double scaleFactor,
        int partNumber,
        int totalParts,
        TupleDomain<ColumnHandle> predicate)
{
    ImmutableList.Builder<TpchColumn<E>> builder = ImmutableList.builder();
    for (ColumnHandle column : columns) {
        String columnName = ((TpchColumnHandle) column).getColumnName();
        if (columnName.equalsIgnoreCase(TpchMetadata.ROW_NUMBER_COLUMN_NAME)) {
            builder.add(new RowNumberTpchColumn<E>());
        }
        else {
            builder.add(table.getColumn(columnName));
        }
    }

    return createTpchRecordSet(table, builder.build(), scaleFactor, partNumber + 1, totalParts, predicate);
}
 
Example 23
Source Project: presto   Source File: CassandraRecordSetProvider.java    License: Apache License 2.0 6 votes vote down vote up
@Override
public RecordSet getRecordSet(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<? extends ColumnHandle> columns)
{
    CassandraSplit cassandraSplit = (CassandraSplit) split;
    CassandraTableHandle cassandraTable = (CassandraTableHandle) table;

    List<CassandraColumnHandle> cassandraColumns = columns.stream()
            .map(column -> (CassandraColumnHandle) column)
            .collect(toList());

    String selectCql = CassandraCqlUtils.selectFrom(cassandraTable, cassandraColumns).getQueryString();
    StringBuilder sb = new StringBuilder(selectCql);
    if (sb.charAt(sb.length() - 1) == ';') {
        sb.setLength(sb.length() - 1);
    }
    sb.append(cassandraSplit.getWhereClause());
    String cql = sb.toString();
    log.debug("Creating record set: %s", cql);

    return new CassandraRecordSet(cassandraSession, cql, cassandraColumns);
}
 
Example 24
Source Project: presto   Source File: KafkaRecordSetProvider.java    License: Apache License 2.0 5 votes vote down vote up
@Override
public RecordSet getRecordSet(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<? extends ColumnHandle> columns)
{
    KafkaSplit kafkaSplit = convertSplit(split);

    List<KafkaColumnHandle> kafkaColumns = columns.stream()
            .map(KafkaHandleResolver::convertColumnHandle)
            .collect(ImmutableList.toImmutableList());

    RowDecoder keyDecoder = decoderFactory.create(
            kafkaSplit.getKeyDataFormat(),
            getDecoderParameters(kafkaSplit.getKeyDataSchemaContents()),
            kafkaColumns.stream()
                    .filter(col -> !col.isInternal())
                    .filter(KafkaColumnHandle::isKeyDecoder)
                    .collect(toImmutableSet()));

    RowDecoder messageDecoder = decoderFactory.create(
            kafkaSplit.getMessageDataFormat(),
            getDecoderParameters(kafkaSplit.getMessageDataSchemaContents()),
            kafkaColumns.stream()
                    .filter(col -> !col.isInternal())
                    .filter(col -> !col.isKeyDecoder())
                    .collect(toImmutableSet()));

    return new KafkaRecordSet(kafkaSplit, consumerFactory, kafkaColumns, keyDecoder, messageDecoder);
}
 
Example 25
Source Project: presto   Source File: PrometheusRecordSetProvider.java    License: Apache License 2.0 5 votes vote down vote up
@Override
public RecordSet getRecordSet(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<? extends ColumnHandle> columns)
{
    PrometheusSplit prometheusSplit = (PrometheusSplit) split;

    ImmutableList.Builder<PrometheusColumnHandle> handles = ImmutableList.builder();
    for (ColumnHandle handle : columns) {
        handles.add((PrometheusColumnHandle) handle);
    }

    return new PrometheusRecordSet(prometheusSplit, handles.build());
}
 
Example 26
Source Project: presto   Source File: TestPrometheusRecordSetProvider.java    License: Apache License 2.0 5 votes vote down vote up
@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);
}
 
Example 27
Source Project: presto   Source File: ExampleRecordSetProvider.java    License: Apache License 2.0 5 votes vote down vote up
@Override
public RecordSet getRecordSet(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<? extends ColumnHandle> columns)
{
    ExampleSplit exampleSplit = (ExampleSplit) split;

    ImmutableList.Builder<ExampleColumnHandle> handles = ImmutableList.builder();
    for (ColumnHandle handle : columns) {
        handles.add((ExampleColumnHandle) handle);
    }

    return new ExampleRecordSet(exampleSplit, handles.build());
}
 
Example 28
Source Project: presto   Source File: RedisRecordSetProvider.java    License: Apache License 2.0 5 votes vote down vote up
@Override
public RecordSet getRecordSet(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<? extends ColumnHandle> columns)
{
    RedisSplit redisSplit = convertSplit(split);

    List<RedisColumnHandle> redisColumns = columns.stream()
            .map(RedisHandleResolver::convertColumnHandle)
            .collect(ImmutableList.toImmutableList());

    RowDecoder keyDecoder = decoderFactory.create(
            redisSplit.getKeyDataFormat(),
            emptyMap(),
            redisColumns.stream()
                    .filter(col -> !col.isInternal())
                    .filter(RedisColumnHandle::isKeyDecoder)
                    .collect(toImmutableSet()));

    RowDecoder valueDecoder = decoderFactory.create(
            redisSplit.getValueDataFormat(),
            emptyMap(),
            redisColumns.stream()
                    .filter(col -> !col.isInternal())
                    .filter(col -> !col.isKeyDecoder())
                    .collect(toImmutableSet()));

    return new RedisRecordSet(redisSplit, jedisManager, redisColumns, keyDecoder, valueDecoder);
}
 
Example 29
Source Project: presto   Source File: TestJmxSplitManager.java    License: Apache License 2.0 5 votes vote down vote up
@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);
                }
            }
        }
    }
}
 
Example 30
Source Project: presto   Source File: TestJmxSplitManager.java    License: Apache License 2.0 5 votes vote down vote up
private RecordSet getRecordSet(SchemaTableName schemaTableName)
        throws Exception
{
    JmxTableHandle tableHandle = metadata.getTableHandle(SESSION, schemaTableName);
    List<ColumnHandle> columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(SESSION, tableHandle).values());

    ConnectorSplitSource splitSource = splitManager.getSplits(JmxTransactionHandle.INSTANCE, SESSION, tableHandle, UNGROUPED_SCHEDULING);
    List<ConnectorSplit> allSplits = getAllSplits(splitSource);
    assertEquals(allSplits.size(), nodes.size());
    ConnectorSplit split = allSplits.get(0);

    return recordSetProvider.getRecordSet(JmxTransactionHandle.INSTANCE, SESSION, split, tableHandle, columnHandles);
}