Migrate get_pandas_df to get_df in provider test

[guan404ming]

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.

[guan404ming]

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]

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

[potiuk]

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]

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]

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]

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]

It works!! 🎉
Force push for the conflict.

[eladkal]

It works!! 🎉 Force push for the conflict.

🎉

[potiuk]

It works!! 🎉 Force push for the conflict.

🎉

Nice!!!!

[guan404ming]

Thanks for the review and suggestions!

[github-actions]

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