feat: implement Kafka-to-Spark ingestion pipeline (Chapter 4)

.gitignore

@@ -21,6 +21,7 @@ pants.requirements.pex
 .DS_Store
 
 output_data
+output_json
 
 # Jupyter Notebooks
 .ipynb_checkpoints/

3rdparty/requirements.txt

@@ -1,3 +1,5 @@
 # Third-party Python dependencies for the data engineering projects.
 pyspark>=3.3.0,<4.0.0
 pre-commit>=3.0.0
+pyyaml
+confluent-kafka[jsonschema]

README.md

@@ -35,13 +35,24 @@ data-engineering/
 │   ├── requirements.txt       # Lists project requirements (pandas, PySpark, etc.)
 │   └── user_reqs.lock         # Pants generated dependency lockfile
 ├── projects/                  # Directory containing all sub-projects
+│   ├── ingestion/             # Chapter 4 Kafka/Spark ingestion project
+│   └── essentials/            # Chapter 2 basic Spark examples
 └── scripts/
     ├── BUILD                  # Configures scripts targets for Pants
     └── ai_pr_reviewer.py      # Python script that runs Gemini AI code reviews
 ```
 
 ---
 
+## 📁 Projects
+
+Each project under the `projects/` directory represents a separate learning milestone with self-contained instructions, docker components, and code:
+
+*   [projects/essentials/](projects/essentials/) — Basic local PySpark processing examples (see [projects/essentials/README.md](projects/essentials/README.md)).
+*   [projects/ingestion/](projects/ingestion/) — Real-time event ingestion using Kafka and Spark Streaming (see [projects/ingestion/README.md](projects/ingestion/README.md)).
+
+---
+
 ## 🚀 Getting Started
 
 ### 1. Prerequisites

agents.md

@@ -26,6 +26,13 @@ This repository is a **Python Data Engineering Monorepo** managed by the **Pants
   * `employee_partition_by_hire_date.py`: Local partitioning Spark script.
   * `input_data/`: Small CSV/txt sample inputs.
   * `output_data/`: Automatically generated Spark output targets (ignored by git).
+* `projects/ingestion/`: Ingestion project (derived implementation from *Hello Modern Data Pipelines*, Chapter 4).
+  * `config/input_config.yml`: Spark Ingestion configuration YAML file.
+  * `docker/docker-compose.yml`: Zookeeper, Kafka, and Schema Registry Compose setup.
+  * `input_data/user_events.json`: Sample event stream dataset.
+  * `json_producer.py`: Kafka JSON message producer script.
+  * `kafka_json_to_file_job.py`: PySpark job ingestion script with Spark SQL Kafka integration.
+  * `BUILD`: Pants build definition for the ingestion project.
 * `scripts/`: Python utility scripts.
   * `ai_pr_reviewer.py`: The AI code reviewer script powered by the Gemini API.
   * `BUILD`: Pants build definition for the scripts directory.

projects/essentials/README.md

@@ -0,0 +1,34 @@
+# Essentials Project (Chapter 2)
+
+This project contains initial Spark processing examples derived from Chapter 2 of *Hello Modern Data Pipelines*. It demonstrates basic local batch data processing using PySpark.
+
+---
+
+## 📁 Project Contents
+
+*   [word_count.py](projects/essentials/word_count.py): Processes unstructured text data to perform a classic word-count calculation.
+*   [employee_partition_by_hire_date.py](projects/essentials/employee_partition_by_hire_date.py): Demonstrates PySpark DataFrame API usage, reading CSV employee data and partitioning the output by hire date.
+*   `input_data/`: Contains sample CSV and text input files, such as [employee_data.csv](projects/essentials/input_data/employee_data.csv), for testing the scripts.
+
+---
+
+## 🚀 How to Run
+
+Before running the scripts, ensure your virtual environment is active:
+```bash
+source .venv/bin/activate
+```
+
+### 1. Run WordCount
+Run the word count script:
+```bash
+python projects/essentials/word_count.py
+```
+This generates the results in `projects/essentials/output_data/word_count/`.
+
+### 2. Run Employee Partitioning
+Run the partitioning script:
+```bash
+python projects/essentials/employee_partition_by_hire_date.py
+```
+This partitions the employee data and writes it to `projects/essentials/output_data/employee_partition/`.

projects/ingestion/BUILD

@@ -0,0 +1,3 @@
+python_sources(
+    name="lib",
+)

projects/ingestion/README.md

@@ -0,0 +1,43 @@
+# Ingestion Project (Chapter 4)
+
+This project implements the advanced data ingestion and integration patterns derived from Chapter 4 of *Hello Modern Data Pipelines*. It demonstrates real-time integration by producing events to a Kafka broker and streaming/reading them into Spark.
+
+---
+
+## 📁 Project Contents
+
+*   [docker/docker-compose.yml](projects/ingestion/docker/docker-compose.yml): Launches local Zookeeper, Kafka, and Schema Registry containers.
+*   [json_producer.py](projects/ingestion/json_producer.py): Python producer script that publishes mock events to the Kafka broker.
+*   [kafka_json_to_file_job.py](projects/ingestion/kafka_json_to_file_job.py): PySpark streaming/batch job that pulls events from Kafka, parses them with a schema, and saves them locally as JSON.
+*   [config/input_config.yml](projects/ingestion/config/input_config.yml): Configuration file specifying broker connections, target topics, and JSON parsing schema.
+*   [input_data/user_events.json](projects/ingestion/input_data/user_events.json): Sample JSON data file with mock events.
+
+---
+
+## 🚀 How to Run
+
+Ensure your virtual environment is active:
+```bash
+source .venv/bin/activate
+```
+
+### 1. Spin up Kafka Infrastucture
+Start the Zookeeper, Kafka, and Schema Registry containers:
+```bash
+cd projects/ingestion/docker
+docker-compose up -d
+cd ../../..
+```
+
+### 2. Produce mock JSON events
+Publish the sample events from `user_events.json` into Kafka:
+```bash
+python projects/ingestion/json_producer.py
+```
+
+### 3. Run Ingestion Spark Job
+Extract the events from Kafka and write them to output directories:
+```bash
+python projects/ingestion/kafka_json_to_file_job.py
+```
+Outputs are written locally to `projects/ingestion/output_json/user_events/`.

projects/ingestion/config/input_config.yml

@@ -0,0 +1,8 @@
+data_sources:
+  - source_id: user_events_kafka
+    source_type: kafka
+    kafka_config:
+      bootstrap_servers: "localhost:9092"
+      topic: "user-events-json"
+      starting_offsets: "earliest"
+    json_schema: "STRUCT<user_id:INT,event_type:STRING,timestamp:STRING>"

projects/ingestion/docker/docker-compose.yml

@@ -0,0 +1,27 @@
+version: '2'
+services:
+  zookeeper:
+    image: confluentinc/cp-zookeeper:7.2.1
+    environment:
+      ZOOKEEPER_CLIENT_PORT: 2181
+
+  kafka:
+    image: confluentinc/cp-kafka:7.2.1
+    ports:
+      - "9092:9092"
+    environment:
+      KAFKA_BROKER_ID: 1
+      KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181
+      KAFKA_LISTENER_SECURITY_PROTOCOL_MAP: PLAINTEXT:PLAINTEXT,PLAINTEXT_HOST:PLAINTEXT
+      KAFKA_ADVERTISED_LISTENERS: PLAINTEXT://kafka:29092,PLAINTEXT_HOST://localhost:9092
+      KAFKA_LISTENERS: PLAINTEXT://0.0.0.0:29092,PLAINTEXT_HOST://0.0.0.0:9092
+      KAFKA_INTER_BROKER_LISTENER_NAME: PLAINTEXT
+      KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: 1
+
+  schema-registry:
+    image: confluentinc/cp-schema-registry:7.2.1
+    ports:
+      - "8081:8081"
+    environment:
+      SCHEMA_REGISTRY_HOST_NAME: schema-registry
+      SCHEMA_REGISTRY_KAFKASTORE_BOOTSTRAP_SERVERS: PLAINTEXT://kafka:29092

projects/ingestion/input_data/user_events.json

@@ -0,0 +1,3 @@
+{"user_id": 101, "event_type": "login", "timestamp": "2025-07-07T09:00:00Z"}
+{"user_id": 102, "event_type": "logout", "timestamp": "2025-07-07T09:05:00Z"}
+{"user_id": 104, "event_type": "login", "timestamp": "2025-07-07T09:05:00Z"}

projects/ingestion/json_producer.py

@@ -0,0 +1,42 @@
+import json
+import os
+import sys
+from confluent_kafka import Producer
+
+# Resolve paths relative to this script
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+INPUT_FILE_PATH = os.path.join(SCRIPT_DIR, "input_data/user_events.json")
+
+# Kafka producer configuration
+producer_conf = {"bootstrap.servers": "localhost:9092"}
+producer = Producer(producer_conf)
+
+# Read events from input file with error handling
+records = []
+try:
+    with open(INPUT_FILE_PATH, "r", encoding="utf-8") as f:
+        for line in f:
+            if line.strip():
+                try:
+                    records.append(json.loads(line))
+                except json.JSONDecodeError as e:
+                    print(
+                        f"Skipping malformed JSON line: {line.strip()} - Error: {e}",
+                        file=sys.stderr,
+                    )
+except FileNotFoundError:
+    print(f"Error: Input file not found at {INPUT_FILE_PATH}", file=sys.stderr)
+    sys.exit(1)
+except IOError as e:
+    print(f"Error reading input file {INPUT_FILE_PATH}: {e}", file=sys.stderr)
+    sys.exit(1)
+
+topic = "user-events-json"
+
+# Produce each record to Kafka
+for record in records:
+    json_str = json.dumps(record)
+    producer.produce(topic=topic, value=json_str)
+    print("Produced:", json_str)
+
+producer.flush()

projects/ingestion/kafka_json_to_file_job.py

@@ -0,0 +1,144 @@
+# %% [markdown]
+# # Kafka JSON to File Ingestion Job
+# Extracts user events from a Kafka topic, parses them according to a schema config, and writes them to the local filesystem.
+
+# %%
+import argparse
+import os
+import yaml
+from pyspark.sql import SparkSession
+from pyspark.sql.functions import col, from_json
+
+
+# %%
+def load_config(path: str) -> dict:
+    """Loads YAML configuration file."""
+    with open(path, "r", encoding="utf-8") as f:
+        return yaml.safe_load(f)
+
+
+# %%
+def run_ingestion_job(config_path: str, output_dir: str, is_streaming: bool = True):
+    """Runs Spark Ingestion job to extract JSON events from Kafka,
+
+    parses them using a JSON schema, and writes them to output_dir.
+
+    Args:
+        config_path (str): Path to the ingestion YAML configuration file.
+        output_dir (str): Target directory to save the output JSON events.
+        is_streaming (bool): If True, run as Structured Streaming job. If False, run as Batch job.
+    """
+    config = load_config(config_path)
+
+    # Configuration Validation
+    if not isinstance(config, dict) or "data_sources" not in config:
+        raise ValueError("Invalid configuration file structure: missing 'data_sources'")
+
+    data_sources = config["data_sources"]
+    if not data_sources or not isinstance(data_sources, list):
+        raise ValueError("'data_sources' must be a non-empty list")
+
+    source_conf = data_sources[0]
+    if "kafka_config" not in source_conf:
+        raise ValueError("Missing 'kafka_config' in configuration source")
+
+    kafka_conf = source_conf["kafka_config"]
+    if "bootstrap_servers" not in kafka_conf or "topic" not in kafka_conf:
+        raise ValueError("Missing 'bootstrap_servers' or 'topic' in 'kafka_config'")
+
+    spark = (
+        SparkSession.builder.appName("KafkaJsonToFile")
+        .master("local[*]")
+        .config(
+            "spark.jars.packages",
+            "org.apache.spark:spark-sql-kafka-0-10_2.12:3.5.8",
+        )
+        .getOrCreate()
+    )
+
+    try:
+        kafka_options = {
+            "kafka.bootstrap.servers": kafka_conf["bootstrap_servers"],
+            "subscribe": kafka_conf["topic"],
+            "startingOffsets": kafka_conf.get("starting_offsets", "earliest"),
+        }
+
+        json_schema_str = source_conf.get("json_schema")
+        if not json_schema_str:
+            raise ValueError("Missing 'json_schema' in config")
+
+        if is_streaming:
+            # Structured Streaming Read
+            df_raw = spark.readStream.format("kafka").options(**kafka_options).load()
+            df_json = df_raw.selectExpr("CAST(value AS STRING) as json_str")
+            df_parsed = df_json.select(
+                from_json(col("json_str"), json_schema_str).alias("data")
+            ).select("data.*")
+
+            checkpoint_path = os.path.join(output_dir, "_checkpoint")
+
+            # Structured Streaming Write (using Append mode)
+            query = (
+                df_parsed.writeStream.outputMode("append")
+                .format("json")
+                .option("path", output_dir)
+                .option("checkpointLocation", checkpoint_path)
+                .trigger(processingTime="10 seconds")
+                .start()
+            )
+            try:
+                # Blocks the main thread, keeping the streaming query active until interrupted
+                query.awaitTermination()
+            except KeyboardInterrupt:
+                # Triggered when the user presses Ctrl+C in the terminal
+                import sys
+
+                print(
+                    "Streaming query interrupted by user. Stopping...",
+                    file=sys.stderr,
+                )
+        else:
+            # Batch Read
+            df_raw = spark.read.format("kafka").options(**kafka_options).load()
+            df_json = df_raw.selectExpr("CAST(value AS STRING) as json_str")
+            df_parsed = df_json.select(
+                from_json(col("json_str"), json_schema_str).alias("data")
+            ).select("data.*")
+
+            # Coalesce to 1 partition for small datasets to avoid many small files.
+            # For larger datasets, remove .coalesce(1) to let Spark manage partitions
+            # or repartition based on business keys.
+            df_parsed.coalesce(1).write.mode("overwrite").json(output_dir)
+    finally:
+        spark.stop()
+
+
+# %%
+if __name__ == "__main__":
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    default_config = os.path.join(script_dir, "config/input_config.yml")
+    default_output = os.path.join(script_dir, "output_json/user_events")
+
+    parser = argparse.ArgumentParser(
+        description="Extract user events from Kafka and save as JSON."
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        default=default_config,
+        help="Path to the input configuration YAML file.",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default=default_output,
+        help="Path to save the output JSON events.",
+    )
+    parser.add_argument(
+        "--batch",
+        action="store_true",
+        help="Run as a one-time batch job instead of a continuous streaming job.",
+    )
+
+    args = parser.parse_args()
+    run_ingestion_job(args.config, args.output, is_streaming=not args.batch)