adding types; move import

Author: nyaapa

python/pyspark/sql/pandas/serializers.py

@@ -2174,8 +2174,9 @@ def load_stream(self, stream):
         from pyspark.sql.streaming.stateful_processor_util import (
             TransformWithStateInPandasFuncMode,
         )
+        from typing import Iterator, Any, Optional, Tuple
 
-        def generate_data_batches(batches):
+        def generate_data_batches(batches) -> Iterator[Tuple[Any, Optional[Any], Optional[Any]]]:
             """
             Deserialize ArrowRecordBatches and return a generator of Row.
             The deserialization logic assumes that Arrow RecordBatches contain the data with the
@@ -2186,7 +2187,7 @@ def generate_data_batches(batches):
              into the data generator.
             """
 
-            def extract_rows(cur_batch, col_name, key_offsets):
+            def extract_rows(cur_batch, col_name, key_offsets) -> Optional[Iterator[Tuple[Any, Any]]]:
                 data_column = cur_batch.column(cur_batch.schema.get_field_index(col_name))
 
                 # Check if the entire column is null

python/pyspark/worker.py

@@ -3069,8 +3069,8 @@ def values_gen():
         ser.init_key_offsets = parsed_offsets[1][0]
         stateful_processor_api_client = StatefulProcessorApiClient(state_server_port, key_schema)
 
-        import pandas as pd
         def mapper(a):
+            import pandas as pd
             mode = a[0]
 
             if mode == TransformWithStateInPandasFuncMode.PROCESS_DATA:
@@ -3233,7 +3233,6 @@ def mapper(a):
 
         parsed_offsets = extract_key_value_indexes(arg_offsets)
 
-        import pandas as pd
         def mapper(a):
             df1_keys = [a[0][o] for o in parsed_offsets[0][0]]
             df1_vals = [a[0][o] for o in parsed_offsets[0][1]]