Dev-iL commented on code in PR #74817:
URL: https://github.com/apache/airflow/pull/74817#discussion_r4307061938
##########
providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py:
##########
@@ -220,7 +221,23 @@ def _strip_suffixes(
return path[: -len(suffix)]
return path
+ def _ensure_parquet_engine(self) -> None:
+ needs_parquet = self.file_format == FILE_FORMAT.PARQUET or (
+ ("dtype_backend", "pyarrow") in self.read_kwargs.items()
+ )
+ if self.df_type != "pandas" or not needs_parquet:
+ return
+ if any(importlib.util.find_spec(engine) for engine in ("pyarrow",
"fastparquet")):
Review Comment:
Confirmed at the source rather than from the repro alone: in pandas 3.0.6,
`pandas/io/sql.py:174-175` does `if dtype_backend == "pyarrow": pa =
import_optional_dependency("pyarrow")`, so fastparquet cannot stand in for it
whatever the output format is.
Split the check in b7c2f189c1. `dtype_backend="pyarrow"` now requires
pyarrow specifically, and only Parquet *output* accepts either engine. Two
cases cover it — `pandas-pyarrow-backend-fastparquet-only` and
`pandas-parquet-pyarrow-backend-fastparquet-only`, the second being the overlap
you flagged, where the output is Parquet and pyarrow is still the only thing
that will do. I checked both fail against the previous helper body (`DID NOT
RAISE`), so they are load-bearing.
The message now names the setting that triggered it, e.g.
`read_kwargs={"dtype_backend": "pyarrow"} needs pyarrow. Install
apache-airflow-providers-amazon[parquet].` The `[parquet]` extra resolves to
pyarrow only, so it stays the right pointer for both branches.
---
Drafted-by: Claude Code (Opus 5) (no human review before posting)
##########
providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py:
##########
@@ -199,6 +199,41 @@ def test_fix_dtypes_not_called(self, mock_s3_hook):
mock_fix_dtypes.assert_not_called()
+ @pytest.mark.parametrize(
+ ("df_type", "file_format", "read_kwargs", "available_engines",
"should_raise"),
+ [
+ pytest.param("pandas", "parquet", {}, set(), True,
id="pandas-parquet-no-engine"),
+ pytest.param("pandas", "parquet", {}, {"fastparquet"}, False,
id="pandas-parquet-fastparquet"),
+ pytest.param(
+ "pandas", "csv", {"dtype_backend": "pyarrow"}, set(), True,
id="pandas-pyarrow-backend"
+ ),
+ pytest.param("pandas", "csv", {}, set(), False,
id="pandas-csv-no-engine"),
+ pytest.param("polars", "parquet", {}, set(), False,
id="polars-parquet-no-engine"),
+ ],
+ )
+ @mock.patch("importlib.util.find_spec")
+ def test_parquet_engine_required(
+ self, mock_find_spec, df_type, file_format, read_kwargs,
available_engines, should_raise
+ ):
+ mock_find_spec.side_effect = lambda name: object() if name in
available_engines else None
+ op = SqlToS3Operator(
+ query="query",
+ s3_bucket="bucket",
+ s3_key="key",
+ sql_conn_id="mysql_conn_id",
+ task_id="task_id",
+ df_type=df_type,
+ file_format=file_format,
+ read_kwargs=read_kwargs,
+ dag=None,
+ )
+
+ if should_raise:
+ with pytest.raises(AirflowOptionalProviderFeatureException,
match=r"\[parquet\]"):
+ op._ensure_parquet_engine()
Review Comment:
Agreed. Added `test_execute_checks_parquet_engine_before_querying`, which
drives the raising path through `op.execute(None)` and asserts the hook's
`get_df` was never called.
I checked it is actually load-bearing rather than merely present: with
`self._ensure_parquet_engine()` deleted from `execute()`, it is the only test
in the file that goes red, and the captured log shows `Data from SQL obtained`
— so the query really was reached. It discriminates on ordering, which was your
point.
---
Drafted-by: Claude Code (Opus 5) (no human review before posting)
##########
providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py:
##########
@@ -220,7 +221,23 @@ def _strip_suffixes(
return path[: -len(suffix)]
return path
+ def _ensure_parquet_engine(self) -> None:
+ needs_parquet = self.file_format == FILE_FORMAT.PARQUET or (
+ ("dtype_backend", "pyarrow") in self.read_kwargs.items()
+ )
+ if self.df_type != "pandas" or not needs_parquet:
+ return
+ if any(importlib.util.find_spec(engine) for engine in ("pyarrow",
"fastparquet")):
+ return
+
+ from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
Review Comment:
Done — folded into the existing top-level `common.compat.sdk` import.
Checking that none was left turned up a second one: `_fix_dtypes` has
carried its own local import of the same exception since 4d3c5ead3b5, and the
module-level import this PR adds makes it dead. Removed that as well, since
your reasoning applies to it unchanged — and it also means that branch no
longer runs an import inside an `except ImportError` handler. Happy to drop
that hunk if you would rather keep it out of this PR.
---
Drafted-by: Claude Code (Opus 5) (no human review before posting)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
