Fix: Fixes to dags in new v2-2024 environment (#760)

* fix: yaml issue

* fix: Flake linter issues

* fix: Flake black issue

* fix: Resolutions to configurations and some code in order to move to new v2-2024 environment.

* Fix: Fixes to census_bureau_acs post push to v2-2024 environment.

* fix: Invalid reference to variable when timeout occurs during GSOD_BY_YEAR pipeline execution.

* fix: Miscellaneous fixes.

* fix: Move GSOD and Storms by year to end of process.

* Fix: Resolve KubernetesPodOperator issues and reduce code size.

* Fix: Chicago Taxi Trips now runs E2E successfully.

* Fix: Lint Issues

* Fix: Comment out section that breaks due to the fact that structural changes in the 'resouce limit' code section of pipeline.yaml in airflow has changed causing issues with validation of KubernetesPodOperator.

* Fix: Comment out code that breaks due to the fact that structural changes in the 'resouce limit' code section of pipeline.yaml in airflow has changed causing issues with validation of KubernetesPodOperator.

* Fix: Comment out code that breaks due to the fact that structural changes in the 'resouce limit' code section of pipeline.yaml in airflow has changed causing issues with validation of KubernetesPodOperator. Resolve import issue.

* Fix: Comment out code that breaks due to the fact that structural changes in the 'resouce limit' code section of pipeline.yaml in airflow has changed causing issues with validation of KubernetesPodOperator. Resolve import issue.

datasets/chicago_taxi_trips/pipelines/_images/run_csv_transform_kub/Dockerfile

@@ -11,27 +11,14 @@
 # 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.
-
-# The base image for this build
 FROM python:3.8
-
-# Allow statements and log messages to appear in Cloud logs
 ENV PYTHONUNBUFFERED True
-
-# Copy the requirements file into the image
 COPY requirements.txt ./
-
-# Install the packages specified in the requirements file
 RUN python3 -m pip install --no-cache-dir -r requirements.txt
-
-# The WORKDIR instruction sets the working directory for any RUN, CMD,
-# ENTRYPOINT, COPY and ADD instructions that follow it in the Dockerfile.
-# If the WORKDIR doesn’t exist, it will be created even if it’s not used in
-# any subsequent Dockerfile instruction
+RUN python3 -m pip install fsspec
+RUN python3 -m pip install gcsfs
+RUN curl -sSL https://sdk.cloud.google.com | bash
+ENV PATH $PATH:/root/google-cloud-sdk/bin
 WORKDIR /custom
-
-# Copy the specific data processing script/s in the image under /custom/*
 COPY ./csv_transform.py .
-
-# Command to run the data processing script when the container is run
 CMD ["python3", "csv_transform.py"]

datasets/chicago_taxi_trips/pipelines/_images/run_csv_transform_kub/csv_transform.py

@@ -15,183 +15,74 @@
 import datetime
 import json
 import logging
-import math
 import os
-import pathlib
+import subprocess
 
 import pandas as pd
-import requests
-from google.cloud import storage
 
 
 def main(
-    pipeline: str,
     source_url: str,
-    source_file: str,
     gcs_bucket: str,
+    csv_headers: list,
     csv_gcs_path: str,
-    source_object_folder: str,
-    rename_mappings: dict,
     non_na_columns: list,
-    output_object_folder: str,
-    data_types: dict,
-    date_cols: list,
-    batch_count: int,
+    data_dtypes: dict,
+    chunksize: int,
 ) -> None:
     logging.info(
         f'Chicago Taxi Trips Dataset pipeline process started at {str(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))}'
     )
-    logging.info(
-        f"Creating './files', './source' & './output' folder under pwd - {os.getcwd()}"
+    # Stage the datafile
+    gcs_datafile = f"gs://{gcs_bucket}/{csv_gcs_path}"
+    logging.info(f"Downloading source file from {source_url}.")
+    cmd = f"wget -q -O - {source_url} |gcloud storage cp - {gcs_datafile}"
+    ps = subprocess.Popen(
+        cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
     )
-    pathlib.Path("./files").mkdir(parents=True, exist_ok=True)
-    pathlib.Path("./source").mkdir(parents=True, exist_ok=True)
-    pathlib.Path("./output").mkdir(parents=True, exist_ok=True)
-    if pipeline:
-        logging.info(f"Pipeline = {pipeline} process started")
-        download_file(source_url, source_file)
-        upload_file_to_gcs(source_file, gcs_bucket, csv_gcs_path)
+    ps.communicate()
+    logging.info(f"Reading and processing source file from {gcs_datafile}.")
+    file_batch_nbr = 1
+    batch_filename = f"./output/taxi_trips_{str.zfill(str(file_batch_nbr), 10)}"
+    for i, chunk in enumerate(
+        pd.read_csv(
+            gcs_datafile,
+            chunksize=chunksize,
+            names=csv_headers,
+            skiprows=1,
+            dtype=data_dtypes,
+        )
+    ):
+        batch_num_pad = str(i + 1).zfill(10)
+        batch_path = os.path.join(os.path.dirname(gcs_datafile), "batch")
+        batch_filename = "taxi_trips_"
+        batch_fullpath = os.path.join(
+            batch_path, f"{batch_filename}{batch_num_pad}.csv"
+        )
         logging.info(
-            f"Chicago Taxi Trips Dataset pipeline {pipeline} process completed at "
-            + str(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
+            f"Processing chunk #{batch_num_pad} and writing to {batch_fullpath}"
         )
-        return
-    logging.info(f"Reading list of files in gs://{gcs_bucket}/{source_object_folder}/*")
-    storage_client = storage.Client(gcs_bucket)
-    blobs = storage_client.list_blobs(
-        bucket_or_name=gcs_bucket, prefix=str(source_object_folder) + "/"
-    )
-    blobs = [blob for blob in blobs if not blob.name.endswith("/")]
-    batch_size = math.ceil(len(blobs) / 6)
-    start, stop = batch_size * (batch_count - 1), batch_size * batch_count
-    logging.info(
-        f"bathch_count = {batch_count}, batch_size = {batch_size}, start= {start}, stop= {stop - 1} Inclusive"
-    )
-    logging.info(
-        f"\t\t\t\tBatch-{batch_count} files list for processing\n {[blob.name for blob in blobs[start:stop]]}"
-    )
-    for blob in blobs[start:stop]:
-        source_gcs_object = blob.name
-        file_name = source_gcs_object.split("/")[-1]
-        target_file = f"./source/{file_name}"
-        download_blob(gcs_bucket, source_gcs_object, target_file)
-        file_count = target_file.split("_")[-1]
-        output_file = f"./output/data_output_{file_count}.csv"
-        chunk_process(
-            target_file,
-            data_types,
-            date_cols,
-            rename_mappings,
-            non_na_columns,
-            output_file,
+        chunk.dropna(subset=non_na_columns, inplace=True)
+        chunk["trip_start_timestamp"] = pd.to_datetime(
+            chunk["trip_start_timestamp"], format="%m/%d/%Y %I:%M:%S %p"
         )
-        target_gcs_path = f"{output_object_folder}/{output_file.split('/')[-1]}"
-        upload_file_to_gcs(output_file, gcs_bucket, target_gcs_path)
-        logging.info(f"Removing {output_file} file")
-        os.remove(output_file)
+        chunk["trip_end_timestamp"] = pd.to_datetime(
+            chunk["trip_end_timestamp"], format="%m/%d/%Y %I:%M:%S %p"
+        )
+        chunk.to_csv(batch_fullpath, index=False)
     logging.info(
         f'Chicago Taxi Trips Dataset pipeline process completed at {str(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))}'
     )
 
 
-def download_file(source_url: str, source_file: str) -> None:
-    logging.info(f"Downloading data from {source_url} to {source_file} .")
-    res = requests.get(source_url, stream=True)
-    if res.status_code == 200:
-        with open(source_file, "wb") as fb:
-            for idx, chunk in enumerate(res):
-                fb.write(chunk)
-                if not idx % 10000000:
-                    file_size = os.stat(source_file).st_size / (1024**3)
-                    logging.info(
-                        f"\t{idx} chunks of data downloaded & current file size is {file_size} GB"
-                    )
-    else:
-        logging.info(f"Couldn't download {source_url}: {res.text}")
-    logging.info(
-        f"Successfully downloaded data from {source_url} into {source_file} - { os.stat(source_file).st_size / (1024 ** 3)} GB"
-    )
-
-
-def download_blob(gcs_bucket: str, source_gcs_object: str, target_file: str) -> None:
-    """Downloads a blob from the bucket."""
-    logging.info(
-        f"Downloading data from gs://{gcs_bucket}/{source_gcs_object} to {target_file} ..."
-    )
-    storage_client = storage.Client()
-    bucket = storage_client.bucket(gcs_bucket)
-    blob = bucket.blob(source_gcs_object)
-    blob.download_to_filename(target_file)
-    logging.info("Downloading Completed.")
-
-
-def chunk_process(
-    target_file: str,
-    data_types: dict,
-    date_cols: list,
-    rename_mappings: dict,
-    non_na_columns: list,
-    output_file: str,
-) -> None:
-    logging.info(f"Process started for {output_file} file")
-    logging.info(f"Reading file {target_file} to pandas dataframe...")
-    chunks = pd.read_csv(
-        target_file,
-        chunksize=500000,
-        dtype=data_types,
-        parse_dates=date_cols,
-        dayfirst=False,
-    )
-    logging.info(f"Removing {target_file} file")
-    os.remove(target_file)
-    for idx, chunk in enumerate(chunks):
-        logging.info("\tRenaming headers")
-        rename_headers(chunk, rename_mappings)
-        logging.info(
-            f"\tDropping null rows from specified columns in list = {non_na_columns}"
-        )
-        drop_null_rows(chunk, non_na_columns)
-        if not idx:
-            logging.info(f"\tWriting data to output file = {output_file}")
-            chunk.to_csv(output_file, mode="w", index=False, header=True)
-        else:
-            logging.info(f"\tAppending data to output file = {output_file}")
-            chunk.to_csv(output_file, mode="a", index=False, header=False)
-    logging.info(f"Process completed for {output_file} file")
-
-
-def rename_headers(df: pd.DataFrame, rename_mappings: dict) -> None:
-    df.rename(columns=rename_mappings, inplace=True)
-
-
-def drop_null_rows(df: pd.DataFrame, null_columns: list) -> None:
-    df.dropna(subset=null_columns, inplace=True)
-
-
-def upload_file_to_gcs(
-    target_csv_file: str, target_gcs_bucket: str, target_gcs_path: str
-) -> None:
-    logging.info(f"Uploading output file to gs://{target_gcs_bucket}/{target_gcs_path}")
-    storage_client = storage.Client()
-    bucket = storage_client.bucket(target_gcs_bucket)
-    blob = bucket.blob(target_gcs_path)
-    blob.upload_from_filename(target_csv_file)
-    logging.info("Successfully uploaded file to gcs bucket.")
-
-
 if __name__ == "__main__":
     logging.getLogger().setLevel(logging.INFO)
     main(
-        pipeline=os.environ.get("PIPELINE", ""),
         source_url=os.environ.get("SOURCE_URL", ""),
-        source_file=os.environ.get("SOURCE_FILE", ""),
         gcs_bucket=os.environ.get("GCS_BUCKET", ""),
         csv_gcs_path=os.environ.get("CSV_GCS_PATH", ""),
-        source_object_folder=os.environ.get("SOURCE_OBJECT_FOLDER", ""),
-        rename_mappings=json.loads(os.environ.get("RENAME_MAPPINGS", "{}")),
+        csv_headers=json.loads(os.environ.get("CSV_HEADERS", "{}")),
         non_na_columns=json.loads(os.environ.get("NON_NA_COLUMNS", "[]")),
-        output_object_folder=os.environ.get("OUTPUT_OBJECT_FOLDER", ""),
-        data_types=json.loads(os.environ.get("DATA_TYPES", "{}")),
-        date_cols=json.loads(os.environ.get("DATE_COLS", "[]")),
-        batch_count=int(os.environ.get("BATCH_COUNT", "0")),
+        data_dtypes=json.loads(os.environ.get("DATA_DTYPES", "{}")),
+        chunksize=int(os.environ.get("CHUNKSIZE", "1000000")),
     )

datasets/chicago_taxi_trips/pipelines/_images/run_csv_transform_kub/requirements.txt

@@ -1,3 +1,5 @@
+fsspec
+gcsfs
 google-cloud-storage
 pandas
 requests