Capture AST for DataFrame.pivot and DataFrame.unpivot

src/snowflake/snowpark/dataframe.py

@@ -104,6 +104,7 @@
     build_expr_from_snowpark_column_or_col_name,
     build_expr_from_snowpark_column_or_sql_str,
     build_expr_from_snowpark_column_or_table_fn,
+    build_proto_from_pivot_values,
     fill_ast_for_column,
     with_src_position,
 )
@@ -2192,7 +2193,12 @@ def pivot(
         """
 
         if _emit_ast:
-            raise NotImplementedError("TODO SNOW-1491297: Add coverage for pivot.")
+            stmt = self._session._ast_batch.assign()
+            ast = with_src_position(stmt.expr.sp_dataframe_pivot, stmt)
+            self.set_ast_ref(ast.df)
+            build_expr_from_snowpark_column_or_col_name(ast.pivot_col, pivot_col)
+            build_proto_from_pivot_values(ast.values, values)
+            build_expr_from_python_val(ast.default_on_null, default_on_null)
 
         target_df, pc, pivot_values, default_on_null = prepare_pivot_arguments(
             self, "DataFrame.pivot", pivot_col, values, default_on_null
@@ -2204,6 +2210,7 @@ def pivot(
             snowflake.snowpark.relational_grouped_dataframe._PivotType(
                 pc[0], pivot_values, default_on_null
             ),
+            ast_stmt=stmt,
         )
 
     @df_api_usage
@@ -2237,19 +2244,32 @@ def unpivot(
             ---------------------------------------------
             <BLANKLINE>
         """
+        # AST.
+        stmt = self._session._ast_batch.assign()
+        ast = with_src_position(stmt.expr.sp_dataframe_unpivot, stmt)
+        self.set_ast_ref(ast.df)
+        ast.value_column = value_column
+        ast.name_column = name_column
+        for c in column_list:
+            build_expr_from_snowpark_column_or_col_name(ast.column_list.add(), c)
+
         column_exprs = self._convert_cols_to_exprs("unpivot()", column_list)
         unpivot_plan = Unpivot(value_column, name_column, column_exprs, self._plan)
 
+        df: DataFrame
         if self._select_statement:
-            return self._with_plan(
+            df = self._with_plan(
                 SelectStatement(
                     from_=SelectSnowflakePlan(
                         unpivot_plan, analyzer=self._session._analyzer
                     ),
                     analyzer=self._session._analyzer,
                 )
             )
-        return self._with_plan(unpivot_plan)
+        else:
+            df = self._with_plan(unpivot_plan)
+        df._ast_id = stmt.var_id.bitfield1
+        return df
 
     @df_api_usage
     def limit(

tests/ast/data/DataFrame.pivot.test

@@ -0,0 +1,38 @@
+## TEST CASE
+
+df = session.create_dataframe(
+    [
+        (1, 10000, "JAN"),
+        (1, 400, "JAN"),
+        (2, 4500, "JAN"),
+        (2, 35000, "JAN"),
+        (1, 5000, "FEB"),
+        (1, 3000, "FEB"),
+        (2, 200, "FEB"),
+    ],
+    schema=["k", "t", "mo"],
+)
+
+df1 = df.pivot("mo", ["JAN", "FEB"]).sum("t").sort("k")
+
+df2 = df.pivot("mo", values=["JAN", "FEB"], default_on_null="Nothing").sum("t").sort("k")
+
+## EXPECTED UNPARSER OUTPUT
+
+df = session.create_dataframe([(1, 10000, "JAN"), (1, 400, "JAN"), (2, 4500, "JAN"), (2, 35000, "JAN"), (1, 5000, "FEB"), (1, 3000, "FEB"), (2, 200, "FEB")], schema=["k", "t", "mo"])
+
+df1 = df.pivot("mo", values=["JAN", "FEB"], default_on_null=None)
+
+df1 = df1.sum("t")
+
+df1 = df1.sort("k")
+
+df2 = df.pivot("mo", values=["JAN", "FEB"], default_on_null="Nothing")
+
+df2 = df2.sum("t")
+
+df2 = df2.sort("k")
+
+## EXPECTED ENCODED AST
+
+CrcICrQICqUI0gWhCAr0BwrxBwqNAZoMiQEKGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaEiGSAh4KGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaEAESIpICHwoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKFoQkE4SJPILIQoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKFoSA0pBTgqNAZoMiQEKGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaEiGSAh4KGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaEAESIpICHwoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKFoQkAMSJPILIQoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKFoSA0pBTgqNAZoMiQEKGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaEiGSAh4KGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaEAISIpICHwoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKFoQlCMSJPILIQoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKFoSA0pBTgqOAZoMigEKGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaEiGSAh4KGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaEAISI5ICIAoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKFoQuJECEiTyCyEKGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaEgNKQU4KjQGaDIkBChoaFlNSQ19QT1NJVElPTl9URVNUX01PREUoWhIhkgIeChoaFlNSQ19QT1NJVElPTl9URVNUX01PREUoWhABEiKSAh8KGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaEIgnEiTyCyEKGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaEgNGRUIKjQGaDIkBChoaFlNSQ19QT1NJVElPTl9URVNUX01PREUoWhIhkgIeChoaFlNSQ19QT1NJVElPTl9URVNUX01PREUoWhABEiKSAh8KGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaELgXEiTyCyEKGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaEgNGRUIKjQGaDIkBChoaFlNSQ19QT1NJVElPTl9URVNUX01PREUoWhIhkgIeChoaFlNSQ19QT1NJVElPTl9URVNUX01PREUoWhACEiKSAh8KGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaEMgBEiTyCyEKGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShaEgNGRUISDAoKCgFrCgF0CgJtbxoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKFoSBAoCZGYYASICCAEK8wEK8AEK4AGiCNwBCh/iAhwKGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShnEgeKAgQKAggBGiPyCyAKGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShnEgJtbyIaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKGcqbxJtCmuiAmgKGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShnEiTyCyEKGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShnEgNKQU4SJPILIQoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKGcSA0ZFQhIFCgNkZjEYAiICCAIKZQpjClT6ClEKA3N1bRImCiLyCx8KGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShnEgF0EAEaBlIECgIIAiIaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKGcSBQoDZGYxGAMiAggDCl8KXQpO+ghLEiLyCx8KGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShnEgFrGAEiB4oCBAoCCAMqGhoWU1JDX1BPU0lUSU9OX1RFU1RfTU9ERShnEgUKA2RmMRgEIgIIBAr8AQr5AQrpAaII5QEKKPILJQoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKGkSB05vdGhpbmcSB4oCBAoCCAEaI/ILIAoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKGkSAm1vIhoaFlNSQ19QT1NJVElPTl9URVNUX01PREUoaSpvEm0Ka6ICaAoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKGkSJPILIQoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKGkSA0pBThIk8gshChoaFlNSQ19QT1NJVElPTl9URVNUX01PREUoaRIDRkVCEgUKA2RmMhgFIgIIBQplCmMKVPoKUQoDc3VtEiYKIvILHwoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKGkSAXQQARoGUgQKAggFIhoaFlNSQ19QT1NJVElPTl9URVNUX01PREUoaRIFCgNkZjIYBiICCAYKXwpdCk76CEsSIvILHwoaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKGkSAWsYASIHigIECgIIBioaGhZTUkNfUE9TSVRJT05fVEVTVF9NT0RFKGkSBQoDZGYyGAciAggHEAEaERIPCg0KBWZpbmFsEAMYCSATIgQQARgV

tests/ast/data/DataFrame.unpivot.test.DISABLED

@@ -0,0 +1,12 @@
+## TEST CASE
+
+df = session.create_dataframe(
+    [(1, "electronics", 100, 200), (2, "clothes", 100, 300)],
+    schema=["empid", "dept", "jan", "feb"],
+)
+df = df.unpivot("sales", "month", ["jan", "feb"]).sort("empid")
+
+
+## EXPECTED UNPARSER OUTPUT
+
+## EXPECTED ENCODED AST