Enforce correct dtypes in migration

equinor · Nov 28, 2024 · 8a5567d · 8a5567d
1 parent 83a673a
commit 8a5567d
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/src/ert/storage/migration/to8.py b/src/ert/storage/migration/to8.py
@@ -77,8 +77,14 @@ def _migrate_responses_from_netcdf_to_parquet(path: Path) -> None:
                             real_dir / f"{ds_name}.nc", engine="scipy"
                         )
 
-                        pandas_df = gen_data_ds.to_dataframe().dropna()
-                        polars_df = polars.from_pandas(pandas_df.reset_index())
+                        pandas_df = gen_data_ds.to_dataframe().dropna().reset_index()
+                        polars_df = polars.from_pandas(
+                            pandas_df,
+                            schema_overrides={
+                                "values": polars.Float32,
+                                "realization": polars.UInt16,
+                            },
+                        )
                         polars_df = polars_df.rename({"name": "response_key"})
 
                         if "time" in polars_df: