HyukjinKwon commented on code in PR #40722:
URL: https://github.com/apache/spark/pull/40722#discussion_r1161531972


##########
python/pyspark/pandas/data_type_ops/num_ops.py:
##########
@@ -82,24 +81,41 @@ def add(self, left: IndexOpsLike, right: Any) -> 
SeriesOrIndex:
             raise TypeError("Addition can not be applied to given types.")
 
         right = transform_boolean_operand_to_numeric(right, 
spark_type=left.spark.data_type)
-        Column: Type[GenericColumn] = ConnectColumn if is_remote() else 
PySparkColumn
-        return column_op(Column.__add__)(left, right)
+        if is_remote():
+            from pyspark.sql.connect.column import Column as ConnectColumn
+
+            Column = ConnectColumn
+        else:
+            Column = PySparkColumn  # type: ignore[assignment]
+        return column_op(Column.__add__)(left, right)  # type: ignore[arg-type]
 
     def sub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
         _sanitize_list_like(right)
         if not is_valid_operand_for_numeric_arithmetic(right):
             raise TypeError("Subtraction can not be applied to given types.")
 
         right = transform_boolean_operand_to_numeric(right, 
spark_type=left.spark.data_type)
-        Column: Type[GenericColumn] = ConnectColumn if is_remote() else 
PySparkColumn
-        return column_op(Column.__sub__)(left, right)
+        if is_remote():
+            from pyspark.sql.connect.column import Column as ConnectColumn
+
+            Column = ConnectColumn
+        else:
+            Column = PySparkColumn  # type: ignore[assignment]
+        return column_op(Column.__sub__)(left, right)  # type: ignore[arg-type]
 
     def mod(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
         _sanitize_list_like(right)
         if not is_valid_operand_for_numeric_arithmetic(right):
             raise TypeError("Modulo can not be applied to given types.")
 
-        def mod(left: GenericColumn, right: Any) -> GenericColumn:
+        if is_remote():
+            from pyspark.sql.connect.column import Column as ConnectColumn
+
+            Column = ConnectColumn
+        else:
+            Column = PySparkColumn  # type: ignore[assignment]
+
+        def mod(left: Column, right: Any) -> Column:  # type: 
ignore[valid-type]

Review Comment:
   If this is only for type hint, you can simply cast to `Column` without 
if-else.



##########
python/pyspark/pandas/data_type_ops/num_ops.py:
##########
@@ -110,13 +126,18 @@ def pow(self, left: IndexOpsLike, right: Any) -> 
SeriesOrIndex:
         if not is_valid_operand_for_numeric_arithmetic(right):
             raise TypeError("Exponentiation can not be applied to given 
types.")
 
-        Column = ConnectColumn if is_remote() else PySparkColumn
+        if is_remote():
+            from pyspark.sql.connect.column import Column as ConnectColumn
 
-        def pow_func(left: GenericColumn, right: Any) -> GenericColumn:
+            Column = ConnectColumn
+        else:
+            Column = PySparkColumn  # type: ignore[assignment]
+
+        def pow_func(left: Column, right: Any) -> Column:  # type: 
ignore[valid-type]

Review Comment:
   ditto



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to