Examples with RecordPageSource com.facebook.presto.spi.RecordPageSource used on opensource projects

Example 11 with RecordPageSource

use of com.facebook.presto.spi.RecordPageSource in project presto by prestodb.

the class TestScanFilterAndProjectOperator method testRecordCursorYield.

@Test
public void testRecordCursorYield() {
    // create a generic long function that yields for projection on every row
    // verify we will yield #row times totally
    // create a table with 15 rows
    int length = 15;
    Page input = SequencePageBuilder.createSequencePage(ImmutableList.of(BIGINT), length, 0);
    DriverContext driverContext = newDriverContext();
    // set up generic long function with a callback to force yield
    Metadata metadata = functionAssertions.getMetadata();
    FunctionAndTypeManager functionAndTypeManager = metadata.getFunctionAndTypeManager();
    functionAndTypeManager.registerBuiltInFunctions(ImmutableList.of(new GenericLongFunction("record_cursor", value -> {
        driverContext.getYieldSignal().forceYieldForTesting();
        return value;
    })));
    ExpressionCompiler expressionCompiler = new ExpressionCompiler(metadata, new PageFunctionCompiler(metadata, 0));
    List<RowExpression> projections = ImmutableList.of(call("generic_long_record_cursor", functionAndTypeManager.lookupFunction("generic_long_record_cursor", fromTypes(BIGINT)), BIGINT, field(0, BIGINT)));
    Supplier<CursorProcessor> cursorProcessor = expressionCompiler.compileCursorProcessor(driverContext.getSession().getSqlFunctionProperties(), Optional.empty(), projections, "key");
    Supplier<PageProcessor> pageProcessor = expressionCompiler.compilePageProcessor(driverContext.getSession().getSqlFunctionProperties(), Optional.empty(), projections);
    ScanFilterAndProjectOperator.ScanFilterAndProjectOperatorFactory factory = new ScanFilterAndProjectOperator.ScanFilterAndProjectOperatorFactory(0, new PlanNodeId("test"), new PlanNodeId("0"), (session, split, table, columns) -> new RecordPageSource(new PageRecordSet(ImmutableList.of(BIGINT), input)), cursorProcessor, pageProcessor, TESTING_TABLE_HANDLE, ImmutableList.of(), ImmutableList.of(BIGINT), Optional.empty(), new DataSize(0, BYTE), 0);
    SourceOperator operator = factory.createOperator(driverContext);
    operator.addSplit(new Split(new ConnectorId("test"), TestingTransactionHandle.create(), TestingSplit.createLocalSplit()));
    operator.noMoreSplits();
    // start driver; get null value due to yield for the first 15 times
    for (int i = 0; i < length; i++) {
        driverContext.getYieldSignal().setWithDelay(SECONDS.toNanos(1000), driverContext.getYieldExecutor());
        assertNull(operator.getOutput());
        driverContext.getYieldSignal().reset();
    }
    // the 16th yield is not going to prevent the operator from producing a page
    driverContext.getYieldSignal().setWithDelay(SECONDS.toNanos(1000), driverContext.getYieldExecutor());
    Page output = operator.getOutput();
    driverContext.getYieldSignal().reset();
    assertNotNull(output);
    assertEquals(toValues(BIGINT, output.getBlock(0)), toValues(BIGINT, input.getBlock(0)));
}

Example 12 with RecordPageSource

use of com.facebook.presto.spi.RecordPageSource in project presto by prestodb.

the class TestHiveFileFormats method testCursorProvider.

