refactor(zql): Update Join and FlippedJoin to always fetch parent with partition constraint.

packages/zql/src/builder/builder.ts

@@ -312,12 +312,13 @@ function buildPipelineInternal(
         end,
         name,
         true,
+        partitionKey,
       );
     }
   }
 
   if (ast.where && (!fullyAppliedFilters || delegate.applyFiltersAnyway)) {
-    end = applyWhere(end, ast.where, delegate, name);
+    end = applyWhere(end, ast.where, delegate, name, partitionKey);
   }
 
   if (ast.limit !== undefined) {
@@ -334,7 +335,15 @@ function buildPipelineInternal(
 
   if (ast.related) {
     for (const csq of ast.related) {
-      end = applyCorrelatedSubQuery(csq, delegate, queryID, end, name, false);
+      end = applyCorrelatedSubQuery(
+        csq,
+        delegate,
+        queryID,
+        end,
+        name,
+        false,
+        partitionKey,
+      );
     }
   }
 
@@ -346,21 +355,23 @@ function applyWhere(
   condition: Condition,
   delegate: BuilderDelegate,
   name: string,
+  partitionKey: CompoundKey | undefined,
 ): Input {
   if (!conditionIncludesFlippedSubqueryAtAnyLevel(condition)) {
     return buildFilterPipeline(input, delegate, filterInput =>
       applyFilter(filterInput, condition, delegate, name),
     );
   }
 
-  return applyFilterWithFlips(input, condition, delegate, name);
+  return applyFilterWithFlips(input, condition, delegate, name, partitionKey);
 }
 
 function applyFilterWithFlips(
   input: Input,
   condition: Condition,
   delegate: BuilderDelegate,
   name: string,
+  partitionKey: CompoundKey | undefined,
 ): Input {
   let end = input;
   assert(condition.type !== 'simple', 'Simple conditions cannot have flips');
@@ -386,7 +397,7 @@ function applyFilterWithFlips(
       }
       assert(withFlipped.length > 0, 'Impossible to have no flips here');
       for (const cond of withFlipped) {
-        end = applyFilterWithFlips(end, cond, delegate, name);
+        end = applyFilterWithFlips(end, cond, delegate, name, partitionKey);
       }
       break;
     }
@@ -419,7 +430,9 @@ function applyFilterWithFlips(
       }
 
       for (const cond of withFlipped) {
-        branches.push(applyFilterWithFlips(end, cond, delegate, name));
+        branches.push(
+          applyFilterWithFlips(end, cond, delegate, name, partitionKey),
+        );
       }
 
       const ufi = new UnionFanIn(ufo, branches);
@@ -439,11 +452,15 @@ function applyFilterWithFlips(
         `${name}.${sq.subquery.alias}`,
         sq.correlation.childField,
       );
+
+      const flippedJoinName = `${name}:flipped-join(${sq.subquery.alias})`;
       const flippedJoin = new FlippedJoin({
         parent: end,
         child,
+        storage: delegate.createStorage(flippedJoinName),
         parentKey: sq.correlation.parentField,
         childKey: sq.correlation.childField,
+        partitionKey,
         relationshipName: must(
           sq.subquery.alias,
           'Subquery must have an alias',
@@ -453,10 +470,7 @@ function applyFilterWithFlips(
       });
       delegate.addEdge(end, flippedJoin);
       delegate.addEdge(child, flippedJoin);
-      end = delegate.decorateInput(
-        flippedJoin,
-        `${name}:flipped-join(${sq.subquery.alias})`,
-      );
+      end = delegate.decorateInput(flippedJoin, flippedJoinName);
       break;
     }
   }
@@ -598,6 +612,7 @@ function applyCorrelatedSubQuery(
   end: Input,
   name: string,
   fromCondition: boolean,
+  partitionKey: CompoundKey | undefined,
 ) {
   // TODO: we only omit the join if the CSQ if from a condition since
   // we want to create an empty array for `related` fields that are `limit(0)`
@@ -621,6 +636,7 @@ function applyCorrelatedSubQuery(
     storage: delegate.createStorage(joinName),
     parentKey: sq.correlation.parentField,
     childKey: sq.correlation.childField,
+    partitionKey,
     relationshipName: sq.subquery.alias,
     hidden: sq.hidden ?? false,
     system: sq.system ?? 'client',

packages/zql/src/ivm/flipped-join.more-fetch.test.ts

@@ -169,7 +169,14 @@ suite('one:many:one', () => {
       ]
     `);
 
-    expect(actualStorage).toMatchInlineSnapshot(`{}`);
+    expect(actualStorage).toMatchInlineSnapshot(`
+      {
+        ".issueLabels_1_0:flipped-join(labels_1)": {},
+        ".issueLabels_2_1:flipped-join(labels_2)": {},
+        ":flipped-join(issueLabels_1_0)": {},
+        ":flipped-join(issueLabels_2_1)": {},
+      }
+    `);
 
     expect(log).toMatchInlineSnapshot(`
       [