Estimate do not need to open the zarr with dask, use chunks=None for …

data_access_service/core/estimation_index.py

@@ -206,6 +206,15 @@ def sidecar_extent_provider(api, output_format: str) -> ExtentProvider:
     def provider(
         uuid: str, key: str
     ) -> Tuple[Optional[pd.Timestamp], Optional[pd.Timestamp]]:
+        if key.endswith(".zarr"):
+            # A zarr extent must not go through get_temporal_extent: that
+            # caches a chunks="auto" open and builds the dask graph the
+            # estimate is trying to avoid.
+            from data_access_service.core.size_estimation import (
+                zarr_extent_for_estimate,
+            )
+
+            return zarr_extent_for_estimate(api, uuid, key)
         if not key.endswith(".parquet") or output_format != OUTPUT_FORMAT_CSV:
             return api.get_temporal_extent(uuid, key)
         meta = require_sidecar(api, uuid, key)

data_access_service/core/size_estimation.py

@@ -48,6 +48,99 @@
 log = logging.getLogger(__name__)
 
 
+def zarr_extent_for_estimate(api, uuid: str, key: str):
+    """Min/max time of a zarr key, without caching a dask-backed open.
+
+    Same day-alignment and future-date clamp as ``API.get_temporal_extent``.
+    The store is closed before this returns.
+    """
+    ds = api._instance.get_dataset(key, chunks=None)
+    try:
+        try:
+            start_date, end_date = ds.get_temporal_extent()
+        except ValueError:
+            return None, None
+        return _clamp_extent(start_date, end_date)
+    finally:
+        _close_zarr(ds)
+
+
+def _clamp_extent(start_date, end_date):
+    """Match ``API.get_temporal_extent`` so an estimate trims the same days."""
+    if start_date is not None:
+        start_date = start_date.replace(
+            hour=0, minute=0, second=0, microsecond=0, nanosecond=0
+        )
+    if end_date is not None:
+        if end_date.tzinfo is None:
+            now_compare = pd.Timestamp.now()
+        else:
+            now_compare = pd.Timestamp.now(tz="UTC")
+            end_date = end_date.tz_convert("UTC")
+        end_date = end_date.replace(
+            hour=23, minute=59, second=59, microsecond=999999, nanosecond=999
+        )
+        if end_date > now_compare:
+            end_date = now_compare
+    return start_date, end_date
+
+
+def _close_zarr(ds) -> None:
+    store = getattr(ds, "zarr_store", None)
+    if store is not None:
+        store.close()
+
+
+def _estimate_zarr_key(
+    api,
+    uuid: str,
+    key: str,
+    resolved_subset_request: ResolvedSubsetRequest,
+    output_format: str,
+    cancellation: Optional[Cancellation],
+) -> Optional[dict]:
+    """Size one zarr key from a chunks=None open that is not cached."""
+    ds = api._instance.get_dataset(key, chunks=None)
+    try:
+        check_datasource_supports_format(api, ds, uuid, key, output_format)
+        if not resolved_subset_request.has_data:
+            return _empty_estimate(uuid, key, output_format)
+        raise_if_client_gone(cancellation)
+
+        def _extent(_uuid: str, _key: str):
+            try:
+                start_date, end_date = ds.get_temporal_extent()
+            except ValueError:
+                return None, None
+            return _clamp_extent(start_date, end_date)
+
+        date_start, date_end = trim_date_range_for_keys(
+            api,
+            uuid,
+            [key],
+            resolved_subset_request.start_date,
+            resolved_subset_request.end_date,
+            extent_provider=_extent,
+        )
+        if date_start is None or date_end is None:
+            return _empty_estimate(uuid, key, output_format)
+        return _estimate_zarr_size(
+            api,
+            ds.zarr_store,
+            uuid,
+            key,
+            date_start,
+            date_end,
+            resolved_subset_request.bboxes,
+            resolved_subset_request.columns,
+            output_format,
+            resolved_subset_request.geometry,
+            cancellation,
+        )
+    finally:
+        _close_zarr(ds)
+
+
 def _bytes_to_mb(num_bytes: int) -> float:
     """Bytes -> decimal MB (1 MB = 1,000,000 bytes), the way browsers report
     file sizes. Rounded to 1 decimal place."""
