Centralize the choice of a default buffer size for apply.

Author: LTLA

src/delayedarray/DelayedArray.py

@@ -822,7 +822,7 @@ def __getitem__(self, subset) -> Union["DelayedArray", ndarray]:
 
 
     # For python-level compute.
-    def sum(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[numpy.dtype] = None, buffer_size: int = 1e8) -> numpy.ndarray:
+    def sum(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[numpy.dtype] = None, buffer_size: Optional[int] = None) -> numpy.ndarray:
         """
         Take the sum of values across the ``DelayedArray``, possibly over a
         given axis or set of axes. If the seed has a ``sum()`` method, that
@@ -840,8 +840,9 @@ def sum(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optiona
                 :py:func:`~numpy.sum` for details.
 
             buffer_size:
-                Buffer size in bytes to use for block processing. Larger values
-                generally improve speed at the cost of memory.
+                Buffer size in bytes to use for block processing.
+                Larger values generally improve speed at the cost of memory.
+                If ``None``, defaults to the value returned by :py:func:`~delayedarray.default_buffer_size.default_buffer_size`. 
 
         Returns:
             A NumPy array containing the sums. If ``axis = None``, this will be
@@ -859,7 +860,7 @@ def sum(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optiona
             )
 
 
-    def mean(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[numpy.dtype] = None, buffer_size: int = 1e8) -> numpy.ndarray:
+    def mean(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[numpy.dtype] = None, buffer_size: Optional[int] = None) -> numpy.ndarray:
         """
         Take the mean of values across the ``DelayedArray``, possibly over a
         given axis or set of axes. If the seed has a ``mean()`` method, that
