guan404ming commented on code in PR #54588:
URL: https://github.com/apache/airflow/pull/54588#discussion_r2280989179
##########
providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py:
##########
@@ -261,38 +261,56 @@ 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 | pl.DataFrame) ->
Iterable[tuple[str, pd.DataFrame]]:
+ def _partition_dataframe(
+ self, df: pd.DataFrame | pl.DataFrame
+ ) -> Iterable[tuple[str, pd.DataFrame | pl.DataFrame]]:
"""Partition dataframe using pandas groupby() method."""
try:
import secrets
import string
import numpy as np
+ import pandas as pd
import polars as pl
except ImportError:
pass
- if isinstance(df, pl.DataFrame):
- df = df.to_pandas()
-
# if max_rows_per_file argument is specified, a temporary column with
a random unusual name will be
# added to the dataframe. This column is used to dispatch the
dataframe into smaller ones using groupby()
-
- random_column_name = ""
- if self.max_rows_per_file and not self.groupby_kwargs:
- random_column_name = "".join(secrets.choice(string.ascii_letters)
for _ in range(20))
- df[random_column_name] = np.arange(len(df)) //
self.max_rows_per_file
- self.groupby_kwargs = {"by": random_column_name}
- if not self.groupby_kwargs:
- yield "", df
- return
- for group_label in (grouped_df :=
df.groupby(**self.groupby_kwargs)).groups:
- yield (
- cast("str", group_label),
- grouped_df.get_group(group_label)
- .drop(random_column_name, axis=1, errors="ignore")
- .reset_index(drop=True),
- )
+ if isinstance(df, pd.DataFrame):
+ random_column_name = ""
+ if self.max_rows_per_file and not self.groupby_kwargs:
+ random_column_name =
"".join(secrets.choice(string.ascii_letters) for _ in range(20))
+ df[random_column_name] = np.arange(len(df)) //
self.max_rows_per_file
+ self.groupby_kwargs = {"by": random_column_name}
+ if not self.groupby_kwargs:
+ yield "", df
+ return
+ for group_label in (grouped_df :=
df.groupby(**self.groupby_kwargs)).groups:
+ yield (
+ cast("str", group_label),
+ grouped_df.get_group(group_label)
+ .drop(random_column_name, axis=1, errors="ignore")
+ .reset_index(drop=True),
+ )
+ elif isinstance(df, pl.DataFrame):
+ random_column_name = ""
+ if self.max_rows_per_file and not self.groupby_kwargs:
+ random_column_name =
"".join(secrets.choice(string.ascii_letters) for _ in range(20))
+ df = df.with_columns(
+ (pl.int_range(pl.len()) //
self.max_rows_per_file).alias(random_column_name)
+ )
+ self.groupby_kwargs = {"by": random_column_name}
+ if not self.groupby_kwargs:
+ yield "", df
+ return
+ for group_label, group_df in df.group_by(**self.groupby_kwargs):
+ if random_column_name:
+ group_df = group_df.drop(random_column_name)
+ yield (
+ cast("str", group_label[0] if isinstance(group_label,
tuple) else group_label),
Review Comment:
Polars group_by() actually does return tuples even for single column
grouping, while pandas doesn't. The tuple handling was needed for polars. Let's
handle it in pandas to keep consistent and extensibility.
##########
providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py:
##########
@@ -261,38 +261,56 @@ 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 | pl.DataFrame) ->
Iterable[tuple[str, pd.DataFrame]]:
+ def _partition_dataframe(
+ self, df: pd.DataFrame | pl.DataFrame
+ ) -> Iterable[tuple[str, pd.DataFrame | pl.DataFrame]]:
"""Partition dataframe using pandas groupby() method."""
try:
import secrets
import string
import numpy as np
+ import pandas as pd
import polars as pl
except ImportError:
pass
- if isinstance(df, pl.DataFrame):
- df = df.to_pandas()
-
# if max_rows_per_file argument is specified, a temporary column with
a random unusual name will be
# added to the dataframe. This column is used to dispatch the
dataframe into smaller ones using groupby()
-
- random_column_name = ""
- if self.max_rows_per_file and not self.groupby_kwargs:
- random_column_name = "".join(secrets.choice(string.ascii_letters)
for _ in range(20))
- df[random_column_name] = np.arange(len(df)) //
self.max_rows_per_file
- self.groupby_kwargs = {"by": random_column_name}
- if not self.groupby_kwargs:
- yield "", df
- return
- for group_label in (grouped_df :=
df.groupby(**self.groupby_kwargs)).groups:
- yield (
- cast("str", group_label),
- grouped_df.get_group(group_label)
- .drop(random_column_name, axis=1, errors="ignore")
- .reset_index(drop=True),
- )
+ if isinstance(df, pd.DataFrame):
+ random_column_name = ""
+ if self.max_rows_per_file and not self.groupby_kwargs:
+ random_column_name =
"".join(secrets.choice(string.ascii_letters) for _ in range(20))
+ df[random_column_name] = np.arange(len(df)) //
self.max_rows_per_file
+ self.groupby_kwargs = {"by": random_column_name}
+ if not self.groupby_kwargs:
+ yield "", df
+ return
+ for group_label in (grouped_df :=
df.groupby(**self.groupby_kwargs)).groups:
+ yield (
+ cast("str", group_label),
+ grouped_df.get_group(group_label)
+ .drop(random_column_name, axis=1, errors="ignore")
+ .reset_index(drop=True),
+ )
+ elif isinstance(df, pl.DataFrame):
+ random_column_name = ""
+ if self.max_rows_per_file and not self.groupby_kwargs:
+ random_column_name =
"".join(secrets.choice(string.ascii_letters) for _ in range(20))
+ df = df.with_columns(
+ (pl.int_range(pl.len()) //
self.max_rows_per_file).alias(random_column_name)
+ )
+ self.groupby_kwargs = {"by": random_column_name}
+ if not self.groupby_kwargs:
+ yield "", df
+ return
+ for group_label, group_df in df.group_by(**self.groupby_kwargs):
+ if random_column_name:
+ group_df = group_df.drop(random_column_name)
+ yield (
+ cast("str", group_label[0] if isinstance(group_label,
tuple) else group_label),
Review Comment:
Polars `group_by()` actually does return tuples even for single column
grouping, while pandas doesn't. The tuple handling was needed for polars. Let's
handle it in pandas to keep consistent and extensibility.
--
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]
