@milaboratories/pl-model-common 1.24.9 → 1.24.10

package/src/drivers/pframe/query/utils.test.ts

@@ -0,0 +1,293 @@
+import { describe, expect, it } from "vitest";
+import type { PObjectId } from "../../../pool";
+import type { SpecQuery, SpecQueryJoinEntry } from "./query_spec";
+import {
+  traverseQuerySpec,
+  mapSpecQueryColumns,
+  collectSpecQueryColumns,
+  sortSpecQuery,
+  isBooleanExpression,
+} from "./utils";
+
+type Q = SpecQuery<string>;
+type JE = SpecQueryJoinEntry<string>;
+
+const col = (c: string): Q => ({ type: "column", column: c });
+const entry = (q: Q): JE => ({ entry: q, qualifications: [] });
+
+/** Typed helpers for sortSpecQuery tests (PObjectId is branded). */
+const pid = (c: string) => c as PObjectId;
+const pcol = (c: string): SpecQuery => ({ type: "column", column: pid(c) });
+const pentry = (q: SpecQuery): SpecQueryJoinEntry => ({ entry: q, qualifications: [] });
+
+describe("traverseQuerySpec", () => {
+  it("transforms column leaf", () => {
+    const result = traverseQuerySpec(col("a"), { column: (c) => c.toUpperCase() });
+    expect(result).toEqual({ type: "column", column: "A" });
+  });
+
+  it("transforms sparseToDenseColumn leaf", () => {
+    const q: Q = { type: "sparseToDenseColumn", column: "a", axesIndices: [0, 1] } as Q;
+    const result = traverseQuerySpec(q, { column: (c) => `${c}!` });
+    expect(result).toEqual({ type: "sparseToDenseColumn", column: "a!", axesIndices: [0, 1] });
+  });
+
+  it("passes inlineColumn through unchanged", () => {
+    const q: Q = {
+      type: "inlineColumn",
+      spec: { id: "x", spec: {} },
+      dataInfo: {},
+    } as unknown as Q;
+    const result = traverseQuerySpec(q, { column: (c) => c.toUpperCase() });
+    expect(result).toBe(q);
+  });
+
+  it("transforms columns inside innerJoin", () => {
+    const q: Q = { type: "innerJoin", entries: [entry(col("a")), entry(col("b"))] };
+    const result = traverseQuerySpec(q, { column: (c) => c.toUpperCase() });
+    expect(result).toEqual({
+      type: "innerJoin",
+      entries: [entry({ type: "column", column: "A" }), entry({ type: "column", column: "B" })],
+    });
+  });
+
+  it("transforms columns inside fullJoin", () => {
+    const q: Q = { type: "fullJoin", entries: [entry(col("a"))] };
+    const result = traverseQuerySpec(q, { column: (c) => c.toUpperCase() });
+    expect(result).toEqual({
+      type: "fullJoin",
+      entries: [entry({ type: "column", column: "A" })],
+    });
+  });
+
+  it("transforms columns inside outerJoin", () => {
+    const q: Q = {
+      type: "outerJoin",
+      primary: entry(col("a")),
+      secondary: [entry(col("b"))],
+    };
+    const result = traverseQuerySpec(q, { column: (c) => c.toUpperCase() });
+    expect(result).toEqual({
+      type: "outerJoin",
+      primary: entry({ type: "column", column: "A" }),
+      secondary: [entry({ type: "column", column: "B" })],
+    });
+  });
+
+  it("transforms columns inside filter", () => {
+    const q: Q = {
+      type: "filter",
+      input: col("a"),
+      predicate: { type: "isNull", input: { type: "columnRef", value: pid("id1") } },
+    } as Q;
+    const result = traverseQuerySpec(q, { column: (c) => c.toUpperCase() });
+    expect(result).toEqual({
+      type: "filter",
+      input: { type: "column", column: "A" },
+      predicate: { type: "isNull", input: { type: "columnRef", value: "id1" } },
+    });
+  });
+
+  it("transforms columns inside sort", () => {
+    const q: Q = {
+      type: "sort",
+      input: col("a"),
+      sortBy: [
+        { expression: { type: "columnRef", value: pid("id1") }, ascending: true, nullsFirst: null },
+      ],
+    } as Q;
+    const result = traverseQuerySpec(q, { column: (c) => c.toUpperCase() });
+    expect(result).toEqual({
+      type: "sort",
+      input: { type: "column", column: "A" },
+      sortBy: [
+        { expression: { type: "columnRef", value: "id1" }, ascending: true, nullsFirst: null },
+      ],
+    });
+  });
+
+  it("transforms nested structure", () => {
+    const q: Q = {
+      type: "filter",
+      input: {
+        type: "outerJoin",
+        primary: entry(col("a")),
+        secondary: [entry(col("b")), entry(col("c"))],
+      },
+      predicate: { type: "isNull", input: { type: "columnRef", value: pid("id1") } },
+    } as Q;
+    const result = traverseQuerySpec(q, { column: (c) => c.toUpperCase() });
+    expect(collectSpecQueryColumns(result)).toEqual(["A", "B", "C"]);
+  });
+
+  it("calls node visitor after children are traversed", () => {
+    const q: Q = { type: "innerJoin", entries: [entry(col("a")), entry(col("b"))] };
+    const visited: string[] = [];
+    traverseQuerySpec(q, {
+      column: (c) => c,
+      node: (node) => {
+        visited.push(node.type);
+        return node;
+      },
+    });
+    expect(visited).toEqual(["column", "column", "innerJoin"]);
+  });
+
+  it("calls joinEntry visitor for each entry", () => {
+    const q: Q = { type: "innerJoin", entries: [entry(col("a")), entry(col("b"))] };
+    const visited: string[] = [];
+    traverseQuerySpec(q, {
+      column: (c) => c,
+      joinEntry: (e) => {
+        if (e.entry.type === "column") visited.push(e.entry.column);
+        return e;
+      },
+    });
+    expect(visited).toEqual(["a", "b"]);
+  });
+});
+
+describe("mapSpecQueryColumns", () => {
+  it("maps all column references", () => {
+    const q: Q = {
+      type: "outerJoin",
+      primary: entry(col("a")),
+      secondary: [entry(col("b"))],
+    };
+    const result = mapSpecQueryColumns(q, (c) => c.toUpperCase());
+    expect(collectSpecQueryColumns(result)).toEqual(["A", "B"]);
+  });
+});
+
+describe("collectSpecQueryColumns", () => {
+  it("collects from a single column", () => {
+    expect(collectSpecQueryColumns(col("a"))).toEqual(["a"]);
+  });
+
+  it("collects from join", () => {
+    const q: Q = { type: "innerJoin", entries: [entry(col("a")), entry(col("b"))] };
+    expect(collectSpecQueryColumns(q)).toEqual(["a", "b"]);
+  });
+
+  it("collects from nested structure", () => {
+    const q: Q = {
+      type: "filter",
+      input: {
+        type: "outerJoin",
+        primary: entry(col("x")),
+        secondary: [entry(col("y")), entry(col("z"))],
+      },
+      predicate: { type: "isNull", input: { type: "columnRef", value: pid("id1") } },
+    } as Q;
+    expect(collectSpecQueryColumns(q)).toEqual(["x", "y", "z"]);
+  });
+
+  it("collects from sparseToDenseColumn", () => {
+    const q: Q = { type: "sparseToDenseColumn", column: "s", axesIndices: [0] } as Q;
+    expect(collectSpecQueryColumns(q)).toEqual(["s"]);
+  });
+
+  it("returns empty for inlineColumn", () => {
+    const q: Q = {
+      type: "inlineColumn",
+      spec: { id: "x", spec: {} },
+      dataInfo: {},
+    } as unknown as Q;
+    expect(collectSpecQueryColumns(q)).toEqual([]);
+  });
+});
+
+describe("sortSpecQuery", () => {
+  it("sorts innerJoin entries by column id", () => {
+    const q: SpecQuery = { type: "innerJoin", entries: [pentry(pcol("b")), pentry(pcol("a"))] };
+    const result = sortSpecQuery(q);
+    expect(result).toEqual({
+      type: "innerJoin",
+      entries: [pentry(pcol("a")), pentry(pcol("b"))],
+    });
+  });
+
+  it("sorts outerJoin secondary entries", () => {
+    const q: SpecQuery = {
+      type: "outerJoin",
+      primary: pentry(pcol("p")),
+      secondary: [pentry(pcol("c")), pentry(pcol("a")), pentry(pcol("b"))],
+    };
+    const result = sortSpecQuery(q);
+    expect(result).toEqual({
+      type: "outerJoin",
+      primary: pentry(pcol("p")),
+      secondary: [pentry(pcol("a")), pentry(pcol("b")), pentry(pcol("c"))],
+    });
+  });
+
+  it("sorts sparseToDenseColumn axesIndices", () => {
+    const q: SpecQuery = {
+      type: "sparseToDenseColumn",
+      column: pid("a"),
+      axesIndices: [2, 0, 1],
+    } as SpecQuery;
+    const result = sortSpecQuery(q);
+    expect(result).toEqual({
+      type: "sparseToDenseColumn",
+      column: "a",
+      axesIndices: [0, 1, 2],
+    });
+  });
+
+  it("sorts join entry qualifications", () => {
+    const q: SpecQuery = {
+      type: "innerJoin",
+      entries: [
+        {
+          entry: pcol("a"),
+          qualifications: [
+            { axis: { name: "z" }, additionalDomains: {} },
+            { axis: { name: "a" }, additionalDomains: {} },
+          ],
+        },
+      ],
+    };
+    const result = sortSpecQuery(q);
+    const entries = (result as { type: "innerJoin"; entries: SpecQueryJoinEntry[] }).entries;
+    expect(entries[0].qualifications[0].axis).toEqual({ name: "a" });
+    expect(entries[0].qualifications[1].axis).toEqual({ name: "z" });
+  });
+
+  it("sorts nested structures recursively", () => {
+    const q: SpecQuery = {
+      type: "filter",
+      input: { type: "innerJoin", entries: [pentry(pcol("b")), pentry(pcol("a"))] },
+      predicate: { type: "isNull", input: { type: "columnRef", value: pid("id1") } },
+    };
+    const result = sortSpecQuery(q);
+    const inner = (result as { type: "filter"; input: SpecQuery }).input;
+    expect(collectSpecQueryColumns(inner)).toEqual(["a", "b"]);
+  });
+
+  it("leaves column leaf unchanged", () => {
+    const q: SpecQuery = pcol("a");
+    expect(sortSpecQuery(q)).toEqual(pcol("a"));
+  });
+});
+
+describe("isBooleanExpression", () => {
+  it("returns true for boolean expression types", () => {
+    expect(
+      isBooleanExpression({ type: "isNull", input: { type: "columnRef", value: pid("a") } }),
+    ).toBe(true);
+    expect(
+      isBooleanExpression({
+        type: "not",
+        input: { type: "isNull", input: { type: "columnRef", value: pid("a") } },
+      }),
+    ).toBe(true);
+    expect(isBooleanExpression({ type: "and", input: [] })).toBe(true);
+  });
+
+  it("returns false for non-boolean expression types", () => {
+    expect(isBooleanExpression({ type: "columnRef", value: pid("a") })).toBe(false);
+    expect(isBooleanExpression({ type: "axisRef", value: { name: "x" } })).toBe(false);
+    expect(isBooleanExpression({ type: "constant", value: 42 })).toBe(false);
+  });
+});

