test: switch from Pandas to Polars

Author: dvp

tests/test_duckdb_dao.py

@@ -2,13 +2,15 @@
 
 from __future__ import annotations
 
+import datetime as dt
+
 from contextlib import closing
 
-import pandas as pd
+import polars as pl
 import pytest
 
 from duckdb import InvalidInputException, connect
-from xpypact.collector import FullDataCollector
+from xpypact.collector import UTC, FullDataCollector
 from xpypact.dao.duckdb import DuckDBDAO as DataAccessObject
 from xpypact.dao.duckdb import create_indices
 from xpypact.dao.duckdb.implementation import save
@@ -45,37 +47,36 @@ def test_save(inventory_with_gamma) -> None:
         dc = FullDataCollector()
         dc.append(inventory_with_gamma, material_id=1, case_id=1)
         save(con, dc.get_result())
-        run_data = dao.load_rundata().df().loc[0]
-        assert run_data["timestamp"] == pd.Timestamp("2022-02-21 01:52:45")
-        assert run_data["run_name"] == "* Material Cu, fluxes 104_2_1_1"
-        nuclides = dao.load_nuclides().df()
-        nuclides = nuclides.set_index(["element", "mass_number", "state"])
-        assert not nuclides.loc["Cu"].empty
+        run_data = dao.load_rundata().pl()
+        assert run_data.select("timestamp").item() == dt.datetime(
+            2022,
+            2,
+            21,
+            1,
+            52,
+            45,
+            tzinfo=UTC,
+        )
+        assert run_data["run_name"].item() == "* Material Cu, fluxes 104_2_1_1"
+        nuclides = dao.load_nuclides().pl()
+        assert not nuclides.filter(pl.col("element").eq("Cu")).is_empty()
         timestep_times = dao.load_timestep_times().pl()
         assert timestep_times.height == 2
-        time_steps = dao.load_time_steps().df()
-        assert not time_steps.empty
-        time_steps = time_steps.set_index("time_step_number")
-        assert not time_steps.loc[2].empty
-        time_step_nuclides = dao.load_time_step_nuclides().filter("material_id=1").df()
-        assert not time_step_nuclides.empty
-        time_step_nuclides = time_step_nuclides.set_index(
-            [
-                "time_step_number",
-                "zai",
-            ],
-        )
-        assert not time_step_nuclides.loc[2, 290630].empty
+        time_steps = dao.load_time_steps().pl()
+        assert not time_steps.is_empty()
+        assert not time_steps.filter(time_step_number=2).is_empty()
+        time_step_nuclides = dao.load_time_step_nuclides().filter("material_id=1").pl()
+        assert not time_step_nuclides.is_empty()
+        assert not time_step_nuclides.filter(time_step_number=2, zai=290630).is_empty()
         _check_gbins(dao)
         create_indices(con)  # check integrity
 
 
-def _check_gbins(dao):
-    gbins = dao.load_gbins().df().set_index("g")
-    assert gbins.loc[0].boundary == pytest.approx(1e-11)
-    gamma = dao.load_gamma().filter("material_id=1").df()
-    assert not gamma.empty
-    gamma = gamma.set_index(["time_step_number", "g"])
-    assert not gamma.loc[2, 1].empty
-    gamma2 = dao.load_gamma(2).df()
-    assert not gamma2.empty
+def _check_gbins(dao: DataAccessObject) -> None:
+    gbins = dao.load_gbins().pl()
+    assert gbins.filter(g=0).select("boundary").item() == pytest.approx(1e-11)
+    gamma = dao.load_gamma().filter("material_id=1").pl()
+    assert not gamma.is_empty()
+    assert not gamma.filter(time_step_number=2, g=1).is_empty()
+    gamma2 = dao.load_gamma(2).pl()
+    assert not gamma2.is_empty()