Migrate `get_pandas_df` to `get_df` in `provider` test by guan404ming · Pull Request #49339 · apache/airflow

guan404ming · 2025-04-16T08:39:54Z

Related Issue

Migrate get_pandas_df to get_df in providers #49334

Why

get_pandas and get_pandas_by_chunks is deprecated in #48875, which would be replaced by get_df. Thus, we need to migrate them in providers.

How

This PR is focus on migration of test migration only case as they are more simple and general. Thus I thought that put them together would be better.

Apache Drill
Apache Druid
Apache Impala
Apache Pinot
Elasticsearch
SQLite
Vertica

In this test, the original get_pandas_df would still be tested since it's implemented like

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

Lines 387 to 400 in 438220f

    
               def get_pandas_df( 
        
                   self, 
        
                   sql, 
        
                   parameters: list | tuple | Mapping[str, Any] | None = None, 
        
                   **kwargs, 
        
               ) -> DataFrame: 
        
                   """ 
        
                   Execute the sql and returns a pandas 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 pandas.io.sql.read_sql method 
        
                   """ 
        
                   return self._get_pandas_df(sql, parameters, **kwargs)

and get_df is like this

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

Lines 433 to 436 in 438220f

    
           if df_type == "pandas": 
        
               return self._get_pandas_df(sql, parameters, **kwargs) 
        
           if df_type == "polars": 
        
               return self._get_polars_df(sql, parameters, **kwargs)

thus it still handled the backwards compatibility.

cc: @eladkal

^ 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 airflow-core/newsfragments.

providers/apache/drill/tests/unit/apache/drill/hooks/test_drill.py

guan404ming · 2025-04-19T11:38:23Z

Hi @eladkal after my investigation, I found that the cursor execution logic is different for pandas and polars. Reference to the polars.read_database source code, there is a abstraction for cursor execution when calling read_database:

https://github.com/pola-rs/polars/blob/5c114f0f32a62bf796629e4ccf01b93146ae22c4/py-polars/polars/io/database/functions.py#L247-L256 and the executer https://github.com/pola-rs/polars/blob/main/py-polars/polars/io/database/_executor.py

Thus we have also need to mock the execution abstraction for polars, while pandas just do cursor.fetchall() directly which meaning that the current implement for test is suitable for pandas. Also, we need to skip the polars test for low_deps and AF2 in ci, so I split them into two functions rather than using @pytest.mark.parametrize to handle the test here.

guan404ming · 2025-04-19T11:53:31Z

I currently removed the test for polars for Elasticsearchin this PR, since Elasticsearch provider has its own ElasticsearchSQLCursor which is not compatable with polars cursor execution. I think the modification which the Elasticsearch need is maybe quite different with the other ones in this PR thus I would try open another one to deal with Elasticsearch and I have added it in todo list in #49334

devel-common/src/tests_common/test_utils/markers.py

potiuk · 2025-04-23T19:33:58Z

I do not think we should base it on polars importability. We should not check for Airflow 3 here, but instead we should put minimum common.sql version (to the new version that will be released ) in all providers that are going to use the new feature. This is a "common.sql" not airflow feature, so theoretically when you install the new provider and they have common.sql >=NEW_VERSION the tests should also work for compatibility tests and polars should be installed - because we are going to install the new common.sql there.

Or am I missing something?.

potiuk · 2025-04-28T08:52:06Z

There is still one problem left - compat tests - because there polars is not present as extra in airflow. But what we can do, we can somewhat modify the scripts to add polars extra to common.sql - cheating a bit.

In this script : https://github.com/apache/airflow/blob/main/scripts/in_container/install_airflow_and_providers.py#L720 we are looping through providers available in /dist folder and install them. So if find apache_airflow_providers_common_sql .whl file there we could add [polars] to it.

guan404ming · 2025-04-28T09:04:38Z

I agree with you.

Initially, I wanted to bump the Apache Airflow version to >= 3.1.0 in the providers. However, it seems like that is not the best approach because it will break compatibility. Considering this is not "really" incompatible but lack of package for test, I thought we can proceed with that.

