SNOW-1830601 UDAFs for the decoder

src/snowflake/snowpark/_internal/ast/batch.py

@@ -72,6 +72,10 @@ def reset_id_gen(self) -> None:
         """Resets the ID generator."""
         self._id_gen = itertools.count(start=1)
 
+    def reset_callables(self) -> None:
+        """Resets the callables."""
+        self._callables = {}
+
     def assign(self, symbol: Optional[str] = None) -> proto.Assign:
         """
         Creates a new assignment statement.

src/snowflake/snowpark/relational_grouped_dataframe.py

@@ -423,7 +423,7 @@ def end_partition(self, pdf: pandas.DataFrame) -> pandas.DataFrame:
         _apply_in_pandas_udtf = self._dataframe._session.udtf.register(
             _ApplyInPandas,
             output_schema=output_schema,
-            _emit_ast=_emit_ast,
+            _emit_ast=False,
             **kwargs,
         )
         partition_by = [Column(expr, _emit_ast=False) for expr in self._grouping_exprs]

src/snowflake/snowpark/types.py

@@ -784,7 +784,7 @@ def json_value(self) -> Dict[str, Any]:
     def _fill_ast(self, ast: proto.SpDataType) -> None:
         ast.sp_struct_type.structured = self.structured
         for field in self.fields:
-            field._fill_ast(ast.sp_struct_type.fields.add())
+            field._fill_ast(ast.sp_struct_type.fields.list.add())
 
 
 class VariantType(DataType):