Apache Airflow® Quickstart

1. Create a new directory for your project and open it:
   mkdir airflow-quickstart-etl && cd airflow-quickstart-etl
   
2. Clone the repository and open it:
   git clone -b etl --single-branch https://github.com/astronomer/airflow-quickstart.git && cd airflow-quickstart/etl
   
   Your directory should have the following structure:
   .
   ├── Dockerfile
   ├── README.md
   ├── dags
   │	└── example_etl_galaxies.py
   ├── include
   │	├── astronomy.db
   │	├── custom_functions
   │	│	└── galaxy_functions.py
   │	└── data
   │		└── galaxy_names.txt
   ├── packages.txt
   ├── requirements.txt
   ├── solutions
   │	└── example_etl_galaxies_solution.py
   └── tests
       └── dags
           └── test_dag_example.py
   

1. Create a new directory for your project and open it:
   mkdir airflow-quickstart-genai && cd airflow-quickstart-genai
   
2. Clone the repository and open it:
   git clone -b generative-ai --single-branch https://github.com/astronomer/airflow-quickstart.git && cd airflow-quickstart/generative-ai
   
   Your directory should have the following structure:
   .
   ├── Dockerfile
   ├── README.md
   ├── airflow_settings.yaml
   ├── dags
   │	└── example_vector_embeddings.py
   ├── include
   │	├── custom_functions
   │	│	└── embedding_func.py
   │	└── data
   │		└── galaxy_names.txt
   ├── packages.txt
   ├── requirements.txt
   ├── solutions
   │	└── example_vector_embeddings_solution.py
   └── tests
       └── dags
           └── test_dag_integrity.py
   

1. Start the project using the Astro CLI:
   astro dev start
   
   The CLI will let you know when all Airflow services are up and running.
2. In your browser, navigate to localhost:8080 and sign in to the Airflow UI using username admin and password admin.
3. Unpause the example_astronauts DAG.
4. Explore the landing page and individual DAG view page to get a sense of the metadata available about the DAG, run, and all task instances. For a deep-dive into the UI’s features, see An introduction to the Airflow UI. For example, the DAGs view will look like this screenshot:
   
   As you start to trigger DAG runs, the graph view will look like this screenshot:
   
   The Gantt chart will look like this screenshot:
   

1. Start the project using the Astro CLI:
   astro dev start
   
   The CLI will let you know when all Airflow services are up and running.
2. In your browser, navigate to localhost:8080 and sign in to the Airflow UI using username admin and password admin.
3. Unpause the example_astronauts DAG.
4. Explore the DAGs view (landing page) and individual DAG view page to get a sense of the metadata available about the DAG, run, and all task instances. For a deep-dive into the UI’s features, see An introduction to the Airflow UI. For example, the DAGs view will look like this screenshot:
   
   As you start to trigger DAG runs, the graph view will look like this screenshot:
   
   The Gantt chart will look like this screenshot:
   

• write responsive pipelines that change based on user inputs.
• perform database operations using SQL.
• access and extract data from local files.
• execute transformations with Pandas.

​

Pipeline structure

example_etl_galaxies

• create_galaxy_table_in_duckdb uses a hook to create a database connection and a SQL query to create a database table.
• extract_galaxy_data returns a dataframe created using a modularized function imported from the project’s include directory.
• transform_galaxy_data gets a user-specified value from the DAG context and uses it to execute a simple data transformation on the dataframe, returning another dataframe.
• load_galaxy_data uses a database hook to load the dataframe into the database. You can load the dataframe directly in the context of a SQL query. No conversion of the dataframe is required.
• print_loaded_galaxies executes a SELECT query on the database and prints the data to the logs.

• orchestrate a generative AI pipeline.
• compute vector embeddings of words using Python’s SentenceTransformers library.
• compare the embeddings of a word of interest to a list of words to find the semantically closest match.

​

Pipeline structure

example_vector_embeddings

• get_words gets a list of words from the context to embed.
• create_embeddings creates embeddings for the list of words.
• create_vector_table creates a table in the DuckDB database and an HNSW index on the embedding vector.
• insert_words_into_db inserts the words and embeddings into the table.
• embed_word embeds a single word and returns the embeddings.
• find_closest_word_match finds the closest match to a word of interest.

​

Create a connection to an external system

1. A Postgres database has already been added to the project for you. Confirm that port 5432 is available, or, if not, modify the external port in the project’s docker-compose.override.yml:
   version: "3.1"
   services:
     example-db:
       image: postgres:14.9
       ports:
         - "5435:5432"
       environment:
         - POSTGRES_USER=example
         - POSTGRES_PASSWORD=example
         - POSTGRES_DB=example
   
2. Create a connection in the UI. Airflow supports a number of ways to create connections, but most users choose the UI. To create a connection in the UI, go to Admin > Connections.
   
   Your connection should look like this screenshot, with the password being example:
   
