use of com.facebook.presto.split.SplitManager in project presto by prestodb.
the class ExtractSpatialJoins method loadKdbTree.
private static KdbTree loadKdbTree(String tableName, Session session, Metadata metadata, SplitManager splitManager, PageSourceManager pageSourceManager) {
QualifiedObjectName name = toQualifiedObjectName(tableName, session.getCatalog().get(), session.getSchema().get());
TableHandle tableHandle = metadata.getTableHandle(session, name).orElseThrow(() -> new PrestoException(INVALID_SPATIAL_PARTITIONING, format("Table not found: %s", name)));
Map<String, ColumnHandle> columnHandles = metadata.getColumnHandles(session, tableHandle);
List<ColumnHandle> visibleColumnHandles = columnHandles.values().stream().filter(handle -> !metadata.getColumnMetadata(session, tableHandle, handle).isHidden()).collect(toImmutableList());
checkSpatialPartitioningTable(visibleColumnHandles.size() == 1, "Expected single column for table %s, but found %s columns", name, columnHandles.size());
ColumnHandle kdbTreeColumn = Iterables.getOnlyElement(visibleColumnHandles);
TableLayoutResult layout = metadata.getLayout(session, tableHandle, Constraint.alwaysTrue(), Optional.of(ImmutableSet.of(kdbTreeColumn)));
TableHandle newTableHandle = layout.getLayout().getNewTableHandle();
Optional<KdbTree> kdbTree = Optional.empty();
try (SplitSource splitSource = splitManager.getSplits(session, newTableHandle, UNGROUPED_SCHEDULING, WarningCollector.NOOP)) {
while (!Thread.currentThread().isInterrupted()) {
SplitBatch splitBatch = getFutureValue(splitSource.getNextBatch(NOT_PARTITIONED, Lifespan.taskWide(), 1000));
List<Split> splits = splitBatch.getSplits();
for (Split split : splits) {
try (ConnectorPageSource pageSource = pageSourceManager.createPageSource(session, split, newTableHandle, ImmutableList.of(kdbTreeColumn))) {
do {
getFutureValue(pageSource.isBlocked());
Page page = pageSource.getNextPage();
if (page != null && page.getPositionCount() > 0) {
checkSpatialPartitioningTable(!kdbTree.isPresent(), "Expected exactly one row for table %s, but found more", name);
checkSpatialPartitioningTable(page.getPositionCount() == 1, "Expected exactly one row for table %s, but found %s rows", name, page.getPositionCount());
String kdbTreeJson = VARCHAR.getSlice(page.getBlock(0), 0).toStringUtf8();
try {
kdbTree = Optional.of(KdbTreeUtils.fromJson(kdbTreeJson));
} catch (IllegalArgumentException e) {
checkSpatialPartitioningTable(false, "Invalid JSON string for KDB tree: %s", e.getMessage());
}
}
} while (!pageSource.isFinished());
} catch (IOException e) {
throw new UncheckedIOException(e);
}
}
if (splitBatch.isLastBatch()) {
break;
}
}
}
checkSpatialPartitioningTable(kdbTree.isPresent(), "Expected exactly one row for table %s, but got none", name);
return kdbTree.get();
}
use of com.facebook.presto.split.SplitManager in project presto by prestodb.
the class ExtractSpatialJoins method tryCreateSpatialJoin.
private static Result tryCreateSpatialJoin(Context context, JoinNode joinNode, RowExpression filter, PlanNodeId nodeId, List<VariableReferenceExpression> outputVariables, CallExpression spatialComparison, Metadata metadata, SplitManager splitManager, PageSourceManager pageSourceManager) {
FunctionMetadata spatialComparisonMetadata = metadata.getFunctionAndTypeManager().getFunctionMetadata(spatialComparison.getFunctionHandle());
checkArgument(spatialComparison.getArguments().size() == 2 && spatialComparisonMetadata.getOperatorType().isPresent());
PlanNode leftNode = joinNode.getLeft();
PlanNode rightNode = joinNode.getRight();
List<VariableReferenceExpression> leftVariables = leftNode.getOutputVariables();
List<VariableReferenceExpression> rightVariables = rightNode.getOutputVariables();
RowExpression radius;
Optional<VariableReferenceExpression> newRadiusVariable;
CallExpression newComparison;
if (spatialComparisonMetadata.getOperatorType().get() == OperatorType.LESS_THAN || spatialComparisonMetadata.getOperatorType().get() == OperatorType.LESS_THAN_OR_EQUAL) {
// ST_Distance(a, b) <= r
radius = spatialComparison.getArguments().get(1);
Set<VariableReferenceExpression> radiusVariables = extractUnique(radius);
if (radiusVariables.isEmpty() || (rightVariables.containsAll(radiusVariables) && containsNone(leftVariables, radiusVariables))) {
newRadiusVariable = newRadiusVariable(context, radius);
newComparison = new CallExpression(spatialComparison.getSourceLocation(), spatialComparison.getDisplayName(), spatialComparison.getFunctionHandle(), spatialComparison.getType(), ImmutableList.of(spatialComparison.getArguments().get(0), mapToExpression(newRadiusVariable, radius)));
} else {
return Result.empty();
}
} else {
// r >= ST_Distance(a, b)
radius = spatialComparison.getArguments().get(0);
Set<VariableReferenceExpression> radiusVariables = extractUnique(radius);
if (radiusVariables.isEmpty() || (rightVariables.containsAll(radiusVariables) && containsNone(leftVariables, radiusVariables))) {
newRadiusVariable = newRadiusVariable(context, radius);
OperatorType flippedOperatorType = flip(spatialComparisonMetadata.getOperatorType().get());
FunctionHandle flippedHandle = getFlippedFunctionHandle(spatialComparison, metadata.getFunctionAndTypeManager());
newComparison = new CallExpression(spatialComparison.getSourceLocation(), // TODO verify if this is the correct function displayName
flippedOperatorType.getOperator(), flippedHandle, spatialComparison.getType(), ImmutableList.of(spatialComparison.getArguments().get(1), mapToExpression(newRadiusVariable, radius)));
} else {
return Result.empty();
}
}
RowExpression newFilter = replaceExpression(filter, ImmutableMap.of(spatialComparison, newComparison));
PlanNode newRightNode = newRadiusVariable.map(variable -> addProjection(context, rightNode, variable, radius)).orElse(rightNode);
JoinNode newJoinNode = new JoinNode(joinNode.getSourceLocation(), joinNode.getId(), joinNode.getType(), leftNode, newRightNode, joinNode.getCriteria(), joinNode.getOutputVariables(), Optional.of(newFilter), joinNode.getLeftHashVariable(), joinNode.getRightHashVariable(), joinNode.getDistributionType(), joinNode.getDynamicFilters());
return tryCreateSpatialJoin(context, newJoinNode, newFilter, nodeId, outputVariables, (CallExpression) newComparison.getArguments().get(0), Optional.of(newComparison.getArguments().get(1)), metadata, splitManager, pageSourceManager);
}
use of com.facebook.presto.split.SplitManager in project presto by prestodb.
the class ExtractSpatialJoins method tryCreateSpatialJoin.
private static Result tryCreateSpatialJoin(Context context, JoinNode joinNode, RowExpression filter, PlanNodeId nodeId, List<VariableReferenceExpression> outputVariables, CallExpression spatialFunction, Optional<RowExpression> radius, Metadata metadata, SplitManager splitManager, PageSourceManager pageSourceManager) {
FunctionAndTypeManager functionAndTypeManager = metadata.getFunctionAndTypeManager();
List<RowExpression> arguments = spatialFunction.getArguments();
verify(arguments.size() == 2);
RowExpression firstArgument = arguments.get(0);
RowExpression secondArgument = arguments.get(1);
// Currently, only inner joins are supported for spherical geometries.
if (joinNode.getType() != INNER && isSphericalJoin(metadata, firstArgument, secondArgument)) {
return Result.empty();
}
Set<VariableReferenceExpression> firstVariables = extractUnique(firstArgument);
Set<VariableReferenceExpression> secondVariables = extractUnique(secondArgument);
if (firstVariables.isEmpty() || secondVariables.isEmpty()) {
return Result.empty();
}
// If either firstArgument or secondArgument is not a
// VariableReferenceExpression, will replace the left/right join node
// with a projection that adds the argument as a variable.
Optional<VariableReferenceExpression> newFirstVariable = newGeometryVariable(context, firstArgument);
Optional<VariableReferenceExpression> newSecondVariable = newGeometryVariable(context, secondArgument);
PlanNode leftNode = joinNode.getLeft();
PlanNode rightNode = joinNode.getRight();
PlanNode newLeftNode;
PlanNode newRightNode;
// Check if the order of arguments of the spatial function matches the order of join sides
int alignment = checkAlignment(joinNode, firstVariables, secondVariables);
if (alignment > 0) {
newLeftNode = newFirstVariable.map(variable -> addProjection(context, leftNode, variable, firstArgument)).orElse(leftNode);
newRightNode = newSecondVariable.map(variable -> addProjection(context, rightNode, variable, secondArgument)).orElse(rightNode);
} else if (alignment < 0) {
newLeftNode = newSecondVariable.map(variable -> addProjection(context, leftNode, variable, secondArgument)).orElse(leftNode);
newRightNode = newFirstVariable.map(variable -> addProjection(context, rightNode, variable, firstArgument)).orElse(rightNode);
} else {
return Result.empty();
}
RowExpression newFirstArgument = mapToExpression(newFirstVariable, firstArgument);
RowExpression newSecondArgument = mapToExpression(newSecondVariable, secondArgument);
// Implement partitioned spatial joins:
// If the session parameter points to a valid spatial partitioning, use
// that to assign to each probe and build rows the partitions that the
// geometry intersects. This is a projection that adds an array of ints
// which is subsequently unnested.
Optional<String> spatialPartitioningTableName = canPartitionSpatialJoin(joinNode) ? getSpatialPartitioningTableName(context.getSession()) : Optional.empty();
Optional<KdbTree> kdbTree = spatialPartitioningTableName.map(tableName -> loadKdbTree(tableName, context.getSession(), metadata, splitManager, pageSourceManager));
Optional<VariableReferenceExpression> leftPartitionVariable = Optional.empty();
Optional<VariableReferenceExpression> rightPartitionVariable = Optional.empty();
if (kdbTree.isPresent()) {
leftPartitionVariable = Optional.of(context.getVariableAllocator().newVariable(newFirstArgument.getSourceLocation(), "pid", INTEGER));
rightPartitionVariable = Optional.of(context.getVariableAllocator().newVariable(newSecondArgument.getSourceLocation(), "pid", INTEGER));
if (alignment > 0) {
newLeftNode = addPartitioningNodes(context, functionAndTypeManager, newLeftNode, leftPartitionVariable.get(), kdbTree.get(), newFirstArgument, Optional.empty());
newRightNode = addPartitioningNodes(context, functionAndTypeManager, newRightNode, rightPartitionVariable.get(), kdbTree.get(), newSecondArgument, radius);
} else {
newLeftNode = addPartitioningNodes(context, functionAndTypeManager, newLeftNode, leftPartitionVariable.get(), kdbTree.get(), newSecondArgument, Optional.empty());
newRightNode = addPartitioningNodes(context, functionAndTypeManager, newRightNode, rightPartitionVariable.get(), kdbTree.get(), newFirstArgument, radius);
}
}
CallExpression newSpatialFunction = new CallExpression(spatialFunction.getSourceLocation(), spatialFunction.getDisplayName(), spatialFunction.getFunctionHandle(), spatialFunction.getType(), ImmutableList.of(newFirstArgument, newSecondArgument));
RowExpression newFilter = RowExpressionNodeInliner.replaceExpression(filter, ImmutableMap.of(spatialFunction, newSpatialFunction));
return Result.ofPlanNode(new SpatialJoinNode(joinNode.getSourceLocation(), nodeId, SpatialJoinNode.Type.fromJoinNodeType(joinNode.getType()), newLeftNode, newRightNode, outputVariables, newFilter, leftPartitionVariable, rightPartitionVariable, kdbTree.map(KdbTreeUtils::toJson)));
}
Aggregations