Record mutate() expression columns in the dataset version schema

[dreadatour]

mutate() columns built from plain SQL expressions (e.g. b=C("a") - 1) were missing from the persisted DatasetVersion.schema. The schema dict is built as {c.name: c.type.to_dict() for c in columns if isinstance(c.type, SQLType)}, and expression columns arrived with raw sqlalchemy types (INTEGER), so the filter silently dropped them. feature_schema and the preview were correct — only the flat schema lost the column, which breaks every consumer that treats it as the column list.

Minimal repro:

import datachain as dc
from datachain import C

dc.read_values(a=[1, 2, 3], name=["x", "y", "z"]).mutate(b=C("a") - 1).save("probe")

ver = catalog.get_dataset("probe", versions=None).get_version("1.0.0")

Before:

ver.schema          # {'sys__id': UInt64, 'sys__rand': UInt64, 'a': Int64, 'name': String}  — no 'b'
ver.feature_schema  # {'a': 'int', 'name': 'str', 'b': 'int'}  — correct

After: ver.schema includes 'b': Int64.

The fix is in the mutate expression branch: type_coerce the expression to its datachain SQLType (python_to_sql(sql_to_python(...))), the same coercion the nullable path already did. Func/window, group_by, and literal columns already carried SQLTypes, so this was the only save path producing raw types. As a side effect the physical column type is now the canonical one (ClickHouse: Int64 instead of Int32; Nullable(Int64) preserved for nullable expressions).

Not every expression type is coercible on main: python_to_sql returns a plain python type for enums, raises for Date/Time, and produces an Array with an invalid item type for ARRAY(Enum(...)). Coercion is therefore guarded — an expression whose type cannot be resolved to a valid SQLType keeps its raw type and stays out of the flat schema, exactly as on main today. Pinned by tests: an enum-cast column (reads work pre-save, save warns, column absent from .schema, signal dropped on reload), a Date cast, and an ARRAY(Enum) cast.

_{🤖 Co-authored by Claude Code and ChatGPT}

[cloudflare-workers-and-pages]

Deploying datachain with    Cloudflare Pages

Latest commit:	7d31e6d
Status:	✅  Deploy successful!
Preview URL:	https://c7dca5e3.datachain-2g6.pages.dev
Branch Preview URL:	https://fix-mutate-column-schema.datachain-2g6.pages.dev

View logs

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[Copilot]

Pull request overview

Ensures mutate() expression columns are persisted in dataset version schemas.

Changes:

• Coerces expression columns to DataChain SQL types.
• Adds schema and cross-operation regression tests.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
src/datachain/lib/dc/datachain.py	Normalizes expression result types during mutation.
tests/unit/lib/test_datachain.py	Adds persistence and composition coverage.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

[dreadatour]

Test failed 5 times in a row:

E       AssertionError: Possible import time regression; took 802ms
E       assert 802 < 800

This is weird, because I've made measurements on localhost. Import time, current branch vs main — 30 interleaved runs each (python -X importtime, same venv, ms):

30 runs each	min	median	mean	stdev
main	182.2	195.1	200.3	20.9
branch	182.0	198.1	201.8	14.7
diff	−0.2	+3.1	+1.6

The diff is inside the noise (stdev ~15–21 ms), and the minimum — which the test asserts on — is identical. The CI failure on 3.14 is environmental: 3.14 imports ~20% slower than 3.10 (214 ms vs 177 ms in clean venvs), putting that runner at ~800–825 ms against the 800 ms cap.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

tests/unit/lib/test_datachain.py:3934

• The signal→column review matrix requires to_records, which has its own SQLType-dependent conversion path (src/datachain/lib/dc/records.py:122-125). Since this PR specifically changes these expressions from raw SQLAlchemy types to SQLType, add a record read-back here so that path is covered too.

    df = chain.to_pandas()
    assert df["minus"].tolist() == [0, 1]

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.