Repository: dask/dask
Pull request #10579 (2023-10-19 17:13:48)

Speed up ``normalize_chunks`` for common case
In certain xarray zarr/kerchunk reads with many variables, a considerable amount of time is spent in dask.array's normalize chunks (see snakeviz output)

<img width="1059" alt="Screenshot 2023-10-19 at 10 51 33" src="https://github.com/dask/dask/assets/6042212/da46bf5a-a6ad-48bc-804e-5820ac4df706">

For the case that a zarr-like chunksize was given, or a full chunks spec was already given, this is all unnecessary work. This PR skips it for such cases. The same part of the profile following the change looks like:

<img width="1059" alt="Screenshot 2023-10-19 at 10 55 21" src="https://github.com/dask/dask/assets/6042212/55409af8-609c-4c28-8610-ad45cf17915f">



martindurant: bump @dask/dask-dev 
jrbourbeau: Nice, thanks @martindurant -- taking a look now 
martindurant: @jrbourbeau , thanks for the review

I implemented your suggestions and answered questions.
There may be a better way to avoid the time spent in this one function, but I didn't want to make changes further up the stack where they might have more consequences. I don't really understand why we have nans here at all, or why they might be something other than np.nan. For the case where we aren't using a chunksize, I don't really know why running this function is necessary at all and wouldn't have thought about it without the profiling.
martindurant: Removing my stuff and simply using
```
    return tuple(tuple(ch) for ch in chunks)
```
on the last line would be much better, but there are two tests specifically passing float("nan") to check the current behaviour.

Patch:
diff --git a/.github/workflows/additional.yml b/.github/workflows/additional.yml
index d4043bc34..9d598038f 100644
--- a/.github/workflows/additional.yml
+++ b/.github/workflows/additional.yml
@@ -22,8 +22,8 @@ jobs:
           miniforge-version: latest
           use-mamba: true
           channel-priority: strict
-          python-version: "3.10"
-          environment-file: continuous_integration/environment-3.10.yaml
+          python-version: "3.11"
+          environment-file: continuous_integration/environment-3.11.yaml
           activate-environment: test-environment
           auto-activate-base: false
 
diff --git a/.github/workflows/upstream.yml b/.github/workflows/upstream.yml
index af83c7a73..e9c129670 100644
--- a/.github/workflows/upstream.yml
+++ b/.github/workflows/upstream.yml
@@ -56,8 +56,8 @@ jobs:
           miniforge-version: latest
           use-mamba: true
           channel-priority: strict
-          python-version: "3.10"
-          environment-file: continuous_integration/environment-3.10.yaml
+          python-version: "3.11"
+          environment-file: continuous_integration/environment-3.11.yaml
           activate-environment: test-environment
           auto-activate-base: false
 
diff --git a/continuous_integration/scripts/install.sh b/continuous_integration/scripts/install.sh
index a5453db02..84a284630 100644
--- a/continuous_integration/scripts/install.sh
+++ b/continuous_integration/scripts/install.sh
@@ -54,7 +54,7 @@ if [[ ${UPSTREAM_DEV} ]]; then
 
     # FIXME https://github.com/mamba-org/mamba/issues/412
     # mamba uninstall --force ...
-    conda uninstall --force crick tiledb tiledb-py
+    # conda uninstall --force crick tiledb tiledb-py
 
 
 fi
