do not fetch original DB primary key when creating training schema

Author: shuangwu5

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:

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)

mostlyai/sdk/_data/pull.py

@@ -60,7 +60,6 @@ def pull(
         _LOG.info(f"tgt: {tgt}")
         _LOG.info(f"model_type: {model_type}")
         _LOG.info(f"max_sample_size: {max_sample_size}")
-        _LOG.info(f"{schema.tables[tgt].primary_key = }")
 
         # initialize progress counter
         tbl_count_rows = 0

mostlyai/sdk/_local/execution/step_pull_training_data.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 from collections.abc import Callable
 from pathlib import Path
 
@@ -22,8 +21,6 @@
 from mostlyai.sdk._data.file.utils import make_data_table_from_container
 from mostlyai.sdk.domain import Connector, Generator, ModelType
 
-_LOG = logging.getLogger(__name__)
-
 
 def execute_step_pull_training_data(
     *,
@@ -35,8 +32,6 @@ def execute_step_pull_training_data(
     update_progress: Callable,
 ) -> tuple[list[str], int]:
     schema = _create_training_schema(generator=generator, connectors=connectors)
-    _LOG.info(f"execute_step_pull_training_data: {generator = }")
-    _LOG.info(f"execute_step_pull_training_data: {schema.tables[target_table_name].primary_key = }")
 
     # fetch total rows
     tgt_table_total_rows = schema.tables[target_table_name].row_count
@@ -72,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: