Feature add get polars df to dbapihook and bigqueryhook #34679

bfeif · 2023-09-29T06:00:00Z

Changes included in this PR:

Added get_polars_df() method to BigQueryHook and DbApiHook.
Added polars to required packages.

Closes #33911.

Note to Reviewer: This is my first PR in Airflow, and my first time encountering mock. Can you provide any guidance or resources on creating tests for these new functions? I looked to create an analogous test to test_get_pandas_df() in airflow/tests/providers/google/cloud/hooks/test_bigquery.py, but wasn't sure how to proceed. #newbiequestion

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

boring-cyborg · 2023-09-29T06:00:05Z

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contribution Guide (https://github.com/apache/airflow/blob/main/CONTRIBUTING.rst)
Here are some useful points:

Pay attention to the quality of your code (ruff, mypy and type annotations). Our pre-commits will help you with that.
In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example DAG that shows how users should use it.
Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
Be sure to read the Airflow Coding style.
Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
Apache Airflow is a community-driven project and together we are making it better 🚀.
In case of doubts contact the developers at:
Mailing List: dev@airflow.apache.org
Slack: https://s.apache.org/airflow-slack

eladkal · 2023-09-29T09:35:26Z

setup.py

    "bcrypt>=2.0.0",
    "flask-bcrypt>=0.7.1",
 ]
+polars = ["polars>=0.19.5"]


I dont think this is right(?)
Polars should be extra dependency of common.sql provider (it should be defined in provider.yaml)

I know we have pandas here but I think its due to legacy/backward compatible (pandas used to be part of Airflow core)

You are probably correct, @eladkal ; I mostly just mirrored what I saw pandas doing.

As for common.sql, you mentioned about it in the issue as well, but I didn't understand what you meant there. Now that you say provider.yaml, I suppose I get it, you mean like here?

Not quite... by having it as dependency of the provider you force every user of the provider to install this lib this is not desired.
What we should do is have it as optional extra for the provider. Thus you need to define it in:

airflow/airflow/providers/google/provider.yaml

Line 1191 in 23e2c95

additional-extras:

then users can have it when they explicitly ask to install the extra :
pip install apache-airflow-providers-google[polars]

But I think this extra dependency needs to be set in the common.sql provider not in google one?

Hm okay. I've removed polars from core dependencies (I think), and added it to airflow/airflow/providers/google/provider.yaml.

Sorry, I'm still not clear on which you file you mean by "the common.sql provider", what's the precise path to which you're referring?

@eladkal can you please advise here?

Taragolis · 2023-09-29T10:34:58Z

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

+    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)


I'm not sure how for polaris, but for pandas the same method is broken. Let me explain:

Pandas expect to get SQLAlchemy connection rather than DBAPI with only one exception for SQLite

get_sqlalchemy_engine is also broken in some cases because get_url is broken, see discussion in dev list: https://lists.apache.org/thread/8rhmz3qh30hvkondct4sfmgk4vd07mn5

Maybe it is not a case for polaris I'm not familiar with this lib

In the issue it was suggested to have get_df (generic function) that accept as parameter if it should be pandas/polars - why did you choose not to have it eventually?

I was suggest that we should not merge something in single implementation especially if no one know how pandas and polaris interface compatible. I just worried that if we merge it now, than it could turned into the something like BackfillJobRunner._backfill_job_runner.py which literally have 170 different statements.

And in additional pandas is broken for at least half implementations: #34679 (comment), so my vision that we should fix it first before we could create some generic method

@Taragolis the polars doc on read_database() states the following:

This function supports a wide range of native database drivers (ranging from local databases such as SQLite to large cloud databases such as Snowflake), as well as generic libraries such as ADBC, SQLAlchemy and various flavours of ODBC.

... and the connection object eventually gets passed in the constructor to the polars ConnectionExecutor here. Does this answer your question?

In the issue it was suggested to have get_df (generic function) that accept as parameter if it should be pandas/polars - why did you choose not to have it eventually?

@eladkal you are right (link to suggestion in issue), this was suggested.

I made this decision because the APIs of polars and pandas are quite syntactically different for the actual querying, so the get_df generic function would have just been a big switch statement, which seemed to me to add no value (this syntactic difference is larger on the BigQueryHook than on the DbApiHook, I do admit).

Furthermore, I'm not sure why we'd want to add get_df without also removing (or at least removing exposure of) the underlying methods get_polars_df and get_pandas_df; having get_df, get_polars_df, and get_pandas. Doing this removal, however, would be a breaking change for a lot of users' code.

It seemed to me like such a get_df redefinition may fall better under the scope of a refactoring ticket.

We don't have to remove anything. We can deprecate and raise deprecation warning.
For such deprecation we probably will give very long time before actually removing it

It still seems to me to be over-engineering, at least for now @eladkal.

github-actions · 2023-12-07T00:11:42Z

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.

bfeif added 2 commits September 28, 2023 22:28

Added get_polars_df to DbApiHook.

d956581

Added get_polars_df to BigQueryHook.

810c066

boring-cyborg bot added area:dev-tools area:providers kind:documentation provider:common-sql provider:google Google (including GCP) related issues labels Sep 29, 2023

bfeif mentioned this pull request Sep 29, 2023

Add get_polars_df() method to BigQueryHook. #33911

Closed

2 tasks

eladkal reviewed Sep 29, 2023

View reviewed changes

eladkal requested review from Taragolis and josh-fell September 29, 2023 09:37

Taragolis reviewed Sep 29, 2023

View reviewed changes

Removed polars from core dependencies.

bb1d61a

github-actions bot added the stale Stale PRs per the .github/workflows/stale.yml policy file label Dec 7, 2023

github-actions bot closed this Dec 12, 2023

guan404ming mentioned this pull request Apr 8, 2025

Integrate polars in get_df, get_df_by_chunks #48875

Merged

eladkal mentioned this pull request Apr 16, 2025

Status of testing Providers that were prepared on April 16, 2025 #49354

Closed

56 tasks

Feature add get polars df to dbapihook and bigqueryhook #34679

Feature add get polars df to dbapihook and bigqueryhook #34679

Uh oh!

Conversation

bfeif commented Sep 29, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

boring-cyborg bot commented Sep 29, 2023

Uh oh!

eladkal Sep 29, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

eladkal Sep 29, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bfeif Sep 29, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bfeif Sep 29, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bfeif Oct 22, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Dec 7, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

bfeif commented Sep 29, 2023 •

edited

Loading

eladkal Sep 29, 2023 •

edited

Loading

eladkal Sep 29, 2023 •

edited

Loading

bfeif Sep 29, 2023 •

edited

Loading

bfeif Sep 29, 2023 •

edited

Loading

bfeif Oct 22, 2023 •

edited

Loading