Repository: dbt-labs/dbt-core

Issue:
[CT-2377] [Spike] Performance improvement to memory footprint
When a macro runs a call statement with `fetch_result` set to `True`, that result may be held in memory longer than it needs to be.


Discussion:
@nathaniel-may Thanks for opening! (I'd been meaning to)

More context in [this internal Slack thread](https://dbt-labs.slack.com/archives/C0131TY7EEA/p1678400613373389)

As a general rule, should we be garbage-collecting the `context` from each NodeRunner, after that node finishes compiling/executing?

Resolving patch:
diff --git a/core/dbt/context/providers.py b/core/dbt/context/providers.py
--- a/core/dbt/context/providers.py
+++ b/core/dbt/context/providers.py
@@ -54,6 +54,7 @@
     LoadAgateTableNotSeedError,
     LoadAgateTableValueError,
     MacroDispatchArgError,
+    MacroResultAlreadyLoadedError,
     MacrosSourcesUnWriteableError,
     MetricArgsError,
     MissingConfigError,
@@ -704,7 +705,7 @@ def __init__(
         self.config: RuntimeConfig
         self.model: Union[Macro, ManifestNode] = model
         super().__init__(config, manifest, model.package_name)
-        self.sql_results: Dict[str, AttrDict] = {}
+        self.sql_results: Dict[str, Optional[AttrDict]] = {}
         self.context_config: Optional[ContextConfig] = context_config
         self.provider: Provider = provider
         self.adapter = get_adapter(self.config)
@@ -732,12 +733,29 @@ def invocation_args_dict(self):
         return args_to_dict(self.config.args)
 
     @contextproperty
-    def _sql_results(self) -> Dict[str, AttrDict]:
+    def _sql_results(self) -> Dict[str, Optional[AttrDict]]:
         return self.sql_results
 
     @contextmember
     def load_result(self, name: str) -> Optional[AttrDict]:
-        return self.sql_results.get(name)
+        if name in self.sql_results:
+            # handle the special case of "main" macro
+            # See: https://github.com/dbt-labs/dbt-core/blob/ada8860e48b32ac712d92e8b0977b2c3c9749981/core/dbt/task/run.py#L228
+            if name == "main":
+                return self.sql_results["main"]
+
+            # handle a None, which indicates this name was populated but has since been loaded
+            elif self.sql_results[name] is None:
+                raise MacroResultAlreadyLoadedError(name)
+
+            # Handle the regular use case
+            else:
+                ret_val = self.sql_results[name]
+                self.sql_results[name] = None
+                return ret_val
+        else:
+            # Handle trying to load a result that was never stored
+            return None
 
     @contextmember
     def store_result(
diff --git a/core/dbt/exceptions.py b/core/dbt/exceptions.py
--- a/core/dbt/exceptions.py
+++ b/core/dbt/exceptions.py
@@ -1019,6 +1019,17 @@ def get_message(self) -> str:
         return msg
 
 
+class MacroResultAlreadyLoadedError(CompilationError):
+    def __init__(self, result_name):
+        self.result_name = result_name
+        super().__init__(msg=self.get_message())
+
+    def get_message(self) -> str:
+        msg = f"The 'statement' result named '{self.result_name}' has already been loaded into a variable"
+
+        return msg
+
+
 # parser level exceptions
 class DictParseError(ParsingError):
     def __init__(self, exc: ValidationError, node):


Test patch:
diff --git a/tests/functional/statements/fixtures.py b/tests/functional/statements/fixtures.py
--- a/tests/functional/statements/fixtures.py
+++ b/tests/functional/statements/fixtures.py
@@ -136,3 +136,47 @@
 union all
 select 'table' as source, {{ table_value }} as value
 """
+
+models__statement_duplicated_load = """
+-- {{ ref('seed') }}
+
+{%- call statement('test_statement', fetch_result=True) -%}
+
+  select
+    count(*) as "num_records"
+
+  from {{ ref('seed') }}
+
+{%- endcall -%}
+
+{% set result = load_result('test_statement') %}
+{% set result = load_result('test_statement') %}
+
+select 1
+"""
+
+models__statement_load_main_twice = """
+-- {{ ref('seed') }}
+
+{%- call statement('main', fetch_result=True) -%}
+
+  select
+    count(*) as "num_records"
+
+  from {{ ref('seed') }}
+
+{%- endcall -%}
+
+{% set result = load_result('main') %}
+{% set result = load_result('main') %}
+
+{% set res_table = result['table'] %}
+{% set res_matrix = result['data'] %}
+
+{% set matrix_value = res_matrix[0][0] %}
+{% set table_value = res_table[0]['num_records'] %}
+
+select 'matrix' as source, {{ matrix_value }} as value
+union all
+select 'table' as source, {{ table_value }} as value
+"""
diff --git a/tests/functional/statements/test_statements.py b/tests/functional/statements/test_statements.py
--- a/tests/functional/statements/test_statements.py
+++ b/tests/functional/statements/test_statements.py
@@ -4,6 +4,8 @@
 from dbt.tests.util import run_dbt, check_relations_equal, write_file
 from tests.functional.statements.fixtures import (
     models__statement_actual,
+    models__statement_duplicated_load,
+    models__statement_load_main_twice,
     seeds__statement_actual,
     seeds__statement_expected,
 )
@@ -21,7 +23,11 @@ def setUp(self, project):
 
     @pytest.fixture(scope="class")
     def models(self):
-        return {"statement_actual.sql": models__statement_actual}
+        return {
+            "statement_actual.sql": models__statement_actual,
+            "statement_duplicated_load.sql": models__statement_duplicated_load,
+            "statement_load_main_twice.sql": models__statement_load_main_twice,
+        }
 
     @pytest.fixture(scope="class")
     def project_config_update(self):
@@ -35,7 +41,23 @@ def project_config_update(self):
     def test_postgres_statements(self, project):
         results = run_dbt(["seed"])
         assert len(results) == 2
-        results = run_dbt()
+        results = run_dbt(["run", "-m", "statement_actual"])
         assert len(results) == 1
 
         check_relations_equal(project.adapter, ["statement_actual", "statement_expected"])
+
+    def test_duplicated_load_statements(self, project):
+        run_dbt(["seed"])
+        results = run_dbt(["run", "-m", "statement_duplicated_load"], False)
+        assert len(results) == 1
+        assert results.results[0].status == "error"
+        assert (
+            "The 'statement' result named 'test_statement' has already been loaded into a variable"
+            in results.results[0].message
+        )
+
+    def test_load_statement_on_main_twice(self, project):
+        run_dbt(["seed"])
+        results = run_dbt(["run", "-m", "statement_load_main_twice"])
+        assert len(results) == 1
+        check_relations_equal(project.adapter, ["statement_load_main_twice", "statement_expected"])