Skip to content
Merged
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
8 changes: 4 additions & 4 deletions sqlmesh/core/test/definition.py
Original file line number Diff line number Diff line change
Expand Up @@ -644,16 +644,16 @@ def _create_df(
return self._execute(query)

rows = values["rows"]
columns_str: t.Optional[t.List[str]] = None
if columns:
columns_str = [str(c) for c in columns]
referenced_columns = list(dict.fromkeys(col for row in rows for col in row))
_raise_if_unexpected_columns(columns, referenced_columns)

if partial:
columns = referenced_columns
columns_str = [c for c in columns_str if c in referenced_columns]

return pd.DataFrame.from_records(
rows, columns=[str(c) for c in columns] if columns else None
)
return pd.DataFrame.from_records(rows, columns=columns_str)

def _add_missing_columns(
self, query: exp.Query, all_columns: t.Optional[t.Collection[str]] = None
Expand Down
36 changes: 36 additions & 0 deletions tests/core/test_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -770,6 +770,42 @@ def test_partial_data_column_order(sushi_context: Context) -> None:
).run()
)

# - output df must differ if sorted by (id, event_date) vs. (event_date, id)
# - output partial must be true
_check_successful_or_raise(
_create_test(
body=load_yaml(
"""
test_foo:
model: sushi.foo
inputs:
sushi.items:
- id: 9876
event_date: 2020-01-01
- id: 1234
name: hello
event_date: 2020-01-02
outputs:
partial: true
query:
- event_date: 2020-01-01
id: 9876
- event_date: 2020-01-02
id: 1234
name: hello
"""
),
test_name="test_foo",
model=sushi_context.upsert_model(
_create_model(
"SELECT id, name, price, event_date FROM sushi.items",
default_catalog=sushi_context.default_catalog,
)
),
context=sushi_context,
).run()
)


def test_partial_data_missing_schemas(sushi_context: Context) -> None:
_check_successful_or_raise(
Expand Down