potiuk · 2025-04-28T09:28:27Z

I agree with you.

Initially, I wanted to bump the Apache Airflow version to >= 3.1.0 in the providers. However, it seems like that is not the best approach because it will break compatibility. Considering this is not "really" incompatible but lack of package for test, I thought we can proceed with that.

Nice.. Looks good. Let's see if CI agrees with this approach :)

guan404ming · 2025-04-28T10:12:02Z

It works!! 🎉
Force push for the conflict.

Co-authored-by: Elad Kalif <[email protected]> Co-authored-by: Jarek Potiuk <[email protected]>

eladkal · 2025-04-28T10:18:39Z

It works!! 🎉 Force push for the conflict.

🎉

potiuk · 2025-04-28T10:36:21Z

It works!! 🎉 Force push for the conflict.

🎉

Nice!!!!

guan404ming · 2025-04-28T11:03:36Z

Thanks for the review and suggestions!

github-actions · 2025-04-28T11:03:58Z

Backport failed to create: v3-0-test. View the failure log Run details

Status	Branch	Result
❌	v3-0-test

You can attempt to backport this manually by running:

cherry_picker c6ca8a2 v3-0-test

This should apply the commit to the v3-0-test branch and leave the commit in conflict state marking
the files that need manual conflict resolution.

After you have resolved the conflicts, you can continue the backport process by running:

cherry_picker --continue

