Repository: dask/dask

Issue:
Cannot compare tz-naive and tz-aware timestamps on concat
**What happened**:

When concatenating two dask dataframes with indices dype=datetime64[ns, UTC], I get a `TypeError: Cannot compare tz-naive and tz-aware timestamps`. One of the the dask dataframe was created with `dd.from_pandas` and the other with `dd.read_parquet`

**What you expected to happen**:

An happy concatenation ;-)

**Minimal Complete Verifiable Example**:

```python
import dask.dataframe as dd
import pandas

filename = "outfile.parq"

# create a DF with tz=UTC and write to parquet file
df = pandas.DataFrame(
    columns=["A"], index=pandas.date_range("2018", "2019", freq="MS", tz="UTC"), data=1.0
)
ddf = dd.from_pandas(df, npartitions=1)
ddf.to_parquet(filename)

# read back this dask df into rddf
rddf = dd.read_parquet(filename)

# print indices
print(ddf.index)
print(rddf.index)


# attempt concatenation of dask df from parquet and new dask df
nddf = dd.concat([rddf, ddf])
# this raises the following traceback
# Traceback (most recent call last):
#   File ".../tst_parquet_tz.py", line 22, in <module>
#     nddf = dd.concat([rddf, ddf])
#   File "...\site-packages\dask\dataframe\multi.py", line 1108, in concat
#     if all(
#   File "...\site-packages\dask\dataframe\multi.py", line 1109, in <genexpr>
#     dfs[i].divisions[-1] < dfs[i + 1].divisions[0]
#   File "pandas\_libs\tslibs\timestamps.pyx", line 281, in pandas._libs.tslibs.timestamps._Timestamp.__richcmp__
#   File "pandas\_libs\tslibs\timestamps.pyx", line 295, in pandas._libs.tslibs.timestamps._Timestamp._assert_tzawareness_compat
# TypeError: Cannot compare tz-naive and tz-aware timestamps
```

**Anything else we need to know?**:

When printing both indices, we see that the dtype of both indices are `datetime64[ns, UTC]` yet the representation of the index coming from read-parquet shows tz-naive dates.
```
Dask Index Structure:
npartitions=1
2018-01-01 00:00:00+00:00    datetime64[ns, UTC]
2019-01-01 00:00:00+00:00                    ...
dtype: datetime64[ns, UTC]
Dask Name: from_pandas, 2 tasks
Dask Index Structure:
npartitions=1
2018-01-01    datetime64[ns, UTC]
2019-01-01                    ...
dtype: datetime64[ns, UTC]
Dask Name: read-parquet, 2 tasks
```

**Environment**:

Dask version: 2.30.0
Python version: 3.8
Operating System: win10
Install method (conda, pip, source): pip

Discussion:

Patch:
diff --git a/dask/array/blockwise.py b/dask/array/blockwise.py
index 72fb9bc2f..b1386dba9 100644
--- a/dask/array/blockwise.py
+++ b/dask/array/blockwise.py
@@ -3,10 +3,132 @@ import warnings
 
 import tlz as toolz
 
+from collections.abc import Mapping
+from functools import reduce
+from itertools import product
+
 from .. import base, utils
 from ..delayed import unpack_collections
 from ..highlevelgraph import HighLevelGraph
