use of io.prestosql.spi.relation.VariableReferenceExpression in project hetu-core by openlookeng.
the class TestDynamicFiltersCollector method TestCollectingGlobalDynamicFilters.
@Test
public void TestCollectingGlobalDynamicFilters() throws InterruptedException {
final QueryId queryId = new QueryId("test_query");
final String filterId = "1";
final String columnName = "column";
final TestingColumnHandle columnHandle = new TestingColumnHandle(columnName);
final Set<String> valueSet = ImmutableSet.of("1", "2", "3");
TaskContext taskContext = mock(TaskContext.class);
Session session = testSessionBuilder().setQueryId(queryId).setSystemProperty(ENABLE_DYNAMIC_FILTERING, "true").setSystemProperty(DYNAMIC_FILTERING_DATA_TYPE, "HASHSET").build();
when(taskContext.getSession()).thenReturn(session);
// set up state store and merged dynamic filters map
Map mockMap = new HashMap<>();
StateStoreProvider stateStoreProvider = mock(StateStoreProvider.class);
StateStore stateStore = mock(StateStore.class);
StateMap stateMap = new MockStateMap<>("test-map", mockMap);
when(stateStoreProvider.getStateStore()).thenReturn(stateStore);
when(stateStore.getStateCollection(any())).thenReturn(stateMap);
when(stateStore.createStateMap(any())).thenReturn(stateMap);
when(stateStore.getOrCreateStateCollection(any(), any())).thenReturn(stateMap);
// set up state store listener and dynamic filter cache
StateStoreListenerManager stateStoreListenerManager = new StateStoreListenerManager(stateStoreProvider);
DynamicFilterCacheManager dynamicFilterCacheManager = new DynamicFilterCacheManager();
stateStoreListenerManager.addStateStoreListener(new DynamicFilterListener(dynamicFilterCacheManager), MERGED_DYNAMIC_FILTERS);
LocalDynamicFiltersCollector collector = new LocalDynamicFiltersCollector(taskContext, Optional.empty(), dynamicFilterCacheManager);
TableScanNode tableScan = mock(TableScanNode.class);
when(tableScan.getAssignments()).thenReturn(ImmutableMap.of(new Symbol(columnName), columnHandle));
List<DynamicFilters.Descriptor> dynamicFilterDescriptors = ImmutableList.of(new DynamicFilters.Descriptor(filterId, new VariableReferenceExpression(columnName, BIGINT)));
collector.initContext(ImmutableList.of(dynamicFilterDescriptors), SymbolUtils.toLayOut(tableScan.getOutputSymbols()));
assertTrue(collector.getDynamicFilters(tableScan).isEmpty(), "there should be no dynamic filter available");
// put some values in state store as a new dynamic filter
// and wait for the listener to process the event
stateMap.put(createKey(DynamicFilterUtils.FILTERPREFIX, filterId, queryId.getId()), valueSet);
TimeUnit.MILLISECONDS.sleep(100);
// get available dynamic filter and verify it
List<Map<ColumnHandle, DynamicFilter>> dynamicFilters = collector.getDynamicFilters(tableScan);
assertEquals(dynamicFilters.size(), 1, "there should be a new dynamic filter");
assertEquals(dynamicFilters.size(), 1);
DynamicFilter dynamicFilter = dynamicFilters.get(0).get(columnHandle);
assertTrue(dynamicFilter instanceof HashSetDynamicFilter, "new dynamic filter should be hashset");
assertEquals(dynamicFilter.getSize(), valueSet.size(), "new dynamic filter should have correct size");
for (String value : valueSet) {
assertTrue(dynamicFilter.contains(value), "new dynamic filter should contain correct values");
}
// clean up when task finishes
collector.removeDynamicFilter(true);
DynamicFilter cachedFilter = dynamicFilterCacheManager.getDynamicFilter(DynamicFilterCacheManager.createCacheKey(filterId, queryId.getId()));
assertNull(cachedFilter, "cached dynamic filter should have been removed");
}
use of io.prestosql.spi.relation.VariableReferenceExpression in project hetu-core by openlookeng.
the class RowExpressionVerifier method visitCast.
@Override
protected Boolean visitCast(Cast expected, RowExpression actual) {
// TODO: clean up cast path
if (actual instanceof ConstantExpression && expected.getExpression() instanceof Literal && expected.getType().equals(actual.getType().toString())) {
Literal literal = (Literal) expected.getExpression();
if (literal instanceof StringLiteral) {
Object value = LiteralInterpreter.evaluate((ConstantExpression) actual);
String actualString = value instanceof Slice ? ((Slice) value).toStringUtf8() : String.valueOf(value);
return ((StringLiteral) literal).getValue().equals(actualString);
}
return getValueFromLiteral(literal).equals(String.valueOf(LiteralInterpreter.evaluate((ConstantExpression) actual)));
}
if (actual instanceof VariableReferenceExpression && expected.getExpression() instanceof SymbolReference && expected.getType().equals(actual.getType().toString())) {
return visitSymbolReference((SymbolReference) expected.getExpression(), actual);
}
if (!(actual instanceof CallExpression) || !functionResolution.isCastFunction(((CallExpression) actual).getFunctionHandle())) {
return false;
}
if (!expected.getType().equalsIgnoreCase(actual.getType().toString()) && !(expected.getType().toLowerCase(ENGLISH).equals(VARCHAR) && actual.getType().getTypeSignature().getBase().equals(VARCHAR))) {
return false;
}
return process(expected.getExpression(), ((CallExpression) actual).getArguments().get(0));
}
use of io.prestosql.spi.relation.VariableReferenceExpression in project hetu-core by openlookeng.
the class TestStarTreeAggregationRule method testSupportedProjectNode.
@Test
public void testSupportedProjectNode() {
// node is projectNode and expression instance of cast
Assignments assignment1 = Assignments.builder().put(columnCustkey, new VariableReferenceExpression(columnCustkey.getName(), custkeyHandle.getType())).build();
Optional<PlanNode> planNode1 = Optional.of(new ProjectNode(newId(), baseTableScan, assignment1));
assertTrue(CubeOptimizerUtil.supportedProjectNode(planNode1));
// not projectNode
ListMultimap<Symbol, Symbol> mappings = ImmutableListMultimap.<Symbol, Symbol>builder().put(output, columnOrderkey).put(output, columnOrderkey).build();
Optional<PlanNode> planNode2 = Optional.of(new UnionNode(newId(), ImmutableList.of(baseTableScan, baseTableScan), mappings, ImmutableList.copyOf(mappings.keySet())));
assertFalse(CubeOptimizerUtil.supportedProjectNode(planNode2));
// expression not instance of Cast
Assignments assignment2 = Assignments.builder().put(columnCustkey, new InputReferenceExpression(1, INTEGER)).build();
Optional<PlanNode> planNode3 = Optional.of(new ProjectNode(newId(), baseTableScan, assignment2));
assertFalse(CubeOptimizerUtil.supportedProjectNode(planNode3));
// expression is instance of SymbolReference OR Literal
Assignments assignment3 = Assignments.builder().put(columnCustkey, // should be INTEGER
new VariableReferenceExpression(columnCustkey.getName(), custkeyHandle.getType())).build();
Optional<PlanNode> planNode4 = Optional.of(new ProjectNode(newId(), baseTableScan, assignment3));
assertTrue(CubeOptimizerUtil.supportedProjectNode(planNode4));
// empty node
Optional<PlanNode> emptyPlanNode = Optional.empty();
assertTrue(CubeOptimizerUtil.supportedProjectNode(emptyPlanNode));
}
use of io.prestosql.spi.relation.VariableReferenceExpression in project hetu-core by openlookeng.
the class CubeStatementGenerator method generate.
public static CubeStatement generate(String fromTable, AggregationNode aggregationNode, Map<String, Object> symbolMappings) {
CubeStatement.Builder builder = CubeStatement.newBuilder();
builder.from(fromTable);
Map<Symbol, AggregationNode.Aggregation> aggregations = aggregationNode.getAggregations();
// Extract selection and aggregation
for (Symbol symbol : aggregationNode.getOutputSymbols()) {
Object column = symbolMappings.get(symbol.getName());
if (column instanceof ColumnHandle) {
// Selection
String columnName = ((ColumnHandle) column).getColumnName();
builder.select(columnName);
} else if (aggregations.containsKey(symbol)) {
// Aggregation
Map<Symbol, AggregationSignature> signatures = Collections.emptyMap();
AggregationNode.Aggregation aggregation = aggregations.get(symbol);
List<RowExpression> arguments = aggregation.getArguments();
if (arguments.isEmpty()) {
signatures = createSignature(aggregation, symbol, null);
} else if (arguments.size() == 1) {
RowExpression argument = arguments.get(0);
if (OriginalExpressionUtils.isExpression(argument)) {
Expression argAsExpr = OriginalExpressionUtils.castToExpression(argument);
if (argAsExpr instanceof SymbolReference) {
signatures = createSignature(aggregation, symbol, symbolMappings.get(((SymbolReference) argAsExpr).getName()));
}
} else if (argument instanceof VariableReferenceExpression) {
signatures = createSignature(aggregation, symbol, symbolMappings.get(((VariableReferenceExpression) argument).getName()));
}
}
if (signatures.isEmpty()) {
throw new IllegalArgumentException("Failed to generate aggregator signature");
}
for (Map.Entry<Symbol, AggregationSignature> aggEntry : signatures.entrySet()) {
builder.aggregate(aggEntry.getValue());
}
} else {
throw new IllegalArgumentException("Column " + column + " is not an actual column or expressions");
}
}
// Extract group by
for (Symbol symbol : aggregationNode.getGroupingKeys()) {
Object column = symbolMappings.get(symbol.getName());
if (column instanceof ColumnHandle) {
builder.groupByAddString(((ColumnHandle) column).getColumnName());
} else {
// Don't know how to handle it
throw new IllegalArgumentException("Column " + symbol + " is not an actual column");
}
}
return builder.build();
}
use of io.prestosql.spi.relation.VariableReferenceExpression in project hetu-core by openlookeng.
the class TableDeleteOperator method getOutput.
@Override
public Page getOutput() {
if (state == State.FINISHED || (sourceLayout.isPresent() && state != State.FINISHING)) {
return null;
}
state = State.FINISHED;
TableHandle tableHandleForDelete = tableHandle;
if (sourceLayout.isPresent()) {
// Replace the values from subqueries in filter predicates
VariablesExtractor.extractUnique(filter.get()).stream().forEach(expr -> symbolExpressionMap.putIfAbsent(new Symbol(expr.getName()), expr));
RowExpression rowExpression = RowExpressionVariableInliner.inlineVariables(node -> symbolExpressionMap.get(new Symbol(node.getName())), filter.get());
// Create the tuple domain based on filter and values from source;
RowExpressionDomainTranslator.ExtractionResult<VariableReferenceExpression> decomposedPredicate = (new RowExpressionDomainTranslator(metadata)).fromPredicate(session.toConnectorSession(), rowExpression, RowExpressionDomainTranslator.BASIC_COLUMN_EXTRACTOR);
TupleDomain<ColumnHandle> tupleDomain = decomposedPredicate.getTupleDomain().transform(variableName -> assignments.get(new Symbol(variableName.getName())));
Constraint constraint = new Constraint(tupleDomain);
// Apply the constraint on the table handle
Optional<TableHandle> tableHandleDelete = metadata.applyDelete(session, this.tableHandle, constraint);
if (tableHandleDelete.isPresent()) {
tableHandleForDelete = tableHandleDelete.get();
}
}
OptionalLong rowsDeletedCount = metadata.executeDelete(session, tableHandleForDelete);
// output page will only be constructed once,
// so a new PageBuilder is constructed (instead of using PageBuilder.reset)
PageBuilder page = new PageBuilder(1, TYPES);
BlockBuilder rowsBuilder = page.getBlockBuilder(0);
page.declarePosition();
if (rowsDeletedCount.isPresent()) {
BIGINT.writeLong(rowsBuilder, rowsDeletedCount.getAsLong());
} else {
rowsBuilder.appendNull();
}
return page.build();
}
Aggregations