…est (apache#49339) * test: migrate get_pandas_df to get_df in provider test Co-authored-by: Elad Kalif <[email protected]> Co-authored-by: Jarek Potiuk <[email protected]> * chore: add polars in deps * fix: install polars in ci --------- Co-authored-by: Elad Kalif <[email protected]> Co-authored-by: Jarek Potiuk <[email protected]> Co-authored-by: Wesley Chiu <[email protected]> (cherry picked from commit c6ca8a2)

…est (#49339) (#49885) * test: migrate get_pandas_df to get_df in provider test * chore: add polars in deps * fix: install polars in ci --------- (cherry picked from commit c6ca8a2) Co-authored-by: Guan Ming(Wesley) Chiu <[email protected]> Co-authored-by: Elad Kalif <[email protected]> Co-authored-by: Wesley Chiu <[email protected]>

…49339) * test: migrate get_pandas_df to get_df in provider test Co-authored-by: Elad Kalif <[email protected]> Co-authored-by: Jarek Potiuk <[email protected]> * chore: add polars in deps * fix: install polars in ci --------- Co-authored-by: Elad Kalif <[email protected]> Co-authored-by: Jarek Potiuk <[email protected]> Co-authored-by: Wesley Chiu <[email protected]>

boring-cyborg bot added area:providers provider:apache-druid provider:apache-drill provider:apache-pinot provider:apache-impala provider:elasticsearch provider:sqlite provider:vertica labels Apr 16, 2025

guan404ming mentioned this pull request Apr 16, 2025

Migrate get_pandas_df to get_df in providers #49334

Closed

16 tasks

eladkal self-requested a review April 17, 2025 05:11

eladkal reviewed Apr 17, 2025

View reviewed changes

providers/apache/drill/tests/unit/apache/drill/hooks/test_drill.py Show resolved Hide resolved

guan404ming force-pushed the miragte-test branch 5 times, most recently from 2449640 to 84979e9 Compare April 19, 2025 11:31

guan404ming requested a review from eladkal April 19, 2025 11:54

eladkal reviewed Apr 19, 2025

View reviewed changes

devel-common/src/tests_common/test_utils/markers.py Outdated Show resolved Hide resolved

guan404ming force-pushed the miragte-test branch 4 times, most recently from 6cf4613 to 5d498dc Compare April 23, 2025 07:08

guan404ming closed this Apr 23, 2025

guan404ming reopened this Apr 23, 2025

guan404ming force-pushed the miragte-test branch from 5d498dc to 314dd76 Compare April 23, 2025 08:04

eladkal requested a review from potiuk April 23, 2025 13:34

potiuk force-pushed the miragte-test branch from 314dd76 to 67f3438 Compare April 23, 2025 15:04

guan404ming force-pushed the miragte-test branch 2 times, most recently from 8383a71 to 3012315 Compare April 28, 2025 07:49

potiuk enabled auto-merge (squash) April 28, 2025 08:30

auto-merge was automatically disabled April 28, 2025 09:24
Head branch was pushed to by a user without write access

guan404ming requested review from ashb and gopidesupavan as code owners April 28, 2025 09:24

guan404ming and others added 3 commits April 28, 2025 18:12

test: migrate get_pandas_df to get_df in provider test

6ac6470

Co-authored-by: Elad Kalif <[email protected]> Co-authored-by: Jarek Potiuk <[email protected]>

chore: add polars in deps

42f1395

fix: install polars in ci

9704c04

guan404ming force-pushed the miragte-test branch from 0c063c6 to 9704c04 Compare April 28, 2025 10:12

potiuk enabled auto-merge (squash) April 28, 2025 10:36

potiuk approved these changes Apr 28, 2025

View reviewed changes

potiuk merged commit c6ca8a2 into apache:main Apr 28, 2025
96 checks passed

guan404ming deleted the miragte-test branch April 28, 2025 11:04

potiuk mentioned this pull request Apr 28, 2025

[v3-0-test] test: migrate get_pandas_df to get_df in provider t… #49885

Merged

guan404ming changed the title ~~test: migrate get_pandas_df to get_df in provider test~~ Migrate get_pandas_df to get_df in provider test Apr 30, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Migrate `get_pandas_df` to `get_df` in `provider` test#49339

Migrate `get_pandas_df` to `get_df` in `provider` test#49339
potiuk merged 3 commits intoapache:mainfrom
guan404ming:miragte-test

guan404ming commented Apr 16, 2025 •

edited

Loading

Uh oh!

Uh oh!

guan404ming commented Apr 19, 2025 •

edited

Loading

Uh oh!

guan404ming commented Apr 19, 2025 •

edited

Loading

Uh oh!

Uh oh!

potiuk commented Apr 23, 2025

Uh oh!

potiuk commented Apr 28, 2025

Uh oh!

guan404ming commented Apr 28, 2025

Uh oh!

potiuk commented Apr 28, 2025

Uh oh!

guan404ming commented Apr 28, 2025 •

edited

Loading

Uh oh!

eladkal commented Apr 28, 2025

Uh oh!

potiuk commented Apr 28, 2025

Uh oh!

Uh oh!

guan404ming commented Apr 28, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Apr 28, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	def get_pandas_df(
	self,
	sql,
	parameters: list \| tuple \| Mapping[str, Any] \| None = None,
	**kwargs,
	) -> DataFrame:
	"""
	Execute the sql and returns a pandas 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 pandas.io.sql.read_sql method
	"""
	return self._get_pandas_df(sql, parameters, **kwargs)

	if df_type == "pandas":
	return self._get_pandas_df(sql, parameters, **kwargs)
	if df_type == "polars":
	return self._get_polars_df(sql, parameters, **kwargs)

Conversation

guan404ming commented Apr 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Related Issue

Why

How

Uh oh!

Uh oh!

guan404ming commented Apr 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

guan404ming commented Apr 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

potiuk commented Apr 23, 2025

Uh oh!

potiuk commented Apr 28, 2025

Uh oh!

guan404ming commented Apr 28, 2025

Uh oh!

potiuk commented Apr 28, 2025

Uh oh!

guan404ming commented Apr 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

eladkal commented Apr 28, 2025

Uh oh!

potiuk commented Apr 28, 2025

Uh oh!

Uh oh!

guan404ming commented Apr 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Apr 28, 2025

Backport failed to create: v3-0-test. View the failure log Run details

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

guan404ming commented Apr 16, 2025 •

edited

Loading

guan404ming commented Apr 19, 2025 •

edited

Loading

guan404ming commented Apr 19, 2025 •

edited

Loading

guan404ming commented Apr 28, 2025 •

edited

Loading

guan404ming commented Apr 28, 2025 •

edited

Loading