@@ -877,8 +878,9 @@ def mean(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Option
                 :py:func:`~numpy.mean` for details.
 
             buffer_size:
-                Buffer size in bytes to use for block processing. Larger values
-                generally improve speed at the cost of memory.
+                Buffer size in bytes to use for block processing.
+                Larger values generally improve speed at the cost of memory.
+                If ``None``, defaults to the value returned by :py:func:`~delayedarray.default_buffer_size.default_buffer_size`. 
 
         Returns:
             A NumPy array containing the means. If ``axis = None``, this will
@@ -896,7 +898,7 @@ def mean(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Option
             )
 
 
-    def var(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[numpy.dtype] = None, ddof: int = 0, buffer_size: int = 1e8) -> numpy.ndarray:
+    def var(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[numpy.dtype] = None, ddof: int = 0, buffer_size: Optional[int] = None) -> numpy.ndarray:
         """
         Take the variances of values across the ``DelayedArray``, possibly over
         a given axis or set of axes. If the seed has a ``var()`` method, that
@@ -918,8 +920,9 @@ def var(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optiona
                 Typically set to 1 to obtain the sample variance.
 
             buffer_size:
-                Buffer size in bytes to use for block processing. Larger values
-                generally improve speed at the cost of memory.
+                Buffer size in bytes to use for block processing.
+                Larger values generally improve speed at the cost of memory.
+                If ``None``, defaults to the value returned by :py:func:`~delayedarray.default_buffer_size.default_buffer_size`. 
 
         Returns:
             A NumPy array containing the variances. If ``axis = None``,
@@ -937,7 +940,7 @@ def var(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optiona
                 masked=is_masked(self),
             )
 
-    def any(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[numpy.dtype] = None, buffer_size: int = 1e8) -> numpy.ndarray:
+    def any(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[numpy.dtype] = None, buffer_size: Optional[int] = None) -> numpy.ndarray:
         """Test whether any array element along a given axis evaluates to True.
 
         Compute this test across the ``DelayedArray``, possibly over a
@@ -956,8 +959,9 @@ def any(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optiona
                 :py:func:`~numpy.any` for details.
 
             buffer_size:
-                Buffer size in bytes to use for block processing. Larger values
-                generally improve speed at the cost of memory.
+                Buffer size in bytes to use for block processing.
+                Larger values generally improve speed at the cost of memory.
+                If ``None``, defaults to the value returned by :py:func:`~delayedarray.default_buffer_size.default_buffer_size`. 
 
         Returns:
             A NumPy array containing the boolean values. If ``axis = None``, this will
@@ -974,7 +978,7 @@ def any(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optiona
                 masked=is_masked(self),
             )
 
-    def all(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[numpy.dtype] = None, buffer_size: int = 1e8) -> numpy.ndarray:
+    def all(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[numpy.dtype] = None, buffer_size: Optional[int] = None) -> numpy.ndarray:
         """Test whether all array elements along a given axis evaluate to True.
 
         Compute this test across the ``DelayedArray``, possibly over a
@@ -993,8 +997,9 @@ def all(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optiona
                 :py:func:`~numpy.all` for details.
 
             buffer_size:
-                Buffer size in bytes to use for block processing. Larger values
-                generally improve speed at the cost of memory.
+                Buffer size in bytes to use for block processing.
+                Larger values generally improve speed at the cost of memory.
+                If ``None``, defaults to the value returned by :py:func:`~delayedarray.default_buffer_size.default_buffer_size`. 
 
         Returns:
             A NumPy array containing the boolean values. If ``axis = None``, this will
@@ -1115,7 +1120,7 @@ def _reduce_SparseNdarray(x: SparseNdarray, multipliers: List[int], axes: List[i
     return        
 
 
-def _reduce(x: DelayedArray, axes: List[int], operation: Callable, buffer_size: int):
+def _reduce(x: DelayedArray, axes: List[int], operation: Callable, buffer_size: Optional[int]):
     multipliers = _create_offset_multipliers(x.shape, axes)
     if is_sparse(x):
         apply_over_blocks(

src/delayedarray/__init__.py

@@ -37,6 +37,7 @@
 from .RegularTicks import RegularTicks
 from .apply_over_dimension import apply_over_dimension
 from .apply_over_blocks import apply_over_blocks
+from .default_buffer_size import default_buffer_size
 
 from .create_dask_array import create_dask_array
 from .is_sparse import is_sparse

src/delayedarray/apply_over_blocks.py

@@ -6,13 +6,14 @@
 from .is_sparse import is_sparse
 from .extract_dense_array import extract_dense_array
 from .extract_sparse_array import extract_sparse_array
+from .default_buffer_size import default_buffer_size
 
 __author__ = "ltla"
 __copyright__ = "ltla"
 __license__ = "MIT"
 
 
-def apply_over_blocks(x, fun: Callable, allow_sparse: bool = False, grid: Optional[AbstractGrid] = None, buffer_size: int = 1e8) -> list:
+def apply_over_blocks(x, fun: Callable, allow_sparse: bool = False, grid: Optional[AbstractGrid] = None, buffer_size: Optional[int] = None) -> list:
     """
     Iterate over an array by blocks. We apply a user-provided function and
     collect the results before proceeding to the next block.
@@ -38,8 +39,9 @@ def apply_over_blocks(x, fun: Callable, allow_sparse: bool = False, grid: Option
             of :py:func:`~delayedarray.chunk_grid.chunk_grid` on ``x``.
 
         buffer_size: 
-            Buffer_size in bytes, to hold a single block per iteration. Larger
-            values generally improve speed at the cost of memory.
+            Buffer_size in bytes, to hold a single block per iteration.
+            Larger values generally improve speed at the cost of memory.
+            If ``None``, defaults to the value returned by :py:func:`~delayedarray.default_buffer_size.default_buffer_size`. 
 
     Returns:
         List containing the output of ``fun`` on each block.
@@ -54,6 +56,9 @@ def apply_over_blocks(x, fun: Callable, allow_sparse: bool = False, grid: Option
 
     dims = (*range(len(x.shape)),)
     collected = []
+
+    if buffer_size is None:
+        buffer_size = default_buffer_size()
     buffer_elements = buffer_size // x.dtype.itemsize
 
     for job in grid.iterate(dims, buffer_elements = buffer_elements):

src/delayedarray/apply_over_dimension.py

@@ -6,13 +6,14 @@
 from .is_sparse import is_sparse
 from .extract_dense_array import extract_dense_array
 from .extract_sparse_array import extract_sparse_array
+from .default_buffer_size import default_buffer_size
 
 __author__ = "ltla"
 __copyright__ = "ltla"
 __license__ = "MIT"
 
 
-def apply_over_dimension(x, dimension: int, fun: Callable, allow_sparse: bool = False, grid: Optional[AbstractGrid] = None, buffer_size: int = 1e8) -> list:
+def apply_over_dimension(x, dimension: int, fun: Callable, allow_sparse: bool = False, grid: Optional[AbstractGrid] = None, buffer_size: Optional[int] = None) -> list:
     """
     Iterate over an array on a certain dimension. At each iteration, the block
     of observations consists of the full extent of all dimensions other than
@@ -42,8 +43,9 @@ def apply_over_dimension(x, dimension: int, fun: Callable, allow_sparse: bool =
             of :py:func:`~delayedarray.chunk_grid.chunk_grid` on ``x``.
 
         buffer_size: 
-            Buffer_size in bytes, to hold a single block per iteration. Larger
-            values generally improve speed at the cost of memory.
+            Buffer_size in bytes, to hold a single block per iteration.
+            Larger values generally improve speed at the cost of memory.
+            If ``None``, defaults to the value returned by :py:func:`~delayedarray.default_buffer_size.default_buffer_size`. 
 
     Returns:
         List containing the output of ``fun`` on each block.
@@ -57,9 +59,11 @@ def apply_over_dimension(x, dimension: int, fun: Callable, allow_sparse: bool =
     else:
         extractor = extract_dense_array
 
-    collected = []
+    if buffer_size is None:
+        buffer_size = default_buffer_size()
     buffer_elements = buffer_size // x.dtype.itemsize
 
+    collected = []
     for job in grid.iterate((dimension,), buffer_elements = buffer_elements):
         subsets = (*(range(s, e) for s, e in job),)
         output = fun(job[dimension], extractor(x, subsets))

src/delayedarray/default_buffer_size.py

@@ -0,0 +1,33 @@
+from typing import Optional
+
+__author__ = "ltla"
+__copyright__ = "ltla"
+__license__ = "MIT"
+
+
+old_buffer_size = 1e8
+
+
+def default_buffer_size(buffer_size: Optional[int] = None) -> int:
+    """
+    Get or set the default buffer size used by :py:func:`~delayedarray.apply_over_blocks.apply_over_blocks`,
+    :py:func:`~delayedarray.apply_over_dimension.apply_over_dimension`, etc.
+
+    Args:
+        buffer_size:
+            Buffer size in bytes.
+            The buffer is typically used to load a block of an array in memory for further processing.
+            Alternatively ``None``.
+
+    Returns:
+        If ``buffer_size = None``, the current default buffer size is returned.
+
+        If ``buffer_size`` is an integer, the default buffer size is set to this value, and the previous buffer size is returned.
+    """
+    global old_buffer_size
+    if buffer_size is None:
+        return old_buffer_size
+
+    previous = old_buffer_size
+    old_buffer_size = buffer_size
+    return previous

tests/test_default_buffer_size.py

@@ -0,0 +1,12 @@
+import delayedarray
+
+
+def test_default_buffer_size():
+    assert delayedarray.default_buffer_size() == 1e8
+
+    old = delayedarray.default_buffer_size(500)
+    assert old == 1e8
+    assert delayedarray.default_buffer_size() == 500
+
+    delayedarray.default_buffer_size(old)
+    assert delayedarray.default_buffer_size() == old