pyspark-connectby 1.1.0__tar.gz → 1.1.1__tar.gz

{pyspark_connectby-1.1.0 → pyspark_connectby-1.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pyspark-connectby
-Version: 1.1.0
+Version: 1.1.1
 Summary: connectby hierarchy query in spark
 Author: Chen, Yu
 Requires-Python: >=3.7,<4.0

{pyspark_connectby-1.1.0 → pyspark_connectby-1.1.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pyspark-connectby"
-version = "1.1.0"
+version = "1.1.1"
 description = "connectby hierarchy query in spark"
 authors = ["Chen, Yu"]
 readme = "README.md"

{pyspark_connectby-1.1.0 → pyspark_connectby-1.1.1}/pyspark_connectby/connectby_query.py

@@ -1,7 +1,3 @@
-__author__ = 'Chen, Yu'
-__date__ = '2024-02'
-__version__ = '0.8'
-
 from dataclasses import dataclass
 from typing import Union, List
 
@@ -30,7 +26,7 @@ class Path:
         return self.steps[-1]
 
     @property
-    def level(self) -> str:
+    def level(self) -> int:
         return len(self.steps)
 
 
@@ -79,7 +75,7 @@ class ConnectByQuery:
         rows = self.df.collect()
         return [Path.path_start_with(r[self.child_col]) for r in rows]
 
-    def __descendants_recursive(self, path: Path) -> []:
+    def __fetch_descendants(self, path: Path) -> []:
         children_nodes: [Node] = self.__children_with_parent(path.end_id)
         is_leaf = len(children_nodes) == 0
         if is_leaf:
@@ -87,30 +83,29 @@ class ConnectByQuery:
             return []
 
         children = [Path(steps=path.steps + [c.node_id]) for c in children_nodes]
-        grandchildren = list(map(lambda c: self.__descendants_recursive(c), children))
+        grandchildren = list(map(lambda c: self.__fetch_descendants(c), children))
 
         descendants = children + grandchildren
         return descendants
 
     @staticmethod
-    def __flatten(nested_list):
+    def __flatten_list(nested_list: []) -> []:
         flat_list = []
         for item in nested_list:
             if isinstance(item, list):
-                flat_list += ConnectByQuery.__flatten(item)
+                flat_list += ConnectByQuery.__flatten_list(item)
             else:
                 flat_list.append(item)
         return flat_list
 
     def __run(self) -> [Path]:
-        descendants = list(map(lambda e: self.__descendants_recursive(e), self.start_paths))
-        descendants_paths: [Path] = self.__flatten(descendants)
+        descendants = list(map(lambda e: self.__fetch_descendants(e), self.start_paths))
+        descendants_paths: [Path] = self.__flatten_list(descendants)
 
         return self.start_paths + descendants_paths
 
     def get_result_df(self) -> DataFrame:
         result_paths: [Path] = self.__run()
-
         schema = f'''
             {COLUMN_START_WITH} string, 
             {self.child_col} string, 
@@ -118,7 +113,7 @@ class ConnectByQuery:
             {COLUMN_CONNECT_BY_ISLEAF} boolean
         '''
         spark = self.df._session
-        result_df = \
-            spark.createDataFrame([(p.start_id, p.end_id, p.level, p.is_leaf) for p in result_paths], schema=schema)
 
+        result_df = spark.createDataFrame([(p.start_id, p.end_id, p.level, p.is_leaf) for p in result_paths],
+                                          schema=schema)
         return result_df.join(self.df, on=self.child_col)