refactor model data loading to handle optional 'param_count' column and improve logging

main.py

@@ -193,24 +193,38 @@ def setup_database():
         )
 
         # Load model data
-        model_df = pl.scan_parquet(
+        model_lazy_df = pl.scan_parquet(
             "hf://datasets/davanstrien/models_with_metadata_and_summaries/data/train-*.parquet"
         )
-        model_row_count = model_df.select(pl.len()).collect().item()
+        model_row_count = model_lazy_df.select(pl.len()).collect().item()
         logger.info(f"Row count of new model data: {model_row_count}")
 
         if model_collection.count() < model_row_count:
-            model_df = model_df.select(
-                [
-                    "modelId",
-                    "summary",
-                    "likes",
-                    "downloads",
-                    "last_modified",
-                    "param_count",
-                ]
-            )
-            model_df = model_df.collect()
+            schema = model_lazy_df.schema
+            select_columns = [
+                "modelId",
+                "summary",
+                "likes",
+                "downloads",
+                "last_modified",
+            ]
+            if "param_count" in schema:
+                logger.info("Found 'param_count' column in model data schema.")
+                select_columns.append("param_count")
+            else:
+                logger.warning(
+                    "'param_count' column not found in model data schema. Will add it with null values."
+                )
+
+            # Select specified columns and then collect
+            model_df = model_lazy_df.select(select_columns).collect()
+
+            # If param_count was not in the original schema, add it now to the collected DataFrame
+            if "param_count" not in model_df.columns:
+                model_df = model_df.with_columns(
+                    pl.lit(None).cast(pl.Int64).alias("param_count")
+                )
+
             total_rows = len(model_df)
 
             for i in range(0, total_rows, BATCH_SIZE):