refactor + add more comments

Author: sfc-gh-aalam

src/snowflake/snowpark/_internal/debug_utils.py

@@ -11,19 +11,20 @@
 UNKNOWN_FILE = "__UNKNOWN_FILE__"
 
 
-class DataFrameTraceNode:
-    """A node in the trace of a tree that represents the lineage of a DataFrame."""
+class DataFrameLineageNode:
+    """A node representing a dataframe operation in the DAG that represents the lineage of a DataFrame."""
 
     def __init__(self, batch_id: int, stmt_cache) -> None:
         self.batch_id = batch_id
         self.stmt_cache = stmt_cache
 
     @cached_property
     def children(self) -> set[int]:
+        """Returns the batch_ids of the children of this node."""
         return get_dependent_bind_ids(self.stmt_cache[self.batch_id])
 
     def get_src(self):
-        """The source Stmt of the DataFrame descried by the batch_id."""
+        """The source Stmt of the DataFrame described by the batch_id."""
         stmt = self.stmt_cache[self.batch_id]
         api_call = stmt.bind.expr.WhichOneof("variant")
         return (
@@ -49,7 +50,7 @@ def _read_file(
             code_lines = [line.rstrip() for line in code_lines]
             return "\n".join(code_lines)
 
-    def get_format_id(self) -> str:
+    def get_source_id(self) -> str:
         """Unique identifier of the location of the DataFrame in the source code."""
         src = self.get_src()
         if src is None:
@@ -62,8 +63,8 @@ def get_format_id(self) -> str:
         end_column = src.end_column
         return f"{fileno}:{start_line}:{start_column}-{end_line}:{end_column}"
 
-    def get_format_src(self) -> str:
-        """The snippet of the source code where the DataFrame was created."""
+    def get_source_snippet(self) -> str:
+        """Read the source file and extract the snippet where the dataframe is created."""
         src = self.get_src()
         if src is None:
             return "No source"
@@ -83,7 +84,7 @@ def get_format_src(self) -> str:
             code_identifier = f"{filename}:{start_line}"
         else:
             code_identifier = (
-                f"{filename}:{start_line}:{start_column}-{end_line}:{end_column}"
+                f"{filename}|{start_line}:{start_column}-{end_line}:{end_column}"
             )
 
         if filename != UNKNOWN_FILE and os.access(filename, os.R_OK):

src/snowflake/snowpark/dataframe.py

@@ -29,7 +29,7 @@
 )
 
 import snowflake.snowpark
-from snowflake.snowpark._internal.debug_utils import DataFrameTraceNode
+from snowflake.snowpark._internal.debug_utils import DataFrameLineageNode
 import snowflake.snowpark._internal.proto.generated.ast_pb2 as proto
 from snowflake.connector.options import installed_pandas, pandas, pyarrow
 
@@ -351,37 +351,52 @@ def _disambiguate(
 
 
 def _get_df_lineage(dataframes_involved: List["DataFrame"]) -> List[str]:
-    curr: List[DataFrameTraceNode] = []
+    """Helper function to get the lineage of dataframes involved in the exception.
+    It gathers the lineage in the following way:
+
+    1. For each dataframe, it checks if it has an AST ID and if so, it creates a
+        DataFrameLineageNode for it.
+    2. We use BFS to traverse the lineage using dataframes from 1. as the first layer.
+    3. During each iteration, we check if the node's source_id has been visited. If not,
+        we add it to the visited set and append its source format to the trace. This step
+        is needed to avoid source_id added multiple times in lineage due to loops.
+    4. We then explore the next layer by adding the children of the current node to the
+        next layer. We check if the child ID has been visited and if not, we add it to the
+        visited set and append the DataFrameLineageNode for it to the next layer.
+    5. We repeat this process until there are no more nodes to explore.
+    """
+    curr: List[DataFrameLineageNode] = []
     visited_batch_id = set()
-    visited_format_id = set()
+    visited_source_id = set()
 
     for df in dataframes_involved:
         if (batch_id := df._ast_id) is not None:
             stmt_cache = df._session._ast_batch._bind_stmt_cache
-            curr.append(DataFrameTraceNode(batch_id, stmt_cache))
+            curr.append(DataFrameLineageNode(batch_id, stmt_cache))
             if batch_id not in visited_batch_id:
                 visited_batch_id.add(batch_id)
 
-    trace = []
+    lineage = []
 
     while curr:
-        next = []
+        next: List[DataFrameLineageNode] = []
         for node in curr:
+            # tracing updates
+            source_id = node.get_source_id()
+            if source_id not in visited_source_id:
+                visited_source_id.add(source_id)
+                lineage.append(node.get_source_snippet())
+
+            # explore next layer
             for child_id in node.children:
-                if child_id in visited_format_id:
+                if child_id in visited_batch_id:
                     continue
-                visited_format_id.add(child_id)
-                next.append(DataFrameTraceNode(child_id, node.stmt_cache))
-
-            # tracing updates
-            format_id = node.get_format_id()
-            if format_id not in visited_format_id:
-                visited_format_id.add(format_id)
-                trace.append(node.get_format_src())
+                visited_batch_id.add(child_id)
+                next.append(DataFrameLineageNode(child_id, node.stmt_cache))
 
         curr = next
 
-    return trace
+    return lineage
 
 
 def dataframe_exception_handler(func):
@@ -402,7 +417,7 @@ def wrapper(*args, **kwargs):
             traceback_lines = traceback.format_exception(error_type, error_value, tb)
             formatted_traceback = "".join(traceback_lines)
 
-            # compute the trace
+            # get the dataframe lineage
             dataframes_involved = []
             for arg in args:
                 if isinstance(arg, DataFrame):
@@ -428,7 +443,7 @@ def wrapper(*args, **kwargs):
                 if lineage_trace_len > show_lineage_len:
                     traceback_with_debug_info.append(
                         f"... and {lineage_trace_len - show_lineage_len} more.\nYou can increase "
-                        "the trace length by setting SNOWPARK_PYTHON_DATAFRAME_LINEAGE_LENGTH_ON_ERROR "
+                        "the lineage length by setting SNOWPARK_PYTHON_DATAFRAME_LINEAGE_LENGTH_ON_ERROR "
                         "environment variable."
                     )