Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/concepts/models/seed_models.md
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,8 @@ MODEL (
```
**Note:** The dataset schema provided in the definition takes precedence over column names defined in the header of the CSV file. This means that the order in which columns are provided in the `MODEL` definition must match the order of columns in the CSV file.

Declaring `columns` also controls how the CSV values are read. Columns with a text type such as `VARCHAR` or `TEXT` keep their values exactly as written in the CSV, so identifiers like ZIP codes or account numbers keep their leading zeros (`00123` is loaded as `00123`, not `123`). Empty values, and the NA markers controlled by `csv_settings` such as `na_values` and `keep_default_na`, are still loaded as `NULL`. Integer columns that contain empty values are loaded as integers with `NULL`s rather than floats. If `columns` is omitted, types are inferred by Pandas as described above.

### Markers

The `$root` marker can be used in the `path` attribute to indicate that the CSV file path is relative to the root of the SQLMesh project:
Expand Down
6 changes: 5 additions & 1 deletion sqlmesh/core/model/definition.py
Original file line number Diff line number Diff line change
Expand Up @@ -1858,7 +1858,11 @@ def _ensure_hydrated(self) -> None:

@cached_property
def _reader(self) -> CsvSeedReader:
return self.seed.reader(dialect=self.dialect, settings=self.kind.csv_settings)
return self.seed.reader(
dialect=self.dialect,
settings=self.kind.csv_settings,
columns_to_types=self.columns_to_types_,
)

@property
def _data_hash_values_no_sql(self) -> t.List[str]:
Expand Down
101 changes: 83 additions & 18 deletions sqlmesh/core/model/seed.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,10 +89,17 @@ def _na_values_validator(cls, v: t.Any) -> t.Optional[NaValues]:


class CsvSeedReader:
def __init__(self, content: str, dialect: str, settings: CsvSettings):
def __init__(
self,
content: str,
dialect: str,
settings: CsvSettings,
columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
):
self.content = content
self.dialect = dialect
self.settings = settings
self._declared_columns_to_types = columns_to_types or {}
self._df: t.Optional[pd.DataFrame] = None

@property
Expand All @@ -117,25 +124,78 @@ def read(self, batch_size: t.Optional[int] = None) -> t.Generator[pd.DataFrame,
batch_start += batch_size

def _get_df(self) -> pd.DataFrame:
import pandas as pd

if self._df is None:
self._df = pd.read_csv(
StringIO(self.content),
index_col=False,
on_bad_lines="error",
low_memory=False,
**{k: v for k, v in self.settings.dict().items() if v is not None},
)
self._df = self._df.rename(
columns={
col: normalize_identifiers(col, dialect=self.dialect).name
for col in self._df.columns
},
)
df = self._read_csv()
if self._declared_columns_to_types:
self._restore_declared_types(df)
self._df = df

return self._df

def _read_csv(self, **kwargs: t.Any) -> pd.DataFrame:
import pandas as pd

df = pd.read_csv(
StringIO(self.content),
index_col=False,
on_bad_lines="error",
low_memory=False,
**kwargs,
**{k: v for k, v in self.settings.dict().items() if v is not None},
)
return df.rename(
columns={
col: normalize_identifiers(col, dialect=self.dialect).name for col in df.columns
},
)

def _restore_declared_types(self, df: pd.DataFrame) -> None:
"""Fixes columns whose declared type was lost by pandas' type inference (in place).

Inference turns TEXT columns like `00123` into integers (dropping leading zeros) and
integer columns with blanks into floats. Such columns are re-read as raw strings: TEXT
columns keep them as is, and integer columns are converted to Python ints (exact and
unbounded), unless a value isn't an integer literal, in which case they're left as is.
"""
import pandas as pd

text_columns = []
int_columns = []
for name, data_type in self._declared_columns_to_types.items():
# Mirror SeedModel.render_seed: exact name first, then the normalized name
# (the columns of `df` are already normalized)
if name not in df:
name = normalize_identifiers(name, dialect=self.dialect).name
if name not in df:
continue
column = df[name]
if data_type.this in exp.DataType.TEXT_TYPES:
# Unless every value is already a string: e.g. 00123 is read as an int, and
# TRUE/FALSE as bools (in an object column when the column has blanks)
if pd.api.types.infer_dtype(column, skipna=True) != "string":
text_columns.append(name)
elif data_type.is_type(*exp.DataType.INTEGER_TYPES):
if pd.api.types.is_float_dtype(column):
int_columns.append(name)

if not text_columns and not int_columns:
return

raw_df = self._read_csv(dtype=str)
for name in text_columns:
df[name] = raw_df[name]

for name in int_columns:
try:
df[name] = pd.Series(
[None if pd.isna(v) else int(v) for v in raw_df[name]],
index=df.index,
dtype=object,
)
except ValueError:
# A non-integer literal (e.g. 1.5, 1e3): keep the floats inferred by pandas
pass


class Seed(PydanticModel):
"""Represents content of a seed.
Expand All @@ -145,8 +205,13 @@ class Seed(PydanticModel):

content: str

def reader(self, dialect: str = "", settings: t.Optional[CsvSettings] = None) -> CsvSeedReader:
return CsvSeedReader(self.content, dialect, settings or CsvSettings())
def reader(
self,
dialect: str = "",
settings: t.Optional[CsvSettings] = None,
columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
) -> CsvSeedReader:
return CsvSeedReader(self.content, dialect, settings or CsvSettings(), columns_to_types)


def create_seed(path: str | Path) -> Seed:
Expand Down
49 changes: 49 additions & 0 deletions tests/core/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -1649,6 +1649,55 @@ def test_seed_model_boolean_nulls_are_preserved(tmp_path):
assert "(4, NULL)" in rendered_sql


def test_seed_model_declared_types_are_applied_when_reading_csv(tmp_path):
model_csv_path = (tmp_path / "model.csv").absolute()
model_csv_path.write_text(
"code,n,amt,flag,d\n00123,1,1.5,true,2026-01-02\n45958,,,false,\n,3,2.25,,2026-03-04\n",
encoding="utf-8",
)

model = create_seed_model(
"test_db.test_model",
SeedKind(path=str(model_csv_path)),
columns={
"code": "text",
"n": "int",
"amt": "double",
"flag": "boolean",
"d": "date",
},
dialect="duckdb",
)

df = next(model.render_seed())

assert df["code"].to_list() == ["00123", "45958", None]
assert df["n"].to_list() == [1, None, 3]
assert [type(v) for v in df["n"]] == [int, type(None), int]
assert df["amt"].to_list() == [1.5, None, 2.25]
assert df["flag"].to_list() == [True, False, None]
assert df["d"].to_list() == [date(2026, 1, 2), None, date(2026, 3, 4)]

context = Context(
config=Config(
default_connection=DuckDBConnectionConfig(),
model_defaults=ModelDefaultsConfig(dialect="duckdb"),
)
)
context.upsert_model(model)
rendered_sql = context.render(model).sql("duckdb")

assert "('00123', 1, 1.5, TRUE, CAST('2026-01-02' AS DATE))" in rendered_sql
assert "('45958', NULL, NULL, FALSE, NULL)" in rendered_sql
assert "(NULL, 3, 2.25, NULL, CAST('2026-03-04' AS DATE))" in rendered_sql

assert context.engine_adapter.fetchall(rendered_sql) == [
("00123", 1, 1.5, True, date(2026, 1, 2)),
("45958", None, None, False, None),
(None, 3, 2.25, None, date(2026, 3, 4)),
]


def test_seed_with_special_characters_in_column(tmp_path, assert_exp_eq):
config = Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))
context = Context(config=config)
Expand Down
104 changes: 104 additions & 0 deletions tests/core/test_seed.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,3 +97,107 @@ def test_column_hashes():
**seed.reader().column_hashes,
"ds": "3396890652",
}


def _declared(**types: str):
return {name: exp.DataType.build(tpe) for name, tpe in types.items()}


def _values(series: pd.Series):
"""Series values with missing entries as None."""
return [None if pd.isna(v) else v for v in series]


def test_declared_text_keeps_leading_zeros():
seed = Seed(content="code,other,undeclared\n00123,1,007\n45958,2,8\n,3,9\n")
df = next(seed.reader(columns_to_types=_declared(code="text", other="int")).read())

assert _values(df["code"]) == ["00123", "45958", None]
assert df["other"].dtype == "int64"
assert df["other"].to_list() == [1, 2, 3]
assert df["undeclared"].to_list() == [7, 8, 9]


@pytest.mark.parametrize("blank", [False, True])
def test_declared_text_keeps_boolean_like_values(blank):
# pandas reads TRUE/FALSE as bools, in an object column when there are blanks
seed = Seed(content="flag,other\nFALSE,1\nTRUE,2\n" + (",3\n" if blank else ""))
df = next(seed.reader(columns_to_types=_declared(flag="text")).read())

assert _values(df["flag"]) == ["FALSE", "TRUE"] + ([None] if blank else [])


def test_declared_int_with_blank_stays_integer():
seed = Seed(content="n,big,ok\n1,9007199254740993,1\n,,2\n3,3,3\n")
df = next(seed.reader(columns_to_types=_declared(n="int", big="bigint", ok="int")).read())

assert df["n"].to_list() == [1, None, 3]
assert [type(v) for v in df["n"]] == [int, type(None), int]
assert df["big"].to_list() == [9007199254740993, None, 3]
assert all(v is None or type(v) is int for v in df["big"])
assert df["ok"].dtype == "int64"


@pytest.mark.parametrize("value", ["1.5", "1.0", "1e3"])
def test_declared_int_with_non_integer_literals_keeps_inferred_floats(value):
seed = Seed(content=f"n,code\n{value},001\n,002\n3,003\n")
df = next(seed.reader(columns_to_types=_declared(n="int", code="text")).read())
expected = next(seed.reader().read())

assert df["n"].dtype == "float64"
assert _values(df["n"]) == _values(expected["n"])
assert df["code"].to_list() == ["001", "002", "003"]


@pytest.mark.parametrize(
"value",
["9223372036854775808", "99999999999999999999"],
)
@pytest.mark.parametrize("blank", [False, True])
def test_declared_int_out_of_int64_range_is_read_as_before(value, blank):
seed = Seed(content=f"n,id\n{value},1\n{'' if blank else 5},2\n")
declared = seed.reader(columns_to_types=_declared(n="bigint", id="int"))
df = next(declared.read())
expected = next(seed.reader().read())

assert df["n"].to_list() == expected["n"].to_list()
assert df["id"].to_list() == [1, 2]


def test_declared_columns_honor_csv_settings():
seed = Seed(content="code|n\n00123|1\nNONE|\n")
settings = CsvSettings(delimiter="|", na_values=["NONE"])
df = next(
seed.reader(settings=settings, columns_to_types=_declared(code="text", n="int")).read()
)

assert _values(df["code"]) == ["00123", None]
assert _values(df["n"]) == [1, None]


def test_declared_columns_match_normalized_csv_headers():
seed = Seed(content="Code,N\n001,1\n,2\n")
reader = seed.reader(dialect="snowflake", columns_to_types=_declared(CODE="text", N="int"))
df = next(reader.read())

assert list(df.columns) == ["CODE", "N"]
assert _values(df["CODE"]) == ["001", None]


def test_column_hashes_only_change_for_seeds_that_were_read_incorrectly():
declared = _declared(id="int", name="text", code="text", n="int")

unaffected = Seed(content="id,name\n1,alice\n2,bob\n")
assert (
unaffected.reader(columns_to_types=declared).column_hashes
== unaffected.reader().column_hashes
)

# These seeds had wrong values (dropped leading zeros, integers read as floats), so their
# hashes change once after upgrade, which is expected.
for content in ["code\n00123\n45958\n", "n,k\n1,a\n,b\n3,c\n"]:
affected = Seed(content=content)
assert (
affected.reader(columns_to_types=declared).column_hashes
!= affected.reader().column_hashes
)