Fix use_arrow_dtype parameter for read_parquet (#2698)

Author: hekaisheng

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)

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):