Examples with ConnectorPageSource io.trino.spi.connector.ConnectorPageSource used on opensource projects

Example 41 with ConnectorPageSource

use of io.trino.spi.connector.ConnectorPageSource in project trino by trinodb.

the class BenchmarkProjectionPushdownHive method readPages.

@Benchmark
public List<Page> readPages(BenchmarkContext context) {
    List<Page> pages = new ArrayList<>(100);
    ConnectorPageSource pageSource = context.createPageSource();
    while (!pageSource.isFinished()) {
        Page page = pageSource.getNextPage();
        if (page != null) {
            pages.add(page.getLoadedPage());
        }
    }
    return pages;
}

Example 42 with ConnectorPageSource

use of io.trino.spi.connector.ConnectorPageSource in project trino by trinodb.

the class AbstractTestHive method testGetRecordsUnpartitioned.

@Test
public void testGetRecordsUnpartitioned() throws Exception {
    try (Transaction transaction = newTransaction()) {
        ConnectorMetadata metadata = transaction.getMetadata();
        ConnectorSession session = newSession();
        metadata.beginQuery(session);
        ConnectorTableHandle tableHandle = getTableHandle(metadata, tableUnpartitioned);
        List<ColumnHandle> columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values());
        Map<String, Integer> columnIndex = indexColumns(columnHandles);
        List<ConnectorSplit> splits = getAllSplits(tableHandle, transaction, session);
        assertThat(splits).hasSameSizeAs(tableUnpartitionedPartitions);
        for (ConnectorSplit split : splits) {
            HiveSplit hiveSplit = (HiveSplit) split;
            assertEquals(hiveSplit.getPartitionKeys(), ImmutableList.of());
            long rowNumber = 0;
            try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, split, tableHandle, columnHandles, DynamicFilter.EMPTY)) {
                assertPageSourceType(pageSource, TEXTFILE);
                MaterializedResult result = materializeSourceDataStream(session, pageSource, getTypes(columnHandles));
                for (MaterializedRow row : result) {
                    rowNumber++;
                    if (rowNumber % 19 == 0) {
                        assertNull(row.getField(columnIndex.get("t_string")));
                    } else if (rowNumber % 19 == 1) {
                        assertEquals(row.getField(columnIndex.get("t_string")), "");
                    } else {
                        assertEquals(row.getField(columnIndex.get("t_string")), "unpartitioned");
                    }
                    assertEquals(row.getField(columnIndex.get("t_tinyint")), (byte) (1 + rowNumber));
                }
            }
            assertEquals(rowNumber, 100);
        }
    }
}

Example 43 with ConnectorPageSource

use of io.trino.spi.connector.ConnectorPageSource in project trino by trinodb.

the class AbstractTestHive method testGetPartialRecords.

@Test
public void testGetPartialRecords() throws Exception {
    try (Transaction transaction = newTransaction()) {
        ConnectorMetadata metadata = transaction.getMetadata();
        ConnectorSession session = newSession();
        metadata.beginQuery(session);
        ConnectorTableHandle tableHandle = getTableHandle(metadata, tablePartitionFormat);
        List<ColumnHandle> columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values());
        Map<String, Integer> columnIndex = indexColumns(columnHandles);
        List<ConnectorSplit> splits = getAllSplits(tableHandle, transaction, session);
        assertEquals(splits.size(), tablePartitionFormatPartitions.size());
        for (ConnectorSplit split : splits) {
            HiveSplit hiveSplit = (HiveSplit) split;
            List<HivePartitionKey> partitionKeys = hiveSplit.getPartitionKeys();
            String ds = partitionKeys.get(0).getValue();
            String fileFormat = partitionKeys.get(1).getValue();
            HiveStorageFormat fileType = HiveStorageFormat.valueOf(fileFormat.toUpperCase(ENGLISH));
            int dummyPartition = Integer.parseInt(partitionKeys.get(2).getValue());
            long rowNumber = 0;
            try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, hiveSplit, tableHandle, columnHandles, DynamicFilter.EMPTY)) {
                assertPageSourceType(pageSource, fileType);
                MaterializedResult result = materializeSourceDataStream(session, pageSource, getTypes(columnHandles));
                for (MaterializedRow row : result) {
                    rowNumber++;
                    assertEquals(row.getField(columnIndex.get("t_double")), 6.2 + rowNumber);
                    assertEquals(row.getField(columnIndex.get("ds")), ds);
                    assertEquals(row.getField(columnIndex.get("file_format")), fileFormat);
                    assertEquals(row.getField(columnIndex.get("dummy")), dummyPartition);
                }
            }
            assertEquals(rowNumber, 100);
        }
    }
}

