GitHub user jtakish added a comment to the discussion: SAP HANA Provider Package
This might not be the place for it, let me know if you want me to move it to a
different thread.
I was looking at the code for the GenericTransfer operator today. It uses
fetchall and exponentially growing LIMIT + OFFSET to batch data. The newest
version in Airflow 3.2.2, uses deferrals and async to run all of the statements
concurrently. Each statement triggers a full table scan, the rows have to get
ordered, and then all rows not in the LIMIT get discarded.
I have a get_rows_by_chunk method in the HANA hook. It yields batches from a
generator.
~~~
def get_records_by_chunks(
self, sql: str, parameters: Iterable[Any] | Mapping[str, Any] | None =
None, chunksize: int = 10000
) -> Iterator[tuple[Any, ...] | list[tuple[Any, ...]]]:
"""
Streams records from SAP HANA, yielding chunks of rows.
This method allows for fetching large datasets without loading them all
into memory. Each record is passed through
``_make_common_data_structure``
to ensure JSON serialization. The ``descriptions`` and
``last_description``
attributes are available immediately after execution.
:param sql: The SQL statement.
:param parameters: Parameters to bind to the SQL statement.
:param chunksize: The number of records per chunk.
:return: A generator yielding lists of tuples (or a single tuple if
chunksize is 1).
"""
self.descriptions = []
conn = None
cur = None
try:
conn = self.get_conn()
cur = conn.cursor()
self._run_command(cur, sql, parameters)
self.descriptions.append(cur.description)
except Exception as e:
if cur:
cur.close()
if conn:
conn.close()
raise e
return chunk_handler(self, conn, cur, chunksize)
~~~
~~~
def chunk_handler(
hook: T, conn: Any, cursor: Any, chunksize: int
) -> Iterator[tuple[Any, ...] | list[tuple[Any, ...]]]:
"""
Yield rows in batches.
This allows for processing large datasets without loading all data into
memory.
The ``descriptions`` and ``last_description`` attributes of the hook are
available immediately after calling this method.
:param hook: The ``DbApiHook`` instance.
:param conn: The database connection object.
:param cursor: The database cursor.
:param chunksize: The number of records to return per chunk.
:return: A generator yielding lists of tuples.
"""
nb_rows = 0
make_common_data_structure = getattr(hook, "_make_common_data_structure")
log = getattr(hook, "log")
try:
while results := make_common_data_structure(fetch_many_handler(cursor,
chunksize)):
nb_rows += len(results)
log.info("Fetched %s rows so far", nb_rows)
yield results
else:
log.info("Done fetching. Fetched %s total rows", nb_rows)
finally:
if cursor:
cursor.close()
if conn:
conn.close()
~~~
I think there is a way to maybe improve the generic transfer using a similar
logic? I realize the generic transfer is intended to be used only for "results
that fit into memory" but that is a little ambiguous. Also using a generator
would essentially remove this restriction.
GitHub link:
https://github.com/apache/airflow/discussions/44768#discussioncomment-17542082
----
This is an automatically sent email for [email protected].
To unsubscribe, please send an email to: [email protected]
