io.prestosql.spi.connector.ColumnHandle Java Examples

public HiveTableHandle applyPartitionResult(HiveTableHandle handle, HivePartitionResult partitions, Optional<Set<ColumnHandle>> columns)
{
    return new HiveTableHandle(
            handle.getSchemaName(),
            handle.getTableName(),
            handle.getTableParameters(),
            ImmutableList.copyOf(partitions.getPartitionColumns()),
            Optional.of(getPartitionsAsList(partitions)),
            partitions.getCompactEffectivePredicate(),
            partitions.getEnforcedConstraint(),
            partitions.getBucketHandle(),
            partitions.getBucketFilter(),
            handle.getAnalyzePartitionValues(),
            handle.getAnalyzeColumnNames(),
            Optionals.combine(handle.getConstraintColumns(), columns,
                    (oldColumns, newColumns) -> Sets.union(oldColumns, newColumns)));
}
 

@Override
public Map<String, ColumnHandle> getColumnHandles(ConnectorSession session, ConnectorTableHandle tableHandle)
{
    RaptorTableHandle raptorTableHandle = (RaptorTableHandle) tableHandle;
    ImmutableMap.Builder<String, ColumnHandle> builder = ImmutableMap.builder();
    for (TableColumn tableColumn : dao.listTableColumns(raptorTableHandle.getTableId())) {
        builder.put(tableColumn.getColumnName(), getRaptorColumnHandle(tableColumn));
    }

    RaptorColumnHandle uuidColumn = shardUuidColumnHandle();
    builder.put(uuidColumn.getColumnName(), uuidColumn);

    if (raptorTableHandle.isBucketed()) {
        RaptorColumnHandle bucketNumberColumn = bucketNumberColumnHandle();
        builder.put(bucketNumberColumn.getColumnName(), bucketNumberColumn);
    }

    return builder.build();
}
 

@BeforeMethod
public void setUp()
{
    scanAssignments = ImmutableMap.<Symbol, ColumnHandle>builder()
            .put(A, new TestingColumnHandle("a"))
            .put(B, new TestingColumnHandle("b"))
            .put(C, new TestingColumnHandle("c"))
            .put(D, new TestingColumnHandle("d"))
            .put(E, new TestingColumnHandle("e"))
            .put(F, new TestingColumnHandle("f"))
            .build();

    Map<Symbol, ColumnHandle> assignments = Maps.filterKeys(scanAssignments, Predicates.in(ImmutableList.of(A, B, C, D, E, F)));
    baseTableScan = TableScanNode.newInstance(
            newId(),
            makeTableHandle(TupleDomain.all()),
            ImmutableList.copyOf(assignments.keySet()),
            assignments);

    expressionNormalizer = new ExpressionIdentityNormalizer();
}
 

@Override
public Optional<ConstraintApplicationResult<ConnectorTableHandle>> applyFilter(
        ConnectorSession session,
        ConnectorTableHandle handle,
        Constraint constraint)
{
    log.debug("applyFilter(session=%s, handle=%s, summary=%s, predicate=%s, columns=%s)",
            session, handle, constraint.getSummary(), constraint.predicate(), constraint.getPredicateColumns());
    BigQueryTableHandle bigQueryTableHandle = (BigQueryTableHandle) handle;

    TupleDomain<ColumnHandle> oldDomain = bigQueryTableHandle.getConstraint();
    TupleDomain<ColumnHandle> newDomain = oldDomain.intersect(constraint.getSummary());
    if (oldDomain.equals(newDomain)) {
        return Optional.empty();
    }

    BigQueryTableHandle updatedHandle = bigQueryTableHandle.withConstraint(newDomain);

    return Optional.of(new ConstraintApplicationResult<>(updatedHandle, constraint.getSummary()));
}
 

