deepline 0.1.243 → 0.1.244

package/dist/bundling-sources/apps/play-runner-workers/src/runtime/dataset-handles.ts

@@ -1,528 +0,0 @@
-import {
-  applyCsvRenameProjection,
-  cloneCsvAliasedRow,
-  stripCsvProjectedFields,
-  type CsvRenameOptions,
-} from '../../../../shared_libs/play-runtime/csv-rename';
-import {
-  createDeferredPlayDataset,
-  isPlayDataset,
-  type PlayDataset,
-  type PlayDatasetInput,
-  type PlayDatasetKind,
-  type PlayDatasetWorkProgressSummary,
-} from '../../../../shared_libs/plays/dataset';
-
-export const WORKER_DATASET_PREVIEW_ROWS = 5;
-export const WORKER_DATASET_IN_MEMORY_ROWS = 100;
-const STREAM_MATERIALIZE_CHUNK_ROWS = 5_000;
-const STREAM_ITERATE_CHUNK_ROWS = 1_000;
-
-export type DatasetRow = Record<string, unknown>;
-export type WorkerDatasetHandle<T extends DatasetRow> = PlayDataset<T>;
-export type WorkerDatasetInput<T extends DatasetRow> = PlayDatasetInput<T>;
-
-type DatasetChunkReader<T extends DatasetRow> = (
-  chunkSize: number,
-) => AsyncIterable<T[]>;
-
-export type WorkerDatasetCapabilities = {
-  storage: 'materialized' | 'runtime_sheet' | 'csv_stream' | 'iterable';
-  chunkReadable: boolean;
-  pageBacked: boolean;
-  replayable: boolean;
-  countHint: number | null;
-  materialization: 'bounded' | 'in_memory';
-};
-
-const datasetChunkReaders = new WeakMap<
-  object,
-  DatasetChunkReader<DatasetRow>
->();
-const datasetCountHints = new WeakMap<object, number | null>();
-const datasetCapabilities = new WeakMap<object, WorkerDatasetCapabilities>();
-
-function cloneRow<T extends DatasetRow>(row: T): T {
-  const cloned: DatasetRow = {};
-  for (const key of Reflect.ownKeys(row)) {
-    const descriptor = Object.getOwnPropertyDescriptor(row, key);
-    if (!descriptor) continue;
-    Object.defineProperty(cloned, key, descriptor);
-  }
-  return cloned as T;
-}
-
-function internalDatasetRow<T extends DatasetRow>(row: T): T {
-  const stripped = cloneCsvAliasedRow(row);
-  const publicRow: DatasetRow = {};
-  for (const key of Reflect.ownKeys(stripped)) {
-    if (
-      typeof key === 'string' &&
-      /^__deepline(?!CsvProjectedValues$)/.test(key)
-    ) {
-      continue;
-    }
-    const descriptor = Object.getOwnPropertyDescriptor(stripped, key);
-    if (!descriptor) continue;
-    Object.defineProperty(publicRow, key, descriptor);
-  }
-  return publicRow as T;
-}
-
-function materializedDatasetRow<T extends DatasetRow>(row: T): T {
-  const stripped = stripCsvProjectedFields(row) as DatasetRow;
-  return Object.fromEntries(
-    Object.entries(stripped).filter(([key]) => !key.startsWith('__deepline')),
-  ) as T;
-}
-
-function registerChunkReader<T extends DatasetRow>(
-  dataset: PlayDataset<T>,
-  reader: DatasetChunkReader<T>,
-  countHint: number | null,
-  capabilities: Omit<WorkerDatasetCapabilities, 'countHint'>,
-): PlayDataset<T> {
-  datasetChunkReaders.set(
-    dataset as object,
-    reader as DatasetChunkReader<DatasetRow>,
-  );
-  datasetCountHints.set(dataset as object, countHint);
-  datasetCapabilities.set(dataset as object, { ...capabilities, countHint });
-  return dataset;
-}
-
-function normalizedChunkSize(chunkSize: number): number {
-  return Math.max(1, Math.floor(chunkSize));
-}
-
-export function createMaterializedDatasetHandle<T extends DatasetRow>(input: {
-  name: string;
-  rows: readonly T[];
-  datasetKind?: PlayDatasetKind;
-  count?: number;
-  sourceLabel?: string | null;
-  workProgress?: PlayDatasetWorkProgressSummary;
-}): WorkerDatasetHandle<T> {
-  const datasetKind = input.datasetKind ?? 'map';
-  const materializedRows = input.rows.map(cloneRow);
-  const dataset = createDeferredPlayDataset({
-    datasetKind,
-    datasetId: `${datasetKind}:${input.name}`,
-    count: Math.max(0, Math.floor(input.count ?? materializedRows.length)),
-    previewRows: materializedRows
-      .slice(0, WORKER_DATASET_PREVIEW_ROWS)
-      .map(cloneRow),
-    sourceLabel: input.sourceLabel ?? null,
-    tableNamespace: input.name,
-    workProgress: input.workProgress,
-    resolvers: {
-      count: async () =>
-        Math.max(0, Math.floor(input.count ?? materializedRows.length)),
-      peek: async (limit) => materializedRows.slice(0, Math.max(0, limit)),
-      materialize: async (limit) =>
-        limit === undefined
-          ? materializedRows.map(cloneRow)
-          : materializedRows.slice(0, Math.max(0, limit)).map(cloneRow),
-      iterate: () =>
-        ({
-          async *[Symbol.asyncIterator]() {
-            for (const row of materializedRows) {
-              yield cloneRow(row);
-            }
-          },
-        }) as AsyncIterable<T>,
-    },
-  });
-  return registerChunkReader(
-    dataset,
-    async function* (chunkSize: number) {
-      const size = normalizedChunkSize(chunkSize);
-      for (let offset = 0; offset < materializedRows.length; offset += size) {
-        yield materializedRows.slice(offset, offset + size).map(cloneRow);
-      }
-    },
-    Math.max(0, Math.floor(input.count ?? materializedRows.length)),
-    {
-      storage: 'materialized',
-      chunkReadable: true,
-      pageBacked: false,
-      replayable: true,
-      materialization: 'in_memory',
-    },
-  );
-}
-
-export function createPersistedDatasetHandle<T extends DatasetRow>(input: {
-  playName: string;
-  name: string;
-  count: number;
-  datasetKind?: PlayDatasetKind;
-  previewRows?: readonly T[];
-  cachedRows?: readonly T[] | null;
-  sourceLabel?: string | null;
-  workProgress?: PlayDatasetWorkProgressSummary;
-  readRows: (input: { limit: number; offset: number }) => Promise<readonly T[]>;
-  trace?: (phase: string, ms: number, extra?: Record<string, unknown>) => void;
-  nowMs: () => number;
-}): WorkerDatasetHandle<T> {
-  const datasetKind = input.datasetKind ?? 'map';
-  const count = Math.max(0, Math.floor(input.count));
-  const previewRows = (input.previewRows ?? [])
-    .slice(0, WORKER_DATASET_PREVIEW_ROWS)
-    .map(materializedDatasetRow);
-  const cachedRows =
-    input.cachedRows && input.cachedRows.length <= WORKER_DATASET_IN_MEMORY_ROWS
-      ? input.cachedRows.map(internalDatasetRow)
-      : null;
-
-  async function loadRows(limit: number, offset: number): Promise<T[]> {
-    if (limit <= 0) {
-      return [];
-    }
-    const normalizedLimit = normalizedChunkSize(limit);
-    const normalizedOffset = Math.max(0, Math.floor(offset));
-    if (
-      cachedRows &&
-      normalizedOffset < cachedRows.length &&
-      normalizedOffset + normalizedLimit <= cachedRows.length
-    ) {
-      return cachedRows
-        .slice(normalizedOffset, normalizedOffset + normalizedLimit)
-        .map(internalDatasetRow);
-    }
-    const startedAt = input.nowMs();
-    const rows = await input.readRows({
-      limit: normalizedLimit,
-      offset: normalizedOffset,
-    });
-    input.trace?.('dataset.read', input.nowMs() - startedAt, {
-      datasetKind,
-      playName: input.playName,
-      tableNamespace: input.name,
-      limit: normalizedLimit,
-      offset: normalizedOffset,
-      rows: rows.length,
-    });
-    return rows.map(internalDatasetRow);
-  }
-
-  async function* readChunks(
-    chunkSize: number,
-  ): AsyncGenerator<T[], void, void> {
-    const size = normalizedChunkSize(chunkSize);
-    let offset = 0;
-    while (offset < count) {
-      const rows = await loadRows(Math.min(size, count - offset), offset);
-      if (rows.length === 0) {
-        return;
-      }
-      yield rows;
-      offset += rows.length;
-    }
-  }
-
-  const dataset = createDeferredPlayDataset({
-    datasetKind,
-    datasetId: `${datasetKind}:${input.name}`,
-    count,
-    backing: {
-      storage: 'neon_sheet',
-      sheet: {
-        playName: input.playName,
-        tableNamespace: input.name,
-      },
-    },
-    previewRows,
-    sourceLabel: input.sourceLabel ?? null,
-    tableNamespace: input.name,
-    workProgress: input.workProgress,
-    resolvers: {
-      count: async () => count,
-      peek: async (limit) =>
-        (await loadRows(Math.max(0, limit), 0)).map(materializedDatasetRow),
-      materialize: async (limit) => {
-        const rows: T[] = [];
-        const maxRows = limit ?? count;
-        for await (const chunk of readChunks(STREAM_MATERIALIZE_CHUNK_ROWS)) {
-          for (const row of chunk) {
-            if (rows.length >= maxRows) return rows;
-            rows.push(materializedDatasetRow(row));
-          }
-        }
-        return rows;
-      },
-      iterate: () =>
-        ({
-          async *[Symbol.asyncIterator]() {
-            for await (const chunk of readChunks(STREAM_ITERATE_CHUNK_ROWS)) {
-              for (const row of chunk) yield row;
-            }
-          },
-        }) as AsyncIterable<T>,
-    },
-  });
-  return registerChunkReader(
-    dataset,
-    (chunkSize) => readChunks(chunkSize),
-    count,
-    {
-      storage: 'runtime_sheet',
-      chunkReadable: true,
-      pageBacked: true,
-      replayable: true,
-      materialization: 'bounded',
-    },
-  );
-}
-
-export function createCsvDatasetHandle<T extends DatasetRow>(input: {
-  name: string;
-  logicalPath: string;
-  expectedBytes?: number | null;
-  renameOptions?: CsvRenameOptions;
-  open: () => Promise<AsyncIterable<Uint8Array> | null>;
-  streamRows: (
-    byteChunks: AsyncIterable<Uint8Array>,
-    chunkSize: number,
-  ) => AsyncIterable<T[]>;
-  trace?: (phase: string, ms: number, extra?: Record<string, unknown>) => void;
-  nowMs: () => number;
-}): WorkerDatasetHandle<T> {
-  const datasetId = `csv:${input.name}`;
-  let cachedCount: number | null = null;
-
-  async function* readChunks(
-    chunkSize: number,
-  ): AsyncGenerator<T[], void, void> {
-    const startedAt = input.nowMs();
-    let yieldedRows = 0;
-    let yieldedChunks = 0;
-    let observedBytes = 0;
-    let observedByteChunks = 0;
-    const body = await input.open();
-    if (!body) {
-      throw new Error(
-        `ctx.csv("${input.logicalPath}"): dataset source is not reachable from the play runner.`,
-      );
-    }
-    const size = normalizedChunkSize(chunkSize);
-    const countedBody = {
-      async *[Symbol.asyncIterator]() {
-        for await (const chunk of body) {
-          observedBytes += chunk.byteLength;
-          observedByteChunks += 1;
-          yield chunk;
-        }
-      },
-    } satisfies AsyncIterable<Uint8Array>;
-    try {
-      for await (const chunk of input.streamRows(countedBody, size)) {
-        yieldedRows += chunk.length;
-        yieldedChunks += 1;
-        yield applyCsvRenameProjection(chunk, input.renameOptions) as T[];
-      }
-    } finally {
-      input.trace?.('csv.stream', input.nowMs() - startedAt, {
-        name: input.name,
-        logicalPath: input.logicalPath,
-        expectedBytes: input.expectedBytes ?? null,
-        observedBytes,
-        observedByteChunks,
-        yieldedRows,
-        yieldedChunks,
-        chunkSize: size,
-      });
-    }
-    if (
-      typeof input.expectedBytes === 'number' &&
-      input.expectedBytes > 0 &&
-      yieldedRows === 0
-    ) {
-      throw new Error(
-        observedBytes === 0
-          ? `ctx.csv("${input.logicalPath}"): dataset source yielded 0 bytes for ` +
-              `a non-empty ${input.expectedBytes} byte staged file.`
-          : `ctx.csv("${input.logicalPath}"): parsed 0 rows after reading ` +
-              `${observedBytes} bytes from a non-empty ${input.expectedBytes} byte dataset source.`,
-      );
-    }
-  }
-
-  const dataset = createDeferredPlayDataset({
-    datasetKind: 'csv',
-    datasetId,
-    count: 0,
-    previewRows: [],
-    tableNamespace: input.name,
-    resolvers: {
-      count: async () => {
-        if (cachedCount !== null) return cachedCount;
-        let total = 0;
-        for await (const chunk of readChunks(STREAM_MATERIALIZE_CHUNK_ROWS)) {
-          total += chunk.length;
-        }
-        cachedCount = total;
-        return total;
-      },
-      peek: async (limit) => {
-        const rows: T[] = [];
-        for await (const chunk of readChunks(Math.max(1, limit))) {
-          for (const row of chunk) {
-            rows.push(row);
-            if (rows.length >= limit) return rows;
-          }
-        }
-        return rows;
-      },
-      materialize: async (limit) => {
-        const rows: T[] = [];
-        const maxRows = limit ?? Number.POSITIVE_INFINITY;
-        for await (const chunk of readChunks(STREAM_MATERIALIZE_CHUNK_ROWS)) {
-          for (const row of chunk) {
-            if (rows.length >= maxRows) return rows;
-            rows.push(row);
-          }
-        }
-        return rows;
-      },
-      iterate: () =>
-        ({
-          async *[Symbol.asyncIterator]() {
-            for await (const chunk of readChunks(STREAM_ITERATE_CHUNK_ROWS)) {
-              for (const row of chunk) yield row;
-            }
-          },
-        }) as AsyncIterable<T>,
-    },
-  });
-  return registerChunkReader(
-    dataset,
-    (chunkSize) => readChunks(chunkSize),
-    null,
-    {
-      storage: 'csv_stream',
-      chunkReadable: true,
-      pageBacked: false,
-      replayable: true,
-      materialization: 'bounded',
-    },
-  );
-}
-
-export function createInlineDatasetHandle<T extends DatasetRow>(
-  rows: readonly T[],
-  metadata?: {
-    name?: string;
-    kind?: PlayDatasetKind;
-    sourceLabel?: string | null;
-  },
-): WorkerDatasetHandle<T> {
-  return createMaterializedDatasetHandle({
-    name: metadata?.name ?? 'csv',
-    rows,
-    datasetKind: metadata?.kind ?? 'csv',
-    sourceLabel: metadata?.sourceLabel ?? null,
-  });
-}
-
-export function datasetRowCountHint<T extends DatasetRow>(
-  input: WorkerDatasetInput<T>,
-): number | null {
-  if (Array.isArray(input)) return input.length;
-  if (isPlayDataset<T>(input)) {
-    const hinted = datasetCountHints.get(input as object);
-    if (hinted !== undefined) return hinted;
-  }
-  return null;
-}
-
-export function datasetHandleCapabilities<T extends DatasetRow>(
-  input: WorkerDatasetInput<T>,
-): WorkerDatasetCapabilities {
-  if (Array.isArray(input)) {
-    return {
-      storage: 'materialized',
-      chunkReadable: true,
-      pageBacked: false,
-      replayable: true,
-      countHint: input.length,
-      materialization: 'in_memory',
-    };
-  }
-  if (isPlayDataset<T>(input)) {
-    const capabilities = datasetCapabilities.get(input as object);
-    if (capabilities) {
-      return capabilities;
-    }
-    return {
-      storage: 'iterable',
-      chunkReadable: false,
-      pageBacked: Boolean(input.backing),
-      replayable: true,
-      countHint: null,
-      materialization: 'bounded',
-    };
-  }
-  return {
-    storage: 'iterable',
-    chunkReadable: false,
-    pageBacked: false,
-    replayable: false,
-    countHint: null,
-    materialization: 'bounded',
-  };
-}
-
-export function isDatasetHandle<T extends DatasetRow>(
-  value: unknown,
-): value is WorkerDatasetHandle<T> {
-  return isPlayDataset<T>(value);
-}
-
-export function iterDatasetChunks<T extends DatasetRow>(
-  input: WorkerDatasetInput<T>,
-  chunkSize: number,
-): AsyncIterable<T[]> {
-  const size = normalizedChunkSize(chunkSize);
-  if (isPlayDataset<T>(input)) {
-    const chunkReader = datasetChunkReaders.get(input as object);
-    if (chunkReader) {
-      return chunkReader(size) as AsyncIterable<T[]>;
-    }
-  }
-  return {
-    async *[Symbol.asyncIterator]() {
-      let chunk: T[] = [];
-      const flush = function* () {
-        if (chunk.length > 0) {
-          yield chunk;
-          chunk = [];
-        }
-      };
-      const push = function* (row: T) {
-        chunk.push(row);
-        if (chunk.length >= size) {
-          yield chunk;
-          chunk = [];
-        }
-      };
-
-      if (isPlayDataset<T>(input)) {
-        for await (const row of input) yield* push(row);
-        yield* flush();
-        return;
-      }
-      if (
-        input != null &&
-        typeof input === 'object' &&
-        Symbol.asyncIterator in input
-      ) {
-        for await (const row of input as AsyncIterable<T>) yield* push(row);
-        yield* flush();
-        return;
-      }
-      for (const row of input as Iterable<T>) yield* push(row);
-      yield* flush();
-    },
-  };
-}