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

Example 1 with RecordPageSource

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

the class SystemPageSourceProvider method createPageSource.

@Override
public ConnectorPageSource createPageSource(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<ColumnHandle> columns, DynamicFilter dynamicFilter) {
    requireNonNull(columns, "columns is null");
    SystemTransactionHandle systemTransaction = (SystemTransactionHandle) transaction;
    SystemSplit systemSplit = (SystemSplit) split;
    SchemaTableName tableName = ((SystemTableHandle) table).getSchemaTableName();
    SystemTable systemTable = tables.getSystemTable(session, tableName).orElseThrow(() -> new TrinoException(NOT_FOUND, format("Table '%s' not found", tableName)));
    List<ColumnMetadata> tableColumns = systemTable.getTableMetadata().getColumns();
    Map<String, Integer> columnsByName = new HashMap<>();
    for (int i = 0; i < tableColumns.size(); i++) {
        ColumnMetadata column = tableColumns.get(i);
        if (columnsByName.put(column.getName(), i) != null) {
            throw new TrinoException(GENERIC_INTERNAL_ERROR, "Duplicate column name: " + column.getName());
        }
    }
    ImmutableList.Builder<Integer> userToSystemFieldIndex = ImmutableList.builder();
    for (ColumnHandle column : columns) {
        String columnName = ((SystemColumnHandle) column).getColumnName();
        Integer index = columnsByName.get(columnName);
        if (index == null) {
            throw new TrinoException(GENERIC_INTERNAL_ERROR, format("Column does not exist: %s.%s", tableName, columnName));
        }
        userToSystemFieldIndex.add(index);
    }
    TupleDomain<ColumnHandle> constraint = systemSplit.getConstraint();
    if (constraint.isNone()) {
        return new EmptyPageSource();
    }
    TupleDomain<Integer> newConstraint = systemSplit.getConstraint().transformKeys(columnHandle -> columnsByName.get(((SystemColumnHandle) columnHandle).getColumnName()));
    try {
        return new MappedPageSource(systemTable.pageSource(systemTransaction.getConnectorTransactionHandle(), session, newConstraint), userToSystemFieldIndex.build());
    } catch (UnsupportedOperationException e) {
        return new RecordPageSource(new MappedRecordSet(toRecordSet(systemTransaction.getConnectorTransactionHandle(), systemTable, session, newConstraint), userToSystemFieldIndex.build()));
    }
}

Example 2 with RecordPageSource

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

the class AbstractTestHive method assertPageSourceType.

protected static void assertPageSourceType(ConnectorPageSource pageSource, HiveStorageFormat hiveStorageFormat) {
    if (pageSource instanceof RecordPageSource) {
        RecordCursor hiveRecordCursor = ((RecordPageSource) pageSource).getCursor();
        hiveRecordCursor = ((HiveRecordCursor) hiveRecordCursor).getRegularColumnRecordCursor();
        if (hiveRecordCursor instanceof HiveBucketValidationRecordCursor) {
            hiveRecordCursor = ((HiveBucketValidationRecordCursor) hiveRecordCursor).delegate();
        }
        if (hiveRecordCursor instanceof HiveCoercionRecordCursor) {
            hiveRecordCursor = ((HiveCoercionRecordCursor) hiveRecordCursor).getRegularColumnRecordCursor();
        }
        assertInstanceOf(hiveRecordCursor, recordCursorType(), hiveStorageFormat.name());
    } else {
        assertInstanceOf(((HivePageSource) pageSource).getPageSource(), pageSourceType(hiveStorageFormat), hiveStorageFormat.name());
    }
}

Example 3 with RecordPageSource

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

the class TestHiveFileFormats method testCursorProvider.

private void testCursorProvider(HiveRecordCursorProvider cursorProvider, FileSplit split, Properties splitProperties, List<TestColumn> testReadColumns, ConnectorSession session, long fileSize, int rowCount) {
    ConnectorPageSource pageSource = createPageSourceFromCursorProvider(cursorProvider, split, splitProperties, fileSize, testReadColumns, session);
    RecordCursor cursor = ((RecordPageSource) pageSource).getCursor();
    checkCursor(cursor, testReadColumns, rowCount);
}

Example 4 with RecordPageSource

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

the class TestScanFilterAndProjectOperator method testRecordCursorSource.

@Test
public void testRecordCursorSource() {
    Page input = SequencePageBuilder.createSequencePage(ImmutableList.of(VARCHAR), 10_000, 0);
    DriverContext driverContext = newDriverContext();
    List<RowExpression> projections = ImmutableList.of(field(0, VARCHAR));
    Supplier<CursorProcessor> cursorProcessor = functionAssertions.getExpressionCompiler().compileCursorProcessor(Optional.empty(), projections, "key");
    Supplier<PageProcessor> pageProcessor = functionAssertions.getExpressionCompiler().compilePageProcessor(Optional.empty(), projections);
    ScanFilterAndProjectOperator.ScanFilterAndProjectOperatorFactory factory = new ScanFilterAndProjectOperator.ScanFilterAndProjectOperatorFactory(0, new PlanNodeId("test"), new PlanNodeId("0"), (session, split, table, columns, dynamicFilter) -> new RecordPageSource(new PageRecordSet(ImmutableList.of(VARCHAR), input)), cursorProcessor, pageProcessor, TEST_TABLE_HANDLE, ImmutableList.of(), DynamicFilter.EMPTY, ImmutableList.of(VARCHAR), DataSize.ofBytes(0), 0);
    SourceOperator operator = factory.createOperator(driverContext);
    operator.addSplit(new Split(new CatalogName("test"), TestingSplit.createLocalSplit(), Lifespan.taskWide()));
    operator.noMoreSplits();
    MaterializedResult expected = toMaterializedResult(driverContext.getSession(), ImmutableList.of(VARCHAR), ImmutableList.of(input));
    MaterializedResult actual = toMaterializedResult(driverContext.getSession(), ImmutableList.of(VARCHAR), toPages(operator));
    assertEquals(actual.getRowCount(), expected.getRowCount());
    assertEquals(actual, expected);
}

Example 5 with RecordPageSource

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

the class TpchTables method getTablePages.

public static Iterator<Page> getTablePages(String tableName, double scaleFactor, DecimalTypeMapping decimalTypeMapping) {
    TpchTable<?> table = TpchTable.getTable(tableName);
    ConnectorPageSource pageSource = new RecordPageSource(createTpchRecordSet(table, decimalTypeMapping, scaleFactor, 1, 1, TupleDomain.all()));
    return new AbstractIterator<>() {

        @Override
        protected Page computeNext() {
            if (pageSource.isFinished()) {
                return endOfData();
            }
            Page page = pageSource.getNextPage();
            if (page == null) {
                return computeNext();
            }
            return page.getLoadedPage();
        }
    };
}