Got UDTFs to work, changed the UDTF test, got RelationalGroupedDataFrame to work, added a way to reset the recorded callables

Author: sfc-gh-vbudati

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/_internal/udf_utils.py

@@ -595,6 +595,8 @@ def extract_return_input_types(
                 return True, False, res_return_type, res_input_types
         elif isinstance(
             return_type_from_type_hints, PandasDataFrameType
+        ) and not isinstance(
+            return_type, PandasDataFrameType
         ):  # vectorized UDTF
             return_type = PandasDataFrameType(
                 [x.datatype for x in return_type], [x.name for x in return_type]

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):