Merge pull request #620 from DagsHub/bug/duplicate-tree-composition

Bug: composing duplicate queries throws an error

Author: kbolashev

dagshub/data_engine/model/query.py

@@ -1,6 +1,7 @@
 import datetime
 import enum
 import logging
+import uuid
 from typing import Optional, Union, Dict
 
 import pytz
@@ -190,10 +191,26 @@ def compose(self, op: str, other: Optional[Union[str, int, float, "QueryFilterTr
                 return
             composite_tree = Tree()
             root_node = composite_tree.create_node(op)
-            composite_tree.paste(root_node.identifier, self._operand_tree)
-            composite_tree.paste(root_node.identifier, other._operand_tree)
+
+            # Use deep clones to avoid potentially composing two of the same tree together
+            # and hitting duplicate identifiers
+            composite_tree.paste(root_node.identifier, self._clone_operand_tree())
+            composite_tree.paste(root_node.identifier, other._clone_operand_tree())
             self._operand_tree = composite_tree
 
+    def _clone_operand_tree(self) -> Tree:
+        """
+        Returns a deep copy of the operand tree, also changing all node identifiers.
+        This allows to compose two of the same tree together.
+        """
+        if self._operand_tree.root is None:
+            return Tree()
+        new_tree = Tree(self._operand_tree.subtree(self._operand_tree.root), deep=True)
+        node_ids = [node.identifier for node in new_tree.all_nodes_itr()]
+        for node_id in node_ids:
+            new_tree.update_node(node_id, identifier=str(uuid.uuid4()))
+        return new_tree
+
     @property
     def _column_filter_node(self) -> Node:
         return next(self._operand_tree.filter_nodes(lambda n: n.tag == UNFILLED_NODE_TAG), None)

tests/data_engine/test_querying.py

@@ -685,3 +685,40 @@ def test_datetime_is_null(ds):
     expected = {"query": {"filter": {"comparator": "IS_NULL", "key": "x", "value": "0", "valueType": "DATETIME"}}}
 
     assert q.serialize_gql_query_input() == expected
+
+
+def test_duplicate_subquery(ds):
+    add_int_fields(ds, "x")
+    add_int_fields(ds, "y")
+    ds2 = ds["x"] > 5
+
+    ds3 = (ds2["y"] < 10) | (ds2["y"] > 20)
+
+    expected = {
+        "or": {
+            "children": [
+                {
+                    "and": {
+                        "children": [
+                            {"gt": {"data": {"field": "x", "value": 5}}},
+                            {"lt": {"data": {"field": "y", "value": 10}}},
+                        ],
+                        "data": None,
+                    }
+                },
+                {
+                    "and": {
+                        "children": [
+                            {"gt": {"data": {"field": "x", "value": 5}}},
+                            {"gt": {"data": {"field": "y", "value": 20}}},
+                        ],
+                        "data": None,
+                    }
+                },
+            ],
+            "data": None,
+        }
+    }
+
+    actual = ds3.get_query().filter.tree_to_dict()
+    assert actual == expected