3. Add a DAG to the project. Create a new Python file in the dags directory and paste this DAG code there:
   from airflow.decorators import (
       dag,
       task,
   )  # This DAG uses the TaskFlow API. See: https://www.astronomer.io/docs/airflow/airflow-decorators
   from pendulum import datetime, duration
   from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
   from include.custom_functions.load_functions import get_sql_query
   from airflow.models.dataset import Dataset
   
   _DUCKDB_INSTANCE_NAME = os.getenv("DUCKDB_INSTANCE_NAME", "include/astronomy.db")
   _TABLE_NAME = os.getenv("_TABLE_NAME", "galaxy_data")
   _DUCKDB_TABLE_NAME = os.getenv("DUCKDB_TABLE_NAME", "galaxy_data")
   _DUCKDB_TABLE_URI = f"duckdb://{_DUCKDB_INSTANCE_NAME}/{_DUCKDB_TABLE_NAME}"
   
   @dag(
       start_date=datetime(2024, 7, 1),  # Date after which the DAG can be scheduled
       schedule=[Dataset(_DUCKDB_TABLE_URI)],  # See: https://www.astronomer.io/docs/airflow/scheduling-in-airflow for options
       catchup=False,  # See: https://www.astronomer.io/docs/airflow/rerunning-dags#catchup
       max_consecutive_failed_dag_runs=5,  # auto-pauses the DAG after 5 consecutive failed runs, experimental
       max_active_runs=1,  # Allow only one concurrent run of this DAG, prevents parallel DuckDB calls
       doc_md=__doc__,  # Add DAG Docs in the UI, see https://www.astronomer.io/docs/airflow/custom-airflow-ui-docs-tutorial
       default_args={
           "owner": "Astro",  # owner of this DAG in the Airflow UI
           "retries": 3,  # tasks retry 3 times before they fail
           "retry_delay": duration(seconds=30),  # tasks wait 30s in between retries
       },  # default_args are applied to all tasks in a DAG
       tags=["example", "ETL"],  # Add tags in the UI
       # Warning - in-memory DuckDB is not a persistent database between workers. To move this workflow into production, use a
       # cloud-based database and, based on concurrency capabilities, adjust the two parameters below.
       concurrency=1, # allow only a single task execution at a time, prevents parallel DuckDB calls
       is_paused_upon_creation=False, # start running the DAG as soon as it's created
   )
   def example_etl_galaxies_load():  # By default, the dag_id is the name of the decorated function
   
       @task
       def extract_galaxy_data_duckdb(
           duckdb_instance_name: str = _DUCKDB_INSTANCE_NAME,
           table_name: str = _TABLE_NAME,
       ):
           cursor = duckdb.connect(duckdb_instance_name)
           galaxy_data_df = cursor.sql(f"SELECT * FROM {table_name};").df()
   
           return galaxy_data_df
   
       @task
       def create_sql_query(df):
           sql_str = get_sql_query(df, _TABLE_NAME)
           return sql_str
   
       create_galaxy_table_postgres = SQLExecuteQueryOperator(
           task_id="create_galaxy_table_postgres",
           conn_id="postgres_default",
           sql=f"""
               DROP TABLE IF EXISTS {_TABLE_NAME};
               CREATE TABLE {_TABLE_NAME} (
                   name VARCHAR PRIMARY KEY,
                   distance_from_milkyway INT,
                   distance_from_solarsystem INT,
                   type_of_galaxy VARCHAR,
                   characteristics VARCHAR
               );
               """,
           )
   
       create_sql_query_obj = create_sql_query(extract_galaxy_data_duckdb())
   
       load_galaxy_data_postgres = SQLExecuteQueryOperator(
           task_id = "load_galaxy_data_postgres",
           conn_id = "postgres_default",
           sql = create_sql_query_obj,
       )
   
       create_galaxy_table_postgres >> load_galaxy_data_postgres
   
   example_etl_galaxies_load()
   
   

   See all 76 lines
   This DAG extracts data from the project’s DuckDB database, creates a table in the project’s Postgres database, and loads the data into the table. Using an Airflow Dataset trigger, it will run when example_etl_galaxies updates the galaxy_data dataset. Connection parameters vary between operators. In the case of SQLExecuteQueryOperator, it is conn_id:
   load_galaxy_data_postgres = SQLExecuteQueryOperator(
       task_id = "load_galaxy_data_postgres",
       conn_id = "postgres_default",
       sql = create_sql_query_obj,
   )
   
4. Trigger your new DAG! Trigger the example_etl_galaxies DAG and not the new one you just added. Your new DAG will run after the load_galaxy_task in example_etl_galaxies runs successfully. If all goes well, the graph view will look like this screenshot:
   

• configure a DAG to use different LMs.
• use the Airflow UI to compare the performance of the models you select.

​

Experiment with different LMs to compare performance

1. Start your experiment by using a different model. Find the _LM variable definition in the get_embeddings_one_word function close to the top of the example_vector_embeddings DAG and replace the model string with distiluse-base-multilingual-cased-v2:
   _LM = os.getenv("LM", "distiluse-base-multilingual-cased-v2")
   
   The default is very fast, but this one is slower and lower-performing overall, so the results should be noticeably different. You could also try a model with higher overall performance, such as all-mpnet-base-v2. For a list of possible models to choose from, see SBERT’s Pretrained models list.
2. Next, find the dimensions of the model in the SBERT docs. For example, the distiluse-base-multilingual-cased-v2 model has dimensions of 512:
   
3. Use this number to redefine another top-level variable, _LM_DIMENSIONS:
   _LM_DIMENSIONS = os.getenv("LM_DIMS", "512")
   
   This value is used in the vector column type definition in the create_vector_table task and the select query in the find_closest_word_match task.
4. Rerun the DAG. Depending on the models you choose, you might see large differences in the performance of the create_embeddings task. For example, using the default all-MiniLM-L6-v2 model should result in runtimes of around 4 seconds:
   
   By contrast, using the distiluse-base-multilingual-cased-v2 model might result in runtimes three times as long or longer:
   
5. Check the log output from the find_closest_word_match task and look for differences between the search result sets. For example, the faster LM all-MiniLM-L6-v2 returns sun, planet, light:
   
   The more performant LM all-mpnet-base-v2 returns sun, rocket, planet: