refactor: DAG constructor to not rely on pandas + handle reserved keywords cases + add test cases

kayjan · kayjan · commit c19d83c129c9 · 2024-04-13T00:00:48.000+08:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,11 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+### Changed:
+- DAG Constructor: `list_to_dag` and `dict_to_dag` does not rely on `dataframe_to_dag` as pandas dataframe operation
+is phased out.
+### Fixed:
+- DAG Constructor: Handle cases where reserved keywords are part of attribute upon creation and throw error accordingly.
 
 ## [0.17.0] - 2024-04-04
 ### Added
diff --git a/bigtree/dag/construct.py b/bigtree/dag/construct.py
@@ -6,6 +6,7 @@
 from bigtree.utils.assertions import (
     assert_dataframe_no_duplicate_attribute,
     assert_dataframe_not_empty,
+    assert_key_not_in_dict_or_df,
     assert_length_not_empty,
     filter_attributes,
     isnull,
@@ -20,7 +21,6 @@
 __all__ = ["list_to_dag", "dict_to_dag", "dataframe_to_dag"]
 
 
-@optional_dependencies_pandas
 def list_to_dag(
     relations: List[Tuple[str, str]],
     node_type: Type[DAGNode] = DAGNode,
@@ -44,13 +44,26 @@ def list_to_dag(
     """
     assert_length_not_empty(relations, "Input list", "relations")
 
-    relation_data = pd.DataFrame(relations, columns=["parent", "child"])
-    return dataframe_to_dag(
-        relation_data, child_col="child", parent_col="parent", node_type=node_type
-    )
+    node_dict: Dict[str, DAGNode] = dict()
+    parent_node = DAGNode()
+
+    for parent_name, child_name in relations:
+        if parent_name not in node_dict:
+            parent_node = node_type(parent_name)
+            node_dict[parent_name] = parent_node
+        else:
+            parent_node = node_dict[parent_name]
+        if child_name not in node_dict:
+            child_node = node_type(child_name)
+            node_dict[child_name] = child_node
+        else:
+            child_node = node_dict[child_name]
+
+        child_node.parents = [parent_node]
+
+    return parent_node
 
 
-@optional_dependencies_pandas
 def dict_to_dag(
     relation_attrs: Dict[str, Any],
     parent_key: str = "parents",
@@ -84,20 +97,34 @@ def dict_to_dag(
     """
     assert_length_not_empty(relation_attrs, "Dictionary", "relation_attrs")
 
-    # Convert dictionary to dataframe
-    data = pd.DataFrame(relation_attrs).T.rename_axis("_tmp_child").reset_index()
-    if parent_key not in data:
+    node_dict: Dict[str, DAGNode] = dict()
+    parent_node: DAGNode | None = None
+
+    for child_name, node_attrs in relation_attrs.items():
+        node_attrs = node_attrs.copy()
+        parent_names: List[str] = []
+        if parent_key in node_attrs:
+            parent_names = node_attrs.pop(parent_key)
+        assert_key_not_in_dict_or_df(node_attrs, ["parent", "parents", "children"])
+
+        if child_name in node_dict:
+            child_node = node_dict[child_name]
+            child_node.set_attrs(node_attrs)
+        else:
+            child_node = node_type(child_name, **node_attrs)
+            node_dict[child_name] = child_node
+
+        for parent_name in parent_names:
+            parent_node = node_dict.get(parent_name, node_type(parent_name))
+            node_dict[parent_name] = parent_node
+            child_node.parents = [parent_node]
+
+    if parent_node is None:
         raise ValueError(
             f"Parent key {parent_key} not in dictionary, check `relation_attrs` and `parent_key`"
         )
 
-    data = data.explode(parent_key)
-    return dataframe_to_dag(
-        data,
-        child_col="_tmp_child",
-        parent_col=parent_key,
-        node_type=node_type,
-    )
+    return parent_node
 
 
 @optional_dependencies_pandas
@@ -163,6 +190,7 @@ def dataframe_to_dag(
         attribute_cols = list(data.columns)
         attribute_cols.remove(child_col)
         attribute_cols.remove(parent_col)
+    assert_key_not_in_dict_or_df(attribute_cols, ["parent", "parents", "children"])
 
     data = data[[child_col, parent_col] + attribute_cols].copy()
 
diff --git a/bigtree/utils/assertions.py b/bigtree/utils/assertions.py
@@ -58,6 +58,23 @@ def assert_str_in_list(
         )
 
 
+def assert_key_not_in_dict_or_df(
+    parameter_dict: Union[Dict[str, Any], pd.DataFrame],
+    not_accepted_parameters: List[str],
+) -> None:
+    """Raise ValueError is parameter is in key of dictionary
+
+    Args:
+        parameter_dict (Dict[str, Any]/pd.DataFrame): argument input for parameter
+        not_accepted_parameters (List[str]): list of not accepted parameters
+    """
+    for parameter in parameter_dict:
+        if parameter in not_accepted_parameters:
+            raise ValueError(
+                f"Invalid input, check `{parameter}` is not a valid key as it is a reserved keyword"
+            )
+
+
 def assert_key_in_dict(
     parameter_name: str,
     parameter: Any,
diff --git a/tests/dag/test_construct.py b/tests/dag/test_construct.py
@@ -80,10 +80,48 @@ def test_dict_to_dag_empty_error(self):
             parameter="relation_attrs"
         )
 
-    def test_dict_to_dag_parent_key_error(self):
+    @staticmethod
+    def test_dict_to_dag_parent_key_error():
+        relation_dict = {
+            "a": {"age": 90},
+            "b": {"age": 65},
+            "c": {"parent1": ["a", "b"], "age": 60},
+            "d": {"parent1": ["a", "c"], "age": 40},
+            "e": {"parent1": ["d"], "age": 35},
+            "f": {"parent1": ["c", "d"], "age": 38},
+            "g": {"parent1": ["c"], "age": 10},
+            "h": {"parent1": ["g"], "age": 6},
+        }
+        with pytest.raises(ValueError) as exc_info:
+            dict_to_dag(relation_dict)
+        assert str(exc_info.value) == Constants.ERROR_DAG_DICT_PARENT_KEY.format(
+            parent_key="parents"
+        )
+
+    def test_dict_to_dag_parent_key_reserved_keyword_parents_error(self):
         with pytest.raises(ValueError) as exc_info:
             dict_to_dag(self.relation_dict, parent_key="parent")
-        assert str(exc_info.value) == Constants.ERROR_DAG_DICT_PARENT_KEY
+        assert str(exc_info.value) == Constants.ERROR_DAG_DICT_INVALID_KEY.format(
+            parameter="parents"
+        )
+
+    @staticmethod
+    def test_dict_to_dag_parent_key_reserved_keyword_parent_error():
+        relation_dict = {
+            "a": {"age": 90},
+            "b": {"age": 65},
+            "c": {"parent": ["a", "b"], "age": 60},
+            "d": {"parent": ["a", "c"], "age": 40},
+            "e": {"parent": ["d"], "age": 35},
+            "f": {"parent": ["c", "d"], "age": 38},
+            "g": {"parent": ["c"], "age": 10},
+            "h": {"parent": ["g"], "age": 6},
+        }
+        with pytest.raises(ValueError) as exc_info:
+            dict_to_dag(relation_dict)
+        assert str(exc_info.value) == Constants.ERROR_DAG_DICT_INVALID_KEY.format(
+            parameter="parent"
+        )
 
     def test_dict_to_dag_node_type(self):
         dag = dict_to_dag(self.relation_dict, node_type=DAGNodeA)
@@ -178,6 +216,54 @@ def test_dataframe_to_dag_parent_col_error(self):
             parent_col=parent_col
         )
 
+    @staticmethod
+    def test_dataframe_to_dag_parent_col_reserved_keyword_parents_error():
+        data = pd.DataFrame(
+            [
+                ["h", "g", "a", 6],
+                ["g", "c", "a", 10],
+                ["f", "d", "a", 38],
+                ["f", "c", "a", 38],
+                ["e", "d", "a", 35],
+                ["d", "c", "a", 40],
+                ["d", "a", "a", 40],
+                ["c", "b", "a", 60],
+                ["c", "a", "a", 60],
+                ["a", None, None, 90],
+                ["b", None, None, 65],
+            ],
+            columns=["child", "parent", "parents", "age"],
+        )
+        with pytest.raises(ValueError) as exc_info:
+            dataframe_to_dag(data, parent_col="parent")
+        assert str(exc_info.value) == Constants.ERROR_DAG_DICT_INVALID_KEY.format(
+            parameter="parents"
+        )
+
+    @staticmethod
+    def test_dataframe_to_dag_parent_col_reserved_keyword_parent_error():
+        data = pd.DataFrame(
+            [
+                ["h", "g", "a", 6],
+                ["g", "c", "a", 10],
+                ["f", "d", "a", 38],
+                ["f", "c", "a", 38],
+                ["e", "d", "a", 35],
+                ["d", "c", "a", 40],
+                ["d", "a", "a", 40],
+                ["c", "b", "a", 60],
+                ["c", "a", "a", 60],
+                ["a", None, None, 90],
+                ["b", None, None, 65],
+            ],
+            columns=["child", "parent", "parents", "age"],
+        )
+        with pytest.raises(ValueError) as exc_info:
+            dataframe_to_dag(data, parent_col="parents")
+        assert str(exc_info.value) == Constants.ERROR_DAG_DICT_INVALID_KEY.format(
+            parameter="parent"
+        )
+
     def test_dataframe_to_dag_attribute_cols(self):
         dag = dataframe_to_dag(self.data, attribute_cols=["age"])
         assert_dag_structure_root(dag)
diff --git a/tests/test_constants.py b/tests/test_constants.py
@@ -20,9 +20,8 @@ class Constants:
     ERROR_CUSTOM_TYPE = "Node type is not `{type}`"
 
     # dag/construct
-    ERROR_DAG_DICT_PARENT_KEY = (
-        "Parent key parent not in dictionary, check `relation_attrs` and `parent_key`"
-    )
+    ERROR_DAG_DICT_INVALID_KEY = "Invalid input, check `{parameter}` is not a valid key as it is a reserved keyword"
+    ERROR_DAG_DICT_PARENT_KEY = "Parent key {parent_key} not in dictionary, check `relation_attrs` and `parent_key`"
     ERROR_DAG_DATAFRAME_EMPTY_CHILD = (
         "Child name cannot be empty, check column: {child_col}"
     )
