diff --git a/docs/concepts/models/seed_models.md b/docs/concepts/models/seed_models.md index 6f14960182..2dad95fea7 100644 --- a/docs/concepts/models/seed_models.md +++ b/docs/concepts/models/seed_models.md @@ -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: diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py index e8e122dece..b5b8f3b4bb 100644 --- a/sqlmesh/core/model/definition.py +++ b/sqlmesh/core/model/definition.py @@ -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]: diff --git a/sqlmesh/core/model/seed.py b/sqlmesh/core/model/seed.py index ff12085690..694897ddc8 100644 --- a/sqlmesh/core/model/seed.py +++ b/sqlmesh/core/model/seed.py @@ -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 @@ -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. @@ -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: diff --git a/tests/core/test_model.py b/tests/core/test_model.py index c5d1a71bee..7dc6312385 100644 --- a/tests/core/test_model.py +++ b/tests/core/test_model.py @@ -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) diff --git a/tests/core/test_seed.py b/tests/core/test_seed.py index b6f335b0c6..d531184eb9 100644 --- a/tests/core/test_seed.py +++ b/tests/core/test_seed.py @@ -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 + )