io.prestosql.spi.connector.RecordSet Java Examples

The following examples show how to use io.prestosql.spi.connector.RecordSet. 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 File: PrestoThriftTypeUtils.java    From presto with 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 #2
Source File: PrestoThriftBlock.java    From presto with 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 #3
Source File: TestJdbcRecordSetProvider.java    From presto with 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 #4
Source File: TestJdbcRecordSetProvider.java    From presto with 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 #5
Source File: FieldSetFilteringRecordSet.java    From presto with 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 #6
Source File: TestClassLoaderSafeWrappers.java    From presto with 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 #7
Source File: SystemPageSourceProvider.java    From presto with 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 #8
Source File: TestJdbcRecordSet.java    From presto with 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 #9
Source File: JdbcRecordSetProvider.java    From presto with 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 File: TpchRecordSetProvider.java    From presto with 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 #11
Source File: TestJmxSplitManager.java    From presto with 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 #12
Source File: TestJmxSplitManager.java    From presto with 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 #13
Source File: PrestoThriftTypeUtils.java    From presto with 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 #14
Source File: JmxRecordSetProvider.java    From presto with 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 #15
Source File: CassandraRecordSetProvider.java    From presto with 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 #16
Source File: TestExampleRecordSet.java    From presto with 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 #17
Source File: TestExampleRecordSet.java    From presto with 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 #18
Source File: TestExampleRecordSet.java    From presto with 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 #19
Source File: TestExampleRecordSetProvider.java    From presto with 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 #20
Source File: TpchIndexedData.java    From presto with 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 #21
Source File: TestPrometheusIntegrationTests1.java    From presto with 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 #22
Source File: IndexSourceOperator.java    From presto with 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 #23
Source File: TpchIndexedData.java    From presto with 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 #24
Source File: ThriftIndexedTpchService.java    From presto with Apache License 2.0 5 votes vote down vote up
/**
 * Get lookup result and re-map output columns based on requested order.
 */
private static RecordSet lookupIndexKeys(RecordSet keys, IndexedTable table, List<String> outputColumnNames, List<String> lookupColumnNames)
{
    RecordSet allColumnsOutputRecordSet = table.lookupKeys(
            new MappedRecordSet(
                    keys,
                    computeRemap(lookupColumnNames, table.getKeyColumns())));
    List<Integer> outputRemap = computeRemap(table.getOutputColumns(), outputColumnNames);
    return new MappedRecordSet(allColumnsOutputRecordSet, outputRemap);
}
 
Example #25
Source File: TpchIndexedData.java    From presto with Apache License 2.0 5 votes vote down vote up
public TpchIndexedData(TpchIndexSpec tpchIndexSpec)
{
    requireNonNull(tpchIndexSpec, "tpchIndexSpec is null");

    TpchMetadata tpchMetadata = new TpchMetadata();
    TpchRecordSetProvider tpchRecordSetProvider = new TpchRecordSetProvider();

    ImmutableMap.Builder<Set<TpchScaledColumn>, IndexedTable> indexedTablesBuilder = ImmutableMap.builder();

    Set<TpchScaledTable> tables = tpchIndexSpec.listIndexedTables();
    for (TpchScaledTable table : tables) {
        SchemaTableName tableName = new SchemaTableName("sf" + table.getScaleFactor(), table.getTableName());
        TpchTableHandle tableHandle = tpchMetadata.getTableHandle(null, tableName);
        Map<String, ColumnHandle> columnHandles = new LinkedHashMap<>(tpchMetadata.getColumnHandles(null, tableHandle));
        for (Set<String> columnNames : tpchIndexSpec.getColumnIndexes(table)) {
            List<String> keyColumnNames = ImmutableList.copyOf(columnNames); // Finalize the key order
            Set<TpchScaledColumn> keyColumns = keyColumnNames.stream()
                    .map(name -> new TpchScaledColumn(table, name))
                    .collect(toImmutableSet());

            TpchTable<?> tpchTable = TpchTable.getTable(table.getTableName());
            RecordSet recordSet = tpchRecordSetProvider.getRecordSet(tpchTable, ImmutableList.copyOf(columnHandles.values()), table.getScaleFactor(), 0, 1, TupleDomain.all());
            IndexedTable indexedTable = indexTable(recordSet, ImmutableList.copyOf(columnHandles.keySet()), keyColumnNames);
            indexedTablesBuilder.put(keyColumns, indexedTable);
        }
    }

    indexedTables = indexedTablesBuilder.build();
}
 
