Dealing with a change request to optimize the queries for ancestor re…

…lationships. This is achieved via an optimized query in the Django QueryManager.

aiida/backends/djsite/queries.py

@@ -270,6 +270,34 @@ def get_bands_and_parents_structure(self, args):
 
         return entry_list
 
+    def get_all_parents(self, node_pks, return_values=['id']):
+        """
+        Get all the parents of given nodes
+        :param node_pks: one node pk or an iterable of node pks
+        :return: a list of aiida objects with all the parents of the nodes
+        """
+        from aiida.backends.djsite.db import models
+        from aiida.common.links import LinkType
+
+        try:
+            the_node_pks = list(node_pks)
+        except TypeError:
+            the_node_pks = [node_pks]
+
+        parents = models.DbNode.objects.none()
+        q_inputs = models.DbNode.aiidaobjects.filter(
+                outputs__pk__in=the_node_pks,
+                output_links__type__in=(LinkType.CREATE.value, LinkType.INPUT.value)).distinct()
+
+        while q_inputs.count() > 0:
+            inputs = list(q_inputs)
+            parents = q_inputs | parents.all()
+            q_inputs = models.DbNode.aiidaobjects.filter(
+                outputs__in=inputs,
+                output_links__type__in=(LinkType.CREATE.value, LinkType.INPUT.value)).distinct()
+
+        return parents.values_list(*return_values)
+
 
 def get_closest_parents(pks,*args,**kwargs):
     """

aiida/backends/general/abstractqueries.py

@@ -298,5 +298,18 @@ def get_bands_and_parents_structure(self, args):
             entry_list.append([str(bid), str(formula),
                                bdate.strftime('%d %b %Y'), blabel])
 
-        print entry_list
         return entry_list
+
+    def get_all_parents(self, node_pks, return_values=['id']):
+        """
+        Get all the parents of given nodes
+        :param node_pks: one node pk or an iterable of node pks
+        :return: a list of aiida objects with all the parents of the nodes
+        """
+        from aiida.orm.querybuilder import QueryBuilder
+        from aiida.orm import Node
+        qb = QueryBuilder()
+        qb.append(Node, tag='low_node',
+                  filters={'id': {'in': node_pks}})
+        qb.append(Node, ancestor_of='low_node', project=return_values)
+        return qb.all()

aiida/backends/tests/query.py

@@ -691,7 +691,9 @@ def test_query_path(self):
 
         from aiida.orm.querybuilder import QueryBuilder
         from aiida.orm import Node
+        from aiida.backends.utils import QueryFactory
 
+        q = QueryFactory()()
         n1 = Node()
         n1.label='n1'
         n1.store()
@@ -734,6 +736,15 @@ def test_query_path(self):
 
 
 
+        # There are no parents to n9, checking that
+        self.assertEqual(set([]), set(q.get_all_parents([n9.pk])))
+        # There is one parent to n6
+        self.assertEqual(set([(_,) for _ in (n6.pk,)]), set([tuple(_) for _ in q.get_all_parents([n7.pk])]))
+        # There are several parents to n4
+        self.assertEqual(set([(_.pk,) for _ in (n1,n2)]), set([tuple(_) for _ in q.get_all_parents([n4.pk])]))
+        # There are several parents to n5
+        self.assertEqual(set([(_.pk,) for _ in (n1,n2,n3,n4)]), set([tuple(_) for _ in q.get_all_parents([n5.pk])]))
+
 
         # Yet, no links from 1 to 8
         self.assertEquals(
@@ -871,7 +882,7 @@ def test_create_node_and_query(self):
         self.assertEqual(idx,99)
         self.assertTrue(len(QueryBuilder().append(Node,project=['id','label']).all(batch_size=10)) > 99)
 
-class TestStatisticsQuery(AiidaTestCase):
+class TestManager(AiidaTestCase):
     def test_statistics(self):
         """
         Test if the statistics query works properly.

aiida/orm/importexport.py

@@ -1643,7 +1643,7 @@ def fill_in_query(partial_query, originating_entity_str, current_entity_str,
 
 def export_tree(what, folder, also_parents=True, also_calc_outputs=True,
                      allowed_licenses=None, forbidden_licenses=None,
-                     silent=False):
+                     silent=False, use_querybuilder_ancestors=False):
     """
     Export the DB entries passed in the 'what' list to a file tree.
 
@@ -1673,7 +1673,7 @@ def export_tree(what, folder, also_parents=True, also_calc_outputs=True,
     from aiida.common.links import LinkType
     from aiida.common.folders import RepositoryFolder
     from aiida.orm.querybuilder import QueryBuilder
-
+    from aiida.backends.utils import QueryFactory
     if not silent:
         print "STARTING EXPORT..."
 
@@ -1702,11 +1702,16 @@ def export_tree(what, folder, also_parents=True, also_calc_outputs=True,
         if given_node_entry_ids:
             # Also add the parents (to any level) to the query
             # This is done via the ancestor relationship.
-            qb = QueryBuilder()
-            qb.append(Node, tag='low_node',
-                      filters={'id': {'in': given_node_entry_ids}})
-            qb.append(Node, ancestor_of='low_node', project=['id'])
-            additional_ids = [_ for _, in qb.all()]
+            if use_querybuilder_ancestors:
+                qb = QueryBuilder()
+                qb.append(Node, tag='low_node',
+                          filters={'id': {'in': given_node_entry_ids}})
+                qb.append(Node, ancestor_of='low_node', project=['id'])
+                additional_ids = [_ for _, in qb.all()]
+            else:
+                q = QueryFactory()()
+                additional_ids = [_ for _, in q.get_all_parents(given_node_entry_ids, return_values=['id'])]
+
             given_node_entry_ids = given_node_entry_ids.union(additional_ids)
 
     if also_calc_outputs:

docs/source/installation/index.rst

@@ -746,13 +746,13 @@ Updating from 0.9.* Django to 0.10.0 Django
 
         > psql AIIDADB
         AIIDADB=# CREATE TABLE db_dbpath (
-            id integer NOT NULL,
-            depth integer NOT NULL,
-            entry_edge_id integer,
-            direct_edge_id integer,
-            exit_edge_id integer,
-            child_id integer NOT NULL,
-            parent_id integer NOT NULL
+        id integer NOT NULL,
+        depth integer NOT NULL,
+        entry_edge_id integer,
+        direct_edge_id integer,
+        exit_edge_id integer,
+        child_id integer NOT NULL,
+        parent_id integer NOT NULL
         );