@Override
public Void visitTableWriter(TableWriterNode node, Void context) {
List<String> columns = new ArrayList<>();
for (int i = 0; i < node.getColumnNames().size(); i++) {
columns.add(node.getColumnNames().get(i) + " := " + node.getColumns().get(i));
}
printNode(
node,
format("TableWriter[%s]", Joiner.on(", ").join(columns)),
NODE_COLORS.get(NodeType.TABLE_WRITER));
return node.getSource().accept(this, context);
}
@Override
public SubPlanBuilder visitTableWriter(TableWriterNode node, Void context) {
// TODO: create table in pre-execution step, not here
// Part of the plan should be an Optional<StateChangeListener<QueryState>> and this
// callback can create the table and abort the table creation if the query fails.
OutputTableHandle target =
metadata.beginCreateTable(node.getCatalog(), node.getTableMetadata());
SubPlanBuilder current = node.getSource().accept(this, context);
current.setRoot(
new TableWriterNode(
node.getId(),
current.getRoot(),
target,
node.getColumns(),
node.getColumnNames(),
node.getOutputSymbols(),
node.getSampleWeightSymbol()));
return current;
}
@Override
public PlanNode visitTableWriter(TableWriterNode node, RewriteContext<Set<Symbol>> context) {
ImmutableSet.Builder<Symbol> expectedInputs =
ImmutableSet.<Symbol>builder().addAll(node.getColumns());
if (node.getSampleWeightSymbol().isPresent()) {
expectedInputs.add(node.getSampleWeightSymbol().get());
}
if (node.getPartitionFunction().isPresent()) {
PartitionFunctionBinding functionBinding = node.getPartitionFunction().get();
functionBinding
.getPartitionFunctionArguments()
.stream()
.filter(PartitionFunctionArgumentBinding::isVariable)
.map(PartitionFunctionArgumentBinding::getColumn)
.forEach(expectedInputs::add);
functionBinding.getHashColumn().ifPresent(expectedInputs::add);
}
PlanNode source = context.rewrite(node.getSource(), expectedInputs.build());
return new TableWriterNode(
node.getId(),
source,
node.getTarget(),
node.getColumns(),
node.getColumnNames(),
node.getOutputSymbols(),
node.getSampleWeightSymbol(),
node.getPartitionFunction());
}