Update polars to fix PanicException

services/worker/pyproject.toml

@@ -50,7 +50,7 @@ transformers = "^4.36.0"
 typer = "^0.4.2"
 uvicorn = "^0.20.0"
 wget = "^3.2"
-polars = "^0.19.15"
+polars = ">=0.20.0"
 
 [tool.poetry.group.dev.dependencies]
 bandit = "^1.7.4"

services/worker/src/worker/job_runners/split/descriptive_statistics.py

@@ -193,7 +193,7 @@ def compute_histogram(
         hist_df_reverted = df.with_columns(pl.col(column_name).mul(-1).alias("reverse"))["reverse"].hist(
             bins=bins_edges_reverted
         )
-        hist_reverted = hist_df_reverted["reverse_count"].cast(int).to_list()
+        hist_reverted = hist_df_reverted["count"].cast(int).to_list()
         hist = hist_reverted[::-1]
         hist = [hist[0] + hist[1]] + hist[2:-2] + [hist[-2] + hist[-1]]
     else:
@@ -219,9 +219,7 @@ def compute_histogram(
     )
 
 
-def min_max_median_std_nan_count_proportion(
-    data: pl.DataFrame, column_name: str, n_samples: int
-) -> tuple[float, float, float, float, float, int, float]:
+def min_max_mean_median_std(data: pl.DataFrame, column_name: str) -> tuple[float, float, float, float, float]:
     """
     Compute minimum, maximum, median, standard deviation, number of nan samples and their proportion in column data.
     """
@@ -231,7 +229,6 @@ def min_max_median_std_nan_count_proportion(
         mean=pl.all().mean(),
         median=pl.all().median(),
         std=pl.all().std(),
-        nan_count=pl.all().null_count(),
     )
     stats_names = pl.Series(col_stats.keys())
     stats_expressions = [pl.struct(stat) for stat in col_stats.values()]
@@ -240,26 +237,19 @@ def min_max_median_std_nan_count_proportion(
         .select(name=stats_names, stats=pl.concat_list(stats_expressions).flatten())
         .unnest("stats")
     )
-    minimum, maximum, mean, median, std, nan_count = stats[column_name].to_list()
+    minimum, maximum, mean, median, std = stats[column_name].to_list()
     if any(statistic is None for statistic in [minimum, maximum, mean, median, std]):
         # this should be possible only if all values are none
         if not all(statistic is None for statistic in [minimum, maximum, mean, median, std]):
             raise StatisticsComputationError(
                 f"Unexpected result for {column_name=}: "
                 f"Some measures among {minimum=}, {maximum=}, {mean=}, {median=}, {std=} are None but not all of them. "
             )
-        if nan_count != n_samples:
-            raise StatisticsComputationError(
-                f"Unexpected result for {column_name=}: "
-                f"{minimum=}, {maximum=}, {mean=}, {median=}, {std=} are None but not all values in column are None. "
-            )
-        return minimum, maximum, mean, median, std, nan_count, 1.0
+        return minimum, maximum, mean, median, std
 
     minimum, maximum, mean, median, std = np.round([minimum, maximum, mean, median, std], DECIMALS).tolist()
-    nan_proportion = np.round(nan_count / n_samples, DECIMALS).item() if nan_count else 0.0
-    nan_count = int(nan_count)
 
-    return minimum, maximum, mean, median, std, nan_count, nan_proportion
+    return minimum, maximum, mean, median, std
 
 
 def value_counts(data: pl.DataFrame, column_name: str) -> dict[Any, Any]:
@@ -385,10 +375,7 @@ def _compute_statistics(
         data: pl.DataFrame, column_name: str, n_samples: int, n_bins: int
     ) -> NumericalStatisticsItem:
         logging.info(f"Compute statistics for float column {column_name} with polars. ")
-        minimum, maximum, mean, median, std, nan_count, nan_proportion = min_max_median_std_nan_count_proportion(
-            data, column_name, n_samples
-        )
-        logging.debug(f"{minimum=}, {maximum=}, {mean=}, {median=}, {std=}, {nan_count=} {nan_proportion=}")
+        nan_count, nan_proportion = nan_count_proportion(data, column_name, n_samples)
         if nan_count == n_samples:  # all values are None
             return NumericalStatisticsItem(
                 nan_count=n_samples,
@@ -400,6 +387,8 @@ def _compute_statistics(
                 std=None,
                 histogram=None,
             )
+        minimum, maximum, mean, median, std = min_max_mean_median_std(data, column_name)
+        logging.debug(f"{minimum=}, {maximum=}, {mean=}, {median=}, {std=}, {nan_count=} {nan_proportion=}")
 
         hist = compute_histogram(
             data,
@@ -442,11 +431,8 @@ def _compute_statistics(
         data: pl.DataFrame, column_name: str, n_samples: int, n_bins: int
     ) -> NumericalStatisticsItem:
         logging.info(f"Compute statistics for integer column {column_name} with polars. ")
-        minimum, maximum, mean, median, std, nan_count, nan_proportion = min_max_median_std_nan_count_proportion(
-            data, column_name, n_samples
-        )
-        logging.debug(f"{minimum=}, {maximum=}, {mean=}, {median=}, {std=}, {nan_count=} {nan_proportion=}")
-        if nan_count == n_samples:  # all values are None
+        nan_count, nan_proportion = nan_count_proportion(data, column_name, n_samples=n_samples)
+        if nan_count == n_samples:
             return NumericalStatisticsItem(
                 nan_count=n_samples,
                 nan_proportion=1.0,
@@ -458,6 +444,9 @@ def _compute_statistics(
                 histogram=None,
             )
 
+        minimum, maximum, mean, median, std = min_max_mean_median_std(data, column_name)
+        logging.debug(f"{minimum=}, {maximum=}, {mean=}, {median=}, {std=}, {nan_count=} {nan_proportion=}")
+
         minimum, maximum = int(minimum), int(maximum)
         hist = compute_histogram(
             data,