Repository: dask/dask

Issue:
`read_parquet` with `arrow_to_pandas` can cause incorrect `_meta`
**Describe the issue**:

If users provide a `arrow_to_pandas` argument that affects the dtypes, the Dask DataFrame's `_meta` will be incorrect.

**Minimal Complete Verifiable Example**:

```python
>>> import pandas as pd
>>> import dask.dataframe as dd
>>> df = pd.DataFrame({"A": [pd.Timestamp("2000-01-01")]})
>>> df.to_parquet("/tmp/test.parquet")
>>> ddf = dd.read_parquet("/tmp/test.parquet", arrow_to_pandas={"timestamp_as_object": True})
>>> ddf.A.dtype
dtype('<M8[ns]')
>>> ddf.compute().A.dtype
dtype('O')
```

I don't know exactly how `_meta` is constructed from the parquet / arrow metadata, but if there's call to `.to_pandas()` maybe fixing this is as easy as including the user-provided argument.


Discussion:
Thanks @TomAugspurger! Yeah, my guess is we're missing `arrow_to_pandas` somewhere along the way. @rjzamora do you have bandwidth to look into this one? 
II think this is something which came up last year: https://github.com/dask/dask/issues/7219 .   We could do some kind of additional reseting of `meta` in https://github.com/dask/dask/blob/08db480f9387aab3c93c4f256b7c4c726a7314fb/dask/dataframe/io/parquet/core.py#L473-L495 
We can certainly fix this by updating `_create_dd_meta` for the arrow engine to use `_arrow_table_to_pandas`

Resolving patch:
diff --git a/dask/dataframe/io/parquet/arrow.py b/dask/dataframe/io/parquet/arrow.py
--- a/dask/dataframe/io/parquet/arrow.py
+++ b/dask/dataframe/io/parquet/arrow.py
@@ -16,7 +16,6 @@
     Engine,
     _get_aggregation_depth,
     _normalize_index_columns,
-    _parse_pandas_metadata,
     _process_open_file_options,
     _row_groups_to_parts,
     _set_gather_statistics,
@@ -24,12 +23,7 @@
     _sort_and_analyze_paths,
     _split_user_options,
 )
-from dask.dataframe.io.utils import (
-    _get_pyarrow_dtypes,
-    _is_local_fs,
-    _meta_from_dtypes,
-    _open_input_files,
-)
+from dask.dataframe.io.utils import _get_pyarrow_dtypes, _is_local_fs, _open_input_files
 from dask.dataframe.utils import clear_known_categories
 from dask.delayed import Delayed
 from dask.utils import getargspec, natural_sort_key
@@ -978,16 +972,9 @@ def _create_dd_meta(cls, dataset_info):
         physical_column_names = dataset_info.get("physical_schema", schema).names
         columns = None
 
-        # Set index and column names using
-        # pandas metadata (when available)
+        # Use pandas metadata to update categories
         pandas_metadata = _get_pandas_metadata(schema)
         if pandas_metadata:
-            (
-                index_names,
-                column_names,
-                storage_name_mapping,
-                column_index_names,
-            ) = _parse_pandas_metadata(pandas_metadata)
             if categories is None:
                 categories = []
                 for col in pandas_metadata["columns"]:
@@ -995,16 +982,30 @@ def _create_dd_meta(cls, dataset_info):
                         col["name"] not in categories
                     ):
                         categories.append(col["name"])
-        else:
-            # No pandas metadata implies no index, unless selected by the user
-            index_names = []
-            column_names = physical_column_names
-            storage_name_mapping = {k: k for k in column_names}
-            column_index_names = [None]
+
+        # Use _arrow_table_to_pandas to generate meta
+        arrow_to_pandas = dataset_info["kwargs"].get("arrow_to_pandas", {}).copy()
+        meta = cls._arrow_table_to_pandas(
+            schema.empty_table(),
+            categories,
+            arrow_to_pandas=arrow_to_pandas,
+        )
+        index_names = list(meta.index.names)
+        column_names = list(meta.columns)
+        if index_names and index_names != [None]:
+            # Reset the index if non-null index name
+            meta.reset_index(inplace=True)
+
+        # Use index specified in the pandas metadata if
+        # the index column was not specified by the user
         if index is None and index_names:
-            # Pandas metadata has provided the index name for us
             index = index_names
 
+        # Set proper index for meta
+        index_cols = index or ()
+        if index_cols and index_cols != [None]:
+            meta.set_index(index_cols, inplace=True)
+
         # Ensure that there is no overlap between partition columns
         # and explicit column storage
         if partitions:
@@ -1023,25 +1024,20 @@ def _create_dd_meta(cls, dataset_info):
                     )
                 )
 
+        # Get all available column names
         column_names, index_names = _normalize_index_columns(
             columns, column_names + partitions, index, index_names
         )
-
         all_columns = index_names + column_names
 
-        # Check that categories are included in columns
-        if categories and not set(categories).intersection(all_columns):
-            raise ValueError(
-                "categories not in available columns.\n"
-                "categories: {} | columns: {}".format(categories, list(all_columns))
-            )
-
-        dtypes = _get_pyarrow_dtypes(schema, categories)
-        dtypes = {storage_name_mapping.get(k, k): v for k, v in dtypes.items()}
-
-        index_cols = index or ()
-        meta = _meta_from_dtypes(all_columns, dtypes, index_cols, column_index_names)
         if categories:
+            # Check that categories are included in columns
+            if not set(categories).intersection(all_columns):
+                raise ValueError(
+                    "categories not in available columns.\n"
+                    "categories: {} | columns: {}".format(categories, list(all_columns))
+                )
+
             # Make sure all categories are set to "unknown".
             # Cannot include index names in the `cols` argument.
             meta = clear_known_categories(
@@ -1049,7 +1045,7 @@ def _create_dd_meta(cls, dataset_info):
             )
 
         if partition_obj:
-
+            # Update meta dtypes for partitioned columns
             for partition in partition_obj:
                 if isinstance(index, list) and partition.name == index[0]:
                     # Index from directory structure


Test patch:
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -3130,6 +3130,23 @@ def _arrow_table_to_pandas(
     dd.read_parquet(str(tmpdir), engine=ArrowEngineWithTimestampClamp).compute()
 
 
+@PYARROW_MARK
+def test_arrow_to_pandas(tmpdir, engine):
+    # Test that dtypes are correct when arrow_to_pandas is used
+    # (See: https://github.com/dask/dask/issues/9664)
+
+    df = pd.DataFrame({"A": [pd.Timestamp("2000-01-01")]})
+    path = str(tmpdir.join("test.parquet"))
+    df.to_parquet(path, engine=engine)
+
+    arrow_to_pandas = {"timestamp_as_object": True}
+    expect = pq.ParquetFile(path).read().to_pandas(**arrow_to_pandas)
+    got = dd.read_parquet(path, engine="pyarrow", arrow_to_pandas=arrow_to_pandas)
+
+    assert_eq(expect, got)
+    assert got.A.dtype == got.compute().A.dtype
+
+
 @pytest.mark.parametrize(
     "write_cols",
     [["part", "col"], ["part", "kind", "col"]],