From 551656698e65f0f7d2262b7b74fa2c227f8656e1 Mon Sep 17 00:00:00 2001 From: Suhas Narreddy Date: Wed, 11 Nov 2026 03:29:18 +0100 Subject: [PATCH] Cypher: implement coll.sortMulti --- query/AST/FunctionDecls.cpp | 14 ++ query/ir/codegen/DBProgramGenerator.cpp | 26 +++ query/ir/codegen/DBProgramGenerator.h | 1 + query/ir/dialect/db/DBHeldTags.cpp | 44 ++++- query/ir/dialect/db/DBHeldTags.h | 9 + query/ir/dialect/db/DBOps.cpp | 73 +++++++ query/ir/dialect/db/DBOps.td | 37 ++++ query/ir/dialect/nl/NLOps.cpp | 33 ++++ query/ir/dialect/nl/NLOps.td | 34 ++++ query/ir/exec/NLExecutor.cpp | 164 ++++++++++++++++ query/ir/exec/NLExecutor.h | 4 + query/ir/exec/NLProgram.h | 45 +++++ query/ir/exec/NLTranslator.cpp | 26 +++ query/ir/exec/NLTranslator.h | 1 + query/ir/lowering/DBLowering.cpp | 122 ++++++++---- query/ir/lowering/DBLowering.h | 2 + storage/list/ListBuffer.cpp | 17 +- storage/list/ListBuffer.h | 9 + test/query/ir/journey/JourneyParetoTest.cpp | 5 +- test/query/ir/list/CMakeLists.txt | 1 + test/query/ir/list/CollSortMultiTest.cpp | 183 ++++++++++++++++++ test/query/ir/passes/CMakeLists.txt | 2 + .../passes/SortMultiSpeculatabilityTest.cpp | 181 +++++++++++++++++ 23 files changed, 982 insertions(+), 51 deletions(-) create mode 100644 test/query/ir/list/CollSortMultiTest.cpp create mode 100644 test/query/ir/passes/SortMultiSpeculatabilityTest.cpp diff --git a/query/AST/FunctionDecls.cpp b/query/AST/FunctionDecls.cpp index c52a34bc63..18c553f240 100644 --- a/query/AST/FunctionDecls.cpp +++ b/query/AST/FunctionDecls.cpp @@ -988,6 +988,20 @@ void FunctionDecls::initCollectionFunctions() { indexOf->setReturnTypes({{EvaluatedType::Integer}}); indexOf->setSearchesItsList(true); } + + const PossibleArgTypes listOrCell {EvaluatedType::List, EvaluatedType::Any}; + const PossibleArgTypes integerOrCell {EvaluatedType::Integer, EvaluatedType::Any}; + + std::vector sortMultis; + coll->createOverloads("sortMulti", {listOrCell, listOrCell, integerOrCell, integerOrCell}, sortMultis); + + for (FunctionSignature* sortMulti : sortMultis) { + const bool sortsAList = sortMulti->argumentTypes().front().getType() == EvaluatedType::List; + + sortMulti->setRequiredArgCount(1); + sortMulti->setReturnTypes({{EvaluatedType::List}}); + sortMulti->setReturnsItsArgumentShape(sortsAList); + } } void FunctionDecls::initSpatialFunctions() { diff --git a/query/ir/codegen/DBProgramGenerator.cpp b/query/ir/codegen/DBProgramGenerator.cpp index 33b0f53180..3daa3484c2 100644 --- a/query/ir/codegen/DBProgramGenerator.cpp +++ b/query/ir/codegen/DBProgramGenerator.cpp @@ -10445,6 +10445,9 @@ void DBProgramGenerator::translateFunctionExpr(const Expr* expr, } else if (funcName == "replace") { translateReplace(expr, args); return; + } else if (funcName == "coll.sortMulti") { + translateSortMulti(expr, args); + return; } else if (funcName == "point.distance") { translatePointDistance(expr, args); return; @@ -10623,6 +10626,29 @@ void DBProgramGenerator::translateSubstring(const Expr* expr, const ExprChain* a _part._exprMap[expr] = substring.getResult(); } +void DBProgramGenerator::translateSortMulti(const Expr* expr, const ExprChain* args) { + bioassert(args && !args->empty() && args->size() <= 4, "coll.sortMulti() takes 1 to 4 arguments."); + + const ExprChain::ExprVector& argExprs = args->getExprs(); + const size_t argCount = argExprs.size(); + + const mlir::Value list = translateArg(argExprs[0]); + const mlir::Value fields = argCount > 1 ? translateArg(argExprs[1]) : mlir::Value {}; + const mlir::Value limit = argCount > 2 ? translateArg(argExprs[2]) : mlir::Value {}; + const mlir::Value skip = argCount > 3 ? translateArg(argExprs[3]) : mlir::Value {}; + + const mlir::db::ColumnType noneType = allocColumnType(mlir::NoneType::get(_mlirCtxt)); + + mlir::db::SortMulti sortMulti = _opBuilder.create(_opBuilder.getUnknownLoc(), + noneType, + list, + fields, + limit, + skip); + + _part._exprMap[expr] = sortMulti.getResult(); +} + void DBProgramGenerator::translateSplit(const Expr* expr, const ExprChain* args) { bioassert(args && args->size() == 2, "split() takes 2 arguments."); diff --git a/query/ir/codegen/DBProgramGenerator.h b/query/ir/codegen/DBProgramGenerator.h index 209100f1e1..48e32bc6d6 100644 --- a/query/ir/codegen/DBProgramGenerator.h +++ b/query/ir/codegen/DBProgramGenerator.h @@ -1391,6 +1391,7 @@ class DBProgramGenerator { void translateRange(const Expr* expr, const ExprChain* args); void translateSubstring(const Expr* expr, const ExprChain* args); + void translateSortMulti(const Expr* expr, const ExprChain* args); void translateSplit(const Expr* expr, const ExprChain* args); void translateReplace(const Expr* expr, const ExprChain* args); void translatePointDistance(const Expr* expr, const ExprChain* args); diff --git a/query/ir/dialect/db/DBHeldTags.cpp b/query/ir/dialect/db/DBHeldTags.cpp index 87f8ded931..9c7da62797 100644 --- a/query/ir/dialect/db/DBHeldTags.cpp +++ b/query/ir/dialect/db/DBHeldTags.cpp @@ -8,6 +8,7 @@ #include "mlir/IR/BuiltinTypes.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" #include "DBColumnElement.h" @@ -29,6 +30,15 @@ TagMask tagBit(CellTag tag) { return TagMask {1} << static_cast(tag); } +TagMask tagMaskOf(std::initializer_list tags) { + TagMask mask = 0; + for (const CellTag tag : tags) { + mask |= tagBit(tag); + } + + return mask; +} + CellTag integerTag(mlir::IntegerType integer) { if (integer.getWidth() == 1) { return CellTag::Bool; @@ -133,8 +143,13 @@ class HeldTagWalk { _steps.push_back(ReadStep {}); } - bool holdsOnlyTags(mlir::Value column) { - reach(column, NO_READ); + bool holdsOnlyTags(mlir::Value column, llvm::ArrayRef reads) { + size_t path = NO_READ; + for (const mlir::StringAttr read : llvm::reverse(reads)) { + path = prepend(read, path); + } + + reach(column, path); while (!_worklist.empty()) { const Reach next = _worklist.pop_back_val(); @@ -246,6 +261,8 @@ class HeldTagWalk { return walkMadeMap(map, path); } else if (mlir::db::MergeMaps merge = mlir::dyn_cast_or_null(made)) { return walkMergedMaps(merge, path); + } else if (mlir::db::SortMulti sort = mlir::dyn_cast_or_null(made)) { + return walkSortedList(sort, path); } else { return walkType(element, path); } @@ -270,6 +287,16 @@ class HeldTagWalk { return true; } + // A sort holds some of the elements of the list it sorts + bool walkSortedList(mlir::db::SortMulti sort, size_t path) { + if (path == NO_READ) { + return holds(CellTag::ListView); + } + + reach(sort.getList(), path); + return true; + } + // A key missing from the map reads null, and the map iterated is a list of one bool walkMadeMap(mlir::db::MakeMap map, size_t path) { if (path == NO_READ) { @@ -377,12 +404,13 @@ class HeldTagWalk { } bool db::holdsOnlyNullOr(mlir::Value column, std::initializer_list tags) { - TagMask mask = 0; - for (const CellTag tag : tags) { - mask |= tagBit(tag); - } + return readsHoldOnlyNullOr(column, {}, tags); +} - HeldTagWalk walk(mask); +bool db::readsHoldOnlyNullOr(mlir::Value column, + llvm::ArrayRef reads, + std::initializer_list tags) { + HeldTagWalk walk(tagMaskOf(tags)); - return walk.holdsOnlyTags(column); + return walk.holdsOnlyTags(column, reads); } diff --git a/query/ir/dialect/db/DBHeldTags.h b/query/ir/dialect/db/DBHeldTags.h index 419182ed97..ee1d70ec4f 100644 --- a/query/ir/dialect/db/DBHeldTags.h +++ b/query/ir/dialect/db/DBHeldTags.h @@ -2,10 +2,13 @@ #include +#include "llvm/ADT/ArrayRef.h" + #include "CellTag.h" namespace mlir { class Value; +class StringAttr; } namespace db { @@ -15,4 +18,10 @@ namespace db { // carry it, and from a key or an element read off it to the value it was read out of. bool holdsOnlyNullOr(mlir::Value column, std::initializer_list tags); +// The same of what @param reads, taken in order off each row's value, arrive at: the value +// under a key of a map, or each element of a list where the key is null +bool readsHoldOnlyNullOr(mlir::Value column, + llvm::ArrayRef reads, + std::initializer_list tags); + } diff --git a/query/ir/dialect/db/DBOps.cpp b/query/ir/dialect/db/DBOps.cpp index 9d8c38d29a..e9d0a68fcd 100644 --- a/query/ir/dialect/db/DBOps.cpp +++ b/query/ir/dialect/db/DBOps.cpp @@ -2547,6 +2547,24 @@ LogicalResult Range::verify() { return success(); } +LogicalResult SortMulti::verify() { + const Value fields = getFields(); + const Value limit = getLimit(); + const Value skip = getSkip(); + + const bool readsLists = isColumnOf(getList().getType()) + && (!fields || isColumnOf(fields.getType())); + const bool readsIntegerBounds = (!limit || isIntegerColumn(limit)) && (!skip || isIntegerColumn(skip)); + + if (!readsLists) { + return emitOpError("sorts a column of lists on a column of lists of fields"); + } else if (!readsIntegerBounds) { + return emitOpError("limit and skip must be integer columns"); + } + + return success(); +} + LogicalResult Replace::verify() { const bool readsStrings = llvm::all_of(getOperation()->getOperands(), [](Value operand) { return isColumnOf(operand.getType()); @@ -3287,6 +3305,61 @@ Speculation::Speculatability ListSlice::getSpeculatability() { return speculatableWhen(holdsOnlyLists(getList()) && readsIntegerBounds); } +// A key is sorted on in the order <=> gives cells, which raises on an embedding, and a list +// or a map may hold one at any depth. With no field the list is kept as it is. +Speculation::Speculatability SortMulti::getSpeculatability() { + const Value list = getList(); + const Value fields = getFields(); + const Value limit = getLimit(); + const Value skip = getSkip(); + + int64_t skipped = 0; + const bool skipsALiteralCount = !skip || (matchIntegerLiteral(skip, skipped) && skipped >= 0); + const bool readsIntegerLimit = !limit || holdsOnlyIntegers(limit); + + ConstantOp fieldsLiteral = fields ? fields.getDefiningOp() : ConstantOp {}; + const ArrayAttr names = fieldsLiteral ? dyn_cast(fieldsLiteral.getValue()) : ArrayAttr {}; + + if (!skipsALiteralCount || !readsIntegerLimit || !holdsOnlyLists(list)) { + return Speculation::NotSpeculatable; + } else if (!fields || (names && names.empty())) { + return Speculation::Speculatable; + } else if (!names || !::db::readsHoldOnlyNullOr(list, {StringAttr {}}, {::db::CellTag::MapView})) { + return Speculation::NotSpeculatable; + } + + for (const Attribute name : names) { + const StringAttr field = dyn_cast(name); + if (!field) { + return Speculation::NotSpeculatable; + } + + const StringRef text = field.getValue(); + const StringAttr key = StringAttr::get(getContext(), text.starts_with("^") ? text.drop_front() : text); + const StringAttr keyOfEachElement[] = {StringAttr {}, key}; + + const bool keyHasAnOrder = ::db::readsHoldOnlyNullOr(list, + keyOfEachElement, + {::db::CellTag::Int, + ::db::CellTag::UInt, + ::db::CellTag::Double, + ::db::CellTag::String, + ::db::CellTag::Bool, + ::db::CellTag::DateTime, + ::db::CellTag::Duration, + ::db::CellTag::Point, + ::db::CellTag::NodeID, + ::db::CellTag::EdgeID, + ::db::CellTag::Path}); + + if (!keyHasAnOrder) { + return Speculation::NotSpeculatable; + } + } + + return Speculation::Speculatable; +} + Speculation::Speculatability ListIndex::getSpeculatability() { return speculatableWhen(holdsOnlyLists(getLhs())); } diff --git a/query/ir/dialect/db/DBOps.td b/query/ir/dialect/db/DBOps.td index 9b5dc7e988..fd6318e5cf 100644 --- a/query/ir/dialect/db/DBOps.td +++ b/query/ir/dialect/db/DBOps.td @@ -684,6 +684,43 @@ def ListSlice : TuringOp<"list_slice", [MayRaise, ConstantThroughOperands, AttrS }]; } +/** +* SortMulti +*/ +def SortMulti : TuringOp<"sort_multi", [MayRaise, ConstantThroughOperands, AttrSizedOperandSegments, DeclareOpInterfaceMethods]> { + let summary = "Sort each row's list of maps on the values of some keys"; + + let description = [{ + `coll.sortMulti(maps, fields, limit, skip)`, APOC's `apoc.coll.sortMulti`. The op reads + one list column and produces one list cell per row: + + %sorted = db.sort_multi(%maps) fields %fields limit %limit skip %skip + : (!db.column, !db.column>, + !db.column, !db.column) -> !db.column + + Row r's maps are sorted on the keys row r's fields name, the first one first: a field + starting with `^` sorts ascending on the key after it, any other descending. Two values + compare as ORDER BY compares them, and a missing key is a null, which APOC orders before + every value: first ascending, last descending. Maps that tie keep their order. The sorted maps past the first `skip` are + kept, at most `limit` of them, all where `limit` is negative. + + An operand the query left out is absent: no field keeps the list as it is, whatever it + holds, no limit keeps every map and no skip keeps the first. A row with a null operand + gets a null cell. Sorted on a field, a null element holds no key, as `null.a` is null, + and an element that is no map raises, as do a field that is no string and a negative skip. + The result element type is the input's and is resolved during lowering. + }]; + + let arguments = (ins Column:$list, Optional:$fields, Optional:$limit, Optional:$skip); + let results = (outs Column:$result); + let hasVerifier = 1; + + let assemblyFormat = [{ + `(` $list `)` (`fields` $fields^)? (`limit` $limit^)? (`skip` $skip^)? attr-dict `:` + functional-type(operands, results) + }]; +} + /** * PercentileCont / PercentileDisc */ diff --git a/query/ir/dialect/nl/NLOps.cpp b/query/ir/dialect/nl/NLOps.cpp index 33c506cfa9..5442d6c733 100644 --- a/query/ir/dialect/nl/NLOps.cpp +++ b/query/ir/dialect/nl/NLOps.cpp @@ -1013,6 +1013,39 @@ LogicalResult DynamicMapKey::verify() { return success(); } +LogicalResult SortMulti::verify() { + const Value fields = getFields(); + const Value limit = getLimit(); + const Value skip = getSkip(); + + const auto readsAList = [](Value chunk) { + return isa(valueElementOf(chunk.getType())); + }; + + const auto readsAnInteger = [](Value chunk) { + return isa(valueElementOf(chunk.getType())); + }; + + const Type list = valueElementOf(getList().getType()); + const Type sorted = isa(list) ? storage::ListType::get(getContext(), list) : list; + const Type result = cast(getResult().getType()).getElementType(); + const auto nullableResult = dyn_cast(result); + + const bool readsLists = readsAList(getList()) && (!fields || readsAList(fields)); + const bool readsIntegerBounds = (!limit || readsAnInteger(limit)) && (!skip || readsAnInteger(skip)); + const bool answersTheInputList = nullableResult && nullableResult.getValueType() == sorted; + + if (!readsLists) { + return emitOpError("sorts a chunk of lists or of tagged cells on one of lists of fields, optionally nullable"); + } else if (!readsIntegerBounds) { + return emitOpError("reads its limit and its skip out of chunks of integers or of tagged cells"); + } else if (!answersTheInputList) { + return emitOpError("result must be a chunk of nullable lists of the input list's elements"); + } + + return success(); +} + LogicalResult ContainerIndex::verify() { const Type base = valueElementOf(getBase().getType()); if (!isa(base)) { diff --git a/query/ir/dialect/nl/NLOps.td b/query/ir/dialect/nl/NLOps.td index 920a91a72a..75eb81f60d 100644 --- a/query/ir/dialect/nl/NLOps.td +++ b/query/ir/dialect/nl/NLOps.td @@ -516,6 +516,40 @@ def ListSlice : NLOp<"list_slice", [Pure, RowAlignedThroughOperands, AttrSizedOp }]; } +/** +* SortMulti +*/ +def SortMulti : NLOp<"sort_multi", [Pure, RowAlignedThroughOperands, AttrSizedOperandSegments]> { + let summary = "Row-wise sort of a list of maps on the values of some keys"; + + let description = [{ + Imperative counterpart of db.sort_multi. Row r of the result holds the maps of row r's + list in the order row r's fields sort them, cut by its skip and its limit. A sorted list + lives in the step's own buffer, which the next run of the step clears; with no field the + row is a slice of its input list, which copies nothing. + + %sorted = nl.sort_multi %maps fields %fields limit %limit skip %skip + : (!nl.chunk>>, + !nl.chunk>>, + !nl.chunk>, + !nl.chunk>) + -> !nl.chunk>> + + Every operand carries rows - lowering lays a constant one out over them first. The + limit and the skip are read as nullable value chunks, or as chunks of type-erased + cells whose integer each row reads out of its cell. + }]; + + let arguments = (ins NLChunk:$list, Optional:$fields, Optional:$limit, Optional:$skip); + let results = (outs NLChunk:$result); + let hasVerifier = 1; + + let assemblyFormat = [{ + $list (`fields` $fields^)? (`limit` $limit^)? (`skip` $skip^)? attr-dict `:` + functional-type(operands, results) + }]; +} + /** * PercentileCont / PercentileDisc */ diff --git a/query/ir/exec/NLExecutor.cpp b/query/ir/exec/NLExecutor.cpp index d759eab413..a8da13d88d 100644 --- a/query/ir/exec/NLExecutor.cpp +++ b/query/ir/exec/NLExecutor.cpp @@ -5944,6 +5944,65 @@ std::optional optConstTaggedListRead(const Column* input, size_t) { return cellListRead(*cell); } +struct SortField { + std::string_view _key; + bool _ascending {false}; +}; + +void readSortFields(const ListView names, std::vector& fields) { + fields.clear(); + + for (const CellView name : names) { + if (name.getTag() != CellTag::String) { + throw IRException("coll.sortMulti sorts on fields named by strings"); + } + + const std::string_view text = name.getAs(); + const bool ascending = text.starts_with('^'); + + fields.push_back(SortField {._key = ascending ? text.substr(1) : text, ._ascending = ascending}); + } +} + +std::optional sortKey(const MapView map, const std::string_view key) { + MapEntryView entry; + if (!findMapEntry(map, key, entry)) { + return std::nullopt; + } + + const CellView value = entry.getValue(); + if (value.getTag() == CellTag::Null) { + return std::nullopt; + } + + return value; +} + +// A missing key is a null, which APOC's sortMulti orders before every value +std::strong_ordering orderOnKeys(std::span> lhs, + std::span> rhs, + std::span fields) { + for (size_t field = 0; field < fields.size(); field++) { + const std::optional& lhsKey = lhs[field]; + const std::optional& rhsKey = rhs[field]; + + std::strong_ordering order = std::strong_ordering::equal; + if (lhsKey && rhsKey) { + order = *lhsKey <=> *rhsKey; + } else if (lhsKey) { + order = std::strong_ordering::greater; + } else if (rhsKey) { + order = std::strong_ordering::less; + } + + if (order != 0) { + return fields[field]._ascending ? order : 0 <=> order; + } + } + + return std::strong_ordering::equal; +} + template NLContainer listContainer(const Column* input, size_t row) { const std::optional list = (*static_cast(input))[row]; @@ -9760,6 +9819,111 @@ void NLExecutor::runListSlice(NLExecutionContext*, NLFunctionData* data) { } } +void NLExecutor::runSortMulti(NLExecutionContext*, NLFunctionData* data) { + const NLSortMultiData* sortMulti = static_cast(data); + + const Column* const list = sortMulti->getList(); + const NLListReadFunction listRead = sortMulti->getListRead(); + const Column* const fields = sortMulti->getFields(); + const NLListReadFunction fieldsRead = sortMulti->getFieldsRead(); + const NLRangeBound& limit = sortMulti->getLimit(); + const NLRangeBound& skip = sortMulti->getSkip(); + + NLStepBuffers* const stepBuffers = sortMulti->getStepBuffers(); + stepBuffers->clear(); + + ListBuffer<>& listBuffer = stepBuffers->listBuffer(); + + const size_t rowCount = list->size(); + + std::vector>& outputRaw = + static_cast*>(sortMulti->getResult())->getRaw(); + outputRaw.resize(rowCount); + + std::vector sortFields; + std::span parsedNames; + std::vector> sortKeys; + std::vector order; + std::vector sorted; + + for (size_t rowIndex = 0; rowIndex < rowCount; rowIndex++) { + const std::optional elements = listRead(list, rowIndex); + const std::optional names = fields ? fieldsRead(fields, rowIndex) : std::optional {ListView {}}; + + const std::optional kept = + limit._column ? limit._read(limit._column, rowIndex) : std::optional {-1}; + + const std::optional skipped = + skip._column ? skip._read(skip._column, rowIndex) : std::optional {0}; + + const bool readsANull = !elements || !names || !kept || !skipped; + if (readsANull) { + outputRaw[rowIndex] = std::nullopt; + continue; + } else if (*skipped < 0) { + throw IRException("coll.sortMulti cannot skip a negative count"); + } + + const std::span nameElements = names->elements(); + const bool namesOtherFields = nameElements.data() != parsedNames.data() || nameElements.size() != parsedNames.size(); + if (namesOtherFields) { + readSortFields(*names, sortFields); + parsedNames = nameElements; + } + + const std::span maps = elements->elements(); + const size_t size = maps.size(); + const size_t fieldCount = sortFields.size(); + + const size_t first = std::min(size, static_cast(*skipped)); + const size_t remaining = size - first; + const size_t count = *kept < 0 ? remaining : std::min(remaining, static_cast(*kept)); + + if (fieldCount == 0) { + outputRaw[rowIndex] = elements->slice(first, count); + continue; + } + + sortKeys.resize(size * fieldCount); + + for (size_t position = 0; position < size; position++) { + const CellView element = maps[position]; + const CellTag tag = element.getTag(); + const std::span> keysOfElement {sortKeys.data() + position * fieldCount, fieldCount}; + + // A key read off null is null, as `null.a` is + if (tag == CellTag::Null) { + std::ranges::fill(keysOfElement, std::nullopt); + continue; + } else if (tag != CellTag::MapView) { + throw IRException(fmt::format("coll.sortMulti sorts a list of maps, not of '{}'", queryTypeName(tag))); + } + + const MapView map = element.getAs(); + for (size_t field = 0; field < fieldCount; field++) { + keysOfElement[field] = sortKey(map, sortFields[field]._key); + } + } + + order.resize(size); + std::iota(order.begin(), order.end(), size_t {0}); + + const std::span> keys {sortKeys}; + std::ranges::stable_sort(order, [&keys, &sortFields, fieldCount](size_t lhs, size_t rhs) { + return orderOnKeys(keys.subspan(lhs * fieldCount, fieldCount), + keys.subspan(rhs * fieldCount, fieldCount), + sortFields) < 0; + }); + + sorted.resize(count); + for (size_t position = 0; position < count; position++) { + sorted[position] = maps[order[first + position]]; + } + + outputRaw[rowIndex] = listBuffer.gather(sorted); + } +} + void NLExecutor::runPercentileCont(NLExecutionContext*, NLFunctionData* data) { const auto interpolate = [](const std::vector& sorted, double percentile) { return std::optional(continuousPercentile(sorted, percentile)); diff --git a/query/ir/exec/NLExecutor.h b/query/ir/exec/NLExecutor.h index 11592a2f16..542900be94 100644 --- a/query/ir/exec/NLExecutor.h +++ b/query/ir/exec/NLExecutor.h @@ -460,6 +460,10 @@ class NLExecutor { // its bounds, which copies nothing - a slice of a list is a run of the same elements. static void runListSlice(NLExecutionContext* context, NLFunctionData* data); + // Sort one list of maps per row (nl.sort_multi) on the keys its fields name, into the + // step's own buffer, which the next run clears. With no field the row is a slice of its list. + static void runSortMulti(NLExecutionContext* context, NLFunctionData* data); + // Read one percentile per row (nl.percentile_cont / nl.percentile_disc) of the numbers // its list holds. The discrete one answers a value of the list as @tparam Result: a // nullable number where the list names its number type, a tagged cell where it does not. diff --git a/query/ir/exec/NLProgram.h b/query/ir/exec/NLProgram.h index c4c8bd541e..e531dc4fae 100644 --- a/query/ir/exec/NLProgram.h +++ b/query/ir/exec/NLProgram.h @@ -4778,6 +4778,51 @@ class NLListSliceData : public NLFunctionData { Column* _result {nullptr}; }; +// nl.sort_multi data: the list column and the fields column, each with the read telling an +// absent cell from a list, the limit and the skip with the read their column shape takes, +// the buffers the sorted lists are written into and the column the step fills. An operand +// the query left out has no column. +class NLSortMultiData : public NLFunctionData { +public: + NLSortMultiData(const Column* list, + NLListReadFunction listRead, + const Column* fields, + NLListReadFunction fieldsRead, + const NLRangeBound& limit, + const NLRangeBound& skip, + NLStepBuffers* stepBuffers, + Column* result) + : _list(list), + _listRead(listRead), + _fields(fields), + _fieldsRead(fieldsRead), + _limit(limit), + _skip(skip), + _stepBuffers(stepBuffers), + _result(result) + { + } + + const Column* getList() const { return _list; } + NLListReadFunction getListRead() const { return _listRead; } + const Column* getFields() const { return _fields; } + NLListReadFunction getFieldsRead() const { return _fieldsRead; } + const NLRangeBound& getLimit() const { return _limit; } + const NLRangeBound& getSkip() const { return _skip; } + NLStepBuffers* getStepBuffers() const { return _stepBuffers; } + Column* getResult() const { return _result; } + +private: + const Column* _list {nullptr}; + NLListReadFunction _listRead {nullptr}; + const Column* _fields {nullptr}; + NLListReadFunction _fieldsRead {nullptr}; + NLRangeBound _limit; + NLRangeBound _skip; + NLStepBuffers* _stepBuffers {nullptr}; + Column* _result {nullptr}; +}; + // 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 { diff --git a/query/ir/exec/NLTranslator.cpp b/query/ir/exec/NLTranslator.cpp index c81c3e4c9b..9e0035cc6b 100644 --- a/query/ir/exec/NLTranslator.cpp +++ b/query/ir/exec/NLTranslator.cpp @@ -1432,6 +1432,8 @@ void NLTranslator::translateBlock(mlir::Block& block, NLStmtContainer* body) { translateMakeList(makeList, body); } else if (nl::ListSlice listSlice = mlir::dyn_cast(operation)) { translateListSlice(listSlice, body); + } else if (nl::SortMulti sortMulti = mlir::dyn_cast(operation)) { + translateSortMulti(sortMulti, body); } else if (mlir::isa(operation)) { translatePercentile(&operation, body); } else if (nl::MakeMap makeMap = mlir::dyn_cast(operation)) { @@ -4341,6 +4343,30 @@ void NLTranslator::translateListSlice(nl::ListSlice slice, NLStmtContainer* body body->emplaceStmt(&NLExecutor::runListSlice, data); } +void NLTranslator::translateSortMulti(nl::SortMulti sortMulti, NLStmtContainer* body) { + const mlir::Value resultValue = sortMulti.getResult(); + + Column* const result = allocColumnForChunkType(resultValue.getType()); + _valueSlots[resultValue] = result; + + const Column* const list = getColumn(sortMulti.getList()); + + const mlir::Value fieldsValue = sortMulti.getFields(); + const Column* const fields = fieldsValue ? getColumn(fieldsValue) : nullptr; + const NLListReadFunction fieldsRead = fields ? NLExecutor::selectListRead(fields) : nullptr; + + NLSortMultiData* data = _program->allocFunctionData(list, + NLExecutor::selectListRead(list), + fields, + fieldsRead, + rangeBound(sortMulti.getLimit()), + rangeBound(sortMulti.getSkip()), + _program->allocStepBuffers(), + result); + + body->emplaceStmt(&NLExecutor::runSortMulti, data); +} + void NLTranslator::translatePercentile(mlir::Operation* operation, NLStmtContainer* body) { const mlir::Value resultValue = operation->getResult(0); const mlir::Value percentileValue = operation->getOperand(1); diff --git a/query/ir/exec/NLTranslator.h b/query/ir/exec/NLTranslator.h index eb215a5bb6..b5d2defdc4 100644 --- a/query/ir/exec/NLTranslator.h +++ b/query/ir/exec/NLTranslator.h @@ -1157,6 +1157,7 @@ class NLTranslator { // Translate an nl.list_slice: allocate the column of views the step fills, and bind // the reads its list and its bounds are taken through void translateListSlice(mlir::nl::ListSlice slice, NLStmtContainer* body); + void translateSortMulti(mlir::nl::SortMulti sortMulti, NLStmtContainer* body); void translatePercentile(mlir::Operation* operation, NLStmtContainer* body); // The read one element column of an nl.make_list contributes its cell through, chosen diff --git a/query/ir/lowering/DBLowering.cpp b/query/ir/lowering/DBLowering.cpp index cfebb0f10c..085890b127 100644 --- a/query/ir/lowering/DBLowering.cpp +++ b/query/ir/lowering/DBLowering.cpp @@ -50,6 +50,31 @@ namespace storage = mlir::storage; namespace { +void collectPresentChunks(llvm::ArrayRef chunks, llvm::SmallVectorImpl& present) { + present.clear(); + + for (const mlir::Value chunk : chunks) { + if (chunk) { + present.push_back(chunk); + } + } +} + +// A list over the elements @param listChunk's lists hold, read as nullable: a row with no +// list gets none. A type-erased cell holds a list of cells in its turn. +nl::ChunkType listOfSameElementsChunkType(mlir::MLIRContext* context, mlir::Value listChunk) { + const mlir::Type listElement = mlir::cast(listChunk.getType()).getElementType(); + const auto nullable = mlir::dyn_cast(listElement); + const mlir::Type value = nullable ? nullable.getValueType() : listElement; + + const mlir::Type list = + mlir::isa(value) + ? storage::ListType::get(context, storage::AnyType::get(context)) + : value; + + return nl::ChunkType::get(context, storage::NullableType::get(context, list)); +} + bool isDefinedWithin(mlir::Value value, mlir::Operation* operation) { mlir::Operation* const definingOp = value.getDefiningOp(); mlir::Operation* const definer = definingOp ? definingOp : value.getParentBlock()->getParentOp(); @@ -1571,6 +1596,8 @@ void DBLowering::lowerOperationByKind(mlir::Operation& operation) { lowerPointFunction(&operation, _builder.getI1Type()); } else if (mlir::db::ListSlice listSlice = mlir::dyn_cast(operation)) { lowerListSlice(listSlice); + } else if (mlir::db::SortMulti sortMulti = mlir::dyn_cast(operation)) { + lowerSortMulti(sortMulti); } else if (mlir::isa(operation)) { lowerPercentile(&operation); } else if (mlir::db::ListComprehension listComprehension = mlir::dyn_cast(operation)) { @@ -2687,71 +2714,92 @@ void DBLowering::lowerMakeList(mlir::db::MakeList makeList) { _valueMap[makeList.getResult()] = lists.getResult(); } +void DBLowering::rowAlignOptionalChunks(llvm::MutableArrayRef chunks) { + llvm::SmallVector present; + collectPresentChunks(chunks, present); + + const mlir::Value cardinality = cardinalityDriver(present); + + for (mlir::Value& chunk : chunks) { + if (chunk) { + chunk = rowAlignedChunk(chunk, cardinality); + } + } +} + void DBLowering::lowerListSlice(mlir::db::ListSlice slice) { const mlir::Location loc = _builder.getUnknownLoc(); mlir::MLIRContext* const context = _builder.getContext(); - mlir::Value listChunk = mapValue(slice.getList()); + const mlir::Value from = slice.getFrom(); + const mlir::Value to = slice.getTo(); - llvm::SmallVector operands {listChunk}; + // A list or a bound holding one value for every row rather than one per row is laid + // out over the rows the others carry, as lowerRange lays its bounds out + std::array chunks { + mapValue(slice.getList()), + from ? mapValue(from) : mlir::Value {}, + to ? mapValue(to) : mlir::Value {}, + }; + rowAlignOptionalChunks(chunks); - mlir::Value fromChunk = slice.getFrom() ? mapValue(slice.getFrom()) : mlir::Value {}; - mlir::Value toChunk = slice.getTo() ? mapValue(slice.getTo()) : mlir::Value {}; + auto& [listChunk, fromChunk, toChunk] = chunks; if (fromChunk) { - operands.push_back(fromChunk); + fromChunk = numberOperandChunk(fromChunk); } if (toChunk) { - operands.push_back(toChunk); + toChunk = numberOperandChunk(toChunk); } - // A list or a bound holding one value for every row rather than one per row is laid - // out over the rows the others carry, as lowerRange lays its bounds out - const mlir::Value cardinality = cardinalityDriver(operands); + const nl::ChunkType resultType = listOfSameElementsChunkType(context, listChunk); - listChunk = rowAlignedChunk(listChunk, cardinality); + // The op goes where the chunks it reads stand, which are the laid-out ones + llvm::SmallVector operands; + collectPresentChunks(chunks, operands); + setInsertionForNaryOp(operands); - if (fromChunk) { - fromChunk = numberOperandChunk(rowAlignedChunk(fromChunk, cardinality)); - } + nl::ListSlice run = _builder.create(loc, resultType, listChunk, fromChunk, toChunk); - if (toChunk) { - toChunk = numberOperandChunk(rowAlignedChunk(toChunk, cardinality)); - } + _valueMap[slice.getResult()] = run.getResult(); +} - // The slice holds the elements the sliced value holds, and is read as nullable: a row - // whose list or whose bound is null gets no slice. Slicing a type-erased cell answers - // the list its tag says it holds, whose elements are cells in their turn - const mlir::Type listElement = mlir::cast(listChunk.getType()).getElementType(); - const auto nullable = mlir::dyn_cast(listElement); - const mlir::Type value = nullable ? nullable.getValueType() : listElement; +void DBLowering::lowerSortMulti(mlir::db::SortMulti sortMulti) { + const mlir::Location loc = _builder.getUnknownLoc(); + mlir::MLIRContext* const context = _builder.getContext(); - const mlir::Type sliced = - mlir::isa(value) - ? storage::ListType::get(context, storage::AnyType::get(context)) - : value; + const mlir::Value fields = sortMulti.getFields(); + const mlir::Value limit = sortMulti.getLimit(); + const mlir::Value skip = sortMulti.getSkip(); - const nl::ChunkType resultType = nl::ChunkType::get(context, - storage::NullableType::get(context, sliced)); + std::array chunks { + mapValue(sortMulti.getList()), + fields ? mapValue(fields) : mlir::Value {}, + limit ? mapValue(limit) : mlir::Value {}, + skip ? mapValue(skip) : mlir::Value {}, + }; + rowAlignOptionalChunks(chunks); - // The op goes where the chunks it reads stand, which are the laid-out ones - operands.clear(); - operands.push_back(listChunk); + auto& [listChunk, fieldsChunk, limitChunk, skipChunk] = chunks; - if (fromChunk) { - operands.push_back(fromChunk); + if (limitChunk) { + limitChunk = numberOperandChunk(limitChunk); } - if (toChunk) { - operands.push_back(toChunk); + if (skipChunk) { + skipChunk = numberOperandChunk(skipChunk); } + const nl::ChunkType resultType = listOfSameElementsChunkType(context, listChunk); + + llvm::SmallVector operands; + collectPresentChunks(chunks, operands); setInsertionForNaryOp(operands); - nl::ListSlice run = _builder.create(loc, resultType, listChunk, fromChunk, toChunk); + nl::SortMulti sorted = _builder.create(loc, resultType, listChunk, fieldsChunk, limitChunk, skipChunk); - _valueMap[slice.getResult()] = run.getResult(); + _valueMap[sortMulti.getResult()] = sorted.getResult(); } void DBLowering::lowerPercentile(mlir::Operation* operation) { diff --git a/query/ir/lowering/DBLowering.h b/query/ir/lowering/DBLowering.h index 49f49804e0..26bd5790c5 100644 --- a/query/ir/lowering/DBLowering.h +++ b/query/ir/lowering/DBLowering.h @@ -377,7 +377,9 @@ class DBLowering { mlir::Type accumulatedElement(mlir::Type accumulated, mlir::Type chunkType); mlir::nl::ChunkType accumulatorChunkType(mlir::Type accumulated); + void rowAlignOptionalChunks(llvm::MutableArrayRef chunks); void lowerListSlice(mlir::db::ListSlice slice); + void lowerSortMulti(mlir::db::SortMulti sortMulti); void lowerPercentile(mlir::Operation* operation); void lowerPatternComprehension(mlir::db::PatternComprehension comprehension); void lowerScanEdges(mlir::db::ScanEdges scanEdges); diff --git a/storage/list/ListBuffer.cpp b/storage/list/ListBuffer.cpp index b86c2f2c11..8a3344effe 100644 --- a/storage/list/ListBuffer.cpp +++ b/storage/list/ListBuffer.cpp @@ -103,15 +103,26 @@ ListView ListBuffer::concatenate(ListView a, ListView b) { template ListView ListBuffer::reverse(ListView list) { + return writeElements(std::views::reverse(list.elements())); +} + +template +ListView ListBuffer::gather(std::span elements) { + return writeElements(elements); +} + +template +template +ListView ListBuffer::writeElements(const Elements& elements) { size_t valueBytes = 0; - for (const CellView element : list) { + for (const CellView element : elements) { valueBytes += elementValueSize(element); } - ListWriteCursor cursor = reserveList(list.size(), valueBytes); + ListWriteCursor cursor = reserveList(std::ranges::size(elements), valueBytes); constexpr size_t tagSize = decltype(_elements)::tagSize(); - for (const CellView element : std::views::reverse(list.elements())) { + for (const CellView element : elements) { const size_t numBytes = tagSize + elementValueSize(element); cursor.writeRaw(element.getData(), numBytes); } diff --git a/storage/list/ListBuffer.h b/storage/list/ListBuffer.h index 081c9e0843..dcb05b8d12 100644 --- a/storage/list/ListBuffer.h +++ b/storage/list/ListBuffer.h @@ -81,6 +81,12 @@ class ListBuffer { */ ListView reverse(ListView list); + /** + * @brief Stores @param elements, which may come from several lists, in their order as + * one new list in contiguous storage and returns a @ref ListView over it. + */ + ListView gather(std::span elements); + /** * @brief Stores the elements of @param list - the elements of a nested list among them * - as one new list in this buffer and returns a @ref ListView over the copy. @@ -107,6 +113,9 @@ class ListBuffer { /// Returns the value @param element holds, a nested list copied into this buffer ListItemVariant ownedItem(CellView element); + + template + ListView writeElements(const Elements& elements); }; // Default size; alias to avoid ugly empty template diff --git a/test/query/ir/journey/JourneyParetoTest.cpp b/test/query/ir/journey/JourneyParetoTest.cpp index e910d48394..76c9d1968e 100644 --- a/test/query/ir/journey/JourneyParetoTest.cpp +++ b/test/query/ir/journey/JourneyParetoTest.cpp @@ -172,13 +172,12 @@ TEST_F(JourneyParetoTest, collectsTheLegsAsMaps) { } // sortMulti sorts descending on a key, and ascending on one written with a ^. -// Disabled: apoc.coll.sortMulti does not exist (TUR-298). -TEST_F(JourneyParetoTest, DISABLED_sortsTheLegsWithSortMulti) { +TEST_F(JourneyParetoTest, sortsTheLegsWithSortMulti) { expectRows("WITH [{service: 'a', nlc: 'A', departure: 482, arrival: 497}, " "{service: 'b', nlc: 'B', departure: 470, arrival: 500}, " "{service: 'c', nlc: 'A', departure: 490, arrival: 502}, " "{service: 'd', nlc: 'A', departure: 470, arrival: 500}] AS OTTlist " - "RETURN [leg IN apoc.coll.sortMulti(OTTlist, ['nlc', '^arrival', 'departure']) | leg.service]", + "RETURN [leg IN coll.sortMulti(OTTlist, ['nlc', '^arrival', 'departure']) | leg.service]", {{"[b, a, d, c]"}}); } diff --git a/test/query/ir/list/CMakeLists.txt b/test/query/ir/list/CMakeLists.txt index ad8a589574..3cbc10c78d 100644 --- a/test/query/ir/list/CMakeLists.txt +++ b/test/query/ir/list/CMakeLists.txt @@ -168,6 +168,7 @@ add_turing_test(test_query_ir_slice_by_type_erased_bound SliceByTypeErasedBoundT add_call_v3_test(test_query_ir_list_of_paths ListOfPathsTest.cpp) add_call_v3_test(test_query_ir_coll_index_of CollIndexOfTest.cpp) +add_call_v3_test(test_query_ir_coll_sort_multi CollSortMultiTest.cpp) add_turing_test(test_query_ir_index_or_slice_a_non_list_cell IndexOrSliceANonListCellTest.cpp) diff --git a/test/query/ir/list/CollSortMultiTest.cpp b/test/query/ir/list/CollSortMultiTest.cpp new file mode 100644 index 0000000000..000ded6bce --- /dev/null +++ b/test/query/ir/list/CollSortMultiTest.cpp @@ -0,0 +1,183 @@ +#include + +#include +#include + +#include "CallV3Test.h" +#include "StringRowSink.h" + +using namespace turing::test; + +namespace { + +using Rows = std::vector; + +} + +class CollSortMultiTest : public CallV3Test { +protected: + void expectRows(std::string_view query, const Rows& expected) { + StringRowSink sink; + runQuery(query, sink); + + EXPECT_EQ(sink.getRows(), expected) << query; + } +}; + +TEST_F(CollSortMultiTest, sortsAscendingOnAFieldNamedWithACaret) { + expectRows("RETURN coll.sortMulti([{name: 'b'}, {name: 'c'}, {name: 'a'}], ['^name'])", + {{"{name: a}, {name: b}, {name: c}"}}); +} + +TEST_F(CollSortMultiTest, sortsDescendingOnAFieldNamedWithoutACaret) { + expectRows("RETURN coll.sortMulti([{name: 'b'}, {name: 'c'}, {name: 'a'}], ['name'])", + {{"{name: c}, {name: b}, {name: a}"}}); +} + +TEST_F(CollSortMultiTest, breaksATieOnTheNextField) { + expectRows("RETURN coll.sortMulti([{name: 'graphs', age: 32}, {name: 'are', age: 21}, " + "{name: 'everywhere', age: 42}, {name: 'are', age: 50}], ['^name', 'age'])", + {{"{age: 50, name: are}, {age: 21, name: are}, {age: 42, name: everywhere}, {age: 32, name: graphs}"}}); +} + +TEST_F(CollSortMultiTest, keepsTheOrderOfMapsThatTie) { + expectRows("RETURN coll.sortMulti([{a: 1, b: 'x'}, {a: 0, b: 'y'}, {a: 1, b: 'z'}, {a: 0, b: 'w'}], ['^a'])", + {{"{a: 0, b: y}, {a: 0, b: w}, {a: 1, b: x}, {a: 1, b: z}"}}); +} + +TEST_F(CollSortMultiTest, keepsTheListWithoutAField) { + expectRows("RETURN coll.sortMulti([{a: 2}, {a: 1}]), coll.sortMulti([{a: 2}, {a: 1}], [])", + {{"{a: 2}, {a: 1}", "{a: 2}, {a: 1}"}}); +} + +TEST_F(CollSortMultiTest, keepsAsManyMapsAsTheLimit) { + expectRows("RETURN coll.sortMulti([{a: 3}, {a: 1}, {a: 2}], ['^a'], 2), " + "coll.sortMulti([{a: 3}, {a: 1}, {a: 2}], ['^a'], -1), " + "coll.sortMulti([{a: 3}, {a: 1}, {a: 2}], ['^a'], 0)", + {{"{a: 1}, {a: 2}", "{a: 1}, {a: 2}, {a: 3}", ""}}); +} + +TEST_F(CollSortMultiTest, skipsTheFirstMaps) { + expectRows("RETURN coll.sortMulti([{a: 3}, {a: 1}, {a: 2}, {a: 4}], ['^a'], 2, 1), " + "coll.sortMulti([{a: 3}, {a: 1}, {a: 2}], ['^a'], -1, 2), " + "coll.sortMulti([{a: 3}, {a: 1}, {a: 2}], ['^a'], 5, 4)", + {{"{a: 2}, {a: 3}", "{a: 3}", ""}}); +} + +TEST_F(CollSortMultiTest, sortsAMissingKeyBeforeEveryValue) { + expectRows("RETURN coll.sortMulti([{a: 2}, {b: 1}, {a: null}, {a: 1}], ['^a']), " + "coll.sortMulti([{a: 2}, {b: 1}, {a: 1}], ['a'])", + {{"{b: 1}, {a: null}, {a: 1}, {a: 2}", "{a: 2}, {a: 1}, {b: 1}"}}); +} + +TEST_F(CollSortMultiTest, ordersValuesOfDifferentTypes) { + expectRows("RETURN coll.sortMulti([{v: 2.5}, {v: 'a'}, {v: 1}, {v: true}], ['^v'])", + {{"{v: a}, {v: true}, {v: 1}, {v: 2.5}"}}); +} + +TEST_F(CollSortMultiTest, answersNullOverANullArgument) { + expectRows("RETURN coll.sortMulti(null, ['^a']), coll.sortMulti([{a: 1}], null), " + "coll.sortMulti([{a: 1}], ['^a'], null), coll.sortMulti([{a: 1}], ['^a'], 1, null)", + {{"null", "null", "null", "null"}}); +} + +TEST_F(CollSortMultiTest, sortsTheListOfEveryRow) { + expectRows("UNWIND [[{a: 2}, {a: 1}], [{a: 3}, {a: 4}], null] AS xs RETURN coll.sortMulti(xs, ['a'])", + {{"{a: 2}, {a: 1}"}, {"{a: 4}, {a: 3}"}, {"null"}}); +} + +TEST_F(CollSortMultiTest, sortsAConstantListOnceForEveryRow) { + StringRowSink sink; + runQuery("EXPLAIN (nl) MATCH (n:Person) RETURN coll.sortMulti([{a: 2}, {a: 1}], ['^a'])", sink); + + std::string_view nl; + for (const StringRowSink::Row& row : sink.getRows()) { + if (row.front() == "nl") { + nl = row.back(); + } + } + + const size_t sort = nl.find("nl.sort_multi"); + ASSERT_NE(sort, std::string_view::npos) << nl; + EXPECT_LT(sort, nl.find("nl.for")) << nl; + + expectRows("MATCH (n:Person {name: 'Remy'}) RETURN coll.sortMulti([{a: 2}, {a: 1}], ['^a'])", + {{"{a: 1}, {a: 2}"}}); +} + +TEST_F(CollSortMultiTest, sortsAListHeldInACell) { + expectRows("UNWIND [[{a: 2}, {a: 1}], 'x'] AS xs WITH xs WHERE xs <> 'x' RETURN coll.sortMulti(xs, ['^a'])", + {{"{a: 1}, {a: 2}"}}); + + runQueryExpectingError("UNWIND [[{a: 2}, {a: 1}], 1] AS xs RETURN coll.sortMulti(xs, ['a'])", + "reads a list, not 'Integer'"); +} + +TEST_F(CollSortMultiTest, sortsTheListOfEveryRowAcrossChunks) { + expectRows("UNWIND range(1, 70000) AS i " + "WITH i, coll.sortMulti([{a: i}, {a: 0}, {a: -i}], ['a']) AS sorted " + "RETURN count(*), sum(sorted[0].a), sum(sorted[2].a), " + "sum(size(coll.sortMulti([{a: i}, {a: 0}, {a: -i}], ['a'], -1, i % 3)))", + {{"70000", "2450035000", "-2450035000", "140000"}}); +} + +TEST_F(CollSortMultiTest, readsTheFieldsAndBoundsOfEveryRow) { + expectRows("UNWIND [['^a', 1], ['a', 0]] AS spec " + "RETURN coll.sortMulti([{a: 1}, {a: 3}, {a: 2}], [spec[0]], 2, spec[1])", + {{"{a: 2}, {a: 3}"}, {"{a: 3}, {a: 2}"}}); +} + +TEST_F(CollSortMultiTest, sortsCollectedMaps) { + expectRows("MATCH (p:Person) WITH collect({name: p.name}) AS people " + "RETURN coll.sortMulti(people, ['^name'], 3)", + {{"{name: Adam}, {name: Cyrus}, {name: Doruk}"}}); +} + +TEST_F(CollSortMultiTest, sortsAStoredList) { + runWrite("MATCH (n:Person {name: 'Remy'}) SET n.items = [{a: 2}, {a: 1}, {a: 3}]"); + + expectRows("MATCH (n:Person {name: 'Remy'}) RETURN coll.sortMulti(n.items, ['^a'])", + {{"{a: 1}, {a: 2}, {a: 3}"}}); +} + +TEST_F(CollSortMultiTest, unwindsTheSortedMaps) { + expectRows("UNWIND coll.sortMulti([{a: 2}, {a: 1}, {a: 3}], ['^a']) AS m RETURN m.a", + {{"1"}, {"2"}, {"3"}}); +} + +TEST_F(CollSortMultiTest, readsTheNameWhateverItsCase) { + expectRows("RETURN COLL.SORTMULTI([{a: 1}, {a: 2}], ['a'])", {{"{a: 2}, {a: 1}"}}); +} + +TEST_F(CollSortMultiTest, keepsAnyListWithoutAField) { + expectRows("RETURN coll.sortMulti([{a: 1}, null, 2], [], 2), " + "CASE WHEN false THEN coll.sortMulti([{a: 1}, null], ['a']) END", + {{"{a: 1}, null", "null"}}); +} + +TEST_F(CollSortMultiTest, sortsANullElementAsAMapHoldingNoKey) { + expectRows("RETURN coll.sortMulti([{a: 1}, null, {a: 0}], ['^a']), coll.sortMulti([{a: 1}, null, {a: 0}], ['a']), " + "CASE WHEN false THEN coll.sortMulti([{a: 1}, 2], ['a']) END", + {{"null, {a: 0}, {a: 1}", "{a: 1}, {a: 0}, null", "null"}}); +} + +TEST_F(CollSortMultiTest, raisesOnAnElementThatIsNoMap) { + runQueryExpectingError("RETURN coll.sortMulti([{a: 1}, 2], ['^a'])", "coll.sortMulti sorts a list of maps"); +} + +TEST_F(CollSortMultiTest, raisesOnAFieldThatIsNoString) { + runQueryExpectingError("RETURN coll.sortMulti([{a: 1}], [1])", "coll.sortMulti sorts on fields named by strings"); +} + +TEST_F(CollSortMultiTest, raisesOnANegativeSkip) { + runQueryExpectingError("RETURN coll.sortMulti([{a: 1}], ['^a'], -1, -1)", "coll.sortMulti cannot skip a negative count"); +} + +TEST_F(CollSortMultiTest, rejectsAValueThatIsNoList) { + runQueryExpectingError("RETURN coll.sortMulti(1)", "Invalid arguments for function 'coll.sortMulti'"); + runQueryExpectingError("RETURN coll.sortMulti([{a: 1}], 'a')", "Invalid arguments for function 'coll.sortMulti'"); +} + +int main(int argc, char** argv) { + return turing::test::turingTestMain(argc, argv); +} diff --git a/test/query/ir/passes/CMakeLists.txt b/test/query/ir/passes/CMakeLists.txt index 89f48d9f23..41e92671e3 100644 --- a/test/query/ir/passes/CMakeLists.txt +++ b/test/query/ir/passes/CMakeLists.txt @@ -172,6 +172,8 @@ add_ir_tests(test_query_ir_map_key_read_speculatability MapKeyReadSpeculatabilit add_ir_tests(test_query_ir_list_read_speculatability ListReadSpeculatabilityTest.cpp) +add_ir_tests(test_query_ir_sort_multi_speculatability SortMultiSpeculatabilityTest.cpp) + add_turing_test(test_query_ir_unwound_value_bound_by_an_equality UnwoundValueBoundByAnEqualityTest.cpp) add_turing_test(test_query_ir_fuse_unwind_equality_stand_in FuseUnwindEqualityStandInTest.cpp) diff --git a/test/query/ir/passes/SortMultiSpeculatabilityTest.cpp b/test/query/ir/passes/SortMultiSpeculatabilityTest.cpp new file mode 100644 index 0000000000..b895100fb0 --- /dev/null +++ b/test/query/ir/passes/SortMultiSpeculatabilityTest.cpp @@ -0,0 +1,181 @@ +#include + +#include +#include + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OwningOpRef.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Parser/Parser.h" + +#include "DBDialect.h" +#include "DBOps.h" +#include "StorageDialect.h" + +using mlir::Speculation::NotSpeculatable; +using mlir::Speculation::Speculatable; + +// A sort raises on an element that is no map, a field that is no string, a negative skip and +// a sort key holding an embedding, so it is speculatable only over operands that hold none +class SortMultiSpeculatabilityTest : public ::testing::Test { +protected: + SortMultiSpeculatabilityTest() { + _context.getOrLoadDialect(); + _context.getOrLoadDialect(); + _context.getOrLoadDialect(); + } + + mlir::Speculation::Speculatability speculatabilityOfTheSort(std::string_view ops) { + std::string program = R"mlir( +!integers = !db.column +!cells = !db.column +!unresolved = !db.column +!fields = !db.column> +!maps = !db.column> +!mixed = !db.column> + +func.func @main(%integer: !integers, %cell: !cells, %names: !fields, %map: !db.column) { + )mlir"; + + program += ops; + program += "\n return\n}\n"; + + mlir::OwningOpRef module = mlir::parseSourceString(program, mlir::ParserConfig(&_context)); + EXPECT_TRUE(module) << ops; + if (!module) { + return NotSpeculatable; + } + + mlir::Operation* sort = nullptr; + module->walk([&sort](mlir::db::SortMulti op) { + sort = op; + }); + + EXPECT_TRUE(sort) << ops; + if (!sort) { + return NotSpeculatable; + } + + return mlir::cast(sort).getSpeculatability(); + } + + mlir::MLIRContext _context; +}; + +TEST_F(SortMultiSpeculatabilityTest, sortsALiteralListOfMapsOnLiteralFields) { + EXPECT_EQ(speculatabilityOfTheSort(R"mlir( + %maps = db.constant([{a = 2 : i64}, {a = 1 : i64}, {b = "x" : !storage.string}]) + %fields = db.constant(["^a" : !storage.string, "b" : !storage.string]) + %r = db.sort_multi(%maps) fields %fields limit %integer : (!maps, !fields, !integers) -> !unresolved +)mlir"), + Speculatable); + + EXPECT_EQ(speculatabilityOfTheSort(R"mlir( + %maps = db.constant([{a = 2 : i64}, {a = 1 : i64}]) + %r = db.sort_multi(%maps) : (!maps) -> !unresolved +)mlir"), + Speculatable); +} + +TEST_F(SortMultiSpeculatabilityTest, raisesOnAnElementThatIsNoMap) { + EXPECT_EQ(speculatabilityOfTheSort(R"mlir( + %maps = db.constant([{a = 2 : i64}, 1 : i64]) + %fields = db.constant(["a" : !storage.string]) + %r = db.sort_multi(%maps) fields %fields : (!mixed, !fields) -> !unresolved +)mlir"), + NotSpeculatable); + + EXPECT_EQ(speculatabilityOfTheSort("%r = db.sort_multi(%cell) : (!cells) -> !unresolved"), NotSpeculatable); +} + +TEST_F(SortMultiSpeculatabilityTest, sortsANullElement) { + EXPECT_EQ(speculatabilityOfTheSort(R"mlir( + %maps = db.constant([{a = 2 : i64}, unit]) + %fields = db.constant(["a" : !storage.string]) + %r = db.sort_multi(%maps) fields %fields : (!maps, !fields) -> !unresolved +)mlir"), + Speculatable); + + EXPECT_EQ(speculatabilityOfTheSort(R"mlir( + %one = db.constant(1 : i64) + %m = db.make_map(%one) keys ["a"] : !integers + %null = db.constant("" : !storage.nullable) + %maps = db.make_list(%m, %null) : (!db.column, !db.column>) -> !maps + %fields = db.constant(["a" : !storage.string]) + %r = db.sort_multi(%maps) fields %fields : (!maps, !fields) -> !unresolved +)mlir"), + Speculatable); +} + +TEST_F(SortMultiSpeculatabilityTest, keepsAnyListWithoutAField) { + EXPECT_EQ(speculatabilityOfTheSort(R"mlir( + %maps = db.constant([{a = 2 : i64}, unit, 1 : i64]) + %r = db.sort_multi(%maps) : (!mixed) -> !unresolved +)mlir"), + Speculatable); + + EXPECT_EQ(speculatabilityOfTheSort(R"mlir( + %maps = db.constant([{a = 2 : i64}, unit, 1 : i64]) + %fields = db.constant([]) + %r = db.sort_multi(%maps) fields %fields : (!mixed, !db.column>) -> !unresolved +)mlir"), + Speculatable); +} + +TEST_F(SortMultiSpeculatabilityTest, sortsTheResultOfAnotherSort) { + EXPECT_EQ(speculatabilityOfTheSort(R"mlir( + %maps = db.constant([{a = 2 : i64}, {a = 1 : i64}]) + %fields = db.constant(["a" : !storage.string]) + %inner = db.sort_multi(%maps) : (!maps) -> !unresolved + %r = db.sort_multi(%inner) fields %fields : (!unresolved, !fields) -> !unresolved +)mlir"), + Speculatable); +} + +TEST_F(SortMultiSpeculatabilityTest, raisesOnAKeyHoldingAValueWithNoOrder) { + EXPECT_EQ(speculatabilityOfTheSort(R"mlir( + %maps = db.constant([{a = [1 : i64]}, {a = 1 : i64}]) + %fields = db.constant(["a" : !storage.string]) + %r = db.sort_multi(%maps) fields %fields : (!maps, !fields) -> !unresolved +)mlir"), + NotSpeculatable); +} + +TEST_F(SortMultiSpeculatabilityTest, raisesOnFieldsItCannotRead) { + EXPECT_EQ(speculatabilityOfTheSort(R"mlir( + %maps = db.constant([{a = 2 : i64}]) + %r = db.sort_multi(%maps) fields %names : (!maps, !fields) -> !unresolved +)mlir"), + NotSpeculatable); + + EXPECT_EQ(speculatabilityOfTheSort(R"mlir( + %maps = db.constant([{a = 2 : i64}]) + %fields = db.constant([1 : i64]) + %r = db.sort_multi(%maps) fields %fields : (!maps, !db.column>) -> !unresolved +)mlir"), + NotSpeculatable); +} + +TEST_F(SortMultiSpeculatabilityTest, raisesOnASkipThatMayBeNegative) { + EXPECT_EQ(speculatabilityOfTheSort(R"mlir( + %maps = db.constant([{a = 2 : i64}]) + %two = db.constant(2 : i64) + %r = db.sort_multi(%maps) skip %two : (!maps, !integers) -> !unresolved +)mlir"), + Speculatable); + + EXPECT_EQ(speculatabilityOfTheSort(R"mlir( + %maps = db.constant([{a = 2 : i64}]) + %minus = db.constant(-1 : i64) + %r = db.sort_multi(%maps) skip %minus : (!maps, !integers) -> !unresolved +)mlir"), + NotSpeculatable); + + EXPECT_EQ(speculatabilityOfTheSort(R"mlir( + %maps = db.constant([{a = 2 : i64}]) + %r = db.sort_multi(%maps) skip %integer : (!maps, !integers) -> !unresolved +)mlir"), + NotSpeculatable); +}