private void createScanFilterAndProjectOperatorFactories(List<Page> inputPages, PageProcessor pageProcessor, CursorProcessor cursorProcessor, List<ColumnHandle> columnHandles, List<Type> types)
{
    operatorFactory = new ScanFilterAndProjectOperatorFactory(
            0,
            new PlanNodeId("test"),
            new PlanNodeId("test_source"),
            (session, split, table, columns, dynamicFilter) -> new FixedPageSource(inputPages),
            () -> cursorProcessor,
            () -> pageProcessor,
            TEST_TABLE_HANDLE,
            columnHandles,
            TupleDomain::all,
            types,
            FILTER_AND_PROJECT_MIN_OUTPUT_PAGE_SIZE,
            FILTER_AND_PROJECT_MIN_OUTPUT_PAGE_ROW_COUNT);
}
 

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

private TableStatistics toTableStatistics(Map<String, ColumnHandle> columnHandles, TableStatisticsData statisticsData)
{
    long rowCount = statisticsData.getRowCount();
    TableStatistics.Builder tableStatistics = TableStatistics.builder()
            .setRowCount(Estimate.of(rowCount));

    if (rowCount > 0) {
        Map<String, ColumnStatisticsData> columnsData = statisticsData.getColumns();
        for (Map.Entry<String, ColumnHandle> entry : columnHandles.entrySet()) {
            TpcdsColumnHandle columnHandle = (TpcdsColumnHandle) entry.getValue();
            tableStatistics.setColumnStatistics(entry.getValue(), toColumnStatistics(columnsData.get(entry.getKey()), columnHandle.getType(), rowCount));
        }
    }

    return tableStatistics.build();
}
 

private Map<String, ColumnHandle> getColumnHandles(LocalFileTableHandle tableHandle)
{
    ImmutableMap.Builder<String, ColumnHandle> columnHandles = ImmutableMap.builder();
    int index = 0;
    for (ColumnMetadata column : localFileTables.getColumns(tableHandle)) {
        int ordinalPosition;
        if (column.getName().equals(SERVER_ADDRESS_COLUMN_NAME)) {
            ordinalPosition = SERVER_ADDRESS_ORDINAL_POSITION;
        }
        else {
            ordinalPosition = index;
            index++;
        }
        columnHandles.put(column.getName(), new LocalFileColumnHandle(column.getName(), column.getType(), ordinalPosition));
    }
    return columnHandles.build();
}
 

@Override
public Optional<ConstraintApplicationResult<ConnectorTableHandle>> applyFilter(ConnectorSession session, ConnectorTableHandle table, Constraint constraint)
{
    KuduTableHandle handle = (KuduTableHandle) table;

    TupleDomain<ColumnHandle> oldDomain = handle.getConstraint();
    TupleDomain<ColumnHandle> newDomain = oldDomain.intersect(constraint.getSummary());
    if (oldDomain.equals(newDomain)) {
        return Optional.empty();
    }

    handle = new KuduTableHandle(
            handle.getSchemaTableName(),
            handle.getTable(clientSession),
            newDomain,
            handle.getDesiredColumns(),
            handle.isDeleteHandle());

    return Optional.of(new ConstraintApplicationResult<>(handle, constraint.getSummary()));
}
 

private static List<String> resolvePinotColumns(SchemaTableName schemaTableName, List<String> prestoColumnNames, Map<String, ColumnHandle> columnHandles)
{
    ImmutableList.Builder<String> pinotColumnNamesBuilder = ImmutableList.builder();
    for (String prestoColumnName : prestoColumnNames) {
        if (prestoColumnName.equals(WILDCARD)) {
            pinotColumnNamesBuilder.addAll(columnHandles.values().stream().map(handle -> ((PinotColumnHandle) handle).getColumnName()).collect(toImmutableList()));
        }
        else {
            PinotColumnHandle columnHandle = (PinotColumnHandle) columnHandles.get(prestoColumnName);
            if (columnHandle == null) {
                throw new ColumnNotFoundException(schemaTableName, prestoColumnName);
            }
            pinotColumnNamesBuilder.add(columnHandle.getColumnName());
        }
    }
    return pinotColumnNamesBuilder.build();
}
 