private void testCursorProvider(HiveRecordCursorProvider cursorProvider, FileSplit split, HiveStorageFormat storageFormat, List<TestColumn> testColumns, ConnectorSession session, int rowCount) {
    List<HivePartitionKey> partitionKeys = testColumns.stream().filter(TestColumn::isPartitionKey).map(TestColumn::toHivePartitionKey).collect(toList());
    List<HiveColumnHandle> partitionKeyColumnHandles = getColumnHandles(testColumns.stream().filter(TestColumn::isPartitionKey).collect(toImmutableList()));
    List<Column> tableDataColumns = testColumns.stream().filter(column -> !column.isPartitionKey()).map(column -> new Column(column.getName(), HiveType.valueOf(column.getType()), Optional.empty(), Optional.empty())).collect(toImmutableList());
    Configuration configuration = new Configuration();
    configuration.set("io.compression.codecs", LzoCodec.class.getName() + "," + LzopCodec.class.getName());
    Optional<ConnectorPageSource> pageSource = HivePageSourceProvider.createHivePageSource(ImmutableSet.of(cursorProvider), ImmutableSet.of(), configuration, session, split.getPath(), OptionalInt.empty(), split.getStart(), split.getLength(), split.getLength(), Instant.now().toEpochMilli(), new Storage(StorageFormat.create(storageFormat.getSerDe(), storageFormat.getInputFormat(), storageFormat.getOutputFormat()), "location", Optional.empty(), false, ImmutableMap.of(), ImmutableMap.of()), TupleDomain.all(), getColumnHandles(testColumns), ImmutableMap.of(), partitionKeys, DateTimeZone.getDefault(), FUNCTION_AND_TYPE_MANAGER, new SchemaTableName("schema", "table"), partitionKeyColumnHandles, tableDataColumns, ImmutableMap.of(), tableDataColumns.size(), TableToPartitionMapping.empty(), Optional.empty(), false, DEFAULT_HIVE_FILE_CONTEXT, TRUE_CONSTANT, false, ROW_EXPRESSION_SERVICE, Optional.empty(), ImmutableMap.of());
    RecordCursor cursor = ((RecordPageSource) pageSource.get()).getCursor();
    checkCursor(cursor, testColumns, rowCount);
}

Example 13 with RecordPageSource

use of com.facebook.presto.spi.RecordPageSource in project presto by prestodb.

the class SystemPageSourceProvider method createPageSource.

@Override
public ConnectorPageSource createPageSource(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorSplit split, List<ColumnHandle> columns, SplitContext splitContext) {
    requireNonNull(columns, "columns is null");
    SystemTransactionHandle systemTransaction = (SystemTransactionHandle) transactionHandle;
    SystemSplit systemSplit = (SystemSplit) split;
    SchemaTableName tableName = systemSplit.getTableHandle().getSchemaTableName();
    SystemTable systemTable = tables.getSystemTable(session, tableName).orElseThrow(() -> new PrestoException(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 PrestoException(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 PrestoException(GENERIC_INTERNAL_ERROR, format("Column does not exist: %s.%s", tableName, columnName));
        }
        userToSystemFieldIndex.add(index);
    }
    TupleDomain<ColumnHandle> constraint = systemSplit.getConstraint();
    ImmutableMap.Builder<Integer, Domain> newConstraints = ImmutableMap.builder();
    for (Map.Entry<ColumnHandle, Domain> entry : constraint.getDomains().get().entrySet()) {
        String columnName = ((SystemColumnHandle) entry.getKey()).getColumnName();
        newConstraints.put(columnsByName.get(columnName), entry.getValue());
    }
    TupleDomain<Integer> newContraint = withColumnDomains(newConstraints.build());
    try {
        return new MappedPageSource(systemTable.pageSource(systemTransaction.getConnectorTransactionHandle(), session, newContraint), userToSystemFieldIndex.build());
    } catch (UnsupportedOperationException e) {
        return new RecordPageSource(new MappedRecordSet(toRecordSet(systemTransaction.getConnectorTransactionHandle(), systemTable, session, newContraint), userToSystemFieldIndex.build()));
    }
}

Example 14 with RecordPageSource

use of com.facebook.presto.spi.RecordPageSource in project presto by prestodb.

the class TpchConnectorIndex method lookup.

@Override
public ConnectorPageSource lookup(RecordSet rawInputRecordSet) {
    // convert the input record set from the column ordering in the query to
    // match the column ordering of the index
    RecordSet inputRecordSet = keyFormatter.apply(rawInputRecordSet);
    // lookup the values in the index
    RecordSet rawOutputRecordSet = indexedTable.lookupKeys(inputRecordSet);
    // expect by the query
    return new RecordPageSource(outputFormatter.apply(rawOutputRecordSet));
}