-
Notifications
You must be signed in to change notification settings - Fork 13.7k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
AIP-58: Add Airflow ObjectStore (AFS) (#34729)
This adds the ObjectStorage and ObjectStorePath APIs per AIP-58. ObjectStorePath is a pathlib.Pathlib like interface for objects residing on object storage.
- Loading branch information
1 parent
85f0ef3
commit 04e2fbd
Showing
57 changed files
with
2,876 additions
and
68 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -48,6 +48,7 @@ body: | |
| - celery | ||
| - cloudant | ||
| - cncf-kubernetes | ||
| - common-io | ||
| - common-sql | ||
| - daskexecutor | ||
| - databricks | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,135 @@ | ||
| # | ||
| # 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 | ||
|
|
||
| # [START tutorial] | ||
| # [START import_module] | ||
| import pendulum | ||
| import requests | ||
|
|
||
| from airflow.decorators import dag, task | ||
| from airflow.io.store.path import ObjectStoragePath | ||
|
|
||
| # [END import_module] | ||
|
|
||
| API = "https://opendata.fmi.fi/timeseries" | ||
|
|
||
| aq_fields = { | ||
| "fmisid": "int32", | ||
| "time": "datetime64[ns]", | ||
| "AQINDEX_PT1H_avg": "float64", | ||
| "PM10_PT1H_avg": "float64", | ||
| "PM25_PT1H_avg": "float64", | ||
| "O3_PT1H_avg": "float64", | ||
| "CO_PT1H_avg": "float64", | ||
| "SO2_PT1H_avg": "float64", | ||
| "NO2_PT1H_avg": "float64", | ||
| "TRSC_PT1H_avg": "float64", | ||
| } | ||
|
|
||
| # [START create_object_storage_path] | ||
| base = ObjectStoragePath("s3://airflow-tutorial-data/", conn_id="aws_default") | ||
| # [END create_object_storage_path] | ||
|
|
||
|
|
||
| # [START instantiate_dag] | ||
| @dag( | ||
| schedule=None, | ||
| start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), | ||
| catchup=False, | ||
| tags=["example"], | ||
| ) | ||
| def tutorial_objectstorage(): | ||
| """ | ||
| ### Object Storage Tutorial Documentation | ||
| This is a tutorial DAG to showcase the usage of the Object Storage API. | ||
| Documentation that goes along with the Airflow Object Storage tutorial is | ||
| located | ||
| [here](https://airflow.apache.org/docs/apache-airflow/stable/tutorial/objectstorage.html) | ||
| """ | ||
| # [END instantiate_dag] | ||
| import duckdb | ||
| import pandas as pd | ||
|
|
||
| # [START get_air_quality_data] | ||
| @task | ||
| def get_air_quality_data(**kwargs) -> ObjectStoragePath: | ||
| """ | ||
| #### Get Air Quality Data | ||
| This task gets air quality data from the Finnish Meteorological Institute's | ||
| open data API. The data is saved as parquet. | ||
| """ | ||
| execution_date = kwargs["logical_date"] | ||
| start_time = kwargs["data_interval_start"] | ||
|
|
||
| params = { | ||
| "format": "json", | ||
| "precision": "double", | ||
| "groupareas": "0", | ||
| "producer": "airquality_urban", | ||
| "area": "Uusimaa", | ||
| "param": ",".join(aq_fields.keys()), | ||
| "starttime": start_time.isoformat(timespec="seconds"), | ||
| "endtime": execution_date.isoformat(timespec="seconds"), | ||
| "tz": "UTC", | ||
| } | ||
|
|
||
| response = requests.get(API, params=params) | ||
| response.raise_for_status() | ||
|
|
||
| # ensure the bucket exists | ||
| base.mkdir(exists_ok=True) | ||
|
|
||
| formatted_date = execution_date.format("YYYYMMDD") | ||
| path = base / f"air_quality_{formatted_date}.parquet" | ||
|
|
||
| df = pd.DataFrame(response.json()).astype(aq_fields) | ||
| with path.open("wb") as file: | ||
| df.to_parquet(file) | ||
|
|
||
| return path | ||
|
|
||
| # [END get_air_quality_data] | ||
|
|
||
| # [START analyze] | ||
| @task | ||
| def analyze(path: ObjectStoragePath, **kwargs): | ||
| """ | ||
| #### Analyze | ||
| This task analyzes the air quality data, prints the results | ||
| """ | ||
| conn = duckdb.connect(database=":memory:") | ||
| conn.register_filesystem(path.fs) | ||
| conn.execute(f"CREATE OR REPLACE TABLE airquality_urban AS SELECT * FROM read_parquet('{path}')") | ||
|
|
||
| df2 = conn.execute("SELECT * FROM airquality_urban").fetchdf() | ||
|
|
||
| print(df2.head()) | ||
|
|
||
| # [END analyze] | ||
|
|
||
| # [START main_flow] | ||
| obj_path = get_air_quality_data() | ||
| analyze(obj_path) | ||
| # [END main_flow] | ||
|
|
||
|
|
||
| # [START dag_invocation] | ||
| tutorial_objectstorage() | ||
| # [END dag_invocation] | ||
| # [END tutorial] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,90 @@ | ||
| # 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, | ||
| Callable, | ||
| ) | ||
|
|
||
| from fsspec.implementations.local import LocalFileSystem | ||
|
|
||
| from airflow.compat.functools import cache | ||
| from airflow.providers_manager import ProvidersManager | ||
| from airflow.stats import Stats | ||
| from airflow.utils.module_loading import import_string | ||
|
|
||
| if TYPE_CHECKING: | ||
| from fsspec import AbstractFileSystem | ||
|
|
||
| log = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| def _file(_: str | None) -> LocalFileSystem: | ||
| return LocalFileSystem() | ||
|
|
||
|
|
||
| # builtin supported filesystems | ||
| _BUILTIN_SCHEME_TO_FS: dict[str, Callable[[str | None], AbstractFileSystem]] = { | ||
| "file": _file, | ||
| } | ||
|
|
||
|
|
||
| @cache | ||
| def _register_filesystems() -> dict[str, Callable[[str | None], AbstractFileSystem]]: | ||
| scheme_to_fs = _BUILTIN_SCHEME_TO_FS.copy() | ||
| with Stats.timer("airflow.io.load_filesystems") as timer: | ||
| manager = ProvidersManager() | ||
| for fs_module_name in manager.filesystem_module_names: | ||
| fs_module = import_string(fs_module_name) | ||
| for scheme in getattr(fs_module, "schemes", []): | ||
| if scheme in scheme_to_fs: | ||
| log.warning("Overriding scheme %s for %s", scheme, fs_module_name) | ||
|
|
||
| method = getattr(fs_module, "get_fs", None) | ||
| if method is None: | ||
| raise ImportError(f"Filesystem {fs_module_name} does not have a get_fs method") | ||
| scheme_to_fs[scheme] = method | ||
|
|
||
| log.debug("loading filesystems from providers took %.3f seconds", timer.duration) | ||
| return scheme_to_fs | ||
|
|
||
|
|
||
| def get_fs(scheme: str, conn_id: str | None = None) -> AbstractFileSystem: | ||
| """ | ||
| Get a filesystem by scheme. | ||
| :param scheme: the scheme to get the filesystem for | ||
| :return: the filesystem method | ||
| :param conn_id: the airflow connection id to use | ||
| """ | ||
| filesystems = _register_filesystems() | ||
| try: | ||
| return filesystems[scheme](conn_id) | ||
| except KeyError: | ||
| raise ValueError(f"No filesystem registered for scheme {scheme}") | ||
|
|
||
|
|
||
| def has_fs(scheme: str) -> bool: | ||
| """ | ||
| Check if a filesystem is available for a scheme. | ||
| :param scheme: the scheme to check | ||
| :return: True if a filesystem is available for the scheme | ||
| """ | ||
| return scheme in _register_filesystems() |
Oops, something went wrong.