@Test
public void doesNotFireOnNonDeterministicPredicate()
{
    ColumnHandle columnHandle = new TpchColumnHandle("nationkey", BIGINT);
    tester().assertThat(removeRedundantTableScanPredicate)
            .on(p -> p.filter(
                    new ComparisonExpression(
                            EQUAL,
                            new FunctionCallBuilder(tester().getMetadata())
                                    .setName(QualifiedName.of("rand"))
                                    .build(),
                            new GenericLiteral("BIGINT", "42")),
                    p.tableScan(
                            nationTableHandle,
                            ImmutableList.of(p.symbol("nationkey", BIGINT)),
                            ImmutableMap.of(p.symbol("nationkey", BIGINT), columnHandle),
                            TupleDomain.all())))
            .doesNotFire();
}
 

private Map<PlanNodeId, SplitSource> visitScanAndFilter(TableScanNode node, Optional<FilterNode> filter)
{
    List<DynamicFilters.Descriptor> dynamicFilters = filter
            .map(FilterNode::getPredicate)
            .map(DynamicFilters::extractDynamicFilters)
            .map(DynamicFilters.ExtractResult::getDynamicConjuncts)
            .orElse(ImmutableList.of());

    Supplier<TupleDomain<ColumnHandle>> dynamicFilterSupplier = TupleDomain::all;
    if (!dynamicFilters.isEmpty()) {
        log.debug("Dynamic filters: %s", dynamicFilters);
        dynamicFilterSupplier = dynamicFilterService.createDynamicFilterSupplier(session.getQueryId(), dynamicFilters, node.getAssignments());
    }

    // get dataSource for table
    SplitSource splitSource = splitManager.getSplits(
            session,
            node.getTable(),
            stageExecutionDescriptor.isScanGroupedExecution(node.getId()) ? GROUPED_SCHEDULING : UNGROUPED_SCHEDULING,
            dynamicFilterSupplier);

    splitSources.add(splitSource);

    return ImmutableMap.of(node.getId(), splitSource);
}
 

@Test
public void testPredicatePushdown()
        throws Exception
{
    for (Node node : nodes) {
        String nodeIdentifier = node.getNodeIdentifier();
        TupleDomain<ColumnHandle> nodeTupleDomain = TupleDomain.fromFixedValues(ImmutableMap.of(columnHandle, NullableValue.of(createUnboundedVarcharType(), utf8Slice(nodeIdentifier))));
        JmxTableHandle tableHandle = new JmxTableHandle(new SchemaTableName("schema", "tableName"), ImmutableList.of("objectName"), ImmutableList.of(columnHandle), true, nodeTupleDomain);

        ConnectorSplitSource splitSource = splitManager.getSplits(JmxTransactionHandle.INSTANCE, SESSION, tableHandle, UNGROUPED_SCHEDULING);
        List<ConnectorSplit> allSplits = getAllSplits(splitSource);

        assertEquals(allSplits.size(), 1);
        assertEquals(allSplits.get(0).getAddresses().size(), 1);
        assertEquals(allSplits.get(0).getAddresses().get(0).getHostText(), nodeIdentifier);
    }
}
 

@Test
public void testRenameColumn()
{
    assertNull(metadata.getTableHandle(SESSION, DEFAULT_TEST_ORDERS));
    metadata.createTable(SESSION, getOrdersTable(), false);
    ConnectorTableHandle tableHandle = metadata.getTableHandle(SESSION, DEFAULT_TEST_ORDERS);
    assertInstanceOf(tableHandle, RaptorTableHandle.class);

    RaptorTableHandle raptorTableHandle = (RaptorTableHandle) tableHandle;
    ColumnHandle columnHandle = metadata.getColumnHandles(SESSION, tableHandle).get("orderkey");

    metadata.renameColumn(SESSION, raptorTableHandle, columnHandle, "orderkey_renamed");

    assertNull(metadata.getColumnHandles(SESSION, tableHandle).get("orderkey"));
    assertNotNull(metadata.getColumnHandles(SESSION, tableHandle).get("orderkey_renamed"));
}
 

