@toloka-platform/sdk 0.12.0 → 0.12.2

package/src/cli/commands/dataset.ts

@@ -3,6 +3,7 @@ import type { Command } from "commander";
 import {
   collectValues,
   columnList,
+  fieldList,
   positiveInteger,
   resolveContext,
   splitParts,
@@ -18,6 +19,7 @@ import {
   loadDataset,
   uploadDataset,
 } from "../data";
+import { addDatasetFields } from "../fields";
 import { assertProjectSelected } from "../project-context";
 import { createDatasetSlice, listDatasetSlices, showDatasetSlice } from "../slices";
 import { listDatasetColumns, splitDataset, type SplitRatios } from "../splits";
@@ -46,8 +48,11 @@ export const registerDatasetCommands = (program: Command, host: CommandHost): vo
       const result = await createDataset(client, { projectId, name });
       context.io.stdout(result.output);
     });
-  dataset
+  const fields = dataset
     .command("fields")
+    .description("print field declarations inferred from uploaded rows, or add dataset fields from them");
+  fields
+    .command("declarations", { isDefault: true })
     .argument("<name>", "dataset name or id")
     .description("print field declarations inferred from uploaded rows")
     .action(async (name: string) => {
@@ -59,20 +64,57 @@ export const registerDatasetCommands = (program: Command, host: CommandHost): vo
       });
       context.io.stdout(result.output);
     });
+  fields
+    .command("add")
+    .argument("<name>", "dataset name or id")
+    .argument("[fields...]", "top-level fields of the uploaded rows to add as dataset fields")
+    .option("--all", "add every top-level field of the uploaded rows that is not a dataset field yet", false)
+    .option("--no-wait", "return as soon as the fields are added, while their values are still being copied")
+    .option("--timeout <ms>", "how long to wait for the values to be copied", positiveInteger)
+    .option("--json", "print JSON", false)
+    .description(
+      "add dataset fields from fields of the uploaded rows, so filters, --balance-by and split columns can use their values; uploads after this fill them too",
+    )
+    .action(
+      async (
+        name: string,
+        names: string[],
+        options: { all: boolean; wait: boolean; timeout?: number; json: boolean },
+      ) => {
+        const { sourceRoot, client, projectId } = await resolveContext(globals(), context);
+        const result = await addDatasetFields(client, {
+          sourceRoot,
+          datasetName: name,
+          fields: names,
+          all: options.all,
+          wait: options.wait,
+          json: options.json,
+          ...(projectId === undefined ? {} : { projectId }),
+          ...(options.timeout === undefined ? {} : { waitFor: { timeoutMs: options.timeout } }),
+        });
+        context.io.stdout(result.output);
+      },
+    );
   dataset
     .command("load")
     .argument("<name>", "dataset name or id")
     .requiredOption("--out <file>", "JSONL file to write")
+    .option(
+      "--fields <fields>",
+      "comma-separated dataset fields to write (default: every field); uploaded holds each uploaded row whole, as one value",
+      fieldList,
+    )
     .option("--limit <n>", "how many rows to keep after download", positiveInteger)
     .option("--timeout <ms>", "how long to wait for the export", positiveInteger)
