Merge pull request #318 from dclong/dev

Merge dev into main

Author: dclong

dsutil/__init__.py

@@ -3,4 +3,4 @@
 from . import git
 from . import poetry
 
-__version__ = "0.68.3"
+__version__ = "0.69.0"

dsutil/hadoop/repart_hdfs.py

@@ -0,0 +1,49 @@
+"""Repartition a HDFS path of the Parquet format. 
+"""
+from typing import Optional
+from argparse import ArgumentParser, Namespace
+from pyspark.sql import SparkSession
+from .utils import repart_hdfs
+
+spark = SparkSession.builder.appName("Repart_HDFS").enableHiveSupport().getOrCreate()
+
+
+def parse_args(args=None, namespace=None) -> Namespace:
+    """Parse command-line arguments.
+
+    :param args: The arguments to parse.
+        If None, the arguments from command-line are parsed.
+    :param namespace: An inital Namespace object.
+    :return: A namespace object containing parsed options.
+    """
+    parser = ArgumentParser(
+        description="Repartition a HDFS path which is of the Parquet format."
+    )
+    parser.add_argument(
+        "-p",
+        "--path",
+        "--hdfs-path",
+        dest="hdfs_path",
+        type=str,
+        help="The HDFS path (of the Parquet format) to repartition."
+    )
+    parser.add_argument(
+        "-n",
+        "--num-parts",
+        dest="num_parts",
+        type=int,
+        help="The new number of partitions."
+    )
+    return parser.parse_args(args=args, namespace=namespace)
+
+
+def main(args: Optional[Namespace] = None):
+    """The main function for script usage.
+    """
+    if args is None:
+        args = parse_args()
+    repart_hdfs(spark, path=args.hdfs_path, num_parts=args.num_parts)
+
+
+if __name__ == "__main__":
+    main()

dsutil/hadoop/utils.py

@@ -2,6 +2,8 @@
 """
 from __future__ import annotations
 from typing import Union
+import sys
+import datetime
 from pyspark.sql import DataFrame, Window
 from pyspark.sql.functions import col, spark_partition_id, rank, coalesce, lit, max, sum
 
@@ -71,3 +73,27 @@ def calc_global_rank(frame: DataFrame, order_by: Union[str, list[str]]) -> DataF
         ["part_id"],
     ).withColumn("rank",
                  col("local_rank") + col("sum_factor"))
+
+
+def repart_hdfs(spark, path: str, num_parts: int) -> None:
+    """Repartition a HDFS path of the Parquet format.
+
+    :param spark: A SparkSession object. 
+    :param path: The HDFS path to repartition. 
+    :param num_parts: The new number of partitions. 
+    """
+    path = path.rstrip("/")
+    ts = datetime.datetime.now().strftime("%Y%m%d%H%M%S%f")
+    path_tmp = path + f"_repart_tmp_{ts}"
+    spark.read.parquet(path).repartition(num_parts) \
+        .write.mode("overwrite").parquet(path_tmp)
+    sc = spark.sparkContext
+    fs = sc._jvm.org.apache.hadoop.fs.FileSystem.get(sc._jsc.hadoopConfiguration())  # pylint: disable=W0212
+    if fs.delete(sc._jvm.org.apache.hadoop.fs.Path(path), True):  # pylint: disable=W0212
+        if not fs.rename(
+            sc._jvm.org.apache.hadoop.fs.Path(path_tmp),  # pylint: disable=W0212
+            sc._jvm.org.apache.hadoop.fs.Path(path),  # pylint: disable=W0212
+        ):
+            sys.exit(f"Failed to rename the HDFS path {path_tmp} to {path}!")
+    else:
+        sys.exit(f"Failed to remove the (old) HDFS path: {path}!")

pyproject.toml

@@ -1,11 +1,12 @@
 [tool.poetry]
 name = "dsutil"
-version = "0.68.3"
+version = "0.69.0"
 description = "A utils Python package for data scientists."
 authors = ["Benjamin Du <[email protected]>"]
 
 [tool.poetry.scripts]
 logf = "dsutil.hadoop:logf.main"
+repart_hdfs = "dsutil.hadoop:repart_hdfs.main"
 pyspark_submit = "dsutil.hadoop:pyspark_submit.main"
 pykinit = "dsutil.hadoop:kerberos.main"
 match_memory = "dsutil:memory.main"

readme.md

@@ -25,7 +25,7 @@ Currently, Python 3.7 and 3.8 are supported.
 
 You can download a copy of the latest release and install it using pip.
 ```bash
-pip3 install --user -U https://github.com/dclong/dsutil/releases/download/v0.68.3/dsutil-0.68.3-py3-none-any.whl
+pip3 install --user -U https://github.com/dclong/dsutil/releases/download/v0.69.0/dsutil-0.69.0-py3-none-any.whl
 ```
 Or you can use the following command to install the latest master branch
 if you have pip 20.0+.
@@ -35,7 +35,7 @@ pip3 install --user -U git+https://github.com/dclong/dsutil@main
 Use one of the following commands if you want to install all components of dsutil. 
 Available additional components are `cv`, `docker`, `pdf`, `jupyter`, `admin` and `all`.
 ```bash
-pip3 install "dsutil[cv] @ https://github.com/dclong/dsutil/releases/download/v0.68.3/dsutil-0.68.3-py3-none-any.whl"
+pip3 install "dsutil[cv] @ https://github.com/dclong/dsutil/releases/download/v0.69.0/dsutil-0.69.0-py3-none-any.whl"
 # or
 pip3 install --user -U "dsutil[all] @ git+https://github.com/dclong/dsutil@main"
 ```