diff --git a/scripts/sp2xr_pipeline.py b/scripts/sp2xr_pipeline.py index 9bec6ce..7308bac 100644 --- a/scripts/sp2xr_pipeline.py +++ b/scripts/sp2xr_pipeline.py @@ -27,7 +27,12 @@ from sp2xr.distribution import ( make_hist_meta, ) from sp2xr.concentrations import add_concentrations -from sp2xr.schema import CANONICAL_DTYPES, DEFAULT_FLOAT, enforce_schema +from sp2xr.schema import ( + cast_and_arrow, + CANONICAL_DTYPES, + DEFAULT_FLOAT, + enforce_schema, +) def main(): @@ -334,7 +339,7 @@ def main(): ddf_conc = ddf_pbp_hk_dt.map_partitions( add_concentrations, dt=run_config["dt"], meta=meta_conc - ).map_partitions(_cast_and_arrow, meta=meta_conc) + ).map_partitions(cast_and_arrow, meta=meta_conc) idx_target = "datetime64[ns]" ddf_conc = ddf_conc.map_partitions( @@ -481,7 +486,7 @@ def main(): BC_type=run_config["BC_type"], t=1, meta=meta_hist, - ).map_partitions(_cast_and_arrow, meta=meta_hist) + ).map_partitions(cast_and_arrow, meta=meta_hist) results.append(ddf_out) # --- Scattering histogram @@ -512,7 +517,7 @@ def main(): BC_type=None, t=1, meta=meta_hist, - ).map_partitions(_cast_and_arrow, meta=meta_hist) + ).map_partitions(cast_and_arrow, meta=meta_hist) results.append(ddf_scatt) # --- Timelag histogram diff --git a/src/sp2xr/schema.py b/src/sp2xr/schema.py index d199c05..ff576a8 100644 --- a/src/sp2xr/schema.py +++ b/src/sp2xr/schema.py @@ -176,7 +176,7 @@ def enforce_schema( return ddf -def _cast_and_arrow(pdf: pd.DataFrame) -> pd.DataFrame: +def cast_and_arrow(pdf: pd.DataFrame) -> pd.DataFrame: """ Cast every column in *this* partition to the canonical dtype (or DEFAULT_FLOAT), then switch the frame to pyarrow-backed dtypes.