Java Code Examples for io.prestosql.spi.predicate.Domain#create()

The following examples show how to use io.prestosql.spi.predicate.Domain#create() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: TestPrometheusSplit.java    From presto with Apache License 2.0 6 votes vote down vote up 
@Test
public void testPredicatePushDownSetsLowerBoundOnly()
        throws Exception
{
    long predicateLowValue = 1568638171999L - 600000L;
    Range lowRange = Range.greaterThanOrEqual(TimestampType.TIMESTAMP, predicateLowValue);
    ValueSet valueSet = ValueSet.ofRanges(lowRange);
    Domain testDomain = Domain.create(valueSet, false);
    TupleDomain<ColumnHandle> testTupleDomain = TupleDomain.withColumnDomains(ImmutableMap.<ColumnHandle, Domain>of(
            (ColumnHandle) new PrometheusColumnHandle("timestamp", TimestampType.TIMESTAMP, 2), testDomain));
    prometheusTableHandle.setPredicate(Optional.of(testTupleDomain));
    io.airlift.units.Duration maxQueryRangeDuration = new io.airlift.units.Duration(120, TimeUnit.SECONDS);
    io.airlift.units.Duration queryChunkSizeDuration = new io.airlift.units.Duration(30, TimeUnit.SECONDS);
    LocalDateTime now = ofInstant(Instant.ofEpochMilli(1568638171999L), ZoneId.systemDefault());
    PrometheusTimeMachine.useFixedClockAt(now);
    TemporalAmount maxQueryAsTime = java.time.Duration.ofMillis(maxQueryRangeDuration.toMillis());
    List<String> splitTimes = PrometheusSplitManager.generateTimesForSplits(now, maxQueryRangeDuration, queryChunkSizeDuration, prometheusTableHandle);

    String earliestSplit = splitTimes.get(0);
    ZonedDateTime earliestSplitAsTime = ZonedDateTime.ofInstant(Instant.ofEpochMilli(longFromDecimalSecondString(earliestSplit)), ZoneId.systemDefault());
    TemporalAmount queryChunkAsTime = java.time.Duration.ofMillis(queryChunkSizeDuration.toMillis());
    ZonedDateTime startOfQuery = earliestSplitAsTime.minus(queryChunkAsTime);
    assertNotEquals(startOfQuery.toLocalDateTime(), now.minus(maxQueryAsTime).minus(java.time.Duration.ofMillis((splitTimes.size() - 1) * OFFSET_MILLIS)));
    assertEquals(startOfQuery.toInstant().toEpochMilli(), Instant.ofEpochMilli(predicateLowValue).toEpochMilli() - ((splitTimes.size() - 1) * OFFSET_MILLIS));
}
Example 2
Source File: TupleDomainOrcPredicate.java    From presto with Apache License 2.0 6 votes vote down vote up 
private static <F, T extends Comparable<T>> Domain createDomain(Type type, boolean hasNullValue, RangeStatistics<F> rangeStatistics, Function<F, T> function)
{
    F min = rangeStatistics.getMin();
    F max = rangeStatistics.getMax();

    if (min != null && max != null) {
        return Domain.create(ValueSet.ofRanges(Range.range(type, function.apply(min), true, function.apply(max), true)), hasNullValue);
    }
    if (max != null) {
        return Domain.create(ValueSet.ofRanges(Range.lessThanOrEqual(type, function.apply(max))), hasNullValue);
    }
    if (min != null) {
        return Domain.create(ValueSet.ofRanges(Range.greaterThanOrEqual(type, function.apply(min))), hasNullValue);
    }
    return Domain.create(ValueSet.all(type), hasNullValue);
}
Example 3
Source File: AtopSplitManager.java    From presto with Apache License 2.0 6 votes vote down vote up 
@Override
public ConnectorSplitSource getSplits(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorTableHandle table, SplitSchedulingStrategy splitSchedulingStrategy)
{
    AtopTableHandle tableHandle = (AtopTableHandle) table;

    List<ConnectorSplit> splits = new ArrayList<>();
    ZonedDateTime end = ZonedDateTime.now(timeZone);
    for (Node node : nodeManager.getWorkerNodes()) {
        ZonedDateTime start = end.minusDays(maxHistoryDays - 1).withHour(0).withMinute(0).withSecond(0).withNano(0);
        while (start.isBefore(end)) {
            ZonedDateTime splitEnd = start.withHour(23).withMinute(59).withSecond(59).withNano(0);
            Domain splitDomain = Domain.create(ValueSet.ofRanges(Range.range(TIMESTAMP_WITH_TIME_ZONE, 1000 * start.toEpochSecond(), true, 1000 * splitEnd.toEpochSecond(), true)), false);
            if (tableHandle.getStartTimeConstraint().overlaps(splitDomain) && tableHandle.getEndTimeConstraint().overlaps(splitDomain)) {
                splits.add(new AtopSplit(node.getHostAndPort(), start.toEpochSecond(), start.getZone()));
            }
            start = start.plusDays(1).withHour(0).withMinute(0).withSecond(0).withNano(0);
        }
    }

    return new FixedSplitSource(splits);
}
Example 4
Source File: TupleDomainParquetPredicate.java    From presto with Apache License 2.0 6 votes vote down vote up 
private static <F, T> Domain createDomain(Type type, boolean hasNullValue, ParquetRangeStatistics<F> rangeStatistics, Function<F, T> function)
{
    F min = rangeStatistics.getMin();
    F max = rangeStatistics.getMax();

    if (min != null && max != null) {
        return Domain.create(ValueSet.ofRanges(Range.range(type, function.apply(min), true, function.apply(max), true)), hasNullValue);
    }
    if (max != null) {
        return Domain.create(ValueSet.ofRanges(Range.lessThanOrEqual(type, function.apply(max))), hasNullValue);
    }
    if (min != null) {
        return Domain.create(ValueSet.ofRanges(Range.greaterThanOrEqual(type, function.apply(min))), hasNullValue);
    }
    return Domain.create(ValueSet.all(type), hasNullValue);
}
Example 5
Source File: TestPrometheusSplit.java    From presto with Apache License 2.0 5 votes vote down vote up 
@Test
public void testPredicatePushDownLowerBoundDirect()
        throws Exception
{
    Range lowRange = Range.greaterThanOrEqual(TimestampType.TIMESTAMP, 1570460709643L);
    ValueSet valueSet = ValueSet.ofRanges(lowRange);
    Domain testDomain = Domain.create(valueSet, false);
    TupleDomain<ColumnHandle> testTupleDomain = TupleDomain.withColumnDomains(ImmutableMap.<ColumnHandle, Domain>of(
            (ColumnHandle) new PrometheusColumnHandle("timestamp", TimestampType.TIMESTAMP, 2), testDomain));
    prometheusTableHandle.setPredicate(Optional.of(testTupleDomain));
    Optional<PrometheusPredicateTimeInfo> predicateTimes = PrometheusSplitManager.determinePredicateTimes(testTupleDomain);
    ZonedDateTime expected = ZonedDateTime.ofInstant(Instant.ofEpochMilli(1570460709643L), ZoneId.systemDefault());
    assertEquals(predicateTimes.get().predicateLowerTimeBound.get(), expected);
}
Example 6
Source File: TestPrometheusSplit.java    From presto with Apache License 2.0 5 votes vote down vote up 
@Test
public void testPredicatePushDownSetsUpperBoundOnly()
        throws Exception
{
    long predicateHighValue = 1568638171999L;
    Range highRange = Range.lessThanOrEqual(TimestampType.TIMESTAMP, predicateHighValue);
    ValueSet valueSet = ValueSet.ofRanges(highRange);

    Domain testDomain = Domain.create(valueSet, false);
    TupleDomain<ColumnHandle> testTupleDomain = TupleDomain.withColumnDomains(ImmutableMap.<ColumnHandle, Domain>of(
            (ColumnHandle) new PrometheusColumnHandle("timestamp", TimestampType.TIMESTAMP, 2), testDomain));
    prometheusTableHandle.setPredicate(Optional.of(testTupleDomain));
    io.airlift.units.Duration maxQueryRangeDuration = new io.airlift.units.Duration(120, TimeUnit.SECONDS);
    io.airlift.units.Duration queryChunkSizeDuration = new io.airlift.units.Duration(30, TimeUnit.SECONDS);
    LocalDateTime now = ofInstant(Instant.ofEpochMilli(1568638171999L + 600000L), ZoneId.systemDefault());
    PrometheusTimeMachine.useFixedClockAt(now);

    List<String> splitTimes = PrometheusSplitManager.generateTimesForSplits(now, maxQueryRangeDuration, queryChunkSizeDuration, prometheusTableHandle);

    TemporalAmount expectedMaxQueryAsTime = java.time.Duration.ofMillis(maxQueryRangeDuration.toMillis() +
            ((splitTimes.size() - 1) * OFFSET_MILLIS));
    String lastSplit = splitTimes.get(splitTimes.size() - 1);
    ZonedDateTime lastSplitAsTime = ZonedDateTime.ofInstant(Instant.ofEpochMilli(longFromDecimalSecondString(lastSplit)), ZoneId.systemDefault());
    String earliestSplit = splitTimes.get(0);
    ZonedDateTime earliestSplitAsTime = ZonedDateTime.ofInstant(Instant.ofEpochMilli(longFromDecimalSecondString(earliestSplit)), ZoneId.systemDefault());
    TemporalAmount queryChunkAsTime = java.time.Duration.ofMillis(queryChunkSizeDuration.toMillis());
    java.time.Duration actualMaxDuration = Duration.between(earliestSplitAsTime
            .minus(queryChunkAsTime), lastSplitAsTime);

    assertEquals(lastSplitAsTime.toInstant().toEpochMilli(), 1568638171999L);
    assertEquals(actualMaxDuration, expectedMaxQueryAsTime);
}
Example 7
Source File: TestPrometheusSplit.java    From presto with Apache License 2.0 5 votes vote down vote up 
@Test
public void testPredicatePushDownSetsUpperAndLowerBound()
        throws Exception
{
    long predicateHighValue = 1568638171999L;
    Range highRange = Range.equal(TimestampType.TIMESTAMP, predicateHighValue);
    long predicateLowValue = 1568638171999L - 600000L;
    Range lowRange = Range.equal(TimestampType.TIMESTAMP, predicateLowValue);
    ValueSet valueSet = ValueSet.ofRanges(lowRange, highRange);

    Domain testDomain = Domain.create(valueSet, false);
    TupleDomain<ColumnHandle> testTupleDomain = TupleDomain.withColumnDomains(ImmutableMap.<ColumnHandle, Domain>of(
            (ColumnHandle) new PrometheusColumnHandle("timestamp", TimestampType.TIMESTAMP, 2), testDomain));
    prometheusTableHandle.setPredicate(Optional.of(testTupleDomain));
    io.airlift.units.Duration maxQueryRangeDuration = new io.airlift.units.Duration(120, TimeUnit.SECONDS);
    io.airlift.units.Duration queryChunkSizeDuration = new io.airlift.units.Duration(30, TimeUnit.SECONDS);
    LocalDateTime now = ofInstant(Instant.ofEpochMilli(1568638171999L + 1200000L), ZoneId.systemDefault());
    PrometheusTimeMachine.useFixedClockAt(now);

    List<String> splitTimes = PrometheusSplitManager.generateTimesForSplits(now, maxQueryRangeDuration, queryChunkSizeDuration, prometheusTableHandle);

    TemporalAmount expectedMaxQueryAsTime = java.time.Duration.ofMillis(new io.airlift.units.Duration(10, TimeUnit.MINUTES).toMillis() +
            ((splitTimes.size() - 1) * OFFSET_MILLIS));
    String lastSplit = splitTimes.get(splitTimes.size() - 1);
    ZonedDateTime lastSplitAsTime = ZonedDateTime.ofInstant(Instant.ofEpochMilli(longFromDecimalSecondString(lastSplit)), ZoneId.systemDefault());
    String earliestSplit = splitTimes.get(0);
    ZonedDateTime earliestSplitAsTime = ZonedDateTime.ofInstant(Instant.ofEpochMilli(longFromDecimalSecondString(earliestSplit)), ZoneId.systemDefault());
    TemporalAmount queryChunkAsTime = java.time.Duration.ofMillis(queryChunkSizeDuration.toMillis());
    java.time.Duration actualMaxDuration = Duration.between(earliestSplitAsTime
            .minus(queryChunkAsTime), lastSplitAsTime);

    assertEquals(lastSplitAsTime.toInstant().toEpochMilli(), 1568638171999L);
    assertEquals(actualMaxDuration, expectedMaxQueryAsTime);
}
Example 8
Source File: DynamicFilterSourceOperator.java    From presto with Apache License 2.0 5 votes vote down vote up 
private Domain convertToDomain(Type type, Block block)
{
    ImmutableList.Builder<Object> values = ImmutableList.builder();
    for (int position = 0; position < block.getPositionCount(); ++position) {
        Object value = TypeUtils.readNativeValue(type, block, position);
        if (value != null) {
            values.add(value);
        }
    }
    // Inner and right join doesn't match rows with null key column values.
    return Domain.create(ValueSet.copyOf(type, values.build()), false);
}
Example 9
Source File: DomainTranslator.java    From presto with Apache License 2.0 5 votes vote down vote up 
private static Domain extractEquatableDomain(ComparisonExpression.Operator comparisonOperator, Type type, Object value, boolean complement)
{
    checkArgument(value != null);
    switch (comparisonOperator) {
        case EQUAL:
            return Domain.create(complementIfNecessary(ValueSet.of(type, value), complement), false);
        case NOT_EQUAL:
            return Domain.create(complementIfNecessary(ValueSet.of(type, value).complement(), complement), false);
        case IS_DISTINCT_FROM:
            // Need to potential complement the whole domain for IS_DISTINCT_FROM since it is null-aware
            return complementIfNecessary(Domain.create(ValueSet.of(type, value).complement(), true), complement);
        default:
            throw new AssertionError("Unhandled operator: " + comparisonOperator);
    }
}
Example 10
Source File: TupleDomainOrcPredicate.java    From presto with Apache License 2.0 4 votes vote down vote up 
@VisibleForTesting
public static Domain getDomain(Type type, long rowCount, ColumnStatistics columnStatistics)
{
    if (rowCount == 0) {
        return Domain.none(type);
    }

    if (columnStatistics == null) {
        return Domain.all(type);
    }

    if (columnStatistics.hasNumberOfValues() && columnStatistics.getNumberOfValues() == 0) {
        return Domain.onlyNull(type);
    }

    boolean hasNullValue = columnStatistics.getNumberOfValues() != rowCount;

    if (type.getJavaType() == boolean.class && columnStatistics.getBooleanStatistics() != null) {
        BooleanStatistics booleanStatistics = columnStatistics.getBooleanStatistics();

        boolean hasTrueValues = (booleanStatistics.getTrueValueCount() != 0);
        boolean hasFalseValues = (columnStatistics.getNumberOfValues() != booleanStatistics.getTrueValueCount());
        if (hasTrueValues && hasFalseValues) {
            return Domain.all(BOOLEAN);
        }
        if (hasTrueValues) {
            return Domain.create(ValueSet.of(BOOLEAN, true), hasNullValue);
        }
        if (hasFalseValues) {
            return Domain.create(ValueSet.of(BOOLEAN, false), hasNullValue);
        }
    }
    else if (isShortDecimal(type) && columnStatistics.getDecimalStatistics() != null) {
        return createDomain(type, hasNullValue, columnStatistics.getDecimalStatistics(), value -> rescale(value, (DecimalType) type).unscaledValue().longValue());
    }
    else if (isLongDecimal(type) && columnStatistics.getDecimalStatistics() != null) {
        return createDomain(type, hasNullValue, columnStatistics.getDecimalStatistics(), value -> encodeUnscaledValue(rescale(value, (DecimalType) type).unscaledValue()));
    }
    else if (isCharType(type) && columnStatistics.getStringStatistics() != null) {
        return createDomain(type, hasNullValue, columnStatistics.getStringStatistics(), value -> truncateToLengthAndTrimSpaces(value, type));
    }
    else if (isVarcharType(type) && columnStatistics.getStringStatistics() != null) {
        return createDomain(type, hasNullValue, columnStatistics.getStringStatistics());
    }
    else if (type instanceof DateType && columnStatistics.getDateStatistics() != null) {
        return createDomain(type, hasNullValue, columnStatistics.getDateStatistics(), value -> (long) value);
    }
    else if (type.getJavaType() == long.class && columnStatistics.getIntegerStatistics() != null) {
        return createDomain(type, hasNullValue, columnStatistics.getIntegerStatistics());
    }
    else if (type.getJavaType() == double.class && columnStatistics.getDoubleStatistics() != null) {
        return createDomain(type, hasNullValue, columnStatistics.getDoubleStatistics());
    }
    else if (REAL.equals(type) && columnStatistics.getDoubleStatistics() != null) {
        return createDomain(type, hasNullValue, columnStatistics.getDoubleStatistics(), value -> (long) floatToRawIntBits(value.floatValue()));
    }
    return Domain.create(ValueSet.all(type), hasNullValue);
}
Example 11
Source File: TestDomainTranslator.java    From presto with Apache License 2.0 4 votes vote down vote up 
@Test
public void testInOptimization()
{
    Domain testDomain = Domain.create(
            ValueSet.all(BIGINT)
                    .subtract(ValueSet.ofRanges(
                            Range.equal(BIGINT, 1L), Range.equal(BIGINT, 2L), Range.equal(BIGINT, 3L))), false);

    TupleDomain<Symbol> tupleDomain = withColumnDomains(ImmutableMap.<Symbol, Domain>builder().put(C_BIGINT, testDomain).build());
    assertEquals(toPredicate(tupleDomain), not(in(C_BIGINT, ImmutableList.of(1L, 2L, 3L))));

    testDomain = Domain.create(
            ValueSet.ofRanges(
                    Range.lessThan(BIGINT, 4L)).intersect(
                    ValueSet.all(BIGINT)
                            .subtract(ValueSet.ofRanges(Range.equal(BIGINT, 1L), Range.equal(BIGINT, 2L), Range.equal(BIGINT, 3L)))), false);

    tupleDomain = withColumnDomains(ImmutableMap.<Symbol, Domain>builder().put(C_BIGINT, testDomain).build());
    assertEquals(toPredicate(tupleDomain), and(lessThan(C_BIGINT, bigintLiteral(4L)), not(in(C_BIGINT, ImmutableList.of(1L, 2L, 3L)))));

    testDomain = Domain.create(ValueSet.ofRanges(
            Range.range(BIGINT, 1L, true, 3L, true),
            Range.range(BIGINT, 5L, true, 7L, true),
            Range.range(BIGINT, 9L, true, 11L, true)),
            false);

    tupleDomain = withColumnDomains(ImmutableMap.<Symbol, Domain>builder().put(C_BIGINT, testDomain).build());
    assertEquals(toPredicate(tupleDomain),
            or(between(C_BIGINT, bigintLiteral(1L), bigintLiteral(3L)), between(C_BIGINT, bigintLiteral(5L), bigintLiteral(7L)), between(C_BIGINT, bigintLiteral(9L), bigintLiteral(11L))));

    testDomain = Domain.create(
            ValueSet.ofRanges(
                    Range.lessThan(BIGINT, 4L))
                    .intersect(ValueSet.all(BIGINT)
                            .subtract(ValueSet.ofRanges(Range.equal(BIGINT, 1L), Range.equal(BIGINT, 2L), Range.equal(BIGINT, 3L))))
                    .union(ValueSet.ofRanges(Range.range(BIGINT, 7L, true, 9L, true))), false);

    tupleDomain = withColumnDomains(ImmutableMap.<Symbol, Domain>builder().put(C_BIGINT, testDomain).build());
    assertEquals(toPredicate(tupleDomain), or(and(lessThan(C_BIGINT, bigintLiteral(4L)), not(in(C_BIGINT, ImmutableList.of(1L, 2L, 3L)))), between(C_BIGINT, bigintLiteral(7L), bigintLiteral(9L))));

    testDomain = Domain.create(
            ValueSet.ofRanges(Range.lessThan(BIGINT, 4L))
                    .intersect(ValueSet.all(BIGINT)
                            .subtract(ValueSet.ofRanges(Range.equal(BIGINT, 1L), Range.equal(BIGINT, 2L), Range.equal(BIGINT, 3L))))
                    .union(ValueSet.ofRanges(Range.range(BIGINT, 7L, false, 9L, false), Range.range(BIGINT, 11L, false, 13L, false))), false);

    tupleDomain = withColumnDomains(ImmutableMap.<Symbol, Domain>builder().put(C_BIGINT, testDomain).build());
    assertEquals(toPredicate(tupleDomain), or(
            and(lessThan(C_BIGINT, bigintLiteral(4L)), not(in(C_BIGINT, ImmutableList.of(1L, 2L, 3L)))),
            and(greaterThan(C_BIGINT, bigintLiteral(7L)), lessThan(C_BIGINT, bigintLiteral(9L))),
            and(greaterThan(C_BIGINT, bigintLiteral(11L)), lessThan(C_BIGINT, bigintLiteral(13L)))));
}
Example 12
Source File: TestDatabaseShardManager.java    From presto with Apache License 2.0 4 votes vote down vote up 
private static Domain createDomain(Range first, Range... ranges)
{
    return Domain.create(ValueSet.ofRanges(first, ranges), false);
}
Example 13
Source File: TestPulsarSplitManager.java    From pulsar with Apache License 2.0 4 votes vote down vote up 
@Test(dataProvider = "rewriteNamespaceDelimiter", singleThreaded = true)
public void testPublishTimePredicatePushdown(String delimiter) throws Exception {
    updateRewriteNamespaceDelimiterIfNeeded(delimiter);
    TopicName topicName = TOPIC_1;

    setup();
    log.info("!----- topic: %s -----!", topicName);
    PulsarTableHandle pulsarTableHandle = new PulsarTableHandle(pulsarConnectorId.toString(),
            topicName.getNamespace(),
            topicName.getLocalName(),
            topicName.getLocalName());


    Map<ColumnHandle, Domain> domainMap = new HashMap<>();
    Domain domain = Domain.create(ValueSet.ofRanges(Range.range(TIMESTAMP, currentTimeMs + 1L, true,
            currentTimeMs + 50L, true)), false);
    domainMap.put(PulsarInternalColumn.PUBLISH_TIME.getColumnHandle(pulsarConnectorId.toString(), false), domain);
    TupleDomain<ColumnHandle> tupleDomain = TupleDomain.withColumnDomains(domainMap);

    PulsarTableLayoutHandle pulsarTableLayoutHandle = new PulsarTableLayoutHandle(pulsarTableHandle, tupleDomain);

    final ResultCaptor<Collection<PulsarSplit>> resultCaptor = new ResultCaptor<>();
    doAnswer(resultCaptor).when(this.pulsarSplitManager)
            .getSplitsNonPartitionedTopic(anyInt(), any(), any(), any(), any(), any());

    ConnectorSplitSource connectorSplitSource = this.pulsarSplitManager.getSplits(
            mock(ConnectorTransactionHandle.class), mock(ConnectorSession.class),
            pulsarTableLayoutHandle, null);


    verify(this.pulsarSplitManager, times(1))
            .getSplitsNonPartitionedTopic(anyInt(), any(), any(), any(), any(), any());

    int totalSize = 0;
    int initalStart = 1;
    for (PulsarSplit pulsarSplit : resultCaptor.getResult()) {
        assertEquals(pulsarSplit.getConnectorId(), pulsarConnectorId.toString());
        assertEquals(pulsarSplit.getSchemaName(), topicName.getNamespace());
        assertEquals(pulsarSplit.getTableName(), topicName.getLocalName());
        assertEquals(pulsarSplit.getSchema(),
                new String(topicsToSchemas.get(topicName.getSchemaName()).getSchema()));
        assertEquals(pulsarSplit.getSchemaType(), topicsToSchemas.get(topicName.getSchemaName()).getType());
        assertEquals(pulsarSplit.getStartPositionEntryId(), initalStart);
        assertEquals(pulsarSplit.getStartPositionLedgerId(), 0);
        assertEquals(pulsarSplit.getStartPosition(), PositionImpl.get(0, initalStart));
        assertEquals(pulsarSplit.getEndPositionLedgerId(), 0);
        assertEquals(pulsarSplit.getEndPositionEntryId(), initalStart + pulsarSplit.getSplitSize());
        assertEquals(pulsarSplit.getEndPosition(), PositionImpl.get(0, initalStart + pulsarSplit
                .getSplitSize()));

        initalStart += pulsarSplit.getSplitSize();
        totalSize += pulsarSplit.getSplitSize();
    }
    assertEquals(totalSize, 49);

}
Example 14
Source File: TestPulsarSplitManager.java    From pulsar with Apache License 2.0 4 votes vote down vote up 
@Test(dataProvider = "rewriteNamespaceDelimiter", singleThreaded = true)
public void testPublishTimePredicatePushdownPartitionedTopic(String delimiter) throws Exception {
    updateRewriteNamespaceDelimiterIfNeeded(delimiter);
    TopicName topicName = PARTITIONED_TOPIC_1;

    setup();
    log.info("!----- topic: %s -----!", topicName);
    PulsarTableHandle pulsarTableHandle = new PulsarTableHandle(pulsarConnectorId.toString(),
            topicName.getNamespace(),
            topicName.getLocalName(),
            topicName.getLocalName());


    Map<ColumnHandle, Domain> domainMap = new HashMap<>();
    Domain domain = Domain.create(ValueSet.ofRanges(Range.range(TIMESTAMP, currentTimeMs + 1L, true,
            currentTimeMs + 50L, true)), false);
    domainMap.put(PulsarInternalColumn.PUBLISH_TIME.getColumnHandle(pulsarConnectorId.toString(), false), domain);
    TupleDomain<ColumnHandle> tupleDomain = TupleDomain.withColumnDomains(domainMap);

    PulsarTableLayoutHandle pulsarTableLayoutHandle = new PulsarTableLayoutHandle(pulsarTableHandle, tupleDomain);

    final ResultCaptor<Collection<PulsarSplit>> resultCaptor = new ResultCaptor<>();
    doAnswer(resultCaptor).when(this.pulsarSplitManager)
            .getSplitsPartitionedTopic(anyInt(), any(), any(), any(), any(), any());

    ConnectorSplitSource connectorSplitSource = this.pulsarSplitManager.getSplits(
            mock(ConnectorTransactionHandle.class), mock(ConnectorSession.class),
            pulsarTableLayoutHandle, null);


    verify(this.pulsarSplitManager, times(1))
            .getSplitsPartitionedTopic(anyInt(), any(), any(), any(), any(), any());


    int partitions = partitionedTopicsToPartitions.get(topicName.toString());
    for (int i = 0; i < partitions; i++) {
        List<PulsarSplit> splits = getSplitsForPartition(topicName.getPartition(i), resultCaptor.getResult());
        int totalSize = 0;
        int initialStart = 1;
        for (PulsarSplit pulsarSplit : splits) {
            assertEquals(pulsarSplit.getConnectorId(), pulsarConnectorId.toString());
            assertEquals(pulsarSplit.getSchemaName(), topicName.getNamespace());
            assertEquals(pulsarSplit.getTableName(), topicName.getPartition(i).getLocalName());
            assertEquals(pulsarSplit.getSchema(),
                    new String(topicsToSchemas.get(topicName.getSchemaName()).getSchema()));
            assertEquals(pulsarSplit.getSchemaType(), topicsToSchemas.get(topicName.getSchemaName()).getType());
            assertEquals(pulsarSplit.getStartPositionEntryId(), initialStart);
            assertEquals(pulsarSplit.getStartPositionLedgerId(), 0);
            assertEquals(pulsarSplit.getStartPosition(), PositionImpl.get(0, initialStart));
            assertEquals(pulsarSplit.getEndPositionLedgerId(), 0);
            assertEquals(pulsarSplit.getEndPositionEntryId(), initialStart + pulsarSplit.getSplitSize());
            assertEquals(pulsarSplit.getEndPosition(), PositionImpl.get(0, initialStart + pulsarSplit.getSplitSize()));

            initialStart += pulsarSplit.getSplitSize();
            totalSize += pulsarSplit.getSplitSize();
        }

        assertEquals(totalSize, 49);
    }
}