This is an automated email from the ASF dual-hosted git repository.
dabla pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 70a78162261 Page SqlToS3Operator like GenericTransfer (#75730)
70a78162261 is described below
commit 70a781622611f09ec34f0bbc34c4f05c59ad4c67
Author: David Blain <[email protected]>
AuthorDate: Tue Oct 27 18:45:10 2026 +0100
Page SqlToS3Operator like GenericTransfer (#75730)
* Page SqlToS3Operator like GenericTransfer
A result that does not fit in memory could not be transferred with
SqlToS3Operator: it reads the whole query into one DataFrame and
max_rows_per_file only splits that frame afterwards. GenericTransfer
already pages a source with page_size and paginated_sql_statement_clause;
this gives SqlToS3Operator the same two parameters, one query per page
through the clause and one file per page (<key>_0, <key>_1, ...), so one
page is in memory at a time and the two operators are configured alike.
The per-DataFrame write and upload and the read with its dtype fix move
out of execute() into methods the paged and the unpaged path share.
* Convert a templated page_size before paging and raise ValueError for
argument clashes
A templated page_size reaches execute() as a string, and "0" was truthy
there:
the operator paged with a page size of zero and never ended. page_size is
now
converted after rendering, in render_template_fields() as GenericTransfer
does,
so a rendered "0" means unpaged; a negative value is refused.
The exclusion of page_size with max_rows_per_file or groupby_kwargs raises
ValueError, as the other Amazon transfers do for an argument clash, and
tests
the argument with "is not None" so the constructor applies no logic to a
template field.
The paged path is a generator that yields (suffix, DataFrame) pairs like
_partition_dataframe(), so execute() runs one loop over either source and a
page is released before the next one is read.
* Release a page before the next one is read
The generator dropped its reference to a page before reading the next one,
but the loop variable in execute() still held the previous page until the
next pair was unpacked, which happens after the read. Two pages were alive
at
the peak while the docs promise one. The loop now releases the page right
after its upload, and a test with weak references checks that no earlier
page is alive when a page is read.
---
providers/amazon/docs/transfer/sql_to_s3.rst | 19 +++
.../providers/amazon/aws/transfers/sql_to_s3.py | 130 ++++++++++----
.../tests/system/amazon/aws/example_sql_to_s3.py | 13 ++
.../unit/amazon/aws/transfers/test_sql_to_s3.py | 188 +++++++++++++++++++++
4 files changed, 319 insertions(+), 31 deletions(-)
diff --git a/providers/amazon/docs/transfer/sql_to_s3.rst
b/providers/amazon/docs/transfer/sql_to_s3.rst
index 9aa9db59384..54cccb06314 100644
--- a/providers/amazon/docs/transfer/sql_to_s3.rst
+++ b/providers/amazon/docs/transfer/sql_to_s3.rst
@@ -67,6 +67,25 @@ Example usage:
:start-after: [START howto_transfer_sql_to_s3_with_groupby_param]
:end-before: [END howto_transfer_sql_to_s3_with_groupby_param]
+Paging
+------
+
+A result that does not fit in memory is read in pages with the ``page_size``
param, the way
+``GenericTransfer`` pages a source: one query per page through
``paginated_sql_statement_clause``
+(``"{} LIMIT {} OFFSET {}"`` by default, formatted with the query, the page
size and the offset), one file
+per page (``<key>_0``, ``<key>_1``, ...), and only one page in memory at a
time. The transfer stops at the
+first page that comes back short. On most databases the query needs an ``ORDER
BY`` for the pages to be
+stable; a hook that serves pages itself, such as one streaming a single
result, needs none.
+``page_size`` cannot be combined with ``max_rows_per_file`` or
``groupby_kwargs``.
+
+Example usage:
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_sql_to_s3.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_transfer_sql_to_s3_with_page_size]
+ :end-before: [END howto_transfer_sql_to_s3_with_page_size]
+
Reference
---------
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py
b/providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py
index 5313aa26c4e..90ab13c0516 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py
@@ -36,6 +36,7 @@ from airflow.providers.common.compat.sdk import (
)
if TYPE_CHECKING:
+ import jinja2
import pandas as pd
import polars as pl
@@ -96,6 +97,14 @@ class SqlToS3Operator(BaseOperator):
:param df_kwargs: arguments to include in DataFrame ``.to_parquet()``,
``.to_json()`` or ``.to_csv()``,
or in ``.write_parquet()``, ``.write_json()`` or ``.write_csv()`` when
``df_type`` is ``'polars'``.
:param groupby_kwargs: argument to include in DataFrame ``groupby()``.
+ :param page_size: (optional) number of rows per page. The query is then
paged the way
+ ``GenericTransfer`` pages it, one query per page through
``paginated_sql_statement_clause``,
+ and every page is written as a file of its own (``<key>_0``,
``<key>_1``, ...), so only one
+ page is held in memory at a time. Cannot be combined with
``max_rows_per_file`` or
+ ``groupby_kwargs``.
+ :param paginated_sql_statement_clause: SQL statement clause used for
pagination, formatted with the
+ query, the page size and the offset (default ``"{} LIMIT {} OFFSET
{}"``). Only used with
+ ``page_size``.
"""
template_fields: Sequence[str] = (
@@ -104,6 +113,8 @@ class SqlToS3Operator(BaseOperator):
"query",
"sql_conn_id",
"aws_conn_id",
+ "page_size",
+ "paginated_sql_statement_clause",
)
template_ext: Sequence[str] = (".sql",)
template_fields_renderers = {
@@ -133,6 +144,8 @@ class SqlToS3Operator(BaseOperator):
df_kwargs: dict | None = None,
pd_kwargs: dict | None = None,
groupby_kwargs: dict | None = None,
+ page_size: int | None = None,
+ paginated_sql_statement_clause: str | None = None,
**kwargs,
) -> None:
super().__init__(**kwargs)
@@ -148,6 +161,8 @@ class SqlToS3Operator(BaseOperator):
self.groupby_kwargs = groupby_kwargs or {}
self.sql_hook_params = sql_hook_params
self.df_type = df_type
+ self.page_size = page_size
+ self.paginated_sql_statement_clause = paginated_sql_statement_clause
or "{} LIMIT {} OFFSET {}"
if read_pd_kwargs is not None:
warnings.warn(
@@ -175,6 +190,12 @@ class SqlToS3Operator(BaseOperator):
"can not be both specified. Please choose one."
)
+ if self.page_size is not None and (self.max_rows_per_file or
self.groupby_kwargs):
+ raise ValueError(
+ "SqlToS3Operator argument page_size can not be combined with "
+ "max_rows_per_file or groupby_kwargs: every page is already
written as a file of its own."
+ )
+
try:
self.file_format = FILE_FORMAT[file_format.upper()]
except KeyError:
@@ -244,49 +265,96 @@ class SqlToS3Operator(BaseOperator):
f"{trigger} needs {' or '.join(engines)}. Install
apache-airflow-providers-amazon[parquet]."
)
+ def get_paginated_sql(self, offset: int) -> str:
+ """Format the paginated SQL statement using the current format."""
+ return self.paginated_sql_statement_clause.format(self.query,
self.page_size, offset)
+
+ def render_template_fields(
+ self,
+ context: Context,
+ jinja_env: jinja2.Environment | None = None,
+ ) -> None:
+ super().render_template_fields(context=context, jinja_env=jinja_env)
+ # A templated page_size renders as a string; the paging decision needs
the number.
+ if self.page_size is not None:
+ self.page_size = int(self.page_size)
+ if self.page_size < 0:
+ raise ValueError(f"SqlToS3Operator argument page_size can not
be negative: {self.page_size}")
+
def execute(self, context: Context) -> None:
self._ensure_parquet_engine()
sql_hook = self._get_hook()
s3_conn = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
- data_df = sql_hook.get_df(
- sql=self.query, parameters=self.parameters, df_type=self.df_type,
**self.read_kwargs
+ pages = (
+ self._pages(sql_hook)
+ if self.page_size
+ else self._partition_dataframe(df=self._read_dataframe(sql_hook,
self.query))
)
+ for suffix, df in pages:
+ self._upload_dataframe(s3_conn, df, suffix)
+ del df # release the page before the next one is read; the loop
variable would hold it
+
+ def _pages(self, sql_hook: DbApiHook) -> Iterable[tuple[str, pd.DataFrame
| pl.DataFrame]]:
+ """
+ Read the query one page at a time, each page keyed by its number.
+
+ The first page is always yielded, so the key exists even for an empty
result; a later empty
+ page only tells that the previous one was the last. A page is released
before the next is read.
+ """
+ if TYPE_CHECKING:
+ assert self.page_size is not None # execute() pages only when
page_size is set
+ offset = 0
+ page = 0
+ while True:
+ paginated_sql = self.get_paginated_sql(offset)
+ self.log.info("Executing: \n %s", paginated_sql)
+ df = self._read_dataframe(sql_hook, paginated_sql)
+ if page == 0 or len(df):
+ yield str(page), df
+ if len(df) < self.page_size:
+ self.log.info("Last page read; ending transfer after %d
page(s).", page + 1)
+ return
+ del df
+ offset += self.page_size
+ page += 1
+
+ def _read_dataframe(self, sql_hook: DbApiHook, sql: str) -> pd.DataFrame |
pl.DataFrame:
+ """Run *sql* through the hook's DataFrame API and apply the pandas
dtype fixes."""
+ df = sql_hook.get_df(sql=sql, parameters=self.parameters,
df_type=self.df_type, **self.read_kwargs)
self.log.info("Data from SQL obtained")
# Only apply dtype fixes to pandas DataFrames since Polars doesn't
have the same NaN/None inconsistencies as panda
if ("dtype_backend", "pyarrow") not in self.read_kwargs.items() and
self.df_type == "pandas":
- self._fix_dtypes(data_df, self.file_format) # type:
ignore[arg-type]
+ self._fix_dtypes(df, self.file_format) # type: ignore[arg-type]
+ return df
+
+ def _upload_dataframe(self, s3_conn: S3Hook, df: pd.DataFrame |
pl.DataFrame, suffix: str) -> None:
+ """Write *df* in the configured file format and upload it as
``<key>_<suffix>``."""
file_options = FILE_OPTIONS_MAP[self.file_format]
write_function = file_options.polars_function if self.df_type ==
"polars" else file_options.function
+ buf = io.BytesIO()
+ self.log.info("Writing data to in-memory buffer")
+ clean_key = self._strip_suffixes(self.s3_key)
+ object_key = (
+ f"{clean_key}_{suffix}{file_options.suffix}" if suffix else
f"{clean_key}{file_options.suffix}"
+ )
- for group_name, df in self._partition_dataframe(df=data_df):
- buf = io.BytesIO()
- self.log.info("Writing data to in-memory buffer")
- clean_key = self._strip_suffixes(self.s3_key)
- object_key = (
- f"{clean_key}_{group_name}{file_options.suffix}"
- if group_name
- else f"{clean_key}{file_options.suffix}"
- )
-
- if self.file_format != FILE_FORMAT.PARQUET and
self.df_kwargs.get("compression") == "gzip":
- object_key += ".gz"
- df_kwargs = {k: v for k, v in self.df_kwargs.items() if k !=
"compression"}
- with gzip.GzipFile(fileobj=buf, mode="wb",
filename=object_key) as gz:
- getattr(df, write_function)(gz, **df_kwargs)
+ if self.file_format != FILE_FORMAT.PARQUET and
self.df_kwargs.get("compression") == "gzip":
+ object_key += ".gz"
+ df_kwargs = {k: v for k, v in self.df_kwargs.items() if k !=
"compression"}
+ with gzip.GzipFile(fileobj=buf, mode="wb", filename=object_key) as
gz:
+ getattr(df, write_function)(gz, **df_kwargs)
+ else:
+ if self.file_format == FILE_FORMAT.PARQUET:
+ getattr(df, write_function)(buf, **self.df_kwargs)
else:
- if self.file_format == FILE_FORMAT.PARQUET:
- getattr(df, write_function)(buf, **self.df_kwargs)
- else:
- text_buf = io.TextIOWrapper(buf, encoding="utf-8",
write_through=True)
- getattr(df, write_function)(text_buf, **self.df_kwargs)
- text_buf.flush()
-
- buf.seek(0)
-
- self.log.info("Uploading data to S3")
- s3_conn.load_file_obj(
- file_obj=buf, key=object_key, bucket_name=self.s3_bucket,
replace=self.replace
- )
+ text_buf = io.TextIOWrapper(buf, encoding="utf-8",
write_through=True)
+ getattr(df, write_function)(text_buf, **self.df_kwargs)
+ text_buf.flush()
+
+ buf.seek(0)
+
+ self.log.info("Uploading data to S3")
+ s3_conn.load_file_obj(file_obj=buf, key=object_key,
bucket_name=self.s3_bucket, replace=self.replace)
def _partition_dataframe(
self, df: pd.DataFrame | pl.DataFrame
diff --git a/providers/amazon/tests/system/amazon/aws/example_sql_to_s3.py
b/providers/amazon/tests/system/amazon/aws/example_sql_to_s3.py
index 103501e14aa..3205529768c 100644
--- a/providers/amazon/tests/system/amazon/aws/example_sql_to_s3.py
+++ b/providers/amazon/tests/system/amazon/aws/example_sql_to_s3.py
@@ -186,6 +186,18 @@ with DAG(
)
# [END howto_transfer_sql_to_s3_with_groupby_param]
+ # [START howto_transfer_sql_to_s3_with_page_size]
+ sql_to_s3_task_with_page_size = SqlToS3Operator(
+ task_id="sql_to_s3_with_page_size_task",
+ sql_conn_id=conn_id_name,
+ query=SQL_QUERY,
+ s3_bucket=bucket_name,
+ s3_key=key,
+ replace=True,
+ page_size=2,
+ )
+ # [END howto_transfer_sql_to_s3_with_page_size]
+
delete_bucket = S3DeleteBucketOperator(
task_id="delete_bucket",
bucket_name=bucket_name,
@@ -213,6 +225,7 @@ with DAG(
wait_cluster_available_before_transfer,
sql_to_s3_task,
sql_to_s3_task_with_groupby,
+ sql_to_s3_task_with_page_size,
# TEST TEARDOWN
delete_bucket,
delete_cluster,
diff --git a/providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py
b/providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py
index b97ee952b42..81890fdcabc 100644
--- a/providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py
+++ b/providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py
@@ -20,6 +20,7 @@ from __future__ import annotations
import gzip
import importlib.util
import io
+import weakref
from unittest import mock
import pandas as pd
@@ -31,6 +32,7 @@ from airflow.exceptions import
AirflowProviderDeprecationWarning
from airflow.models import Connection
from airflow.providers.amazon.aws.transfers.sql_to_s3 import SqlToS3Operator
from airflow.providers.common.compat.sdk import AirflowException,
AirflowOptionalProviderFeatureException
+from airflow.providers.common.sql.hooks.sql import DbApiHook
PANDAS_3_PLUS = Version(pd.__version__).major >= 3
@@ -500,6 +502,192 @@ class TestSqlToS3Operator:
)
)
+ @mock.patch("airflow.providers.amazon.aws.transfers.sql_to_s3.S3Hook")
+ def test_execute_with_page_size_writes_one_file_per_page(self,
mock_s3_hook):
+ """Paged like GenericTransfer: one query per page, one file per page,
stop on a short page."""
+ pages = [pd.DataFrame({"a": [1, 2]}), pd.DataFrame({"a": [3, 4]}),
pd.DataFrame({"a": [5]})]
+ mock_dbapi_hook = mock.Mock(return_value=mock.Mock(spec=DbApiHook))
+ get_df_mock = mock_dbapi_hook.return_value.get_df
+ get_df_mock.side_effect = pages
+ uploaded: dict[str, bytes] = {}
+ # Read the buffer at upload time: the operator's text wrapper closes
it once the upload is done.
+ mock_s3_hook.return_value.load_file_obj.side_effect = lambda file_obj,
key, **kwargs: uploaded.update(
+ {key: file_obj.getvalue()}
+ )
+
+ op = SqlToS3Operator(
+ query="SELECT * FROM t",
+ s3_bucket="bucket",
+ s3_key="key.csv",
+ sql_conn_id="mysql_conn_id",
+ aws_conn_id="aws_conn_id",
+ task_id="task_id",
+ replace=True,
+ df_kwargs={"index": False},
+ page_size=2,
+ dag=None,
+ )
+ op._get_hook = mock_dbapi_hook
+ op.execute(None)
+
+ assert [call.kwargs["sql"] for call in get_df_mock.call_args_list] == [
+ "SELECT * FROM t LIMIT 2 OFFSET 0",
+ "SELECT * FROM t LIMIT 2 OFFSET 2",
+ "SELECT * FROM t LIMIT 2 OFFSET 4",
+ ]
+ uploads = mock_s3_hook.return_value.load_file_obj.call_args_list
+ assert list(uploaded) == ["key_0.csv", "key_1.csv", "key_2.csv"]
+ assert all(call.kwargs["replace"] is True for call in uploads)
+ assert uploaded["key_2.csv"] == b"a\n5\n"
+
+ @pytest.mark.parametrize(
+ ("pages", "expected_keys"),
+ [
+ pytest.param(
+ [pd.DataFrame({"a": [1, 2]}), pd.DataFrame({"a": [3, 4]}),
pd.DataFrame({"a": []})],
+ ["key_0.csv", "key_1.csv"],
+ id="exact-multiple-ends-with-an-empty-page",
+ ),
+ pytest.param(
+ [pd.DataFrame({"a": []})], ["key_0.csv"],
id="empty-result-still-writes-the-first-page"
+ ),
+ ],
+ )
+ @mock.patch("airflow.providers.amazon.aws.transfers.sql_to_s3.S3Hook")
+ def test_execute_with_page_size_empty_pages(self, mock_s3_hook, pages,
expected_keys):
+ mock_dbapi_hook = mock.Mock(return_value=mock.Mock(spec=DbApiHook))
+ mock_dbapi_hook.return_value.get_df.side_effect = pages
+
+ op = SqlToS3Operator(
+ query="query",
+ s3_bucket="bucket",
+ s3_key="key",
+ sql_conn_id="mysql_conn_id",
+ aws_conn_id="aws_conn_id",
+ task_id="task_id",
+ page_size=2,
+ dag=None,
+ )
+ op._get_hook = mock_dbapi_hook
+ op.execute(None)
+
+ assert mock_dbapi_hook.return_value.get_df.call_count == len(pages)
+ uploads = mock_s3_hook.return_value.load_file_obj.call_args_list
+ assert [call.kwargs["key"] for call in uploads] == expected_keys
+
+ @mock.patch("airflow.providers.amazon.aws.transfers.sql_to_s3.S3Hook")
+ def test_execute_with_page_size_uses_the_pagination_clause(self,
mock_s3_hook):
+ """The clause is formatted with the query, the page size and the
offset, as GenericTransfer does."""
+ mock_dbapi_hook = mock.Mock(return_value=mock.Mock(spec=DbApiHook))
+ mock_dbapi_hook.return_value.get_df.side_effect = [pd.DataFrame({"a":
[1]})]
+
+ op = SqlToS3Operator(
+ query="SELECT * FROM t",
+ s3_bucket="bucket",
+ s3_key="key",
+ sql_conn_id="mysql_conn_id",
+ aws_conn_id="aws_conn_id",
+ task_id="task_id",
+ page_size="3",
+ paginated_sql_statement_clause="{0} OFFSET {2} ROWS FETCH NEXT {1}
ROWS ONLY",
+ dag=None,
+ )
+ op._get_hook = mock_dbapi_hook
+ op.render_template_fields({})
+ op.execute(None)
+
+ mock_dbapi_hook.return_value.get_df.assert_called_once_with(
+ sql="SELECT * FROM t OFFSET 0 ROWS FETCH NEXT 3 ROWS ONLY",
parameters=None, df_type="pandas"
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.transfers.sql_to_s3.S3Hook")
+ def
test_execute_with_page_size_releases_a_page_before_reading_the_next(self,
mock_s3_hook):
+ """Only one page is in memory at a time: the previous page is gone
when the next is read."""
+ pages_alive_at_read: list[int] = []
+ page_refs: list[weakref.ref] = []
+
+ def get_df(**kwargs):
+ pages_alive_at_read.append(sum(ref() is not None for ref in
page_refs))
+ df = pd.DataFrame({"a": [1, 2]}) if len(page_refs) < 2 else
pd.DataFrame({"a": [1]})
+ page_refs.append(weakref.ref(df))
+ return df
+
+ mock_dbapi_hook = mock.Mock(return_value=mock.Mock(spec=DbApiHook))
+ mock_dbapi_hook.return_value.get_df.side_effect = get_df
+
+ op = SqlToS3Operator(
+ query="query",
+ s3_bucket="bucket",
+ s3_key="key",
+ sql_conn_id="mysql_conn_id",
+ aws_conn_id="aws_conn_id",
+ task_id="task_id",
+ page_size=2,
+ dag=None,
+ )
+ op._get_hook = mock_dbapi_hook
+ op.execute(None)
+
+ assert pages_alive_at_read == [0, 0, 0]
+ assert mock_s3_hook.return_value.load_file_obj.call_count == 3
+
+ @pytest.mark.parametrize(
+ "kwargs",
+ [
+ pytest.param({"max_rows_per_file": 10}, id="max_rows_per_file"),
+ pytest.param({"groupby_kwargs": {"by": "a"}}, id="groupby_kwargs"),
+ ],
+ )
+ def test_page_size_excludes_the_other_ways_to_split_the_output(self,
kwargs):
+ with pytest.raises(ValueError, match="page_size can not be combined"):
+ SqlToS3Operator(
+ query="query",
+ s3_bucket="bucket",
+ s3_key="key",
+ sql_conn_id="mysql_conn_id",
+ task_id="task_id",
+ page_size=10,
+ **kwargs,
+ )
+
+ @pytest.mark.parametrize(
+ ("page_size", "expected"),
+ [
+ pytest.param("0", 0, id="rendered-zero-means-unpaged"),
+ pytest.param(3, 3, id="int-stays"),
+ pytest.param(None, None, id="unset-stays"),
+ ],
+ )
+ def test_render_template_fields_converts_page_size(self, page_size,
expected):
+ """A templated page_size renders as a string; "0" must mean unpaged,
not an endless loop."""
+ op = SqlToS3Operator(
+ query="query",
+ s3_bucket="bucket",
+ s3_key="key",
+ sql_conn_id="mysql_conn_id",
+ task_id="task_id",
+ page_size=page_size,
+ dag=None,
+ )
+
+ op.render_template_fields({})
+
+ assert op.page_size == expected
+
+ def test_render_template_fields_rejects_a_negative_page_size(self):
+ op = SqlToS3Operator(
+ query="query",
+ s3_bucket="bucket",
+ s3_key="key",
+ sql_conn_id="mysql_conn_id",
+ task_id="task_id",
+ page_size="-1",
+ dag=None,
+ )
+
+ with pytest.raises(ValueError, match="page_size can not be negative"):
+ op.render_template_fields({})
+
@mock.patch("airflow.providers.common.sql.operators.sql.BaseHook.get_connection")
def test_hook_params(self, mock_get_conn):
mock_get_conn.return_value = Connection(conn_id="postgres_test",
conn_type="postgres")
