diff --git a/.changeset/filter-anchor-partitions.md b/.changeset/filter-anchor-partitions.md new file mode 100644 index 0000000000..7e596b9691 --- /dev/null +++ b/.changeset/filter-anchor-partitions.md @@ -0,0 +1,5 @@ +--- +"@platforma-sdk/model": minor +--- + +Add `filterToAnchorPartitions` option to `getAnchoredPColumns`. When set, restricts each returned column's visible axis-0 key range to the partition keys present in the resolved anchor column(s) by injecting a `pl7.app/axisKeys/0` annotation onto the returned column specs. diff --git a/sdk/model/src/render/api.test.ts b/sdk/model/src/render/api.test.ts new file mode 100644 index 0000000000..06731965a4 --- /dev/null +++ b/sdk/model/src/render/api.test.ts @@ -0,0 +1,170 @@ +import { AnchoredIdDeriver, createPlRef } from "@milaboratories/pl-model-common"; +import type { PColumn, PColumnSpec, PlRef } from "@milaboratories/pl-model-common"; +import { beforeEach, expect, test, vi } from "vitest"; +import { ResultPool } from "./api"; + +// Prevent getCfgRenderCtx from throwing during ResultPool construction +vi.mock("../internal", () => ({ + getCfgRenderCtx: vi.fn().mockReturnValue({}), +})); + +// vi.hoisted ensures mockGetColumns is available inside the vi.mock factory +const { mockGetColumns } = vi.hoisted(() => ({ mockGetColumns: vi.fn() })); + +// Mock PColumnCollection so getAnchoredPColumns doesn't need a real render context. +// Must use a regular class (not an arrow function) as the constructor mock. +vi.mock("./util/column_collection", () => ({ + PColumnCollection: class MockPColumnCollection { + addColumnProvider() { + return this; + } + addAxisLabelProvider() { + return this; + } + getColumns(...args: unknown[]) { + return mockGetColumns(...args); + } + }, +})); + +beforeEach(() => { + mockGetColumns.mockReset(); +}); + +function makePartitionedData(keys: string[]) { + return { + type: "JsonPartitioned" as const, + parts: keys.map((k) => ({ key: [k], value: undefined })), + partitionKeyLength: 1, + }; +} + +function makeColumn(data: unknown): PColumn { + return { + id: "col-id", + spec: { + kind: "PColumn" as const, + name: "test/name", + valueType: "String", + axesSpec: [], + annotations: {}, + }, + data, + } as PColumn; +} + +class TestResultPool extends ResultPool { + private readonly mockCols = new Map>(); + + setCol(ref: PlRef, col: PColumn) { + this.mockCols.set(`${ref.blockId}/${ref.name}`, col); + } + + override getPColumnByRef(ref: PlRef): any { + return this.mockCols.get(`${ref.blockId}/${ref.name}`); + } + + /** Expose private deriveAnchorPartitionKeys for testing */ + deriveKeys(anchors: AnchoredIdDeriver | Record) { + return (this as any).deriveAnchorPartitionKeys(anchors); + } +} + +// ─── deriveAnchorPartitionKeys ─────────────────────────────────────────────── + +test("deriveAnchorPartitionKeys: returns [] for AnchoredIdDeriver input", () => { + const pool = new TestResultPool(); + expect(pool.deriveKeys(new AnchoredIdDeriver({}))).toEqual([]); +}); + +test("deriveAnchorPartitionKeys: returns undefined when anchor column is not found", () => { + const pool = new TestResultPool(); + const ref = createPlRef("block-1", "missing"); + // No column registered — getPColumnByRef returns undefined + expect(pool.deriveKeys({ main: ref })).toBeUndefined(); +}); + +test("deriveAnchorPartitionKeys: returns undefined when anchor column has no data", () => { + const pool = new TestResultPool(); + const ref = createPlRef("block-1", "anchor"); + pool.setCol(ref, makeColumn(undefined)); + expect(pool.deriveKeys({ main: ref })).toBeUndefined(); +}); + +test("deriveAnchorPartitionKeys: returns partition keys from a PlRef anchor", () => { + const pool = new TestResultPool(); + const ref = createPlRef("block-1", "anchor"); + pool.setCol(ref, makeColumn(makePartitionedData(["sample-A", "sample-B"]))); + expect(pool.deriveKeys({ main: ref })).toEqual(["sample-A", "sample-B"]); +}); + +test("deriveAnchorPartitionKeys: deduplicates and merges keys from multiple PlRef anchors", () => { + const pool = new TestResultPool(); + const ref1 = createPlRef("block-1", "anchor1"); + const ref2 = createPlRef("block-1", "anchor2"); + pool.setCol(ref1, makeColumn(makePartitionedData(["s1", "s2"]))); + pool.setCol(ref2, makeColumn(makePartitionedData(["s2", "s3"]))); + + const keys = pool.deriveKeys({ a: ref1, b: ref2 }); + expect(keys).toHaveLength(3); + expect(keys).toEqual(expect.arrayContaining(["s1", "s2", "s3"])); +}); + +test("deriveAnchorPartitionKeys: skips PColumnSpec values (not PlRef)", () => { + const pool = new TestResultPool(); + const spec: PColumnSpec = { + kind: "PColumn", + name: "test/spec", + valueType: "Int", + axesSpec: [], + annotations: {}, + }; + expect(pool.deriveKeys({ main: spec })).toEqual([]); +}); + +// ─── getAnchoredPColumns filterToAnchorPartitions ──────────────────────────── + +test("getAnchoredPColumns: injects pl7.app/axisKeys/0 annotation when filterToAnchorPartitions is true", () => { + const pool = new TestResultPool(); + const ref = createPlRef("block-1", "anchor"); + pool.setCol(ref, makeColumn(makePartitionedData(["s1", "s2"]))); + + mockGetColumns.mockReturnValue([makeColumn(makePartitionedData([]))]); + vi.spyOn(pool, "resolveAnchorCtx").mockReturnValue(new AnchoredIdDeriver({}) as any); + + const result = pool.getAnchoredPColumns({ main: ref }, [{ axes: [{ anchor: "main", idx: 0 }] }], { + filterToAnchorPartitions: true, + }); + + expect(result).toHaveLength(1); + expect(result![0].spec.annotations["pl7.app/axisKeys/0"]).toBe(JSON.stringify(["s1", "s2"])); +}); + +test("getAnchoredPColumns: returns undefined when anchor data is not ready", () => { + const pool = new TestResultPool(); + const ref = createPlRef("block-1", "anchor"); + pool.setCol(ref, makeColumn(undefined)); // data not ready + + mockGetColumns.mockReturnValue([makeColumn(makePartitionedData([]))]); + vi.spyOn(pool, "resolveAnchorCtx").mockReturnValue(new AnchoredIdDeriver({}) as any); + + const result = pool.getAnchoredPColumns({ main: ref }, [{ axes: [{ anchor: "main", idx: 0 }] }], { + filterToAnchorPartitions: true, + }); + + expect(result).toBeUndefined(); +}); + +test("getAnchoredPColumns: does not inject annotation when filterToAnchorPartitions is absent", () => { + const pool = new TestResultPool(); + const ref = createPlRef("block-1", "anchor"); + pool.setCol(ref, makeColumn(makePartitionedData(["s1"]))); + + mockGetColumns.mockReturnValue([makeColumn(makePartitionedData([]))]); + vi.spyOn(pool, "resolveAnchorCtx").mockReturnValue(new AnchoredIdDeriver({}) as any); + + const result = pool.getAnchoredPColumns({ main: ref }, [{ axes: [{ anchor: "main", idx: 0 }] }]); + + expect(result).toHaveLength(1); + expect(result![0].spec.annotations?.["pl7.app/axisKeys/0"]).toBeUndefined(); +}); diff --git a/sdk/model/src/render/api.ts b/sdk/model/src/render/api.ts index 01da03f87f..3726707c0d 100644 --- a/sdk/model/src/render/api.ts +++ b/sdk/model/src/render/api.ts @@ -69,7 +69,7 @@ import type { LabelDerivationOps } from "./util/label"; import { deriveLabels } from "./util/label"; import type { APColumnSelectorWithSplit } from "./util/split_selectors"; import { patchInSetFilters } from "./util/pframe_upgraders"; -import { allPColumnsReady } from "./util/pcolumn_data"; +import { allPColumnsReady, getUniquePartitionKeys } from "./util/pcolumn_data"; import type { PColumnDataUniversal } from "./internal"; /** @@ -112,6 +112,15 @@ type UniversalPColumnOpts = { labelOps?: LabelDerivationOps; dontWaitAllData?: boolean; exclude?: AnchoredPColumnSelector | AnchoredPColumnSelector[]; + /** + * When true, restricts each returned column's visible axis-0 key range to the + * partition keys present in the resolved anchor column(s), by injecting a + * `pl7.app/axisKeys/0` annotation onto each returned column spec. + * Graph-maker respects this annotation when building sample dropdowns via + * `findLabelsForColumnAxis`, so only dataset-scoped samples are shown. + * Returns undefined if any anchor column's data is not yet available. + */ + filterToAnchorPartitions?: boolean; } & ResolveAnchorsOptions; type GetOptionsOpts = { @@ -205,6 +214,27 @@ export class ResultPool implements ColumnProvider, AxisLabelProvider { return new AnchoredIdDeriver(resolvedAnchors); } + /** + * Collects unique axis-0 partition key values from all PlRef anchors. + * Returns undefined if any anchor's data is not yet ready (caller should propagate undefined). + * Returns an empty array when anchorsOrCtx is an AnchoredIdDeriver (no refs to resolve). + */ + private deriveAnchorPartitionKeys( + anchorsOrCtx: AnchoredIdDeriver | Record, + ): (string | number)[] | undefined { + if (anchorsOrCtx instanceof AnchoredIdDeriver) return []; + const allKeys = new Set(); + for (const value of Object.values(anchorsOrCtx)) { + if (!isPlRef(value)) continue; + const col = this.getPColumnByRef(value); + if (!col?.data) return undefined; + const uniqueKeys = getUniquePartitionKeys(col.data); + if (uniqueKeys === undefined) return undefined; + for (const key of uniqueKeys[0] ?? []) allKeys.add(key); + } + return [...allKeys]; + } + /** * Returns columns that match the provided anchors and selectors. It applies axis filters and label derivation. * @@ -223,13 +253,27 @@ export class ResultPool implements ColumnProvider, AxisLabelProvider { ): PColumn[] | undefined { const anchorCtx = this.resolveAnchorCtx(anchorsOrCtx); if (!anchorCtx) return undefined; - return new PColumnCollection() + const columns = new PColumnCollection() .addColumnProvider(this) .addAxisLabelProvider(this) .getColumns(predicateOrSelectors, { ...opts, anchorCtx, }); + if (!columns || !opts?.filterToAnchorPartitions) return columns; + const allowedKeys = this.deriveAnchorPartitionKeys(anchorsOrCtx); + if (allowedKeys === undefined) return undefined; + const keysJson = JSON.stringify(allowedKeys); + return columns.map((col) => ({ + ...col, + spec: { + ...col.spec, + annotations: { + ...col.spec.annotations, + "pl7.app/axisKeys/0": keysJson, + }, + }, + })); } /**