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

@Test
public void testNoPredicate()
        throws Exception
{
    JmxTableHandle tableHandle = new JmxTableHandle(new SchemaTableName("schema", "tableName"), ImmutableList.of("objectName"), ImmutableList.of(columnHandle), true, TupleDomain.all());
    ConnectorSplitSource splitSource = splitManager.getSplits(JmxTransactionHandle.INSTANCE, SESSION, tableHandle, UNGROUPED_SCHEDULING);
    List<ConnectorSplit> allSplits = getAllSplits(splitSource);
    assertEquals(allSplits.size(), nodes.size());

    Set<String> actualNodes = nodes.stream().map(Node::getNodeIdentifier).collect(toSet());
    Set<String> expectedNodes = new HashSet<>();
    for (ConnectorSplit split : allSplits) {
        List<HostAddress> addresses = split.getAddresses();
        assertEquals(addresses.size(), 1);
        expectedNodes.add(addresses.get(0).getHostText());
    }
    assertEquals(actualNodes, expectedNodes);
}
 

public HiveTableHandle(
        String schemaName,
        String tableName,
        Map<String, String> tableParameters,
        List<HiveColumnHandle> partitionColumns,
        Optional<HiveBucketHandle> bucketHandle)
{
    this(
            schemaName,
            tableName,
            Optional.of(tableParameters),
            partitionColumns,
            Optional.empty(),
            TupleDomain.all(),
            TupleDomain.all(),
            bucketHandle,
            Optional.empty(),
            Optional.empty(),
            Optional.empty(),
            Optional.empty());
}
 

@Test
public void testInformationSchemaPredicatePushdownWithConstraintPredicateOnViewsTable()
{
    TransactionId transactionId = transactionManager.beginTransaction(false);

    // predicate on non columns enumerating table should not cause tables to be enumerated
    Constraint constraint = new Constraint(TupleDomain.all(), TestInformationSchemaMetadata::testConstraint);
    ConnectorSession session = createNewSession(transactionId);
    ConnectorMetadata metadata = new InformationSchemaMetadata("test_catalog", this.metadata);
    InformationSchemaTableHandle tableHandle = (InformationSchemaTableHandle)
            metadata.getTableHandle(session, new SchemaTableName("information_schema", "views"));
    tableHandle = metadata.applyFilter(session, tableHandle, constraint)
            .map(ConstraintApplicationResult::getHandle)
            .map(InformationSchemaTableHandle.class::cast)
            .orElseThrow(AssertionError::new);

    assertEquals(tableHandle.getPrefixes(), ImmutableSet.of(new QualifiedTablePrefix("test_catalog", "test_schema")));
}
 

@BeforeMethod
public void setUp()
{
    symbolAllocator = new SymbolAllocator();
    columnA = symbolAllocator.newSymbol("a", BIGINT);
    columnB = symbolAllocator.newSymbol("b", INTEGER);
    columnC = symbolAllocator.newSymbol("c", DOUBLE);
    columnD = symbolAllocator.newSymbol("d", DATE);
    columnE = symbolAllocator.newSymbol("e", VarcharType.createVarcharType(3));  // varchar(3), to test type only coercion

    Map<Symbol, ColumnHandle> assignments = ImmutableMap.<Symbol, ColumnHandle>builder()
            .put(columnA, new TestingColumnHandle("a"))
            .put(columnB, new TestingColumnHandle("b"))
            .put(columnC, new TestingColumnHandle("c"))
            .put(columnD, new TestingColumnHandle("d"))
            .put(columnE, new TestingColumnHandle("e"))
            .build();

    baseTableScan = new TableScanNode(
            newId(),
            TEST_TABLE_HANDLE,
            ImmutableList.copyOf(assignments.keySet()),
            assignments,
            TupleDomain.all());
}
 

private TableScanNode tableScan(String id, String... symbols)
{
    List<Symbol> symbolsList = Arrays.stream(symbols).map(Symbol::new).collect(toImmutableList());
    ImmutableMap.Builder<Symbol, ColumnHandle> assignments = ImmutableMap.builder();

    for (Symbol symbol : symbolsList) {
        assignments.put(symbol, new TpchColumnHandle("orderkey", BIGINT));
    }

    TpchTableHandle tableHandle = new TpchTableHandle("orders", 1.0);
    return new TableScanNode(
            new PlanNodeId(id),
            new TableHandle(new CatalogName("tpch"), tableHandle, INSTANCE, Optional.of(new TpchTableLayoutHandle(tableHandle, TupleDomain.all()))),
            symbolsList,
            assignments.build(),
            TupleDomain.all());
}
 