@Override
public Optional<ConstraintApplicationResult<ConnectorTableHandle>> applyFilter(ConnectorSession session, ConnectorTableHandle table, Constraint constraint)
{
    PinotTableHandle handle = (PinotTableHandle) table;
    TupleDomain<ColumnHandle> oldDomain = handle.getConstraint();
    TupleDomain<ColumnHandle> newDomain = oldDomain.intersect(constraint.getSummary());
    if (oldDomain.equals(newDomain)) {
        return Optional.empty();
    }

    handle = new PinotTableHandle(
            handle.getSchemaName(),
            handle.getTableName(),
            newDomain,
            handle.getLimit(),
            handle.getQuery());
    return Optional.of(new ConstraintApplicationResult<>(handle, constraint.getSummary()));
}
 

@Test
public void consumesDeterministicPredicateIfNewDomainIsSame()
{
    ColumnHandle columnHandle = new TpchColumnHandle("nationkey", BIGINT);
    tester().assertThat(removeRedundantTableScanPredicate)
            .on(p -> p.filter(expression("nationkey = BIGINT '44'"),
                    p.tableScan(
                            nationTableHandle,
                            ImmutableList.of(p.symbol("nationkey", BIGINT)),
                            ImmutableMap.of(p.symbol("nationkey", BIGINT), columnHandle),
                            TupleDomain.fromFixedValues(ImmutableMap.of(
                                    columnHandle, NullableValue.of(BIGINT, (long) 44))))))
            .matches(constrainedTableScanWithTableLayout(
                    "nation",
                    ImmutableMap.of("nationkey", singleValue(BIGINT, (long) 44)),
                    ImmutableMap.of("nationkey", "nationkey")));
}
 

@Test
public void testCreateTable()
{
    assertNull(metadata.getTableHandle(SESSION, DEFAULT_TEST_ORDERS));

    metadata.createTable(SESSION, getOrdersTable(), false);

    ConnectorTableHandle tableHandle = metadata.getTableHandle(SESSION, DEFAULT_TEST_ORDERS);
    assertInstanceOf(tableHandle, RaptorTableHandle.class);
    assertEquals(((RaptorTableHandle) tableHandle).getTableId(), 1);

    ConnectorTableMetadata table = metadata.getTableMetadata(SESSION, tableHandle);
    assertTableEqual(table, getOrdersTable());

    ColumnHandle columnHandle = metadata.getColumnHandles(SESSION, tableHandle).get("orderkey");
    assertInstanceOf(columnHandle, RaptorColumnHandle.class);
    assertEquals(((RaptorColumnHandle) columnHandle).getColumnId(), 1);

    ColumnMetadata columnMetadata = metadata.getColumnMetadata(SESSION, tableHandle, columnHandle);
    assertNotNull(columnMetadata);
    assertEquals(columnMetadata.getName(), "orderkey");
    assertEquals(columnMetadata.getType(), BIGINT);
}
 

@JsonCreator
public KuduTableHandle(
        @JsonProperty("schemaTableName") SchemaTableName schemaTableName,
        @JsonProperty("constraint") TupleDomain<ColumnHandle> constraint,
        @JsonProperty("desiredColumns") Optional<List<ColumnHandle>> desiredColumns,
        @JsonProperty("isDeleteHandle") boolean isDeleteHandle)
{
    this(schemaTableName, null, constraint, desiredColumns, isDeleteHandle);
}
 

