SNOW-1897976: adding pipe method to DataFrame

src/snowflake/snowpark/dataframe.py

@@ -7027,6 +7027,33 @@ def print_schema(self, level: Optional[int] = None) -> None:
     # naturalJoin = natural_join
     # withColumns = with_columns
 
+    def pipe(self, f, *args, **kwargs):
+        """
+        Applies a function to the DataFrame and returns the result.
+
+        Args:
+            f: A function to apply to the DataFrame.
+            *args: Additional positional arguments to pass to the function.
+            **kwargs: Additional keyword arguments to pass to the function.
+
+        Returns:
+            The result of applying the function to the DataFrame.
+
+        Examples::
+            >>> df = session.create_dataframe([(1, "a"), (2, "b")], schema=["col1", "col2"])
+            >>> def add_column(df):
+            ...     return df.with_column("col3", lit(3))
+            >>> new_df = df.pipe(add_column)
+            >>> new_df.show()
+            --------------------
+            |"COL1"  |"COL2"  |"COL3"  |
+            --------------------
+            |1        |a       |3       |
+            |2        |b       |3       |
+            --------------------
+        """
+        return f(self, *args, **kwargs)
+
 
 def map(
     dataframe: DataFrame,

tests/unit/test_dataframe.py

@@ -409,6 +409,24 @@ def test_dataFrame_printSchema(capfd, mock_server_connection):
     )
 
 
+def test_dataFrame_pipe(mock_server_connection):
+    session = snowflake.snowpark.session.Session(mock_server_connection)
+    df = session.create_dataframe([[1, ""], [3, None]])
+    df._plan._metadata = PlanMetadata(
+        attributes=[
+            Attribute("A", IntegerType(), False),
+            Attribute("B", StringType()),
+        ],
+        quoted_identifiers=None,
+    )
+
+    def test_func(df):
+        return df
+
+    result_df, expected_result = df.pipe(test_func), test_func(df)
+    assert result_df == expected_result
+
+
 def test_session():
     fake_session = mock.create_autospec(Session, _session_id=123456)
     fake_session._analyzer = mock.Mock()