use of io.prestosql.spi.connector.RecordPageSource in project hetu-core by openlookeng.
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 localMetadata = functionAssertions.getMetadata();
localMetadata.getFunctionAndTypeManager().registerBuiltInFunctions(ImmutableList.of(new GenericLongFunction("record_cursor", value -> {
driverContext.getYieldSignal().forceYieldForTesting();
return value;
})));
ExpressionCompiler compiler = new ExpressionCompiler(localMetadata, new PageFunctionCompiler(localMetadata, 0));
List<RowExpression> projections = ImmutableList.of(call(QualifiedObjectName.valueOfDefaultFunction("generic_long_record_cursor").toString(), new BuiltInFunctionHandle(internalScalarFunction(QualifiedObjectName.valueOfDefaultFunction("generic_long_record_cursor"), BIGINT.getTypeSignature(), ImmutableList.of(BIGINT.getTypeSignature()))), BIGINT, field(0, BIGINT)));
Supplier<CursorProcessor> cursorProcessor = compiler.compileCursorProcessor(Optional.empty(), projections, "key");
Supplier<PageProcessor> pageProcessor = compiler.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(BIGINT), input)), cursorProcessor, pageProcessor, TEST_TABLE_HANDLE, ImmutableList.of(), null, ImmutableList.of(BIGINT), new DataSize(0, BYTE), 0, ReuseExchangeOperator.STRATEGY.REUSE_STRATEGY_DEFAULT, new UUID(0, 0), false, Optional.empty(), 0, 0);
SourceOperator operator = factory.createOperator(driverContext);
operator.addSplit(new Split(new CatalogName("test"), TestingSplit.createLocalSplit(), Lifespan.taskWide()));
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)));
}
use of io.prestosql.spi.connector.RecordPageSource in project hetu-core by openlookeng.
the class HBasePageSourceProvider method createPageSource.
@Override
public ConnectorPageSource createPageSource(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<ColumnHandle> columns) {
// if delete rows, we should replace $rowId -> real rowkey name
List<ColumnHandle> columnsReplaceRowKey = new ArrayList<>();
columns.forEach(ch -> {
if (Constants.HBASE_ROWID_NAME.equals(ch.getColumnName())) {
if (ch instanceof HBaseColumnHandle && table instanceof HBaseTableHandle) {
HBaseColumnHandle rowColumnHandle = (HBaseColumnHandle) ch;
columnsReplaceRowKey.add(new HBaseColumnHandle(((HBaseTableHandle) table).getRowId(), rowColumnHandle.getFamily(), rowColumnHandle.getQualifier(), rowColumnHandle.getType(), rowColumnHandle.getOrdinal(), rowColumnHandle.getComment(), rowColumnHandle.isIndexed()));
}
} else {
columnsReplaceRowKey.add(ch);
}
});
RecordSet recordSet = recordSetProvider.getRecordSet(transactionHandle, session, split, table, columnsReplaceRowKey);
HBaseRecordSet hbaseRecordSet = null;
if (recordSet instanceof HBaseRecordSet) {
hbaseRecordSet = (HBaseRecordSet) recordSet;
}
if (columnsReplaceRowKey.stream().anyMatch(ch -> (ch instanceof HBaseColumnHandle) && (table instanceof HBaseTableHandle) && ((HBaseColumnHandle) ch).getOrdinal() == ((HBaseTableHandle) table).getRowIdOrdinal())) {
return new HBaseUpdatablePageSource(hbaseRecordSet, hbaseConnection);
} else {
return new RecordPageSource(recordSet);
}
}
use of io.prestosql.spi.connector.RecordPageSource in project hetu-core by openlookeng.
the class TestScanFilterAndProjectOperator method testRecordCursorSource.
@Test
public void testRecordCursorSource() {
final Page input = SequencePageBuilder.createSequencePage(ImmutableList.of(VARCHAR), 10_000, 0);
DriverContext driverContext = newDriverContext();
List<RowExpression> projections = ImmutableList.of(field(0, VARCHAR));
Supplier<CursorProcessor> cursorProcessor = expressionCompiler.compileCursorProcessor(Optional.empty(), projections, "key");
Supplier<PageProcessor> pageProcessor = expressionCompiler.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(), null, ImmutableList.of(VARCHAR), new DataSize(0, BYTE), 0, ReuseExchangeOperator.STRATEGY.REUSE_STRATEGY_DEFAULT, new UUID(0, 0), false, Optional.empty(), 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);
}
use of io.prestosql.spi.connector.RecordPageSource in project hetu-core by openlookeng.
the class JdbcPageSourceProvider method createPageSource.
@Override
public ConnectorPageSource createPageSource(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<ColumnHandle> columns) {
JdbcTableHandle jdbcTableHandle = (JdbcTableHandle) table;
RecordSet recordSet = recordSetProvider.getRecordSet(transaction, session, split, table, columns);
if (jdbcTableHandle.getDeleteOrUpdate()) {
return new JdbcUpdatablePageSource(recordSet, session, table, jdbcClient, config, (JdbcSplit) split);
} else {
return new RecordPageSource(recordSet);
}
}
use of io.prestosql.spi.connector.RecordPageSource in project hetu-core by openlookeng.
the class AbstractTestHive method assertPageSourceType.
protected static void assertPageSourceType(ConnectorPageSource connectorPageSource, HiveStorageFormat hiveStorageFormat) {
ConnectorPageSource pageSource = connectorPageSource;
if (pageSource instanceof OrcConcatPageSource) {
pageSource = ((OrcConcatPageSource) pageSource).getConnectorPageSource();
}
if (pageSource instanceof RecordPageSource) {
RecordCursor hiveRecordCursor = ((RecordPageSource) pageSource).getCursor();
hiveRecordCursor = ((HiveRecordCursor) hiveRecordCursor).getRegularColumnRecordCursor();
if (hiveRecordCursor instanceof HiveCoercionRecordCursor) {
hiveRecordCursor = ((HiveCoercionRecordCursor) hiveRecordCursor).getRegularColumnRecordCursor();
}
assertInstanceOf(hiveRecordCursor, recordCursorType(hiveStorageFormat), hiveStorageFormat.name());
} else {
assertInstanceOf(((HivePageSource) pageSource).getPageSource(), pageSourceType(hiveStorageFormat), hiveStorageFormat.name());
}
}
Aggregations