Skip to content

Commit e689df5

Browse files
authored
Fix seed boolean null handling (#5870)
Signed-off-by: nkwork9999 <143652584+nkwork9999@users.noreply.github.com>
1 parent 9460918 commit e689df5

2 files changed

Lines changed: 40 additions & 3 deletions

File tree

sqlmesh/core/model/definition.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1712,6 +1712,7 @@ def render(
17121712

17131713
def render_seed(self) -> t.Iterator[QueryOrDF]:
17141714
import numpy as np
1715+
import pandas as pd
17151716

17161717
self._ensure_hydrated()
17171718

@@ -1752,8 +1753,6 @@ def render_seed(self) -> t.Iterator[QueryOrDF]:
17521753

17531754
# convert all date/time types to native pandas timestamp
17541755
for column in [*date_columns, *datetime_columns]:
1755-
import pandas as pd
1756-
17571756
df[column] = pd.to_datetime(df[column], infer_datetime_format=True, errors="ignore") # type: ignore
17581757

17591758
# extract datetime.date from pandas timestamp for DATE columns
@@ -1769,7 +1768,7 @@ def render_seed(self) -> t.Iterator[QueryOrDF]:
17691768
)
17701769

17711770
for column in bool_columns:
1772-
df[column] = df[column].apply(lambda i: str_to_bool(str(i)))
1771+
df[column] = df[column].apply(lambda i: None if pd.isna(i) else str_to_bool(str(i)))
17731772

17741773
df.loc[:, string_columns] = df[string_columns].mask(
17751774
cond=lambda x: x.notna(), # type: ignore

tests/core/test_model.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1541,6 +1541,44 @@ def test_seed_model_custom_types(tmp_path):
15411541
assert df["empty_date"].iloc[0] is None
15421542

15431543

1544+
def test_seed_model_boolean_nulls_are_preserved(tmp_path):
1545+
model_csv_path = (tmp_path / "model.csv").absolute()
1546+
1547+
with open(model_csv_path, "w", encoding="utf-8") as fd:
1548+
fd.write("id,test_ind\n")
1549+
fd.write("1,null\n")
1550+
fd.write("2,false\n")
1551+
fd.write("3,true\n")
1552+
fd.write("4,null\n")
1553+
1554+
model = create_seed_model(
1555+
"test_db.test_model",
1556+
SeedKind(path=str(model_csv_path)),
1557+
columns={
1558+
"id": "int",
1559+
"test_ind": "boolean",
1560+
},
1561+
dialect="databricks",
1562+
)
1563+
1564+
df = next(model.render_seed())
1565+
1566+
assert df["test_ind"].to_list() == [None, False, True, None]
1567+
1568+
context = Context(
1569+
config=Config(
1570+
default_connection=DuckDBConnectionConfig(),
1571+
model_defaults=ModelDefaultsConfig(dialect="databricks"),
1572+
)
1573+
)
1574+
context.upsert_model(model)
1575+
1576+
rendered_sql = context.render(model).sql("databricks")
1577+
1578+
assert "CAST(NULL AS BOOLEAN)" in rendered_sql
1579+
assert "(4, NULL)" in rendered_sql
1580+
1581+
15441582
def test_seed_with_special_characters_in_column(tmp_path, assert_exp_eq):
15451583
config = Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))
15461584
context = Context(config=config)

0 commit comments

Comments
 (0)