diff --git a/lisa/analysis/status.py b/lisa/analysis/status.py index d1679869d901ea7674b3c212f2800010f4d84de7..72b9e5ca5ecfaf9f213c3ed5cfc4532b524be3ba 100644 --- a/lisa/analysis/status.py +++ b/lisa/analysis/status.py @@ -65,7 +65,7 @@ class StatusAnalysis(TraceAnalysisBase): 1, # We want to select the first row, so make sure the filter # evaluates to true at that index. - fill_value=pl.col('overutilized').not_(), + fill_value=pl.col('overutilized').not_().first(), ) ) df = df_refit_index(df, window=trace.window) diff --git a/lisa/trace.py b/lisa/trace.py index 3820c537dc7a68eb2f6c6f65c15144c30713fdce..801d9c60266c22304f305013fa53a13a9090ea51 100644 --- a/lisa/trace.py +++ b/lisa/trace.py @@ -4552,9 +4552,15 @@ class _TraceCache(Loggable): elif isinstance(data, pl.LazyFrame): with pl.StringCache(): try: - data.sink_parquet(path, **kwargs) - # Some LazyFrame cannot be sunk lazily to a parquet file - except polars.exceptions.InvalidOperationError: + # TOOD: revisit when polars streaming engine is complete + # and it does not raise a DeprecationWarning anymore. + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=DeprecationWarning) + data.sink_parquet(path, **kwargs) + # The streaming engine may have issues with some LazyFrames, so + # fall back on collecting. + except Exception: + path.unlink(missing_ok=True) data.collect().write_parquet(path, **kwargs) else: data.to_parquet(path, **kwargs) @@ -4697,7 +4703,7 @@ class _TraceCache(Loggable): except KeyError: swap_entry = _CacheDataSwapEntry(cache_desc_nf) - data_path = os.path.join(swap_dir, swap_entry.data_filename) + data_path = Path(swap_dir, swap_entry.data_filename) # If that would make the swap dir too large, try to do some cleanup if self._estimate_data_swap_size(data) + self._swap_size > self.max_swap_size: @@ -4737,7 +4743,7 @@ class _TraceCache(Loggable): self._swap_content[swap_entry.cache_desc_nf] = swap_entry try: - data_swapped_size = os.stat(data_path).st_size + data_swapped_size = data_path.stat().st_size except FileNotFoundError: data_swapped_size = 0