Repository: mars-project/mars

Issue:
[BUG] `use_arrow_dtype` doesn't work when read a directory path for `read_parquet`
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
`use_arrow_dtype` doesn't work when read a directory path for `read_parquet`

**To Reproduce**
``` Python
In [26]: df = md.read_parquet(tempdir, use_arrow_dtype=True)

In [27]: df.dtypes
Out[27]: 
a      int64
b     object
c    float64
dtype: object

```
Expected dtypes should be:
``` Python
In [14]: df.dtypes
Out[14]: 
a            int64
b    Arrow[string]
c          float64
dtype: object
```


**Expected behavior**
A clear and concise description of what you expected to happen.

**Additional context**
Add any other context about the problem here.



Discussion:

Resolving patch:
diff --git a/mars/dataframe/datasource/read_parquet.py b/mars/dataframe/datasource/read_parquet.py
--- a/mars/dataframe/datasource/read_parquet.py
+++ b/mars/dataframe/datasource/read_parquet.py
@@ -48,7 +48,12 @@
 from ...utils import is_object_dtype
 from ..arrays import ArrowStringDtype
 from ..operands import OutputType
-from ..utils import parse_index, to_arrow_dtypes, contain_arrow_dtype
+from ..utils import (
+    parse_index,
+    to_arrow_dtypes,
+    contain_arrow_dtype,
+    arrow_table_to_pandas_dataframe,
+)
 from .core import (
     IncrementalIndexDatasource,
     ColumnPruneSupportedDataSourceMixin,
@@ -351,7 +356,7 @@ def _execute_partitioned(cls, ctx, op: "DataFrameReadParquet"):
         table = piece.read(partitions=partitions)
         if op.nrows is not None:
             table = table.slice(0, op.nrows)
-        ctx[out.key] = table.to_pandas()
+        ctx[out.key] = arrow_table_to_pandas_dataframe(table, op.use_arrow_dtype)
 
     @classmethod
     def execute(cls, ctx, op: "DataFrameReadParquet"):
@@ -500,10 +505,10 @@ def read_parquet(
         if columns:
             dtypes = dtypes[columns]
 
-        if use_arrow_dtype is None:
-            use_arrow_dtype = options.dataframe.use_arrow_dtype
-        if use_arrow_dtype:
-            dtypes = to_arrow_dtypes(dtypes)
+    if use_arrow_dtype is None:
+        use_arrow_dtype = options.dataframe.use_arrow_dtype
+    if use_arrow_dtype:
+        dtypes = to_arrow_dtypes(dtypes)
 
     index_value = parse_index(pd.RangeIndex(-1))
     columns_value = parse_index(dtypes.index, store_data=True)


Test patch:
diff --git a/mars/dataframe/datasource/tests/test_datasource_execution.py b/mars/dataframe/datasource/tests/test_datasource_execution.py
--- a/mars/dataframe/datasource/tests/test_datasource_execution.py
+++ b/mars/dataframe/datasource/tests/test_datasource_execution.py
@@ -1050,6 +1050,12 @@ def test_read_parquet_arrow(setup):
             r = mdf.execute().fetch()
             pd.testing.assert_frame_equal(df, r.sort_values("a").reset_index(drop=True))
 
+            # test `use_arrow_dtype=True`
+            mdf = md.read_parquet(f"{tempdir}/*.parquet", use_arrow_dtype=True)
+            result = mdf.execute().fetch()
+            assert isinstance(mdf.dtypes.iloc[1], md.ArrowStringDtype)
+            assert isinstance(result.dtypes.iloc[1], md.ArrowStringDtype)
+
             mdf = md.read_parquet(
                 f"{tempdir}/*.parquet",
                 groups_as_chunks=True,
@@ -1058,6 +1064,23 @@ def test_read_parquet_arrow(setup):
             r = mdf.execute().fetch()
             pd.testing.assert_frame_equal(df, r.sort_values("a").reset_index(drop=True))
 
+    # test partitioned
+    with tempfile.TemporaryDirectory() as tempdir:
+        df = pd.DataFrame(
+            {
+                "a": np.random.rand(300),
+                "b": [f"s{i}" for i in range(300)],
+                "c": np.random.choice(["a", "b", "c"], (300,)),
+            }
+        )
+        df.to_parquet(tempdir, partition_cols=["c"])
+        mdf = md.read_parquet(tempdir)
+        r = mdf.execute().fetch().astype(df.dtypes)
+        pd.testing.assert_frame_equal(
+            df.sort_values("a").reset_index(drop=True),
+            r.sort_values("a").reset_index(drop=True),
+        )
+
 
 @pytest.mark.skipif(fastparquet is None, reason="fastparquet not installed")
 def test_read_parquet_fast_parquet(setup):