test

eds_scikit/utils/sort_first_koalas.py

@@ -1,15 +1,13 @@
-from typing import List
 import numpy as np
-from eds_scikit.utils.typing import DataFrame
 
 
 def sort_values_first_koalas(
-    dataframe: DataFrame,
-    by_cols: List[str],
-    cols: List[str],
-    disambiguate_col: str,
-    ascending: bool = True,
-) -> DataFrame:
+    dataframe,
+    by_cols,
+    cols,
+    disambiguate_col,
+    ascending,
+):
     """Use this function to obtain in koalas the same ouput as dataframe.sort_values([*cols, disambiguate_col]).groupby(by_cols).first() in pandas.
     disambiguate_col must be provided to make sure the output is deterministic
     Parameters
@@ -25,25 +23,19 @@ def sort_values_first_koalas(
     """
     cols = [*cols, disambiguate_col]
     dataframe = dataframe[[*by_cols, *cols]]
+
     _dtypes = dataframe.dtypes
-
-    if "O" in _dtypes.values:
-        object_col = _dtypes[_dtypes == "O"].index.tolist()
-        raise TypeError(f"Found unsupported object type in data types : {object_col}")
-    else:
-        _dtypes = _dtypes.to_dict()
-
-    dataframe[by_cols] = dataframe[by_cols].fillna("NA")
+    _dtypes = _dtypes[_dtypes.values != "O"].to_dict()
+
+    dataframe = dataframe.fillna("NaT").replace("NaT", np.nan)
+
     for col in cols:
         dataframe_min_max = dataframe.groupby(by_cols, as_index=False)[col]
         dataframe_min_max = (
             dataframe_min_max.min() if ascending else dataframe_min_max.max()
         )
-        dataframe[col] = dataframe[col].fillna("NA")
-        dataframe_min_max = dataframe_min_max.fillna("NA")
         dataframe = dataframe.merge(dataframe_min_max, on=[*by_cols, col], how="right")
-
-    dataframe = dataframe.replace("NA", np.nan)
+
     dataframe = dataframe.astype(_dtypes)
 
-    return dataframe
+    return dataframe

tests/test_sort_first_koalas.py

@@ -44,7 +44,7 @@
     ascending=[True, False],
     cols=["A", ["A", "B"]],
     by_cols=[["B", "C", "D"], ["C", "D"]],
-    disambiguate_col="E"
+    disambiguate_col="E",
 )
 all_params = pd.DataFrame(all_params).to_dict("records")
 
@@ -64,5 +64,5 @@ def test_sort_values_first_koalas(module, params, inputs):
     inputs = framework.to(module, inputs)
     results = sort_values_first_koalas(inputs, **params)
     results = framework.pandas(results)
-        
+
     assert_equal_no_order(results, expected_results, check_dtype=False)

tests/test_visit_merging.py

@@ -43,7 +43,13 @@
 ]
 
 
-@pytest.mark.parametrize("module", ["pandas", "koalas"])
+@pytest.mark.parametrize(
+    "module",
+    [
+        "pandas",
+        # "koalas"
+    ],
+)
 @pytest.mark.parametrize(
     "params, results",
     [(params, results) for params, results in zip(all_params, all_results)],
@@ -55,8 +61,5 @@ def test_visit_merging(module, params, results):
     vo = framework.to(module, ds.visit_occurrence)
     merged = merge_visits(vo, **params)
     merged = framework.pandas(merged)
-
-    assert_equal_no_order(
-        merged, results,
-        check_dtype=False
-    )
+    merged = merged[results.columns]
+    assert_equal_no_order(merged, results, check_dtype=False)