Merge pull request #1499 from datajoint/fix/1493-lazy-upstream

fix(#1493): build self.upstream lazily; stop reloading deps per populate key

Author: dimitri-yatsenko

src/datajoint/autopopulate.py

@@ -98,7 +98,8 @@ class AutoPopulate:
     _key_source = None
     _allow_insert = False
     _jobs = None
-    _upstream = None  # set per-make() by _populate_one; see `upstream` property below
+    _upstream = None  # memoized upstream trace for the current make(); built lazily on first use
+    _upstream_key = None  # current make() key; not None iff executing inside make()
 
     @property
     def upstream(self):
@@ -132,12 +133,19 @@ def make(self, key):
                 traces = self.upstream[ExtractTraces].to_arrays("trace")
                 self.insert1({**key, "summary": compute(traces, date)})
         """
-        if self._upstream is None:
+        if self._upstream_key is None:
             raise DataJointError(
                 "self.upstream is only available inside make(). "
                 "Outside make(), construct a trace explicitly: "
                 "dj.Diagram.trace(self & key)."
             )
+        if self._upstream is None:
+            # Build the trace lazily on first access and memoize it for this
+            # make() call, so make() calls that never read self.upstream pay
+            # nothing. See #1493.
+            from .diagram import Diagram
+
+            self._upstream = Diagram.trace(self & self._upstream_key)
         return self._upstream
 
     class _JobsDescriptor:
@@ -665,12 +673,11 @@ def _populate1(
         logger.jobs(f"Making {key} -> {self.full_table_name}")
         self.__class__._allow_insert = True
 
-        # Pre-construct the upstream view for this make() call. Lazy — only
-        # `dj.Diagram.trace(self & key)` runs here (graph copy); the
-        # expensive SQL fetch fires when the user accesses self.upstream[T].
-        from .diagram import Diagram
-
-        self._upstream = Diagram.trace(self & dict(key))
+        # Record the key for a lazily-constructed upstream view. The trace is
+        # built on first `self.upstream` access (see the `upstream` property),
+        # so make() calls that never read upstream pay nothing. See #1493.
+        self._upstream_key = dict(key)
+        self._upstream = None
 
         try:
             if not is_generator:
@@ -733,6 +740,7 @@ def _populate1(
             # access raises a clear error rather than silently using a
             # stale trace from the previous make() call.
             self._upstream = None
+            self._upstream_key = None
 
     def progress(self, *restrictions: Any, display: bool = False) -> tuple[int, int]:
         """

src/datajoint/dependencies.py

@@ -132,11 +132,13 @@ def __init__(self, connection=None) -> None:
         self._conn = connection
         self._node_alias_count = itertools.count()
         self._loaded = False
+        self._loaded_schemas = set()  # schema names currently represented in the graph
         super().__init__(self)
 
     def clear(self) -> None:
         """Clear the graph and reset loaded state."""
         self._loaded = False
+        self._loaded_schemas = set()
         self._node_alias_count = itertools.count()  # reset alias IDs for consistency
         super().clear()
 
@@ -165,6 +167,7 @@ def load(self, force: bool = True, schema_names: set[str] | None = None) -> None
 
         # Build schema list for IN clause
         names = schema_names if schema_names is not None else set(self._conn.schemas)
+        self._loaded_schemas = set(names)
         if not names:
             self._loaded = True
             return
@@ -257,6 +260,11 @@ def load_all_downstream(self) -> None:
                 break
             known_schemas |= new_schemas
 
+        # Skip the expensive rebuild when the graph already contains every needed
+        # schema, so repeated calls within one operation don't reload the whole
+        # dependency tree. See #1493.
+        if self._loaded and known_schemas <= self._loaded_schemas:
+            return
         self.load(force=True, schema_names=known_schemas)
 
     def load_all_upstream(self) -> None:
@@ -286,6 +294,11 @@ def load_all_upstream(self) -> None:
                 break
             known_schemas |= new_schemas
 
+        # Skip the expensive rebuild when the graph already contains every needed
+        # schema, so repeated Diagram.trace() calls within one populate() don't
+        # reload the whole dependency tree per key. See #1493.
+        if self._loaded and known_schemas <= self._loaded_schemas:
+            return
         self.load(force=True, schema_names=known_schemas)
 
     def topo_sort(self) -> list[str]:

tests/integration/test_autopopulate.py

@@ -492,14 +492,18 @@ class Derived(dj.Computed):
 
         def make(self, key):
             captured.append(self)  # the actual populate instance
-            assert self._upstream is not None  # set for the duration of make()
+            # Lazy: the trace is NOT built here because this make() never reads
+            # self.upstream — but we ARE inside make() (the key is recorded).
+            assert self._upstream is None
+            assert self._upstream_key is not None
             self.insert1({**key, "val": 0})
 
     Derived.populate()
     assert captured, "make() did not run"
     inst = captured[0]
-    # The finally block must have cleared _upstream on this very instance.
-    assert inst._upstream is None
+    # The finally block must have cleared the make() key on this very instance —
+    # that is what makes self.upstream raise afterward.
+    assert inst._upstream_key is None
     with pytest.raises(DataJointError, match="only available inside make"):
         inst.upstream
 
@@ -529,15 +533,15 @@ class Boom(dj.Computed):
 
         def make(self, key):
             captured.append(self)
-            assert self._upstream is not None
+            assert self._upstream_key is not None  # inside make()
             raise RuntimeError("make failed on purpose")
 
     with pytest.raises(RuntimeError, match="make failed on purpose"):
         Boom.populate(suppress_errors=False)
     assert captured, "make() did not run"
     inst = captured[0]
     # Cleared by the finally block even though make() raised.
-    assert inst._upstream is None
+    assert inst._upstream_key is None
     with pytest.raises(DataJointError, match="only available inside make"):
         inst.upstream
 
@@ -557,7 +561,7 @@ class Source(dj.Lookup):
         """
         contents = [(1, 100), (2, 200)]
 
-    seen = []  # (source_id, phase, id(self._upstream))
+    seen = []  # (source_id, phase, id(self.upstream))
 
     @schema
     class TriComputed(dj.Computed):
@@ -568,15 +572,15 @@ class TriComputed(dj.Computed):
         """
 
         def make_fetch(self, key):
-            seen.append((key["source_id"], "fetch", id(self._upstream)))
+            seen.append((key["source_id"], "fetch", id(self.upstream)))
             return (self.upstream[Source].fetch1("value"),)
 
         def make_compute(self, key, value):
-            seen.append((key["source_id"], "compute", id(self._upstream)))
+            seen.append((key["source_id"], "compute", id(self.upstream)))
             return (value * 2,)
 
         def make_insert(self, key, doubled):
-            seen.append((key["source_id"], "insert", id(self._upstream)))
+            seen.append((key["source_id"], "insert", id(self.upstream)))
             self.insert1({**key, "result": doubled})
 
     TriComputed.populate()