-    .description("download live dataset rows into a JSONL file")
-    .action(async (name: string, options: { out: string; limit?: number; timeout?: number }) => {
+    .description("download live dataset rows into a JSONL file, one object of the chosen fields per row")
+    .action(async (name: string, options: { out: string; fields?: string[]; limit?: number; timeout?: number }) => {
       const { sourceRoot, client, projectId } = await resolveContext(globals(), context);
       const result = await loadDataset(client, {
         sourceRoot,
         datasetName: name,
         outPath: resolve(options.out),
         ...(projectId === undefined ? {} : { projectId }),
+        ...(options.fields === undefined ? {} : { fields: options.fields }),
         ...(options.limit === undefined ? {} : { limit: options.limit }),
         ...(options.timeout === undefined ? {} : { timeoutMs: options.timeout }),
       });

package/src/cli/context.ts

@@ -145,16 +145,22 @@ export const splitRatios = (raw: string): { train: number; validate: number; tes
 
 export const collectValues = (value: string, previous: readonly string[] = []): string[] => [...previous, value];
 
-export const columnList = (raw: string): string[] => {
-  const columns = raw
-    .split(",")
-    .map((column) => column.trim())
-    .filter((column) => column.length > 0);
-  if (columns.length === 0) {
-    throw new InvalidArgumentError("Expected one or more column names separated by commas, for example text,intent.");
-  }
-  return columns;
-};
+const nameList =
+  (noun: string) =>
+  (raw: string): string[] => {
+    const names = raw
+      .split(",")
+      .map((name) => name.trim())
+      .filter((name) => name.length > 0);
+    if (names.length === 0) {
+      throw new InvalidArgumentError(`Expected one or more ${noun} names separated by commas, for example text,intent.`);
+    }
+    return names;
+  };
+
+export const columnList = nameList("column");
+
+export const fieldList = nameList("field");
 
 const SPLIT_PART_ROLES = ["train", "validate", "test"] as const;
 

package/src/cli/data.ts

@@ -92,6 +92,7 @@ export type LoadDatasetOptions = {
   readonly datasetName: string;
   readonly outPath: string;
   readonly projectId?: string;
+  readonly fields?: readonly string[];
   readonly limit?: number;
   readonly timeoutMs?: number;
   readonly fetch?: typeof globalThis.fetch;
@@ -106,7 +107,7 @@ type DatasetRef = {
 const isJsonObject = (value: unknown): value is Record<string, unknown> =>
   typeof value === "object" && value !== null && !Array.isArray(value);
 
-const summarizeJsonSchema = (schema: unknown): string => {
+export const summarizeJsonSchema = (schema: unknown): string => {
   if (!isJsonObject(schema)) return "unknown";
   if (Array.isArray(schema.anyOf)) {
     const variants = [...new Set(schema.anyOf.map(summarizeJsonSchema))];
@@ -156,6 +157,13 @@ const declaredDatasetFields = async (options: {
 
 const formatDatasetChoice = (dataset: Dataset): string => `${dataset.name} (${dataset.id})`;
 
+const MAX_LISTED_NAMES = 10;
+
+export const listedNames = (names: readonly string[]): string =>
+  names.length <= MAX_LISTED_NAMES
+    ? names.join(", ")
+    : `${names.slice(0, MAX_LISTED_NAMES).join(", ")} and ${names.length - MAX_LISTED_NAMES} more`;
+
 const collectLiveDatasets = async (client: PotatoClient, projectId: string): Promise<Dataset[]> => {
   const datasets: Dataset[] = [];
   for await (const dataset of client.datasets.listAll({ projectId })) datasets.push(dataset);
@@ -210,8 +218,10 @@ export const resolveLiveDataset = async (
     createWhenMissing && datasetId === undefined,
     live.length === 0
       ? `Dataset "${token}" does not exist in project ${projectId}. Run toloka dataset create ${JSON.stringify(token)} first, or toloka apply for a dataset a pipeline declares.`
-      : `Dataset "${token}" does not exist in project ${projectId}. This project has: ${live.map(formatDatasetChoice).join(", ")}.`,
-    { hint: { disabled: true } },
+      : `Dataset "${token}" does not exist in project ${projectId}. This project has: ${listedNames(live.map(formatDatasetChoice))}.`,
+    live.length <= MAX_LISTED_NAMES
+      ? { hint: { disabled: true } }
+      : { hint: { description: "Run toloka dataset list to see every dataset of the project." } },
   );
 
   assertCli(
@@ -440,22 +450,35 @@ export async function uploadDataset(
   const uploadedOutput = `Uploaded ${completed.result.itemsIngested} rows to dataset ${resolved.name}.`;
   let inferredFields: ReadonlyArray<InferredField> = [];
   let inferenceReport: string | undefined;
-  let warning: string | undefined;
+  const warnings: string[] = [];
   try {
     const inferred = await client.datasets.getInferredFields({ datasetId });
     inferredFields = inferred.fields;
     inferenceReport = formatInferredFields(inferredFields);
   } catch (error) {
     const reason = error instanceof Error ? error.message : String(error);
-    warning = `Upload succeeded, but inferred fields could not be fetched: ${reason}`;
+    warnings.push(`Upload succeeded, but inferred fields could not be fetched: ${reason}`);
+  }
+  if (fields.length > 0) {
+    try {
+      await client.waitFor({ type: "datasetFields", datasetId }, options.waitFor);
+    } catch (error) {
+      if (!(error instanceof PotatoApiError && error.code === "wait_timeout")) throw error;
+      warnings.push(
+        `Upload succeeded, but its rows are still being copied into the dataset fields; check them with toloka dataset columns ${fieldsHint} --json.`,
+      );
+    }
   }
+  const addFieldsHint =
+    fields.length === 0 && inferredFields.some((field) => field.depth === 1)
+      ? `To filter, balance or split by these fields, add them as dataset fields: toloka dataset fields add ${fieldsHint} <field>… (or --all).`
+      : undefined;
   return {
     datasetId,
     rowsUploaded: completed.result.itemsIngested,
     inferredFields,
-    output:
-      inferenceReport === undefined ? uploadedOutput : `${uploadedOutput}\n${inferenceReport}`,
-    ...(warning === undefined ? {} : { warning }),
+    output: [uploadedOutput, inferenceReport, addFieldsHint].filter((line) => line !== undefined).join("\n"),
+    ...(warnings.length === 0 ? {} : { warning: warnings.join("\n") }),
   };
 }
 
@@ -539,16 +562,19 @@ export async function loadDataset(
   options: LoadDatasetOptions,
 ): Promise<DatasetCommandResult> {
   const resolved = await resolveLiveDataset(client, options, false);
-  const fieldNames: string[] = [];
+  const datasetRef = resolved.nameIsUnique ? resolved.name : resolved.datasetId;
+  const visibleFields: string[] = [];
   for await (const field of client.datasets.listAllFields({ datasetId: resolved.datasetId })) {
-    if (field.internal || field.name === RAW_UPLOAD_FIELD_NAME) continue;
-    fieldNames.push(field.name);
+    if (!field.internal) visibleFields.push(field.name);
   }
-  assertCli(
-    fieldNames.length > 0,
-    `Dataset ${resolved.name} has no fields to export.`,
-    { hint: { description: "Upload rows first with toloka dataset upload <name> <file>." } },
-  );
+  assertCli(visibleFields.length > 0, `Dataset ${resolved.name} has no rows yet.`, {
+    hint: { description: `Upload rows with toloka dataset upload ${datasetRef} <file>.` },
+  });
+  const unknown = (options.fields ?? []).filter((name) => !visibleFields.includes(name));
+  assertCli(unknown.length === 0, `Dataset ${resolved.name} has no field ${unknown.join(", ")}.`, {
+    hint: { description: `Its fields are ${listedNames(visibleFields)}.` },
+  });
+  const fieldNames = options.fields === undefined ? visibleFields : [...options.fields];
 
   const started = await client.datasets.startExport({
     datasetId: resolved.datasetId,
@@ -587,7 +613,9 @@ export async function loadDataset(
     options.outPath,
     limited.length === 0 ? "" : `${limited.map((row) => JSON.stringify(row)).join("\n")}\n`,
   );
-  return { output: `Wrote ${limited.length} rows to ${options.outPath}.` };
+  return {
+    output: `Wrote ${limited.length} rows with the fields ${listedNames(fieldNames)} to ${options.outPath}.`,
+  };
 }
 
 export async function setProjectSecret(

package/src/cli/fields.ts

@@ -0,0 +1,80 @@
+import type { DatasetField, PotatoClient } from "../client";
+import type { WaitForOptions } from "../wait-for";
+import { assertCli } from "./cli-error";
+import { listedNames, resolveLiveDataset, summarizeJsonSchema } from "./data";
+import { waitForFilledFields } from "./slices";
+import { formatTable } from "./table";
+import { uploadedRowFieldNames } from "./uploaded-fields";
+
+export type AddDatasetFieldsOptions = {
+  readonly sourceRoot: string;
+  readonly datasetName: string;
+  readonly projectId?: string;
+  readonly fields: readonly string[];
+  readonly all: boolean;
+  readonly wait: boolean;
+  readonly waitFor?: WaitForOptions;
+  readonly json: boolean;
+};
+
+const FIELDS_HINT = {
+  description: "Name fields of the uploaded rows, for example label split, or pass --all; toloka dataset upload prints them.",
+};
+
+const FIELD_NAME = /^[a-z]([a-z0-9_]{0,126}[a-z0-9])?$/u;
+
+const FIELD_NAME_RULE =
+  "a dataset field name takes lowercase letters, numbers and underscores, starts with a letter, ends with a letter or number, and has at most 128 characters";
+
+export async function addDatasetFields(
+  client: PotatoClient,
+  options: AddDatasetFieldsOptions,
+): Promise<{ readonly output: string }> {
+  assertCli(!(options.all && options.fields.length > 0), "Name the fields or pass --all, not both.", { hint: FIELDS_HINT });
+  assertCli(options.all || options.fields.length > 0, "Name the fields of the uploaded rows to add, or pass --all.", {
+    hint: FIELDS_HINT,
+  });
+  const dataset = await resolveLiveDataset(client, options, false);
+  const datasetRef = dataset.nameIsUnique ? dataset.name : dataset.datasetId;
+  const existing = new Set<string>();
+  for await (const field of client.datasets.listAllFields({ datasetId: dataset.datasetId })) existing.add(field.name);
+  const candidates = options.all
+    ? (await uploadedRowFieldNames(client, dataset.datasetId)).filter((name) => !existing.has(name))
+    : [...options.fields];
+  const unnamable = candidates.filter((name) => !FIELD_NAME.test(name));
+  assertCli(options.all || unnamable.length === 0, `${unnamable.join(", ")} cannot be a dataset field name: ${FIELD_NAME_RULE}.`, {
+    hint: { description: "Rename the field in the file before uploading it, or pick fields with names a dataset field can take." },
+  });
+  const names = candidates.filter((name) => FIELD_NAME.test(name));
+  const skipped = unnamable.length === 0 ? [] : [`Skipped ${listedNames(unnamable)}: ${FIELD_NAME_RULE}.`];
+
+  if (names.length === 0) {
+    return {
+      output: [`Every field of the rows uploaded to ${dataset.name} that can be a dataset field is one already.`, ...skipped].join(
+        "\n",
+      ),
+    };
+  }
+
+  const created = await client.datasets.createFieldsFromUploadedRows({ datasetId: dataset.datasetId, names });
+  const fields: DatasetField[] = options.wait
+    ? (await waitForFilledFields(client, { datasetId: dataset.datasetId, names, waitFor: options.waitFor })).filter(
+        (field) => names.includes(field.name),
+      )
+    : created;
+
+  if (options.json) return { output: JSON.stringify(fields, null, 2) };
+
+  return {
+    output: [
+      formatTable([
+        ["FIELD", "ID", "TYPE"],
+        ...fields.map((field) => [field.name, field.id, summarizeJsonSchema(field.jsonSchema)]),
+      ]),
+      options.wait
+        ? `Copied the values of ${listedNames(names)} from the uploaded rows; later uploads fill them too.`
+        : `Copying the values of ${listedNames(names)} from the uploaded rows; check with toloka dataset columns ${datasetRef} --json.`,
+      ...skipped,
+    ].join("\n"),
+  };
+}

package/src/cli/filters.ts

@@ -1,5 +1,6 @@
 import type { PotatoClient } from "../client";
-import { assertCli } from "./cli-error";
+import { assertCli, REPORT_HINT } from "./cli-error";
+import { uploadedFieldHint } from "./uploaded-fields";
 
 export type FieldValueFilter = {
   fieldValues: {
@@ -32,19 +33,27 @@ const splitFilter = (filter: string): { fieldName: string; rawValue: string } =>
 
 export const resolveFieldValueFilters = async (
   client: Pick<PotatoClient, "datasets">,
-  datasetId: string,
+  dataset: { readonly datasetId: string; readonly name: string; readonly nameIsUnique: boolean },
   filters: readonly string[],
 ): Promise<FieldValueFilter> => {
+  const { datasetId } = dataset;
+  const datasetRef = dataset.nameIsUnique ? dataset.name : datasetId;
   const parsed = filters.map(splitFilter);
   const fieldIdsByName = new Map<string, string[]>();
   for await (const field of client.datasets.listAllFields({ datasetId })) {
     fieldIdsByName.set(field.name, [...(fieldIdsByName.get(field.name) ?? []), field.id]);
   }
+  const unresolved = parsed.find(({ fieldName }) => fieldIdsByName.get(fieldName)?.length !== 1)?.fieldName;
+  const uploadedHint =
+    unresolved === undefined || fieldIdsByName.has(unresolved)
+      ? undefined
+      : await uploadedFieldHint(client, { datasetId, datasetRef, name: unresolved });
+  assertCli(unresolved === undefined, `Dataset field ${unresolved} was not found or is ambiguous.`, {
+    hint: uploadedHint ?? { description: `Run toloka dataset columns ${datasetRef} to see the column names.` },
+  });
   const fieldValues = parsed.map(({ fieldName, rawValue }) => {
-    const [fieldId, ...others] = fieldIdsByName.get(fieldName) ?? [];
-    assertCli(fieldId && others.length === 0, `Dataset field ${fieldName} was not found or is ambiguous.`, {
-      hint: { description: "Run toloka dataset columns <dataset> to see the column names." },
-    });
+    const [fieldId] = fieldIdsByName.get(fieldName) ?? [];
+    assertCli(fieldId, `Dataset field ${fieldName} has no resolved ID.`, { hint: REPORT_HINT });
     const filter: FieldValueFilter["fieldValues"][number] = {
       fieldId,
       jsonPath: "$",

package/src/cli/slices.ts

@@ -1,4 +1,4 @@
-import type { DatasetSlice, PotatoClient } from "../client";
+import type { DatasetField, DatasetSlice, PotatoClient } from "../client";
 import { PotatoApiError } from "../errors";
 import type { WaitForOptions } from "../wait-for";
 import { assertCli } from "./cli-error";
@@ -76,6 +76,19 @@ export const waitUntilDone = async <Resource>(
   }
 };
 
+export const waitForFilledFields = (
+  client: Pick<PotatoClient, "waitFor">,
+  {
+    datasetId,
+    names,
+    waitFor,
+  }: { readonly datasetId: string; readonly names?: readonly string[]; readonly waitFor?: WaitForOptions },
+): Promise<DatasetField[]> =>
+  waitUntilDone(
+    () => client.waitFor({ type: "datasetFields", datasetId, ...(names === undefined ? {} : { names }) }, waitFor ?? {}),
+    { what: `the fields of dataset ${datasetId} to finish filling`, checkCommand: `dataset columns ${datasetId} --json` },
+  );
+
 const render = async (client: SliceApi, slice: DatasetSlice, json: boolean): Promise<DatasetSliceCommandResult> => ({
   output: json ? JSON.stringify(slice, null, 2) : formatSlice(slice, await fieldNamesOf(client, slice.datasetId)),
 });
@@ -91,6 +104,7 @@ export async function createDatasetSlice(
   );
   const dataset = await resolveLiveDataset(client, options, false);
   const frozenColumns = options.freeze && [...options.freeze];
+  if (options.wait) await waitForFilledFields(client, { datasetId: dataset.datasetId, waitFor: options.waitFor });
   const created = await client.datasets.createSlice(
     options.from === undefined
       ? {
@@ -98,7 +112,7 @@ export async function createDatasetSlice(
           name: options.name,
           ...(options.filters === undefined
             ? {}
-            : { itemFilter: await resolveFieldValueFilters(client, dataset.datasetId, options.filters) }),
+            : { itemFilter: await resolveFieldValueFilters(client, dataset, options.filters) }),
           frozenColumns,
         }
       : { datasetId: dataset.datasetId, name: options.name, fromSlice: options.from, frozenColumns },

package/src/cli/splits.ts

@@ -4,8 +4,9 @@ import { assertCli } from "./cli-error";
 import type { SplitPartSpec } from "./context";
 import { resolveLiveDataset } from "./data";
 import { resolveFieldValueFilters } from "./filters";
-import { fieldNamesOf, formatCount, waitUntilDone } from "./slices";
+import { fieldNamesOf, formatCount, waitForFilledFields, waitUntilDone } from "./slices";
 import { formatTable } from "./table";
+import { uploadedFieldHint } from "./uploaded-fields";
 
 type SplitApi = Pick<PotatoClient, "datasets" | "waitFor">;
 
@@ -30,6 +31,8 @@ export type SplitDatasetOptions = {
   readonly json: boolean;
 };
 
+export const DEFAULT_SPLIT_COLUMN = "split";
+
 const formatPercent = (share: number): string => `${Math.round(share * 100)}%`;
 
 const presetParts = (ratios: SplitRatios, validateOnTest: boolean): SplitPartSpec[] => {
@@ -127,21 +130,73 @@ export const waitForSplitColumn = async (
   return done;
 };
 
+const listFields = async (client: Pick<PotatoClient, "datasets">, datasetId: string): Promise<DatasetField[]> => {
+  const fields: DatasetField[] = [];
+  for await (const field of client.datasets.listAllFields({ datasetId })) fields.push(field);
+  return fields;
+};
+
+const assertFieldsToReadExist = async (
+  client: Pick<PotatoClient, "datasets">,
+  {
+    dataset,
+    fields,
+    column,
+    balanceBy,
+  }: {
+    dataset: { readonly datasetId: string; readonly name: string; readonly nameIsUnique: boolean };
+    fields: readonly DatasetField[];
+    column: string;
+    balanceBy: string | undefined;
+  },
+): Promise<void> => {
+  const datasetRef = dataset.nameIsUnique ? dataset.name : dataset.datasetId;
+  const isField = (name: string) => fields.some((field) => field.name === name || field.id === name);
+  if (!isField(column)) {
+    const uploaded = await uploadedFieldHint(client, { datasetId: dataset.datasetId, datasetRef, name: column });
+    assertCli(
+      uploaded === undefined,
+      `Dataset ${dataset.name} has no field ${column}, but its uploaded rows have one.`,
+      {
+        hint: {
+          description: `Add it as a dataset field with toloka dataset fields add ${datasetRef} ${column}, so the split takes its values as parts, or pass another --column for a new split column.`,
+        },
+      },
+    );
+  }
+  if (balanceBy !== undefined && !isField(balanceBy)) {
+    const uploaded = await uploadedFieldHint(client, { datasetId: dataset.datasetId, datasetRef, name: balanceBy });
+    assertCli(false, `Dataset ${dataset.name} has no field ${balanceBy} to balance by.`, {
+      hint: uploaded ?? { description: `List its fields with toloka dataset columns ${datasetRef}.` },
+    });
+  }
+};
+
 export async function splitDataset(client: PotatoClient, options: SplitDatasetOptions): Promise<DatasetSplitCommandResult> {
   assertCli(options.slice === undefined || options.filters === undefined, "Pass --slice or --filter, not both.", {
     hint: { description: "--slice fills the rows of a slice; --filter fills the rows that match now." },
   });
+  const parts = requestedParts(options);
   const dataset = await resolveLiveDataset(client, options, false);
+  const fields = options.wait
+    ? await waitForFilledFields(client, { datasetId: dataset.datasetId, waitFor: options.waitFor })
+    : await listFields(client, dataset.datasetId);
+  await assertFieldsToReadExist(client, {
+    dataset,
+    fields,
+    column: options.column ?? DEFAULT_SPLIT_COLUMN,
+    balanceBy: options.balanceBy,
+  });
   const field = await client.datasets.split({
     datasetId: dataset.datasetId,
     column: options.column,
-    parts: requestedParts(options),
+    parts,
     balanceBy: options.balanceBy,
     seed: options.seed,
     ...(options.slice === undefined ? {} : { slice: options.slice }),
     ...(options.filters === undefined
       ? {}
-      : { itemFilter: await resolveFieldValueFilters(client, dataset.datasetId, options.filters) }),
+      : { itemFilter: await resolveFieldValueFilters(client, dataset, options.filters) }),
   });
   const done = options.wait ? await waitForSplitColumn(client, { datasetId: dataset.datasetId, field, waitFor: options.waitFor }) : field;
   return {

package/src/cli/trainings.ts

@@ -18,8 +18,9 @@ import type { BaseModelId, TrainingConfigurationInput } from "../../openapi/gen/
 import { assertCli, RETRY_OR_REPORT_HINT, TolokaCliError } from "./cli-error";
 import type { CliErrorHint } from "./cli-error";
 import { resolveLiveDataset } from "./data";
-import { formatCount, waitUntilDone } from "./slices";
-import { waitForSplitColumn } from "./splits";
+import { formatCount, waitForFilledFields, waitUntilDone } from "./slices";
+import { DEFAULT_SPLIT_COLUMN, waitForSplitColumn } from "./splits";
+import { uploadedFieldHint } from "./uploaded-fields";
 import { formatTable, type TableIo as TrainingIo, writeOutput } from "./table";
 
 const ESTIMATE_FILE_HINT: CliErrorHint = {
@@ -161,10 +162,14 @@ const formatCostEstimate = (estimate: TrainingCostEstimate): string => {
     ].join("\n");
   }
 
+  const leftOutAsTest = estimate.inputs.training.heldOutRows ?? 0;
+
   return formatTable([
     ["Training", estimate.trainingId],
     ["Training rows", String(estimate.inputs.training.report?.validCount ?? 0)],
-    ["Held-out test rows", String(estimate.inputs.training.heldOutRows ?? 0)],
+    ...(leftOutAsTest === 0
+      ? []
+      : [["Left out as test", `${formatCount(leftOutAsTest)} train rows that a split column now marks as test`]]),
     ["Training tokens", String(estimate.inputs.training.report?.tokens.total ?? 0)],
     ["Maximum GPU seconds", String(estimate.budget.maximumGpuSeconds)],
     ["Maximum charge", `$${estimate.maximumCharge} ${estimate.currency}`],
@@ -295,10 +300,14 @@ const splitColumnId = async (
 
 const fillTrainingSplitColumn = async (
   client: TrainingDraftApi,
-  { datasetId, datasetName, splitColumn }: { datasetId: string; datasetName: string; splitColumn: string | undefined },
+  {
+    datasetId,
+    datasetName,
+    datasetRef,
+    splitColumn,
+  }: { datasetId: string; datasetName: string; datasetRef: string; splitColumn: string | undefined },
 ): Promise<DatasetField> => {
-  const fields: DatasetField[] = [];
-  for await (const field of client.datasets.listAllFields({ datasetId })) fields.push(field);
+  const fields = await waitForFilledFields(client, { datasetId, waitFor: AUTO_WAIT });
   const splitColumns = fields.filter((field) => field.split !== null);
   assertCli(
     splitColumn !== undefined || splitColumns.length <= 1,
@@ -312,9 +321,23 @@ const fillTrainingSplitColumn = async (
 
   if (existing?.split?.recipe === null) return existing;
 
+  const column = existing?.name ?? splitColumn ?? DEFAULT_SPLIT_COLUMN;
+  if (existing === undefined) {
+    const uploaded = await uploadedFieldHint(client, { datasetId, datasetRef, name: column });
+    assertCli(
+      uploaded === undefined,
+      `Dataset ${datasetName} has no field ${column}, but its uploaded rows have one.`,
+      {
+        hint: {
+          description: `Add it as a dataset field with toloka dataset fields add ${datasetRef} ${column} to train on its parts, or pass --split-column <name> to sample a new split column.`,
+        },
+      },
+    );
+  }
+
   return waitForSplitColumn(client, {
     datasetId,
-    field: await client.datasets.split({ datasetId, column: existing?.name ?? splitColumn ?? "split" }),
+    field: await client.datasets.split({ datasetId, column }),
     waitFor: AUTO_WAIT,
   });
 };
@@ -348,6 +371,7 @@ const splitAndSlice = async (
   const column = await fillTrainingSplitColumn(client, {
     datasetId: dataset.datasetId,
     datasetName: dataset.name,
+    datasetRef: dataset.nameIsUnique ? dataset.name : dataset.datasetId,
     splitColumn,
   });
   const slice = await freezeTrainingSlice(client, { datasetId: dataset.datasetId, name: trainingName });

package/src/cli/uploaded-fields.ts

@@ -0,0 +1,20 @@
+import type { PotatoClient } from "../client";
+import type { CliErrorHint } from "./cli-error";
+
+export const uploadedRowFieldNames = async (
+  client: Pick<PotatoClient, "datasets">,
+  datasetId: string,
+): Promise<string[]> =>
+  (await client.datasets.getInferredFields({ datasetId })).fields
+    .filter((field) => field.depth === 1)
+    .map((field) => field.name);
+
+export const uploadedFieldHint = async (
+  client: Pick<PotatoClient, "datasets">,
+  { datasetId, datasetRef, name }: { readonly datasetId: string; readonly datasetRef: string; readonly name: string },
+): Promise<CliErrorHint | undefined> =>
+  (await uploadedRowFieldNames(client, datasetId)).includes(name)
+    ? {
+        description: `${name} is a field of the uploaded rows but not a dataset field yet. Add it with toloka dataset fields add ${datasetRef} ${name}, then run this again.`,
+      }
+    : undefined;