This is an automated email from the ASF dual-hosted git repository.
shahar1 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 8e0ccacbb2d ElasticsearchSQLHook: add chunked Polars DataFrame support
via custom SQL reader (#68411)
8e0ccacbb2d is described below
commit 8e0ccacbb2d3f1aa43226d2abc860d06dac43749
Author: SameerMesiah97 <[email protected]>
AuthorDate: Mon Jun 29 13:40:22 2026 +0100
ElasticsearchSQLHook: add chunked Polars DataFrame support via custom SQL
reader (#68411)
Implement _get_polars_df_by_chunks using Elasticsearch SQL cursor-based
pagination and add unit tests covering chunking behavior, pagination
across cursor pages, and cursor cleanup.
---
.../providers/elasticsearch/hooks/elasticsearch.py | 32 +++-
.../airflow/providers/elasticsearch/utils/sql.py | 88 ++++++++++-
.../unit/elasticsearch/hooks/test_elasticsearch.py | 25 +++
.../tests/unit/elasticsearch/utils/test_sql.py | 176 ++++++++++++++++++++-
4 files changed, 316 insertions(+), 5 deletions(-)
diff --git
a/providers/elasticsearch/src/airflow/providers/elasticsearch/hooks/elasticsearch.py
b/providers/elasticsearch/src/airflow/providers/elasticsearch/hooks/elasticsearch.py
index 7df6ec55916..d66ea9cc1cc 100644
---
a/providers/elasticsearch/src/airflow/providers/elasticsearch/hooks/elasticsearch.py
+++
b/providers/elasticsearch/src/airflow/providers/elasticsearch/hooks/elasticsearch.py
@@ -29,7 +29,7 @@ from elasticsearch import Elasticsearch
from airflow.providers.common.compat.sdk import BaseHook
from airflow.providers.common.sql.hooks.sql import DbApiHook
from airflow.providers.elasticsearch._compat import apply_compat_with
-from airflow.providers.elasticsearch.utils.sql import read_sql_to_polars
+from airflow.providers.elasticsearch.utils.sql import read_sql_to_polars,
read_sql_to_polars_by_chunks
if TYPE_CHECKING:
from elastic_transport import ObjectApiResponse
@@ -283,6 +283,36 @@ class ElasticsearchSQLHook(DbApiHook):
**kwargs,
)
+ def _get_polars_df_by_chunks(
+ self,
+ sql,
+ parameters=None,
+ *,
+ chunksize: int,
+ **kwargs,
+ ):
+ """
+ Execute an Elasticsearch SQL query and return the results as chunked
Polars DataFrames.
+
+ This method uses Elasticsearch SQL cursor-based pagination instead of
DB-API,
+ as Elasticsearch is not fully compatible with polars.read_database.
+
+ :param sql: SQL query string
+ :param parameters: Optional query parameters
+ :param chunksize: Number of rows per yielded DataFrame
+ :param kwargs: Additional arguments passed to the underlying reader
+ :return: Generator of polars.DataFrame objects
+ """
+ client = self.get_conn().es
+
+ yield from read_sql_to_polars_by_chunks(
+ client=client,
+ query=sql,
+ params=parameters,
+ chunksize=chunksize,
+ **kwargs,
+ )
+
class ElasticsearchPythonHook(BaseHook):
"""
diff --git
a/providers/elasticsearch/src/airflow/providers/elasticsearch/utils/sql.py
b/providers/elasticsearch/src/airflow/providers/elasticsearch/utils/sql.py
index 1cc954949ca..c67e0ce3905 100644
--- a/providers/elasticsearch/src/airflow/providers/elasticsearch/utils/sql.py
+++ b/providers/elasticsearch/src/airflow/providers/elasticsearch/utils/sql.py
@@ -19,7 +19,7 @@
from __future__ import annotations
import logging
-from collections.abc import Iterable, Mapping
+from collections.abc import Generator, Iterable, Mapping
from typing import TYPE_CHECKING, Any
from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
@@ -106,3 +106,89 @@ def read_sql_to_polars(
log.debug("Failed to clear Elasticsearch SQL cursor",
exc_info=True)
return pl.DataFrame(rows, schema=columns, orient="row", strict=False)
+
+
+def read_sql_to_polars_by_chunks(
+ client: Elasticsearch,
+ query: str,
+ params: Mapping[str, Any] | Iterable | None = None,
+ fetch_size: int = 1000,
+ chunksize: int = 1000,
+) -> Generator[pl.DataFrame, None, None]:
+ """
+ Execute an Elasticsearch SQL query and return results as chunked Polars
DataFrames.
+
+ This uses Elasticsearch SQL cursor-based pagination instead of DB-API,
+ as Elasticsearch does not provide a fully compliant DB-API interface.
+
+ :param client: Elasticsearch client
+ :param query: SQL query string
+ :param params: Optional query parameters
+ :param fetch_size: Number of rows fetched per Elasticsearch request
+ :param chunksize: Maximum number of rows per yielded DataFrame
+ :return: Generator of Polars DataFrames
+ """
+ body: dict[str, Any] = {
+ "query": query,
+ "fetch_size": fetch_size,
+ }
+
+ try:
+ import polars as pl
+ except ImportError:
+ raise AirflowOptionalProviderFeatureException(
+ "Polars support requires installing the 'polars' extra: "
+ "pip install apache-airflow-providers-elasticsearch[polars]"
+ ) from None
+
+ if params:
+ body["params"] = params
+
+ response = client.sql.query(**body)
+
+ columns_meta = response.get("columns", [])
+ columns = [col["name"] for col in columns_meta]
+
+ cursor = response.get("cursor")
+ last_cursor = cursor
+
+ rows: list[list[Any]] = list(response.get("rows", []))
+
+ try:
+ while rows:
+ if len(rows) >= chunksize:
+ yield pl.DataFrame(
+ rows[:chunksize],
+ schema=columns,
+ orient="row",
+ strict=False,
+ )
+ rows = rows[chunksize:]
+ continue
+
+ if not cursor:
+ yield pl.DataFrame(
+ rows,
+ schema=columns,
+ orient="row",
+ strict=False,
+ )
+ break
+
+ response = client.sql.query(cursor=cursor)
+
+ rows.extend(response.get("rows", []))
+
+ cursor = response.get("cursor")
+ if cursor:
+ last_cursor = cursor
+
+ finally:
+ if last_cursor:
+ try:
+ client.sql.clear_cursor(cursor=last_cursor)
+ except Exception:
+ log.debug(
+ "Failed to clear Elasticsearch SQL cursor",
+ exc_info=True,
+ )
diff --git
a/providers/elasticsearch/tests/unit/elasticsearch/hooks/test_elasticsearch.py
b/providers/elasticsearch/tests/unit/elasticsearch/hooks/test_elasticsearch.py
index 3c3247aeec6..85e56a358fe 100644
---
a/providers/elasticsearch/tests/unit/elasticsearch/hooks/test_elasticsearch.py
+++
b/providers/elasticsearch/tests/unit/elasticsearch/hooks/test_elasticsearch.py
@@ -276,6 +276,31 @@ class TestElasticsearchSQLHook:
max_rows=10,
)
+
@mock.patch("airflow.providers.elasticsearch.hooks.elasticsearch.read_sql_to_polars_by_chunks")
+ def test_get_df_polars_df_by_chunks(self, mock_reader):
+
+ self.conn.es = MagicMock()
+
+ generator = iter(["df"])
+ mock_reader.return_value = generator
+
+ result = self.db_hook.get_df_by_chunks(
+ sql="SELECT 1",
+ df_type="polars",
+ fetch_size=100,
+ chunksize=100,
+ )
+
+ assert list(result) == ["df"]
+
+ mock_reader.assert_called_once_with(
+ client=self.conn.es,
+ query="SELECT 1",
+ params=None,
+ fetch_size=100,
+ chunksize=100,
+ )
+
def test_run(self):
statement = "SELECT * FROM hollywood.actors"
diff --git a/providers/elasticsearch/tests/unit/elasticsearch/utils/test_sql.py
b/providers/elasticsearch/tests/unit/elasticsearch/utils/test_sql.py
index 53a23fedeea..046657fbb17 100644
--- a/providers/elasticsearch/tests/unit/elasticsearch/utils/test_sql.py
+++ b/providers/elasticsearch/tests/unit/elasticsearch/utils/test_sql.py
@@ -22,9 +22,7 @@ from unittest.mock import MagicMock
import pytest
-from airflow.providers.elasticsearch.utils.sql import (
- read_sql_to_polars,
-)
+from airflow.providers.elasticsearch.utils.sql import read_sql_to_polars,
read_sql_to_polars_by_chunks
COLUMNS = [
{"name": "id", "type": "integer"},
@@ -179,3 +177,175 @@ def test_read_sql_to_polars_no_cursor_cleanup():
read_sql_to_polars(es, "SELECT *")
es.sql.clear_cursor.assert_not_called()
+
+
+def test_read_sql_to_polars_by_chunks_single_chunk():
+ es = _mock_es(
+ [
+ {
+ "columns": COLUMNS,
+ "rows": [[1, "a"], [2, "b"]],
+ }
+ ]
+ )
+
+ chunks = list(
+ read_sql_to_polars_by_chunks(
+ es,
+ "SELECT *",
+ chunksize=10,
+ )
+ )
+
+ assert len(chunks) == 1
+
+ df = chunks[0]
+ assert df.shape == (2, 2)
+ assert df.to_dict(as_series=False) == {
+ "id": [1, 2],
+ "name": ["a", "b"],
+ }
+
+
+def test_read_sql_to_polars_by_chunks_single_page_multiple_chunks():
+ es = _mock_es(
+ [
+ {
+ "columns": COLUMNS,
+ "rows": [
+ [1, "a"],
+ [2, "b"],
+ [3, "c"],
+ ],
+ }
+ ]
+ )
+
+ chunks = list(
+ read_sql_to_polars_by_chunks(
+ es,
+ "SELECT *",
+ chunksize=2,
+ )
+ )
+
+ assert [chunk.shape for chunk in chunks] == [
+ (2, 2),
+ (1, 2),
+ ]
+
+ assert chunks[0].to_dict(as_series=False) == {
+ "id": [1, 2],
+ "name": ["a", "b"],
+ }
+
+ assert chunks[1].to_dict(as_series=False) == {
+ "id": [3],
+ "name": ["c"],
+ }
+
+
+def test_read_sql_to_polars_by_chunks_across_cursor_pages():
+ es = _mock_es(
+ [
+ {
+ "columns": COLUMNS,
+ "rows": [[1, "a"], [2, "b"]],
+ "cursor": "cursor_1",
+ },
+ {
+ "rows": [[3, "c"], [4, "d"]],
+ "cursor": None,
+ },
+ ]
+ )
+
+ chunks = list(
+ read_sql_to_polars_by_chunks(
+ es,
+ "SELECT *",
+ chunksize=3,
+ )
+ )
+
+ assert len(chunks) == 2
+
+ assert chunks[0].to_dict(as_series=False) == {
+ "id": [1, 2, 3],
+ "name": ["a", "b", "c"],
+ }
+
+ assert chunks[1].to_dict(as_series=False) == {
+ "id": [4],
+ "name": ["d"],
+ }
+
+
+def test_read_sql_to_polars_by_chunks_clears_cursor():
+ es = _mock_es(
+ [
+ {
+ "columns": COLUMNS,
+ "rows": [[1, "a"]],
+ "cursor": "cursor_1",
+ },
+ {
+ "rows": [[2, "b"]],
+ "cursor": None,
+ },
+ ]
+ )
+
+ list(
+ read_sql_to_polars_by_chunks(
+ es,
+ "SELECT *",
+ chunksize=1,
+ )
+ )
+
+ es.sql.clear_cursor.assert_called_once()
+
+
+def test_read_sql_to_polars_by_chunks_no_cursor_cleanup():
+ es = _mock_es(
+ [
+ {
+ "columns": COLUMNS,
+ "rows": [[1, "a"]],
+ }
+ ]
+ )
+
+ list(
+ read_sql_to_polars_by_chunks(
+ es,
+ "SELECT *",
+ chunksize=1,
+ )
+ )
+
+ es.sql.clear_cursor.assert_not_called()
+
+
+def test_read_sql_to_polars_clears_cursor_on_pagination_error():
+ es = MagicMock()
+
+ es.sql.query.side_effect = [
+ {
+ "columns": COLUMNS,
+ "rows": [[1, "a"]],
+ "cursor": "cursor_1",
+ },
+ RuntimeError("boom"),
+ ]
+
+ with pytest.raises(RuntimeError, match="boom"):
+ list(
+ read_sql_to_polars_by_chunks(
+ es,
+ "SELECT *",
+ )
+ )
+
+ es.sql.clear_cursor.assert_called_once()
