Search in sources :

Example 26 with INTEGER

use of io.trino.spi.type.IntegerType.INTEGER in project trino by trinodb.

the class TestInsert method createLocalQueryRunner.

@Override
protected LocalQueryRunner createLocalQueryRunner() {
    Session.SessionBuilder sessionBuilder = testSessionBuilder().setCatalog("mock").setSchema("schema");
    LocalQueryRunner queryRunner = LocalQueryRunner.create(sessionBuilder.build());
    queryRunner.createCatalog("mock", MockConnectorFactory.builder().withGetTableHandle((session, schemaTableName) -> {
        if (schemaTableName.getTableName().equals("test_table_preferred_partitioning")) {
            return new MockConnectorTableHandle(schemaTableName);
        }
        if (schemaTableName.getTableName().equals("test_table_required_partitioning")) {
            return new MockConnectorTableHandle(schemaTableName);
        }
        return null;
    }).withGetColumns(name -> ImmutableList.of(new ColumnMetadata("column1", INTEGER), new ColumnMetadata("column2", INTEGER))).withGetInsertLayout((session, tableName) -> {
        if (tableName.getTableName().equals("test_table_preferred_partitioning")) {
            return Optional.of(new ConnectorTableLayout(ImmutableList.of("column1")));
        }
        if (tableName.getTableName().equals("test_table_required_partitioning")) {
            return Optional.of(new ConnectorTableLayout(new TpchPartitioningHandle("orders", 10), ImmutableList.of("column1")));
        }
        return Optional.empty();
    }).withGetNewTableLayout((session, tableMetadata) -> {
        if (tableMetadata.getTable().getTableName().equals("new_test_table_preferred_partitioning")) {
            return Optional.of(new ConnectorTableLayout(ImmutableList.of("column1")));
        }
        if (tableMetadata.getTable().getTableName().equals("new_test_table_required_partitioning")) {
            return Optional.of(new ConnectorTableLayout(new TpchPartitioningHandle("orders", 10), ImmutableList.of("column1")));
        }
        if (tableMetadata.getTable().getTableName().equals("new_test_table_unpartitioned")) {
            return Optional.empty();
        }
        return Optional.empty();
    }).build(), ImmutableMap.of());
    return queryRunner;
}
Also used : ConnectorTableLayout(io.trino.spi.connector.ConnectorTableLayout) ColumnMetadata(io.trino.spi.connector.ColumnMetadata) PREFERRED_WRITE_PARTITIONING_MIN_NUMBER_OF_PARTITIONS(io.trino.SystemSessionProperties.PREFERRED_WRITE_PARTITIONING_MIN_NUMBER_OF_PARTITIONS) TableWriterNode(io.trino.sql.planner.plan.TableWriterNode) Test(org.testng.annotations.Test) PlanNode(io.trino.sql.planner.plan.PlanNode) ImmutableList(com.google.common.collect.ImmutableList) MockConnectorFactory(io.trino.connector.MockConnectorFactory) PlanMatchPattern.exchange(io.trino.sql.planner.assertions.PlanMatchPattern.exchange) LocalQueryRunner(io.trino.testing.LocalQueryRunner) LOCAL(io.trino.sql.planner.plan.ExchangeNode.Scope.LOCAL) REPARTITION(io.trino.sql.planner.plan.ExchangeNode.Type.REPARTITION) INTEGER(io.trino.spi.type.IntegerType.INTEGER) BasePlanTest(io.trino.sql.planner.assertions.BasePlanTest) ImmutableSet(com.google.common.collect.ImmutableSet) ImmutableMap(com.google.common.collect.ImmutableMap) PlanMatchPattern.values(io.trino.sql.planner.assertions.PlanMatchPattern.values) StatsProvider(io.trino.cost.StatsProvider) TASK_WRITER_COUNT(io.trino.SystemSessionProperties.TASK_WRITER_COUNT) PlanMatchPattern.node(io.trino.sql.planner.assertions.PlanMatchPattern.node) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) TpchPartitioningHandle(io.trino.plugin.tpch.TpchPartitioningHandle) TestingSession.testSessionBuilder(io.trino.testing.TestingSession.testSessionBuilder) PlanMatchPattern.anyTree(io.trino.sql.planner.assertions.PlanMatchPattern.anyTree) REMOTE(io.trino.sql.planner.plan.ExchangeNode.Scope.REMOTE) MatchResult(io.trino.sql.planner.assertions.MatchResult) USE_PREFERRED_WRITE_PARTITIONING(io.trino.SystemSessionProperties.USE_PREFERRED_WRITE_PARTITIONING) Metadata(io.trino.metadata.Metadata) Matcher(io.trino.sql.planner.assertions.Matcher) Optional(java.util.Optional) SymbolAliases(io.trino.sql.planner.assertions.SymbolAliases) ExchangeNode(io.trino.sql.planner.plan.ExchangeNode) Session(io.trino.Session) ColumnMetadata(io.trino.spi.connector.ColumnMetadata) ConnectorTableLayout(io.trino.spi.connector.ConnectorTableLayout) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) TpchPartitioningHandle(io.trino.plugin.tpch.TpchPartitioningHandle) LocalQueryRunner(io.trino.testing.LocalQueryRunner) Session(io.trino.Session)

