Java Code Examples for io.prestosql.spi.predicate.TupleDomain#fromFixedValues()

@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 testApplyFilterWithConstraint()
{
    JmxTableHandle handle = metadata.getTableHandle(SESSION, new SchemaTableName(JMX_SCHEMA_NAME, "java.lang:*"));

    JmxColumnHandle nodeColumnHandle = new JmxColumnHandle("node", createUnboundedVarcharType());
    NullableValue nodeColumnValue = NullableValue.of(createUnboundedVarcharType(), utf8Slice(localNode.getNodeIdentifier()));

    JmxColumnHandle objectNameColumnHandle = new JmxColumnHandle("object_name", createUnboundedVarcharType());
    NullableValue objectNameColumnValue = NullableValue.of(createUnboundedVarcharType(), utf8Slice("presto.memory:type=MemoryPool,name=reserved"));

    TupleDomain<ColumnHandle> tupleDomain = TupleDomain.fromFixedValues(ImmutableMap.of(nodeColumnHandle, nodeColumnValue, objectNameColumnHandle, objectNameColumnValue));

    Optional<ConstraintApplicationResult<ConnectorTableHandle>> result = metadata.applyFilter(SESSION, handle, new Constraint(tupleDomain));

    assertTrue(result.isPresent());
    assertEquals(result.get().getRemainingFilter(), TupleDomain.fromFixedValues(ImmutableMap.of(objectNameColumnHandle, objectNameColumnValue)));
    assertEquals(((JmxTableHandle) result.get().getHandle()).getNodeFilter(), TupleDomain.fromFixedValues(ImmutableMap.of(nodeColumnHandle, nodeColumnValue)));
}
 

private boolean rowMatchesPredicate()
{
    if (predicate.isAll()) {
        return true;
    }
    if (predicate.isNone()) {
        return false;
    }

    Map<ColumnHandle, NullableValue> rowMap = predicate.getDomains().get().keySet().stream()
            .collect(toImmutableMap(
                    column -> column,
                    column -> {
                        TpchColumnHandle tpchColumnHandle = (TpchColumnHandle) column;
                        Type type = tpchColumnHandle.getType();
                        TpchColumn<E> tpchColumn = table.getColumn(tpchColumnHandle.getColumnName());
                        return NullableValue.of(type, getPrestoObject(tpchColumn, type));
                    }));

    TupleDomain<ColumnHandle> rowTupleDomain = TupleDomain.fromFixedValues(rowMap);

    return predicate.contains(rowTupleDomain);
}
 

@Test
public void testApplyFilterWithSameConstraint()
{
    JmxTableHandle handle = metadata.getTableHandle(SESSION, new SchemaTableName(JMX_SCHEMA_NAME, "java.lang:*"));

    JmxColumnHandle columnHandle = new JmxColumnHandle("node", createUnboundedVarcharType());
    TupleDomain<ColumnHandle> nodeTupleDomain = TupleDomain.fromFixedValues(ImmutableMap.of(columnHandle, NullableValue.of(createUnboundedVarcharType(), utf8Slice(localNode.getNodeIdentifier()))));

    JmxTableHandle newTableHandle = new JmxTableHandle(handle.getTableName(), handle.getObjectNames(), handle.getColumnHandles(), handle.isLiveData(), nodeTupleDomain);

    Optional<ConstraintApplicationResult<ConnectorTableHandle>> result = metadata.applyFilter(SESSION, newTableHandle, new Constraint(nodeTupleDomain));
    assertFalse(result.isPresent());
}
 

@Override
public Optional<ConnectorResolvedIndex> resolveIndex(
        ConnectorSession session,
        ConnectorTableHandle tableHandle,
        Set<ColumnHandle> indexableColumns,
        Set<ColumnHandle> outputColumns,
        TupleDomain<ColumnHandle> tupleDomain)
{
    TpchTableHandle tpchTableHandle = (TpchTableHandle) tableHandle;

    // Keep the fixed values that don't overlap with the indexableColumns
    // Note: technically we could more efficiently utilize the overlapped columns, but this way is simpler for now

    Map<ColumnHandle, NullableValue> fixedValues = TupleDomain.extractFixedValues(tupleDomain).orElse(ImmutableMap.of())
            .entrySet().stream()
            .filter(entry -> !indexableColumns.contains(entry.getKey()))
            .filter(entry -> !entry.getValue().isNull()) // strip nulls since meaningless in index join lookups
            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));

    // determine all columns available for index lookup
    Set<String> lookupColumnNames = ImmutableSet.<String>builder()
            .addAll(handleToNames(ImmutableList.copyOf(indexableColumns)))
            .addAll(handleToNames(ImmutableList.copyOf(fixedValues.keySet())))
            .build();

    // do we have an index?
    if (indexedData.getIndexedTable(tpchTableHandle.getTableName(), tpchTableHandle.getScaleFactor(), lookupColumnNames).isEmpty()) {
        return Optional.empty();
    }

    TupleDomain<ColumnHandle> filteredTupleDomain = tupleDomain.filter((column, domain) -> !fixedValues.containsKey(column));
    TpchIndexHandle indexHandle = new TpchIndexHandle(
            tpchTableHandle.getTableName(),
            tpchTableHandle.getScaleFactor(),
            lookupColumnNames,
            TupleDomain.fromFixedValues(fixedValues));
    return Optional.of(new ConnectorResolvedIndex(indexHandle, filteredTupleDomain));
}
 

