Skip to content

Commit abca80c

Browse files
NopeCharonshuchu
andauthored
fix: Validate entity column types in feast dbt import (#7067)
`feast dbt import` only checked that the entity columns exist. A model whose entity column is declared as FLOAT64, BOOLEAN, BYTES or TIMESTAMP was imported silently, producing an entity keyed on a floating point, boolean, binary or timestamp value. A model with an ARRAY entity column crashed the whole import with an unhandled ValueError from FeatureView. The import now maps each entity column's dbt type through the existing dbt-to-Feast type mapping and: - skips models with an array entity column, with a warning, in the same way as models that miss an entity or timestamp column; - prints a warning for entity columns that map to anything other than String, Int32 or Int64, naming the dbt type and the Feast type it maps to, and still imports the model. This also flags Snowflake's NUMBER(38,0), which the mapping turns into Float64. Columns without a declared data_type default to STRING and do not warn. The list of imported models used for the summary and for --output code generation is now collected while processing models instead of being recomputed afterwards from column presence, so skipped models can no longer leak into the generated file. The dbt how-to documents the check and the new warnings. Fixes #5876 Signed-off-by: NopeCharon <[email protected]> Co-authored-by: Shuchu Han <[email protected]>
1 parent b1f2168 commit abca80c

3 files changed

Lines changed: 228 additions & 10 deletions

File tree

‎docs/how-to-guides/dbt-integration.md‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,8 @@ feast dbt import -m target/manifest.json -e driver_id --tag feast
215215

216216
All specified entity columns must exist in each dbt model being imported. Models missing any entity column will be skipped with a warning.
217217

218+
Entity columns should be string or integer columns. Feast prints a warning for entity columns whose type maps to another Feast type, such as `FLOAT64`, `BOOLEAN`, `BYTES` or `TIMESTAMP`, and still imports the model. Models with an array entity column are skipped with a warning.
219+
218220
### Generated Code
219221

220222
The `--output` flag generates code like:
@@ -431,3 +433,11 @@ Ensure your dbt model includes the entity column specified with `--entity-column
431433
### "Missing timestamp column"
432434

433435
By default, Feast looks for `event_timestamp`. Use `--timestamp-field` to specify a different column name.
436+
437+
### "Entity column ... may not be suitable as an entity key"
438+
439+
The entity column's `data_type` in schema.yml maps to a Feast type other than a string or integer type. Floating point values are unreliable join keys, and boolean, binary and timestamp columns rarely identify an entity. The model is still imported. Use a `STRING` or `INT64` column as the entity, or declare the column's type in schema.yml if it is wrong. `NUMBER` types with a precision above 18, such as Snowflake's `NUMBER(38,0)`, map to `Float64`; declare integer IDs as `INTEGER` or `NUMBER(18,0)` instead.
440+
441+
### "Model ... has array entity column(s)"
442+
443+
Array columns cannot be used as entity keys, so the model is skipped. Choose a scalar column with `--entity-column`.

‎sdk/python/feast/cli/dbt_import.py‎

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -121,8 +121,9 @@ def import_command(
121121
# Generate Python file instead of applying to registry
122122
feast dbt import -m target/manifest.json -e driver_id --tag feast --output features.py
123123
"""
124-
from feast.dbt.mapper import DbtToFeastMapper
125-
from feast.dbt.parser import DbtManifestParser
124+
from feast.dbt.mapper import DbtToFeastMapper, map_dbt_type_to_feast_type
125+
from feast.dbt.parser import DbtManifestParser, DbtModel
126+
from feast.types import Array, Int32, Int64, String
126127

127128
# Parse manifest
128129
click.echo(f"{Fore.CYAN}Parsing dbt manifest: {manifest_path}{Style.RESET_ALL}")
@@ -198,6 +199,7 @@ def import_command(
198199

199200
all_objects: List[Any] = []
200201
entities_created: Dict[str, Any] = {}
202+
valid_models: List[DbtModel] = []
201203

202204
for model in models:
203205
# Validate timestamp field exists
@@ -218,6 +220,32 @@ def import_command(
218220
)
219221
continue
220222

223+
# Validate entity column types: arrays cannot be entity keys, and
224+
# types other than strings and integers make poor join keys
225+
columns_by_name = {c.name: c for c in model.columns}
226+
entity_types = {
227+
e: map_dbt_type_to_feast_type(columns_by_name[e].data_type)
228+
for e in entity_cols
229+
}
230+
array_entities = [e for e, t in entity_types.items() if isinstance(t, Array)]
231+
if array_entities:
232+
click.echo(
233+
f"{Fore.YELLOW}Warning: Model '{model.name}' has array entity "
234+
f"column(s): {', '.join(array_entities)}. Skipping.{Style.RESET_ALL}"
235+
)
236+
continue
237+
238+
for entity_col, feast_type in entity_types.items():
239+
if feast_type not in (String, Int32, Int64):
240+
click.echo(
241+
f"{Fore.YELLOW}Warning: Model '{model.name}' entity column "
242+
f"'{entity_col}' has type "
243+
f"'{columns_by_name[entity_col].data_type}' (Feast type "
244+
f"{feast_type}), which may not be suitable as an entity key. "
245+
f"Recommended types: STRING, INT64. Importing anyway."
246+
f"{Style.RESET_ALL}"
247+
)
248+
221249
# Create or reuse entities (one per entity column)
222250
model_entities: List[Any] = []
223251
for entity_col in entity_cols:
@@ -253,6 +281,7 @@ def import_command(
253281
exclude_columns=excluded,
254282
)
255283
all_objects.append(feature_view)
284+
valid_models.append(model)
256285

257286
click.echo(
258287
f" {Fore.GREEN}✓{Style.RESET_ALL} {model.name}: "
@@ -265,14 +294,6 @@ def import_command(
265294
)
266295
raise SystemExit(0)
267296

268-
# Filter models that were actually processed (have valid columns)
269-
valid_models = [
270-
m
271-
for m in models
272-
if timestamp_field in [c.name for c in m.columns]
273-
and all(e in [c.name for c in m.columns] for e in entity_cols)
274-
]
275-
276297
# Summary
277298
click.echo(f"\n{Fore.CYAN}Summary:{Style.RESET_ALL}")
278299
click.echo(f" Entities: {len(entities_created)}")
Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,187 @@
1+
"""
2+
Unit tests for the entity column type checks of `feast dbt import`.
3+
"""
4+
5+
import json
6+
7+
import pytest
8+
from click.testing import CliRunner
9+
10+
# Skip all tests in this module if dbt-artifacts-parser is not installed
11+
pytest.importorskip("dbt_artifacts_parser", reason="dbt-artifacts-parser not installed")
12+
13+
from feast.cli.cli import cli # noqa: E402
14+
15+
MISSING = object()
16+
17+
18+
def _create_column(name: str, data_type=MISSING):
19+
"""Helper to create a column definition, optionally without a data_type."""
20+
column = {"name": name, "description": "", "tags": [], "meta": {}}
21+
if data_type is not MISSING:
22+
column["data_type"] = data_type
23+
return column
24+
25+
26+
def _create_model_node(name: str, entity_columns: dict):
27+
"""Helper to create a model node with a timestamp, a feature and the given
28+
entity columns."""
29+
return {
30+
"name": name,
31+
"unique_id": f"model.test_project.{name}",
32+
"resource_type": "model",
33+
"package_name": "test_project",
34+
"path": f"models/{name}.sql",
35+
"original_file_path": f"models/{name}.sql",
36+
"fqn": ["test_project", name],
37+
"alias": name,
38+
"checksum": {"name": "sha256", "checksum": "abc123"},
39+
"database": "my_database",
40+
"schema": "analytics",
41+
"description": "",
42+
"columns": {
43+
"event_timestamp": _create_column("event_timestamp", "TIMESTAMP"),
44+
"trip_count": _create_column("trip_count", "INT64"),
45+
**entity_columns,
46+
},
47+
"tags": ["feast"],
48+
"meta": {},
49+
"config": {"enabled": True, "materialized": "table", "tags": []},
50+
"depends_on": {"nodes": [], "macros": []},
51+
}
52+
53+
54+
def _write_manifest(tmp_path, models: dict) -> str:
55+
"""Write a manifest with one model per entry of `models`, which maps model
56+
names to their entity columns, and return its path."""
57+
manifest = {
58+
"metadata": {
59+
"dbt_schema_version": "https://schemas.getdbt.com/dbt/manifest/v9.json",
60+
"dbt_version": "1.5.0",
61+
"generated_at": "2024-01-10T00:00:00Z",
62+
"invocation_id": "test-invocation-id",
63+
"env": {},
64+
"adapter_type": "bigquery",
65+
},
66+
"nodes": {
67+
f"model.test_project.{name}": _create_model_node(name, entity_columns)
68+
for name, entity_columns in models.items()
69+
},
70+
"sources": {},
71+
"macros": {},
72+
"docs": {},
73+
"exposures": {},
74+
"metrics": {},
75+
"groups": {},
76+
"selectors": {},
77+
"disabled": {},
78+
"parent_map": {},
79+
"child_map": {},
80+
}
81+
manifest_path = tmp_path / "manifest.json"
82+
manifest_path.write_text(json.dumps(manifest))
83+
return str(manifest_path)
84+
85+
86+
def _import(manifest_path: str, entity_columns: list, *extra_args: str):
87+
args = ["dbt", "import", "-m", manifest_path, "-d", "file", *extra_args]
88+
for entity_column in entity_columns:
89+
args += ["-e", entity_column]
90+
return CliRunner().invoke(cli, args, catch_exceptions=False)
91+
92+
93+
@pytest.mark.parametrize(
94+
"data_type, feast_type",
95+
[
96+
("FLOAT64", "Float64"),
97+
("NUMBER(38,0)", "Float64"),
98+
("BOOLEAN", "Bool"),
99+
("BYTES", "Bytes"),
100+
("TIMESTAMP", "UnixTimestamp"),
101+
],
102+
)
103+
def test_import_warns_on_unsuitable_entity_column_type(tmp_path, data_type, feast_type):
104+
manifest_path = _write_manifest(
105+
tmp_path,
106+
{"driver_stats": {"driver_id": _create_column("driver_id", data_type)}},
107+
)
108+
109+
result = _import(manifest_path, ["driver_id"], "--dry-run")
110+
111+
assert result.exit_code == 0, result.output
112+
assert (
113+
f"Warning: Model 'driver_stats' entity column 'driver_id' has type "
114+
f"'{data_type}' (Feast type {feast_type}), which may not be suitable "
115+
f"as an entity key." in result.output
116+
)
117+
# The model is still imported
118+
assert "driver_stats: DataSource + FeatureView" in result.output
119+
120+
121+
@pytest.mark.parametrize(
122+
"data_type",
123+
[
124+
"STRING",
125+
"VARCHAR(36)",
126+
"INT64",
127+
"INTEGER",
128+
"NUMBER(10,0)",
129+
pytest.param(MISSING, id="undeclared"),
130+
],
131+
)
132+
def test_import_accepts_string_and_integer_entity_columns(tmp_path, data_type):
133+
manifest_path = _write_manifest(
134+
tmp_path,
135+
{"driver_stats": {"driver_id": _create_column("driver_id", data_type)}},
136+
)
137+
138+
result = _import(manifest_path, ["driver_id"], "--dry-run")
139+
140+
assert result.exit_code == 0, result.output
141+
assert "Warning" not in result.output
142+
assert "driver_stats: DataSource + FeatureView" in result.output
143+
144+
145+
def test_import_warns_only_for_unsuitable_entity_columns(tmp_path):
146+
manifest_path = _write_manifest(
147+
tmp_path,
148+
{
149+
"driver_stats": {
150+
"driver_id": _create_column("driver_id", "INT64"),
151+
"vehicle_id": _create_column("vehicle_id", "FLOAT64"),
152+
}
153+
},
154+
)
155+
156+
result = _import(manifest_path, ["driver_id", "vehicle_id"], "--dry-run")
157+
158+
assert result.exit_code == 0, result.output
159+
assert "entity column 'vehicle_id' has type 'FLOAT64'" in result.output
160+
assert "entity column 'driver_id'" not in result.output
161+
assert "driver_stats: DataSource + FeatureView" in result.output
162+
163+
164+
def test_import_skips_models_with_array_entity_columns(tmp_path):
165+
manifest_path = _write_manifest(
166+
tmp_path,
167+
{
168+
"driver_stats": {"driver_id": _create_column("driver_id", "INT64")},
169+
"driver_lists": {"driver_id": _create_column("driver_id", "ARRAY<INT64>")},
170+
},
171+
)
172+
output_path = tmp_path / "features.py"
173+
174+
result = _import(manifest_path, ["driver_id"], "--output", str(output_path))
175+
176+
assert result.exit_code == 0, result.output
177+
assert (
178+
"Warning: Model 'driver_lists' has array entity column(s): driver_id. "
179+
"Skipping." in result.output
180+
)
181+
assert "driver_stats: DataSource + FeatureView" in result.output
182+
assert "driver_lists: DataSource + FeatureView" not in result.output
183+
assert "FeatureViews: 1" in result.output
184+
185+
generated_code = output_path.read_text()
186+
assert "driver_stats" in generated_code
187+
assert "driver_lists" not in generated_code

0 commit comments

Comments
 (0)