Join failed to get groupBy module on different file structure

api/py/ai/chronon/join.py

@@ -63,7 +63,7 @@ def JoinPart(
     gc.collect()
     group_by_module_name = None
     for ref in gc.get_referrers(group_by):
-        if "__name__" in ref and ref["__name__"].startswith("group_bys"):
+        if isinstance(ref, dict) and "__name__" in ref and "group_bys." in ref["__name__"]:
             group_by_module_name = ref["__name__"]
             break
     if group_by_module_name:

api/py/test/sample/joins/quickstart/training_set.py

@@ -17,9 +17,9 @@
 from ai.chronon.api.ttypes import Source, EventSource
 from ai.chronon.query import Query, select
 
-from group_bys.quickstart.purchases import v1 as purchases_v1
-from group_bys.quickstart.returns import v1 as returns_v1
-from group_bys.quickstart.users import v1 as users
+from sample.group_bys.quickstart.purchases import v1 as purchases_v1
+from sample.group_bys.quickstart.returns import v1 as returns_v1
+from sample.group_bys.quickstart.users import v1 as users
 
 """
 This is the "left side" of the join that will comprise our training set. It is responsible for providing the primary keys

api/py/test/test_sample.py

@@ -0,0 +1,6 @@
+from sample.joins.quickstart.training_set import source
+
+
+def test_source():
+    assert source.events.table == "data.checkouts"
+