Example 27 with INTEGER

use of io.trino.spi.type.IntegerType.INTEGER in project trino by trinodb.

the class BaseRaptorConnectorTest method testBucketNumberHiddenColumn.

@Test
public void testBucketNumberHiddenColumn() {
    assertUpdate("" + "CREATE TABLE test_bucket_number " + "WITH (bucket_count = 50, bucketed_on = ARRAY ['orderkey']) " + "AS SELECT * FROM orders", "SELECT count(*) FROM orders");
    MaterializedResult actualResults = computeActual("SELECT DISTINCT \"$bucket_number\" FROM test_bucket_number");
    assertEquals(actualResults.getTypes(), ImmutableList.of(INTEGER));
    Set<Object> actual = actualResults.getMaterializedRows().stream().map(row -> row.getField(0)).collect(toSet());
    assertEquals(actual, IntStream.range(0, 50).boxed().collect(toSet()));
}
Also used : SkipException(org.testng.SkipException) IntStream(java.util.stream.IntStream) Assertions.assertInstanceOf(io.airlift.testing.Assertions.assertInstanceOf) MaterializedResult(io.trino.testing.MaterializedResult) LocalDateTime(java.time.LocalDateTime) Assertions.assertGreaterThan(io.airlift.testing.Assertions.assertGreaterThan) BOOLEAN(io.trino.spi.type.BooleanType.BOOLEAN) Assert.assertEquals(org.testng.Assert.assertEquals) Test(org.testng.annotations.Test) TestTable(io.trino.testing.sql.TestTable) SHARD_UUID_COLUMN_TYPE(io.trino.plugin.raptor.legacy.RaptorColumnHandle.SHARD_UUID_COLUMN_TYPE) VARCHAR(io.trino.spi.type.VarcharType.VARCHAR) HashMultimap(com.google.common.collect.HashMultimap) ImmutableList(com.google.common.collect.ImmutableList) Assertions.assertThatThrownBy(org.assertj.core.api.Assertions.assertThatThrownBy) Arrays.asList(java.util.Arrays.asList) Map(java.util.Map) TestingConnectorBehavior(io.trino.testing.TestingConnectorBehavior) Assertions.assertLessThan(io.airlift.testing.Assertions.assertLessThan) Assertions.assertGreaterThanOrEqual(io.airlift.testing.Assertions.assertGreaterThanOrEqual) MaterializedRow(io.trino.testing.MaterializedRow) INTEGER(io.trino.spi.type.IntegerType.INTEGER) Assert.assertFalse(org.testng.Assert.assertFalse) TestTable.randomTableSuffix(io.trino.testing.sql.TestTable.randomTableSuffix) Collectors.toSet(java.util.stream.Collectors.toSet) Flaky(io.trino.testng.services.Flaky) Assert.assertNotEquals(org.testng.Assert.assertNotEquals) Language(org.intellij.lang.annotations.Language) Collection(java.util.Collection) Set(java.util.Set) ArrayType(io.trino.spi.type.ArrayType) Iterables.getOnlyElement(com.google.common.collect.Iterables.getOnlyElement) UUID(java.util.UUID) Assert.assertNotNull(org.testng.Assert.assertNotNull) SetMultimap(com.google.common.collect.SetMultimap) String.format(java.lang.String.format) BaseConnectorTest(io.trino.testing.BaseConnectorTest) ImmutableMap.toImmutableMap(com.google.common.collect.ImmutableMap.toImmutableMap) BIGINT(io.trino.spi.type.BigintType.BIGINT) LocalDate(java.time.LocalDate) StringJoiner(java.util.StringJoiner) Function.identity(java.util.function.Function.identity) Optional(java.util.Optional) DATE(io.trino.spi.type.DateType.DATE) MaterializedResult(io.trino.testing.MaterializedResult) Test(org.testng.annotations.Test) BaseConnectorTest(io.trino.testing.BaseConnectorTest)

