Copy source before modifying it in GroupBy

api/py/ai/chronon/group_by.py

@@ -12,6 +12,7 @@
 #     See the License for the specific language governing permissions and
 #     limitations under the License.
 
+import copy
 import inspect
 import json
 import logging
@@ -556,7 +557,7 @@ def _normalize_source(source):
 
     if not isinstance(sources, list):
         sources = [sources]
-    sources = [_sanitize_columns(_normalize_source(source)) for source in sources]
+    sources = [_sanitize_columns(_normalize_source(_copy_source(source))) for source in sources]
 
     deps = [dep for src in sources for dep in utils.get_dependencies(src, dependencies, lag=lag)]
 
@@ -599,3 +600,21 @@ def _normalize_source(source):
     )
     validate_group_by(group_by)
     return group_by
+
+
+def _copy_source(source: ttypes.Source) -> ttypes.Source:
+    # Hold a reference to the join in a join source so that the
+    # module name can be extracted from GC referrers later on
+    join = None
+    if isinstance(source, ttypes.JoinSource):
+        join = source.join
+    elif isinstance(source, ttypes.Source) and source.joinSource:
+        join = source.joinSource.join
+
+    source_copy = copy.deepcopy(source)
+
+    if isinstance(source_copy, ttypes.JoinSource):
+        source_copy.join = join
+    elif isinstance(source_copy, ttypes.Source) and source_copy.joinSource:
+        source_copy.joinSource.join = join
+    return source_copy