diff --git a/weather_mv/loader_pipeline/sinks.py b/weather_mv/loader_pipeline/sinks.py index 0f226fd1..cd882b48 100644 --- a/weather_mv/loader_pipeline/sinks.py +++ b/weather_mv/loader_pipeline/sinks.py @@ -29,7 +29,6 @@ import apache_beam as beam import cfgrib import numpy as np -import pandas as pd import rasterio import rioxarray import xarray as xr @@ -409,19 +408,12 @@ def open_dataset(uri: str, """Open the dataset at 'uri' and return a xarray.Dataset.""" try: if is_zarr: - start_date = end_date = None - if 'start_date' in open_dataset_kwargs: - start_date = open_dataset_kwargs['start_date'] - if open_dataset_kwargs is not None: - open_dataset_kwargs.pop('start_date', None) - if 'end_date' in open_dataset_kwargs: - end_date = open_dataset_kwargs['end_date'] - if open_dataset_kwargs is not None: - open_dataset_kwargs.pop('end_date', None) + if open_dataset_kwargs is not None: + start_date = open_dataset_kwargs.pop('start_date', None) + end_date = open_dataset_kwargs.pop('end_date', None) ds: xr.Dataset = _add_is_normalized_attr(xr.open_dataset(uri, engine='zarr', **open_dataset_kwargs), False) if start_date is not None and end_date is not None: - time_range = pd.date_range(start_date, end_date, periods=25, inclusive='left') - ds = ds.sel(time=time_range) + ds = ds.sel(time=slice(start_date, end_date)) beam.metrics.Metrics.counter('Success', 'ReadNetcdfData').inc() yield ds ds.close()