private JmxTableHandle getJmxTableHandle(SchemaTableName tableName)
{
    try {
        String objectNamePattern = toPattern(tableName.getTableName().toLowerCase(ENGLISH));
        List<ObjectName> objectNames = mbeanServer.queryNames(WILDCARD, null).stream()
                .filter(name -> name.getCanonicalName().toLowerCase(ENGLISH).matches(objectNamePattern))
                .collect(toImmutableList());
        if (objectNames.isEmpty()) {
            return null;
        }
        List<JmxColumnHandle> columns = new ArrayList<>();
        columns.add(new JmxColumnHandle(NODE_COLUMN_NAME, createUnboundedVarcharType()));
        columns.add(new JmxColumnHandle(OBJECT_NAME_NAME, createUnboundedVarcharType()));
        for (ObjectName objectName : objectNames) {
            MBeanInfo mbeanInfo = mbeanServer.getMBeanInfo(objectName);

            getColumnHandles(mbeanInfo).forEach(columns::add);
        }

        // Since this method is being called on all nodes in the cluster, we must ensure (by sorting)
        // that attributes are in the same order on all of them.
        columns = columns.stream()
                .distinct()
                .sorted(comparing(JmxColumnHandle::getColumnName))
                .collect(toImmutableList());

        return new JmxTableHandle(tableName, objectNames.stream().map(ObjectName::toString).collect(toImmutableList()), columns, true, TupleDomain.all());
    }
    catch (JMException | PrestoException e) {
        return null;
    }
}
 

@Test
public void testJsonRoundTrip()
{
    JmxTableHandle table = new JmxTableHandle(SCHEMA_TABLE_NAME, ImmutableList.of("objectName"), COLUMNS, true, TupleDomain.all());

    String json = TABLE_CODEC.toJson(table);
    JmxTableHandle copy = TABLE_CODEC.fromJson(json);
    assertEquals(copy, table);
}
 

public ConnectorTableLayout(ConnectorTableLayoutHandle handle)
{
    this(handle,
            Optional.empty(),
            TupleDomain.all(),
            Optional.empty(),
            Optional.empty(),
            Optional.empty(),
            emptyList());
}
 

public JdbcTableHandle(SchemaTableName schemaTableName, @Nullable String catalogName, @Nullable String schemaName, String tableName)
{
    this(
            schemaTableName,
            catalogName,
            schemaName,
            tableName,
            TupleDomain.all(),
            Optional.empty(),
            OptionalLong.empty(),
            Optional.empty());
}
 

private JmxTableHandle getJmxHistoryTableHandle(SchemaTableName tableName)
{
    JmxTableHandle handle = getJmxTableHandle(tableName);
    if (handle == null) {
        return null;
    }
    ImmutableList.Builder<JmxColumnHandle> builder = ImmutableList.builder();
    builder.add(new JmxColumnHandle(TIMESTAMP_COLUMN_NAME, TIMESTAMP));
    builder.addAll(handle.getColumnHandles());
    return new JmxTableHandle(handle.getTableName(), handle.getObjectNames(), builder.build(), false, TupleDomain.all());
}
 

@Override
public ConnectorTableProperties getTableProperties(ConnectorSession session, ConnectorTableHandle table)
{
    InformationSchemaTableHandle tableHandle = (InformationSchemaTableHandle) table;
    return new ConnectorTableProperties(
            tableHandle.getPrefixes().isEmpty() ? TupleDomain.none() : TupleDomain.all(),
            Optional.empty(),
            Optional.empty(),
            Optional.empty(),
            emptyList());
}
 

@Test
public void testAllRoundTrip()
{
    TupleDomain<Symbol> tupleDomain = TupleDomain.all();
    ExtractionResult result = fromPredicate(toPredicate(tupleDomain));
    assertEquals(result.getRemainingExpression(), TRUE_LITERAL);
    assertEquals(result.getTupleDomain(), tupleDomain);
}
 

public PinotTableHandle(String schemaName, String tableName)
{
    this(schemaName, tableName, TupleDomain.all(), OptionalLong.empty(), Optional.empty());
}
 

