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

Example 16 with RecordSet

use of io.trino.spi.connector.RecordSet in project trino by trinodb.

the class RedisRecordSetProvider method getRecordSet.

@Override
public RecordSet getRecordSet(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<? extends ColumnHandle> columns) {
    RedisSplit redisSplit = (RedisSplit) split;
    List<RedisColumnHandle> redisColumns = columns.stream().map(RedisColumnHandle.class::cast).collect(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 17 with RecordSet

use of io.trino.spi.connector.RecordSet in project trino by trinodb.

the class ThriftIndexedTpchService method createLookupPageSource.

@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 18 with RecordSet

use of io.trino.spi.connector.RecordSet in project trino by trinodb.

the class ThriftIndexedTpchService method lookupIndexKeys.

/**
 * 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 19 with RecordSet

use of io.trino.spi.connector.RecordSet in project trino by trinodb.

the class TpchIndexProvider method getIndex.

@Override
public ConnectorIndex getIndex(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorIndexHandle indexHandle, List<ColumnHandle> lookupSchema, List<ColumnHandle> outputSchema) {
    TpchIndexHandle tpchIndexHandle = (TpchIndexHandle) indexHandle;
    Map<ColumnHandle, NullableValue> fixedValues = TupleDomain.extractFixedValues(tpchIndexHandle.getFixedValues()).get();
    checkArgument(lookupSchema.stream().noneMatch(handle -> fixedValues.keySet().contains(handle)), "Lookup columnHandles are not expected to overlap with the fixed value predicates");
    // Establish an order for the fixedValues
    List<ColumnHandle> fixedValueColumns = ImmutableList.copyOf(fixedValues.keySet());
    // Extract the fixedValues as their raw values and types
    List<Object> rawFixedValues = new ArrayList<>(fixedValueColumns.size());
    List<Type> rawFixedTypes = new ArrayList<>(fixedValueColumns.size());
    for (ColumnHandle fixedValueColumn : fixedValueColumns) {
        rawFixedValues.add(fixedValues.get(fixedValueColumn).getValue());
        rawFixedTypes.add(((TpchColumnHandle) fixedValueColumn).getType());
    }
    // Establish the schema after we append the fixed values to the lookup keys.
    List<ColumnHandle> finalLookupSchema = ImmutableList.<ColumnHandle>builder().addAll(lookupSchema).addAll(fixedValueColumns).build();
    Optional<TpchIndexedData.IndexedTable> indexedTable = indexedData.getIndexedTable(tpchIndexHandle.getTableName(), tpchIndexHandle.getScaleFactor(), tpchIndexHandle.getIndexColumnNames());
    checkState(indexedTable.isPresent());
    TpchIndexedData.IndexedTable table = indexedTable.get();
    // Compute how to map from the final lookup schema to the table index key order
    List<Integer> keyRemap = computeRemap(handleToNames(finalLookupSchema), table.getKeyColumns());
    Function<RecordSet, RecordSet> keyFormatter = key -> new MappedRecordSet(new AppendingRecordSet(key, rawFixedValues, rawFixedTypes), keyRemap);
    // Compute how to map from the output of the indexed data to the expected output schema
    List<Integer> outputRemap = computeRemap(table.getOutputColumns(), handleToNames(outputSchema));
    Function<RecordSet, RecordSet> outputFormatter = output -> new MappedRecordSet(output, outputRemap);
    return new TpchConnectorIndex(keyFormatter, outputFormatter, table);
}

Example 20 with RecordSet

use of io.trino.spi.connector.RecordSet in project trino by trinodb.

the class TpchIndexedData method indexTable.

private static IndexedTable indexTable(RecordSet recordSet, 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());
}