You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "gopidesupavan (via GitHub)" <gi...@apache.org> on 2026/02/20 16:25:16 UTC
[PR] AIP-99: Add AnalyticsOperator [airflow]
gopidesupavan opened a new pull request, #62232:
URL: https://github.com/apache/airflow/pull/62232
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [ ] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on code in PR #62232:
URL: https://github.com/apache/airflow/pull/62232#discussion_r2839720001
##########
providers/common/ai/src/airflow/providers/common/ai/config.py:
##########
@@ -0,0 +1,95 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
Review Comment:
good call moved it.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/mixins.py:
##########
@@ -0,0 +1,91 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
Review Comment:
Yes updated.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on code in PR #62232:
URL: https://github.com/apache/airflow/pull/62232#discussion_r2839731815
##########
providers/common/ai/src/airflow/providers/common/ai/datafusion/engine.py:
##########
@@ -0,0 +1,102 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from datafusion import SessionContext
+
+from airflow.providers.common.ai.datafusion.format_handlers import get_format_handler
+from airflow.providers.common.ai.datafusion.object_storage_provider import ObjectStorageProviderFactory
+from airflow.providers.common.ai.exceptions import ObjectStoreCreationException, QueryExecutionException
+from airflow.providers.common.ai.utils.config import ConnectionConfig, DataSourceConfig
+from airflow.utils.log.logging_mixin import LoggingMixin
+
+
+class DataFusionEngine(LoggingMixin):
+ """Apache DataFusion engine."""
+
+ def __init__(self):
+ super().__init__()
+ # TODO: session context has additional parameters via SessionConfig see what's possible we can use Possible via DataFusionHook ?
+ self.df_ctx = SessionContext()
+ self.registered_tables: dict[str, str] = {}
+
+ @cached_property
+ def session_context(self) -> SessionContext:
+ """Return the session context."""
+ return self.df_ctx
Review Comment:
updated :) i overthink
##########
providers/common/ai/docs/operators/analytics.rst:
##########
@@ -0,0 +1,82 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+Analytics Operator
+===================
+
+The Analytics operator is designed to run analytic queries on data stored in various datastores. It is a generic operator that can query data in S3, GCS, Azure, and Local File System.
+
+The Analytics Operator uses Apache DataFusion as its query engine and supports SQL as the query language. It operates on a single node engine to deliver high-performance analytics on the data. It can be used for various analytics tasks such as data exploration, data aggregation, and more. `<https://datafusion.apache.org/>`_.
+
+
+When to Use Analytics Operator
+------------------------------
+
+Analytics Operator is suitable for running analytics on large volumes of datasets, with performance and efficiency. Under the hood, it uses Apache DataFusion, a high-performance, extensible query engine for Apache Arrow, which enables fast SQL queries on various data formats and storage systems. DataFusion is chosen for its ability to handle large-scale data processing on a single node, providing low-latency analytics without the need for a full database setup and without the need for high compute clusters. For more on Analytics Operator with DataFusion use cases, see `<https://datafusion.apache.org/user-guide/introduction.html#use-cases>`_.
+
+
+Supported Storage Systems
+-------------------------
+- S3
+- Local File System
Review Comment:
updated
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on code in PR #62232:
URL: https://github.com/apache/airflow/pull/62232#discussion_r2836722304
##########
providers/common/ai/src/airflow/providers/common/ai/datafusion/object_storage_provider.py:
##########
@@ -0,0 +1,89 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datafusion.object_store import AmazonS3, LocalFileSystem
+
+from airflow.providers.common.ai.datafusion.base import ObjectStorageProvider
+from airflow.providers.common.ai.exceptions import ObjectStoreCreationException
+from airflow.providers.common.ai.utils.config import ConnectionConfig, StorageType
+
+
+class S3ObjectStorageProvider(ObjectStorageProvider):
Review Comment:
Did some benchmark tests, results are here https://github.com/apache/airflow/pull/62232#issuecomment-3939582305 :)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on code in PR #62232:
URL: https://github.com/apache/airflow/pull/62232#discussion_r2836722747
##########
providers/common/ai/tests/system/common/ai/example_analytics.py:
##########
Review Comment:
yeah good point will move them..
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #62232:
URL: https://github.com/apache/airflow/pull/62232#discussion_r2836698904
##########
providers/common/ai/src/airflow/providers/common/ai/datafusion/object_storage_provider.py:
##########
@@ -0,0 +1,89 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datafusion.object_store import AmazonS3, LocalFileSystem
+
+from airflow.providers.common.ai.datafusion.base import ObjectStorageProvider
+from airflow.providers.common.ai.exceptions import ObjectStoreCreationException
+from airflow.providers.common.ai.utils.config import ConnectionConfig, StorageType
+
+
+class S3ObjectStorageProvider(ObjectStorageProvider):
Review Comment:
Do we not have sufficient base code and classes to read from object storages in apache-airflow-providers-common-io?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on PR #62232:
URL: https://github.com/apache/airflow/pull/62232#issuecomment-3943495167
@jason810496 thanks for review, please check, have updated ..
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on PR #62232:
URL: https://github.com/apache/airflow/pull/62232#issuecomment-3939582305
Benchmarks:
Though we already have datafusion benchamarks here https://datafusion.apache.org/blog/2024/11/18/datafusion-fastest-single-node-parquet-clickbench/ with tpch.
I have tried using trips yello_taxi data with 44.5M records. and run these agregation queries on the data that stored in s3. it took 1 minute 47s
```
"""SELECT COUNT(*) AS total_trips
FROM yellow_taxi;""",
"""SELECT SUM(total_amount) AS total_revenue
FROM yellow_taxi;""",
"""SELECT
AVG(trip_distance) AS avg_distance_miles,
AVG(fare_amount) AS avg_fare
FROM yellow_taxi;""",
"""SELECT
MIN(tip_amount) AS min_tip,
MAX(tip_amount) AS max_tip,
AVG(tip_amount) AS avg_tip
FROM yellow_taxi
WHERE payment_type = 1;""",
"""SELECT
DATE_TRUNC('day', tpep_pickup_datetime) AS pickup_day,
COUNT(*) AS trips,
AVG(passenger_count) AS avg_passengers
FROM yellow_taxi
GROUP BY pickup_day
ORDER BY pickup_day DESC;""",
"""SELECT
DATE_TRUNC('month', tpep_pickup_datetime) AS pickup_month,
SUM(fare_amount) AS total_fares,
SUM(tolls_amount) AS total_tolls
FROM yellow_taxi
GROUP BY pickup_month
ORDER BY pickup_month DESC;"""
```
<img width="1164" height="827" alt="image" src="https://github.com/user-attachments/assets/ed39ffbf-26a7-4a41-8efa-bc50c59b5dcc" />
With HITL approval operator the results can be viewed before the next task:
<img width="1055" height="810" alt="image" src="https://github.com/user-attachments/assets/9b260e2d-8219-409b-a631-7abf3eae040f" />
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #62232:
URL: https://github.com/apache/airflow/pull/62232#discussion_r2836700263
##########
providers/common/ai/tests/system/common/ai/example_analytics.py:
##########
Review Comment:
Should this not go into "example_dags" folder of provider rather than system test?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #62232:
URL: https://github.com/apache/airflow/pull/62232#discussion_r2838857217
##########
providers/common/ai/docs/operators/analytics.rst:
##########
@@ -0,0 +1,82 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+Analytics Operator
+===================
+
+The Analytics operator is designed to run analytic queries on data stored in various datastores. It is a generic operator that can query data in S3, GCS, Azure, and Local File System.
+
+The Analytics Operator uses Apache DataFusion as its query engine and supports SQL as the query language. It operates on a single node engine to deliver high-performance analytics on the data. It can be used for various analytics tasks such as data exploration, data aggregation, and more. `<https://datafusion.apache.org/>`_.
+
+
+When to Use Analytics Operator
+------------------------------
+
+Analytics Operator is suitable for running analytics on large volumes of datasets, with performance and efficiency. Under the hood, it uses Apache DataFusion, a high-performance, extensible query engine for Apache Arrow, which enables fast SQL queries on various data formats and storage systems. DataFusion is chosen for its ability to handle large-scale data processing on a single node, providing low-latency analytics without the need for a full database setup and without the need for high compute clusters. For more on Analytics Operator with DataFusion use cases, see `<https://datafusion.apache.org/user-guide/introduction.html#use-cases>`_.
+
+
+Supported Storage Systems
+-------------------------
+- S3
+- Local File System
Review Comment:
> So here we are registering all those object stores that supports by datafusion. https://datafusion.apache.org/python/autoapi/datafusion/object_store/index.html And few table providers Delta and Iceberg. not yet implemented in this PR
If this is the case, could we mention the https://datafusion.apache.org/python/autoapi/datafusion/object_store/index.html doc here and add note for users that we haven't supported some of the object store that datafusion already supported?
Since we only mentioned `S3` and `Local File System` here, which might be confused for users.
##########
providers/common/ai/docs/operators/analytics.rst:
##########
@@ -0,0 +1,82 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+Analytics Operator
+===================
+
+The Analytics operator is designed to run analytic queries on data stored in various datastores. It is a generic operator that can query data in S3, GCS, Azure, and Local File System.
+
+The Analytics Operator uses Apache DataFusion as its query engine and supports SQL as the query language. It operates on a single node engine to deliver high-performance analytics on the data. It can be used for various analytics tasks such as data exploration, data aggregation, and more. `<https://datafusion.apache.org/>`_.
+
+
+When to Use Analytics Operator
+------------------------------
+
+Analytics Operator is suitable for running analytics on large volumes of datasets, with performance and efficiency. Under the hood, it uses Apache DataFusion, a high-performance, extensible query engine for Apache Arrow, which enables fast SQL queries on various data formats and storage systems. DataFusion is chosen for its ability to handle large-scale data processing on a single node, providing low-latency analytics without the need for a full database setup and without the need for high compute clusters. For more on Analytics Operator with DataFusion use cases, see `<https://datafusion.apache.org/user-guide/introduction.html#use-cases>`_.
Review Comment:
The content of the paragraph seems almost identical to the last paragraph.
##########
providers/common/ai/src/airflow/providers/common/ai/datafusion/engine.py:
##########
@@ -0,0 +1,102 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from datafusion import SessionContext
+
+from airflow.providers.common.ai.datafusion.format_handlers import get_format_handler
+from airflow.providers.common.ai.datafusion.object_storage_provider import ObjectStorageProviderFactory
+from airflow.providers.common.ai.exceptions import ObjectStoreCreationException, QueryExecutionException
+from airflow.providers.common.ai.utils.config import ConnectionConfig, DataSourceConfig
+from airflow.utils.log.logging_mixin import LoggingMixin
+
+
+class DataFusionEngine(LoggingMixin):
+ """Apache DataFusion engine."""
+
+ def __init__(self):
+ super().__init__()
+ # TODO: session context has additional parameters via SessionConfig see what's possible we can use Possible via DataFusionHook ?
+ self.df_ctx = SessionContext()
+ self.registered_tables: dict[str, str] = {}
+
+ @cached_property
+ def session_context(self) -> SessionContext:
+ """Return the session context."""
+ return self.df_ctx
Review Comment:
May I ask why do we need another `cache_property` for `SessionContext`? Since we aren't using Hook here, so the `cache_property` seems unnecessary for me.
##########
providers/common/ai/src/airflow/providers/common/ai/datafusion/engine.py:
##########
@@ -0,0 +1,102 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from datafusion import SessionContext
+
+from airflow.providers.common.ai.datafusion.format_handlers import get_format_handler
+from airflow.providers.common.ai.datafusion.object_storage_provider import ObjectStorageProviderFactory
+from airflow.providers.common.ai.exceptions import ObjectStoreCreationException, QueryExecutionException
+from airflow.providers.common.ai.utils.config import ConnectionConfig, DataSourceConfig
+from airflow.utils.log.logging_mixin import LoggingMixin
+
+
+class DataFusionEngine(LoggingMixin):
+ """Apache DataFusion engine."""
+
+ def __init__(self):
+ super().__init__()
+ # TODO: session context has additional parameters via SessionConfig see what's possible we can use Possible via DataFusionHook ?
Review Comment:
Does it means that for the long-term goal, we will separate DataFusion as a independent provider (since we need hook for DataFusion), and Common AI provider will depend on DataFusion provider?
##########
providers/common/ai/src/airflow/providers/common/ai/utils/config.py:
##########
@@ -0,0 +1,95 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
Review Comment:
Not sure will it be more straightforward for users if we move `config` module to parent module?
```python
from airflow.providers.common.ai.utils.config import DataSourceConfig, FormatType
from airflow.providers.common.ai.config import DataSourceConfig, FormatType
```
##########
providers/common/ai/src/airflow/providers/common/ai/utils/mixins.py:
##########
@@ -0,0 +1,91 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
Review Comment:
How about using lazy import `AwsGenericHook` instead of directly import? Since we mark `amazon` as optional dependency.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/analytics.py:
##########
@@ -0,0 +1,163 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import json
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Literal
+
+from airflow.providers.common.ai.datafusion.engine import DataFusionEngine
+from airflow.providers.common.ai.utils.mixins import CommonAIHookMixin
+from airflow.sdk import BaseOperator, Context
+
+if TYPE_CHECKING:
+ from airflow.providers.common.ai.utils.config import DataSourceConfig
+
+
+class AnalyticsOperator(BaseOperator, CommonAIHookMixin):
+ """
+ Operator to run queries on various datasource's stored in object stores like S3, GCS, Azure, etc.
+
+ :param datasource_configs: List of datasource configurations to register.
+ :param queries: List of SQL queries to execute.
+ :param max_rows_check: Maximum number of rows allowed in query results. Queries exceeding this will be skipped.
+ :param engine: Optional DataFusion engine instance.
+ :param result_output_format: List of output formats for results. Supported: 'tabulate', 'json'. Default is 'tabulate'.
+ """
+
+ template_fields: Sequence[str] = (
+ "datasource_configs",
+ "queries",
+ "max_rows_check",
+ "result_output_format",
+ )
+
+ def __init__(
+ self,
+ datasource_configs: list[DataSourceConfig],
+ queries: list[str],
+ max_rows_check: int = 100,
+ engine: DataFusionEngine | None = None,
+ result_output_format: Literal["tabulate", "json"] = "tabulate",
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.datasource_configs = datasource_configs
+ self.queries = queries
+ self.engine = engine
+ self.max_rows_check = max_rows_check
+ self.result_output_format = result_output_format
+
+ @cached_property
+ def _df_engine(self):
+ if self.engine is None:
+ return DataFusionEngine()
+ return self.engine
+
+ def execute(self, context: Context) -> Any:
Review Comment:
```suggestion
def execute(self, context: Context) -> str:
```
##########
providers/common/ai/src/airflow/providers/common/ai/utils/config.py:
##########
@@ -0,0 +1,95 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+@dataclass(frozen=True)
+class ConnectionConfig:
+ """Configuration for datafusion object store connections."""
+
+ conn_id: str
+ credentials: dict[str, Any] = field(default_factory=dict)
+ extra_config: dict[str, Any] = field(default_factory=dict)
+
+
+class FormatType(str, Enum):
+ """Supported data formats."""
+
+ PARQUET = "parquet"
+ CSV = "csv"
+ JSON = "json"
+ AVRO = "avro"
+
+
+class StorageType(str, Enum):
+ """Storage types for Data Fusion."""
+
+ GCS = "gcs"
+ S3 = "s3"
+ AZURE = "azure"
+ LOCAL = "local"
+ HTTP = "http"
+
+
+@dataclass
Review Comment:
```suggestion
@dataclass(frozen=True)
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on code in PR #62232:
URL: https://github.com/apache/airflow/pull/62232#discussion_r2839730032
##########
providers/common/ai/src/airflow/providers/common/ai/datafusion/engine.py:
##########
@@ -0,0 +1,102 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from datafusion import SessionContext
+
+from airflow.providers.common.ai.datafusion.format_handlers import get_format_handler
+from airflow.providers.common.ai.datafusion.object_storage_provider import ObjectStorageProviderFactory
+from airflow.providers.common.ai.exceptions import ObjectStoreCreationException, QueryExecutionException
+from airflow.providers.common.ai.utils.config import ConnectionConfig, DataSourceConfig
+from airflow.utils.log.logging_mixin import LoggingMixin
+
+
+class DataFusionEngine(LoggingMixin):
+ """Apache DataFusion engine."""
+
+ def __init__(self):
+ super().__init__()
+ # TODO: session context has additional parameters via SessionConfig see what's possible we can use Possible via DataFusionHook ?
Review Comment:
Yes it could be new provider or may be we move this complete datafusion module to common-io provider , ideally this is dealing with storages and few table providers eg: Delta and Iceberg. so we have to see how ti goes for now.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on code in PR #62232:
URL: https://github.com/apache/airflow/pull/62232#discussion_r2836753522
##########
providers/common/ai/src/airflow/providers/common/ai/example_dags/example_analytics.py:
##########
Review Comment:
Updated https://github.com/apache/airflow/pull/62232/changes/b852be5a97f3aba1fab7e16900d901bebf2dd029
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on code in PR #62232:
URL: https://github.com/apache/airflow/pull/62232#discussion_r2839733750
##########
providers/common/ai/docs/operators/analytics.rst:
##########
@@ -0,0 +1,82 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+Analytics Operator
+===================
+
+The Analytics operator is designed to run analytic queries on data stored in various datastores. It is a generic operator that can query data in S3, GCS, Azure, and Local File System.
+
+The Analytics Operator uses Apache DataFusion as its query engine and supports SQL as the query language. It operates on a single node engine to deliver high-performance analytics on the data. It can be used for various analytics tasks such as data exploration, data aggregation, and more. `<https://datafusion.apache.org/>`_.
+
+
+When to Use Analytics Operator
+------------------------------
+
+Analytics Operator is suitable for running analytics on large volumes of datasets, with performance and efficiency. Under the hood, it uses Apache DataFusion, a high-performance, extensible query engine for Apache Arrow, which enables fast SQL queries on various data formats and storage systems. DataFusion is chosen for its ability to handle large-scale data processing on a single node, providing low-latency analytics without the need for a full database setup and without the need for high compute clusters. For more on Analytics Operator with DataFusion use cases, see `<https://datafusion.apache.org/user-guide/introduction.html#use-cases>`_.
Review Comment:
modified slight variation.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan closed pull request #62232: AIP-99: Add AnalyticsOperator
URL: https://github.com/apache/airflow/pull/62232
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan merged PR #62232:
URL: https://github.com/apache/airflow/pull/62232
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on code in PR #62232:
URL: https://github.com/apache/airflow/pull/62232#discussion_r2836721705
##########
providers/common/ai/src/airflow/providers/common/ai/datafusion/object_storage_provider.py:
##########
@@ -0,0 +1,89 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datafusion.object_store import AmazonS3, LocalFileSystem
+
+from airflow.providers.common.ai.datafusion.base import ObjectStorageProvider
+from airflow.providers.common.ai.exceptions import ObjectStoreCreationException
+from airflow.providers.common.ai.utils.config import ConnectionConfig, StorageType
+
+
+class S3ObjectStorageProvider(ObjectStorageProvider):
Review Comment:
Yeah we do have few operators that reads data, but here we are planning to use Apache DataFusion as query engine, as its very fast compare to what we have already in the code base and it can read large amount of data with less time..
So here we are registering all those object stores that supports by datafusion. https://datafusion.apache.org/python/autoapi/datafusion/object_store/index.html
And few table providers Delta and Iceberg. not yet implemented in this PR
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on PR #62232:
URL: https://github.com/apache/airflow/pull/62232#issuecomment-3939592932
https://github.com/user-attachments/assets/2a887d61-05e8-45d7-8bc3-c255e6c3ea83
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #62232:
URL: https://github.com/apache/airflow/pull/62232#discussion_r2849613089
##########
providers/common/ai/src/airflow/providers/common/ai/datafusion/format_handlers.py:
##########
@@ -0,0 +1,113 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.config import FormatType
+from airflow.providers.common.ai.datafusion.base import FormatHandler
+from airflow.providers.common.ai.exceptions import FileFormatRegistrationException
+
+if TYPE_CHECKING:
+ from datafusion import SessionContext
+
+
+class ParquetFormatHandler(FormatHandler):
+ """
+ Parquet format handler.
+
+ :param options: Additional options for the Parquet format.
+ https://datafusion.apache.org/python/autoapi/datafusion/context/index.html#datafusion.context.SessionContext.register_parquet
+ """
+
+ def __init__(self, options: dict[str, Any] | None = None):
+ self.options = options or {}
+
+ @property
+ def get_format(self) -> str:
+ """Return the format type."""
+ return FormatType.PARQUET.value
+
+ def register_data_source_format(self, ctx: SessionContext, table_name: str, path: str):
+ """Register a data source format."""
+ try:
+ ctx.register_parquet(table_name, path, **self.options)
+ except Exception as e:
+ raise FileFormatRegistrationException("Failed to register Parquet data source: %s", e)
+
Review Comment:
`raise FileFormatRegistrationException("Failed to register Parquet data source: %s", e)` passes multiple args to the exception (resulting in an odd tuple-like message) and drops exception chaining. Prefer interpolating the message (or using `from e`) so the resulting exception message is correct and the original stack trace is preserved.
##########
providers/common/ai/src/airflow/providers/common/ai/datafusion/object_storage_provider.py:
##########
@@ -0,0 +1,89 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datafusion.object_store import AmazonS3, LocalFileSystem
+
+from airflow.providers.common.ai.config import ConnectionConfig, StorageType
+from airflow.providers.common.ai.datafusion.base import ObjectStorageProvider
+from airflow.providers.common.ai.exceptions import ObjectStoreCreationException
+
+
+class S3ObjectStorageProvider(ObjectStorageProvider):
+ """S3 Object Storage Provider using DataFusion's AmazonS3."""
+
+ def get_storage_type(self) -> str:
+ """Return the storage type."""
+ return StorageType.S3.value
+
+ def create_object_store(self, path: str, connection_config: ConnectionConfig | None = None):
+ """Create an S3 object store using DataFusion's AmazonS3."""
+ if connection_config is None:
+ raise ValueError("connection_config must be provided")
+
+ try:
+ credentials = connection_config.credentials
+ bucket = self.get_bucket(path)
+
+ s3_store = AmazonS3(**credentials, **connection_config.extra_config, bucket_name=bucket)
+ self.log.info("Created S3 object store for bucket %s", bucket)
+
+ return s3_store
+
+ except Exception as e:
+ raise ObjectStoreCreationException("Failed to create S3 object store", e)
Review Comment:
`raise ObjectStoreCreationException("Failed to create S3 object store", e)` passes the original exception as a second positional arg, producing an odd message and not chaining the exception. Prefer `raise ObjectStoreCreationException(f"Failed to create S3 object store: {e}") from e` (or similar) so the message and traceback are correct.
```suggestion
raise ObjectStoreCreationException(f"Failed to create S3 object store: {e}") from e
```
##########
providers/common/ai/src/airflow/providers/common/ai/datafusion/object_storage_provider.py:
##########
@@ -0,0 +1,89 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datafusion.object_store import AmazonS3, LocalFileSystem
+
+from airflow.providers.common.ai.config import ConnectionConfig, StorageType
+from airflow.providers.common.ai.datafusion.base import ObjectStorageProvider
+from airflow.providers.common.ai.exceptions import ObjectStoreCreationException
+
+
+class S3ObjectStorageProvider(ObjectStorageProvider):
+ """S3 Object Storage Provider using DataFusion's AmazonS3."""
+
+ def get_storage_type(self) -> str:
+ """Return the storage type."""
+ return StorageType.S3.value
+
+ def create_object_store(self, path: str, connection_config: ConnectionConfig | None = None):
+ """Create an S3 object store using DataFusion's AmazonS3."""
+ if connection_config is None:
+ raise ValueError("connection_config must be provided")
+
+ try:
+ credentials = connection_config.credentials
+ bucket = self.get_bucket(path)
+
+ s3_store = AmazonS3(**credentials, **connection_config.extra_config, bucket_name=bucket)
+ self.log.info("Created S3 object store for bucket %s", bucket)
+
+ return s3_store
+
+ except Exception as e:
+ raise ObjectStoreCreationException("Failed to create S3 object store", e)
+
+ def get_scheme(self) -> str:
+ """Return the scheme for S3."""
+ return "s3://"
+
+
+class LocalObjectStorageProvider(ObjectStorageProvider):
+ """Local Object Storage Provider using DataFusion's LocalFileSystem."""
+
+ def get_storage_type(self) -> str:
+ """Return the storage type."""
+ return StorageType.LOCAL.value
+
+ def create_object_store(self, path: str, connection_config: ConnectionConfig | None = None):
+ """Create a Local object store."""
+ return LocalFileSystem()
+
+ def get_scheme(self) -> str:
+ """Return the scheme to a Local file system."""
+ return "file://"
+
+
+class ObjectStorageProviderFactory:
+ """Factory to create object storage providers based on storage type."""
+
+ # TODO: Add support for GCS, Azure, HTTP: https://datafusion.apache.org/python/autoapi/datafusion/object_store/index.html
+ _providers: dict[str, type] = {
+ StorageType.S3: S3ObjectStorageProvider,
+ StorageType.LOCAL: LocalObjectStorageProvider,
+ }
+
+ @classmethod
+ def create_provider(cls, storage_type: StorageType) -> ObjectStorageProvider:
+ """Create a storage provider instance."""
+ if storage_type not in cls._providers:
+ raise ValueError(
+ f"Unsupported storage type: {storage_type}. Supported types: {list(cls._providers.keys())}"
Review Comment:
`ObjectStorageProviderFactory._providers` is typed as `dict[str, type]` but keyed with `StorageType` enum members, and `create_provider()` only accepts `StorageType`. Since callers/tests may reasonably pass strings (e.g. from config/YAML), consider accepting `str | StorageType` and normalizing (or change the mapping keys to `StorageType` consistently and update tests/type hints).
```suggestion
_providers: dict[StorageType, type[ObjectStorageProvider]] = {
StorageType.S3: S3ObjectStorageProvider,
StorageType.LOCAL: LocalObjectStorageProvider,
}
@classmethod
def create_provider(cls, storage_type: StorageType | str) -> ObjectStorageProvider:
"""Create a storage provider instance."""
from airflow.providers.common.ai.config import StorageType as _StorageType
if isinstance(storage_type, str):
try:
storage_type = _StorageType(storage_type)
except ValueError as e:
supported = [st.value for st in _StorageType]
raise ValueError(
f"Unsupported storage type: {storage_type}. Supported types: {supported}"
) from e
if storage_type not in cls._providers:
supported = [st.value for st in cls._providers.keys()]
raise ValueError(
f"Unsupported storage type: {storage_type}. Supported types: {supported}"
```
##########
providers/common/ai/src/airflow/providers/common/ai/utils/mixins.py:
##########
@@ -0,0 +1,92 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.config import ConnectionConfig
+from airflow.providers.common.compat.sdk import BaseHook, Connection
+
+if TYPE_CHECKING:
+ from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+
+class CommonAIHookMixin:
+ """Mixin for Common AI."""
+
+ def get_db_api_hook(self, conn_id: str) -> DbApiHook:
+ """Get the given connection's database hook."""
+ connection = BaseHook.get_connection(conn_id)
+ return connection.get_hook()
+
+ def get_conn_config_from_airflow_connection(self, conn_id: str) -> ConnectionConfig:
+ """Get connection configuration from Airflow connection."""
+ try:
+ airflow_conn = BaseHook.get_connection(conn_id)
+
+ config = self._convert_airflow_connection(airflow_conn)
+
+ log.info("Loaded connection config for: %s", conn_id)
+ return config
+
+ except Exception as e:
+ log.error("Failed to get connection config for %s: %s", conn_id, e)
+ raise
+
+ def _convert_airflow_connection(self, conn: Connection) -> ConnectionConfig:
+ """Convert Airflow connection to ConnectionConfig."""
+ credentials = self._get_credentials(conn)
+
+ extra_config = conn.extra_dejson if conn.extra else {}
+
+ return ConnectionConfig(
+ conn_id=conn.conn_id,
+ credentials=credentials,
+ extra_config=extra_config,
+ )
+
+ @classmethod
+ def remove_none_values(cls, params: dict[str, Any]) -> dict[str, Any]:
+ """Filter out None values from the dictionary."""
+ return {k: v for k, v in params.items() if v is not None}
+
+ def _get_credentials(self, conn: Connection) -> dict[str, Any]:
+ from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+
+ credentials = {}
+
+ match conn.conn_type:
+ case "aws":
+ s3_conn: AwsGenericHook = AwsGenericHook(aws_conn_id=conn.conn_id, client_type="s3")
+ creds = s3_conn.get_credentials()
Review Comment:
`AwsGenericHook` is imported unconditionally inside `_get_credentials()`. If users install this provider without the `amazon` extra, using an `aws` connection will raise a raw `ImportError`. Consider catching `ImportError` here and raising `AirflowOptionalProviderFeatureException` with a clear message about installing `apache-airflow-providers-common-ai[amazon]` (or `apache-airflow-providers-amazon`).
##########
providers/common/ai/src/airflow/providers/common/ai/operators/analytics.py:
##########
@@ -0,0 +1,163 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import json
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Literal
+
+from airflow.providers.common.ai.datafusion.engine import DataFusionEngine
+from airflow.providers.common.ai.utils.mixins import CommonAIHookMixin
+from airflow.sdk import BaseOperator, Context
+
+if TYPE_CHECKING:
+ from airflow.providers.common.ai.config import DataSourceConfig
+
+
+class AnalyticsOperator(BaseOperator, CommonAIHookMixin):
+ """
+ Operator to run queries on various datasource's stored in object stores like S3, GCS, Azure, etc.
+
+ :param datasource_configs: List of datasource configurations to register.
+ :param queries: List of SQL queries to execute.
+ :param max_rows_check: Maximum number of rows allowed in query results. Queries exceeding this will be skipped.
+ :param engine: Optional DataFusion engine instance.
+ :param result_output_format: List of output formats for results. Supported: 'tabulate', 'json'. Default is 'tabulate'.
+ """
+
+ template_fields: Sequence[str] = (
+ "datasource_configs",
+ "queries",
+ "max_rows_check",
+ "result_output_format",
+ )
+
+ def __init__(
+ self,
+ datasource_configs: list[DataSourceConfig],
+ queries: list[str],
+ max_rows_check: int = 100,
+ engine: DataFusionEngine | None = None,
+ result_output_format: Literal["tabulate", "json"] = "tabulate",
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.datasource_configs = datasource_configs
+ self.queries = queries
+ self.engine = engine
+ self.max_rows_check = max_rows_check
+ self.result_output_format = result_output_format
+
+ @cached_property
+ def _df_engine(self):
+ if self.engine is None:
+ return DataFusionEngine()
+ return self.engine
+
+ def execute(self, context: Context) -> str:
+
+ results = []
+ for datasource_config in self.datasource_configs:
+ connection_config = self.get_conn_config_from_airflow_connection(datasource_config.conn_id)
Review Comment:
`execute()` unconditionally calls `get_conn_config_from_airflow_connection(datasource_config.conn_id)`. For local filesystem examples/tests you pass `conn_id=""`, which will call `BaseHook.get_connection("")` and fail at runtime. Consider treating a falsy `conn_id` as “no connection needed” and using an empty `ConnectionConfig` (or only resolving connections for non-local/object-store types that require it).
```suggestion
# Treat a falsy conn_id (e.g. "", None) as "no connection needed"
if getattr(datasource_config, "conn_id", None):
connection_config = self.get_conn_config_from_airflow_connection(datasource_config.conn_id)
else:
connection_config = None
```
##########
providers/common/ai/src/airflow/providers/common/ai/operators/analytics.py:
##########
@@ -0,0 +1,163 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import json
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Literal
+
+from airflow.providers.common.ai.datafusion.engine import DataFusionEngine
+from airflow.providers.common.ai.utils.mixins import CommonAIHookMixin
+from airflow.sdk import BaseOperator, Context
+
+if TYPE_CHECKING:
+ from airflow.providers.common.ai.config import DataSourceConfig
+
+
+class AnalyticsOperator(BaseOperator, CommonAIHookMixin):
+ """
+ Operator to run queries on various datasource's stored in object stores like S3, GCS, Azure, etc.
+
+ :param datasource_configs: List of datasource configurations to register.
+ :param queries: List of SQL queries to execute.
+ :param max_rows_check: Maximum number of rows allowed in query results. Queries exceeding this will be skipped.
+ :param engine: Optional DataFusion engine instance.
+ :param result_output_format: List of output formats for results. Supported: 'tabulate', 'json'. Default is 'tabulate'.
+ """
+
+ template_fields: Sequence[str] = (
+ "datasource_configs",
+ "queries",
+ "max_rows_check",
+ "result_output_format",
+ )
+
+ def __init__(
+ self,
+ datasource_configs: list[DataSourceConfig],
+ queries: list[str],
+ max_rows_check: int = 100,
+ engine: DataFusionEngine | None = None,
+ result_output_format: Literal["tabulate", "json"] = "tabulate",
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.datasource_configs = datasource_configs
+ self.queries = queries
+ self.engine = engine
+ self.max_rows_check = max_rows_check
+ self.result_output_format = result_output_format
+
+ @cached_property
+ def _df_engine(self):
+ if self.engine is None:
+ return DataFusionEngine()
+ return self.engine
+
+ def execute(self, context: Context) -> str:
+
+ results = []
+ for datasource_config in self.datasource_configs:
+ connection_config = self.get_conn_config_from_airflow_connection(datasource_config.conn_id)
+ self._df_engine.register_datasource(datasource_config, connection_config)
+
+ # TODO make it parallel as there is no dependency between queries
+ for query in self.queries:
+ result_dict = self._df_engine.execute_query(query)
+ results.append({"query": query, "data": result_dict})
+
+ match self.result_output_format:
+ case "tabulate":
+ return self._build_tabulate_output(results)
+ case "json":
+ return self._build_json_output(results)
+ case _:
+ raise ValueError(f"Unsupported output format: {self.result_output_format}")
+
+ def _is_result_too_large(self, result_dict: dict[str, Any]) -> tuple[bool, int]:
+ """Check if a result exceeds the max_rows_check limit."""
+ if not result_dict:
+ return False, 0
+ num_rows = len(next(iter(result_dict.values())))
+ max_rows_exceeded = num_rows >= self.max_rows_check
+ if max_rows_exceeded:
+ self.log.warning(
+ "Query returned %s rows, exceeding max_rows_check (%s). Skipping result output as large datasets are unsuitable for return.",
+ num_rows,
+ self.max_rows_check,
+ )
+ return max_rows_exceeded, num_rows
Review Comment:
`_is_result_too_large()` treats a result with exactly `max_rows_check` rows as too large (`>=`). The docs say “more than this number of rows” and the name `max_rows_check` suggests the limit should be inclusive. Consider changing the condition to `num_rows > self.max_rows_check` (and keep the warning/message consistent).
##########
providers/common/ai/src/airflow/providers/common/ai/datafusion/format_handlers.py:
##########
@@ -0,0 +1,113 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.config import FormatType
+from airflow.providers.common.ai.datafusion.base import FormatHandler
+from airflow.providers.common.ai.exceptions import FileFormatRegistrationException
+
+if TYPE_CHECKING:
+ from datafusion import SessionContext
+
+
+class ParquetFormatHandler(FormatHandler):
+ """
+ Parquet format handler.
+
+ :param options: Additional options for the Parquet format.
+ https://datafusion.apache.org/python/autoapi/datafusion/context/index.html#datafusion.context.SessionContext.register_parquet
+ """
+
+ def __init__(self, options: dict[str, Any] | None = None):
+ self.options = options or {}
+
+ @property
+ def get_format(self) -> str:
+ """Return the format type."""
+ return FormatType.PARQUET.value
+
+ def register_data_source_format(self, ctx: SessionContext, table_name: str, path: str):
+ """Register a data source format."""
+ try:
+ ctx.register_parquet(table_name, path, **self.options)
+ except Exception as e:
+ raise FileFormatRegistrationException("Failed to register Parquet data source: %s", e)
+
+
+class CsvFormatHandler(FormatHandler):
+ """
+ CSV format handler.
+
+ :param options: Additional options for the CSV format.
+ https://datafusion.apache.org/python/autoapi/datafusion/context/index.html#datafusion.context.SessionContext.register_csv
+ """
+
+ def __init__(self, options: dict[str, Any] | None = None):
+ self.options = options or {}
+
+ @property
+ def get_format(self) -> str:
+ """Return the format type."""
+ return FormatType.CSV.value
+
+ def register_data_source_format(self, ctx: SessionContext, table_name: str, path: str):
+ """Register a data source format."""
+ try:
+ ctx.register_csv(table_name, path, **self.options)
+ except Exception as e:
+ raise FileFormatRegistrationException("Failed to register csv data source: %s", e)
+
+
+class AvroFormatHandler(FormatHandler):
+ """
+ Avro format handler.
+
+ :param options: Additional options for the Avro format.
+ https://datafusion.apache.org/python/autoapi/datafusion/context/index.html#datafusion.context.SessionContext.register_avro
+ """
+
+ def __init__(self, options: dict[str, Any] | None = None):
+ self.options = options or {}
+
+ @property
+ def get_format(self) -> str:
+ """Return the format type."""
+ return FormatType.AVRO.value
+
+ def register_data_source_format(self, ctx: SessionContext, table_name: str, path: str) -> None:
+ """Register a data source format."""
+ try:
+ ctx.register_avro(table_name, path, **self.options)
+ except Exception as e:
+ raise FileFormatRegistrationException("Failed to register Avro data source: %s", e)
+
Review Comment:
Same issue as above: `FileFormatRegistrationException("Failed to register Avro data source: %s", e)` passes multiple args and loses exception chaining. Prefer a formatted message and `raise ... from e`.
##########
providers/common/ai/tests/unit/common/ai/datafusion/test_object_storage_provider.py:
##########
@@ -0,0 +1,78 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest.mock import patch
+
+import pytest
+
+from airflow.providers.common.ai.config import ConnectionConfig, StorageType
+from airflow.providers.common.ai.datafusion.object_storage_provider import (
+ LocalObjectStorageProvider,
+ ObjectStorageProviderFactory,
+ S3ObjectStorageProvider,
+)
+from airflow.providers.common.ai.exceptions import ObjectStoreCreationException
+
+
+class TestObjectStorageProvider:
+ @patch("airflow.providers.common.ai.datafusion.object_storage_provider.AmazonS3")
+ def test_s3_provider_success(self, mock_s3):
+ provider = S3ObjectStorageProvider()
+ connection_config = ConnectionConfig(
+ conn_id="aws_default",
+ credentials={"access_key_id": "fake_key", "secret_access_key": "fake_secret"},
+ )
+
+ store = provider.create_object_store("s3://demo-data/path", connection_config)
+
+ mock_s3.assert_called_once_with(
+ access_key_id="fake_key", secret_access_key="fake_secret", bucket_name="demo-data"
+ )
+ assert store == mock_s3.return_value
+ assert provider.get_storage_type() == StorageType.S3.value
+ assert provider.get_scheme() == "s3://"
+
+ def test_s3_provider_failure(self):
+ provider = S3ObjectStorageProvider()
+ connection_config = ConnectionConfig(conn_id="aws_default")
+
+ with patch(
+ "airflow.providers.common.ai.datafusion.object_storage_provider.AmazonS3",
+ side_effect=Exception("Error"),
+ ):
+ with pytest.raises(ObjectStoreCreationException, match="Failed to create S3 object store"):
+ provider.create_object_store("s3://demo-data/path", connection_config)
+
+ @patch("airflow.providers.common.ai.datafusion.object_storage_provider.LocalFileSystem")
+ def test_local_provider(self, mock_local):
+ provider = LocalObjectStorageProvider()
+ assert provider.get_storage_type() == StorageType.LOCAL.value
+ assert provider.get_scheme() == "file://"
+ local_store = provider.create_object_store("file://path")
+ assert local_store == mock_local.return_value
+
+ def test_factory_create_provider(self):
+ assert isinstance(
+ ObjectStorageProviderFactory.create_provider(StorageType.S3.value), S3ObjectStorageProvider
+ )
+ assert isinstance(
+ ObjectStorageProviderFactory.create_provider(StorageType.LOCAL.value), LocalObjectStorageProvider
Review Comment:
`ObjectStorageProviderFactory.create_provider()` is called here with `StorageType.S3.value` / `StorageType.LOCAL.value` (strings), but the factory implementation expects a `StorageType` enum and its `_providers` keys are `StorageType` members. This test will fail. Pass the enum values directly (e.g. `StorageType.S3`) or update the factory to accept strings and normalize to `StorageType`.
```suggestion
ObjectStorageProviderFactory.create_provider(StorageType.S3), S3ObjectStorageProvider
)
assert isinstance(
ObjectStorageProviderFactory.create_provider(StorageType.LOCAL), LocalObjectStorageProvider
```
##########
airflow-core/tests/unit/always/test_project_structure.py:
##########
@@ -98,6 +98,8 @@ def test_providers_modules_should_have_tests(self):
"providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_delete_from.py",
"providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_k8s_hashlib_wrapper.py",
"providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_xcom_sidecar.py",
+ "providers/common/ai/tests/unit/common/ai/datafusion/test_base.py",
+ "providers/common/ai/tests/unit/common/ai/test_exceptions.py",
Review Comment:
The project structure test now expects `providers/common/ai/tests/unit/common/ai/datafusion/test_base.py` and `providers/common/ai/tests/unit/common/ai/test_exceptions.py`, but those test files are not present in this PR. This will make `test_providers_modules_should_have_tests` fail. Either add the missing tests or remove these entries (and list the actual test files that exist).
```suggestion
```
##########
providers/common/ai/pyproject.toml:
##########
@@ -59,13 +59,30 @@ requires-python = ">=3.10"
# After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
dependencies = [
"apache-airflow>=3.0.0",
+ "datafusion>=50.0.0"
+]
+
+# The optional dependencies should be modified in place in the generated file
+# Any change in the dependencies is preserved when the file is regenerated
+[project.optional-dependencies]
+"amazon" = [
+ "apache-airflow-providers-amazon"
+]
+"common.compat" = [
+ "apache-airflow-providers-common-compat"
+]
+"common.sql" = [
+ "apache-airflow-providers-common-sql"
]
Review Comment:
`airflow.providers.common.compat.sdk` is imported by this provider at runtime (e.g. in `utils/mixins.py` and `exceptions.py`), so `apache-airflow-providers-common-compat` must be a required dependency (as in other providers), not only an optional extra. Without it, importing the provider will fail for users who install `apache-airflow-providers-common-ai` without extras.
##########
providers/common/ai/src/airflow/providers/common/ai/datafusion/format_handlers.py:
##########
@@ -0,0 +1,113 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.config import FormatType
+from airflow.providers.common.ai.datafusion.base import FormatHandler
+from airflow.providers.common.ai.exceptions import FileFormatRegistrationException
+
+if TYPE_CHECKING:
+ from datafusion import SessionContext
+
+
+class ParquetFormatHandler(FormatHandler):
+ """
+ Parquet format handler.
+
+ :param options: Additional options for the Parquet format.
+ https://datafusion.apache.org/python/autoapi/datafusion/context/index.html#datafusion.context.SessionContext.register_parquet
+ """
+
+ def __init__(self, options: dict[str, Any] | None = None):
+ self.options = options or {}
+
+ @property
+ def get_format(self) -> str:
+ """Return the format type."""
+ return FormatType.PARQUET.value
+
+ def register_data_source_format(self, ctx: SessionContext, table_name: str, path: str):
+ """Register a data source format."""
+ try:
+ ctx.register_parquet(table_name, path, **self.options)
+ except Exception as e:
+ raise FileFormatRegistrationException("Failed to register Parquet data source: %s", e)
+
+
+class CsvFormatHandler(FormatHandler):
+ """
+ CSV format handler.
+
+ :param options: Additional options for the CSV format.
+ https://datafusion.apache.org/python/autoapi/datafusion/context/index.html#datafusion.context.SessionContext.register_csv
+ """
+
+ def __init__(self, options: dict[str, Any] | None = None):
+ self.options = options or {}
+
+ @property
+ def get_format(self) -> str:
+ """Return the format type."""
+ return FormatType.CSV.value
+
+ def register_data_source_format(self, ctx: SessionContext, table_name: str, path: str):
+ """Register a data source format."""
+ try:
+ ctx.register_csv(table_name, path, **self.options)
+ except Exception as e:
+ raise FileFormatRegistrationException("Failed to register csv data source: %s", e)
+
Review Comment:
Same issue as above: `FileFormatRegistrationException("Failed to register csv data source: %s", e)` passes multiple args and loses exception chaining. Prefer a formatted message and `raise ... from e`.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on code in PR #62232:
URL: https://github.com/apache/airflow/pull/62232#discussion_r2839723720
##########
providers/common/ai/src/airflow/providers/common/ai/config.py:
##########
@@ -0,0 +1,95 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+@dataclass(frozen=True)
+class ConnectionConfig:
+ """Configuration for datafusion object store connections."""
+
+ conn_id: str
+ credentials: dict[str, Any] = field(default_factory=dict)
+ extra_config: dict[str, Any] = field(default_factory=dict)
+
+
+class FormatType(str, Enum):
+ """Supported data formats."""
+
+ PARQUET = "parquet"
+ CSV = "csv"
+ JSON = "json"
+ AVRO = "avro"
+
+
+class StorageType(str, Enum):
+ """Storage types for Data Fusion."""
+
+ GCS = "gcs"
+ S3 = "s3"
+ AZURE = "azure"
+ LOCAL = "local"
+ HTTP = "http"
+
+
+@dataclass
Review Comment:
Nah, we are updating few fields based on the URI in the `__post_init__`, example storage_type
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-99: Add AnalyticsOperator [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on code in PR #62232:
URL: https://github.com/apache/airflow/pull/62232#discussion_r2839736924
##########
providers/common/ai/src/airflow/providers/common/ai/datafusion/object_storage_provider.py:
##########
@@ -0,0 +1,89 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datafusion.object_store import AmazonS3, LocalFileSystem
+
+from airflow.providers.common.ai.datafusion.base import ObjectStorageProvider
+from airflow.providers.common.ai.exceptions import ObjectStoreCreationException
+from airflow.providers.common.ai.utils.config import ConnectionConfig, StorageType
+
+
+class S3ObjectStorageProvider(ObjectStorageProvider):
Review Comment:
@jscheffl hope this make sense?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org