public HivePartitionResult getPartitions(SemiTransactionalHiveMetastore metastore, HiveIdentity identity, ConnectorTableHandle tableHandle, Constraint constraint)
{
    HiveTableHandle hiveTableHandle = (HiveTableHandle) tableHandle;
    TupleDomain<ColumnHandle> effectivePredicate = constraint.getSummary()
            .intersect(hiveTableHandle.getEnforcedConstraint());

    SchemaTableName tableName = hiveTableHandle.getSchemaTableName();
    Optional<HiveBucketHandle> hiveBucketHandle = hiveTableHandle.getBucketHandle();
    List<HiveColumnHandle> partitionColumns = hiveTableHandle.getPartitionColumns();

    if (effectivePredicate.isNone()) {
        return new HivePartitionResult(partitionColumns, ImmutableList.of(), none(), none(), none(), hiveBucketHandle, Optional.empty());
    }

    Table table = metastore.getTable(identity, tableName.getSchemaName(), tableName.getTableName())
            .orElseThrow(() -> new TableNotFoundException(tableName));

    Optional<HiveBucketFilter> bucketFilter = getHiveBucketFilter(table, effectivePredicate);
    TupleDomain<HiveColumnHandle> compactEffectivePredicate = effectivePredicate
            .transform(HiveColumnHandle.class::cast)
            .simplify(domainCompactionThreshold);

    if (partitionColumns.isEmpty()) {
        return new HivePartitionResult(
                partitionColumns,
                ImmutableList.of(new HivePartition(tableName)),
                compactEffectivePredicate,
                effectivePredicate,
                TupleDomain.all(),
                hiveBucketHandle,
                bucketFilter);
    }

    List<Type> partitionTypes = partitionColumns.stream()
            .map(HiveColumnHandle::getType)
            .collect(toList());

    Iterable<HivePartition> partitionsIterable;
    Predicate<Map<ColumnHandle, NullableValue>> predicate = constraint.predicate().orElse(value -> true);
    if (hiveTableHandle.getPartitions().isPresent()) {
        partitionsIterable = hiveTableHandle.getPartitions().get().stream()
                .filter(partition -> partitionMatches(partitionColumns, effectivePredicate, predicate, partition))
                .collect(toImmutableList());
    }
    else {
        List<String> partitionNames = getFilteredPartitionNames(metastore, identity, tableName, partitionColumns, effectivePredicate);
        partitionsIterable = () -> partitionNames.stream()
                // Apply extra filters which could not be done by getFilteredPartitionNames
                .map(partitionName -> parseValuesAndFilterPartition(tableName, partitionName, partitionColumns, partitionTypes, effectivePredicate, predicate))
                .filter(Optional::isPresent)
                .map(Optional::get)
                .iterator();
    }

    // All partition key domains will be fully evaluated, so we don't need to include those
    TupleDomain<ColumnHandle> remainingTupleDomain = effectivePredicate.filter((column, domain) -> !partitionColumns.contains(column));
    TupleDomain<ColumnHandle> enforcedTupleDomain = effectivePredicate.filter((column, domain) -> partitionColumns.contains(column));
    return new HivePartitionResult(partitionColumns, partitionsIterable, compactEffectivePredicate, remainingTupleDomain, enforcedTupleDomain, hiveBucketHandle, bucketFilter);
}
 

public LocalFileTableHandle(SchemaTableName schemaTableName, OptionalInt timestampColumn, OptionalInt serverAddressColumn)
{
    this(schemaTableName, timestampColumn, serverAddressColumn, TupleDomain.all());
}
 

@Test(dataProvider = "testPropagateExceptionDataProvider", timeOut = 60_000)
public void testPropagateException(boolean error, int threads)
{
    AtomicBoolean iteratorUsedAfterException = new AtomicBoolean();

    BackgroundHiveSplitLoader backgroundHiveSplitLoader = new BackgroundHiveSplitLoader(
            SIMPLE_TABLE,
            () -> new Iterator<HivePartitionMetadata>()
            {
                private boolean threw;

                @Override
                public boolean hasNext()
                {
                    iteratorUsedAfterException.compareAndSet(false, threw);
                    return !threw;
                }

                @Override
                public HivePartitionMetadata next()
                {
                    iteratorUsedAfterException.compareAndSet(false, threw);
                    threw = true;
                    if (error) {
                        throw new Error("loading error occurred");
                    }
                    throw new RuntimeException("loading error occurred");
                }
            },
            TupleDomain.all(),
            TupleDomain::all,
            TYPE_MANAGER,
            createBucketSplitInfo(Optional.empty(), Optional.empty()),
            SESSION,
            new TestingHdfsEnvironment(TEST_FILES),
            new NamenodeStats(),
            new CachingDirectoryLister(new HiveConfig()),
            EXECUTOR,
            threads,
            false,
            false,
            Optional.empty());

    HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader);
    backgroundHiveSplitLoader.start(hiveSplitSource);

    assertThatThrownBy(() -> drain(hiveSplitSource))
            .hasMessageEndingWith("loading error occurred");

    assertThatThrownBy(hiveSplitSource::isFinished)
            .hasMessageEndingWith("loading error occurred");

    if (threads == 1) {
        assertFalse(iteratorUsedAfterException.get());
    }
}
 

@Override
protected ExtractionResult visitExpression(Expression node, Boolean complement)
{
    // If we don't know how to process this node, the default response is to say that the TupleDomain is "all"
    return new ExtractionResult(TupleDomain.all(), complementIfNecessary(node, complement));
}
 

