deepline 0.3.391 → 0.3.393

package/dist/bundling-sources/shared_libs/play-runtime/scheduler-backends/runtime-event-search.ts

@@ -2040,7 +2040,7 @@ export async function attachRuntimeEventSearchTriggers(
   await client.query(`
     CREATE OR REPLACE FUNCTION ${schema}.project_origin_log_batch(observation ${schema}.outbox, item jsonb, run_org text, play text, compact_logs boolean, canonical_events boolean, at_time timestamptz, doc jsonb, closed boolean)
     RETURNS text[] LANGUAGE plpgsql SECURITY DEFINER SET search_path = pg_catalog AS $function$
-    DECLARE fresh text[]; mismatches bigint; expected_ids text[]; existing_identities jsonb;
+    DECLARE fresh text[]; mismatches bigint; expected_keys jsonb; existing_identities jsonb;
       batch_lines jsonb := item->'lines';
       batch_ids jsonb := item->'lineIds';
       batch_contexts jsonb := item->'lineContexts';
@@ -2110,24 +2110,33 @@ export async function attachRuntimeEventSearchTriggers(
       -- Resolve the whole key set once. A per-line LATERAL lookup can choose
       -- a non-leading event_id index scan when scope/org statistics are skewed,
       -- re-reading the same history for every ID even though the PK is unique.
-      SELECT array_agg('log-origin:' || observation.run_id || ':' || (id #>> '{}') ORDER BY ordinal)
-        INTO expected_ids FROM jsonb_array_elements(batch_ids) WITH ORDINALITY AS ids(id, ordinal);
+      -- Generic plans may scan the scope/org prefix and test ANY(text[]) for
+      -- every row. Use the indexed run prefix and a once-built JSON key set.
+      SELECT jsonb_object_agg('log-origin:' || observation.run_id || ':' || (id #>> '{}'), true)
+        INTO expected_keys FROM jsonb_array_elements(batch_ids) AS ids(id);
       SELECT jsonb_object_agg(event.event_id, jsonb_build_object(
         'runId', event.run_id, 'kind', event.kind, 'digest', event.source_digest))
         INTO existing_identities
         FROM runtime_observability.events event
        WHERE event.runtime_scope = ${quoteLiteral(scope)} AND event.org_id = run_org
-         AND event.event_id = ANY(expected_ids);
+         AND event.run_id = observation.run_id AND expected_keys ? event.event_id;
       -- JSON key lookups preserve one proof per identity without a join whose
       -- generic plan can multiply the two variable-size sets. Tombstones keep
       -- their digest/identity: retry never restores an expired body.
-      SELECT count(*) INTO mismatches
-      FROM jsonb_array_elements(batch_lines) WITH ORDINALITY AS entries(body, ordinal)
-      WHERE existing_identities->expected_ids[ordinal]->>'runId' IS DISTINCT FROM observation.run_id
-        OR existing_identities->expected_ids[ordinal]->>'kind' IS DISTINCT FROM ${RUN_EVENT_KIND_SQL_LITERAL.logAppended}
-        OR existing_identities->expected_ids[ordinal]->>'digest' IS DISTINCT FROM CASE WHEN NULLIF(batch_contexts->(ordinal-1)::int, 'null'::jsonb) IS NULL
+      -- Stream aligned IDs with their lines: positional text[] subscripting
+      -- walks variable-width elements from the start on every lookup.
+      WITH entries AS MATERIALIZED (
+        SELECT body, 'log-origin:' || observation.run_id || ':' || id AS eid,
+          NULLIF(batch_contexts->(ordinal-1)::int, 'null'::jsonb) AS ctx
+        FROM ROWS FROM (jsonb_array_elements(batch_lines), jsonb_array_elements_text(batch_ids))
+          WITH ORDINALITY AS aligned(body, id, ordinal)
+      )
+      SELECT count(*) INTO mismatches FROM entries
+      WHERE existing_identities->eid->>'runId' IS DISTINCT FROM observation.run_id
+        OR existing_identities->eid->>'kind' IS DISTINCT FROM ${RUN_EVENT_KIND_SQL_LITERAL.logAppended}
+        OR existing_identities->eid->>'digest' IS DISTINCT FROM CASE WHEN ctx IS NULL
           THEN encode(sha256(convert_to(body #>> '{}','UTF8')),'hex')
-          ELSE encode(sha256(convert_to(jsonb_build_array(body #>> '{}',batch_contexts->(ordinal-1)::int)::text,'UTF8')),'hex') END;
+          ELSE encode(sha256(convert_to(jsonb_build_array(body #>> '{}',ctx)::text,'UTF8')),'hex') END;
       IF mismatches > 0 THEN
         RAISE EXCEPTION 'Origin log identity changed its content or ingress is closed after terminal acceptance';
       END IF;
@@ -2796,8 +2805,9 @@ export async function attachRuntimeEventSearchTriggers(
         RETURNING last_seq INTO assigned_seq;
         -- Build ordinal lookup once. Generic plans estimate a run and unnest
         -- as tiny, then compare every inserted line against the entire ID list
-        -- in a nested loop. Fresh IDs came only from this transaction's inserts
-        -- for this scope/org/run, so the primary-key set is the complete target.
+        -- in a nested loop or a per-row ANY filter. Fresh IDs came only from
+        -- this transaction's inserts for this scope/org/run, so JSON key
+        -- membership preserves the complete target within the indexed run.
         SELECT jsonb_object_agg(fresh.event_id, fresh.ordinality)
           INTO fresh_positions
           FROM unnest(fresh_event_ids) WITH ORDINALITY AS fresh(event_id, ordinality);
@@ -2805,7 +2815,8 @@ export async function attachRuntimeEventSearchTriggers(
            SET log_seq = assigned_seq - line_count + (fresh_positions->>event.event_id)::bigint
          WHERE event.runtime_scope = ${quoteLiteral(scope)}
            AND event.org_id = run_org
-           AND event.event_id = ANY(fresh_event_ids);
+           AND event.run_id = observation.run_id
+           AND fresh_positions ? event.event_id;
       END IF;
       RETURN;
     END
@@ -3730,11 +3741,13 @@ export async function prepareRuntimeEventRetentionIndex(
     };
     let ready = (await inspect()).rows[0];
     if (ready) {
+      // Even IF NOT EXISTS acquires a table lock. A fully proved existing index
+      // needs no shared DDL; invalid or mismatched definitions still fail loudly.
       assertReady(ready);
     } else {
       await client.query(`CREATE INDEX CONCURRENTLY IF NOT EXISTS ${input.name}
-      ON runtime_observability.events (ingested_at, event_id)
-      WHERE ${input.predicate}`);
+        ON runtime_observability.events (ingested_at, event_id)
+        WHERE ${input.predicate}`);
       ready = (await inspect()).rows[0];
       if (!ready) {
         throw new Error(

package/dist/bundling-sources/shared_libs/play-runtime/tool-execute-retry-policy.ts

@@ -798,10 +798,28 @@ export function classifyToolExecuteHttpFailure(input: {
     executionFenced,
     explicitRetryable,
   });
+  // A status GET may fail after a stateful provider operation was accepted.
+  // Provider pacing remains useful, but retrying the whole tool would repeat
+  // its POST. This explicit recovery proof is stronger than a rejected-call
+  // 429; ordinary pacing denials retain their established retry semantics.
+  const failureBody = parseJsonObject(input.bodyText);
+  const flatPublicDetails = failureBody?.details;
+  const acceptedOperationNeedsRecovery =
+    explicitRetryable === false &&
+    (error instanceof ToolExecutionError
+      ? error.origin === 'provider'
+      : failureBody?.failure_origin === 'provider') &&
+    ((error instanceof ToolExecutionError &&
+      error.publicDetails?.operationAccepted === true) ||
+      (flatPublicDetails !== null &&
+        typeof flatPublicDetails === 'object' &&
+        (flatPublicDetails as Record<string, unknown>).operationAccepted ===
+          true));
   const shouldRetry =
-    providerOwnedRateLimit ||
-    attestedProviderDeadline ||
-    (retryDecision.retryable && input.attempt < retryDecision.attemptCap);
+    !acceptedOperationNeedsRecovery &&
+    (providerOwnedRateLimit ||
+      attestedProviderDeadline ||
+      (retryDecision.retryable && input.attempt < retryDecision.attemptCap));
   // In-attempt retry safety and later-run repairability are different
   // decisions. A non-idempotent 5xx must not be repeated automatically inside
   // this execution attempt, but an explicit later `plays run` must be able to
@@ -862,6 +880,9 @@ export function classifyToolExecuteHttpFailure(input: {
                     : TOOL_EXECUTE_RETRY_DELAY_FALLBACK_MS;
   return {
     ...retryDecision,
+    ...(acceptedOperationNeedsRecovery
+      ? { retryable: false, reason: 'non_retryable_status' as const }
+      : {}),
     error,
     shouldRetry,
     isRateLimit: input.status === 429,

package/dist/bundling-sources/shared_libs/play-runtime/tool-result.ts

@@ -710,7 +710,18 @@ function toV2RawToolOutputPathPreservingProviderData(path: string): string {
 }
 
 function candidateResultPaths(path: string): string[] {
+  // `result.data` is the public envelope. A following `data` belongs to the
+  // provider and must survive (e.g. QuickEnrich's result.data.data rows).
+  const normalized = String(path || '')
+    .trim()
+    .replace(/^\./, '');
+  const providerPath = normalized.startsWith('result.data.')
+    ? `toolResponse.raw.${normalized.slice('result.data.'.length)}`
+    : normalized === 'result.data'
+      ? 'toolResponse.raw'
+      : '';
   const candidates = [
+    providerPath,
     normalizeResultPath(path),
     toV2RawToolOutputPathPreservingProviderData(path),
   ];
@@ -802,9 +813,29 @@ function acceptedTargetValue(
   ) {
     return undefined;
   }
+  if (target === 'email') return emailTargetValue(value) ?? undefined;
   return isMeaningfulValue(value) ? value : undefined;
 }
 
+function emailTargetValue(value: unknown): string | null {
+  if (typeof value === 'string') {
+    const email = value.trim();
+    return /^[^\s@*]+@[^\s@*]+\.[^\s@*]+$/.test(email) ? email : null;
+  }
+  if (Array.isArray(value)) {
+    for (const entry of value) {
+      const email = emailTargetValue(entry);
+      if (email) return email;
+    }
+  } else if (isRecord(value)) {
+    for (const key of ['email', 'email_address', 'address', 'value']) {
+      const email = emailTargetValue(value[key]);
+      if (email) return email;
+    }
+  }
+  return null;
+}
+
 function getFirstMeaningfulValueAtPath(
   root: unknown,
   path: string,

package/dist/bundling-sources/shared_libs/plays/authoring-contract.ts

@@ -12,12 +12,12 @@ import type { ToolExecutionErrorSchemaVersion } from './tool-execution-error';
 import type { PlayDataset, PlayDatasetInput, PlayDatasetRow } from './dataset';
 
 export const LEGACY_PLAY_AUTHORING_CONTRACT_EDITION = 1 as const;
-// New source uses ordinary list values. Historical artifacts keep their pinned
-// edition; deploy the edition-10 reader before promoting this writer.
+// Edition 11 readers retain every CSV source row. Promote this writer in the
+// follow-up PR after those readers are deployed everywhere.
 export const PLAY_AUTHORING_CONTRACT_EDITION = 10 as const;
 export const PLAY_AUTHORING_INPUT_SCHEMA_SNAPSHOT_EDITION = 3 as const;
 export const SUPPORTED_PLAY_AUTHORING_CONTRACT_EDITIONS = [
-  1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+  1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,
 ] as const;
 
 export type PlayAuthoringContractEdition =
@@ -96,6 +96,13 @@ const PLAY_AUTHORING_CONTRACT_CHANGELOG_DETAILS = [
     newWritesEnd: null,
     readerRemoval: null,
   },
+  {
+    changed:
+      'CSV datasets always use source row identity and process every input occurrence, even when key or rowIdentity: "key" is supplied. Editions 1–10 retain their admitted row identity and completed receipts; non-CSV datasets retain explicit key deduplication.',
+    compatibilityOwner: 'Plays Runtime',
+    newWritesEnd: null,
+    readerRemoval: null,
+  },
 ] as const;
 
 /** Authoring types and runtime projection must select the same getter contract. */
@@ -105,6 +112,12 @@ export function playToolListGetterMode(
   return edition >= 10 ? 'array' : 'dataset';
 }
 
+export function playCsvDatasetRowIdentity(
+  edition: PlayAuthoringContractEdition,
+): 'source_row' | 'key' {
+  return edition >= 11 ? 'source_row' : 'key';
+}
+
 export const PLAY_AUTHORING_CONTRACT_CHANGELOG =
   SUPPORTED_PLAY_AUTHORING_CONTRACT_EDITIONS.map((edition, index) => ({
     edition,
@@ -1133,6 +1146,8 @@ export type PlayAuthoringDatasetDefinitionOptions<InputRow extends object> = {
 export type PlayAuthoringDatasetRunOptions<InputRow extends object> = {
   description?: string;
   key?: PlayAuthoringDatasetRowKey<InputRow>;
+  /** CSV defaults to source row identity. Use key only for intentional deduplication. */
+  rowIdentity?: 'source_row' | 'key';
   onRowError?: 'isolate' | 'fail';
   mode?: 'upsert' | 'net_new';
   /**
@@ -1147,6 +1162,7 @@ export type PlayAuthoringDatasetRunOptions<InputRow extends object> = {
 export const PLAY_AUTHORING_DATASET_RUN_OPTION_FIELDS = [
   'description',
   'key',
+  'rowIdentity',
   'onRowError',
   'mode',
   'undrawnColumns',
@@ -2517,6 +2533,23 @@ export const PLAY_AUTHORING_FIELD_REGISTRY = {
     errorMessage:
       'ctx.dataset run key must be a non-empty field, field list, or key function.',
   },
+  'ctx.dataset.run.rowIdentity': {
+    schema: Type.Union([Type.Literal('source_row'), Type.Literal('key')]),
+    fixtures: {
+      valid: 'source_row',
+      invalid: 'dedupe',
+      absent: undefined,
+      unresolved: { expression: 'rowIdentity' },
+      edition1: 'key',
+    },
+    referenceType: "'source_row' | 'key'",
+    required: false,
+    resolution: 'static-when-present',
+    issueCode: 'play_authoring_dataset_option_invalid',
+    description:
+      'Select source_row to process every input occurrence. CSV datasets default to source_row in edition 11; earlier editions retain key identity. Select key only to intentionally deduplicate business keys.',
+    errorMessage: 'ctx.dataset run rowIdentity must be "source_row" or "key".',
+  },
   'ctx.dataset.run.onRowError': {
     schema: Type.Union([Type.Literal('isolate'), Type.Literal('fail')]),
     fixtures: {
@@ -3315,7 +3348,7 @@ export const PLAY_AUTHORING_CLOUD_TYPE_DECLARATIONS = [
   'export type StepProgramOutput<TProgram> = TProgram extends StepProgram<unknown, infer Output, unknown> ? Output : never;',
   'export type DatasetRowKey<InputRow extends object> = (keyof InputRow & string) | readonly (keyof InputRow & string)[] | ((row: InputRow, index: number) => string | number | readonly unknown[]);',
   'export type DatasetDefinitionOptions<InputRow extends object> = { key?: DatasetRowKey<InputRow> };',
-  `export type DatasetRunOptions<InputRow extends object> = { description?: ${cloudReferenceType('ctx.dataset.run.description')}; key?: DatasetRowKey<InputRow>; onRowError?: ${cloudReferenceType('ctx.dataset.run.onRowError')}; mode?: ${cloudReferenceType('ctx.dataset.run.mode')}; undrawnColumns?: ${cloudReferenceType('ctx.dataset.run.undrawnColumns')} };`,
+  `export type DatasetRunOptions<InputRow extends object> = { description?: ${cloudReferenceType('ctx.dataset.run.description')}; key?: DatasetRowKey<InputRow>; rowIdentity?: ${cloudReferenceType('ctx.dataset.run.rowIdentity')}; onRowError?: ${cloudReferenceType('ctx.dataset.run.onRowError')}; mode?: ${cloudReferenceType('ctx.dataset.run.mode')}; undrawnColumns?: ${cloudReferenceType('ctx.dataset.run.undrawnColumns')} };`,
   'export type DatasetBuilder<InputRow extends object, OutputRow extends object> = {',
   '  withColumn<Name extends string, Value>(name: Name, resolver: ColumnResolver<OutputRow, Value>): DatasetBuilder<InputRow, OutputRow & Record<Name, Value>>;',
   '  withColumn<Name extends string, Value>(name: Name, definition: DatasetColumnDefinition<OutputRow, Value> & { readonly runIf: (row: OutputRow, index: number) => boolean | Promise<boolean> }): DatasetBuilder<InputRow, OutputRow & Record<Name, Value | null>>;',

package/dist/bundling-sources/shared_libs/plays/dataset.ts

@@ -4,6 +4,118 @@ import type { PlayDatasetColumnTypes } from '../play-data-plane/contracts';
 const PLAY_DATASET_BRAND = Symbol.for('deepline.play.dataset');
 const NODE_INSPECT_CUSTOM = Symbol.for('nodejs.util.inspect.custom');
 const residentRowsByDataset = new WeakMap<object, readonly unknown[]>();
+// Source ordinals are execution provenance, never customer columns. Weak
+// references keep streaming sources bounded and transforms preserve ordinals.
+type PlayDatasetSourceRowIdentity = Readonly<{
+  datasetId: string;
+  rowIndex: number;
+}>;
+const sourceRowIdentityByRow = new WeakMap<
+  object,
+  PlayDatasetSourceRowIdentity
+>();
+const sourceRowIdentitiesByRows = new WeakMap<
+  readonly unknown[],
+  readonly (PlayDatasetSourceRowIdentity | undefined)[]
+>();
+type PlayDatasetSourceOccurrence<T> = {
+  row: T;
+  sourceIdentity?: PlayDatasetSourceRowIdentity;
+};
+const sourceOccurrencesByDataset = new WeakMap<
+  object,
+  () => AsyncIterable<PlayDatasetSourceOccurrence<unknown>>
+>();
+
+export function setPlayDatasetSourceRowIdentity<T>(
+  row: T,
+  identity: Readonly<{ datasetId: string; rowIndex: number }>,
+): T {
+  if (row && typeof row === 'object') sourceRowIdentityByRow.set(row, identity);
+  return row;
+}
+
+export function getPlayDatasetSourceRowIdentity(
+  row: unknown,
+): Readonly<{ datasetId: string; rowIndex: number }> | undefined {
+  return row && typeof row === 'object'
+    ? sourceRowIdentityByRow.get(row)
+    : undefined;
+}
+
+/** Container provenance distinguishes occurrences even when a mapper returns a shared typed value. */
+export function getPlayDatasetSourceRowIdentityAt(
+  rows: readonly unknown[],
+  index: number,
+): PlayDatasetSourceRowIdentity | undefined {
+  const identities = sourceRowIdentitiesByRows.get(rows);
+  return identities
+    ? identities[index]
+    : getPlayDatasetSourceRowIdentity(rows[index]);
+}
+
+export function copyPlayDatasetSourceRowIdentity<T>(
+  source: unknown,
+  target: T,
+): T {
+  const identity = getPlayDatasetSourceRowIdentity(source);
+  return identity ? setPlayDatasetSourceRowIdentity(target, identity) : target;
+}
+
+/** A mapper may return the same object for several physical source records. */
+function copySourceOccurrenceToMappedValue<T>(
+  identity: PlayDatasetSourceRowIdentity | undefined,
+  mapped: T,
+): T {
+  if (!identity || !mapped || typeof mapped !== 'object') return mapped;
+  const prototype = Object.getPrototypeOf(mapped);
+  // Plain row values can carry provenance on a distinct execution object.
+  // Typed values retain their mapper result and internal slots; their distinct
+  // occurrences travel in the iterator/container sidecar instead.
+  const occurrence = (
+    Array.isArray(mapped)
+      ? Object.defineProperties([], Object.getOwnPropertyDescriptors(mapped))
+      : prototype === Object.prototype || prototype === null
+        ? Object.create(prototype, Object.getOwnPropertyDescriptors(mapped))
+        : mapped
+  ) as T;
+  return setPlayDatasetSourceRowIdentity(occurrence, identity);
+}
+
+async function* iterateSourceOccurrences<T>(
+  input: PlayDatasetInput<T>,
+): AsyncIterable<PlayDatasetSourceOccurrence<T>> {
+  if (isPlayDataset(input)) {
+    const occurrences = sourceOccurrencesByDataset.get(input);
+    if (occurrences) {
+      yield* occurrences() as AsyncIterable<PlayDatasetSourceOccurrence<T>>;
+      return;
+    }
+    for await (const row of input)
+      yield { row, sourceIdentity: getPlayDatasetSourceRowIdentity(row) };
+    return;
+  }
+  if (Array.isArray(input)) {
+    for (let index = 0; index < input.length; index += 1) {
+      yield {
+        row: input[index]!,
+        sourceIdentity: getPlayDatasetSourceRowIdentityAt(input, index),
+      };
+    }
+    return;
+  }
+  if (
+    input != null &&
+    typeof input === 'object' &&
+    Symbol.asyncIterator in input
+  ) {
+    for await (const row of input as AsyncIterable<T>)
+      yield { row, sourceIdentity: getPlayDatasetSourceRowIdentity(row) };
+    return;
+  }
+  for (const row of input as Iterable<T>)
+    yield { row, sourceIdentity: getPlayDatasetSourceRowIdentity(row) };
+}
 const DEFAULT_MATERIALIZE_LIMIT = 10_000;
 export const PLAY_DATASET_EXECUTION_PAGE_ROWS = 1_000;
 export const PLAY_DATASET_EXECUTION_PAGE_BYTES = 64 * 1024 * 1024;
@@ -32,6 +144,14 @@ export type PlayDatasetWorkProgressSummary = {
   pending: number;
   failed: number;
   degraded?: boolean;
+  /** Fixed-shape source/output accounting; does not redefine legacy totals. */
+  cardinality?: {
+    sourceRows: number;
+    uniqueInputRows: number;
+    duplicateOccurrences: number;
+    outputRows: number;
+    rowIdentity: 'source_row' | 'key';
+  };
   duplicates?: {
     exact?: number;
     semantic?: number;
@@ -658,7 +778,9 @@ function createTransformedPlayDataset<T, U>(
     options?.sourceLabel ??
     `${source.sourceLabel ?? source.tableNamespace ?? source.datasetId}.${transform.kind}`;
 
-  const iterate = async function* (): AsyncIterable<U> {
+  const iterateOccurrences = async function* (): AsyncIterable<
+    PlayDatasetSourceOccurrence<U>
+  > {
     if (transform.kind === 'slice') {
       const bounds = normalizeSliceBounds({
         start: transform.start,
@@ -666,10 +788,10 @@ function createTransformedPlayDataset<T, U>(
         count: await source.count(),
       });
       let index = 0;
-      for await (const row of source) {
+      for await (const occurrence of iterateSourceOccurrences(source)) {
         if (index >= bounds.end) break;
         if (index >= bounds.start) {
-          yield row as unknown as U;
+          yield occurrence as unknown as PlayDatasetSourceOccurrence<U>;
         }
         index += 1;
       }
@@ -678,27 +800,42 @@ function createTransformedPlayDataset<T, U>(
 
     let inputIndex = 0;
     let outputIndex = 0;
-    for await (const row of source) {
+    for await (const { row, sourceIdentity } of iterateSourceOccurrences(
+      source,
+    )) {
       if (transform.kind === 'filter') {
         if (await transform.predicate(row, inputIndex)) {
-          yield row as unknown as U;
+          yield { row: row as unknown as U, sourceIdentity };
           outputIndex += 1;
         }
       } else {
-        yield await transform.mapper(row, outputIndex);
+        yield {
+          row: copySourceOccurrenceToMappedValue(
+            sourceIdentity,
+            await transform.mapper(row, outputIndex),
+          ),
+          sourceIdentity,
+        };
         outputIndex += 1;
       }
       inputIndex += 1;
     }
   };
 
+  const iterate = async function* (): AsyncIterable<U> {
+    for await (const { row } of iterateOccurrences()) yield row;
+  };
+
   const collect = async (limit?: number): Promise<U[]> => {
     const rows: U[] = [];
+    const identities: (PlayDatasetSourceRowIdentity | undefined)[] = [];
+    sourceRowIdentitiesByRows.set(rows, identities);
     const boundedLimit =
       limit === undefined ? undefined : Math.max(0, Math.floor(limit));
     if (boundedLimit === 0) return rows;
-    for await (const row of iterate()) {
+    for await (const { row, sourceIdentity } of iterateOccurrences()) {
       rows.push(row);
+      identities.push(sourceIdentity);
       if (boundedLimit !== undefined && rows.length >= boundedLimit) break;
     }
     return rows;
@@ -722,7 +859,7 @@ function createTransformedPlayDataset<T, U>(
     return total;
   };
 
-  return createDeferredPlayDataset({
+  const dataset = createDeferredPlayDataset({
     datasetKind: source.datasetKind,
     datasetId: transformDatasetId({
       source: source as PlayDataset<unknown>,
@@ -747,6 +884,13 @@ function createTransformedPlayDataset<T, U>(
         }) as AsyncIterable<U>,
     },
   });
+  sourceOccurrencesByDataset.set(
+    dataset,
+    iterateOccurrences as () => AsyncIterable<
+      PlayDatasetSourceOccurrence<unknown>
+    >,
+  );
+  return dataset;
 }
 
 export function createDeferredPlayDataset<T>(input: {
@@ -922,10 +1066,15 @@ export async function* iteratePlayDatasetInputPages<T>(
         : DATASET_UTF8_ENCODER.encode(serialized).byteLength;
     });
   let page: T[] = [];
+  let pageIdentities: (PlayDatasetSourceRowIdentity | undefined)[] = [];
+  sourceRowIdentitiesByRows.set(page, pageIdentities);
   let pageBytes = 0;
   let offset = 0;
   let index = 0;
-  const pushRow = (row: T): Array<{ rows: T[]; offset: number }> => {
+  const pushRow = (
+    row: T,
+    sourceIdentity: PlayDatasetSourceRowIdentity | undefined,
+  ): Array<{ rows: T[]; offset: number }> => {
     options?.onRow?.(row, index);
     index += 1;
     const rowBytes = Math.max(0, Math.ceil(estimateRowBytes(row)));
@@ -934,14 +1083,19 @@ export async function* iteratePlayDatasetInputPages<T>(
       ready.push({ rows: page, offset });
       offset += page.length;
       page = [];
+      pageIdentities = [];
+      sourceRowIdentitiesByRows.set(page, pageIdentities);
       pageBytes = 0;
     }
     page.push(row);
+    pageIdentities.push(sourceIdentity);
     pageBytes += rowBytes;
     if (page.length >= pageSize || pageBytes >= maxPageBytes) {
       ready.push({ rows: page, offset });
       offset += page.length;
       page = [];
+      pageIdentities = [];
+      sourceRowIdentitiesByRows.set(page, pageIdentities);
       pageBytes = 0;
     }
     return ready;
@@ -952,48 +1106,15 @@ export async function* iteratePlayDatasetInputPages<T>(
     const finalPage = page;
     const finalOffset = offset;
     page = [];
+    pageIdentities = [];
+    sourceRowIdentitiesByRows.set(page, pageIdentities);
     pageBytes = 0;
     offset += finalPage.length;
     return { rows: finalPage, offset: finalOffset };
   };
 
-  const emitRow = async function* (row: T) {
-    for (const ready of pushRow(row)) yield ready;
-  };
-
-  if (isPlayDataset(input)) {
-    for await (const row of input) {
-      yield* emitRow(row);
-    }
-    const finalPage = flush();
-    if (finalPage) yield finalPage;
-    return;
-  }
-
-  if (Array.isArray(input)) {
-    for (const row of input) {
-      yield* emitRow(row);
-    }
-    const finalPage = flush();
-    if (finalPage) yield finalPage;
-    return;
-  }
-
-  if (
-    input != null &&
-    typeof input === 'object' &&
-    Symbol.asyncIterator in input
-  ) {
-    for await (const row of input as AsyncIterable<T>) {
-      yield* emitRow(row);
-    }
-    const finalPage = flush();
-    if (finalPage) yield finalPage;
-    return;
-  }
-
-  for (const row of input as Iterable<T>) {
-    yield* emitRow(row);
+  for await (const { row, sourceIdentity } of iterateSourceOccurrences(input)) {
+    for (const ready of pushRow(row, sourceIdentity)) yield ready;
   }
   const finalPage = flush();
   if (finalPage) yield finalPage;
@@ -1003,41 +1124,13 @@ export async function materializePlayDatasetInput<T>(
   input: PlayDatasetInput<T>,
   options?: MaterializePlayDatasetInputOptions<T>,
 ): Promise<T[]> {
-  const pushRow = (rows: T[], row: T): void => {
+  const rows: T[] = [];
+  const identities: (PlayDatasetSourceRowIdentity | undefined)[] = [];
+  sourceRowIdentitiesByRows.set(rows, identities);
+  for await (const { row, sourceIdentity } of iterateSourceOccurrences(input)) {
     options?.onRow?.(row, rows.length);
     rows.push(row);
-  };
-
-  if (isPlayDataset(input)) {
-    const rows: T[] = [];
-    for await (const row of input) {
-      pushRow(rows, row);
-    }
-    return rows;
-  }
-
-  if (Array.isArray(input)) {
-    const rows: T[] = [];
-    for (const row of input) {
-      pushRow(rows, row);
-    }
-    return rows;
-  }
-
-  const rows: T[] = [];
-  if (
-    input != null &&
-    typeof input === 'object' &&
-    Symbol.asyncIterator in input
-  ) {
-    for await (const row of input as AsyncIterable<T>) {
-      pushRow(rows, row);
-    }
-    return rows;
-  }
-
-  for (const row of input as Iterable<T>) {
-    pushRow(rows, row);
+    identities.push(sourceIdentity);
   }
   return rows;
 }