@@ -78,6 +171,14 @@ def estimate_single_key_size(
     :raises ClientGoneError: if the client disconnected
     """
     uuid = resolved_subset_request.uuid
+    if key.endswith(".zarr"):
+        # Do not use get_datasource. That caches a chunks="auto" open, and
+        # the dask graph for a fine-grained satellite store is multiple GB.
+        # Downloads keep that open; only this estimate uses chunks=None.
+        return _estimate_zarr_key(
+            api, uuid, key, resolved_subset_request, output_format, cancellation
+        )
+
     ds = api.get_datasource(uuid, key)
     if ds is None:
         return None

data_access_service/models/co_datasource/abstract_data_src.py

@@ -29,12 +29,22 @@ def get_dataset_names(self) -> frozenset[str]:
         return frozenset(self.get_metadata_catalog())
 
     @abstractmethod
-    def get_dataset(self, dataset_name_with_ext: str) -> DataSource:
+    def get_dataset(
+        self, dataset_name_with_ext: str, chunks: str | None = "auto"
+    ) -> DataSource:
+        """Open one dataset.
+
+        ``chunks`` is forwarded to the zarr open. ``"auto"`` is the library
+        default and builds a dask graph. ``None`` leaves the store lazy with
+        no graph, which is what the size estimate uses. Parquet ignores it.
+        """
         if dataset_name_with_ext not in self.get_dataset_names():
             raise DatasetNotFoundError(
                 dataset_name=dataset_name_with_ext, data_source_name=self.get_name()
             )
 
+        if dataset_name_with_ext.endswith(".zarr"):
+            return self.get_data_src().get_dataset(dataset_name_with_ext, chunks=chunks)
         return self.get_data_src().get_dataset(dataset_name_with_ext)
 
     @abstractmethod

data_access_service/models/co_datasource/aodn_data_src.py

@@ -30,8 +30,12 @@ def get_dataset_names(self) -> frozenset[str]:
             self.__dataset_names = frozenset(self.__data_src.list_datasets())
         return self.__dataset_names
 
-    def get_dataset(self, dataset_name_with_ext: str) -> DataSource:
-        return super().get_dataset(dataset_name_with_ext=dataset_name_with_ext)
+    def get_dataset(
+        self, dataset_name_with_ext: str, chunks: str | None = "auto"
+    ) -> DataSource:
+        return super().get_dataset(
+            dataset_name_with_ext=dataset_name_with_ext, chunks=chunks
+        )
 
     def get_data_src(self) -> GetAodn:
         return self.__data_src

data_access_service/models/co_datasource/co_data_registory.py

@@ -109,19 +109,29 @@ def get_metadata(self) -> Metadata:
         before_sleep=log_retry_attempt("get_dataset", log),
         reraise=True,
     )
-    def get_dataset(self, dataset_name_with_ext: str) -> DataSource | None:
-        with self._datasets_lock:
-            cached = self._datasets.get(dataset_name_with_ext)
-            if cached is not None:
-                log.info("Reusing cached %s dataset", dataset_name_with_ext)
-                return cached
+    def get_dataset(
+        self, dataset_name_with_ext: str, chunks: str | None = "auto"
+    ) -> DataSource | None:
+        # Only the default dask open is cached. A chunks=None open is the
+        # size estimate: it must not replace the chunks="auto" store that
+        # downloads reuse, and it must not be kept after the estimate.
+        use_cache = chunks == "auto"
+        if use_cache:
+            with self._datasets_lock:
+                cached = self._datasets.get(dataset_name_with_ext)
+                if cached is not None:
+                    log.info("Reusing cached %s dataset", dataset_name_with_ext)
+                    return cached
 
         for data_src in self.data_source_list:
             try:
                 log.info(
                     f"Getting {dataset_name_with_ext} dataset from {data_src.get_name()}..."
                 )
-                dataset = data_src.get_dataset(dataset_name_with_ext)
+                if chunks == "auto":
+                    dataset = data_src.get_dataset(dataset_name_with_ext)
+                else:
+                    dataset = data_src.get_dataset(dataset_name_with_ext, chunks=chunks)
             except DatasetNotFoundError:
                 # log the exception and continue to try the next data source
                 log.info(
@@ -131,6 +141,8 @@ def get_dataset(self, dataset_name_with_ext: str) -> DataSource | None:
             # Load happened outside the lock so a 4-minute S3 listing does
             # not block other datasets. setdefault keeps the first stored
             # instance if two threads both missed.
+            if not use_cache:
+                return dataset
             with self._datasets_lock:
                 return self._datasets.setdefault(dataset_name_with_ext, dataset)
 

data_access_service/models/co_datasource/csiro/csiro_data_src.py

@@ -401,8 +401,12 @@ def get_metadata(self) -> Metadata:
     def get_metadata_catalog(self) -> dict:
         return self.__metadata_catalog
 
-    def get_dataset(self, dataset_name_with_ext: str) -> DataSource:
-        return super().get_dataset(dataset_name_with_ext=dataset_name_with_ext)
+    def get_dataset(
+        self, dataset_name_with_ext: str, chunks: str | None = "auto"
+    ) -> DataSource:
+        return super().get_dataset(
+            dataset_name_with_ext=dataset_name_with_ext, chunks=chunks
+        )
 
     def get_name(self) -> str:
         return self.name

tests/core/test_api_estimate_size.py

@@ -121,6 +121,38 @@ def _api_with_zarr(dataset: xr.Dataset) -> tuple[API, MagicMock]:
     return api, mock_ds
 
 
+def test_zarr_suffix_does_not_use_the_cached_dask_open(monkeypatch):
+    """A real .zarr key is sized from chunks=None and that open is closed."""
+    from data_access_service.core import size_estimation
+
+    dataset = _make_dataset()
+    opened = MagicMock(spec=ZarrDataSource)
+    opened.zarr_store = dataset
+    opened.get_temporal_extent.return_value = (
+        pd.Timestamp("2020-01-01"),
+        pd.Timestamp("2020-01-05"),
+    )
+    closed = []
+    monkeypatch.setattr(size_estimation, "_close_zarr", lambda ds: closed.append(ds))
+    api = MagicMock()
+    api.resolve_dim_names.return_value = ("LATITUDE", "LONGITUDE", "TIME")
+    api.map_column_names = MagicMock(side_effect=_identity_map)
+    api._instance.get_dataset.return_value = opened
+    api.get_datasource.side_effect = AssertionError("cached chunks=auto open")
+
+    result = estimate_single_key_size(
+        api,
+        "satellite.zarr",
+        _resolved(keys=["satellite.zarr"]),
+        output_format="netcdf",
+    )
+
+    assert result["estimated_uncompressed_bytes"] == dataset.nbytes
+    api.get_datasource.assert_not_called()
+    api._instance.get_dataset.assert_called_with("satellite.zarr", chunks=None)
+    assert closed == [opened]
+
+
 def test_zarr_estimate_basic():
     dataset = _make_dataset()
     api, _ = _api_with_zarr(dataset)

tests/models/co_datasource/test_co_data_source.py

@@ -433,6 +433,19 @@ def test_get_metadata_raises_on_conflicting_dataset_names(self):
         with pytest.raises(Exception, match="Conflicted dataset names"):
             registry.get_metadata()
 
+    def test_chunks_none_is_not_cached(self):
+        registry, mock_aodn_src, _ = _make_registry(
+            aodn_catalog={"satellite.zarr": {}},
+        )
+        opened = MagicMock()
+        mock_aodn_src.get_dataset.return_value = opened
+
+        assert registry.get_dataset("satellite.zarr", chunks=None) is opened
+        assert registry.get_dataset("satellite.zarr", chunks=None) is opened
+        assert mock_aodn_src.get_dataset.call_count == 2
+        mock_aodn_src.get_dataset.assert_called_with("satellite.zarr", chunks=None)
+        assert "satellite.zarr" not in registry._datasets
+
     def test_get_dataset_returns_from_first_matching_source(self):
         registry, mock_aodn_src, _ = _make_registry()
         expected = MagicMock()