@Test
public void testTableScan()
{
    // Effective predicate is True if there is no effective predicate
    Map<Symbol, ColumnHandle> assignments = Maps.filterKeys(scanAssignments, Predicates.in(ImmutableList.of(A, B, C, D)));
    PlanNode node = TableScanNode.newInstance(
            newId(),
            makeTableHandle(TupleDomain.all()),
            ImmutableList.copyOf(assignments.keySet()),
            assignments);
    Expression effectivePredicate = effectivePredicateExtractor.extract(SESSION, node, TypeProvider.empty(), typeAnalyzer);
    assertEquals(effectivePredicate, BooleanLiteral.TRUE_LITERAL);

    node = new TableScanNode(
            newId(),
            makeTableHandle(TupleDomain.none()),
            ImmutableList.copyOf(assignments.keySet()),
            assignments,
            TupleDomain.none());
    effectivePredicate = effectivePredicateExtractor.extract(SESSION, node, TypeProvider.empty(), typeAnalyzer);
    assertEquals(effectivePredicate, FALSE_LITERAL);

    TupleDomain<ColumnHandle> predicate = TupleDomain.withColumnDomains(ImmutableMap.of(scanAssignments.get(A), Domain.singleValue(BIGINT, 1L)));
    node = new TableScanNode(
            newId(),
            makeTableHandle(predicate),
            ImmutableList.copyOf(assignments.keySet()),
            assignments,
            predicate);
    effectivePredicate = effectivePredicateExtractor.extract(SESSION, node, TypeProvider.empty(), typeAnalyzer);
    assertEquals(normalizeConjuncts(effectivePredicate), normalizeConjuncts(equals(bigintLiteral(1L), AE)));

    predicate = TupleDomain.withColumnDomains(ImmutableMap.of(
            scanAssignments.get(A), Domain.singleValue(BIGINT, 1L),
            scanAssignments.get(B), Domain.singleValue(BIGINT, 2L)));
    node = new TableScanNode(
            newId(),
            makeTableHandle(TupleDomain.withColumnDomains(ImmutableMap.of(scanAssignments.get(A), Domain.singleValue(BIGINT, 1L)))),
            ImmutableList.copyOf(assignments.keySet()),
            assignments,
            predicate);
    effectivePredicate = effectivePredicateExtractorWithoutTableProperties.extract(SESSION, node, TypeProvider.empty(), typeAnalyzer);
    assertEquals(normalizeConjuncts(effectivePredicate), normalizeConjuncts(equals(bigintLiteral(2L), BE), equals(bigintLiteral(1L), AE)));

    node = new TableScanNode(
            newId(),
            makeTableHandle(predicate),
            ImmutableList.copyOf(assignments.keySet()),
            assignments,
            TupleDomain.all());
    effectivePredicate = effectivePredicateExtractor.extract(SESSION, node, TypeProvider.empty(), typeAnalyzer);
    assertEquals(effectivePredicate, and(equals(AE, bigintLiteral(1)), equals(BE, bigintLiteral(2))));

    node = new TableScanNode(
            newId(),
            makeTableHandle(predicate),
            ImmutableList.copyOf(assignments.keySet()),
            assignments,
            TupleDomain.withColumnDomains(ImmutableMap.of(
                    scanAssignments.get(A), Domain.multipleValues(BIGINT, ImmutableList.of(1L, 2L, 3L)),
                    scanAssignments.get(B), Domain.multipleValues(BIGINT, ImmutableList.of(1L, 2L, 3L)))));
    effectivePredicate = effectivePredicateExtractor.extract(SESSION, node, TypeProvider.empty(), typeAnalyzer);
    assertEquals(normalizeConjuncts(effectivePredicate), normalizeConjuncts(equals(bigintLiteral(2L), BE), equals(bigintLiteral(1L), AE)));

    node = new TableScanNode(
            newId(),
            makeTableHandle(TupleDomain.all()),
            ImmutableList.copyOf(assignments.keySet()),
            assignments,
            TupleDomain.all());
    effectivePredicate = effectivePredicateExtractor.extract(SESSION, node, TypeProvider.empty(), typeAnalyzer);
    assertEquals(effectivePredicate, BooleanLiteral.TRUE_LITERAL);
}
 

public static SystemTableHandle fromSchemaTableName(SchemaTableName tableName)
{
    requireNonNull(tableName, "tableName is null");
    return new SystemTableHandle(tableName.getSchemaName(), tableName.getTableName(), TupleDomain.all());
}
 

public static BigQueryTableHandle from(TableInfo tableInfo)
{
    TableId tableId = tableInfo.getTableId();
    String type = tableInfo.getDefinition().getType().toString();
    return new BigQueryTableHandle(tableId.getProject(), tableId.getDataset(), tableId.getTable(), type, TupleDomain.all(), Optional.empty(), OptionalLong.empty());
}