[REVIEW] Enable writing to s3 storage in chunked parquet writer

conda/environments/cudf_dev_cuda11.5.yml

@@ -67,6 +67,13 @@ dependencies:
   - pydata-sphinx-theme
   - librdkafka=1.7.0
   - python-confluent-kafka=1.7.0
+  - moto>=3.1.6
+  - boto3>=1.21.21
+  - botocore>=1.24.21
+  - aiobotocore>=2.2.0
+  - s3fs>=2022.3.0
+  - flask
+  - flask_cors
   - pip:
       - git+https://github.com/python-streamz/streamz.git@master
       - pyorc

python/cudf/cudf/io/parquet.py

@@ -1,12 +1,19 @@
 # Copyright (c) 2019-2022, NVIDIA CORPORATION.
 
+import shutil
+import tempfile
 import warnings
 from collections import defaultdict
 from contextlib import ExitStack
 from typing import Dict, List, Tuple
 from uuid import uuid4
 
 import numpy as np
+
+try:
+    import s3fs
+except (ImportError, ModuleNotFoundError):
+    s3fs = None
 from pyarrow import dataset as ds, parquet as pq
 
 import cudf
@@ -206,12 +213,26 @@ def _process_dataset(
         filters = pq._filters_to_expression(filters)
 
     # Initialize ds.FilesystemDataset
-    dataset = ds.dataset(
-        paths,
-        filesystem=fs,
-        format="parquet",
-        partitioning="hive",
-    )
+    if (
+        s3fs is not None
+        and isinstance(fs, s3fs.S3FileSystem)
+        and len(paths) == 1
+        and fs.isdir(paths[0])
+    ):
+        # TODO: Remove this workaround after following bug is fixed:
+        # https://issues.apache.org/jira/browse/ARROW-16438
+        dataset = ds.dataset(
+            "s3://" + paths[0],
+            format="parquet",
+            partitioning="hive",
+        )
+    else:
+        dataset = ds.dataset(
+            paths,
+            filesystem=fs,
+            format="parquet",
+            partitioning="hive",
+        )
     file_list = dataset.files
     if len(file_list) == 0:
         raise FileNotFoundError(f"{paths} could not be resolved to any files")
@@ -724,6 +745,7 @@ def __init__(
         index=None,
         compression=None,
         statistics="ROWGROUP",
+        **kwargs,
     ) -> None:
         """
         Write a parquet file or dataset incrementally
@@ -776,7 +798,12 @@ def __init__(
                     <filename>.parquet
 
         """
-        self.path = path
+        if isinstance(path, str) and path.startswith("s3://"):
+            self.fs_meta = {"is_s3": True, "actual_path": path}
+            self.path = tempfile.TemporaryDirectory().name
+        else:
+            self.fs_meta = {}
+            self.path = path
         self.common_args = {
             "index": index,
             "compression": compression,
@@ -792,6 +819,7 @@ def __init__(
         # in self._chunked_writers for reverse lookup
         self.path_cw_map: Dict[str, int] = {}
         self.filename = None
+        self.kwargs = kwargs
 
     @_cudf_nvtx_annotate
     def write_table(self, df):
@@ -837,18 +865,19 @@ def write_table(self, df):
             ]
             cw.write_table(grouped_df, this_cw_part_info)
 
-        # Create new cw for unhandled paths encountered in this write_table
-        new_paths, part_info, meta_paths = zip(*new_cw_paths)
-        self._chunked_writers.append(
-            (
-                ParquetWriter(new_paths, **self.common_args),
-                new_paths,
-                meta_paths,
+        if new_cw_paths:
+            # Create new cw for unhandled paths encountered in this write_table
+            new_paths, part_info, meta_paths = zip(*new_cw_paths)
+            self._chunked_writers.append(
+                (
+                    ParquetWriter(new_paths, **self.common_args),
+                    new_paths,
+                    meta_paths,
+                )
             )
-        )
-        new_cw_idx = len(self._chunked_writers) - 1
-        self.path_cw_map.update({k: new_cw_idx for k in new_paths})
-        self._chunked_writers[-1][0].write_table(grouped_df, part_info)
+            new_cw_idx = len(self._chunked_writers) - 1
+            self.path_cw_map.update({k: new_cw_idx for k in new_paths})
+            self._chunked_writers[-1][0].write_table(grouped_df, part_info)
 
     @_cudf_nvtx_annotate
     def close(self, return_metadata=False):
@@ -862,6 +891,15 @@ def close(self, return_metadata=False):
             for cw, _, meta_path in self._chunked_writers
         ]
 
+        if self.fs_meta.get("is_s3", False):
+            local_path = self.path
+            s3_path = self.fs_meta["actual_path"]
+            s3_file, _ = ioutils._get_filesystem_and_paths(
+                s3_path, **self.kwargs
+            )
+            s3_file.put(local_path, s3_path, recursive=True)
+            shutil.rmtree(self.path)
+
         if return_metadata:
             return (
                 merge_parquet_filemetadata(metadata)

python/cudf/cudf/tests/test_s3.py

@@ -18,10 +18,12 @@
 import cudf
 from cudf.testing._utils import assert_eq
 
-moto = pytest.importorskip("moto", minversion="1.3.14")
+moto = pytest.importorskip("moto", minversion="3.1.6")
 boto3 = pytest.importorskip("boto3")
 requests = pytest.importorskip("requests")
 s3fs = pytest.importorskip("s3fs")
+flask = pytest.importorskip("flask")
+flask_cors = pytest.importorskip("flask_cors")
 
 
 @contextmanager
@@ -49,6 +51,7 @@ def s3_base(worker_id):
         # system aws credentials, https://github.com/spulec/moto/issues/1793
         os.environ.setdefault("AWS_ACCESS_KEY_ID", "foobar_key")
         os.environ.setdefault("AWS_SECRET_ACCESS_KEY", "foobar_secret")
+        os.environ.setdefault("S3FS_LOGGING_LEVEL", "DEBUG")
 
         # Launching moto in server mode, i.e., as a separate process
         # with an S3 endpoint on localhost
@@ -457,3 +460,42 @@ def test_write_orc(s3_base, s3so, pdf):
             got = pa.orc.ORCFile(f).read().to_pandas()
 
     assert_eq(pdf, got)
+
+
+def test_write_chunked_parquet(s3_base, s3so):
+    df1 = cudf.DataFrame({"b": [10, 11, 12], "a": [1, 2, 3]})
+    df2 = cudf.DataFrame({"b": [20, 30, 50], "a": [3, 2, 1]})
+    dirname = "chunked_writer_directory"
+    bname = "parquet"
+    from cudf.io.parquet import ParquetDatasetWriter
+
+    with s3_context(
+        s3_base=s3_base, bucket=bname, files={dirname: BytesIO()}
+    ) as s3fs:
+        cw = ParquetDatasetWriter(
+            f"s3://{bname}/{dirname}",
+            partition_cols=["a"],
+            storage_options=s3so,
+        )
+        cw.write_table(df1)
+        cw.write_table(df2)
+        cw.close()
+
+        # TODO: Replace following workaround with:
+        # expect = cudf.read_parquet(f"s3://{bname}/{dirname}/",
+        # storage_options=s3so)
+        # after the following bug is fixed:
+        # https://issues.apache.org/jira/browse/ARROW-16438
+
+        dfs = []
+        for folder in {"a=1", "a=2", "a=3"}:
+            assert s3fs.exists(f"s3://{bname}/{dirname}/{folder}")
+            for file in s3fs.ls(f"s3://{bname}/{dirname}/{folder}"):
+                df = cudf.read_parquet("s3://" + file, storage_options=s3so)
+                dfs.append(df)
+
+        actual = cudf.concat(dfs).astype("int64")
+        assert_eq(
+            actual.sort_values(["b"]).reset_index(drop=True),
+            cudf.concat([df1, df2]).sort_values(["b"]).reset_index(drop=True),
+        )

python/dask_cudf/dask_cudf/io/tests/test_s3.py

@@ -1,3 +1,5 @@
+# Copyright (c) 2020-2022, NVIDIA CORPORATION.
+
 import os
 import shlex
 import subprocess
@@ -42,6 +44,7 @@ def s3_base(worker_id):
         # system aws credentials, https://github.com/spulec/moto/issues/1793
         os.environ.setdefault("AWS_ACCESS_KEY_ID", "foobar_key")
         os.environ.setdefault("AWS_SECRET_ACCESS_KEY", "foobar_secret")
+        os.environ.setdefault("S3FS_LOGGING_LEVEL", "DEBUG")
 
         # Launching moto in server mode, i.e., as a separate process
         # with an S3 endpoint on localhost