package/src/drivers/pframe/query/utils.ts

@@ -27,108 +27,127 @@ export function isBooleanExpression(expr: SpecQueryExpression): expr is SpecQuer
   );
 }
 
-/** Collects all column references from a SpecQuery tree. */
-export function collectQueryColumns<C>(query: SpecQuery<C>): C[] {
-  const result: C[] = [];
-  collectQueryColumnsImpl(query, result);
-  return result;
-}
+/**
+ * Recursively traverses a SpecQuery tree bottom-up, applying visitor callbacks.
+ *
+ * Traversal order:
+ * 1. Recurse into child queries
+ * 2. Apply `column` to transform column references in leaf nodes
+ * 3. Apply `joinEntry` to each join entry (with inner query already traversed)
+ * 4. Assemble node with transformed children
+ * 5. Apply `node` to the assembled node
+ */
+export function traverseQuerySpec<C1, C2>(
+  query: SpecQuery<C1>,
+  visitor: {
+    /** Transform column references in leaf nodes (column, sparseToDenseColumn). */
+    column: (c: C1) => C2;
+    /** Visit a node after its children have been traversed. */
+    node?: (node: SpecQuery<C2>) => SpecQuery<C2>;
+    /** Visit a join entry after its inner query has been traversed. */
+    joinEntry?: (entry: SpecQueryJoinEntry<C2>) => SpecQueryJoinEntry<C2>;
+  },
+): SpecQuery<C2> {
+  const traverseEntry = (entry: SpecQueryJoinEntry<C1>): SpecQueryJoinEntry<C2> => {
+    const traversed: SpecQueryJoinEntry<C2> = {
+      ...entry,
+      entry: traverseQuerySpec(entry.entry, visitor),
+    };
+    return visitor.joinEntry ? visitor.joinEntry(traversed) : traversed;
+  };
 
-function collectQueryColumnsImpl<C>(query: SpecQuery<C>, result: C[]): void {
+  let result: SpecQuery<C2>;
   switch (query.type) {
     case "column":
-      result.push(query.column);
+      result = { type: "column", column: visitor.column(query.column) };
       break;
     case "sparseToDenseColumn":
-      result.push(query.column);
+      result = { ...query, column: visitor.column(query.column) };
       break;
     case "inlineColumn":
+      result = query;
       break;
     case "innerJoin":
     case "fullJoin":
-      for (const e of query.entries) collectQueryColumnsImpl(e.entry, result);
+      result = { ...query, entries: query.entries.map(traverseEntry) };
       break;
     case "outerJoin":
-      collectQueryColumnsImpl(query.primary.entry, result);
-      for (const e of query.secondary) collectQueryColumnsImpl(e.entry, result);
+      result = {
+        ...query,
+        primary: traverseEntry(query.primary),
+        secondary: query.secondary.map(traverseEntry),
+      };
       break;
     case "filter":
     case "sort":
     case "sliceAxes":
-      collectQueryColumnsImpl(query.input, result);
+      result = { ...query, input: traverseQuerySpec(query.input, visitor) };
       break;
     default:
       assertNever(query);
   }
+
+  return visitor.node ? visitor.node(result) : result;
 }
 
-export function sortQuerySpec(query: SpecQuery): SpecQuery {
-  switch (query.type) {
-    case "column":
-    case "inlineColumn":
-      return query;
-    case "sparseToDenseColumn": {
-      const sortedAxesIndices = query.axesIndices.toSorted((lhs, rhs) => lhs - rhs);
-      return {
-        ...query,
-        axesIndices: sortedAxesIndices,
-      };
-    }
-    case "innerJoin":
-    case "fullJoin": {
-      const sortedEntries = query.entries.map(sortQueryJoinEntrySpec);
-      sortedEntries.sort(cmpQueryJoinEntrySpec);
-      return {
-        ...query,
-        entries: sortedEntries,
-      };
-    }
-    case "outerJoin": {
-      const sortedSecondary = query.secondary.map(sortQueryJoinEntrySpec);
-      sortedSecondary.sort(cmpQueryJoinEntrySpec);
-      return {
-        ...query,
-        primary: sortQueryJoinEntrySpec(query.primary),
-        secondary: sortedSecondary,
-      };
-    }
-    case "sliceAxes": {
-      const sortedAxisFilters = query.axisFilters.toSorted((lhs, rhs) => {
-        const lhsKey = canonicalizeJson(lhs.axisSelector);
-        const rhsKey = canonicalizeJson(rhs.axisSelector);
-        return lhsKey < rhsKey ? -1 : lhsKey === rhsKey ? 0 : 1;
-      });
-      return {
-        ...query,
-        input: sortQuerySpec(query.input),
-        axisFilters: sortedAxisFilters,
-      };
-    }
-    case "sort":
-      return {
-        ...query,
-        input: sortQuerySpec(query.input),
-      };
-    case "filter":
-      return {
-        ...query,
-        input: sortQuerySpec(query.input),
-      };
-    default:
-      assertNever(query);
-  }
+/** Recursively maps all column references in a SpecQuery tree. */
+export function mapSpecQueryColumns<C1, C2>(
+  query: SpecQuery<C1>,
+  cb: (c: C1) => C2,
+): SpecQuery<C2> {
+  return traverseQuerySpec(query, { column: cb });
 }
 
-function sortQueryJoinEntrySpec(entry: SpecQueryJoinEntry): SpecQueryJoinEntry {
-  const sortedQualifications = entry.qualifications.toSorted((lhs, rhs) => {
-    const lhsKey = canonicalizeJson(lhs.axis);
-    const rhsKey = canonicalizeJson(rhs.axis);
-    return lhsKey < rhsKey ? -1 : lhsKey === rhsKey ? 0 : 1;
+/** Collects all column references from a SpecQuery tree. */
+export function collectSpecQueryColumns<C>(query: SpecQuery<C>): C[] {
+  const result: C[] = [];
+  traverseQuerySpec(query, {
+    column: (c: C) => {
+      result.push(c);
+      return c;
+    },
+  });
+  return result;
+}
+
+export function sortSpecQuery(query: SpecQuery): SpecQuery {
+  return traverseQuerySpec(query, {
+    column: (c) => c,
+    node: (node) => {
+      switch (node.type) {
+        case "sparseToDenseColumn":
+          return { ...node, axesIndices: node.axesIndices.toSorted((a, b) => a - b) };
+        case "innerJoin":
+        case "fullJoin": {
+          const sorted = [...node.entries].sort(cmpQueryJoinEntrySpec);
+          return { ...node, entries: sorted };
+        }
+        case "outerJoin": {
+          const sorted = [...node.secondary].sort(cmpQueryJoinEntrySpec);
+          return { ...node, secondary: sorted };
+        }
+        case "sliceAxes":
+          return {
+            ...node,
+            axisFilters: node.axisFilters.toSorted((a, b) => {
+              const ak = canonicalizeJson(a.axisSelector);
+              const bk = canonicalizeJson(b.axisSelector);
+              return ak < bk ? -1 : ak === bk ? 0 : 1;
+            }),
+          };
+        default:
+          return node;
+      }
+    },
+    joinEntry: (entry) => ({
+      ...entry,
+      qualifications: entry.qualifications.toSorted((a, b) => {
+        const ak = canonicalizeJson(a.axis);
+        const bk = canonicalizeJson(b.axis);
+        return ak < bk ? -1 : ak === bk ? 0 : 1;
+      }),
+    }),
   });
-  return {
-    entry: sortQuerySpec(entry.entry),
-    qualifications: sortedQualifications,
-  };
 }
 
 function cmpQuerySpec(lhs: SpecQuery, rhs: SpecQuery): number {

package/src/drivers/pframe/table_calculate.ts

@@ -4,8 +4,9 @@ import type { PObjectId } from "../../pool";
 import { assertNever } from "../../util";
 import { getAxisId, type PColumn } from "./spec/spec";
 import type { PColumnValues } from "./data_info";
-import type { SpecQuery, SpecQueryJoinEntry } from "./query/query_spec";
+import type { SpecQuery } from "./query/query_spec";
 import { canonicalizeJson } from "../../json";
+import { mapSpecQueryColumns } from "./query";
 
 /** Defines a terminal column node in the join request tree */
 export interface ColumnJoinEntry<Col> {
@@ -423,49 +424,7 @@ export function sortPTableDef(def: PTableDef<PObjectId>): PTableDef<PObjectId> {
 }
 
 export function mapPTableDefV2<C1, C2>(def: PTableDefV2<C1>, cb: (c: C1) => C2): PTableDefV2<C2> {
-  return { query: mapQuerySpec(def.query, cb) };
-}
-
-/** Recursively maps all column references in a SpecQuery tree. */
-export function mapQuerySpec<C1, C2>(query: SpecQuery<C1>, cb: (c: C1) => C2): SpecQuery<C2> {
-  switch (query.type) {
-    case "column":
-      return { type: "column", column: cb(query.column) };
-    case "sparseToDenseColumn":
-      return { ...query, column: cb(query.column) };
-    case "inlineColumn":
-      return query;
-    case "innerJoin":
-    case "fullJoin":
-      return {
-        ...query,
-        entries: query.entries.map((e) => mapQueryJoinEntrySpec(e, cb)),
-      };
-    case "outerJoin":
-      return {
-        ...query,
-        primary: mapQueryJoinEntrySpec(query.primary, cb),
-        secondary: query.secondary.map((e) => mapQueryJoinEntrySpec(e, cb)),
-      };
-    case "filter":
-      return { ...query, input: mapQuerySpec(query.input, cb) };
-    case "sort":
-      return { ...query, input: mapQuerySpec(query.input, cb) };
-    case "sliceAxes":
-      return { ...query, input: mapQuerySpec(query.input, cb) };
-    default:
-      assertNever(query);
-  }
-}
-
-function mapQueryJoinEntrySpec<C1, C2>(
-  entry: SpecQueryJoinEntry<C1>,
-  cb: (c: C1) => C2,
-): SpecQueryJoinEntry<C2> {
-  return {
-    ...entry,
-    entry: mapQuerySpec(entry.entry, cb),
-  };
+  return { query: mapSpecQueryColumns(def.query, cb) };
 }
 
 export function mapJoinEntry<C1, C2>(entry: JoinEntry<C1>, cb: (c: C1) => C2): JoinEntry<C2> {