Example 44 with ConnectorPageSource

use of io.trino.spi.connector.ConnectorPageSource in project trino by trinodb.

the class AbstractTestHive method readTable.

protected MaterializedResult readTable(Transaction transaction, ConnectorTableHandle tableHandle, List<ColumnHandle> columnHandles, ConnectorSession session, TupleDomain<ColumnHandle> tupleDomain, OptionalInt expectedSplitCount, Optional<HiveStorageFormat> expectedStorageFormat) throws Exception {
    tableHandle = applyFilter(transaction.getMetadata(), tableHandle, new Constraint(tupleDomain));
    List<ConnectorSplit> splits = getAllSplits(getSplits(splitManager, transaction, session, tableHandle));
    if (expectedSplitCount.isPresent()) {
        assertEquals(splits.size(), expectedSplitCount.getAsInt());
    }
    ImmutableList.Builder<MaterializedRow> allRows = ImmutableList.builder();
    for (ConnectorSplit split : splits) {
        try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, split, tableHandle, columnHandles, DynamicFilter.EMPTY)) {
            expectedStorageFormat.ifPresent(format -> assertPageSourceType(pageSource, format));
            MaterializedResult result = materializeSourceDataStream(session, pageSource, getTypes(columnHandles));
            allRows.addAll(result.getMaterializedRows());
        }
    }
    return new MaterializedResult(allRows.build(), getTypes(columnHandles));
}

Example 45 with ConnectorPageSource

use of io.trino.spi.connector.ConnectorPageSource in project trino by trinodb.

the class AbstractTestHiveFileSystem method readTable.

protected MaterializedResult readTable(SchemaTableName tableName) throws IOException {
    try (Transaction transaction = newTransaction()) {
        ConnectorMetadata metadata = transaction.getMetadata();
        ConnectorSession session = newSession();
        ConnectorTableHandle table = getTableHandle(metadata, tableName);
        List<ColumnHandle> columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, table).values());
        metadata.beginQuery(session);
        ConnectorSplitSource splitSource = getSplits(splitManager, transaction, session, table);
        List<Type> allTypes = getTypes(columnHandles);
        List<Type> dataTypes = getTypes(columnHandles.stream().filter(columnHandle -> !((HiveColumnHandle) columnHandle).isHidden()).collect(toImmutableList()));
        MaterializedResult.Builder result = MaterializedResult.resultBuilder(session, dataTypes);
        List<ConnectorSplit> splits = getAllSplits(splitSource);
        for (ConnectorSplit split : splits) {
            try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, split, table, columnHandles, DynamicFilter.EMPTY)) {
                MaterializedResult pageSourceResult = materializeSourceDataStream(session, pageSource, allTypes);
                for (MaterializedRow row : pageSourceResult.getMaterializedRows()) {
                    Object[] dataValues = IntStream.range(0, row.getFieldCount()).filter(channel -> !((HiveColumnHandle) columnHandles.get(channel)).isHidden()).mapToObj(row::getField).toArray();
                    result.row(dataValues);
                }
            }
        }
        metadata.cleanupQuery(session);
        return result.build();
    }
}