eladinver commented on code in PR #54195:
URL: https://github.com/apache/airflow/pull/54195#discussion_r2280794972


##########
providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py:
##########
@@ -220,17 +244,23 @@ def execute(self, context: Context) -> None:
                 file_obj=buf, key=object_key, bucket_name=self.s3_bucket, 
replace=self.replace
             )
 
-    def _partition_dataframe(self, df: pd.DataFrame) -> Iterable[tuple[str, 
pd.DataFrame]]:
+    def _partition_dataframe(self, df: pd.DataFrame | pl.DataFrame) -> 
Iterable[tuple[str, pd.DataFrame]]:
         """Partition dataframe using pandas groupby() method."""
         try:
             import secrets
             import string
 
             import numpy as np
+            import polars as pl
         except ImportError:
             pass
+
+        if isinstance(df, pl.DataFrame):
+            df = df.to_pandas()

Review Comment:
   @guan404ming Taking a more closer look. Is this right?
   This means that user who uses `polars` must also install `pandas`.



##########
providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py:
##########
@@ -220,17 +244,23 @@ def execute(self, context: Context) -> None:
                 file_obj=buf, key=object_key, bucket_name=self.s3_bucket, 
replace=self.replace
             )
 
-    def _partition_dataframe(self, df: pd.DataFrame) -> Iterable[tuple[str, 
pd.DataFrame]]:
+    def _partition_dataframe(self, df: pd.DataFrame | pl.DataFrame) -> 
Iterable[tuple[str, pd.DataFrame]]:
         """Partition dataframe using pandas groupby() method."""
         try:
             import secrets
             import string
 
             import numpy as np
+            import polars as pl
         except ImportError:
             pass
+
+        if isinstance(df, pl.DataFrame):
+            df = df.to_pandas()

Review Comment:
   @guan404ming Taking a more closer look. Is this right?
   This means that user who uses `polars` must also install `pandas`.



-- 
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]

Reply via email to