diff --git a/dask/array/core.py b/dask/array/core.py
index e03b9cd9a..35a6a3bc9 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -3097,7 +3097,10 @@ def normalize_chunks(chunks, shape=None, limit=None, dtype=None, previous_chunks
     if shape is not None:
         chunks = tuple(c if c not in {None, -1} else s for c, s in zip(chunks, shape))
 
+    allints = None
     if chunks and shape is not None:
+        # allints: did we start with chunks as a simple tuple of ints?
+        allints = all(isinstance(c, int) for c in chunks)
         chunks = sum(
             (
                 blockdims_from_blockshape((s,), (c,))
@@ -3128,7 +3131,9 @@ def normalize_chunks(chunks, shape=None, limit=None, dtype=None, previous_chunks
                 "Chunks do not add up to shape. "
                 "Got chunks=%s, shape=%s" % (chunks, shape)
             )
-
+    if allints or isinstance(sum(sum(_) for _ in chunks), int):
+        # Fastpath for when we already know chunks contains only integers
+        return tuple(tuple(ch) for ch in chunks)
     return tuple(
         tuple(int(x) if not math.isnan(x) else np.nan for x in c) for c in chunks
     )
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 88ebea8f3..360cc592d 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -849,9 +849,9 @@ Dask Name: {name}, {layers}"""
         Dask DataFrame Structure:
                          name     id        x        y
         npartitions=1
-        2021-01-01     object  int64  float64  float64
+        2021-01-01     string  int64  float64  float64
         2021-01-02        ...    ...      ...      ...
-        Dask Name: get-partition, 2 graph layers
+        Dask Name: get-partition, 3 graph layers
 
         See Also
         --------
diff --git a/dask/dataframe/io/sql.py b/dask/dataframe/io/sql.py
index b1f59a4e1..de4b33290 100644
--- a/dask/dataframe/io/sql.py
+++ b/dask/dataframe/io/sql.py
@@ -539,12 +539,13 @@ def to_sql(
     Dask Name: from_pandas, 2 tasks
 
     >>> from dask.utils import tmpfile
-    >>> from sqlalchemy import create_engine
+    >>> from sqlalchemy import create_engine, text
     >>> with tmpfile() as f:
     ...     db = 'sqlite:///%s' %f
     ...     ddf.to_sql('test', db)
     ...     engine = create_engine(db, echo=False)
-    ...     result = engine.execute("SELECT * FROM test").fetchall()
+    ...     with engine.connect() as conn:
+    ...         result = conn.execute(text("SELECT * FROM test")).fetchall()
     >>> result
     [(0, 0, '00'), (1, 1, '11'), (2, 2, '22'), (3, 3, '33')]
     """
diff --git a/dask/dataframe/multi.py b/dask/dataframe/multi.py
index ade477ecd..82d43f315 100644
--- a/dask/dataframe/multi.py
+++ b/dask/dataframe/multi.py
@@ -1255,7 +1255,7 @@ def concat(
     ...     dd.from_pandas(pd.Series(['a', 'b'], dtype='category'), 1),
     ...     dd.from_pandas(pd.Series(['a', 'c'], dtype='category'), 1),
     ... ], interleave_partitions=True).dtype
-    CategoricalDtype(categories=['a', 'b', 'c'], ordered=False)
+    CategoricalDtype(categories=['a', 'b', 'c'], ordered=False, categories_dtype=object)
     """
 
     if not isinstance(dfs, list):
diff --git a/dask/dataframe/reshape.py b/dask/dataframe/reshape.py
index cb3868a54..58e067570 100644
--- a/dask/dataframe/reshape.py
+++ b/dask/dataframe/reshape.py
@@ -95,16 +95,16 @@ def get_dummies(
     Dask DataFrame Structure:
                        a      b      c
     npartitions=2
-    0              uint8  uint8  uint8
+    0              bool  bool  bool
     2                ...    ...    ...
     3                ...    ...    ...
     Dask Name: get_dummies, 2 graph layers
     >>> dd.get_dummies(s).compute()  # doctest: +ELLIPSIS
-       a  b  c
-    0  1  0  0
-    1  0  1  0
-    2  0  0  1
-    3  1  0  0
+           a      b      c
+    0   True  False  False
+    1  False   True  False
+    2  False  False   True
+    3   True  False  False
 
     See Also
     --------
diff --git a/dask/sizeof.py b/dask/sizeof.py
index bbd80b3a1..5a98d9465 100644
--- a/dask/sizeof.py
+++ b/dask/sizeof.py
@@ -65,11 +65,17 @@ class SimpleSizeof:
 
     Examples
     --------
+    >>> def _get_gc_overhead():
+    ...     class _CustomObject:
+    ...         def __sizeof__(self):
+    ...             return 0
+    ...
+    ...     return sys.getsizeof(_CustomObject())
 
     >>> class TheAnswer(SimpleSizeof):
-    ...         def __sizeof__(self):
-    ...             # Sizeof always add overhead of an object for GC
-    ...             return 42 - sizeof(object())
+    ...     def __sizeof__(self):
+    ...         # Sizeof always add overhead of an object for GC
+    ...         return 42 - _get_gc_overhead()
 
     >>> sizeof(TheAnswer())
     42
diff --git a/docs/source/develop.rst b/docs/source/develop.rst
index b9085c74d..a680895ff 100644
--- a/docs/source/develop.rst
+++ b/docs/source/develop.rst
@@ -108,7 +108,7 @@ pip or conda_
 
 ``conda``::
 
-  conda env create -n dask-dev -f continuous_integration/environment-3.10.yaml
+  conda env create -n dask-dev -f continuous_integration/environment-3.11.yaml
   conda activate dask-dev
   python -m pip install --no-deps -e .
 
@@ -254,7 +254,7 @@ after the line.
 
 .. _numpydoc: https://numpydoc.readthedocs.io/en/latest/format.html#docstring-standard
 
-Docstrings are tested under Python 3.10 on GitHub Actions. You can test
+Docstrings are tested under Python 3.11 on GitHub Actions. You can test
 docstrings with pytest as follows::
 
    py.test dask --doctest-modules


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index b5d57147d..ddd6b4ad8 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -2443,7 +2443,7 @@ def test_repartition_noop(type_ctor):
         (
             "M",
             "MS",
-            (UserWarning, "'M' will be deprecated, please use 'ME' instead")
+            (FutureWarning, "'M' will be deprecated, please use 'ME' instead")
             if PANDAS_GE_220
             else None,
         ),
@@ -2452,7 +2452,7 @@ def test_repartition_noop(type_ctor):
         (
             "2M",
             "2MS",
-            (UserWarning, "'M' will be deprecated, please use 'ME' instead")
+            (FutureWarning, "'M' will be deprecated, please use 'ME' instead")
             if PANDAS_GE_220
             else None,
         ),
diff --git a/dask/tests/test_docs.py b/dask/tests/test_docs.py
index bf5e0c5ec..4869dca24 100644
--- a/dask/tests/test_docs.py
+++ b/dask/tests/test_docs.py
@@ -15,7 +15,7 @@ def test_development_guidelines_matches_ci():
     development_doc_file = root_dir / "docs" / "source" / "develop.rst"
     additional_ci_file = root_dir / ".github" / "workflows" / "additional.yml"
     upstream_ci_file = root_dir / ".github" / "workflows" / "upstream.yml"
-    latest_env = "environment-3.10.yaml"
+    latest_env = "environment-3.11.yaml"
 
     for filename in [development_doc_file, additional_ci_file, upstream_ci_file]:
         with open(filename, encoding="utf8") as f: