io.prestosql.spi.connector.RecordSet Java Examples

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

/**
 * 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);
}
 

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

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

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

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

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

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