private static TupleDomain<ColumnHandle> fixedValueTupleDomain(TpchMetadata tpchMetadata, TpchColumn<?> column, Object value)
{
    requireNonNull(column, "column is null");
    requireNonNull(value, "value is null");
    return TupleDomain.fromFixedValues(
            ImmutableMap.of(tpchMetadata.toColumnHandle(column), new NullableValue(getPrestoType(column), value)));
}
 

private static TupleDomain<ColumnHandle> fixedValueTupleDomain(TpchMetadata tpchMetadata, TpchColumn<?> column1, Object value1, TpchColumn<?> column2, Object value2)
{
    return TupleDomain.fromFixedValues(
            ImmutableMap.of(
                    tpchMetadata.toColumnHandle(column1), new NullableValue(getPrestoType(column1), value1),
                    tpchMetadata.toColumnHandle(column2), new NullableValue(getPrestoType(column2), value2)));
}
 

/**
 * Get the list of partitions matching the given filters on partition keys.
 *
 * @param table the table to get partitions from
 * @param filterPrefixes the list of possible values for each partition key.
 * Order of values should match {@link CassandraTable#getPartitionKeyColumns()}
 * @return list of {@link CassandraPartition}
 */
public List<CassandraPartition> getPartitions(CassandraTable table, List<Set<Object>> filterPrefixes)
{
    List<CassandraColumnHandle> partitionKeyColumns = table.getPartitionKeyColumns();

    if (filterPrefixes.size() != partitionKeyColumns.size()) {
        return ImmutableList.of(CassandraPartition.UNPARTITIONED);
    }

    Iterable<Row> rows;
    if (getCassandraVersion().compareTo(PARTITION_FETCH_WITH_IN_PREDICATE_VERSION) > 0) {
        log.debug("Using IN predicate to fetch partitions.");
        rows = queryPartitionKeysWithInClauses(table, filterPrefixes);
    }
    else {
        log.debug("Using combination of partition values to fetch partitions.");
        rows = queryPartitionKeysLegacyWithMultipleQueries(table, filterPrefixes);
    }

    if (rows == null) {
        // just split the whole partition range
        return ImmutableList.of(CassandraPartition.UNPARTITIONED);
    }

    ByteBuffer buffer = ByteBuffer.allocate(1000);
    HashMap<ColumnHandle, NullableValue> map = new HashMap<>();
    Set<String> uniquePartitionIds = new HashSet<>();
    StringBuilder stringBuilder = new StringBuilder();

    boolean isComposite = partitionKeyColumns.size() > 1;

    ImmutableList.Builder<CassandraPartition> partitions = ImmutableList.builder();
    for (Row row : rows) {
        buffer.clear();
        map.clear();
        stringBuilder.setLength(0);
        for (int i = 0; i < partitionKeyColumns.size(); i++) {
            ByteBuffer component = row.getBytesUnsafe(i);
            if (isComposite) {
                // build composite key
                short len = (short) component.limit();
                buffer.putShort(len);
                buffer.put(component);
                buffer.put((byte) 0);
            }
            else {
                buffer.put(component);
            }
            CassandraColumnHandle columnHandle = partitionKeyColumns.get(i);
            NullableValue keyPart = columnHandle.getCassandraType().getColumnValue(row, i);
            map.put(columnHandle, keyPart);
            if (i > 0) {
                stringBuilder.append(" AND ");
            }
            stringBuilder.append(CassandraCqlUtils.validColumnName(columnHandle.getName()));
            stringBuilder.append(" = ");
            stringBuilder.append(columnHandle.getCassandraType().getColumnValueForCql(row, i));
        }
        buffer.flip();
        byte[] key = new byte[buffer.limit()];
        buffer.get(key);
        TupleDomain<ColumnHandle> tupleDomain = TupleDomain.fromFixedValues(map);
        String partitionId = stringBuilder.toString();
        if (uniquePartitionIds.add(partitionId)) {
            partitions.add(new CassandraPartition(key, partitionId, tupleDomain, false));
        }
    }
    return partitions.build();
}