Example #26
Source File: IndexSourceOperator.java    From presto with Apache License 2.0 5 votes vote down vote up
public IndexSourceOperatorFactory(
        int operatorId,
        PlanNodeId sourceId,
        ConnectorIndex index,
        Function<RecordSet, RecordSet> probeKeyNormalizer)
{
    this.operatorId = operatorId;
    this.sourceId = requireNonNull(sourceId, "sourceId is null");
    this.index = requireNonNull(index, "index is null");
    this.probeKeyNormalizer = requireNonNull(probeKeyNormalizer, "probeKeyNormalizer is null");
}
 
Example #27
Source File: TpchIndexedData.java    From presto with Apache License 2.0 5 votes vote down vote up
public RecordSet lookupKeys(RecordSet recordSet)
{
    checkArgument(recordSet.getColumnTypes().equals(keyTypes), "Input RecordSet keys do not match expected key type");

    Iterable<RecordSet> outputRecordSets = Iterables.transform(tupleIterable(recordSet), key -> {
        for (Object value : key.getValues()) {
            if (value == null) {
                throw new IllegalArgumentException("TPCH index does not support null values");
            }
        }
        return lookupKey(key);
    });

    return new ConcatRecordSet(outputRecordSets, outputTypes);
}
 
Example #28
Source File: ConcatRecordSet.java    From presto with Apache License 2.0 5 votes vote down vote up
@Override
public RecordCursor cursor()
{
    // NOTE: the ConcatRecordCursor implementation relies on the fact that the
    // cursor creation in the Iterable is lazy so DO NOT materialize this into
    // an ImmutableList
    Iterable<RecordCursor> recordCursors = transform(recordSets, RecordSet::cursor);
    return new ConcatRecordCursor(recordCursors.iterator(), types);
}
 
Example #29
Source File: ThriftIndexedTpchService.java    From presto with Apache License 2.0 5 votes vote down vote up
@Override
protected ConnectorPageSource createLookupPageSource(SplitInfo splitInfo, List<String> outputColumnNames)
{
    IndexedTable indexedTable = indexedData.getIndexedTable(
            splitInfo.getTableName(),
            schemaNameToScaleFactor(splitInfo.getSchemaName()),
            ImmutableSet.copyOf(splitInfo.getLookupColumnNames()))
            .orElseThrow(() -> new IllegalArgumentException(format("No such index: %s%s", splitInfo.getTableName(), splitInfo.getLookupColumnNames())));
    List<Type> lookupColumnTypes = types(splitInfo.getTableName(), splitInfo.getLookupColumnNames());
    RecordSet keyRecordSet = new ListBasedRecordSet(splitInfo.getKeys(), lookupColumnTypes);
    RecordSet outputRecordSet = lookupIndexKeys(keyRecordSet, indexedTable, outputColumnNames, splitInfo.getLookupColumnNames());
    return new RecordPageSource(outputRecordSet);
}
 
Example #30
Source File: IndexSourceOperator.java    From presto with Apache License 2.0 5 votes vote down vote up
public IndexSourceOperator(
        OperatorContext operatorContext,
        PlanNodeId planNodeId,
        ConnectorIndex index,
        Function<RecordSet, RecordSet> probeKeyNormalizer)
{
    this.operatorContext = requireNonNull(operatorContext, "operatorContext is null");
    this.planNodeId = requireNonNull(planNodeId, "planNodeId is null");
    this.index = requireNonNull(index, "index is null");
    this.probeKeyNormalizer = requireNonNull(probeKeyNormalizer, "probeKeyNormalizer is null");
}