All tests passing

Author: charles-turner-1

intake_esm/cat.py

@@ -245,41 +245,60 @@ def load(
                 data['last_updated'] = None
             cat = cls.model_validate(data)
             if cat.catalog_file:
-                if _mapper.fs.exists(cat.catalog_file):
-                    csv_path = cat.catalog_file
-                else:
-                    csv_path = f'{os.path.dirname(_mapper.root)}/{cat.catalog_file}'
-                cat.catalog_file = csv_path
-                converters = read_csv_kwargs.pop('converters', {})  # Hack
-                pl_df = (
-                    pl.scan_csv(  # See https://github.com/pola-rs/polars/issues/13040 - can't use read_csv.
-                        cat.catalog_file,
-                        storage_options=storage_options,
-                        **read_csv_kwargs,
-                    )
-                    .with_columns(
-                        [
-                            pl.col(colname)
-                            .str.replace('^.', '[')  # Replace first/last chars with [ or ].
-                            .str.replace('.$', ']')  # set/tuple => list
-                            .str.replace_all(
-                                "'",
-                                '"',
-                            )
-                            .str.json_decode()  # This is to do with the way polars reads json - single versus double quotes
-                            for colname in converters.keys()
-                        ]
-                    )
-                    .collect()
+                cat._df, cat._pl_df = cat._df_from_file(
+                    cat, _mapper, storage_options, read_csv_kwargs
                 )
             else:
-                pl_df = pl.DataFrame(cat.catalog_dict)
+                cat._pl_df = pl.DataFrame(cat.catalog_dict)
+                cat._df = cat._pl_df.to_pandas()
 
-            cat._df = pl_df.to_pandas()
-            cat._pl_df = pl_df
             cat._cast_agg_columns_with_iterables()
             return cat
 
+    def _df_from_file(
+        self, cat, _mapper, storage_options, read_csv_kwargs
+    ) -> tuple[pd.DataFrame, pl.DataFrame]:
+        """
+        Reading the catalog from disk is a bit messy right now, as polars doesn't support reading
+        bz2 compressed files directly. So we need to screw around a bit to get what we want.
+        """
+        if _mapper.fs.exists(cat.catalog_file):
+            csv_path = cat.catalog_file
+        else:
+            csv_path = f'{os.path.dirname(_mapper.root)}/{cat.catalog_file}'
+        cat.catalog_file = csv_path
+        converters = read_csv_kwargs.pop('converters', {})  # Hack
+        if cat.catalog_file.endswith('.csv.bz2'):
+            df = pd.read_csv(
+                cat.catalog_file,
+                storage_options=storage_options,
+                **read_csv_kwargs,
+            )
+            return df, pl.from_pandas(df)
+        else:
+            pl_df = (
+                pl.scan_csv(  # See https://github.com/pola-rs/polars/issues/13040 - can't use read_csv.
+                    cat.catalog_file,
+                    storage_options=storage_options,
+                    **read_csv_kwargs,
+                )
+                .with_columns(
+                    [
+                        pl.col(colname)
+                        .str.replace('^.', '[')  # Replace first/last chars with [ or ].
+                        .str.replace('.$', ']')  # set/tuple => list
+                        .str.replace_all(
+                            "'",
+                            '"',
+                        )
+                        .str.json_decode()  # This is to do with the way polars reads json - single versus double quotes
+                        for colname in converters.keys()
+                    ]
+                )
+                .collect()
+            )
+            return pl_df.to_pandas(), pl_df
+
     @property
     def columns_with_iterables(self) -> set[str]:
         """Return a set of columns that have iterables."""

intake_esm/source.py

@@ -83,7 +83,7 @@ def _open_dataset(
         if preprocess is not None:
             ds = preprocess(ds)
 
-    if varname and isinstance(varname, str):
+    if varname is not None and isinstance(varname, str):
         varname = [varname]
 
     if requested_variables:
@@ -102,7 +102,7 @@ def _open_dataset(
         ds = ds.set_coords(scalar_variables)
         ds = ds[variables]
         ds.attrs[OPTIONS['vars_key']] = variables
-    elif varname:
+    elif varname is not None:
         ds.attrs[OPTIONS['vars_key']] = varname
 
     ds = _expand_dims(expand_dims, ds)

tests/test_core.py

@@ -4,6 +4,7 @@
 import intake
 import packaging.version
 import pandas as pd
+import polars as pl
 import pydantic
 import pytest
 import xarray as xr
@@ -262,10 +263,24 @@ def test_catalog_serialize(catalog_type, to_csv_kwargs, json_dump_kwargs, direct
     if directory is None:
         directory = os.getcwd()
     cat = intake.open_esm_datastore(f'{directory}/{name}.json')
+    subset_df = cat_subset.esmcat._pl_df.with_columns(
+        [
+            pl.col(colname).cast(pl.Null)
+            for colname in cat_subset.esmcat._pl_df.columns
+            if cat_subset.esmcat._pl_df.get_column(colname).is_null().all()
+        ]
+    )
+
+    df = cat.esmcat._pl_df.with_columns(
+        [
+            pl.col(colname).cast(pl.Null)
+            for colname in cat.esmcat._pl_df.columns
+            if cat.esmcat._pl_df.get_column(colname).is_null().all()
+        ]
+    )
     assert_frame_equal_pl(
-        cat_subset.esmcat._pl_df,
-        cat.esmcat._pl_df,
-        check_dtypes=False,  # lNull types are chaging ty
+        subset_df,
+        df,
     )
     assert cat.esmcat.id == name