diff --git a/query/ir/codegen/DBPassPipeline.cpp b/query/ir/codegen/DBPassPipeline.cpp index 85e7d92e27..7c13d5dcec 100644 --- a/query/ir/codegen/DBPassPipeline.cpp +++ b/query/ir/codegen/DBPassPipeline.cpp @@ -10,7 +10,7 @@ namespace { using DBPassFactory = std::unique_ptr (*)(const mlir::db::DBPassContext*); -constexpr size_t DB_PASS_COUNT = 39; +constexpr size_t DB_PASS_COUNT = 41; // The optimisation pipeline every query runs through, in order. An EXPLAIN prefix // reporting on a pass walks the same table one pass at a time, which is what keeps the @@ -18,6 +18,7 @@ constexpr size_t DB_PASS_COUNT = 39; // the context; only the join's cost model and the metadata count read it, the rewrites // beside them answering off the IR alone. const std::array dbPassPipeline = { + [](const mlir::db::DBPassContext*) { return mlir::db::createUnnestUnwindComprehension(); }, [](const mlir::db::DBPassContext*) { return mlir::db::createHoistHopConjuncts(); }, [](const mlir::db::DBPassContext*) { return mlir::db::createSinkMakePath(); }, [](const mlir::db::DBPassContext*) { return mlir::db::createFuseLabelPredicates(); }, @@ -57,6 +58,7 @@ const std::array dbPassPipeline = { [](const mlir::db::DBPassContext*) { return mlir::db::createFuseExploreDistinctEnds(); }, [](const mlir::db::DBPassContext* context) { return mlir::db::createCountFromMetadata(context); }, [](const mlir::db::DBPassContext*) { return mlir::db::createFuseUnwindRange(); }, + [](const mlir::db::DBPassContext*) { return mlir::db::createSinkUnwindSource(); }, }; } diff --git a/query/ir/dialect/db/DBOps.cpp b/query/ir/dialect/db/DBOps.cpp index 35ad0f5031..56a0401d3e 100644 --- a/query/ir/dialect/db/DBOps.cpp +++ b/query/ir/dialect/db/DBOps.cpp @@ -1884,13 +1884,43 @@ LogicalResult Unwind::verify() { return verifyCarriedResults(getOperation(), getColumnsToFilter(), getCarried()); } +// The source block takes one argument per carried column, of its type, and ends naming +// the source. +LogicalResult UnwindComputed::verify() { + Block& sourceBlock = getSource().front(); + + if (sourceBlock.empty() || !isa(sourceBlock.back())) { + return emitOpError("source region must end with a db.unwind_yield"); + } + + const OperandRange carried = getColumnsToFilter(); + if (sourceBlock.getNumArguments() != carried.size()) { + return emitOpError("source region takes one argument per carried column, ") + << "expected " << carried.size() << " but has " << sourceBlock.getNumArguments(); + } + + for (size_t carriedIndex = 0; carriedIndex < carried.size(); carriedIndex++) { + const Type argumentType = sourceBlock.getArgument(static_cast(carriedIndex)).getType(); + + if (argumentType != carried[carriedIndex].getType()) { + return emitOpError("source argument ") << carriedIndex + << " must have the type of carried column " + << carriedIndex; + } + } + + return verifyCarriedResults(getOperation(), carried, getCarried()); +} + LogicalResult UnwindRange::verify() { const ColumnType elementColumn = dyn_cast(getElement().getType()); - const storage::NullableType nullable = elementColumn ? dyn_cast(elementColumn.getType()) - : storage::NullableType {}; + const Type element = elementColumn ? elementColumn.getType() : Type {}; + const storage::NullableType nullable = dyn_cast_or_null(element); - if (!nullable || !isa(nullable.getValueType())) { - return emitOpError("element must be a column of nullable integers"); + const bool nullableInteger = nullable && isa(nullable.getValueType()); + const bool resolvedInLowering = isa_and_nonnull(element); + if (!nullableInteger && !resolvedInLowering) { + return emitOpError("element must be a column of nullable integers, or of none for lowering to resolve"); } return verifyCarriedResults(getOperation(), getColumnsToFilter(), getCarried()); diff --git a/query/ir/dialect/db/DBOps.td b/query/ir/dialect/db/DBOps.td index cc6303beb5..5b49f7e959 100644 --- a/query/ir/dialect/db/DBOps.td +++ b/query/ir/dialect/db/DBOps.td @@ -385,7 +385,7 @@ def Unwind : TuringOp<"unwind", [Pure]> { /** * UnwindRange */ -def UnwindRange : TuringOp<"unwind_range", [NoMemoryEffect, DeclareOpInterfaceMethods]> { +def UnwindRange : TuringOp<"unwind_range", [NoMemoryEffect, AttrSizedOperandSegments, DeclareOpInterfaceMethods]> { let summary = "Produce one row per integer a row's range bounds span, without building the list"; let description = [{ @@ -402,15 +402,83 @@ def UnwindRange : TuringOp<"unwind_range", [NoMemoryEffect, DeclareOpInterfaceMe `end`, and the error db.range raises for a step of 0. The step is always an operand, so a range the query gave none counts by a constant 1. Nothing is laid out, so no length limit applies. `columns_to_filter` is the carry set, replicated as db.unwind's - is, and the element column is the one db.unwind gives a list of integers. + is, and the element column is the one db.unwind gives a list of integers, or none + where the list above the range was left for lowering to type, as db.reverse's is. + + `from` and `to` restrict each row to the run of its range that db.list_slice would + keep, with the same bounds, and `reversed` counts that run out backwards. Both are + where db.reverse, db.list_slice and db.tail over a db.range fold, so the list is not + built for those either: + + %x = db.unwind_range(%from, %to, %by, from %a, {}) reversed : ... + + reads `UNWIND reverse(range(from, to, by)[a..]) AS x`. + }]; + + let arguments = (ins Column:$start, + Column:$end, + Column:$step, + Optional:$from, + Optional:$to, + Variadic:$columns_to_filter, + UnitAttr:$reversed); + let results = (outs Column:$element, Variadic:$carried); + let hasVerifier = 1; + + let assemblyFormat = [{ + `(` $start `,` $end `,` $step (`,` `from` $from^)? (`,` `to` $to^)? `,` `{` $columns_to_filter `}` `)` + (`reversed` $reversed^)? attr-dict `:` functional-type(operands, results) + }]; +} + +/** +* UnwindComputed +*/ +def UnwindComputed : TuringOp<"unwind_computed", [Pure]> { + let summary = "Produce one row per element of a value its region computes for each input row"; + + let description = [{ + db.unwind whose source is computed inside the op rather than before it, so lowering + can compute it over a few input rows at a time and never hold the lists of a whole + chunk. sink_unwind_source moves the ops computing a db.unwind's source in here: + + %x, %n = db.unwind_computed({%node}) ({ + ^bb0(%nodeIn: !db.column): + %name = db.get_node_properties(%nodeIn) {property = "name"} : ... + %parts = db.split(%name, %comma) : ... + db.unwind_yield %parts : !db.column> + }) : (!db.column) -> (!db.column, !db.column) + + The region's entry block takes one argument per carried column and reads the input + rows through them; a constant it reads stands for every row and needs no carrying. + The yield names the source, whose cells unwind as db.unwind's do. `columns_to_filter` + is the carry set, replicated as db.unwind's is, and the element type is resolved the + same way. }]; - let arguments = (ins Column:$start, Column:$end, Column:$step, Variadic:$columns_to_filter); + let arguments = (ins Variadic:$columns_to_filter); let results = (outs Column:$element, Variadic:$carried); + + // Exactly one block, ending in a db.unwind_yield. + let regions = (region SizedRegion<1>:$source); + let hasVerifier = 1; let assemblyFormat = [{ - `(` $start `,` $end `,` $step `,` `{` $columns_to_filter `}` `)` attr-dict `:` functional-type(operands, results) + `(` `{` $columns_to_filter `}` `)` $source attr-dict `:` functional-type(operands, results) + }]; +} + +/** +* UnwindYield +*/ +def UnwindYield : TuringOp<"unwind_yield", [Pure, Terminator, HasParent<"UnwindComputed">]> { + let summary = "Name the source a db.unwind_computed's region computes"; + + let arguments = (ins Column:$source); + + let assemblyFormat = [{ + $source attr-dict `:` qualified(type($source)) }]; } diff --git a/query/ir/dialect/db/DBPasses.cpp b/query/ir/dialect/db/DBPasses.cpp index 3b8c1dc813..bee4cfa7c9 100644 --- a/query/ir/dialect/db/DBPasses.cpp +++ b/query/ir/dialect/db/DBPasses.cpp @@ -43,6 +43,8 @@ namespace mlir::db { #define GEN_PASS_DEF_PUSHDOWNFILTERS #define GEN_PASS_DEF_FUSEUNWINDEQUALITY #define GEN_PASS_DEF_FUSEUNWINDRANGE +#define GEN_PASS_DEF_UNNESTUNWINDCOMPREHENSION +#define GEN_PASS_DEF_SINKUNWINDSOURCE #define GEN_PASS_DEF_PLACEINFACTORS #define GEN_PASS_DEF_FUSESCANBYNODEIDS #define GEN_PASS_DEF_FUSESCANBYPROPERTYVALUE @@ -1350,12 +1352,74 @@ void eraseIfUnused(Operation* op, mlir::RewriterBase& rewriter); Value appendCarriedColumn(Operation* op, Value column, mlir::RewriterBase& rewriter); Value seedColumnOf(ExplorePaths exploration, Value column); +// What an unwound list is when it is a run of a range: the range, the one slice taken of it +// in its own order - db.tail being the slice from 1 - and whether the run is reversed. The +// chain holds the ops above the range, from the unwind down. +struct RangeRun { + Range _range; + Value _from; + Value _to; + bool _sliced {false}; + bool _fromOne {false}; + bool _reversed {false}; + llvm::SmallVector _chain; +}; + +bool matchRangeRun(Value list, RangeRun& run) { + bool pendingReversal = false; + + Value current = list; + while (Operation* const op = current.getDefiningOp()) { + if (Range range = dyn_cast(op)) { + // A slice is cut in the range's own order, so a reversal under it is no run + if (run._sliced && pendingReversal) { + return false; + } + + run._range = range; + run._reversed = run._reversed != pendingReversal; + return true; + } else if (Reverse reverse = dyn_cast(op)) { + pendingReversal = !pendingReversal; + current = reverse.getInput(); + } else if (ListSlice slice = dyn_cast(op)) { + if (run._sliced) { + return false; + } + + run._sliced = true; + run._from = slice.getFrom(); + run._to = slice.getTo(); + run._reversed = pendingReversal; + pendingReversal = false; + current = slice.getList(); + } else if (Tail tail = dyn_cast(op)) { + if (run._sliced) { + return false; + } + + run._sliced = true; + run._fromOne = true; + run._reversed = pendingReversal; + pendingReversal = false; + current = tail.getInput(); + } else { + return false; + } + + run._chain.push_back(op); + } + + return false; +} + bool fuseUnwindRange(Unwind unwind, mlir::RewriterBase& rewriter) { - Range range = unwind.getSource().getDefiningOp(); - if (!range) { + RangeRun run; + if (!matchRangeRun(unwind.getSource(), run)) { return false; } + Range range = run._range; const Location loc = unwind.getLoc(); rewriter.setInsertionPoint(unwind); @@ -1364,14 +1428,27 @@ bool fuseUnwindRange(Unwind unwind, mlir::RewriterBase& rewriter) { step = rewriter.create(loc, rewriter.getI64IntegerAttr(1)).getResult(); } + Value from = run._from; + if (run._fromOne) { + from = rewriter.create(loc, rewriter.getI64IntegerAttr(1)).getResult(); + } + UnwindRange fused = rewriter.create(loc, unwind->getResultTypes(), range.getStart(), range.getEnd(), step, - unwind.getColumnsToFilter()); + from, + run._to, + unwind.getColumnsToFilter(), + run._reversed); rewriter.replaceOp(unwind, fused->getResults()); + + for (Operation* const op : run._chain) { + eraseIfUnused(op, rewriter); + } + eraseIfUnused(range.getOperation(), rewriter); return true; @@ -1385,6 +1462,299 @@ struct FuseUnwindRange : public impl::FuseUnwindRangeBase { } }; +// The values a comprehension's row tag reaches: the tag and what each WHERE cuts it to. False +// where the tag reaches anything but those cuts and the yield, which an unwind has no row +// tag to give. +bool collectRowTags(Block& body, ComprehensionYield yield, llvm::SmallPtrSetImpl& rowTags) { + llvm::SmallVector worklist {body.getArgument(1)}; + + while (!worklist.empty()) { + const Value rowTag = worklist.pop_back_val(); + rowTags.insert(rowTag); + + for (OpOperand& use : rowTag.getUses()) { + Operation* const user = use.getOwner(); + FilterOp filter = dyn_cast(user); + const unsigned operandNumber = use.getOperandNumber(); + + const bool cutByAFilter = filter && filter->getBlock() == &body && operandNumber != 0; + if (cutByAFilter) { + worklist.push_back(filter->getResult(operandNumber - 1)); + } else if (user != yield.getOperation()) { + return false; + } + } + } + + return rowTags.contains(yield.getRowTags()); +} + +// Re-emits a WHERE of the body over the unwound rows. The tag is dropped from what it cuts, +// and every column of the rows it does not cut yet is cut alongside, so @param rows, the +// columns of the rows in flight, stays row-aligned with the elements. +void cutUnwoundRows(FilterOp filter, + const llvm::SmallPtrSetImpl& rowTags, + mlir::IRMapping& mapping, + llvm::SmallVectorImpl& rows, + mlir::RewriterBase& rewriter) { + const OperandRange filteredColumns = filter.getColumnsToFilter(); + + llvm::SmallVector cutColumns; + for (const Value column : filteredColumns) { + if (!rowTags.contains(column)) { + cutColumns.push_back(mapping.lookupOrDefault(column)); + } + } + + for (const Value column : rows) { + if (!llvm::is_contained(cutColumns, column)) { + cutColumns.push_back(column); + } + } + + llvm::SmallVector cutTypes; + for (const Value column : cutColumns) { + cutTypes.push_back(column.getType()); + } + + FilterOp cut = rewriter.create(filter.getLoc(), + cutTypes, + mapping.lookupOrDefault(filter.getMask()), + cutColumns); + + size_t cutIndex = 0; + for (const auto [column, result] : llvm::zip(filteredColumns, filter.getResults())) { + if (!rowTags.contains(column)) { + mapping.map(result, cut.getResult(cutIndex)); + cutIndex++; + } + } + + for (Value& column : rows) { + const size_t position = llvm::find(cutColumns, column) - cutColumns.begin(); + column = cut.getResult(position); + } +} + +// UNWIND [x IN xs WHERE p | f] AS y is UNWIND xs AS x WITH x WHERE p, then f as y: the body +// of the comprehension runs over the unwound rows, so no list is built for it. +bool unnestUnwindComprehension(Unwind unwind, mlir::RewriterBase& rewriter) { + ListComprehension comprehension = unwind.getSource().getDefiningOp(); + if (!comprehension || !comprehension->hasOneUse()) { + return false; + } + + Block& body = comprehension.getBody().front(); + ComprehensionYield yield = cast(body.getTerminator()); + + // A value standing for every element would stand for one row once it is unwound + if (::db::yieldsConstantColumn(yield.getValue())) { + return false; + } + + llvm::SmallPtrSet rowTags; + if (!collectRowTags(body, yield, rowTags)) { + return false; + } + + const OperandRange bodyColumns = comprehension.getColumnsToFilter(); + const OperandRange unwoundColumns = unwind.getColumnsToFilter(); + + llvm::SmallVector carried(bodyColumns.begin(), bodyColumns.end()); + for (const Value column : unwoundColumns) { + if (!llvm::is_contained(carried, column)) { + carried.push_back(column); + } + } + + llvm::SmallVector resultTypes {unwind.getElement().getType()}; + for (const Value column : carried) { + resultTypes.push_back(column.getType()); + } + + rewriter.setInsertionPoint(unwind); + Unwind unnested = rewriter.create(unwind.getLoc(), resultTypes, comprehension.getSource(), carried); + + const ResultRange unnestedCarried = unnested.getCarried(); + llvm::SmallVector rows(unnestedCarried.begin(), unnestedCarried.end()); + + mlir::IRMapping mapping; + mapping.map(body.getArgument(0), unnested.getElement()); + + for (size_t carriedIndex = 0; carriedIndex < bodyColumns.size(); carriedIndex++) { + mapping.map(body.getArgument(carriedIndex + 2), rows[carriedIndex]); + } + + for (Operation& op : body.without_terminator()) { + if (FilterOp filter = dyn_cast(op)) { + cutUnwoundRows(filter, rowTags, mapping, rows, rewriter); + } else { + rewriter.clone(op, mapping); + } + } + + llvm::SmallVector replacements {mapping.lookupOrDefault(yield.getValue())}; + for (const Value column : unwoundColumns) { + const size_t position = llvm::find(carried, column) - carried.begin(); + replacements.push_back(rows[position]); + } + + rewriter.replaceOp(unwind, replacements); + rewriter.eraseOp(comprehension); + + return true; +} + +struct UnnestUnwindComprehension : public impl::UnnestUnwindComprehensionBase { + void runOnOperation() override { + runWorklist(getOperation(), [](Operation* op, mlir::RewriterBase& rewriter) { + return unnestUnwindComprehension(cast(op), rewriter); + }); + } +}; + +// The ops computing one value per row out of the row's columns, through a region of their +// own included, which compute the same values over any run of the rows. +bool computesAValuePerRow(Operation* op) { + if (op->hasTrait()) { + return false; + } + + return op->hasTrait() + || computesPerRow(op) + || isa(op); +} + +// The ops computing an unwind's source that nothing else reads, in block order, and the +// columns of the rows they read from outside them. +struct UnwindSourceCone { + llvm::SmallVector _ops; + llvm::SmallVector _inputs; +}; + +bool collectUnwindSourceCone(Unwind unwind, UnwindSourceCone& cone) { + const Value source = unwind.getSource(); + Operation* const root = source.getDefiningOp(); + Block* const block = unwind->getBlock(); + + // A source read from storage is a view of what the graph holds, so no list is built for it + const bool computedHere = root && root->getBlock() == block && computesAValuePerRow(root) + && !isa(root); + if (!computedHere || !source.hasOneUse() || root->getNumResults() != 1 || ::db::yieldsConstantColumn(source)) { + return false; + } + + llvm::SmallPtrSet inCone {root}; + cone._ops.push_back(root); + + // Every reader of an op comes after it, so walking back from the root decides each op + // once all of its readers are decided + for (Operation* op = root->getPrevNode(); op; op = op->getPrevNode()) { + if (!computesAValuePerRow(op)) { + continue; + } + + bool readInCone = false; + bool readOutside = false; + for (Operation* const user : op->getUsers()) { + Operation* const ancestor = block->findAncestorOpInBlock(*user); + if (ancestor && inCone.contains(ancestor)) { + readInCone = true; + } else { + readOutside = true; + } + } + + if (readInCone && !readOutside) { + inCone.insert(op); + cone._ops.push_back(op); + } + } + + std::reverse(cone._ops.begin(), cone._ops.end()); + + const auto definedInCone = [block, &inCone](Value value) { + Operation* const definer = value.getDefiningOp() ? value.getDefiningOp() : value.getParentBlock()->getParentOp(); + Operation* const ancestor = definer ? block->findAncestorOpInBlock(*definer) : nullptr; + return ancestor && inCone.contains(ancestor); + }; + + for (Operation* const op : cone._ops) { + op->walk([&](Operation* nested) { + for (const Value operand : nested->getOperands()) { + const bool readFromOutside = !definedInCone(operand) && !::db::yieldsConstantColumn(operand); + if (readFromOutside && !llvm::is_contained(cone._inputs, operand)) { + cone._inputs.push_back(operand); + } + } + }); + } + + return !cone._inputs.empty(); +} + +// Moves what an unwind's source is computed from into a db.unwind_computed, so lowering can +// compute it a few input rows at a time. +bool sinkUnwindSource(Unwind unwind, mlir::RewriterBase& rewriter) { + UnwindSourceCone cone; + if (!collectUnwindSourceCone(unwind, cone)) { + return false; + } + + const OperandRange unwoundColumns = unwind.getColumnsToFilter(); + + llvm::SmallVector carried(unwoundColumns.begin(), unwoundColumns.end()); + for (const Value input : cone._inputs) { + if (!llvm::is_contained(carried, input)) { + carried.push_back(input); + } + } + + llvm::SmallVector carriedTypes; + llvm::SmallVector carriedLocations; + for (const Value column : carried) { + carriedTypes.push_back(column.getType()); + carriedLocations.push_back(column.getLoc()); + } + + llvm::SmallVector resultTypes {unwind.getElement().getType()}; + resultTypes.append(carriedTypes.begin(), carriedTypes.end()); + + const Location loc = unwind.getLoc(); + rewriter.setInsertionPoint(unwind); + UnwindComputed computed = rewriter.create(loc, resultTypes, carried); + + Block* const sourceBlock = rewriter.createBlock(&computed.getSource(), {}, carriedTypes, carriedLocations); + + mlir::IRMapping mapping; + for (size_t carriedIndex = 0; carriedIndex < carried.size(); carriedIndex++) { + mapping.map(carried[carriedIndex], sourceBlock->getArgument(static_cast(carriedIndex))); + } + + rewriter.setInsertionPointToEnd(sourceBlock); + for (Operation* const op : cone._ops) { + rewriter.clone(*op, mapping); + } + + rewriter.create(loc, mapping.lookup(unwind.getSource())); + + rewriter.replaceOp(unwind, computed->getResults().take_front(unwind->getNumResults())); + + for (Operation* const op : llvm::reverse(cone._ops)) { + rewriter.eraseOp(op); + } + + return true; +} + +struct SinkUnwindSource : public impl::SinkUnwindSourceBase { + void runOnOperation() override { + runWorklist(getOperation(), [](Operation* op, mlir::RewriterBase& rewriter) { + return sinkUnwindSource(cast(op), rewriter); + }); + } +}; + bool computesFromTheRow(Operation* op) { if (op->getNumRegions() != 0) { return false; diff --git a/query/ir/dialect/db/DBPasses.td b/query/ir/dialect/db/DBPasses.td index 4a1c04f555..969079ada4 100644 --- a/query/ir/dialect/db/DBPasses.td +++ b/query/ir/dialect/db/DBPasses.td @@ -24,7 +24,17 @@ def FuseUnwindEquality : Pass<"fuse_unwind_equality"> { } def FuseUnwindRange : Pass<"fuse_unwind_range"> { - let summary = "Fuse an unwind of a range into an unwind_range, which counts the integers out instead of building the list"; + let summary = "Fuse an unwind of a range, or of its reversal or a slice of it, into an unwind_range, which counts the integers out instead of building the list"; + let dependentDialects = ["mlir::db::DB"]; +} + +def SinkUnwindSource : Pass<"sink_unwind_source"> { + let summary = "Move the ops computing an unwind's source into a db.unwind_computed, which computes it a few rows at a time"; + let dependentDialects = ["mlir::db::DB"]; +} + +def UnnestUnwindComprehension : Pass<"unnest_unwind_comprehension"> { + let summary = "Unwind a list comprehension's source and run its body over the unwound rows, instead of building its lists"; let dependentDialects = ["mlir::db::DB"]; } diff --git a/query/ir/dialect/nl/NLOps.td b/query/ir/dialect/nl/NLOps.td index 7dbb89183b..a1a682eb43 100644 --- a/query/ir/dialect/nl/NLOps.td +++ b/query/ir/dialect/nl/NLOps.td @@ -240,7 +240,7 @@ def Unwind : NLOp<"unwind", [Pure]> { * UnwindRange */ // Pure: counting out a row's integers is deterministic; the per-integer fan-out is in the nl.for. -def UnwindRange : NLOp<"unwind_range", [Pure]> { +def UnwindRange : NLOp<"unwind_range", [Pure, AttrSizedOperandSegments]> { let summary = "Create a database iterator that counts out the integers each row's bounds span"; let description = [{ Imperative counterpart of db.unwind_range. Each step of an enclosing nl.for fills one @@ -254,14 +254,67 @@ def UnwindRange : NLOp<"unwind_range", [Pure]> { Every bound carries rows - lowering lays a constant one out over them first - and is read as a nullable value chunk. The integers are computed as the chunk is filled, so - the iterator holds one chunk of them however many a row spans. + the iterator holds one chunk of them however many a row spans. `from`, `to` and + `reversed` are db.unwind_range's, and the slice bounds are read as its range bounds are. }]; - let arguments = (ins NLChunk:$start, NLChunk:$end, NLChunk:$step, Variadic:$columns_to_filter); + let arguments = (ins NLChunk:$start, + NLChunk:$end, + NLChunk:$step, + Optional:$from, + Optional:$to, + Variadic:$columns_to_filter, + UnitAttr:$reversed); + let results = (outs NLIterator:$result); + + let assemblyFormat = [{ + `(` $start `,` $end `,` $step (`,` `from` $from^)? (`,` `to` $to^)? `,` `{` $columns_to_filter `}` `)` + (`reversed` $reversed^)? attr-dict `:` functional-type(operands, results) + }]; +} + +/** +* UnwindComputed +*/ +def UnwindComputed : NLOp<"unwind_computed", [Pure]> { + let summary = "Create a database iterator that computes its source over a few input rows at a time and unwinds it"; + let description = [{ + Imperative counterpart of db.unwind_computed. The region's arguments are the carried + chunks cut to a run of input rows, and its yield the source cells of that run. Each + step of an enclosing nl.for fills one chunk of elements, as nl.unwind does, plus one + chunk per carried column gathered from the whole input; once a run's cells are spent + the region runs again over the next, sized so it holds about one chunk of elements: + + %rows = nl.unwind_computed({%names}) ({ + ^bb0(%namesIn: !nl.chunk>): + %parts = nl.split %namesIn, %comma : ... + nl.unwind_yield %parts : !nl.chunk>> + }) : (!nl.chunk>) + -> !nl.iter>, + !nl.chunk>> + }]; + + let arguments = (ins Variadic:$columns_to_filter); let results = (outs NLIterator:$result); + // Exactly one block, ending in an nl.unwind_yield. + let regions = (region SizedRegion<1>:$source); + + let assemblyFormat = [{ + `(` `{` $columns_to_filter `}` `)` $source attr-dict `:` functional-type(operands, results) + }]; +} + +/** +* UnwindYield +*/ +def UnwindYield : NLOp<"unwind_yield", [Pure, Terminator, HasParent<"UnwindComputed">]> { + let summary = "Name the source cells an nl.unwind_computed's region computes for a run of rows"; + + let arguments = (ins NLChunk:$source); + let assemblyFormat = [{ - `(` $start `,` $end `,` $step `,` `{` $columns_to_filter `}` `)` attr-dict `:` functional-type(operands, results) + $source attr-dict `:` qualified(type($source)) }]; } diff --git a/query/ir/exec/NLExecutor.cpp b/query/ir/exec/NLExecutor.cpp index fdadc88566..ae77ce1396 100644 --- a/query/ir/exec/NLExecutor.cpp +++ b/query/ir/exec/NLExecutor.cpp @@ -221,6 +221,8 @@ class NLRangeSpan { throw IRException("range() cannot count by a step of 0"); } + _windowFirst = 0; + _reversed = false; _ascending = *by > 0; _empty = _ascending ? *from > *to : *from < *to; _first = static_cast(*from); @@ -230,13 +232,70 @@ class NLRangeSpan { return true; } + // Narrows the span to the run a slice of its list keeps, as runListSlice cuts it, and + // counts that run out backwards where @param window is reversed. Returns false where a + // slice bound of @param row holds no value. + bool keepWindow(const NLRangeWindowBounds& window, size_t row) { + const NLRangeBound& from = window._from; + const NLRangeBound& to = window._to; + + const std::optional lower = + from._column ? from._read(from._column, row) : std::optional {0}; + const std::optional upper = + to._column ? to._read(to._column, row) : std::optional {}; + + const bool boundsPresent = lower.has_value() && (!to._column || upper.has_value()); + if (!boundsPresent) { + return false; + } + + _reversed = window._reversed; + + if (_empty) { + return true; + } + + // The list holds _steps + 1 integers, up to 2^64, so the run is located by its first + // and last positions, and a bound counting from the end by how far it stops short of it + uint64_t first = 0; + if (*lower >= 0) { + first = static_cast(*lower); + } else { + const uint64_t shortOfLast = static_cast(-(*lower + 1)); + first = shortOfLast >= _steps ? 0 : _steps - shortOfLast; + } + + const types::Int64::Primitive upperBound = upper.value_or(0); + + uint64_t last = 0; + bool upperKeepsAny = true; + if (!to._column) { + last = _steps; + } else if (upperBound > 0) { + last = std::min(static_cast(upperBound) - 1, _steps); + } else if (upperBound == 0) { + upperKeepsAny = false; + } else { + const uint64_t shortOfLast = static_cast(-(upperBound + 1)); + upperKeepsAny = shortOfLast < _steps; + last = upperKeepsAny ? _steps - shortOfLast - 1 : 0; + } + + _empty = !upperKeepsAny || first > last; + _windowFirst = first; + _steps = _empty ? 0 : last - first; + + return true; + } + bool isEmpty() const { return _empty; } uint64_t getSteps() const { return _steps; } // The stride can be wider than an int64 holds, so every integer is offset from the // first unsigned types::Int64::Primitive at(uint64_t position) const { - const uint64_t offset = position * _stride; + const uint64_t listPosition = _reversed ? _windowFirst + _steps - position : _windowFirst + position; + const uint64_t offset = listPosition * _stride; return static_cast(_ascending ? _first + offset : _first - offset); } @@ -244,8 +303,10 @@ class NLRangeSpan { uint64_t _first {0}; uint64_t _stride {0}; uint64_t _steps {0}; + uint64_t _windowFirst {0}; bool _ascending {true}; bool _empty {true}; + bool _reversed {false}; }; // The cursor an unwind_range loop walks: every (row, integer) pair the bounds span, in row @@ -257,10 +318,16 @@ class NLRangeCursor { : _start(loopData->getStart()), _end(loopData->getEnd()), _step(loopData->getStep()), + _window(loopData->getWindow()), _rowCount(_start._column->size()) { - bioassert(_end._column->size() == _rowCount && _step._column->size() == _rowCount, - "Bound columns of an unwind_range are not row-aligned."); + const auto isRowAligned = [this](const NLRangeBound& bound) { + return !bound._column || bound._column->size() == _rowCount; + }; + + const bool boundsRowAligned = isRowAligned(_end) && isRowAligned(_step) + && isRowAligned(_window._from) && isRowAligned(_window._to); + bioassert(boundsRowAligned, "Bound columns of an unwind_range are not row-aligned."); openNextRow(); } @@ -305,6 +372,7 @@ class NLRangeCursor { NLRangeBound _start; NLRangeBound _end; NLRangeBound _step; + NLRangeWindowBounds _window; size_t _rowCount {0}; size_t _row {0}; uint64_t _position {0}; @@ -314,7 +382,9 @@ class NLRangeCursor { _position = 0; while (_row < _rowCount) { - const bool spansIntegers = _span.read(_start, _end, _step, _row) && !_span.isEmpty(); + const bool spansIntegers = _span.read(_start, _end, _step, _row) + && _span.keepWindow(_window, _row) + && !_span.isEmpty(); if (spansIntegers) { return; } @@ -6692,6 +6762,80 @@ void NLExecutor::runUnwindRangeLoop(NLExecutionContext* context, NLFunctionData* } } +void NLExecutor::runUnwindComputedLoop(NLExecutionContext* context, NLFunctionData* data) { + NLUnwindComputedLoopData* loopData = static_cast(data); + const NLStmtContainer* loopBody = loopData->getStmts(); + const NLStmtContainer* sourceStmts = loopData->getSourceStmts(); + const size_t chunkSize = context->getChunkSize(); + + // A null limit leaves the loop unbounded, exactly as in runUnwindLoop. + const NLLimitState* limit = loopData->getLimit(); + const auto withinLimit = [limit]() { + return !limit || limit->getRemaining() > 0; + }; + + const NLUnwindElementCountFunction elementCount = loopData->getElementCountFunc(); + const NLUnwindElementEmitFunction elementEmit = loopData->getElementEmitFunc(); + const NLCarriedColumn* elementCell = loopData->getElementCell(); + const Column* source = loopData->getSource(); + + ColumnVector* runRows = loopData->getRunRows(); + ColumnVector* rows = loopData->getRows(); + ColumnVector* positions = loopData->getPositions(); + std::vector& runRaw = runRows->getRaw(); + std::vector& rowsRaw = rows->getRaw(); + + const size_t inputRows = loopData->getInputRows(); + + // A run of input rows holds about one chunk of elements once the first run has shown + // how many a row spans, so the source never holds the elements of a whole chunk of rows + size_t runStart = 0; + size_t runLength = 1; + + while (runStart < inputRows && withinLimit()) { + runLength = std::min(runLength, inputRows - runStart); + + runRaw.resize(runLength); + std::iota(runRaw.begin(), runRaw.end(), runStart); + + for (const NLCarriedColumn& argument : loopData->sourceArguments()) { + const auto gatherFunc = argument.getGatherFunc(); + gatherFunc(argument.getInput(), runRows, argument.getOutput()); + } + + runBody(context, sourceStmts); + + NLElementCursor cursor(source, elementCount); + size_t elements = 0; + + while (!cursor.exhausted() && withinLimit()) { + cursor.fillChunk(chunkSize, rowsRaw, elementEmit ? &positions->getRaw() : nullptr); + elements += rowsRaw.size(); + + if (elementEmit) { + elementEmit(source, rows, positions, loopData->getElementOutput()); + } else if (elementCell) { + const auto gatherFunc = elementCell->getGatherFunc(); + gatherFunc(elementCell->getInput(), rows, elementCell->getOutput()); + } + + std::ranges::transform(rowsRaw, rowsRaw.begin(), [runStart](size_t row) { + return row + runStart; + }); + + for (const NLCarriedColumn& carriedColumn : loopData->carriedColumns()) { + const auto gatherFunc = carriedColumn.getGatherFunc(); + gatherFunc(carriedColumn.getInput(), rows, carriedColumn.getOutput()); + } + + runBody(context, loopBody); + } + + runStart += runLength; + runLength = elements == 0 ? chunkSize : std::max(1, chunkSize * runLength / elements); + } +} + void NLExecutor::runScanEdgesLoop(NLExecutionContext* context, NLFunctionData* data) { NLScanEdgesLoopData* loopData = static_cast(data); const NLStmtContainer* loopBody = loopData->getStmts(); diff --git a/query/ir/exec/NLExecutor.h b/query/ir/exec/NLExecutor.h index 8cec5d841b..97e7528f0d 100644 --- a/query/ir/exec/NLExecutor.h +++ b/query/ir/exec/NLExecutor.h @@ -71,6 +71,7 @@ class NLExecutor { // Count out the integers each row's bounds span one chunk at a time, with the carry // set gathered by the row each integer came from. A null limit leaves it unbounded. static void runUnwindRangeLoop(NLExecutionContext* context, NLFunctionData* data); + static void runUnwindComputedLoop(NLExecutionContext* context, NLFunctionData* data); // Walk the loop data's columns one row at a time: each step gathers row k of every // column into its one-row output and runs the body over it. A null limit leaves it diff --git a/query/ir/exec/NLProgram.h b/query/ir/exec/NLProgram.h index 2196d69935..d6cde46e5f 100644 --- a/query/ir/exec/NLProgram.h +++ b/query/ir/exec/NLProgram.h @@ -4401,9 +4401,17 @@ class NLListSliceData : public NLFunctionData { Column* _result {nullptr}; }; -// nl.unwind_range loop data: the bounds each row is counted out by, the element chunk a -// step fills with the next integers, and the carry set gathered by the row each integer -// came from. +// The run of each row's range an nl.unwind_range keeps: the slice bounds, absent where the +// query gave none, and whether the run is counted out backwards. +struct NLRangeWindowBounds { + NLRangeBound _from; + NLRangeBound _to; + bool _reversed {false}; +}; + +// nl.unwind_range loop data: the bounds each row is counted out by and the run of it kept, +// the element chunk a step fills with the next integers, and the carry set gathered by the +// row each integer came from. class NLUnwindRangeLoopData : public NLFunctionData { public: using CarriedColumns = std::vector; @@ -4411,10 +4419,12 @@ class NLUnwindRangeLoopData : public NLFunctionData { NLUnwindRangeLoopData(const NLRangeBound& start, const NLRangeBound& end, const NLRangeBound& step, + const NLRangeWindowBounds& window, Column* elementOutput) : _start(start), _end(end), _step(step), + _window(window), _elementOutput(elementOutput) { } @@ -4422,6 +4432,7 @@ class NLUnwindRangeLoopData : public NLFunctionData { const NLRangeBound& getStart() const { return _start; } const NLRangeBound& getEnd() const { return _end; } const NLRangeBound& getStep() const { return _step; } + const NLRangeWindowBounds& getWindow() const { return _window; } Column* getElementOutput() const { return _elementOutput; } const CarriedColumns& carriedColumns() const { return _carriedColumns; } @@ -4442,13 +4453,97 @@ class NLUnwindRangeLoopData : public NLFunctionData { NLRangeBound _start; NLRangeBound _end; NLRangeBound _step; + NLRangeWindowBounds _window; + Column* _elementOutput {nullptr}; + + NLLimitState* _limit {nullptr}; + CarriedColumns _carriedColumns; + NLStmtContainer _stmts; + + ColumnVector _rows; +}; + +// nl.unwind_computed loop data: what an nl.unwind's loop data holds, plus the region +// computing the source over a run of input rows - its statements, the carried inputs it +// reads cut to the run, and the rows of the run. The source, and how its cells unwind, are +// only known once the region is translated. The carried columns are gathered from the +// whole input chunk, whose rows @param input holds. +class NLUnwindComputedLoopData : public NLFunctionData { +public: + using CarriedColumns = std::vector; + + explicit NLUnwindComputedLoopData(const Column* input) + : _input(input) + { + } + + // @param elementEmit and @param elementOutput are null when the source's cells are + // themselves the elements, as on NLUnwindLoopData; they are then gathered through + // @param elementCell, which is null where nothing reads them. + void setSource(const Column* source, + NLUnwindElementCountFunction elementCount, + NLUnwindElementEmitFunction elementEmit, + Column* elementOutput, + const NLCarriedColumn* elementCell) { + _source = source; + _elementCount = elementCount; + _elementEmit = elementEmit; + _elementOutput = elementOutput; + + if (elementCell) { + _elementCell.emplace(*elementCell); + } + } + + size_t getInputRows() const { return _input->size(); } + const Column* getSource() const { return _source; } + NLUnwindElementCountFunction getElementCountFunc() const { return _elementCount; } + NLUnwindElementEmitFunction getElementEmitFunc() const { return _elementEmit; } + Column* getElementOutput() const { return _elementOutput; } + const NLCarriedColumn* getElementCell() const { return _elementCell ? &*_elementCell : nullptr; } + + const CarriedColumns& sourceArguments() const { return _sourceArguments; } + + void addSourceArgument(const NLCarriedColumn& argument) { + _sourceArguments.push_back(argument); + } + + const CarriedColumns& carriedColumns() const { return _carriedColumns; } + + void addCarriedColumn(const NLCarriedColumn& carried) { + _carriedColumns.push_back(carried); + } + + ColumnVector* getRunRows() { return &_runRows; } + ColumnVector* getRows() { return &_rows; } + ColumnVector* getPositions() { return &_positions; } + + NLLimitState* getLimit() const { return _limit; } + void setLimit(NLLimitState* limit) { _limit = limit; } + + NLStmtContainer* getSourceStmts() { return &_sourceStmts; } + const NLStmtContainer* getSourceStmts() const { return &_sourceStmts; } + + NLStmtContainer* getStmts() { return &_stmts; } + const NLStmtContainer* getStmts() const { return &_stmts; } + +private: + const Column* _input {nullptr}; + const Column* _source {nullptr}; + NLUnwindElementCountFunction _elementCount {nullptr}; + NLUnwindElementEmitFunction _elementEmit {nullptr}; Column* _elementOutput {nullptr}; + std::optional _elementCell; NLLimitState* _limit {nullptr}; + CarriedColumns _sourceArguments; CarriedColumns _carriedColumns; + NLStmtContainer _sourceStmts; NLStmtContainer _stmts; + ColumnVector _runRows; ColumnVector _rows; + ColumnVector _positions; }; // What an op computing over the elements of one row's list holds - nl.list_comprehension diff --git a/query/ir/exec/NLTranslator.cpp b/query/ir/exec/NLTranslator.cpp index ce80ed5ac8..6f345eeca3 100644 --- a/query/ir/exec/NLTranslator.cpp +++ b/query/ir/exec/NLTranslator.cpp @@ -906,12 +906,26 @@ void NLTranslator::translateBlock(mlir::Block& block, NLStmtContainer* body) { } else if (nl::UnwindRange unwindRange = mlir::dyn_cast(operation)) { IteratorConfig config; config._kind = IteratorKind::UnwindRange; - config._rangeBounds = {unwindRange.getStart(), unwindRange.getEnd(), unwindRange.getStep()}; + config._rangeBounds = {unwindRange.getStart(), + unwindRange.getEnd(), + unwindRange.getStep(), + unwindRange.getFrom(), + unwindRange.getTo()}; + config._rangeReversed = unwindRange.getReversed(); const mlir::OperandRange carriedColumns = unwindRange.getColumnsToFilter(); config._carriedColumns.assign(carriedColumns.begin(), carriedColumns.end()); _iteratorConfigs[unwindRange.getResult()] = config; + } else if (nl::UnwindComputed unwindComputed = mlir::dyn_cast(operation)) { + IteratorConfig config; + config._kind = IteratorKind::UnwindComputed; + config._unwindSource = &unwindComputed.getSource(); + + const mlir::OperandRange carriedColumns = unwindComputed.getColumnsToFilter(); + config._carriedColumns.assign(carriedColumns.begin(), carriedColumns.end()); + + _iteratorConfigs[unwindComputed.getResult()] = config; } else if (nl::OptionalDrain optionalDrain = mlir::dyn_cast(operation)) { IteratorConfig config; config._kind = IteratorKind::OptionalDrain; @@ -1196,7 +1210,7 @@ void NLTranslator::translateBlock(mlir::Block& block, NLStmtContainer* body) { translateProcedure(procedureOp, body); } else if (nl::Output output = mlir::dyn_cast(operation)) { translateOutput(output, body); - } else if (mlir::isa(operation)) { + } else if (mlir::isa(operation)) { // Structural terminators carry no behavior } else if (!_systemTranslator->translate(operation, body)) { throw IRException(fmt::format("NLTranslator cannot translate operation '{}'", @@ -1268,6 +1282,8 @@ void NLTranslator::translateFor(nl::For forLoop, NLStmtContainer* body) { translateUnwindLoop(config, loopBody, limit, body); } else if (config._kind == IteratorKind::UnwindRange) { translateUnwindRangeLoop(config, loopBody, limit, body); + } else if (config._kind == IteratorKind::UnwindComputed) { + translateUnwindComputedLoop(config, loopBody, limit, body); } else if (config._kind == IteratorKind::OptionalDrain) { translateOptionalDrainLoop(config, loopBody, limit, body); } else if (config._kind == IteratorKind::RowDrain) { @@ -1765,15 +1781,88 @@ void NLTranslator::translateUnwindLoop(const IteratorConfig& config, translateBlock(loopBody, loopData->getStmts()); } +void NLTranslator::translateUnwindComputedLoop(const IteratorConfig& config, + mlir::Block& loopBody, + NLLimitState* limit, + NLStmtContainer* body) { + const llvm::ArrayRef carriedValues = config._carriedColumns; + NLUnwindComputedLoopData* loopData = _program->allocFunctionData(getColumn(carriedValues.front())); + loopData->setLimit(limit); + + const size_t chunkSize = _program->getChunkSize(); + loopData->getRunRows()->reserve(chunkSize); + loopData->getRows()->reserve(chunkSize); + loopData->getPositions()->reserve(chunkSize); + + mlir::Block& sourceBlock = config._unwindSource->front(); + for (size_t carriedIndex = 0; carriedIndex < carriedValues.size(); carriedIndex++) { + const mlir::Value argument = sourceBlock.getArgument(static_cast(carriedIndex)); + if (argument.use_empty()) { + continue; + } + + const NLCarriedColumn sourceArgument(getColumn(carriedValues[carriedIndex]), + allocColumn(argument), + selectGatherForChunkType(argument.getType())); + loopData->addSourceArgument(sourceArgument); + } + + translateBlock(sourceBlock, loopData->getSourceStmts()); + + // The op's verifier guarantees the terminator, so a region without one here means + // unverified IR. + nl::UnwindYield yield = mlir::dyn_cast(sourceBlock.back()); + if (!yield) { + throw IRException("The source region of an nl.unwind_computed does not end with an nl.unwind_yield"); + } + + const mlir::Value sourceValue = yield.getSource(); + const Column* source = getColumn(sourceValue); + const nl::ChunkType sourceChunk = mlir::cast(sourceValue.getType()); + + const mlir::Value elementValue = loopBody.getArgument(0); + + NLUnwindElementCountFunction elementCount = nullptr; + NLUnwindElementEmitFunction elementEmit = nullptr; + selectElementDrain(sourceChunk.getElementType(), elementValue.getType(), elementCount, elementEmit); + + const bool readsTheElements = !elementValue.use_empty(); + if (!readsTheElements) { + elementEmit = nullptr; + } + + Column* const elementOutput = elementEmit ? allocColumn(elementValue) : nullptr; + + if (readsTheElements && !elementEmit) { + const NLCarriedColumn elementCell(source, allocColumn(elementValue), selectGatherForChunkType(sourceChunk)); + loopData->setSource(source, elementCount, elementEmit, elementOutput, &elementCell); + } else { + loopData->setSource(source, elementCount, elementEmit, elementOutput, nullptr); + } + + bindCarriedColumns(carriedValues, loopBody, 1, loopData); + + body->emplaceStmt(&NLExecutor::runUnwindComputedLoop, loopData); + + translateBlock(loopBody, loopData->getStmts()); +} + void NLTranslator::translateUnwindRangeLoop(const IteratorConfig& config, mlir::Block& loopBody, NLLimitState* limit, NLStmtContainer* body) { - const llvm::SmallVector& bounds = config._rangeBounds; + const llvm::SmallVector& bounds = config._rangeBounds; + + const NLRangeWindowBounds window { + ._from = rangeBound(bounds[3]), + ._to = rangeBound(bounds[4]), + ._reversed = config._rangeReversed, + }; NLUnwindRangeLoopData* loopData = _program->allocFunctionData(rangeBound(bounds[0]), rangeBound(bounds[1]), rangeBound(bounds[2]), + window, allocColumnIfUsed(loopBody.getArgument(0))); loopData->setLimit(limit); loopData->getRows()->reserve(_program->getChunkSize()); diff --git a/query/ir/exec/NLTranslator.h b/query/ir/exec/NLTranslator.h index c99c9f6ded..ca72c60faa 100644 --- a/query/ir/exec/NLTranslator.h +++ b/query/ir/exec/NLTranslator.h @@ -87,6 +87,7 @@ class NLTranslator { VectorSearch, Unwind, UnwindRange, + UnwindComputed, ProcedureInit, OptionalDrain, RowDrain, @@ -206,9 +207,14 @@ class NLTranslator { // cells; null for the other kinds. mlir::Value _source; - // The start, end and step an UnwindRange iterator counts each row out by; empty - // for the other kinds. - llvm::SmallVector _rangeBounds; + // The start, end and step an UnwindRange iterator counts each row out by, then the + // slice bounds of the run it keeps, null where absent; empty for the other kinds. + llvm::SmallVector _rangeBounds; + bool _rangeReversed {false}; + + // The region an UnwindComputed iterator computes its source in; null for the other + // kinds. + mlir::Region* _unwindSource {nullptr}; }; NLProgram* _program {nullptr}; @@ -407,6 +413,14 @@ class NLTranslator { // came from, so every one of them must hold the input's rows static void throwIfCarriedNotRowAligned(const IteratorConfig& config); + // Translate the nl.for over an nl.unwind_computed iterator: translate the region + // computing the source over a run of rows, bind the run's arguments, the element loop + // variable and one per carried column, and record the loop statement. + void translateUnwindComputedLoop(const IteratorConfig& config, + mlir::Block& loopBody, + NLLimitState* limit, + NLStmtContainer* body); + // Translate the nl.for over an nl.unwind_range iterator: allocate the element loop // variable and one per carried column, bind the reads its bounds are taken through, // and record the counting loop statement. diff --git a/query/ir/lowering/DBLowering.cpp b/query/ir/lowering/DBLowering.cpp index bc7d802784..d0587bf258 100644 --- a/query/ir/lowering/DBLowering.cpp +++ b/query/ir/lowering/DBLowering.cpp @@ -935,6 +935,7 @@ bool opensSourceLoop(mlir::Operation* operation) { mlir::db::VectorSearch, mlir::db::Unwind, mlir::db::UnwindRange, + mlir::db::UnwindComputed, mlir::db::ScanNodesByLabel, mlir::db::ScanNodesByPropertyValue, mlir::db::ScanEdges, @@ -1335,6 +1336,8 @@ void DBLowering::lowerOperationByKind(mlir::Operation& operation) { lowerUnwind(unwind); } else if (mlir::db::UnwindRange unwindRange = mlir::dyn_cast(operation)) { lowerUnwindRange(unwindRange); + } else if (mlir::db::UnwindComputed unwindComputed = mlir::dyn_cast(operation)) { + lowerUnwindComputed(unwindComputed); } else if (mlir::db::MakeList makeList = mlir::dyn_cast(operation)) { lowerMakeList(makeList); } else if (mlir::db::MakeMap makeMap = mlir::dyn_cast(operation)) { @@ -1758,9 +1761,20 @@ void DBLowering::lowerUnwindRange(mlir::db::UnwindRange unwindRange) { carriedChunks.push_back(mapValue(carriedColumn)); } - std::array bounds {mapValue(unwindRange.getStart()), - mapValue(unwindRange.getEnd()), - mapValue(unwindRange.getStep())}; + const mlir::Value from = unwindRange.getFrom(); + const mlir::Value to = unwindRange.getTo(); + + llvm::SmallVector bounds {mapValue(unwindRange.getStart()), + mapValue(unwindRange.getEnd()), + mapValue(unwindRange.getStep())}; + + if (from) { + bounds.push_back(mapValue(from)); + } + + if (to) { + bounds.push_back(mapValue(to)); + } llvm::SmallVector operandChunks(carriedChunks.begin(), carriedChunks.end()); operandChunks.append(bounds.begin(), bounds.end()); @@ -1787,15 +1801,91 @@ void DBLowering::lowerUnwindRange(mlir::db::UnwindRange unwindRange) { const nl::IteratorType iteratorType = nl::IteratorType::get(context, chunkTypes); + const mlir::Value fromChunk = from ? bounds[3] : mlir::Value {}; + const mlir::Value toChunk = to ? bounds.back() : mlir::Value {}; + nl::UnwindRange rows = _builder.create(_builder.getUnknownLoc(), iteratorType, bounds[0], bounds[1], bounds[2], - carriedChunks); + fromChunk, + toChunk, + carriedChunks, + unwindRange.getReversed()); buildLoopForSource(rows.getResult(), unwindRange.getOperation()); } +void DBLowering::lowerUnwindComputed(mlir::db::UnwindComputed unwind) { + const mlir::Location loc = _builder.getUnknownLoc(); + mlir::MLIRContext* const context = _builder.getContext(); + + llvm::SmallVector carriedChunks; + llvm::SmallVector carriedTypes; + llvm::SmallVector carriedLocations; + for (const mlir::Value carriedColumn : unwind.getColumnsToFilter()) { + const mlir::Value carriedChunk = mapValue(carriedColumn); + carriedChunks.push_back(carriedChunk); + carriedTypes.push_back(carriedChunk.getType()); + carriedLocations.push_back(loc); + } + + setInsertionInto(deepestOwnerBlock(carriedChunks, _rootBlock)); + + // The element type is the source's to decide, and the source is only known once the + // region is lowered, so the iterator is typed after it + nl::UnwindComputed rows = _builder.create(loc, + nl::IteratorType::get(context, carriedTypes), + carriedChunks); + + mlir::Block& dbSource = unwind.getSource().front(); + mlir::Block* const nlSource = &rows.getSource().emplaceBlock(); + nlSource->addArguments(carriedTypes, carriedLocations); + + for (size_t argumentIndex = 0; argumentIndex < carriedChunks.size(); argumentIndex++) { + const unsigned index = static_cast(argumentIndex); + _valueMap[dbSource.getArgument(index)] = nlSource->getArgument(index); + } + + const mlir::Value runCardinality = nlSource->getArgument(0); + + _builder.setInsertionPointToEnd(nlSource); + mlir::Operation* const placeholderYield = _builder.create(loc, runCardinality); + + mlir::Block* const previousInnermostLoopBody = _innermostLoopBody; + const mlir::Value previousInnermostCardinality = _innermostCardinality; + _innermostLoopBody = nlSource; + _innermostCardinality = runCardinality; + + mlir::Value sourceChunk; + for (mlir::Operation& operation : dbSource) { + mlir::db::UnwindYield yield = mlir::dyn_cast(operation); + if (yield) { + sourceChunk = rowAlignedChunk(mapValue(yield.getSource()), runCardinality); + } else { + lowerOperation(operation); + } + } + + _innermostLoopBody = previousInnermostLoopBody; + _innermostCardinality = previousInnermostCardinality; + + _builder.setInsertionPointToEnd(nlSource); + _builder.create(loc, sourceChunk); + placeholderYield->erase(); + + const mlir::Type sourceElement = mlir::cast(sourceChunk.getType()).getElementType(); + + llvm::SmallVector chunkTypes; + chunkTypes.push_back(nl::ChunkType::get(context, ::db::unwoundElementType(context, sourceElement))); + chunkTypes.append(carriedTypes.begin(), carriedTypes.end()); + + rows.getResult().setType(nl::IteratorType::get(context, chunkTypes)); + + _builder.setInsertionPointAfter(rows); + buildLoopForSource(rows.getResult(), unwind.getOperation()); +} + void DBLowering::lowerElementOperands(mlir::Value source, mlir::OperandRange carried, ElementOperands& operands) { const mlir::Location loc = _builder.getUnknownLoc(); mlir::MLIRContext* const context = _builder.getContext(); diff --git a/query/ir/lowering/DBLowering.h b/query/ir/lowering/DBLowering.h index 8655d9cbf3..3688b54f4c 100644 --- a/query/ir/lowering/DBLowering.h +++ b/query/ir/lowering/DBLowering.h @@ -260,6 +260,7 @@ class DBLowering { void lowerVectorSearch(mlir::db::VectorSearch vectorSearch); void lowerUnwind(mlir::db::Unwind unwind); void lowerUnwindRange(mlir::db::UnwindRange unwindRange); + void lowerUnwindComputed(mlir::db::UnwindComputed unwind); void lowerMakeList(mlir::db::MakeList makeList); void lowerMakeMap(mlir::db::MakeMap makeMap); void lowerRange(mlir::db::Range range); diff --git a/test/query/ir/ProductPlacement.cpp b/test/query/ir/ProductPlacement.cpp index fe9646d289..2860e97f21 100644 --- a/test/query/ir/ProductPlacement.cpp +++ b/test/query/ir/ProductPlacement.cpp @@ -93,5 +93,9 @@ size_t turing::test::countUnwindsOverAProduct(mlir::ModuleOp module) { count += carriesAProductColumn(unwindRange.getColumnsToFilter()) ? 1 : 0; }); + module.walk([&count](mlir::db::UnwindComputed unwindComputed) { + count += carriesAProductColumn(unwindComputed.getColumnsToFilter()) ? 1 : 0; + }); + return count; } diff --git a/test/query/ir/unwind/CMakeLists.txt b/test/query/ir/unwind/CMakeLists.txt index 38d4c43d4a..4b3f4343b5 100644 --- a/test/query/ir/unwind/CMakeLists.txt +++ b/test/query/ir/unwind/CMakeLists.txt @@ -50,3 +50,15 @@ add_call_v3_test(test_query_ir_unwind_unread_element UnwindUnreadElementTest.cpp add_turing_test(test_query_ir_unwind_range_codegen UnwindRangeCodegenTest.cpp) target_include_directories(test_query_ir_unwind_range_codegen PRIVATE ${IR_TEST_DIR}) + +add_call_v3_test(test_query_ir_unwind_comprehension UnwindComprehensionTest.cpp) + +add_call_v3_test(test_query_ir_unwind_range_window UnwindRangeWindowTest.cpp) + +add_turing_test(test_query_ir_unwind_list_producer_codegen UnwindListProducerCodegenTest.cpp) +target_include_directories(test_query_ir_unwind_list_producer_codegen PRIVATE ${IR_TEST_DIR}) + +add_call_v3_test(test_query_ir_unwind_computed UnwindComputedTest.cpp) + +add_turing_test(test_query_ir_unwind_computed_codegen UnwindComputedCodegenTest.cpp) +target_include_directories(test_query_ir_unwind_computed_codegen PRIVATE ${IR_TEST_DIR}) diff --git a/test/query/ir/unwind/UnwindComprehensionTest.cpp b/test/query/ir/unwind/UnwindComprehensionTest.cpp new file mode 100644 index 0000000000..9e9ffd4ee5 --- /dev/null +++ b/test/query/ir/unwind/UnwindComprehensionTest.cpp @@ -0,0 +1,133 @@ +#include + +#include +#include + +#include "CallV3Test.h" +#include "StringRowSink.h" + +using namespace turing::test; + +class UnwindComprehensionTest : public CallV3Test { +protected: + void expectRows(std::string_view query, const std::vector& expected) { + StringRowSink sink; + runQuery(query, sink); + + EXPECT_EQ(sink.getRows(), expected) << "query: " << query; + } + + void expectSortedRows(std::string_view query, const std::vector& expected) { + StringRowSink sink; + runQuery(query, sink); + + std::vector rows; + sink.sortedRows(rows); + + EXPECT_EQ(rows, expected) << "query: " << query; + } +}; + +TEST_F(UnwindComprehensionTest, unwindsAProjection) { + expectRows("UNWIND [x IN range(1, 5) | x * 10] AS y RETURN y", {{"10"}, {"20"}, {"30"}, {"40"}, {"50"}}); +} + +TEST_F(UnwindComprehensionTest, unwindsAFilter) { + expectRows("UNWIND [x IN range(1, 10) WHERE x % 3 = 0] AS y RETURN y", {{"3"}, {"6"}, {"9"}}); +} + +TEST_F(UnwindComprehensionTest, unwindsAFilterAndAProjection) { + expectRows("UNWIND [x IN range(1, 10) WHERE x % 2 = 0 | x * x] AS y RETURN y", + {{"4"}, {"16"}, {"36"}, {"64"}, {"100"}}); +} + +TEST_F(UnwindComprehensionTest, projectsNoElementTheFilterDrops) { + expectRows("UNWIND [x IN [2, 0, 4] WHERE x <> 0 | 100 / x] AS y RETURN y", {{"50"}, {"25"}}); +} + +TEST_F(UnwindComprehensionTest, unwindsNothingWhereTheFilterDropsEveryElement) { + expectRows("UNWIND [x IN range(1, 3) WHERE x > 5] AS y RETURN count(*)", {{"0"}}); +} + +TEST_F(UnwindComprehensionTest, unwindsNothingForANullList) { + expectRows("WITH null AS xs UNWIND [x IN xs | x] AS y RETURN count(*)", {{"0"}}); +} + +TEST_F(UnwindComprehensionTest, keepsTheElementOrder) { + expectRows("UNWIND [x IN [3, 'a', 2.5, 1] | x] AS y RETURN y", {{"3"}, {"a"}, {"2.5"}, {"1"}}); +} + +TEST_F(UnwindComprehensionTest, countsOverTheOuterElement) { + expectRows("UNWIND range(1, 3) AS i UNWIND [x IN range(1, i) | x + i] AS y RETURN i, y", + {{"1", "2"}, {"2", "3"}, {"2", "4"}, {"3", "4"}, {"3", "5"}, {"3", "6"}}); +} + +// Only Remy and Adam carry an age, 32 each: x < null drops every element of the six others +TEST_F(UnwindComprehensionTest, readsTheMatchedRowInTheFilterAndTheProjection) { + expectSortedRows("MATCH (n:Person) UNWIND [x IN range(1, 3) WHERE x < n.age | n.name + toString(x)] AS y RETURN n.name, y", + {{"Adam", "Adam1"}, + {"Adam", "Adam2"}, + {"Adam", "Adam3"}, + {"Remy", "Remy1"}, + {"Remy", "Remy2"}, + {"Remy", "Remy3"}}); +} + +TEST_F(UnwindComprehensionTest, filtersAgainstAMatchedProperty) { + expectSortedRows("MATCH (n:Person) UNWIND [x IN range(30, 34) WHERE x >= n.age] AS y RETURN n.name, y", + {{"Adam", "32"}, + {"Adam", "33"}, + {"Adam", "34"}, + {"Remy", "32"}, + {"Remy", "33"}, + {"Remy", "34"}}); +} + +TEST_F(UnwindComprehensionTest, repeatsTheListForEveryMatchedRow) { + expectRows("MATCH (n:Person) UNWIND [x IN range(1, 4) WHERE x % 2 = 0] AS y RETURN count(*)", {{"16"}}); +} + +TEST_F(UnwindComprehensionTest, filtersRowsLongerThanAChunk) { + expectSortedRows("UNWIND range(1, 3) AS i UNWIND [x IN range(1, 5000) WHERE x % 2 = 0 | x * i] AS y RETURN i, count(y), sum(y)", + {{"1", "2500", "6252500"}, {"2", "2500", "12505000"}, {"3", "2500", "18757500"}}); +} + +TEST_F(UnwindComprehensionTest, unwindsANestedComprehension) { + expectRows("UNWIND [x IN range(1, 3) | [y IN range(1, x) | y * 10]] AS l RETURN l", + {{"10"}, {"10, 20"}, {"10, 20, 30"}}); +} + +TEST_F(UnwindComprehensionTest, keepsTheListAnotherProjectionReads) { + expectRows("WITH [x IN range(1, 3) | x * 2] AS xs UNWIND xs AS y RETURN y, xs", + {{"2", "2, 4, 6"}, {"4", "2, 4, 6"}, {"6", "2, 4, 6"}}); +} + +TEST_F(UnwindComprehensionTest, stopsAtTheLimit) { + expectRows("UNWIND range(1, 65536) AS i UNWIND [x IN range(i, i + 999) | x * 2] AS j RETURN i, j LIMIT 2", + {{"1", "2"}, {"1", "4"}}); +} + +TEST_F(UnwindComprehensionTest, countsTheElementsOfEveryRowOfAChunk) { + expectRows("UNWIND range(1, 65536) AS i UNWIND [x IN range(i, i + 999) WHERE x % 2 = 0 | x] AS j RETURN count(*)", + {{"32768000"}}); +} + +TEST_F(UnwindComprehensionTest, unwindsAConstantProjectionOncePerElement) { + expectRows("UNWIND [x IN range(1, 3) | 7] AS y RETURN y", {{"7"}, {"7"}, {"7"}}); +} + +TEST_F(UnwindComprehensionTest, unwindsAConstantProjectionOncePerKeptElement) { + expectRows("UNWIND [x IN range(1, 3) WHERE x > 1 | 'k'] AS y RETURN y", {{"k"}, {"k"}}); +} + +TEST_F(UnwindComprehensionTest, limitsTheFilteredElements) { + expectRows("UNWIND [x IN range(1, 10) WHERE x > 2] AS y RETURN y LIMIT 2", {{"3"}, {"4"}}); +} + +TEST_F(UnwindComprehensionTest, skipsPastAChunkOfFilteredElements) { + expectRows("UNWIND [x IN range(1, 10000) WHERE x % 2 = 1] AS y RETURN y SKIP 2048 LIMIT 2", {{"4097"}, {"4099"}}); +} + +TEST_F(UnwindComprehensionTest, unwindsAComprehensionOverAComprehension) { + expectRows("UNWIND [x IN [y IN range(1, 5) | y * 2] WHERE x > 4] AS z RETURN z", {{"6"}, {"8"}, {"10"}}); +} diff --git a/test/query/ir/unwind/UnwindComputedCodegenTest.cpp b/test/query/ir/unwind/UnwindComputedCodegenTest.cpp new file mode 100644 index 0000000000..903fefbb08 --- /dev/null +++ b/test/query/ir/unwind/UnwindComputedCodegenTest.cpp @@ -0,0 +1,124 @@ +#include + +#include +#include +#include + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OwningOpRef.h" + +#include "DBDialect.h" +#include "DBOps.h" +#include "DBProgramGenerator.h" +#include "NLDialect.h" +#include "StorageDialect.h" + +#include "CypherAST.h" +#include "CypherAnalyzer.h" +#include "CypherParser.h" +#include "Graph.h" +#include "SimpleGraph.h" +#include "SystemAccessor.h" +#include "SystemManager.h" +#include "versioning/Transaction.h" +#include "views/GraphView.h" + +#include "TuringTest.h" +#include "TuringTestEnv.h" + +#include "IRTestOps.h" + +using namespace db; +using namespace turing::test; + +class UnwindComputedCodegenTest : public TuringTest { +public: + void initialize() override { + _env = TuringTestEnv::create(fs::Path {_outDir} / "turing"); + + SystemAccessor system = _env->getSystemManager().accessUnique(); + _graph = system.createGraph(_graphName); + SimpleGraph::createSimpleGraph(_graph); + + _context.getOrLoadDialect(); + _context.getOrLoadDialect(); + _context.getOrLoadDialect(); + _context.getOrLoadDialect(); + } + +protected: + mlir::OwningOpRef generate(std::string_view query) { + SystemAccessor system = _env->getSystemManager().accessUnique(); + const ProcedureManager* procedures = system.getProcedures(); + + const FrozenCommitTx transaction = _graph->openTransaction(); + const GraphView view = transaction.viewGraph(); + + CypherAST ast(procedures, query); + + CypherParser parser(&ast); + parser.parse(query); + + CypherAnalyzer analyzer(&ast, view); + analyzer.analyze(); + + mlir::OpBuilder builder(&_context); + mlir::OwningOpRef owningModule = mlir::ModuleOp::create(builder.getUnknownLoc()); + mlir::ModuleOp module = owningModule.get(); + + DBProgramGenerator generator(&module); + generator.generate(&ast); + + return owningModule; + } + +private: + const std::string _graphName {"simpledb"}; + std::unique_ptr _env; + Graph* _graph {nullptr}; + mlir::MLIRContext _context; +}; + +TEST_F(UnwindComputedCodegenTest, aSplitPerRowIsComputedInTheUnwind) { + const mlir::OwningOpRef module = generate("UNWIND range(1, 3) AS i UNWIND split(toString(i) + ',a,b', ',') AS j RETURN i, j"); + + EXPECT_EQ(countOps(*module), 1u); + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 1u); +} + +TEST_F(UnwindComputedCodegenTest, aReversedComprehensionIsComputedInTheUnwind) { + const mlir::OwningOpRef module = generate("UNWIND range(1, 65536) AS i UNWIND reverse([x IN range(i, i + 999) | x * 2]) AS j RETURN count(*)"); + + EXPECT_EQ(countOps(*module), 1u); + EXPECT_EQ(countOps(*module), 0u); +} + +TEST_F(UnwindComputedCodegenTest, aRangeIsCountedOutRatherThanComputed) { + const mlir::OwningOpRef module = generate("UNWIND range(1, 3) AS i UNWIND range(i, i + 9) AS j RETURN count(*)"); + + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 2u); +} + +TEST_F(UnwindComputedCodegenTest, aConstantListIsNotComputedInTheUnwind) { + const mlir::OwningOpRef module = generate("UNWIND range(1, 3) AS i UNWIND [1, 2, 3] AS j RETURN count(*)"); + + EXPECT_EQ(countOps(*module), 0u); +} + +TEST_F(UnwindComputedCodegenTest, aStoredPropertyIsNotComputedInTheUnwind) { + const mlir::OwningOpRef module = generate("MATCH (n:Person) UNWIND n.name AS x RETURN x"); + + EXPECT_EQ(countOps(*module), 0u); +} + +TEST_F(UnwindComputedCodegenTest, aListAnotherProjectionReadsIsNotComputedInTheUnwind) { + const mlir::OwningOpRef module = generate("UNWIND range(1, 2) AS i WITH i, reverse(range(1, i) + [0]) AS l UNWIND l AS j RETURN i, j, l"); + + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 1u); +} diff --git a/test/query/ir/unwind/UnwindComputedTest.cpp b/test/query/ir/unwind/UnwindComputedTest.cpp new file mode 100644 index 0000000000..08444151a9 --- /dev/null +++ b/test/query/ir/unwind/UnwindComputedTest.cpp @@ -0,0 +1,116 @@ +#include + +#include +#include + +#include "CallV3Test.h" +#include "StringRowSink.h" + +using namespace turing::test; + +class UnwindComputedTest : public CallV3Test { +protected: + void expectRows(std::string_view query, const std::vector& expected) { + StringRowSink sink; + runQuery(query, sink); + + EXPECT_EQ(sink.getRows(), expected) << "query: " << query; + } + + void expectSortedRows(std::string_view query, const std::vector& expected) { + StringRowSink sink; + runQuery(query, sink); + + std::vector rows; + sink.sortedRows(rows); + + EXPECT_EQ(rows, expected) << "query: " << query; + } +}; + +TEST_F(UnwindComputedTest, unwindsASplitStringOfEveryRow) { + expectRows("UNWIND range(1, 3) AS i UNWIND split(toString(i) + ',a,b', ',') AS j RETURN i, j", + {{"1", "1"}, {"1", "a"}, {"1", "b"}, + {"2", "2"}, {"2", "a"}, {"2", "b"}, + {"3", "3"}, {"3", "a"}, {"3", "b"}}); +} + +TEST_F(UnwindComputedTest, unwindsAConcatenation) { + expectRows("UNWIND range(1, 2) AS i UNWIND range(i, i + 1) + [10 * i] AS j RETURN i, j", + {{"1", "1"}, {"1", "2"}, {"1", "10"}, {"2", "2"}, {"2", "3"}, {"2", "20"}}); +} + +TEST_F(UnwindComputedTest, unwindsAReversedComprehension) { + expectRows("UNWIND range(1, 2) AS i UNWIND reverse([x IN range(1, 3) | x * i]) AS j RETURN i, j", + {{"1", "3"}, {"1", "2"}, {"1", "1"}, {"2", "6"}, {"2", "4"}, {"2", "2"}}); +} + +TEST_F(UnwindComputedTest, unwindsASliceOfAComprehension) { + expectRows("UNWIND range(1, 2) AS i UNWIND [x IN range(1, 4) | x + i][1..3] AS j RETURN i, j", + {{"1", "3"}, {"1", "4"}, {"2", "4"}, {"2", "5"}}); +} + +TEST_F(UnwindComputedTest, unwindsACollectSubquery) { + expectRows("UNWIND range(1, 2) AS i UNWIND COLLECT { UNWIND range(1, 3) AS k RETURN k + i } AS j RETURN i, j", + {{"1", "2"}, {"1", "3"}, {"1", "4"}, {"2", "3"}, {"2", "4"}, {"2", "5"}}); +} + +TEST_F(UnwindComputedTest, unwindsAReducedList) { + expectRows("UNWIND range(1, 2) AS i UNWIND reduce(acc = [], k IN range(1, 3) | acc + [k * i]) AS j RETURN i, j", + {{"1", "1"}, {"1", "2"}, {"1", "3"}, {"2", "2"}, {"2", "4"}, {"2", "6"}}); +} + +TEST_F(UnwindComputedTest, unwindsAScalarComputedPerRow) { + expectRows("UNWIND range(1, 3) AS i UNWIND toString(i) + '!' AS j RETURN j", {{"1!"}, {"2!"}, {"3!"}}); +} + +TEST_F(UnwindComputedTest, unwindsNothingForANullList) { + expectRows("UNWIND [1, null, 3] AS i UNWIND range(1, i) + [] AS j RETURN i, j", + {{"1", "1"}, {"3", "1"}, {"3", "2"}, {"3", "3"}}); +} + +// Only Remy and Adam carry an age, 32 each +TEST_F(UnwindComputedTest, carriesTheMatchedRow) { + expectSortedRows("MATCH (n:Person) WHERE n.age IS NOT NULL UNWIND reverse(range(1, n.age)[..2] + [n.name]) AS x RETURN n.name, x", + {{"Adam", "1"}, {"Adam", "2"}, {"Adam", "Adam"}, {"Remy", "1"}, {"Remy", "2"}, {"Remy", "Remy"}}); +} + +TEST_F(UnwindComputedTest, keepsTheListAnotherProjectionReads) { + expectRows("UNWIND range(1, 2) AS i WITH i, reverse(range(1, i) + [0]) AS l UNWIND l AS j RETURN i, j, l", + {{"1", "0", "0, 1"}, {"1", "1", "0, 1"}, {"2", "0", "0, 2, 1"}, {"2", "2", "0, 2, 1"}, {"2", "1", "0, 2, 1"}}); +} + +TEST_F(UnwindComputedTest, unwindsAComputedListPerElementOfAnother) { + expectRows("UNWIND range(1, 2) AS i UNWIND split(toString(i) + ',x', ',') AS s UNWIND split(s + ';' + s, ';') AS t RETURN t", + {{"1"}, {"1"}, {"x"}, {"x"}, {"2"}, {"2"}, {"x"}, {"x"}}); +} + +TEST_F(UnwindComputedTest, unwindsRowsLongerThanAChunk) { + expectSortedRows("UNWIND range(1, 3) AS i UNWIND reverse([x IN range(1, 5000) | x * i]) AS j RETURN i, count(j), sum(j)", + {{"1", "5000", "12502500"}, {"2", "5000", "25005000"}, {"3", "5000", "37507500"}}); +} + +TEST_F(UnwindComputedTest, countsTheElementsOfEveryRowOfAChunk) { + expectRows("UNWIND range(1, 65536) AS i UNWIND reverse([x IN range(i, i + 999) | x * 2]) AS j RETURN count(*), sum(j)", + {{"65536000", "4360503296000"}}); +} + +TEST_F(UnwindComputedTest, stopsAtTheLimit) { + expectRows("UNWIND range(1, 65536) AS i UNWIND reverse([x IN range(i, i + 999) | x]) AS j RETURN i, j LIMIT 2", + {{"1", "1000"}, {"1", "999"}}); +} + +TEST_F(UnwindComputedTest, skipsAcrossRows) { + expectRows("UNWIND range(1, 100) AS i UNWIND split(toString(i) + ',x', ',') AS j RETURN j SKIP 7 LIMIT 3", + {{"x"}, {"5"}, {"x"}}); +} + +TEST_F(UnwindComputedTest, unwindsRowsMostOfWhichHoldNoElement) { + expectRows("UNWIND range(1, 10000) AS i UNWIND range(1, i % 1000 - 998) + [] AS j RETURN collect(i)", + {{"999, 1999, 2999, 3999, 4999, 5999, 6999, 7999, 8999, 9999"}}); +} + +TEST_F(UnwindComputedTest, unwindsRowsWhoseListsGrow) { + expectRows("UNWIND range(1, 2000) AS i UNWIND reverse(range(1, i) + []) AS j RETURN count(*), sum(j)", + {{"2001000", "1335334000"}}); +} diff --git a/test/query/ir/unwind/UnwindListProducerCodegenTest.cpp b/test/query/ir/unwind/UnwindListProducerCodegenTest.cpp new file mode 100644 index 0000000000..e4443a7837 --- /dev/null +++ b/test/query/ir/unwind/UnwindListProducerCodegenTest.cpp @@ -0,0 +1,150 @@ +#include + +#include +#include +#include + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OwningOpRef.h" + +#include "DBDialect.h" +#include "DBOps.h" +#include "DBProgramGenerator.h" +#include "NLDialect.h" +#include "StorageDialect.h" + +#include "CypherAST.h" +#include "CypherAnalyzer.h" +#include "CypherParser.h" +#include "Graph.h" +#include "SimpleGraph.h" +#include "SystemAccessor.h" +#include "SystemManager.h" +#include "versioning/Transaction.h" +#include "views/GraphView.h" + +#include "TuringTest.h" +#include "TuringTestEnv.h" + +#include "IRTestOps.h" + +using namespace db; +using namespace turing::test; + +class UnwindListProducerCodegenTest : public TuringTest { +public: + void initialize() override { + _env = TuringTestEnv::create(fs::Path {_outDir} / "turing"); + + SystemAccessor system = _env->getSystemManager().accessUnique(); + _graph = system.createGraph(_graphName); + SimpleGraph::createSimpleGraph(_graph); + + _context.getOrLoadDialect(); + _context.getOrLoadDialect(); + _context.getOrLoadDialect(); + _context.getOrLoadDialect(); + } + +protected: + mlir::OwningOpRef generate(std::string_view query) { + SystemAccessor system = _env->getSystemManager().accessUnique(); + const ProcedureManager* procedures = system.getProcedures(); + + const FrozenCommitTx transaction = _graph->openTransaction(); + const GraphView view = transaction.viewGraph(); + + CypherAST ast(procedures, query); + + CypherParser parser(&ast); + parser.parse(query); + + CypherAnalyzer analyzer(&ast, view); + analyzer.analyze(); + + mlir::OpBuilder builder(&_context); + mlir::OwningOpRef owningModule = mlir::ModuleOp::create(builder.getUnknownLoc()); + mlir::ModuleOp module = owningModule.get(); + + DBProgramGenerator generator(&module); + generator.generate(&ast); + + return owningModule; + } + +private: + const std::string _graphName {"simpledb"}; + std::unique_ptr _env; + Graph* _graph {nullptr}; + mlir::MLIRContext _context; +}; + +TEST_F(UnwindListProducerCodegenTest, aComprehensionOverARangeBuildsNoList) { + const mlir::OwningOpRef module = generate("UNWIND range(1, 65536) AS i UNWIND [x IN range(i, i + 999) | x * 2] AS j RETURN count(*)"); + + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 2u); +} + +TEST_F(UnwindListProducerCodegenTest, aFilteringComprehensionBuildsNoList) { + const mlir::OwningOpRef module = generate("MATCH (n:Person) UNWIND [x IN range(1, 3) WHERE x < n.age | n.name + toString(x)] AS y RETURN n.name, y"); + + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 0u); +} + +TEST_F(UnwindListProducerCodegenTest, aComprehensionAnotherProjectionReadsIsKept) { + const mlir::OwningOpRef module = generate("WITH [x IN range(1, 3) | x * 2] AS xs UNWIND xs AS y RETURN y, xs"); + + EXPECT_EQ(countOps(*module), 1u); +} + +TEST_F(UnwindListProducerCodegenTest, aReversedRangeBuildsNoList) { + const mlir::OwningOpRef module = generate("UNWIND range(1, 65536) AS i UNWIND reverse(range(i, i + 999)) AS j RETURN count(*)"); + + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 2u); +} + +TEST_F(UnwindListProducerCodegenTest, aSliceOfARangeBuildsNoList) { + const mlir::OwningOpRef module = generate("UNWIND range(1, 65536) AS i UNWIND range(i, i + 999)[1..] AS j RETURN count(*)"); + + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 2u); +} + +TEST_F(UnwindListProducerCodegenTest, theTailOfARangeBuildsNoList) { + const mlir::OwningOpRef module = generate("UNWIND range(1, 65536) AS i UNWIND tail(range(i, i + 999)) AS j RETURN count(*)"); + + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 0u); +} + +TEST_F(UnwindListProducerCodegenTest, aReversedSliceBuildsNoList) { + const mlir::OwningOpRef module = generate("UNWIND reverse(range(1, 10)[2..5]) AS x RETURN x"); + + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 0u); +} + +TEST_F(UnwindListProducerCodegenTest, aSliceOfAReversedRangeIsBuilt) { + const mlir::OwningOpRef module = generate("UNWIND reverse(range(1, 10))[2..5] AS x RETURN x"); + + EXPECT_EQ(countOps(*module), 1u); + EXPECT_EQ(countOps(*module), 0u); +} + +TEST_F(UnwindListProducerCodegenTest, aComprehensionOverAComprehensionBuildsNoList) { + const mlir::OwningOpRef module = generate("UNWIND [x IN [y IN range(1, 5) | y * 2] WHERE x > 4] AS z RETURN z"); + + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 0u); +} diff --git a/test/query/ir/unwind/UnwindRangeWindowTest.cpp b/test/query/ir/unwind/UnwindRangeWindowTest.cpp new file mode 100644 index 0000000000..6758424849 --- /dev/null +++ b/test/query/ir/unwind/UnwindRangeWindowTest.cpp @@ -0,0 +1,166 @@ +#include + +#include +#include + +#include "CallV3Test.h" +#include "StringRowSink.h" + +using namespace turing::test; + +class UnwindRangeWindowTest : public CallV3Test { +protected: + void expectRows(std::string_view query, const std::vector& expected) { + StringRowSink sink; + runQuery(query, sink); + + EXPECT_EQ(sink.getRows(), expected) << "query: " << query; + } + + void expectSortedRows(std::string_view query, const std::vector& expected) { + StringRowSink sink; + runQuery(query, sink); + + std::vector rows; + sink.sortedRows(rows); + + EXPECT_EQ(rows, expected) << "query: " << query; + } +}; + +TEST_F(UnwindRangeWindowTest, unwindsAReversedRange) { + expectRows("UNWIND reverse(range(1, 5)) AS x RETURN x", {{"5"}, {"4"}, {"3"}, {"2"}, {"1"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsAReversedRangeCountingDown) { + expectRows("UNWIND reverse(range(10, 1, -3)) AS x RETURN x", {{"1"}, {"4"}, {"7"}, {"10"}}); +} + +TEST_F(UnwindRangeWindowTest, startsAReversedRangeAtItsLastElement) { + expectRows("UNWIND reverse(range(1, 10, 4)) AS x RETURN x", {{"9"}, {"5"}, {"1"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsNothingForAReversedEmptyRange) { + expectRows("UNWIND reverse(range(5, 1)) AS x RETURN count(*)", {{"0"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsNothingForAReversedNullRange) { + expectRows("UNWIND reverse(range(1, null)) AS x RETURN count(*)", {{"0"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsTwiceReversedRangeInOrder) { + expectRows("UNWIND reverse(reverse(range(1, 3))) AS x RETURN x", {{"1"}, {"2"}, {"3"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsASlice) { + expectRows("UNWIND range(1, 10)[2..5] AS x RETURN x", {{"3"}, {"4"}, {"5"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsASliceFromTheStart) { + expectRows("UNWIND range(1, 10)[..3] AS x RETURN x", {{"1"}, {"2"}, {"3"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsASliceToTheEnd) { + expectRows("UNWIND range(1, 10)[7..] AS x RETURN x", {{"8"}, {"9"}, {"10"}}); +} + +TEST_F(UnwindRangeWindowTest, countsANegativeLowerBoundFromTheEnd) { + expectRows("UNWIND range(1, 10)[-3..] AS x RETURN x", {{"8"}, {"9"}, {"10"}}); +} + +TEST_F(UnwindRangeWindowTest, countsANegativeUpperBoundFromTheEnd) { + expectRows("UNWIND range(1, 10)[..-8] AS x RETURN x", {{"1"}, {"2"}}); +} + +TEST_F(UnwindRangeWindowTest, clampsALowerBoundBeforeTheStart) { + expectRows("UNWIND range(1, 10)[-20..2] AS x RETURN x", {{"1"}, {"2"}}); +} + +TEST_F(UnwindRangeWindowTest, clampsAnUpperBoundPastTheEnd) { + expectRows("UNWIND range(1, 10)[5..100] AS x RETURN x", {{"6"}, {"7"}, {"8"}, {"9"}, {"10"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsNothingForCrossedBounds) { + expectRows("UNWIND range(1, 10)[5..2] AS x RETURN count(*)", {{"0"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsNothingForAnUpperBoundOfZero) { + expectRows("UNWIND range(1, 10)[0..0] AS x RETURN count(*)", {{"0"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsNothingForALowerBoundAtTheEnd) { + expectRows("UNWIND range(1, 10)[10..] AS x RETURN count(*)", {{"0"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsTheFirstElementForTheWholeNegativeSpan) { + expectRows("UNWIND range(1, 10)[-10..1] AS x RETURN x", {{"1"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsNothingForANullLowerBound) { + expectRows("UNWIND range(1, 10)[null..3] AS x RETURN count(*)", {{"0"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsNothingForANullUpperBound) { + expectRows("UNWIND range(1, 10)[1..null] AS x RETURN count(*)", {{"0"}}); +} + +TEST_F(UnwindRangeWindowTest, slicesASteppedRange) { + expectRows("UNWIND range(0, 100, 10)[2..4] AS x RETURN x", {{"20"}, {"30"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsTheTail) { + expectRows("UNWIND tail(range(1, 4)) AS x RETURN x", {{"2"}, {"3"}, {"4"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsNothingForTheTailOfOneElement) { + expectRows("UNWIND tail(range(1, 1)) AS x RETURN count(*)", {{"0"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsNothingForTheTailOfAnEmptyRange) { + expectRows("UNWIND tail(range(1, 0)) AS x RETURN count(*)", {{"0"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsAReversedSlice) { + expectRows("UNWIND reverse(range(1, 10)[2..5]) AS x RETURN x", {{"5"}, {"4"}, {"3"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsAReversedTail) { + expectRows("UNWIND reverse(tail(range(1, 4))) AS x RETURN x", {{"4"}, {"3"}, {"2"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsASliceOfAReversedRange) { + expectRows("UNWIND reverse(range(1, 10))[2..5] AS x RETURN x", {{"8"}, {"7"}, {"6"}}); +} + +TEST_F(UnwindRangeWindowTest, unwindsASliceOfASlice) { + expectRows("UNWIND range(1, 10)[2..5][1..] AS x RETURN x", {{"4"}, {"5"}}); +} + +TEST_F(UnwindRangeWindowTest, slicesByTheOuterElement) { + expectRows("UNWIND range(1, 3) AS i UNWIND range(1, 5)[..i] AS x RETURN i, x", + {{"1", "1"}, {"2", "1"}, {"2", "2"}, {"3", "1"}, {"3", "2"}, {"3", "3"}}); +} + +TEST_F(UnwindRangeWindowTest, walksTheWholeInt64RangeBackwardsUpToALimit) { + expectRows("UNWIND reverse(range(-9223372036854775808, 9223372036854775807)) AS x RETURN x LIMIT 2", + {{"9223372036854775807"}, {"9223372036854775806"}}); +} + +TEST_F(UnwindRangeWindowTest, slicesTheEndOfTheWholeInt64Range) { + expectRows("UNWIND range(-9223372036854775808, 9223372036854775807)[-2..] AS x RETURN x", + {{"9223372036854775806"}, {"9223372036854775807"}}); +} + +TEST_F(UnwindRangeWindowTest, slicesTheMiddleOfTheWholeInt64Range) { + expectRows("UNWIND range(-9223372036854775808, 9223372036854775807)[9223372036854775807..] AS x RETURN x LIMIT 1", + {{"-1"}}); +} + +TEST_F(UnwindRangeWindowTest, skipsPastAChunkOfAReversedRange) { + expectRows("UNWIND reverse(range(1, 5000)) AS x RETURN x SKIP 2048 LIMIT 2", {{"2952"}, {"2951"}}); +} + +TEST_F(UnwindRangeWindowTest, reversesASliceLongerThanAChunk) { + expectSortedRows("UNWIND range(1, 3) AS i UNWIND reverse(range(1, 5000)[1000..]) AS x RETURN i, count(x), sum(x)", + {{"1", "4000", "12002000"}, {"2", "4000", "12002000"}, {"3", "4000", "12002000"}}); +}