Fix #1308: Scope dask.config.set() to specific operations instead of module import

Move the `split_native_chunks` config override from module-level in ag3.py
to context managers within the specific methods that require it. This
prevents importing malariagen_data from silently modifying global dask
configuration, which could degrade performance for unrelated dask
workloads in the same Python session.

Affected operations:
- util._da_compress(): wraps da.compress() call
- snp_data.snp_genotypes(): wraps da.compress() and da.take() calls
- snp_data._locate_site_class(): wraps da.take() call

Author: khushthecoder

malariagen_data/ag3.py

@@ -1,14 +1,10 @@
 import sys
 
-import dask
 import pandas as pd  # type: ignore
 import plotly.express as px  # type: ignore
 import malariagen_data
 from .anopheles import AnophelesDataResource
 
-# silence dask performance warnings
-dask.config.set(**{"array.slicing.split_native_chunks": False})  # type: ignore
-
 MAJOR_VERSION_NUMBER = 3
 MAJOR_VERSION_PATH = "v3"
 CONFIG_PATH = "v3-config.json"

malariagen_data/anoph/snp_data.py

@@ -5,6 +5,7 @@
 
 import allel  # type: ignore
 import bokeh
+import dask
 import dask.array as da
 import numpy as np
 import pandas as pd
@@ -72,13 +73,13 @@ def __init__(
         self._cache_snp_genotypes: Dict[
             base_params.sample_set, zarr.hierarchy.Group
         ] = dict()
-        self._cache_site_filters: Dict[
-            base_params.site_mask, zarr.hierarchy.Group
-        ] = dict()
+        self._cache_site_filters: Dict[base_params.site_mask, zarr.hierarchy.Group] = (
+            dict()
+        )
         self._cache_site_annotations: Optional[zarr.hierarchy.Group] = None
-        self._cache_locate_site_class: OrderedDict[
-            Tuple[Any, ...], np.ndarray
-        ] = OrderedDict()
+        self._cache_locate_site_class: OrderedDict[Tuple[Any, ...], np.ndarray] = (
+            OrderedDict()
+        )
 
         # Create the SNP-calls cache as a per-instance lru_cache wrapping the
         # bound method.  Storing it on the instance (rather than using a
@@ -253,8 +254,7 @@ def _site_filters_for_contig(
         else:
             if contig not in self.contigs:
                 raise ValueError(
-                    f"Contig {contig!r} not found. "
-                    f"Available contigs: {self.contigs}"
+                    f"Contig {contig!r} not found. Available contigs: {self.contigs}"
                 )
             root = self.open_site_filters(mask=mask)
             z = root[f"{contig}/variants/{field}"]
@@ -359,8 +359,7 @@ def _snp_sites_for_contig(
         else:
             if contig not in self.contigs:
                 raise ValueError(
-                    f"Contig {contig!r} not found. "
-                    f"Available contigs: {self.contigs}"
+                    f"Contig {contig!r} not found. Available contigs: {self.contigs}"
                 )
             root = self.open_snp_sites()
             z = root[f"{contig}/variants/{field}"]
@@ -488,8 +487,7 @@ def _snp_genotypes_for_contig(
         else:
             if contig not in self.contigs:
                 raise ValueError(
-                    f"Contig {contig!r} not found. "
-                    f"Available contigs: {self.contigs}"
+                    f"Contig {contig!r} not found. Available contigs: {self.contigs}"
                 )
             root = self.open_snp_genotypes(sample_set=sample_set)
             z = root[f"{contig}/calldata/{field}"]
@@ -612,12 +610,14 @@ def snp_genotypes(
             )
 
             # Filter the Dask array using the boolean array.
-            d = da.compress(loc_samples, d, axis=1)
+            with dask.config.set(**{"array.slicing.split_native_chunks": False}):
+                d = da.compress(loc_samples, d, axis=1)
 
         # Apply the sample_indices, if there are any.
         # Note: this might need to apply to the result of an internal sample_query, e.g. `is_surveillance == True`.
         if sample_indices is not None:
-            d = da.take(d, sample_indices, axis=1)
+            with dask.config.set(**{"array.slicing.split_native_chunks": False}):
+                d = da.take(d, sample_indices, axis=1)
 
         return d
 
@@ -648,8 +648,7 @@ def _snp_variants_for_contig(
         else:
             if contig not in self.contigs:
                 raise ValueError(
-                    f"Contig {contig!r} not found. "
-                    f"Available contigs: {self.contigs}"
+                    f"Contig {contig!r} not found. Available contigs: {self.contigs}"
                 )
             coords = dict()
             data_vars = dict()
@@ -1021,7 +1020,8 @@ def _locate_site_class(
                 chunks=chunks,
             )
             idx = (pos - 1).compute()
-            loc_ann = da.take(loc_ann, idx, axis=0)
+            with dask.config.set(**{"array.slicing.split_native_chunks": False}):
+                loc_ann = da.take(loc_ann, idx, axis=0)
 
             # Compute site selection.
             with self._dask_progress(desc=f"Locate {site_class} sites"):
@@ -1066,8 +1066,7 @@ def _snp_calls_for_contig(
         else:
             if contig not in self.contigs:
                 raise ValueError(
-                    f"Contig {contig!r} not found. "
-                    f"Available contigs: {self.contigs}"
+                    f"Contig {contig!r} not found. Available contigs: {self.contigs}"
                 )
 
             coords = dict()

malariagen_data/util.py

@@ -20,6 +20,7 @@
     colab = None
 
 import allel  # type: ignore
+import dask
 import dask.array as da
 from dask.utils import parse_bytes
 import numba  # type: ignore
@@ -407,8 +408,11 @@ def _da_compress(
     else:
         indexer = da.from_array(indexer, chunks=(axis_old_chunks,))
 
-    # Apply the indexing operation.
-    v = da.compress(indexer, data, axis=axis)
+    # Apply the indexing operation, suppressing the dask performance warning
+    # about split_native_chunks. This config is scoped here rather than at
+    # module level to avoid silently modifying global dask configuration.
+    with dask.config.set(**{"array.slicing.split_native_chunks": False}):
+        v = da.compress(indexer, data, axis=axis)
 
     # Need to compute chunks sizes in order to know dimension sizes;
     # would normally do v.compute_chunk_sizes() but that is slow for