From 86d1127031926c4e34975d07fb5f76ac9d3ac7c9 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Tue, 4 Apr 2023 20:03:34 +0530 Subject: [PATCH 1/7] Add partition_kwargs for partition of dataframes --- .../amazon/aws/transfers/sql_to_s3.py | 31 ++++-- .../amazon/aws/transfers/test_sql_to_s3.py | 94 +++++++++++++++++++ 2 files changed, 116 insertions(+), 9 deletions(-) diff --git a/airflow/providers/amazon/aws/transfers/sql_to_s3.py b/airflow/providers/amazon/aws/transfers/sql_to_s3.py index a8b5a9cd1c53f..bd38bf14d6fc2 100644 --- a/airflow/providers/amazon/aws/transfers/sql_to_s3.py +++ b/airflow/providers/amazon/aws/transfers/sql_to_s3.py @@ -107,6 +107,7 @@ def __init__( verify: bool | str | None = None, file_format: Literal["csv", "json", "parquet"] = "csv", pd_kwargs: dict | None = None, + partition_kwargs: dict | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -119,6 +120,7 @@ def __init__( self.replace = replace self.pd_kwargs = pd_kwargs or {} self.parameters = parameters + self.partition_kwargs = partition_kwargs or {} if "path_or_buf" in self.pd_kwargs: raise AirflowException("The argument path_or_buf is not allowed, please remove it") @@ -170,15 +172,26 @@ def execute(self, context: Context) -> None: self._fix_dtypes(data_df, self.file_format) file_options = FILE_OPTIONS_MAP[self.file_format] - with NamedTemporaryFile(mode=file_options.mode, suffix=file_options.suffix) as tmp_file: - - self.log.info("Writing data to temp file") - getattr(data_df, file_options.function)(tmp_file.name, **self.pd_kwargs) - - self.log.info("Uploading data to S3") - s3_conn.load_file( - filename=tmp_file.name, key=self.s3_key, bucket_name=self.s3_bucket, replace=self.replace - ) + for group_name, df in self._partition_dataframe(df=data_df): + with NamedTemporaryFile(mode=file_options.mode, suffix=file_options.suffix) as tmp_file: + + self.log.info("Writing data to temp file") + getattr(data_df, file_options.function)(tmp_file.name, **self.pd_kwargs) + + self.log.info("Uploading data to S3") + object_key = f"{self.s3_key}_{group_name}" if group_name else self.s3_key + s3_conn.load_file( + filename=tmp_file.name, key=object_key, bucket_name=self.s3_bucket, replace=self.replace + ) + + def _partition_dataframe(self, df: DataFrame) -> Iterable[str, DataFrame]: + """Partition dataframe using pandas groupby() method""" + if not self.partition_kwargs: + yield "", df + else: + grouped_df = df.groupby(**self.partition_kwargs) + for group_label in grouped_df.groups.keys(): + yield group_label, grouped_df.get_group(group_label).reset_index(drop=True) def _get_hook(self) -> DbApiHook: self.log.debug("Get connection for %s", self.sql_conn_id) diff --git a/tests/providers/amazon/aws/transfers/test_sql_to_s3.py b/tests/providers/amazon/aws/transfers/test_sql_to_s3.py index 84db615eebccf..9dc923a469f86 100644 --- a/tests/providers/amazon/aws/transfers/test_sql_to_s3.py +++ b/tests/providers/amazon/aws/transfers/test_sql_to_s3.py @@ -175,3 +175,97 @@ def test_invalid_file_format(self): file_format="invalid_format", dag=None, ) + + def test_with_partition_kwarg(self): + """ + Test operator when the partition_kwargs is specified + """ + query = "query" + s3_bucket = "bucket" + s3_key = "key" + + op = SqlToS3Operator( + query=query, + s3_bucket=s3_bucket, + s3_key=s3_key, + sql_conn_id="mysql_conn_id", + aws_conn_id="aws_conn_id", + task_id="task_id", + replace=True, + pd_kwargs={"index": False, "header": False}, + partition_kwargs={"by": "Team"}, + dag=None, + ) + example = { + "Team": ["Australia", "Australia", "India", "India"], + "Player": ["Ricky", "David Warner", "Virat Kohli", "Rohit Sharma"], + "Runs": [345, 490, 672, 560], + } + + df = pd.DataFrame(example) + data = [] + for group_name, df in op._partition_dataframe(df): + data.append((group_name, df)) + data.sort(key=lambda d: d[0]) + team, df = data[0] + assert df.equals( + pd.DataFrame( + { + "Team": ["Australia", "Australia"], + "Player": ["Ricky", "David Warner"], + "Runs": [345, 490], + } + ) + ) + team, df = data[1] + assert df.equals( + pd.DataFrame( + { + "Team": ["India", "India"], + "Player": ["Virat Kohli", "Rohit Sharma"], + "Runs": [672, 560], + } + ) + ) + + def test_without_partition_kwarg(self): + """ + Test operator when the partition_kwargs is not specified + """ + query = "query" + s3_bucket = "bucket" + s3_key = "key" + + op = SqlToS3Operator( + query=query, + s3_bucket=s3_bucket, + s3_key=s3_key, + sql_conn_id="mysql_conn_id", + aws_conn_id="aws_conn_id", + task_id="task_id", + replace=True, + pd_kwargs={"index": False, "header": False}, + dag=None, + ) + example = { + "Team": ["Australia", "Australia", "India", "India"], + "Player": ["Ricky", "David Warner", "Virat Kohli", "Rohit Sharma"], + "Runs": [345, 490, 672, 560], + } + + df = pd.DataFrame(example) + data = [] + for group_name, df in op._partition_dataframe(df): + data.append((group_name, df)) + + assert len(data) == 1 + team, df = data[0] + assert df.equals( + pd.DataFrame( + { + "Team": ["Australia", "Australia", "India", "India"], + "Player": ["Ricky", "David Warner", "Virat Kohli", "Rohit Sharma"], + "Runs": [345, 490, 672, 560], + } + ) + ) From 9f3e7c5203565f21d8ce6147647afe9a8f237399 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Tue, 4 Apr 2023 20:19:45 +0530 Subject: [PATCH 2/7] Update the name of param --- airflow/providers/amazon/aws/transfers/sql_to_s3.py | 8 ++++---- tests/providers/amazon/aws/transfers/test_sql_to_s3.py | 10 +++++----- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/airflow/providers/amazon/aws/transfers/sql_to_s3.py b/airflow/providers/amazon/aws/transfers/sql_to_s3.py index bd38bf14d6fc2..f8713b629f7b3 100644 --- a/airflow/providers/amazon/aws/transfers/sql_to_s3.py +++ b/airflow/providers/amazon/aws/transfers/sql_to_s3.py @@ -107,7 +107,7 @@ def __init__( verify: bool | str | None = None, file_format: Literal["csv", "json", "parquet"] = "csv", pd_kwargs: dict | None = None, - partition_kwargs: dict | None = None, + groupby_kwargs: dict | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -120,7 +120,7 @@ def __init__( self.replace = replace self.pd_kwargs = pd_kwargs or {} self.parameters = parameters - self.partition_kwargs = partition_kwargs or {} + self.groupby_kwargs = groupby_kwargs or {} if "path_or_buf" in self.pd_kwargs: raise AirflowException("The argument path_or_buf is not allowed, please remove it") @@ -186,10 +186,10 @@ def execute(self, context: Context) -> None: def _partition_dataframe(self, df: DataFrame) -> Iterable[str, DataFrame]: """Partition dataframe using pandas groupby() method""" - if not self.partition_kwargs: + if not self.groupby_kwargs: yield "", df else: - grouped_df = df.groupby(**self.partition_kwargs) + grouped_df = df.groupby(**self.groupby_kwargs) for group_label in grouped_df.groups.keys(): yield group_label, grouped_df.get_group(group_label).reset_index(drop=True) diff --git a/tests/providers/amazon/aws/transfers/test_sql_to_s3.py b/tests/providers/amazon/aws/transfers/test_sql_to_s3.py index 9dc923a469f86..a0e4e6f603edb 100644 --- a/tests/providers/amazon/aws/transfers/test_sql_to_s3.py +++ b/tests/providers/amazon/aws/transfers/test_sql_to_s3.py @@ -176,9 +176,9 @@ def test_invalid_file_format(self): dag=None, ) - def test_with_partition_kwarg(self): + def test_with_groupby_kwarg(self): """ - Test operator when the partition_kwargs is specified + Test operator when the groupby_kwargs is specified """ query = "query" s3_bucket = "bucket" @@ -193,7 +193,7 @@ def test_with_partition_kwarg(self): task_id="task_id", replace=True, pd_kwargs={"index": False, "header": False}, - partition_kwargs={"by": "Team"}, + groupby_kwargs={"by": "Team"}, dag=None, ) example = { @@ -228,9 +228,9 @@ def test_with_partition_kwarg(self): ) ) - def test_without_partition_kwarg(self): + def test_without_groupby_kwarg(self): """ - Test operator when the partition_kwargs is not specified + Test operator when the groupby_kwargs is not specified """ query = "query" s3_bucket = "bucket" From 51826720b20a61fd2986f109075608f9d6fccbfe Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Thu, 6 Apr 2023 09:46:33 +0530 Subject: [PATCH 3/7] Add document for grouping feature --- airflow/providers/amazon/aws/transfers/sql_to_s3.py | 2 +- .../transfer/sql_to_s3.rst | 13 +++++++++++++ .../providers/amazon/aws/example_sql_to_s3.py | 12 ++++++++++++ 3 files changed, 26 insertions(+), 1 deletion(-) diff --git a/airflow/providers/amazon/aws/transfers/sql_to_s3.py b/airflow/providers/amazon/aws/transfers/sql_to_s3.py index f8713b629f7b3..204fd55e6ed24 100644 --- a/airflow/providers/amazon/aws/transfers/sql_to_s3.py +++ b/airflow/providers/amazon/aws/transfers/sql_to_s3.py @@ -184,7 +184,7 @@ def execute(self, context: Context) -> None: filename=tmp_file.name, key=object_key, bucket_name=self.s3_bucket, replace=self.replace ) - def _partition_dataframe(self, df: DataFrame) -> Iterable[str, DataFrame]: + def _partition_dataframe(self, df: DataFrame) -> Iterable[tuple[str, DataFrame]]: """Partition dataframe using pandas groupby() method""" if not self.groupby_kwargs: yield "", df diff --git a/docs/apache-airflow-providers-amazon/transfer/sql_to_s3.rst b/docs/apache-airflow-providers-amazon/transfer/sql_to_s3.rst index 31223be0f67aa..482b4fb805b27 100644 --- a/docs/apache-airflow-providers-amazon/transfer/sql_to_s3.rst +++ b/docs/apache-airflow-providers-amazon/transfer/sql_to_s3.rst @@ -50,6 +50,19 @@ Example usage: :start-after: [START howto_transfer_sql_to_s3] :end-before: [END howto_transfer_sql_to_s3] +Grouping +-------- + +We can group the data in the table by passing the ``groupby_kwargs`` param. This param accepts a ``dict`` which will be passed to pandas :doc:`goupby() ` as kwargs. + +Example usage: + +.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_sql_to_s3.py + :language: python + :dedent: 4 + :start-after: [START howto_transfer_sql_to_s3_with_groupby_param] + :end-before: [END howto_transfer_sql_to_s3_with_groupby_param] + Reference --------- diff --git a/tests/system/providers/amazon/aws/example_sql_to_s3.py b/tests/system/providers/amazon/aws/example_sql_to_s3.py index f983881a3f707..192e21f8d53ca 100644 --- a/tests/system/providers/amazon/aws/example_sql_to_s3.py +++ b/tests/system/providers/amazon/aws/example_sql_to_s3.py @@ -173,6 +173,18 @@ def delete_security_group(sec_group_id: str, sec_group_name: str): ) # [END howto_transfer_sql_to_s3] + # [START howto_transfer_sql_to_s3_with_groupby_param] + sql_to_s3_task_with_groupby = SqlToS3Operator( + task_id="sql_to_s3_task", + sql_conn_id=conn_id_name, + query=SQL_QUERY, + s3_bucket=bucket_name, + s3_key=key, + replace=True, + groupby_kwargs={"by": "color"}, + ) + # [END howto_transfer_sql_to_s3_with_groupby_param] + delete_bucket = S3DeleteBucketOperator( task_id="delete_bucket", bucket_name=bucket_name, From 19ac924e20ee31dcefb1bc2ec4268f9fc0b19bf8 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 7 Apr 2023 18:47:06 +0530 Subject: [PATCH 4/7] Fix link --- docs/apache-airflow-providers-amazon/transfer/sql_to_s3.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/apache-airflow-providers-amazon/transfer/sql_to_s3.rst b/docs/apache-airflow-providers-amazon/transfer/sql_to_s3.rst index 482b4fb805b27..e01d415942934 100644 --- a/docs/apache-airflow-providers-amazon/transfer/sql_to_s3.rst +++ b/docs/apache-airflow-providers-amazon/transfer/sql_to_s3.rst @@ -53,7 +53,7 @@ Example usage: Grouping -------- -We can group the data in the table by passing the ``groupby_kwargs`` param. This param accepts a ``dict`` which will be passed to pandas :doc:`goupby() ` as kwargs. +We can group the data in the table by passing the ``groupby_kwargs`` param. This param accepts a ``dict`` which will be passed to pandas `groupby() `_ as kwargs. Example usage: From e53162f79910c7fc1771e98bfe37297f9f0caf2c Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Mon, 10 Apr 2023 09:16:17 +0530 Subject: [PATCH 5/7] Fix typo --- airflow/providers/amazon/aws/transfers/sql_to_s3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/amazon/aws/transfers/sql_to_s3.py b/airflow/providers/amazon/aws/transfers/sql_to_s3.py index 204fd55e6ed24..2ddb0058dbed9 100644 --- a/airflow/providers/amazon/aws/transfers/sql_to_s3.py +++ b/airflow/providers/amazon/aws/transfers/sql_to_s3.py @@ -176,7 +176,7 @@ def execute(self, context: Context) -> None: with NamedTemporaryFile(mode=file_options.mode, suffix=file_options.suffix) as tmp_file: self.log.info("Writing data to temp file") - getattr(data_df, file_options.function)(tmp_file.name, **self.pd_kwargs) + getattr(df, file_options.function)(tmp_file.name, **self.pd_kwargs) self.log.info("Uploading data to S3") object_key = f"{self.s3_key}_{group_name}" if group_name else self.s3_key From 998e2f74da4d7ea7c680c8ee31e106f20ff360fb Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Wed, 12 Apr 2023 10:37:27 +0530 Subject: [PATCH 6/7] Update docstring --- airflow/providers/amazon/aws/transfers/sql_to_s3.py | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow/providers/amazon/aws/transfers/sql_to_s3.py b/airflow/providers/amazon/aws/transfers/sql_to_s3.py index 2ddb0058dbed9..8cee9b6cffb15 100644 --- a/airflow/providers/amazon/aws/transfers/sql_to_s3.py +++ b/airflow/providers/amazon/aws/transfers/sql_to_s3.py @@ -80,6 +80,7 @@ class SqlToS3Operator(BaseOperator): CA cert bundle than the one used by botocore. :param file_format: the destination file format, only string 'csv', 'json' or 'parquet' is accepted. :param pd_kwargs: arguments to include in DataFrame ``.to_parquet()``, ``.to_json()`` or ``.to_csv()``. + :param groupby_kwargs: argument to include in DataFrame ``groupby()``. """ template_fields: Sequence[str] = ( From 340470fd82d566320adf0a9114254cc1a66eda7f Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Wed, 12 Apr 2023 10:40:14 +0530 Subject: [PATCH 7/7] Add operator to chain --- tests/system/providers/amazon/aws/example_sql_to_s3.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/system/providers/amazon/aws/example_sql_to_s3.py b/tests/system/providers/amazon/aws/example_sql_to_s3.py index 192e21f8d53ca..bda1e4c9dfbe8 100644 --- a/tests/system/providers/amazon/aws/example_sql_to_s3.py +++ b/tests/system/providers/amazon/aws/example_sql_to_s3.py @@ -175,7 +175,7 @@ def delete_security_group(sec_group_id: str, sec_group_name: str): # [START howto_transfer_sql_to_s3_with_groupby_param] sql_to_s3_task_with_groupby = SqlToS3Operator( - task_id="sql_to_s3_task", + task_id="sql_to_s3_with_groupby_task", sql_conn_id=conn_id_name, query=SQL_QUERY, s3_bucket=bucket_name, @@ -214,6 +214,7 @@ def delete_security_group(sec_group_id: str, sec_group_name: str): insert_data, # TEST BODY sql_to_s3_task, + sql_to_s3_task_with_groupby, # TEST TEARDOWN delete_bucket, delete_cluster,