Repository: intake/intake
Pull request #720 (2023-03-08 02:44:12)

Divorce dask in npy source
Doing a source.read() should not require dask, and it would be inefficient for this anyway.


Patch:
diff --git a/README.md b/README.md
index 87b0798..5b0ad1b 100644
--- a/README.md
+++ b/README.md
@@ -17,6 +17,8 @@ Intake helps you:
 
 Documentation is available at [Read the Docs](http://intake.readthedocs.io/en/latest).
 
+Status of intake and related packages is available at [Status Dashboard](https://intake.github.io/status)
+
 Weekly news about this repo and other related projects can be found on the
 [wiki](https://github.com/intake/intake/wiki/Community-News)
 
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 72f0ff0..278f2db 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -61,7 +61,7 @@ author = "Anaconda"
 # built documents.
 #
 # The short X.Y version.
-version = intake.__version__.split("+")[0]
+version = intake.__version__
 # The full version, including alpha/beta/rc tags.
 release = intake.__version__
 
@@ -110,9 +110,6 @@ if not on_rtd:
 #
 # html_theme_options = {}
 
-# Default title is "<project> v<release> documentation"
-html_title = "Intake documentation"
-
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
diff --git a/docs/source/transforms.rst b/docs/source/transforms.rst
index 53034a5..8aa0def 100644
--- a/docs/source/transforms.rst
+++ b/docs/source/transforms.rst
@@ -29,7 +29,7 @@ Consider the following *target* dataset, which loads some simple facts
 about US states from a CSV file. This  example is taken from the Intake
 test suite.
 
-.. code-block:: yaml
+.. code-block::yaml
 
     sources:
       input_data:
@@ -50,7 +50,7 @@ data source class, and the parameters passed are specific to the transform type.
 Note that the driver is referred to by it's fully-qualified name in the
 Intake package.
 
-.. code-block:: yaml
+.. code-block::yaml
 
       derive_cols:
         driver: intake.source.derived.Columns
diff --git a/intake/source/npy.py b/intake/source/npy.py
index 488a8b1..60b6ca8 100644
--- a/intake/source/npy.py
+++ b/intake/source/npy.py
@@ -37,8 +37,8 @@ class NPySource(DataSource):
         shape: tuple of int
             If known, the length of each axis
         chunks: int
-            Size of chunks within a file along biggest dimension - need not
-            be an exact factor of the length of that dimension
+            Size of chunks within a file along biggest dimension - must exactly divide
+            each file, or None for one partition per file.
         storage_options: dict
             Passed to file-system backend.
         """
@@ -46,53 +46,70 @@ class NPySource(DataSource):
         self.shape = shape
         self.dtype = dtype
         self.storage = storage_options or {}
-        self._chunks = chunks if chunks is not None else -1
-        self.chunks = None
+        self.files = None
+        self._chunks = chunks
         self._arrs = None
         self._arr = None
         super(NPySource, self).__init__(metadata=metadata)
 
     def _get_schema(self):
-        import dask.array as da
         from fsspec import open_files
 
-        if self._arr is None:
+        if self.files is None:
             path = self._get_cache(self.path)[0]
 
-            files = open_files(path, "rb", compression=None, **self.storage)
-            if self.shape is None:
-                arr = NumpyAccess(files[0])
+            self.files = open_files(path, "rb", compression=None, **self.storage)
+            if self.dtype is None or self.shape is None:
+                arr = NumpyAccess(self.files[0])
                 self.shape = arr.shape
                 self.dtype = arr.dtype
-                arrs = [arr] + [NumpyAccess(f, self.shape, self.dtype, offset=arr.offset) for f in files[1:]]
-            else:
-                arrs = [NumpyAccess(f, self.shape, self.dtype) for f in files]
-            self.chunks = (self._chunks,) + (-1,) * (len(self.shape) - 1)
-            self._arrs = [da.from_array(arr, self.chunks) for arr in arrs]
-
-            if len(self._arrs) > 1:
-                self._arr = da.stack(self._arrs)
-            else:
-                self._arr = self._arrs[0]
-            self.chunks = self._arr.chunks
-        return Schema(dtype=str(self.dtype), shape=self.shape, extra_metadata=self.metadata, npartitions=self._arr.npartitions, chunks=self.chunks)
+            self.chunks = self._chunks or self.shape[0]
+        shape = list(self.shape)
+        parts_per_file = (shape[0] // self.chunks) + (shape[0] % self.chunks > 0)
+        if len(self.files) > 1:
+            shape = [len(self.files)] + shape
+            chunks = ((1,) * len(self.files),) + ((self.chunks,) * parts_per_file,) + tuple((s,) for s in self.shape[1:])
+        else:
+            shape = shape
+            chunks = ((self.chunks,) * parts_per_file,) + tuple((s,) for s in self.shape[1:])
+
+        print()
+        return Schema(dtype=str(self.dtype), shape=tuple(shape), extra_metadata=self.metadata, npartitions=parts_per_file * len(self.files), chunks=chunks)
 
     def _get_partition(self, i):
-        if isinstance(i, list):
-            i = tuple(i)
-        return self._arr.blocks[i].compute()
+        if isinstance(i, (list, tuple)):
+            i = i[0]
+        parts_per_file = (self.shape[0] // self.chunks) + (self.shape[0] % self.chunks > 0)
+        nfile = i // parts_per_file
+        arr = NumpyAccess(self.files[nfile])
+        infile = i % parts_per_file
+        return arr[self.chunks * infile : self.chunks * (infile + 1)]
 
     def read_partition(self, i):
-        self._get_schema()
         return self._get_partition(i)
 
     def to_dask(self):
+        import dask.array as da
+
         self._get_schema()
+        arrs = [NumpyAccess(f, self.shape, self.dtype) for f in self.files]
+
+        if len(arrs) > 1:
+            chunks = (self._chunks or -1,) + (-1,) * (len(self.shape) - 1)
+        else:
+            chunks = (self._chunks or -1,) + (-1,) * (len(self.shape) - 2)
+        self._arrs = [da.from_array(arr, chunks) for arr in arrs]
+
+        self._arr = da.stack(self._arrs)
         return self._arr
 
     def read(self):
+        import numpy as np
+
         self._get_schema()
-        return self._arr.compute()
+        with self.files as ff:
+            arrs = [np.load(f) for f in ff]
+        return np.stack(arrs)
 
     def _close(self):
         self._arrs = None
diff --git a/scripts/ci/environment-py39.yml b/scripts/ci/environment-py39.yml
index ed383f8..87511b3 100644
--- a/scripts/ci/environment-py39.yml
+++ b/scripts/ci/environment-py39.yml
@@ -21,7 +21,6 @@ dependencies:
   - panel
   - hvplot
   - bokeh
-  - dask
   - intake
   - pip
   - pytest


Test patch:
diff --git a/intake/source/tests/test_npy.py b/intake/source/tests/test_npy.py
index 8be3dd1..14062f3 100644
--- a/intake/source/tests/test_npy.py
+++ b/intake/source/tests/test_npy.py
@@ -50,18 +50,30 @@ def test_multi_file(tempdir, shape):
     np.save(fn1, data1)
     data = np.stack([data0, data1])
     fn = [fn0, fn1]
+
     s = NPySource(fn)
     out = s.read()
     assert (out == data).all()
+    assert (s.to_dask().compute() == data).all()
+
     s = NPySource(fn, chunks=1)
     out = s.read()
     assert (out == data).all()
+    assert (s.to_dask().compute() == data).all()
+
     s = NPySource(fn, shape=shape, dtype="int", chunks=1)
     out = s.read()
     assert (out == data).all()
+    da = s.to_dask()
+    assert (da.compute() == data).all()
+    schema = s.discover()
+    assert schema["shape"] == data.shape
+    assert schema["npartitions"] == da.npartitions
+
     s = NPySource(os.path.join(tempdir, "out*.npy"))
     out = s.read()
     assert (out == data).all()
+    assert (s.to_dask().compute() == data).all()
 
 
 def test_zarr_minimal():