Example 28 with INTEGER

use of io.trino.spi.type.IntegerType.INTEGER in project trino by trinodb.

the class OracleClient method toColumnMapping.

@Override
public Optional<ColumnMapping> toColumnMapping(ConnectorSession session, Connection connection, JdbcTypeHandle typeHandle) {
    String jdbcTypeName = typeHandle.getJdbcTypeName().orElseThrow(() -> new TrinoException(JDBC_ERROR, "Type name is missing: " + typeHandle));
    Optional<ColumnMapping> mappingToVarchar = getForcedMappingToVarchar(typeHandle);
    if (mappingToVarchar.isPresent()) {
        return mappingToVarchar;
    }
    if (jdbcTypeName.equalsIgnoreCase("date")) {
        return Optional.of(ColumnMapping.longMapping(TIMESTAMP_SECONDS, oracleTimestampReadFunction(), trinoTimestampToOracleDateWriteFunction(), FULL_PUSHDOWN));
    }
    switch(typeHandle.getJdbcType()) {
        case Types.SMALLINT:
            return Optional.of(ColumnMapping.longMapping(SMALLINT, ResultSet::getShort, smallintWriteFunction(), FULL_PUSHDOWN));
        case OracleTypes.BINARY_FLOAT:
            return Optional.of(ColumnMapping.longMapping(REAL, (resultSet, columnIndex) -> floatToRawIntBits(resultSet.getFloat(columnIndex)), oracleRealWriteFunction(), FULL_PUSHDOWN));
        case OracleTypes.BINARY_DOUBLE:
        case OracleTypes.FLOAT:
            return Optional.of(ColumnMapping.doubleMapping(DOUBLE, ResultSet::getDouble, oracleDoubleWriteFunction(), FULL_PUSHDOWN));
        case OracleTypes.NUMBER:
            int actualPrecision = typeHandle.getRequiredColumnSize();
            int decimalDigits = typeHandle.getRequiredDecimalDigits();
            // Map negative scale to decimal(p+s, 0).
            int precision = actualPrecision + max(-decimalDigits, 0);
            int scale = max(decimalDigits, 0);
            Optional<Integer> numberDefaultScale = getNumberDefaultScale(session);
            RoundingMode roundingMode = getNumberRoundingMode(session);
            if (precision < scale) {
                if (roundingMode == RoundingMode.UNNECESSARY) {
                    break;
                }
                scale = min(Decimals.MAX_PRECISION, scale);
                precision = scale;
            } else if (numberDefaultScale.isPresent() && precision == PRECISION_OF_UNSPECIFIED_NUMBER) {
                precision = Decimals.MAX_PRECISION;
                scale = numberDefaultScale.get();
            } else if (precision > Decimals.MAX_PRECISION || actualPrecision <= 0) {
                break;
            }
            DecimalType decimalType = createDecimalType(precision, scale);
            // JDBC driver can return BigDecimal with lower scale than column's scale when there are trailing zeroes
            if (decimalType.isShort()) {
                return Optional.of(ColumnMapping.longMapping(decimalType, shortDecimalReadFunction(decimalType, roundingMode), shortDecimalWriteFunction(decimalType), FULL_PUSHDOWN));
            }
            return Optional.of(ColumnMapping.objectMapping(decimalType, longDecimalReadFunction(decimalType, roundingMode), longDecimalWriteFunction(decimalType), FULL_PUSHDOWN));
        case OracleTypes.CHAR:
        case OracleTypes.NCHAR:
            CharType charType = createCharType(typeHandle.getRequiredColumnSize());
            return Optional.of(ColumnMapping.sliceMapping(charType, charReadFunction(charType), oracleCharWriteFunction(), FULL_PUSHDOWN));
        case OracleTypes.VARCHAR:
        case OracleTypes.NVARCHAR:
            return Optional.of(ColumnMapping.sliceMapping(createVarcharType(typeHandle.getRequiredColumnSize()), (varcharResultSet, varcharColumnIndex) -> utf8Slice(varcharResultSet.getString(varcharColumnIndex)), varcharWriteFunction(), FULL_PUSHDOWN));
        case OracleTypes.CLOB:
        case OracleTypes.NCLOB:
            return Optional.of(ColumnMapping.sliceMapping(createUnboundedVarcharType(), (resultSet, columnIndex) -> utf8Slice(resultSet.getString(columnIndex)), varcharWriteFunction(), DISABLE_PUSHDOWN));
        // Oracle's RAW(n)
        case OracleTypes.VARBINARY:
        case OracleTypes.BLOB:
            return Optional.of(ColumnMapping.sliceMapping(VARBINARY, (resultSet, columnIndex) -> wrappedBuffer(resultSet.getBytes(columnIndex)), varbinaryWriteFunction(), DISABLE_PUSHDOWN));
        case OracleTypes.TIMESTAMP:
            return Optional.of(ColumnMapping.longMapping(TIMESTAMP_MILLIS, oracleTimestampReadFunction(), trinoTimestampToOracleTimestampWriteFunction(), FULL_PUSHDOWN));
        case OracleTypes.TIMESTAMPTZ:
            return Optional.of(oracleTimestampWithTimeZoneColumnMapping());
    }
    if (getUnsupportedTypeHandling(session) == CONVERT_TO_VARCHAR) {
        return mapToUnboundedVarchar(typeHandle);
    }
    return Optional.empty();
}
Also used : ZonedDateTime(java.time.ZonedDateTime) StandardColumnMappings.bigintWriteFunction(io.trino.plugin.jdbc.StandardColumnMappings.bigintWriteFunction) NANOSECONDS_PER_MICROSECOND(io.trino.spi.type.Timestamps.NANOSECONDS_PER_MICROSECOND) CharType.createCharType(io.trino.spi.type.CharType.createCharType) SecureRandom(java.security.SecureRandom) Slices.wrappedBuffer(io.airlift.slice.Slices.wrappedBuffer) NOT_SUPPORTED(io.trino.spi.StandardErrorCode.NOT_SUPPORTED) ResultSet(java.sql.ResultSet) Map(java.util.Map) Slices.utf8Slice(io.airlift.slice.Slices.utf8Slice) ZoneOffset(java.time.ZoneOffset) ENGLISH(java.util.Locale.ENGLISH) SMALLINT(io.trino.spi.type.SmallintType.SMALLINT) TypeHandlingJdbcSessionProperties.getUnsupportedTypeHandling(io.trino.plugin.jdbc.TypeHandlingJdbcSessionProperties.getUnsupportedTypeHandling) DateTimeEncoding.packDateTimeWithZone(io.trino.spi.type.DateTimeEncoding.packDateTimeWithZone) LongWriteFunction(io.trino.plugin.jdbc.LongWriteFunction) Set(java.util.Set) PreparedStatement(java.sql.PreparedStatement) SchemaTableName(io.trino.spi.connector.SchemaTableName) ZoneId(java.time.ZoneId) LongReadFunction(io.trino.plugin.jdbc.LongReadFunction) StandardColumnMappings.charReadFunction(io.trino.plugin.jdbc.StandardColumnMappings.charReadFunction) StandardColumnMappings.smallintWriteFunction(io.trino.plugin.jdbc.StandardColumnMappings.smallintWriteFunction) StandardColumnMappings.longDecimalWriteFunction(io.trino.plugin.jdbc.StandardColumnMappings.longDecimalWriteFunction) ConnectionFactory(io.trino.plugin.jdbc.ConnectionFactory) CONVERT_TO_VARCHAR(io.trino.plugin.jdbc.UnsupportedTypeHandling.CONVERT_TO_VARCHAR) DoubleWriteFunction(io.trino.plugin.jdbc.DoubleWriteFunction) DateTimeEncoding.unpackMillisUtc(io.trino.spi.type.DateTimeEncoding.unpackMillisUtc) JdbcTableHandle(io.trino.plugin.jdbc.JdbcTableHandle) DATE(io.trino.spi.type.DateType.DATE) REAL(io.trino.spi.type.RealType.REAL) JoinCondition(io.trino.spi.connector.JoinCondition) TIMESTAMP_MILLIS(io.trino.spi.type.TimestampType.TIMESTAMP_MILLIS) LocalDateTime(java.time.LocalDateTime) BOOLEAN(io.trino.spi.type.BooleanType.BOOLEAN) Float.floatToRawIntBits(java.lang.Float.floatToRawIntBits) SQLException(java.sql.SQLException) TIMESTAMP_TZ_MILLIS(io.trino.spi.type.TimestampWithTimeZoneType.TIMESTAMP_TZ_MILLIS) FULL_PUSHDOWN(io.trino.plugin.jdbc.PredicatePushdownController.FULL_PUSHDOWN) StandardColumnMappings.charWriteFunction(io.trino.plugin.jdbc.StandardColumnMappings.charWriteFunction) VARBINARY(io.trino.spi.type.VarbinaryType.VARBINARY) Math.floorDiv(java.lang.Math.floorDiv) TIMESTAMP_SECONDS(io.trino.spi.type.TimestampType.TIMESTAMP_SECONDS) DecimalType.createDecimalType(io.trino.spi.type.DecimalType.createDecimalType) QueryBuilder(io.trino.plugin.jdbc.QueryBuilder) ConnectorSession(io.trino.spi.connector.ConnectorSession) DOUBLE(io.trino.spi.type.DoubleType.DOUBLE) IdentifierMapping(io.trino.plugin.jdbc.mapping.IdentifierMapping) CharType(io.trino.spi.type.CharType) TINYINT(io.trino.spi.type.TinyintType.TINYINT) StandardColumnMappings.shortDecimalWriteFunction(io.trino.plugin.jdbc.StandardColumnMappings.shortDecimalWriteFunction) StandardColumnMappings.varbinaryWriteFunction(io.trino.plugin.jdbc.StandardColumnMappings.varbinaryWriteFunction) VarcharType.createVarcharType(io.trino.spi.type.VarcharType.createVarcharType) WriteMapping(io.trino.plugin.jdbc.WriteMapping) JDBC_ERROR(io.trino.plugin.jdbc.JdbcErrorCode.JDBC_ERROR) BaseJdbcConfig(io.trino.plugin.jdbc.BaseJdbcConfig) Connection(java.sql.Connection) MICROSECONDS_PER_MILLISECOND(io.trino.spi.type.Timestamps.MICROSECONDS_PER_MILLISECOND) OracleSessionProperties.getNumberDefaultScale(io.trino.plugin.oracle.OracleSessionProperties.getNumberDefaultScale) Math.abs(java.lang.Math.abs) DateTimeEncoding.unpackZoneKey(io.trino.spi.type.DateTimeEncoding.unpackZoneKey) ColumnMapping(io.trino.plugin.jdbc.ColumnMapping) INTEGER(io.trino.spi.type.IntegerType.INTEGER) RoundingMode(java.math.RoundingMode) ImmutableSet(com.google.common.collect.ImmutableSet) ImmutableMap(com.google.common.collect.ImmutableMap) MICROSECONDS_PER_SECOND(io.trino.spi.type.Timestamps.MICROSECONDS_PER_SECOND) TrinoException(io.trino.spi.TrinoException) Math.min(java.lang.Math.min) Instant(java.time.Instant) String.format(java.lang.String.format) List(java.util.List) JdbcTypeHandle(io.trino.plugin.jdbc.JdbcTypeHandle) StandardColumnMappings.longDecimalReadFunction(io.trino.plugin.jdbc.StandardColumnMappings.longDecimalReadFunction) BIGINT(io.trino.spi.type.BigintType.BIGINT) Decimals(io.trino.spi.type.Decimals) Optional(java.util.Optional) Math.max(java.lang.Math.max) DecimalType(io.trino.spi.type.DecimalType) OracleTypes(oracle.jdbc.OracleTypes) MAX_RADIX(java.lang.Character.MAX_RADIX) Types(java.sql.Types) SliceWriteFunction(io.trino.plugin.jdbc.SliceWriteFunction) Math.floorMod(java.lang.Math.floorMod) StandardColumnMappings.varcharWriteFunction(io.trino.plugin.jdbc.StandardColumnMappings.varcharWriteFunction) StandardColumnMappings.tinyintWriteFunction(io.trino.plugin.jdbc.StandardColumnMappings.tinyintWriteFunction) Type(io.trino.spi.type.Type) OracleSessionProperties.getNumberRoundingMode(io.trino.plugin.oracle.OracleSessionProperties.getNumberRoundingMode) VarcharType.createUnboundedVarcharType(io.trino.spi.type.VarcharType.createUnboundedVarcharType) Float.intBitsToFloat(java.lang.Float.intBitsToFloat) Inject(javax.inject.Inject) VarcharType(io.trino.spi.type.VarcharType) OraclePreparedStatement(oracle.jdbc.OraclePreparedStatement) ImmutableList(com.google.common.collect.ImmutableList) Verify.verify(com.google.common.base.Verify.verify) Objects.requireNonNull(java.util.Objects.requireNonNull) DAYS(java.util.concurrent.TimeUnit.DAYS) BaseJdbcClient(io.trino.plugin.jdbc.BaseJdbcClient) Math.toIntExact(java.lang.Math.toIntExact) JdbcJoinCondition(io.trino.plugin.jdbc.JdbcJoinCondition) JdbcColumnHandle(io.trino.plugin.jdbc.JdbcColumnHandle) StandardColumnMappings.integerWriteFunction(io.trino.plugin.jdbc.StandardColumnMappings.integerWriteFunction) DISABLE_PUSHDOWN(io.trino.plugin.jdbc.PredicatePushdownController.DISABLE_PUSHDOWN) DateTimeFormatter(java.time.format.DateTimeFormatter) StandardColumnMappings.shortDecimalReadFunction(io.trino.plugin.jdbc.StandardColumnMappings.shortDecimalReadFunction) RoundingMode(java.math.RoundingMode) OracleSessionProperties.getNumberRoundingMode(io.trino.plugin.oracle.OracleSessionProperties.getNumberRoundingMode) TrinoException(io.trino.spi.TrinoException) DecimalType.createDecimalType(io.trino.spi.type.DecimalType.createDecimalType) DecimalType(io.trino.spi.type.DecimalType) CharType.createCharType(io.trino.spi.type.CharType.createCharType) CharType(io.trino.spi.type.CharType) ColumnMapping(io.trino.plugin.jdbc.ColumnMapping)