@Test
public void testPruneColumns()
{
    PruneTableScanColumns rule = new PruneTableScanColumns(tester().getMetadata());

    ThriftColumnHandle nationKeyColumn = new ThriftColumnHandle("nationKey", VARCHAR, "", false);
    ThriftColumnHandle nameColumn = new ThriftColumnHandle("name", VARCHAR, "", false);

    tester().assertThat(rule)
            .on(p -> {
                Symbol nationKey = p.symbol(nationKeyColumn.getColumnName(), VARCHAR);
                Symbol name = p.symbol(nameColumn.getColumnName(), VARCHAR);
                return p.project(
                        Assignments.of(
                                p.symbol("expr", VARCHAR),
                                nationKey.toSymbolReference()),
                        p.tableScan(
                                NATION_TABLE,
                                ImmutableList.of(nationKey, name),
                                ImmutableMap.<Symbol, ColumnHandle>builder()
                                        .put(nationKey, nationKeyColumn)
                                        .put(name, nameColumn)
                                        .build()));
            })
            .withSession(SESSION)
            .matches(project(
                    ImmutableMap.of("expr", expression(new SymbolReference(nationKeyColumn.getColumnName()))),
                    tableScan(
                            equalTo(new ThriftTableHandle(
                                    TINY_SCHEMA,
                                    "nation",
                                    TupleDomain.all(),
                                    Optional.of(ImmutableSet.of(nationKeyColumn)))),
                            TupleDomain.all(),
                            ImmutableMap.of(nationKeyColumn.getColumnName(), equalTo(nationKeyColumn)))));
}
 

HBaseRecordSet(HBaseSplit split, List<ColumnHandle> columnHandles, HBaseClientManager clientManager) {
    this.hBaseSplit = Objects.requireNonNull(split, "split is null");
    Objects.requireNonNull(clientManager, "clientManager is null");
    this.config = clientManager.getConfig();

    Objects.requireNonNull(columnHandles, "column handles is null");
    this.columnHandles = columnHandles.stream().map(ch -> (HBaseColumnHandle) ch).collect(Collectors.toList());
    this.initFieldIndexMap(this.columnHandles);

    this.columnTypes = columnHandles.stream().map(ch -> ((HBaseColumnHandle) ch).getColumnType())
            .collect(Collectors.toList());

    this.connection = clientManager.createConnection();
}
 

@Override
public void dropColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle column)
{
    IcebergTableHandle icebergTableHandle = (IcebergTableHandle) tableHandle;
    IcebergColumnHandle handle = (IcebergColumnHandle) column;
    org.apache.iceberg.Table icebergTable = getIcebergTable(metastore, hdfsEnvironment, session, icebergTableHandle.getSchemaTableName());
    icebergTable.updateSchema().deleteColumn(handle.getName()).commit();
}
 

public LookupSymbolResolver(Map<Symbol, ColumnHandle> assignments, Map<ColumnHandle, NullableValue> bindings)
{
    requireNonNull(assignments, "assignments is null");
    requireNonNull(bindings, "bindings is null");

    this.assignments = ImmutableMap.copyOf(assignments);
    this.bindings = ImmutableMap.copyOf(bindings);
}
 

@Override
public Optional<ConstraintApplicationResult<ConnectorTableHandle>> applyFilter(ConnectorSession session, ConnectorTableHandle handle, Constraint constraint)
{
    SystemTableHandle table = (SystemTableHandle) handle;

    TupleDomain<ColumnHandle> oldDomain = table.getConstraint();
    TupleDomain<ColumnHandle> newDomain = oldDomain.intersect(constraint.getSummary());
    if (oldDomain.equals(newDomain) && constraint.predicate().isEmpty()) {
        return Optional.empty();
    }

    SystemTable systemTable = checkAndGetTable(session, table);
    if (systemTable instanceof JdbcTable) {
        TupleDomain<ColumnHandle> filtered = ((JdbcTable) systemTable).applyFilter(session, new Constraint(newDomain, constraint.predicate(), constraint.getColumns()));
        newDomain = newDomain.intersect(filtered);
    }

    if (oldDomain.equals(newDomain)) {
        return Optional.empty();
    }

    if (newDomain.isNone()) {
        // TODO (https://github.com/prestosql/presto/issues/3647) indicate the table scan is empty
    }
    table = new SystemTableHandle(table.getSchemaName(), table.getTableName(), newDomain);
    return Optional.of(new ConstraintApplicationResult<>(table, constraint.getSummary()));
}
 

