Merge branch 'release-3.0.0' of https://github.com/awslabs/aws-data-wrangler into release-3.0.0

Author: jaidisido

awswrangler/_arrow.py

@@ -8,6 +8,8 @@
 import pandas as pd
 import pyarrow as pa
 
+from awswrangler._data_types import athena2pyarrow
+
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
@@ -80,3 +82,22 @@ def _table_to_df(
         _logger.debug("metadata: %s", metadata)
         df = _apply_timezone(df=df, metadata=metadata)
     return df
+
+
+def _df_to_table(
+    df: pd.DataFrame,
+    schema: Optional[pa.Schema] = None,
+    index: Optional[bool] = None,
+    dtype: Optional[Dict[str, str]] = None,
+    cpus: Optional[int] = None,
+) -> pa.Table:
+    table: pa.Table = pa.Table.from_pandas(df=df, schema=schema, nthreads=cpus, preserve_index=index, safe=True)
+    if dtype:
+        for col_name, col_type in dtype.items():
+            if col_name in table.column_names:
+                col_index = table.column_names.index(col_name)
+                pyarrow_dtype = athena2pyarrow(col_type)
+                field = pa.field(name=col_name, type=pyarrow_dtype)
+                table = table.set_column(col_index, field, table.column(col_name).cast(pyarrow_dtype))
+                _logger.debug("Casting column %s (%s) to %s (%s)", col_name, col_index, col_type, pyarrow_dtype)
+    return table

awswrangler/distributed/ray/datasources/pandas_file_based_datasource.py

@@ -43,7 +43,7 @@ def _get_write_path_for_block(
 class PandasFileBasedDatasource(FileBasedDatasource):  # pylint: disable=abstract-method
     """Pandas file based datasource, for reading and writing Pandas blocks."""
 
-    _FILE_EXTENSION: Optional[str] = None
+    _FILE_EXTENSION: str
 
     def __init__(self) -> None:
         super().__init__()
@@ -109,7 +109,7 @@ def write_block(write_path: str, block: pd.DataFrame) -> str:
 
         write_block_fn = cached_remote_fn(write_block).options(**ray_remote_args)
 
-        file_format = self._FILE_EXTENSION
+        file_suffix = self._get_file_suffix(self._FILE_EXTENSION, compression)
         write_tasks = []
 
         for block_idx, block in enumerate(blocks):
@@ -119,13 +119,16 @@ def write_block(write_path: str, block: pd.DataFrame) -> str:
                 dataset_uuid=dataset_uuid,
                 block=block,
                 block_index=block_idx,
-                file_format=f"{file_format}{_COMPRESSION_2_EXT.get(compression)}",
+                file_format=file_suffix,
             )
             write_task = write_block_fn.remote(write_path, block)
             write_tasks.append(write_task)
 
         return write_tasks
 
+    def _get_file_suffix(self, file_format: str, compression: Optional[str]) -> str:
+        return f"{file_format}{_COMPRESSION_2_EXT.get(compression)}"
+
     def _write_block(
         self,
         f: "pyarrow.NativeFile",

awswrangler/distributed/ray/datasources/parquet_datasource.py

@@ -17,8 +17,9 @@
     _SerializedPiece,
 )
 
-from awswrangler._arrow import _add_table_partitions
+from awswrangler._arrow import _add_table_partitions, _df_to_table
 from awswrangler.distributed.ray.datasources.pandas_file_based_datasource import PandasFileBasedDatasource
+from awswrangler.s3._write import _COMPRESSION_2_EXT
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -102,9 +103,25 @@ def _write_block(  # type: ignore  # pylint: disable=arguments-differ, arguments
         **writer_args: Any,
     ) -> None:
         """Write a block to S3."""
-        import pyarrow.parquet as pq  # pylint: disable=import-outside-toplevel,redefined-outer-name,reimported
+        import pyarrow as pa  # pylint: disable=import-outside-toplevel,redefined-outer-name,reimported
+
+        schema: pa.Schema = writer_args.get("schema", None)
+        dtype: Optional[Dict[str, str]] = writer_args.get("dtype", None)
+        index: bool = writer_args.get("index", False)
+        compression: Optional[str] = writer_args.get("compression", None)
+        pyarrow_additional_kwargs: Optional[Dict[str, Any]] = writer_args.get("pyarrow_additional_kwargs", {})
+
+        pa.parquet.write_table(
+            _df_to_table(block.to_pandas(), schema=schema, index=index, dtype=dtype),
+            f,
+            compression=compression,
+            **pyarrow_additional_kwargs,
+        )
 
-        pq.write_table(block.to_arrow(), f)
+    def _get_file_suffix(self, file_format: str, compression: Optional[str]) -> str:
+        if compression is not None:
+            return f"{_COMPRESSION_2_EXT.get(compression)[1:]}.{file_format}"  # type: ignore
+        return file_format
 
 
 def _resolve_kwargs(kwargs_fn: Callable[[], Dict[str, Any]], **kwargs: Any) -> Dict[str, Any]:

awswrangler/distributed/ray/modin/s3/_read_parquet.py

@@ -36,4 +36,4 @@ def _read_parquet_distributed(  # pylint: disable=unused-argument
         dataset_kwargs=dataset_kwargs,
         path_root=path_root,
     )
-    return _to_modin(dataset=dataset, to_pandas_kwargs=arrow_kwargs)
+    return _to_modin(dataset=dataset, to_pandas_kwargs=arrow_kwargs, ignore_index=bool(path_root))

awswrangler/distributed/ray/modin/s3/_write_parquet.py

@@ -35,11 +35,11 @@ def _to_parquet_distributed(  # pylint: disable=unused-argument
 ) -> List[str]:
     if bucketing:
         # Add bucket id to the prefix
-        path = f"{path_root}{filename_prefix}_bucket-{df.name:05d}.parquet"
+        path = f"{path_root}{filename_prefix}_bucket-{df.name:05d}{compression_ext}.parquet"
     # Create Ray Dataset
     ds = from_modin(df) if isinstance(df, ModinDataFrame) else from_pandas(df)
     # Repartition into a single block if or writing into a single key or if bucketing is enabled
-    if ds.count() > 0 and (path or bucketing):
+    if ds.count() > 0 and (path or bucketing) and not max_rows_by_file:
         _logger.warning(
             "Repartitioning frame to single partition as a strict path was defined: %s. "
             "This operation is inefficient for large datasets.",
@@ -57,7 +57,11 @@ def _to_parquet_distributed(  # pylint: disable=unused-argument
         # If user has provided a single key, use that instead of generating a path per block
         # The dataset will be repartitioned into a single block
         block_path_provider=UserProvidedKeyBlockWritePathProvider()
-        if path and not path.endswith("/")
+        if path and not path.endswith("/") and not max_rows_by_file
         else DefaultBlockWritePathProvider(),
+        index=index,
+        dtype=dtype,
+        compression=compression,
+        pyarrow_additional_kwargs=pyarrow_additional_kwargs,
     )
     return datasource.get_write_paths()

awswrangler/s3/_write_parquet.py

@@ -13,6 +13,7 @@
 import pyarrow.parquet
 
 from awswrangler import _data_types, _utils, catalog, exceptions, lakeformation
+from awswrangler._arrow import _df_to_table
 from awswrangler._config import apply_configs
 from awswrangler._distributed import engine
 from awswrangler.s3._delete import delete_objects
@@ -59,20 +60,13 @@ def _get_chunk_file_path(file_counter: int, file_path: str) -> str:
 def _new_writer(
     file_path: str,
     compression: Optional[str],
-    pyarrow_additional_kwargs: Optional[Dict[str, str]],
+    pyarrow_additional_kwargs: Dict[str, str],
     schema: pa.Schema,
     boto3_session: boto3.Session,
     s3_additional_kwargs: Optional[Dict[str, str]],
     use_threads: Union[bool, int],
 ) -> Iterator[pyarrow.parquet.ParquetWriter]:
     writer: Optional[pyarrow.parquet.ParquetWriter] = None
-    if not pyarrow_additional_kwargs:
-        pyarrow_additional_kwargs = {}
-    if not pyarrow_additional_kwargs.get("coerce_timestamps"):
-        pyarrow_additional_kwargs["coerce_timestamps"] = "ms"
-    if "flavor" not in pyarrow_additional_kwargs:
-        pyarrow_additional_kwargs["flavor"] = "spark"
-
     with open_s3_object(
         path=file_path,
         mode="wb",
@@ -100,7 +94,7 @@ def _write_chunk(
     boto3_session: Optional[boto3.Session],
     s3_additional_kwargs: Optional[Dict[str, str]],
     compression: Optional[str],
-    pyarrow_additional_kwargs: Optional[Dict[str, str]],
+    pyarrow_additional_kwargs: Dict[str, str],
     table: pa.Table,
     offset: int,
     chunk_size: int,
@@ -124,7 +118,7 @@ def _to_parquet_chunked(
     boto3_session: Optional[boto3.Session],
     s3_additional_kwargs: Optional[Dict[str, str]],
     compression: Optional[str],
-    pyarrow_additional_kwargs: Optional[Dict[str, Any]],
+    pyarrow_additional_kwargs: Dict[str, Any],
     table: pa.Table,
     max_rows_by_file: int,
     num_of_rows: int,
@@ -158,7 +152,7 @@ def _to_parquet(
     index: bool,
     compression: Optional[str],
     compression_ext: str,
-    pyarrow_additional_kwargs: Optional[Dict[str, Any]],
+    pyarrow_additional_kwargs: Dict[str, Any],
     cpus: int,
     dtype: Dict[str, str],
     boto3_session: Optional[boto3.Session],
@@ -173,15 +167,7 @@ def _to_parquet(
         path_root=path_root, path=path, filename_prefix=filename_prefix, compression_ext=compression_ext
     )
     _logger.debug("file_path: %s", file_path)
-    table: pa.Table = pyarrow.Table.from_pandas(df=df, schema=schema, nthreads=cpus, preserve_index=index, safe=True)
-    for col_name, col_type in dtype.items():
-        if col_name in table.column_names:
-            col_index = table.column_names.index(col_name)
-            pyarrow_dtype = _data_types.athena2pyarrow(col_type)
-            field = pa.field(name=col_name, type=pyarrow_dtype)
-            table = table.set_column(col_index, field, table.column(col_name).cast(pyarrow_dtype))
-            _logger.debug("Casting column %s (%s) to %s (%s)", col_name, col_index, col_type, pyarrow_dtype)
-
+    table: pa.Table = _df_to_table(df, schema, index, dtype)
     if max_rows_by_file is not None and max_rows_by_file > 0:
         paths: List[str] = _to_parquet_chunked(
             file_path=file_path,
@@ -209,7 +195,6 @@ def _to_parquet(
     return paths
 
 
-@engine.dispatch_on_engine
 @apply_configs
 def to_parquet(  # pylint: disable=too-many-arguments,too-many-locals,too-many-branches,too-many-statements
     df: pd.DataFrame,
@@ -568,6 +553,13 @@ def to_parquet(  # pylint: disable=too-many-arguments,too-many-locals,too-many-b
     filename_prefix = filename_prefix + uuid.uuid4().hex if filename_prefix else uuid.uuid4().hex
     cpus: int = _utils.ensure_cpu_count(use_threads=use_threads)
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
+    # Pyarrow defaults
+    if not pyarrow_additional_kwargs:
+        pyarrow_additional_kwargs = {}
+    if not pyarrow_additional_kwargs.get("coerce_timestamps"):
+        pyarrow_additional_kwargs["coerce_timestamps"] = "ms"
+    if "flavor" not in pyarrow_additional_kwargs:
+        pyarrow_additional_kwargs["flavor"] = "spark"
 
     # Sanitize table to respect Athena's standards
     if (sanitize_columns is True) or (database is not None and table is not None):

tests/_utils.py

@@ -7,13 +7,15 @@
 
 import boto3
 import botocore.exceptions
+from pandas import DataFrame as PandasDataFrame
 
 import awswrangler as wr
 from awswrangler._distributed import EngineEnum, MemoryFormatEnum
 from awswrangler._utils import try_it
 
 if wr.engine.get() == EngineEnum.RAY and wr.memory_format.get() == MemoryFormatEnum.MODIN:
     import modin.pandas as pd
+    from modin.pandas import DataFrame as ModinDataFrame
 else:
     import pandas as pd
 
@@ -433,3 +435,22 @@ def create_workgroup(wkg_name, config):
             Description=f"AWS SDK for pandas Test - {wkg_name}",
         )
     return wkg_name
+
+
+def to_pandas(df: pd.DataFrame) -> PandasDataFrame:
+    """
+    Convert Modin data frames to pandas for comparison
+    """
+    if isinstance(df, PandasDataFrame):
+        return df
+    elif wr.memory_format.get() == MemoryFormatEnum.MODIN and isinstance(df, ModinDataFrame):
+        return df._to_pandas()
+    raise ValueError("Unknown data frame type %s", type(df))
+
+
+def pandas_equals(df1: pd.DataFrame, df2: pd.DataFrame) -> bool:
+    """
+    Check data frames for equality converting them to pandas first
+    """
+    df1, df2 = to_pandas(df1), to_pandas(df2)
+    return df1.equals(df2)

tests/unit/test_s3_parquet.py

@@ -5,16 +5,23 @@
 
 import boto3
 import numpy as np
-import pandas as pd
 import pyarrow as pa
 import pytest
 
 import awswrangler as wr
+from awswrangler._distributed import EngineEnum, MemoryFormatEnum
 
-from .._utils import ensure_data_types, get_df_list
+if wr.engine.get() == EngineEnum.RAY and wr.memory_format.get() == MemoryFormatEnum.MODIN:
+    import modin.pandas as pd
+else:
+    import pandas as pd
+
+from .._utils import ensure_data_types, get_df_list, pandas_equals, to_pandas
 
 logging.getLogger("awswrangler").setLevel(logging.DEBUG)
 
+pytestmark = pytest.mark.distributed
+
 
 @pytest.mark.parametrize("partition_cols", [None, ["c2"], ["c1", "c2"]])
 def test_parquet_metadata_partitions_dataset(path, partition_cols):
@@ -367,12 +374,15 @@ def test_range_index_recovery_simple(path, use_threads):
 @pytest.mark.parametrize("use_threads", [True, False, 2])
 @pytest.mark.parametrize("name", [None, "foo"])
 def test_range_index_recovery_pandas(path, use_threads, name):
+    # Import pandas because modin.to_parquet does not preserve index.name when writing parquet
+    import pandas as pd
+
     df = pd.DataFrame({"c0": np.arange(10, 15, 1)}, dtype="Int64", index=pd.RangeIndex(start=5, stop=30, step=5))
     df.index.name = name
     path_file = f"{path}0.parquet"
     df.to_parquet(path_file)
-    df2 = wr.s3.read_parquet([path_file], use_threads=use_threads, pyarrow_additional_kwargs={"ignore_metadata": False})
-    assert df.reset_index(level=0).equals(df2.reset_index(level=0))
+    df2 = wr.s3.read_parquet(path_file, use_threads=use_threads, pyarrow_additional_kwargs={"ignore_metadata": False})
+    assert pandas_equals(df.reset_index(level=0), df2.reset_index(level=0))
 
 
 @pytest.mark.xfail(raises=AssertionError, reason="Index equality regression")
@@ -409,7 +419,7 @@ def test_index_columns(path, use_threads, name, pandas):
         df.to_parquet(path_file, index=True)
     else:
         wr.s3.to_parquet(df, path_file, index=True)
-    df2 = wr.s3.read_parquet([path_file], columns=["c0"], use_threads=use_threads)
+    df2 = wr.s3.read_parquet(path_file, columns=["c0"], use_threads=use_threads)
     assert df[["c0"]].equals(df2)
 
 
@@ -419,19 +429,21 @@ def test_index_columns(path, use_threads, name, pandas):
 @pytest.mark.parametrize("pandas", [True, False])
 @pytest.mark.parametrize("drop", [True, False])
 def test_range_index_columns(path, use_threads, name, pandas, drop):
+    if wr.memory_format.get() == MemoryFormatEnum.MODIN and not pandas:
+        pytest.skip("Skip due to Modin data frame index not saved as a named column")
     df = pd.DataFrame({"c0": [0, 1], "c1": [2, 3]}, dtype="Int64", index=pd.RangeIndex(start=5, stop=7, step=1))
     df.index.name = name
     path_file = f"{path}0.parquet"
     if pandas:
+        # Convert to pandas because modin.to_parquet() does not preserve index name
+        df = to_pandas(df)
         df.to_parquet(path_file, index=True)
     else:
         wr.s3.to_parquet(df, path_file, index=True)
-
     name = "__index_level_0__" if name is None else name
     columns = ["c0"] if drop else [name, "c0"]
-    df2 = wr.s3.read_parquet([path_file], columns=columns, use_threads=use_threads)
-
-    assert df[["c0"]].reset_index(level=0, drop=drop).equals(df2.reset_index(level=0, drop=drop))
+    df2 = wr.s3.read_parquet(path_file, columns=columns, use_threads=use_threads)
+    assert pandas_equals(df[["c0"]].reset_index(level=0, drop=drop), df2.reset_index(level=0, drop=drop))
 
 
 def test_to_parquet_dataset_sanitize(path):
@@ -483,10 +495,13 @@ def test_timezone_raw_values(path):
     df["c4"] = pd.to_datetime(datetime(2011, 11, 4, 0, 5, 23, tzinfo=timezone(timedelta(hours=-8))))
     wr.s3.to_parquet(partition_cols=["par"], df=df, path=path, dataset=True, sanitize_columns=False)
     df2 = wr.s3.read_parquet(path, dataset=True, use_threads=False)
-    df3 = pd.read_parquet(path)
+    # Use pandas to read because of Modin "Internal Error: Internal and external indices on axis 1 do not match."
+    import pandas
+
+    df3 = pandas.read_parquet(path)
     df2["par"] = df2["par"].astype("string")
     df3["par"] = df3["par"].astype("string")
-    assert df2.equals(df3)
+    assert pandas_equals(df2, df3)
 
 
 @pytest.mark.parametrize("partition_cols", [None, ["a"], ["a", "b"]])
@@ -505,7 +520,7 @@ def test_empty_column(path, use_threads):
     wr.s3.to_parquet(df, path, dataset=True, partition_cols=["par"])
     df2 = wr.s3.read_parquet(path, dataset=True, use_threads=use_threads)
     df2["par"] = df2["par"].astype("string")
-    assert df.equals(df2)
+    assert pandas_equals(df, df2)
 
 
 def test_mixed_types_column(path) -> None:
@@ -535,7 +550,7 @@ def test_empty_file(path, use_threads):
     boto3.client("s3").put_object(Body=b"", Bucket=bucket, Key=key)
     df2 = wr.s3.read_parquet(path, dataset=True, use_threads=use_threads)
     df2["par"] = df2["par"].astype("string")
-    assert df.equals(df2)
+    assert pandas_equals(df, df2)
 
 
 def test_read_chunked(path):