Repository: mars-project/mars

Issue:
Support reading parquet from partitioned datasets for fastparquet engine
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

Now we support reading parquet from partitioned datasets for `pyarrow` engine, the `fastparquet` engine for partitioned datasets is needed as well.


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
@@ -16,6 +16,7 @@
 
 import os
 import pickle
+from typing import List, Tuple
 from urllib.parse import urlparse
 
 import numpy as np
@@ -34,7 +35,7 @@
 
 from ... import opcodes as OperandDef
 from ...config import options
-from ...lib.filesystem import file_size, get_fs, glob, open_file
+from ...lib.filesystem import file_size, get_fs, glob, open_file, FileSystem
 from ...serialization.serializables import (
     AnyField,
     BoolField,
@@ -52,7 +53,6 @@
     parse_index,
     to_arrow_dtypes,
     contain_arrow_dtype,
-    arrow_table_to_pandas_dataframe,
 )
 from .core import (
     IncrementalIndexDatasource,
@@ -112,6 +112,64 @@ def read_group_to_pandas(
     ):
         raise NotImplementedError
 
+    def read_partitioned_to_pandas(
+        self,
+        f,
+        partitions: pq.ParquetPartitions,
+        partition_keys: List[Tuple],
+        columns=None,
+        nrows=None,
+        use_arrow_dtype=None,
+        **kwargs,
+    ):
+        raw_df = self.read_to_pandas(
+            f, columns=columns, nrows=nrows, use_arrow_dtype=use_arrow_dtype, **kwargs
+        )
+        for i, (name, index) in enumerate(partition_keys):
+            dictionary = partitions[i].dictionary
+            value = dictionary[index]
+            raw_df[name] = pd.Categorical(
+                np.full(raw_df.shape[0], value.as_py()), categories=dictionary.tolist()
+            )
+        return raw_df
+
+    def read_partitioned_dtypes(self, fs: FileSystem, directory, storage_options):
+        # As ParquetDataset will iterate all files,
+        # here we just find one file to infer dtypes
+        current_path = directory
+        partition_cols = []
+        while fs.isdir(current_path):
+            _, dirs, files = next(fs.walk(current_path))
+            dirs = [d for d in dirs if not d.startswith(".")]
+            files = [f for f in files if not f.startswith(".")]
+            if len(files) == 0:
+                # directory as partition
+                partition_cols.append(dirs[0].split("=", 1)[0])
+                current_path = os.path.join(current_path, dirs[0])
+            elif len(dirs) == 0:
+                # parquet files in deepest directory
+                current_path = os.path.join(current_path, files[0])
+            else:  # pragma: no cover
+                raise ValueError(
+                    "Files and directories are mixed in an intermediate directory"
+                )
+
+        # current path is now a parquet file
+        with open_file(current_path, storage_options=storage_options) as f:
+            dtypes = self.read_dtypes(f)
+        for partition in partition_cols:
+            dtypes[partition] = pd.CategoricalDtype()
+        return dtypes
+
+
+def _parse_prefix(path):
+    path_prefix = ""
+    if isinstance(path, str):
+        parsed_path = urlparse(path)
+        if parsed_path.scheme:
+            path_prefix = f"{parsed_path.scheme}://{parsed_path.netloc}"
+    return path_prefix
+
 
 class ArrowEngine(ParquetEngine):
     def get_row_num(self, f):
@@ -150,7 +208,7 @@ def read_group_to_pandas(
 class FastpaquetEngine(ParquetEngine):
     def get_row_num(self, f):
         file = fastparquet.ParquetFile(f)
-        return file.count
+        return file.count()
 
     def read_dtypes(self, f, **kwargs):
         file = fastparquet.ParquetFile(f)
@@ -161,12 +219,11 @@ def read_to_pandas(
         self, f, columns=None, nrows=None, use_arrow_dtype=None, **kwargs
     ):
         file = fastparquet.ParquetFile(f)
-        df = file.to_pandas(**kwargs)
+        df = file.to_pandas(columns, **kwargs)
         if nrows is not None:
             df = df.head(nrows)
         if use_arrow_dtype:
             df = df.astype(to_arrow_dtypes(df.dtypes).to_dict())
-
         return df
 
 
@@ -186,6 +243,7 @@ class DataFrameReadParquet(
     read_kwargs = DictField("read_kwargs")
     incremental_index = BoolField("incremental_index")
     storage_options = DictField("storage_options")
+    is_partitioned = BoolField("is_partitioned")
     merge_small_files = BoolField("merge_small_files")
     merge_small_file_options = DictField("merge_small_file_options")
     # for chunk
@@ -221,11 +279,7 @@ def _tile_partitioned(cls, op: "DataFrameReadParquet"):
         dtypes = cls._to_arrow_dtypes(out_df.dtypes, op)
         dataset = pq.ParquetDataset(op.path)
 
-        parsed_path = urlparse(op.path)
-        if not os.path.exists(op.path) and parsed_path.scheme:
-            path_prefix = f"{parsed_path.scheme}://{parsed_path.netloc}"
-        else:
-            path_prefix = ""
+        path_prefix = _parse_prefix(op.path)
 
         chunk_index = 0
         out_chunks = []
@@ -272,14 +326,20 @@ def _tile_no_partitioned(cls, op: "DataFrameReadParquet"):
         dtypes = cls._to_arrow_dtypes(out_df.dtypes, op)
         shape = (np.nan, out_df.shape[1])
 
-        paths = (
-            op.path
-            if isinstance(op.path, (tuple, list))
-            else glob(op.path, storage_options=op.storage_options)
-        )
+        path_prefix = ""
+        if isinstance(op.path, (tuple, list)):
+            paths = op.path
+        elif get_fs(op.path, op.storage_options).isdir(op.path):
+            parsed_path = urlparse(op.path)
+            if parsed_path.scheme.lower() == "hdfs":
+                path_prefix = f"{parsed_path.scheme}://{parsed_path.netloc}"
+            paths = get_fs(op.path, op.storage_options).ls(op.path)
+        else:
+            paths = glob(op.path, storage_options=op.storage_options)
 
         first_chunk_row_num, first_chunk_raw_bytes = None, None
         for i, pth in enumerate(paths):
+            pth = path_prefix + pth
             if i == 0:
                 with open_file(pth, storage_options=op.storage_options) as f:
                     first_chunk_row_num = get_engine(op.engine).get_row_num(f)
@@ -336,7 +396,7 @@ def _tile_no_partitioned(cls, op: "DataFrameReadParquet"):
 
     @classmethod
     def _tile(cls, op: "DataFrameReadParquet"):
-        if get_fs(op.path, op.storage_options).isdir(op.path):
+        if op.is_partitioned:
             tiled = cls._tile_partitioned(op)
         else:
             tiled = cls._tile_no_partitioned(op)
@@ -349,14 +409,18 @@ def _tile(cls, op: "DataFrameReadParquet"):
     @classmethod
     def _execute_partitioned(cls, ctx, op: "DataFrameReadParquet"):
         out = op.outputs[0]
+        engine = get_engine(op.engine)
         partitions = pickle.loads(op.partitions)
-        piece = pq.ParquetDatasetPiece(
-            op.path, partition_keys=op.partition_keys, open_file_func=open_file
-        )
-        table = piece.read(partitions=partitions)
-        if op.nrows is not None:
-            table = table.slice(0, op.nrows)
-        ctx[out.key] = arrow_table_to_pandas_dataframe(table, op.use_arrow_dtype)
+        with open_file(op.path, storage_options=op.storage_options) as f:
+            ctx[out.key] = engine.read_partitioned_to_pandas(
+                f,
+                partitions,
+                op.partition_keys,
+                columns=op.columns,
+                nrows=op.nrows,
+                use_arrow_dtype=op.use_arrow_dtype,
+                **op.read_kwargs or dict(),
+            )
 
     @classmethod
     def execute(cls, ctx, op: "DataFrameReadParquet"):
@@ -483,16 +547,18 @@ def read_parquet(
     engine_type = check_engine(engine)
     engine = get_engine(engine_type)
 
-    if get_fs(path, storage_options).isdir(path):
-        # If path is a directory, we will read as a partitioned datasets.
-        if engine_type != "pyarrow":
-            raise TypeError(
-                "Only support pyarrow engine when reading from" "partitioned datasets."
-            )
-        dataset = pq.ParquetDataset(path)
-        dtypes = dataset.schema.to_arrow_schema().empty_table().to_pandas().dtypes
-        for partition in dataset.partitions:
-            dtypes[partition.name] = pd.CategoricalDtype()
+    single_path = path[0] if isinstance(path, list) else path
+    fs = get_fs(single_path, storage_options)
+    is_partitioned = False
+    if fs.isdir(single_path):
+        paths = fs.ls(path)
+        if all(fs.isdir(p) for p in paths):
+            # If all are directories, it is read as a partitioned dataset.
+            dtypes = engine.read_partitioned_dtypes(fs, path, storage_options)
+            is_partitioned = True
+        else:
+            with fs.open(paths[0], mode="rb") as f:
+                dtypes = engine.read_dtypes(f)
     else:
         if not isinstance(path, list):
             file_path = glob(path, storage_options=storage_options)[0]
@@ -502,8 +568,8 @@ def read_parquet(
         with open_file(file_path, storage_options=storage_options) as f:
             dtypes = engine.read_dtypes(f)
 
-        if columns:
-            dtypes = dtypes[columns]
+    if columns:
+        dtypes = dtypes[columns]
 
     if use_arrow_dtype is None:
         use_arrow_dtype = options.dataframe.use_arrow_dtype
@@ -521,6 +587,7 @@ def read_parquet(
         read_kwargs=kwargs,
         incremental_index=incremental_index,
         storage_options=storage_options,
+        is_partitioned=is_partitioned,
         memory_scale=memory_scale,
         merge_small_files=merge_small_files,
         merge_small_file_options=merge_small_file_options,


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
@@ -1005,8 +1005,18 @@ def test_date_range_execution(setup):
     pd.testing.assert_index_equal(result, expected)
 
 
-@pytest.mark.skipif(pa is None, reason="pyarrow not installed")
-def test_read_parquet_arrow(setup):
+parquet_engines = ["auto"]
+if pa is not None:
+    parquet_engines.append("pyarrow")
+if fastparquet is not None:
+    parquet_engines.append("fastparquet")
+
+
+@pytest.mark.skipif(
+    len(parquet_engines) == 1, reason="pyarrow and fastparquet are not installed"
+)
+@pytest.mark.parametrize("engine", parquet_engines)
+def test_read_parquet_arrow(setup, engine):
     test_df = pd.DataFrame(
         {
             "a": np.arange(10).astype(np.int64, copy=False),
@@ -1019,36 +1029,41 @@ def test_read_parquet_arrow(setup):
         file_path = os.path.join(tempdir, "test.csv")
         test_df.to_parquet(file_path)
 
-        df = md.read_parquet(file_path)
+        df = md.read_parquet(file_path, engine=engine)
         result = df.execute().fetch()
         pd.testing.assert_frame_equal(result, test_df)
         # size_res = self.executor.execute_dataframe(df, mock=True)
         # assert sum(s[0] for s in size_res) > test_df.memory_usage(deep=True).sum()
 
-    with tempfile.TemporaryDirectory() as tempdir:
-        file_path = os.path.join(tempdir, "test.parquet")
-        test_df.to_parquet(file_path, row_group_size=3)
+    if engine != "fastparquet":
+        with tempfile.TemporaryDirectory() as tempdir:
+            file_path = os.path.join(tempdir, "test.parquet")
+            test_df.to_parquet(file_path, row_group_size=3)
 
-        df = md.read_parquet(file_path, groups_as_chunks=True, columns=["a", "b"])
-        result = df.execute().fetch()
-        pd.testing.assert_frame_equal(
-            result.reset_index(drop=True), test_df[["a", "b"]]
-        )
+            df = md.read_parquet(
+                file_path, groups_as_chunks=True, columns=["a", "b"], engine=engine
+            )
+            result = df.execute().fetch()
+            pd.testing.assert_frame_equal(
+                result.reset_index(drop=True), test_df[["a", "b"]]
+            )
 
-    with tempfile.TemporaryDirectory() as tempdir:
-        file_path = os.path.join(tempdir, "test.parquet")
-        test_df.to_parquet(file_path, row_group_size=5)
+    if engine != "fastparquet":
+        with tempfile.TemporaryDirectory() as tempdir:
+            file_path = os.path.join(tempdir, "test.parquet")
+            test_df.to_parquet(file_path, row_group_size=5)
 
-        df = md.read_parquet(
-            file_path,
-            groups_as_chunks=True,
-            use_arrow_dtype=True,
-            incremental_index=True,
-        )
-        result = df.execute().fetch()
-        assert isinstance(df.dtypes.iloc[1], md.ArrowStringDtype)
-        assert isinstance(result.dtypes.iloc[1], md.ArrowStringDtype)
-        pd.testing.assert_frame_equal(arrow_array_to_objects(result), test_df)
+            df = md.read_parquet(
+                file_path,
+                groups_as_chunks=True,
+                use_arrow_dtype=True,
+                incremental_index=True,
+                engine=engine,
+            )
+            result = df.execute().fetch()
+            assert isinstance(df.dtypes.iloc[1], md.ArrowStringDtype)
+            assert isinstance(result.dtypes.iloc[1], md.ArrowStringDtype)
+            pd.testing.assert_frame_equal(arrow_array_to_objects(result), test_df)
 
     # test wildcards in path
     for merge_small_file_option in [{"n_sample_file": 1}, None]:
@@ -1066,23 +1081,38 @@ def test_read_parquet_arrow(setup):
             df[100:200].to_parquet(file_paths[1], row_group_size=30)
             df[200:].to_parquet(file_paths[2])
 
-            mdf = md.read_parquet(f"{tempdir}/*.parquet")
+            mdf = md.read_parquet(f"{tempdir}/*.parquet", engine=engine)
+            r = mdf.execute().fetch()
+            pd.testing.assert_frame_equal(df, r.sort_values("a").reset_index(drop=True))
+
+            mdf = md.read_parquet(f"{tempdir}", engine=engine)
+            r = mdf.execute().fetch()
+            pd.testing.assert_frame_equal(df, r.sort_values("a").reset_index(drop=True))
+
+            file_list = [os.path.join(tempdir, name) for name in os.listdir(tempdir)]
+            mdf = md.read_parquet(file_list, engine=engine)
             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)
+            mdf = md.read_parquet(
+                f"{tempdir}/*.parquet", engine=engine, 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,
-                merge_small_file_options=merge_small_file_option,
-            )
-            r = mdf.execute().fetch()
-            pd.testing.assert_frame_equal(df, r.sort_values("a").reset_index(drop=True))
+            if engine != "fastparquet":
+                mdf = md.read_parquet(
+                    f"{tempdir}/*.parquet",
+                    groups_as_chunks=True,
+                    engine=engine,
+                    merge_small_file_options=merge_small_file_option,
+                )
+                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:
@@ -1094,7 +1124,7 @@ def test_read_parquet_arrow(setup):
             }
         )
         df.to_parquet(tempdir, partition_cols=["c"])
-        mdf = md.read_parquet(tempdir)
+        mdf = md.read_parquet(tempdir, engine=engine)
         r = mdf.execute().fetch().astype(df.dtypes)
         pd.testing.assert_frame_equal(
             df.sort_values("a").reset_index(drop=True),