diff --git a/mostlyai/sdk/_data/base.py b/mostlyai/sdk/_data/base.py index c9138f07..057f33a6 100644 --- a/mostlyai/sdk/_data/base.py +++ b/mostlyai/sdk/_data/base.py @@ -189,7 +189,7 @@ def get_context_key(self, table_name: str) -> DataIdentifier | None: def get_primary_key(self, table_name: str) -> DataIdentifier | None: primary_key = None - if self.tables[table_name].primary_key is not None: + if self.tables[table_name].primary_key: # first, check primary key on table primary_key = DataIdentifier(table_name, self.tables[table_name].primary_key) elif context_relations := self.get_child_context_relations(table_name): @@ -260,7 +260,7 @@ def get_scp_relations(self, table: str) -> list[ContextRelation]: def _update_key_encoding_types(self) -> None: for tbl_name, tbl_table in self.tables.items(): - if tbl_table.primary_key is not None: + if tbl_table.primary_key: if tbl_table.primary_key in tbl_table.encoding_types: del tbl_table.encoding_types[tbl_table.primary_key] for rel in self.get_relations_to_table(tbl_name): diff --git a/mostlyai/sdk/_data/db/base.py b/mostlyai/sdk/_data/db/base.py index 418e2a6e..8eb24153 100644 --- a/mostlyai/sdk/_data/db/base.py +++ b/mostlyai/sdk/_data/db/base.py @@ -876,6 +876,7 @@ class SqlAlchemyTable(DBTable, abc.ABC): def __init__(self, *args, **kwargs): self.is_view = kwargs.get("is_view", False) + self.lazy_fetch_primary_key = kwargs.get("lazy_fetch_primary_key", True) super().__init__(*args, **kwargs) def __repr__(self): @@ -1065,7 +1066,7 @@ def _lazy_fetch(self, item: str) -> None: if item == "columns": self.columns = [c.name for c in self._sa_table.columns] elif item == "primary_key": - self.primary_key = self._get_primary_key() + self.primary_key = self._get_primary_key() if self.lazy_fetch_primary_key else None elif item == "dtypes": self.dtypes = self._get_dtypes() else: diff --git a/mostlyai/sdk/_data/file/utils.py b/mostlyai/sdk/_data/file/utils.py index 140cd10c..00115126 100644 --- a/mostlyai/sdk/_data/file/utils.py +++ b/mostlyai/sdk/_data/file/utils.py @@ -88,7 +88,9 @@ def _fetch_file_data_table( return data_table -def make_data_table_from_container(container: DataContainer, is_output=False) -> DataTable: +def make_data_table_from_container( + container: DataContainer, is_output: bool = False, lazy_fetch_primary_key: bool = True +) -> DataTable: if isinstance(container, SqlAlchemyContainer): # handle DB containers data_table_class = container.table_class() @@ -100,4 +102,4 @@ def make_data_table_from_container(container: DataContainer, is_output=False) -> data_table_class = read_data_table_from_path(container, return_class=True) else: raise RuntimeError(f"Unknown container type: {type(container)}") - return data_table_class(container=container, is_output=is_output) + return data_table_class(container=container, is_output=is_output, lazy_fetch_primary_key=lazy_fetch_primary_key) diff --git a/mostlyai/sdk/_data/pull_utils.py b/mostlyai/sdk/_data/pull_utils.py index 3a7239d9..52e0922f 100644 --- a/mostlyai/sdk/_data/pull_utils.py +++ b/mostlyai/sdk/_data/pull_utils.py @@ -448,7 +448,7 @@ def fetch_table_data( ) n_fetched_rows = 0 for idx, chunk_df in enumerate(iterator): - if deduplicate_pks and primary_key is not None and primary_key.column in chunk_df.columns: + if deduplicate_pks and primary_key and primary_key.column in chunk_df.columns: # consider only the first occurrence of each primary key if not chunk_df[primary_key.column].is_unique: drop_idx = chunk_df[primary_key.column].duplicated() @@ -568,7 +568,7 @@ def fetch_target_table( key_fraction_df[FRACTION] = keys[MAX_TGT_ROWS_PER_CTX_KEY] / schema.tables[tgt].row_count key_fraction_df = key_fraction_df.rename(columns={tgt_context_key.ref_name(): tgt_context_key.column}) key_fraction_df = key_fraction_df.drop(columns=[MAX_TGT_ROWS_PER_CTX_KEY]) - elif tgt_primary_key is not None: + elif tgt_primary_key: # flat setup with primary key # sample target table by target primary key # use provided keys to filter and order the target table (max_sample_size=None, do_shuffle=False) diff --git a/mostlyai/sdk/_local/execution/step_pull_training_data.py b/mostlyai/sdk/_local/execution/step_pull_training_data.py index 16cedb09..511b666f 100644 --- a/mostlyai/sdk/_local/execution/step_pull_training_data.py +++ b/mostlyai/sdk/_local/execution/step_pull_training_data.py @@ -67,7 +67,7 @@ def _create_training_schema(generator: Generator, connectors: list[Connector]) - container = create_container_from_connector(connector) container.set_location(table.location) # create DataTable - data_table = make_data_table_from_container(container) + data_table = make_data_table_from_container(container, lazy_fetch_primary_key=False) data_table.name = table.name data_table.primary_key = table.primary_key if table.columns: diff --git a/mostlyai/sdk/domain.py b/mostlyai/sdk/domain.py index f5120705..bc271e99 100644 --- a/mostlyai/sdk/domain.py +++ b/mostlyai/sdk/domain.py @@ -2911,7 +2911,7 @@ def validate_keys_exists_in_columns(cls, values): if values.columns: column_names = {col.name for col in values.columns} pk = values.primary_key - if pk is not None and pk not in column_names: + if pk and pk not in column_names: raise ValueError(f"Primary key column '{pk}' does not exist in the table's columns.") for fk in values.foreign_keys or []: if fk.column not in column_names: @@ -2923,7 +2923,7 @@ def validate_keys_exists_in_columns(cls, values): def validate_pk_and_fks_are_not_overlapping(cls, values): primary_key = values.primary_key foreign_keys = [fk.column for fk in values.foreign_keys or []] - if primary_key is not None and primary_key in foreign_keys: + if primary_key and primary_key in foreign_keys: raise ValueError(f"Column '{primary_key}' is both a primary key and a foreign key.") return values diff --git a/tests/_data/unit/test_pull.py b/tests/_data/unit/test_pull.py index 448df580..1712264a 100644 --- a/tests/_data/unit/test_pull.py +++ b/tests/_data/unit/test_pull.py @@ -2120,7 +2120,15 @@ def test_pull_empty_sequences(self, tmp_path): class TestPullWithDB: - def test_excluded_primary_key(self, tmp_path): + @pytest.mark.parametrize( + "columns,lazy_fetch_primary_key,expected_keys", + [ + (["id", "username", "email"], True, {"primary_key": "id"}), + (["id", "username", "email"], False, {}), + (["username", "email"], True, {}), + ], + ) + def test_excluded_primary_key(self, tmp_path, columns, lazy_fetch_primary_key, expected_keys): db_path = tmp_path / "database.db" with sqlite3.connect(str(db_path)) as conn: cursor = conn.cursor() @@ -2137,14 +2145,18 @@ def test_excluded_primary_key(self, tmp_path): "users": SqliteTable( name="users", container=SqliteContainer(dbname=str(db_path)), - primary_key=None, # we specifically exclude the existing primary key in this test - columns=["username", "email"], + primary_key=None, + columns=columns, + lazy_fetch_primary_key=lazy_fetch_primary_key, ), } schema = Schema(tables=tables) pull(tgt="users", schema=schema, workspace_dir=tmp_path) + keys = read_json(tmp_path / "OriginalData" / "tgt-meta" / "keys.json") + assert keys == expected_keys + df = pd.read_parquet(tmp_path / "OriginalData" / "tgt-data") assert len(df) == 2 - assert set(df.columns) == {"username", "email"} + assert set(df.columns) == set(columns) diff --git a/tools/custom_template/pydantic_v2/BaseModel.jinja2 b/tools/custom_template/pydantic_v2/BaseModel.jinja2 index 64958290..4668fd91 100644 --- a/tools/custom_template/pydantic_v2/BaseModel.jinja2 +++ b/tools/custom_template/pydantic_v2/BaseModel.jinja2 @@ -786,7 +786,7 @@ class {{ class_name }}({{ base_class }}):{% if comment is defined %} # {{ comme if values.columns: column_names = {col.name for col in values.columns} pk = values.primary_key - if pk is not None and pk not in column_names: + if pk and pk not in column_names: raise ValueError(f"Primary key column '{pk}' does not exist in the table's columns.") for fk in values.foreign_keys or []: if fk.column not in column_names: @@ -798,7 +798,7 @@ class {{ class_name }}({{ base_class }}):{% if comment is defined %} # {{ comme def validate_pk_and_fks_are_not_overlapping(cls, values): primary_key = values.primary_key foreign_keys = [fk.column for fk in values.foreign_keys or []] - if primary_key is not None and primary_key in foreign_keys: + if primary_key and primary_key in foreign_keys: raise ValueError(f"Column '{primary_key}' is both a primary key and a foreign key.") return values {%- endif %}{%- if class_name == "SyntheticTableConfiguration" %} diff --git a/tools/model.py b/tools/model.py index f050247f..f7436f92 100644 --- a/tools/model.py +++ b/tools/model.py @@ -602,7 +602,7 @@ def validate_keys_exists_in_columns(cls, values): if values.columns: column_names = {col.name for col in values.columns} pk = values.primary_key - if pk is not None and pk not in column_names: + if pk and pk not in column_names: raise ValueError(f"Primary key column '{pk}' does not exist in the table's columns.") for fk in values.foreign_keys or []: if fk.column not in column_names: @@ -614,7 +614,7 @@ def validate_keys_exists_in_columns(cls, values): def validate_pk_and_fks_are_not_overlapping(cls, values): primary_key = values.primary_key foreign_keys = [fk.column for fk in values.foreign_keys or []] - if primary_key is not None and primary_key in foreign_keys: + if primary_key and primary_key in foreign_keys: raise ValueError(f"Column '{primary_key}' is both a primary key and a foreign key.") return values