Feature add get polars df to dbapihook and bigqueryhook

airflow/providers/common/sql/hooks/sql.py

@@ -218,6 +218,26 @@ def get_pandas_df(self, sql, parameters: Iterable | Mapping[str, Any] | None = N
         with closing(self.get_conn()) as conn:
             return psql.read_sql(sql, con=conn, params=parameters, **kwargs)
 
+    def get_polars_df(self, sql, **kwargs):
+        """
+        Executes the sql and returns a polars dataframe.
+
+        :param sql: the sql statement to be executed (str) or a list of
+            sql statements to execute
+        :param parameters: The parameters to render the SQL query with.
+        :param kwargs: (optional) passed into polars.read_database method
+        """
+        try:
+            import polars as pl
+        except ImportError:
+            raise Exception(
+                "polars library not installed, run: pip install "
+                "'apache-airflow-providers-common-sql[polars]'."
+            )
+
+        with closing(self.get_conn()) as conn:
+            return pl.read_database(sql, connection=conn, **kwargs)
+
     def get_pandas_df_by_chunks(
         self, sql, parameters: Iterable | Mapping[str, Any] | None = None, *, chunksize: int | None, **kwargs
     ):

airflow/providers/google/cloud/hooks/bigquery.py

@@ -42,13 +42,15 @@
     QueryJob,
     SchemaField,
     UnknownJob,
+    QueryJobConfig,
 )
 from google.cloud.bigquery.dataset import AccessEntry, Dataset, DatasetListItem, DatasetReference
 from google.cloud.bigquery.table import EncryptionConfiguration, Row, RowIterator, Table, TableReference
 from google.cloud.exceptions import NotFound
 from googleapiclient.discovery import Resource, build
 from pandas_gbq import read_gbq
 from pandas_gbq.gbq import GbqConnector  # noqa
+import polars as pl
 from requests import Session
 from sqlalchemy import create_engine
 
@@ -270,6 +272,36 @@ def get_pandas_df(
 
         return read_gbq(sql, project_id=project_id, dialect=dialect, credentials=credentials, **kwargs)
 
+    def get_polars_df(
+        self,
+        sql: str,
+        parameters: Iterable | Mapping[str, Any] | None = None,
+        dialect: str | None = None,
+        **kwargs,
+    ) -> pd.DataFrame:
+        """Get a Polars DataFrame for the BigQuery results.
+
+        :param sql: The BigQuery SQL to execute.
+        :param parameters: The parameters to render the SQL query with (not
+            used, leave to override superclass method)
+        :param dialect: Dialect of BigQuery SQL – legacy SQL or standard SQL
+            defaults to use `self.use_legacy_sql` if not specified
+        :param kwargs: (optional) passed into polars.from_arrow method
+        """
+        if dialect is None:
+            dialect = "legacy" if self.use_legacy_sql else "standard"
+
+        project_id = self.get_project_id()
+        client = self.get_client(project_id=project_id)
+
+        job_config = QueryJobConfig(dialect=dialect)  # Specify the SQL dialect here
+        query_job = client.query(sql, job_config=job_config)  # API request
+        rows = query_job.result()  # Waits for query to finish
+
+        df = pl.from_arrow(rows.to_arrow(), **kwargs)
+
+        return df
+
     @GoogleBaseHook.fallback_to_default_project_id
     def table_exists(self, dataset_id: str, table_id: str, project_id: str) -> bool:
         """Check if a table exists in Google BigQuery.

airflow/providers/google/common/hooks/base_google.py

@@ -296,6 +296,11 @@ def get_credentials(self) -> google.auth.credentials.Credentials:
         credentials, _ = self.get_credentials_and_project_id()
         return credentials
 
+    def get_project_id(self) -> str:
+        """Returns the project_id str for Google API."""
+        _, project_id = self.get_credentials_and_project_id()
+        return project_id
+
     def _get_access_token(self) -> str:
         """Returns a valid access token from Google API Credentials."""
         credentials = self.get_credentials()

airflow/providers/google/provider.yaml

@@ -1207,6 +1207,9 @@ additional-extras:
   - name: amazon
     dependencies:
       - apache-airflow-providers-amazon>=2.6.0
+  - name: polars
+    dependencies:
+      - polars>=0.19.5
 
 secrets-backends:
   - airflow.providers.google.cloud.secrets.secret_manager.CloudSecretManagerBackend