Programmer-RD-AI commented on code in PR #54588:
URL: https://github.com/apache/airflow/pull/54588#discussion_r2280944725
##########
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."""
Review Comment:
this should be updated to specify the polars support
##########
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),
Review Comment:
i think it would be easier to read if split this up a bit more similar to
the polars dataframe handling
##########
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 = ""
Review Comment:
could we move this to the main function and set it to `None` since its used
by both
##########
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}
Review Comment:
part of this (making the random column name based on a condition) can be
also moved to the top-level of the function body, then we can check if the
`random_column_name` exists then we can add the random column name, this will
reduce duplicate as well
##########
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:
we do not need to handle tuple's here since we only are grouping by a
singular column, but if we need to handle it we should handle it in the pandas
implementation as well
##########
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
Review Comment:
this part also can be moved to after the else if condition's since its the
same to both of the conditions
--
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]
