TUR-279: MLIR: unwind a list built per row without building a whole chunk of lists

query/ir/codegen/DBPassPipeline.cpp

@@ -10,14 +10,15 @@ namespace {
 
 using DBPassFactory = std::unique_ptr<mlir::Pass> (*)(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
 // pipeline it reports on and the pipeline that runs the same one. Every factory is handed
 // 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<DBPassFactory, DB_PASS_COUNT> 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<DBPassFactory, DB_PASS_COUNT> 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(); },
 };
 
 }

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<UnwindYield>(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<unsigned>(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<ColumnType>(getElement().getType());
-    const storage::NullableType nullable = elementColumn ? dyn_cast<storage::NullableType>(elementColumn.getType())
-                                                         : storage::NullableType {};
+    const Type element = elementColumn ? elementColumn.getType() : Type {};
+    const storage::NullableType nullable = dyn_cast_or_null<storage::NullableType>(element);
 
-    if (!nullable || !isa<IntegerType>(nullable.getValueType())) {
-        return emitOpError("element must be a column of nullable integers");
+    const bool nullableInteger = nullable && isa<IntegerType>(nullable.getValueType());
+    const bool resolvedInLowering = isa_and_nonnull<NoneType>(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());

query/ir/dialect/db/DBOps.td

@@ -385,7 +385,7 @@ def Unwind : TuringOp<"unwind", [Pure]> {
 /**
 * UnwindRange
 */
-def UnwindRange : TuringOp<"unwind_range", [NoMemoryEffect, DeclareOpInterfaceMethods<ConditionallySpeculatable>]> {
+def UnwindRange : TuringOp<"unwind_range", [NoMemoryEffect, AttrSizedOperandSegments, DeclareOpInterfaceMethods<ConditionallySpeculatable>]> {
   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<Column>:$from,
+                       Optional<Column>:$to,
+                       Variadic<Column>:$columns_to_filter,
+                       UnitAttr:$reversed);
+  let results   = (outs Column:$element, Variadic<Column>:$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<!storage.node_id>):
+          %name = db.get_node_properties(%nodeIn) {property = "name"} : ...
+          %parts = db.split(%name, %comma) : ...
+          db.unwind_yield %parts : !db.column<!storage.list<!storage.string>>
+      }) : (!db.column<!storage.node_id>) -> (!db.column<none>, !db.column<!storage.node_id>)
+
+    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<Column>:$columns_to_filter);
+  let arguments = (ins Variadic<Column>:$columns_to_filter);
   let results   = (outs Column:$element, Variadic<Column>:$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))
   }];
 }