@Test
public void testAll()
{
    assertTrue(TupleDomain.all().isAll());
    assertEquals(TupleDomain.<ColumnHandle>all(),
            TupleDomain.withColumnDomains(ImmutableMap.of(
                    A, Domain.all(BIGINT))));
    assertEquals(TupleDomain.<ColumnHandle>all(),
            TupleDomain.withColumnDomains(ImmutableMap.<ColumnHandle, Domain>of()));
}
 

public TableScanNode tableScan(List<Symbol> symbols, Map<Symbol, ColumnHandle> assignments)
{
    return tableScan(
            new TableHandle(new CatalogName("testConnector"), new TestingTableHandle(), TestingTransactionHandle.create(), Optional.of(TestingHandle.INSTANCE)),
            symbols,
            assignments);
}
 

@Override
public ColumnMetadata getColumnMetadata(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle columnHandle)
{
    ConnectorTableMetadata tableMetadata = getTableMetadata(session, tableHandle);
    String columnName = ((TpcdsColumnHandle) columnHandle).getColumnName();

    for (ColumnMetadata column : tableMetadata.getColumns()) {
        if (column.getName().equals(columnName)) {
            return column;
        }
    }
    throw new IllegalArgumentException(format("Table '%s' does not have column '%s'", tableMetadata.getTable(), columnName));
}
 

public TableScanOperatorFactory(
        int operatorId,
        PlanNodeId sourceId,
        PageSourceProvider pageSourceProvider,
        TableHandle table,
        Iterable<ColumnHandle> columns,
        Supplier<TupleDomain<ColumnHandle>> dynamicFilter)
{
    this.operatorId = operatorId;
    this.sourceId = requireNonNull(sourceId, "sourceId is null");
    this.pageSourceProvider = requireNonNull(pageSourceProvider, "pageSourceProvider is null");
    this.table = requireNonNull(table, "table is null");
    this.columns = ImmutableList.copyOf(requireNonNull(columns, "columns is null"));
    this.dynamicFilter = requireNonNull(dynamicFilter, "dynamicFilter is null");
}
 

@Override
public ConnectorSplitSource getSplits(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorTableHandle table, SplitSchedulingStrategy splitSchedulingStrategy, Supplier<TupleDomain<ColumnHandle>> dynamicFilter)
{
    try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) {
        return delegate.getSplits(transaction, session, table, splitSchedulingStrategy, dynamicFilter);
    }
}
 

@JsonCreator
public TableInfo(
        @JsonProperty("tableName") QualifiedObjectName tableName,
        @JsonProperty("predicate") TupleDomain<ColumnHandle> predicate)
{
    this.tableName = requireNonNull(tableName, "tableName is null");
    this.predicate = requireNonNull(predicate, "predicate is null");
}
 

@Override
public Optional<Symbol> getAssignedSymbol(PlanNode node, Session session, Metadata metadata, SymbolAliases symbolAliases)
{
    TableHandle tableHandle;
    Map<Symbol, ColumnHandle> assignments;

    if (node instanceof TableScanNode) {
        TableScanNode tableScanNode = (TableScanNode) node;
        tableHandle = tableScanNode.getTable();
        assignments = tableScanNode.getAssignments();
    }
    else if (node instanceof IndexSourceNode) {
        IndexSourceNode indexSourceNode = (IndexSourceNode) node;
        tableHandle = indexSourceNode.getTableHandle();
        assignments = indexSourceNode.getAssignments();
    }
    else {
        return Optional.empty();
    }

    TableMetadata tableMetadata = metadata.getTableMetadata(session, tableHandle);
    String actualTableName = tableMetadata.getTable().getTableName();

    // Wrong table -> doesn't match.
    if (!tableName.equalsIgnoreCase(actualTableName)) {
        return Optional.empty();
    }

    Optional<ColumnHandle> columnHandle = getColumnHandle(tableHandle, session, metadata);

    checkState(columnHandle.isPresent(), "Table %s doesn't have column %s. Typo in test?", tableName, columnName);

    return getAssignedSymbol(assignments, columnHandle.get());
}