Example 29 with INTEGER

use of io.trino.spi.type.IntegerType.INTEGER in project trino by trinodb.

the class TestThriftIndexPageSource method testGetNextPageTwoConcurrentRequests.

@Test
public void testGetNextPageTwoConcurrentRequests() throws Exception {
    final int splits = 3;
    final int lookupRequestsConcurrency = 2;
    final int rowsPerSplit = 1;
    List<SettableFuture<TrinoThriftPageResult>> futures = IntStream.range(0, splits).mapToObj(i -> SettableFuture.<TrinoThriftPageResult>create()).collect(toImmutableList());
    List<CountDownLatch> signals = IntStream.range(0, splits).mapToObj(i -> new CountDownLatch(1)).collect(toImmutableList());
    TestingThriftService client = new TestingThriftService(rowsPerSplit, false, false) {

        @Override
        public ListenableFuture<TrinoThriftPageResult> getRows(TrinoThriftId splitId, List<String> columns, long maxBytes, TrinoThriftNullableToken nextToken) {
            int key = Ints.fromByteArray(splitId.getId());
            signals.get(key).countDown();
            return futures.get(key);
        }
    };
    ThriftConnectorStats stats = new ThriftConnectorStats();
    long pageSizeReceived = 0;
    ThriftIndexPageSource pageSource = new ThriftIndexPageSource((context, headers) -> client, ImmutableMap.of(), stats, new ThriftIndexHandle(new SchemaTableName("default", "table1"), TupleDomain.all()), ImmutableList.of(column("a", INTEGER)), ImmutableList.of(column("b", INTEGER)), new InMemoryRecordSet(ImmutableList.of(INTEGER), generateKeys(0, splits)), MAX_BYTES_PER_RESPONSE, lookupRequestsConcurrency);
    assertNull(pageSource.getNextPage());
    assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), 0);
    signals.get(0).await(1, SECONDS);
    signals.get(1).await(1, SECONDS);
    signals.get(2).await(1, SECONDS);
    assertEquals(signals.get(0).getCount(), 0, "first request wasn't sent");
    assertEquals(signals.get(1).getCount(), 0, "second request wasn't sent");
    assertEquals(signals.get(2).getCount(), 1, "third request shouldn't be sent");
    // at this point first two requests were sent
    assertFalse(pageSource.isFinished());
    assertNull(pageSource.getNextPage());
    assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), 0);
    // completing the second request
    futures.get(1).set(pageResult(20, null));
    Page page = pageSource.getNextPage();
    pageSizeReceived += page.getSizeInBytes();
    assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), pageSizeReceived);
    assertNotNull(page);
    assertEquals(page.getPositionCount(), 1);
    assertEquals(page.getBlock(0).getInt(0, 0), 20);
    // not complete yet
    assertFalse(pageSource.isFinished());
    // once one of the requests completes the next one should be sent
    signals.get(2).await(1, SECONDS);
    assertEquals(signals.get(2).getCount(), 0, "third request wasn't sent");
    // completing the first request
    futures.get(0).set(pageResult(10, null));
    page = pageSource.getNextPage();
    assertNotNull(page);
    pageSizeReceived += page.getSizeInBytes();
    assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), pageSizeReceived);
    assertEquals(page.getPositionCount(), 1);
    assertEquals(page.getBlock(0).getInt(0, 0), 10);
    // still not complete
    assertFalse(pageSource.isFinished());
    // completing the third request
    futures.get(2).set(pageResult(30, null));
    page = pageSource.getNextPage();
    assertNotNull(page);
    pageSizeReceived += page.getSizeInBytes();
    assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), pageSizeReceived);
    assertEquals(page.getPositionCount(), 1);
    assertEquals(page.getBlock(0).getInt(0, 0), 30);
    // finished now
    assertTrue(pageSource.isFinished());
    // after completion
    assertNull(pageSource.getNextPage());
    pageSource.close();
}
Also used : SettableFuture(com.google.common.util.concurrent.SettableFuture) IntStream(java.util.stream.IntStream) Collections.shuffle(java.util.Collections.shuffle) TrinoThriftBlock.integerData(io.trino.plugin.thrift.api.TrinoThriftBlock.integerData) ListenableFuture(com.google.common.util.concurrent.ListenableFuture) Assert.assertNull(org.testng.Assert.assertNull) TrinoThriftTupleDomain(io.trino.plugin.thrift.api.TrinoThriftTupleDomain) Type(io.trino.spi.type.Type) Page(io.trino.spi.Page) Assert.assertEquals(org.testng.Assert.assertEquals) Test(org.testng.annotations.Test) CompletableFuture(java.util.concurrent.CompletableFuture) SettableFuture(com.google.common.util.concurrent.SettableFuture) ArrayList(java.util.ArrayList) ImmutableList(com.google.common.collect.ImmutableList) TrinoThriftId(io.trino.plugin.thrift.api.TrinoThriftId) Block(io.trino.spi.block.Block) TrinoThriftService(io.trino.plugin.thrift.api.TrinoThriftService) TrinoThriftNullableToken(io.trino.plugin.thrift.api.TrinoThriftNullableToken) TrinoThriftSplit(io.trino.plugin.thrift.api.TrinoThriftSplit) INTEGER(io.trino.spi.type.IntegerType.INTEGER) Assert.assertFalse(org.testng.Assert.assertFalse) Futures.immediateFuture(com.google.common.util.concurrent.Futures.immediateFuture) TrinoThriftPageResult(io.trino.plugin.thrift.api.TrinoThriftPageResult) ImmutableMap(com.google.common.collect.ImmutableMap) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) TrinoThriftInteger(io.trino.plugin.thrift.api.datatypes.TrinoThriftInteger) TrinoThriftSplitBatch(io.trino.plugin.thrift.api.TrinoThriftSplitBatch) TrinoThriftNullableTableMetadata(io.trino.plugin.thrift.api.TrinoThriftNullableTableMetadata) Assert.assertNotNull(org.testng.Assert.assertNotNull) TupleDomain(io.trino.spi.predicate.TupleDomain) Ints(com.google.common.primitives.Ints) InMemoryRecordSet(io.trino.spi.connector.InMemoryRecordSet) SchemaTableName(io.trino.spi.connector.SchemaTableName) CountDownLatch(java.util.concurrent.CountDownLatch) List(java.util.List) TrinoThriftSchemaTableName(io.trino.plugin.thrift.api.TrinoThriftSchemaTableName) TrinoThriftNullableColumnSet(io.trino.plugin.thrift.api.TrinoThriftNullableColumnSet) Assert.assertTrue(org.testng.Assert.assertTrue) TrinoThriftNullableSchemaName(io.trino.plugin.thrift.api.TrinoThriftNullableSchemaName) Collections(java.util.Collections) SECONDS(java.util.concurrent.TimeUnit.SECONDS) Page(io.trino.spi.Page) CountDownLatch(java.util.concurrent.CountDownLatch) TrinoThriftId(io.trino.plugin.thrift.api.TrinoThriftId) SchemaTableName(io.trino.spi.connector.SchemaTableName) TrinoThriftSchemaTableName(io.trino.plugin.thrift.api.TrinoThriftSchemaTableName) InMemoryRecordSet(io.trino.spi.connector.InMemoryRecordSet) TrinoThriftPageResult(io.trino.plugin.thrift.api.TrinoThriftPageResult) TrinoThriftNullableToken(io.trino.plugin.thrift.api.TrinoThriftNullableToken) ArrayList(java.util.ArrayList) ImmutableList(com.google.common.collect.ImmutableList) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) List(java.util.List) Test(org.testng.annotations.Test)

Aggregations

INTEGER (io.trino.spi.type.IntegerType.INTEGER)29 ImmutableList (com.google.common.collect.ImmutableList)27 Optional (java.util.Optional)24 Type (io.trino.spi.type.Type)22 List (java.util.List)20 ImmutableMap (com.google.common.collect.ImmutableMap)19 BIGINT (io.trino.spi.type.BigintType.BIGINT)19 DOUBLE (io.trino.spi.type.DoubleType.DOUBLE)16 Map (java.util.Map)16 Test (org.testng.annotations.Test)16 String.format (java.lang.String.format)14 ImmutableSet (com.google.common.collect.ImmutableSet)12 VARCHAR (io.trino.spi.type.VarcharType.VARCHAR)11 ImmutableList.toImmutableList (com.google.common.collect.ImmutableList.toImmutableList)10 Page (io.trino.spi.Page)10 Slice (io.airlift.slice.Slice)9 Assert.assertEquals (org.testng.Assert.assertEquals)9 Assert.assertFalse (org.testng.Assert.assertFalse)9 RowType (io.trino.spi.type.RowType)8 Session (io.trino.Session)7