-from ..blockwise import blockwise as core_blockwise
+from ..blockwise import BlockwiseIO, blockwise as core_blockwise
+
+
+class CreateArraySubgraph(Mapping):
+    """
+    Subgraph for numpy array creation utilities
+
+    Parameters
+    ----------
+    name: string
+        The output name.
+    func : callable
+        Function to apply to populate individual blocks. This function should take
+        an iterable containing the dimensions of the given block.
+    shape: iterable
+        Iterable containing the overall shape of the array.
+    chunks: iterable
+        Iterable containing the chunk sizes along each dimension of array.
+    """
+
+    def __init__(
+        self,
+        name,
+        func,
+        shape,
+        chunks,
+    ):
+        self.name = name
+        self.func = func
+        self.shape = shape
+        self.chunks = chunks
+        # Compute total number of blocks
+        self._len = reduce(lambda acc, x: len(x) * acc, chunks, 1)
+
+    def __getitem__(self, key):
+        """
+        Get a task for a given key.
+
+        The tasks are constructed on demand.
+        """
+        try:
+            name, *idx = key
+        except ValueError:
+            # too many / few values to unpack
+            raise KeyError(key) from None
+
+        if name != self.name:
+            raise KeyError(key)
+
+        # Check for invalid indices
+        if any(i < 0 or i >= len(c) for i, c in zip(idx, self.chunks)):
+            raise KeyError(key)
+
+        # Get this chunk shape
+        ndim = len(self.chunks)
+        chunk = tuple(self.chunks[i][idx[i]] for i in range(ndim))
+
+        return (self.func, chunk)
+
+    def __len__(self):
+        return self._len
+
+    def __iter__(self):
+        """
+        Produce keys for each chunk lazily.
+        """
+        ind = tuple(tuple(range(len(c))) for c in self.chunks)
+        for chunk in product(*ind):
+            yield (self.name, *chunk)
+
+
+class BlockwiseCreateArray(BlockwiseIO):
+    """
+    Specialized Blockwise Layer for array creation routines.
+
+    Enables HighLevelGraph optimizations.
+
+    Parameters
+    ----------
+    name: string
+        The output name.
+    func : callable
+        Function to apply to populate individual blocks. This function should take
+        an iterable containing the dimensions of the given block.
+    shape: iterable
+        Iterable containing the overall shape of the array.
+    chunks: iterable
+        Iterable containing the chunk sizes along each dimension of array.
+    """
+
+    def __init__(
+        self,
+        name,
+        func,
+        shape,
+        chunks,
+    ):
+        self.name = name
+        self.func = func
+        self.shape = shape
+        self.chunks = chunks
+        self.io_name = "blockwise-create-" + name
+        dsk_io = CreateArraySubgraph(
+            self.io_name,
+            func,
+            shape,
+            chunks,
+        )
+        out_ind = tuple(range(len(shape)))
+
+        nchunks = tuple(len(chunk) for chunk in self.chunks)
+        super().__init__(
+            {self.io_name: dsk_io},
+            self.name,
+            out_ind,
+            None,
+            [(self.io_name, out_ind)],
+            {self.io_name: nchunks},
+        )
 
 
 def blockwise(
diff --git a/dask/array/wrap.py b/dask/array/wrap.py
index 6be5405da..790171254 100644
--- a/dask/array/wrap.py
+++ b/dask/array/wrap.py
@@ -7,6 +7,7 @@ from tlz import curry
 
 from ..base import tokenize
 from ..utils import funcname
+from .blockwise import BlockwiseCreateArray
 from .core import Array, normalize_chunks
 from .utils import (
     meta_from_array,
@@ -68,14 +69,15 @@ def wrap_func_shape_as_first_arg(func, *args, **kwargs):
     chunks = parsed["chunks"]
     name = parsed["name"]
     kwargs = parsed["kwargs"]
-
-    keys = product([name], *[range(len(bd)) for bd in chunks])
-    shapes = product(*chunks)
     func = partial(func, dtype=dtype, **kwargs)
-    vals = ((func,) + (s,) + args for s in shapes)
 
-    dsk = dict(zip(keys, vals))
-    return Array(dsk, name, chunks, dtype=dtype, meta=kwargs.get("meta", None))
+    graph = BlockwiseCreateArray(
+        name,
+        func,
+        shape,
+        chunks,
+    )
+    return Array(graph, name, chunks, dtype=dtype, meta=kwargs.get("meta", None))
 
 
 def wrap_func_like(func, *args, **kwargs):
diff --git a/dask/blockwise.py b/dask/blockwise.py
index 164d89549..2fdcb55a4 100644
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -1,5 +1,4 @@
 import itertools
-import warnings
 
 import numpy as np
 
@@ -572,6 +571,9 @@ class BlockwiseIO(Blockwise):
         # BlockwiseIO requires `io_deps` inputs
         self.io_deps = io_deps
 
+    def __repr__(self):
+        return "BlockwiseIO<{} -> {}>".format(self.indices, self.output)
+
     @property
     def _dict(self):
         if hasattr(self, "_cached_dict"):
@@ -658,21 +660,13 @@ def _inject_io_tasks(dsk, io_deps, output_indices, new_axes):
     # Loop through graph, and replace IO-placeholder tasks
     # with the actual underlying IO function
     for k in dsk:
-        for io_name, io_subgraph in io_deps.items():
-            # Leave out `new_axes` in key
-            io_key = (io_name,) + tuple(
-                [
-                    k[i + 1]
-                    for i, idx in enumerate(output_indices)
-                    if idx not in new_axes
-                ]
-            )
-            if io_key in dsk[k]:
-                # Inject IO-function arguments into the blockwise graph
-                # as a single (packed) tuple.
+        for i in range(1, len(dsk[k])):
+            io_key = dsk[k][i]
+            if isinstance(io_key, tuple) and len(io_key) and io_key[0] in io_deps:
+                io_subgraph = io_deps.get(io_key[0])
                 io_item = io_subgraph.get(io_key)
                 io_item = list(io_item[1:]) if len(io_item) > 1 else []
-                new_task = [io_item if v == io_key else v for v in dsk[k]]
+                new_task = [io_item if j == i else v for j, v in enumerate(dsk[k])]
                 dsk[k] = tuple(new_task)
 
     return dsk
@@ -1072,26 +1066,10 @@ def optimize_blockwise(graph, keys=()):
     --------
     rewrite_blockwise
     """
-    with warnings.catch_warnings():
-        # In some cases, rewrite_blockwise (called internally) will do a bad
-        # thing like `string in array[int].
-        # See dask/array/tests/test_atop.py::test_blockwise_numpy_arg for
-        # an example. NumPy currently raises a warning that 'a' == array([1, 2])
-        # will change from returning `False` to `array([False, False])`.
-        #
-        # Users shouldn't see those warnings, so we filter them.
-        # We set the filter here, rather than lower down, to avoid having to
-        # create and remove the filter many times inside a tight loop.
-
-        # https://github.com/dask/dask/pull/4805#discussion_r286545277 explains
-        # why silencing this warning shouldn't cause issues.
-        warnings.filterwarnings(
-            "ignore", "elementwise comparison failed", Warning
-        )  # FutureWarning or DeprecationWarning
+    out = _optimize_blockwise(graph, keys=keys)
+    while out.dependencies != graph.dependencies:
+        graph = out
         out = _optimize_blockwise(graph, keys=keys)
-        while out.dependencies != graph.dependencies:
-            graph = out
-            out = _optimize_blockwise(graph, keys=keys)
     return out
 
 
diff --git a/dask/dataframe/io/parquet/fastparquet.py b/dask/dataframe/io/parquet/fastparquet.py
index 8aa17aa6a..1b36226ca 100644
--- a/dask/dataframe/io/parquet/fastparquet.py
+++ b/dask/dataframe/io/parquet/fastparquet.py
@@ -322,8 +322,9 @@ class FastParquetEngine(Engine):
                             skip_cols.add(col)
                             continue
                         if isinstance(cs_min, np.datetime64):
-                            cs_min = pd.Timestamp(cs_min)
-                            cs_max = pd.Timestamp(cs_max)
+                            tz = getattr(dtypes[col], "tz", None)
+                            cs_min = pd.Timestamp(cs_min, tz=tz)
+                            cs_max = pd.Timestamp(cs_max, tz=tz)
                         d.update(
                             {
                                 "min": cs_min,


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 832f2e9ca..ac1721ac1 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -1528,7 +1528,7 @@ def test_map_blocks_optimize_blockwise(func):
     dsk = c.__dask_graph__()
     optimized = optimize_blockwise(dsk)
     # The two additions and the map_blocks should be fused together
-    assert len(optimized.layers) == len(dsk.layers) - 2
+    assert len(optimized.layers) == len(dsk.layers) - 6
 
 
 def test_repr():
diff --git a/dask/array/tests/test_optimization.py b/dask/array/tests/test_optimization.py
index 8a47ddd4c..c7bac40e6 100644
--- a/dask/array/tests/test_optimization.py
+++ b/dask/array/tests/test_optimization.py
@@ -9,7 +9,13 @@ from dask.highlevelgraph import HighLevelGraph
 from dask.optimization import fuse
 from dask.utils import SerializableLock
 from dask.array.core import getter, getter_nofancy
-from dask.array.optimization import getitem, optimize, optimize_slices, fuse_slice
+from dask.array.optimization import (
+    getitem,
+    optimize,
+    optimize_blockwise,
+    optimize_slices,
+    fuse_slice,
+)
 from dask.array.utils import assert_eq
 
 
@@ -388,6 +394,22 @@ def test_disable_lowlevel_fusion():
         assert_eq(y, [1] * 3)
 
 
+def test_array_creation_blockwise_fusion():
+    """
+    Check that certain array creation routines work with blockwise and can be
+    fused with other blockwise operations.
+    """
+    x = da.ones(3, chunks=(3,))
+    y = da.zeros(3, chunks=(3,))
+    z = da.full(3, fill_value=2, chunks=(3,))
+    a = x + y + z
+    dsk1 = a.__dask_graph__()
+    assert len(dsk1) == 5
+    dsk2 = optimize_blockwise(dsk1)
+    assert len(dsk2) == 1
+    assert_eq(a, np.full(3, 3))
+
+
 def test_gh3937():
     # test for github issue #3937
     x = da.from_array([1, 2, 3.0], (2,))
@@ -428,7 +450,7 @@ def test_fuse_roots_annotations():
 
     z = (x + 1) + (2 * y)
     hlg = dask.blockwise.optimize_blockwise(z.dask)
-    assert len(hlg.layers) == 4
+    assert len(hlg.layers) == 3
     assert {"foo": "bar"} in [l.annotations for l in hlg.layers.values()]
     za = da.Array(hlg, z.name, z.chunks, z.dtype)
     assert_eq(za, z)
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
index 154bca8c0..432c02ebf 100644
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -873,6 +873,9 @@ def test_read_parquet_custom_columns(tmpdir, engine):
         ),
         (pd.DataFrame({"x": [3, 2, 1]}).astype("M8[us]"), {}, {}),
         (pd.DataFrame({"x": [3, 2, 1]}).astype("M8[ms]"), {}, {}),
+        (pd.DataFrame({"x": [3000, 2000, 1000]}).astype("datetime64[ns]"), {}, {}),
+        (pd.DataFrame({"x": [3000, 2000, 1000]}).astype("datetime64[ns, UTC]"), {}, {}),
+        (pd.DataFrame({"x": [3000, 2000, 1000]}).astype("datetime64[ns, CET]"), {}, {}),
         (pd.DataFrame({"x": [3, 2, 1]}).astype("uint16"), {}, {}),
         (pd.DataFrame({"x": [3, 2, 1]}).astype("float32"), {}, {}),
         (pd.DataFrame({"x": [3, 1, 2]}, index=[3, 2, 1]), {}, {}),