# Spark/PySpark vs Geneva: Image Embedding & Captioning Pipeline

This file compares Apache Spark/PySpark with Geneva for building
production image embedding pipelines.

## Key Insight

Spark is the industry standard for big data. But it was designed for
SQL analytics and ETL, not GPU-heavy ML inference. The UDF system
shows its age when running modern deep learning workloads.

Think of it as:
- **Spark** = Big data SQL/ETL engine (with ML bolted on)
- **Geneva** = Feature engineering platform (built for ML from day one)

## Summary

| Approach | Lines of Code |
|----------|---------------|
| Spark | ~150 lines for computation + ~200 lines for production ops |
| Geneva | ~80 lines total, production ops included automatically |

---

## Spark/PySpark Approach

### Spark Excels At

- SQL analytics at massive scale
- ETL pipelines (extract, transform, load)
- Mature ecosystem (Hive, Delta Lake, etc.)
- Enterprise adoption and support

### Challenges for ML Inference

- UDF serialization constraints (can't pickle arbitrary objects)
- Model broadcasting/caching is manual and error-prone
- GPU support requires RAPIDS or external systems (Triton)
- Checkpointing only for Structured Streaming, not batch jobs
- JVM/Python serialization overhead

---

### Basic pandas_udf Approach (Inefficient)

```python
from pyspark.sql import SparkSession
from pyspark.sql.functions import pandas_udf, col
from pyspark.sql.types import ArrayType, FloatType, StringType
import pandas as pd
import numpy as np
from typing import Iterator


@pandas_udf(ArrayType(FloatType()))
def compute_embedding_spark_basic(images: pd.Series) -> pd.Series:
    """
    CLIP embedding using Spark's pandas_udf.

    PROBLEM: Model loads EVERY BATCH (not even per-partition).
    This is extremely inefficient for large models.
    """
    import open_clip
    import torch
    from PIL import Image
    import io

    # Model loads every batch! Very inefficient.
    model, _, preprocess = open_clip.create_model_and_transforms(
        'ViT-B-32', pretrained='laion2b_s34b_b79k'
    )
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model.to(device).eval()

    results = []
    for img_bytes in images:
        if img_bytes is None:
            results.append([0.0] * 512)
            continue
        try:
            img = Image.open(io.BytesIO(img_bytes)).convert('RGB')
            tensor = preprocess(img).unsqueeze(0).to(device)
            with torch.no_grad():
                embedding = model.encode_image(tensor)
            results.append(embedding.cpu().numpy().flatten().tolist())
        except Exception:
            results.append([0.0] * 512)

    return pd.Series(results)
```

### Iterator pandas_udf (Better, but complex)

```python
def create_embedding_udf_with_caching():
    """
    Use mapInPandas with iterator to cache model across batches.

    This is the recommended pattern but it's complex and fragile.
    """

    def compute_embeddings_iterator(
        iterator: Iterator[pd.DataFrame]
    ) -> Iterator[pd.DataFrame]:
        """
        Iterator-based UDF that loads model once per partition.

        Complexity:
        - Must handle iterator protocol correctly
        - Model loading happens once per partition (not per worker)
        - Error handling is tricky (exceptions kill the partition)
        - GPU memory management is manual
        """
        import open_clip
        import torch
        from PIL import Image
        import io

        # Load model once per partition
        model, _, preprocess = open_clip.create_model_and_transforms(
            'ViT-B-32', pretrained='laion2b_s34b_b79k'
        )
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model = model.to(device).eval()

        for batch_df in iterator:
            embeddings = []
            for img_bytes in batch_df["image"]:
                if img_bytes is None:
                    embeddings.append([0.0] * 512)
                    continue
                try:
                    img = Image.open(io.BytesIO(img_bytes)).convert('RGB')
                    tensor = preprocess(img).unsqueeze(0).to(device)
                    with torch.no_grad():
                        embedding = model.encode_image(tensor)
                    embeddings.append(embedding.cpu().numpy().flatten().tolist())
                except Exception:
                    embeddings.append([0.0] * 512)

            batch_df = batch_df.copy()
            batch_df["embedding"] = embeddings
            yield batch_df

    return compute_embeddings_iterator


def create_caption_udf_with_caching():
    """BLIP captioning with iterator pattern."""

    def compute_captions_iterator(
        iterator: Iterator[pd.DataFrame]
    ) -> Iterator[pd.DataFrame]:
        from transformers import BlipProcessor, BlipForConditionalGeneration
        import torch
        from PIL import Image
        import io

        device = "cuda" if torch.cuda.is_available() else "cpu"
        processor = BlipProcessor.from_pretrained(
            "Salesforce/blip-image-captioning-base"
        )
        model = BlipForConditionalGeneration.from_pretrained(
            "Salesforce/blip-image-captioning-base"
        ).to(device).eval()

        for batch_df in iterator:
            captions = []
            for img_bytes in batch_df["image"]:
                if img_bytes is None:
                    captions.append("")
                    continue
                try:
                    img = Image.open(io.BytesIO(img_bytes)).convert('RGB')
                    inputs = processor(img, return_tensors="pt")
                    inputs = {k: v.to(device) for k, v in inputs.items()}
                    with torch.no_grad():
                        output = model.generate(**inputs, max_new_tokens=50)
                    captions.append(processor.decode(output[0], skip_special_tokens=True))
                except Exception:
                    captions.append("")

            batch_df = batch_df.copy()
            batch_df["caption"] = captions
            yield batch_df

    return compute_captions_iterator
```

### Spark 3.4+ predict_batch_udf (Newest API)

```python
def create_predict_batch_udf():
    """
    Spark 3.4+ predict_batch_udf - simplest API but still has issues.

    Limitations:
    - Still loads model per task (not per executor)
    - GPU memory tuning required (tasks per executor)
    - No built-in checkpointing
    """
    from pyspark.ml.functions import predict_batch_udf

    def load_model():
        import open_clip
        import torch

        model, _, preprocess = open_clip.create_model_and_transforms(
            'ViT-B-32', pretrained='laion2b_s34b_b79k'
        )
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model = model.to(device).eval()
        return model, preprocess, device

    def predict_fn(model_state):
        model, preprocess, device = model_state
        import torch
        from PIL import Image
        import io

        def predict(inputs):
            results = []
            for img_bytes in inputs:
                if img_bytes is None:
                    results.append(np.zeros(512, dtype=np.float32))
                    continue
                try:
                    img = Image.open(io.BytesIO(img_bytes)).convert('RGB')
                    tensor = preprocess(img).unsqueeze(0).to(device)
                    with torch.no_grad():
                        embedding = model.encode_image(tensor)
                    results.append(embedding.cpu().numpy().flatten())
                except Exception:
                    results.append(np.zeros(512, dtype=np.float32))
            return np.array(results)

        return predict

    return predict_batch_udf(
        load_model,
        predict_fn,
        return_type=ArrayType(FloatType()),
        batch_size=32,
    )
```

### Basic Pipeline

```python
def run_spark_pipeline_basic():
    """
    Basic Spark pipeline for image embeddings.

    Note: No checkpointing for batch jobs. If it fails,
    you restart from scratch.
    """
    from pyspark.sql.types import StructType, StructField, BinaryType

    spark = SparkSession.builder \
        .appName("ImageEmbeddings") \
        .config("spark.sql.execution.arrow.pyspark.enabled", "true") \
        .config("spark.sql.execution.arrow.maxRecordsPerBatch", "1000") \
        .getOrCreate()

    # Read data
    df = spark.read.parquet("s3://my-bucket/images/")

    # Define output schema for mapInPandas
    output_schema = StructType([
        StructField("id", StringType()),
        StructField("image", BinaryType()),
        StructField("embedding", ArrayType(FloatType())),
    ])

    # Compute embeddings using iterator pattern
    df = df.mapInPandas(
        create_embedding_udf_with_caching(),
        schema=output_schema,
    )

    # Add captions
    output_schema_with_caption = StructType([
        StructField("id", StringType()),
        StructField("image", BinaryType()),
        StructField("embedding", ArrayType(FloatType())),
        StructField("caption", StringType()),
    ])

    df = df.mapInPandas(
        create_caption_udf_with_caching(),
        schema=output_schema_with_caption,
    )

    # Write results
    df.write.mode("overwrite").parquet("s3://my-bucket/embeddings/")

    spark.stop()
```

### Production Pipeline (with custom ops)

For production with Spark batch jobs, you need to build:
1. **Checkpointing** - Spark Streaming has it, batch jobs DO NOT
2. **Incremental processing** - Manual partition tracking
3. **Progress tracking** - Custom logging
4. **Failure recovery** - Manual restart logic

**Key limitation**: Spark checkpointing is for Structured Streaming only.
Batch jobs have no built-in checkpoint/resume capability.

```python
import json
from pathlib import Path
from datetime import datetime


class SparkCheckpointManager:
    """
    Custom checkpoint manager for Spark batch jobs.

    You build this yourself - Spark batch jobs don't checkpoint.

    Note: Spark Structured Streaming has checkpointing, but that's
    for streaming, not batch inference jobs.
    """

    def __init__(self, checkpoint_dir: str):
        self.checkpoint_dir = Path(checkpoint_dir)
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        self.state_file = self.checkpoint_dir / "state.json"
        self.state = self._load_state()

    def _load_state(self) -> dict:
        if self.state_file.exists():
            return json.loads(self.state_file.read_text())
        return {"processed_partitions": [], "rows_processed": 0}

    def _save_state(self):
        self.state["last_checkpoint"] = datetime.now().isoformat()
        self.state_file.write_text(json.dumps(self.state, indent=2))

    def mark_partition_processed(self, partition_path: str, row_count: int):
        self.state["processed_partitions"].append(partition_path)
        self.state["rows_processed"] += row_count
        self._save_state()

    def is_partition_processed(self, partition_path: str) -> bool:
        return partition_path in self.state["processed_partitions"]

    def get_stats(self) -> dict:
        return self.state.copy()


class SparkIncrementalProcessor:
    """
    Incremental processing wrapper for Spark batch jobs.

    Implements partition-level checkpointing since Spark batch
    jobs don't have built-in checkpoint/resume.
    """

    def __init__(
        self,
        input_path: str,
        output_path: str,
        checkpoint_dir: str,
    ):
        self.input_path = input_path
        self.output_path = output_path
        self.checkpoint = SparkCheckpointManager(checkpoint_dir)
        self.spark = None

    def _init_spark(self):
        if self.spark is None:
            self.spark = SparkSession.builder \
                .appName("ImageEmbeddings") \
                .config("spark.sql.execution.arrow.pyspark.enabled", "true") \
                .getOrCreate()

    def _list_input_partitions(self) -> list[str]:
        """List all parquet files/partitions."""
        import subprocess
        result = subprocess.run(
            ["aws", "s3", "ls", self.input_path, "--recursive"],
            capture_output=True, text=True
        )
        return [
            line.split()[-1]
            for line in result.stdout.strip().split("\n")
            if line.endswith(".parquet")
        ]

    def _get_pending_partitions(self) -> list[str]:
        all_parts = self._list_input_partitions()
        return [p for p in all_parts if not self.checkpoint.is_partition_processed(p)]

    def run(self):
        """Run incremental processing with checkpointing."""
        self._init_spark()

        pending = self._get_pending_partitions()
        print(f"Partitions to process: {len(pending)}")

        if not pending:
            print("Nothing to process!")
            return

        output_schema = StructType([
            StructField("id", StringType()),
            StructField("image", BinaryType()),
            StructField("embedding", ArrayType(FloatType())),
            StructField("caption", StringType()),
        ])

        # Process partition by partition for checkpointing
        for i, partition_path in enumerate(pending):
            full_path = f"{self.input_path}/{partition_path}"
            print(f"Processing [{i+1}/{len(pending)}]: {partition_path}")

            try:
                df = self.spark.read.parquet(full_path)

                # Apply transformations
                df = df.mapInPandas(
                    create_embedding_udf_with_caching(),
                    schema=output_schema,
                )
                df = df.mapInPandas(
                    create_caption_udf_with_caching(),
                    schema=output_schema,
                )

                # Count and write
                row_count = df.count()
                output_file = f"{self.output_path}/batch_{i:06d}.parquet"
                df.write.mode("overwrite").parquet(output_file)

                self.checkpoint.mark_partition_processed(partition_path, row_count)
                print(f"  Completed: {row_count} rows")

            except Exception as e:
                print(f"  FAILED: {e}")
                print("  Re-run to resume from this partition.")
                raise

        print(f"Complete! Total rows: {self.checkpoint.get_stats()['rows_processed']}")
        self.spark.stop()


def run_spark_pipeline_production():
    """Production Spark pipeline with custom checkpointing."""
    processor = SparkIncrementalProcessor(
        input_path="s3://my-bucket/images",
        output_path="s3://my-bucket/embeddings",
        checkpoint_dir="/tmp/spark_checkpoints",
    )
    processor.run()
```

### GPU Considerations

**Spark GPU support options:**

1. **NVIDIA RAPIDS Accelerator**
   - Accelerates Spark SQL/DataFrame operations
   - Does NOT help with Python UDFs
   - Requires specific Spark configuration

2. **NVIDIA Triton Inference Server**
   - Offload inference to separate GPU service
   - Decouples Spark from GPU management
   - Additional infrastructure to manage

3. **Direct GPU in UDFs (what we show above)**
   - Each task loads model independently
   - GPU memory contention issues
   - Must tune spark.task.resource.gpu.amount

**GPU memory management:**
- Set spark.executor.resource.gpu.amount = 1
- Set spark.task.resource.gpu.amount = 0.5 (2 tasks per GPU)
- Tune based on model size and GPU memory
- Very error-prone and cluster-specific

### Dependency Management

Options for managing dependencies in Spark (all require significant DevOps work):

1. **Pre-install on cluster (recommended)**
   - All nodes must have same packages
   - Version mismatches cause silent failures
   - Requires cluster restart for updates

2. **spark-submit --py-files**
   - Only works for pure Python packages
   - No support for compiled extensions (torch, etc.)

3. **Conda/virtualenv on cluster**
   - Complex setup
   - Easy to have version drift

4. **Docker/Kubernetes**
   - Most reliable but complex
   - Requires custom Spark images

---

## Geneva Approach

Geneva provides:
- Stateful UDFs (model loads once per worker)
- Fragment-level AND batch-level checkpointing
- Automatic incremental processing
- Per-UDF dependency isolation
- Table versioning
- No JVM/serialization overhead

### UDF Definitions

```python
import geneva


@geneva.udf(num_gpus=1, version="0.1")
class ComputeEmbeddingGeneva:
    """
    CLIP embedding UDF for Geneva.

    Key differences from Spark:
    - setup() loads model ONCE per worker (not per task/batch)
    - No serialization constraints
    - No JVM overhead
    - GPU scheduling automatic
    - Type annotation -> schema
    """

    def setup(self):
        import open_clip
        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
            'ViT-B-32', pretrained='laion2b_s34b_b79k'
        )
        self.model = self.model.cuda().eval()

    def __call__(self, image: bytes) -> list[float]:
        from PIL import Image
        import torch
        import io

        img = Image.open(io.BytesIO(image)).convert('RGB')
        tensor = self.preprocess(img).unsqueeze(0).cuda()
        with torch.no_grad():
            embedding = self.model.encode_image(tensor)
        return embedding.cpu().numpy().flatten().tolist()


@geneva.udf(num_gpus=1, version="0.1")
class GenerateCaptionGeneva:
    """BLIP captioning UDF for Geneva."""

    def setup(self):
        from transformers import BlipProcessor, BlipForConditionalGeneration
        self.processor = BlipProcessor.from_pretrained(
            "Salesforce/blip-image-captioning-base"
        )
        self.model = BlipForConditionalGeneration.from_pretrained(
            "Salesforce/blip-image-captioning-base"
        ).cuda()

    def __call__(self, image: bytes) -> str:
        from PIL import Image
        import torch
        import io

        img = Image.open(io.BytesIO(image)).convert('RGB')
        inputs = self.processor(img, return_tensors="pt")
        inputs = {k: v.cuda() for k, v in inputs.items()}
        with torch.no_grad():
            output = self.model.generate(**inputs, max_new_tokens=50)
        return self.processor.decode(output[0], skip_special_tokens=True)
```

### Pipeline

```python
def run_geneva_pipeline():
    """
    Geneva pipeline - no JVM, no serialization issues, production-ready.
    """
    conn = geneva.connect("s3://my-bucket/geneva-workspace")
    table = conn.open_table("images")

    table.add_columns({
        "embedding": ComputeEmbeddingGeneva,
        "caption": GenerateCaptionGeneva,
    })

    # Automatic checkpointing at fragment AND batch level
    table.backfill("embedding", batch_size=100, commit_granularity=10)
    table.backfill("caption", batch_size=50, commit_granularity=10)

    print(f"Processed {table.count_rows()} images")
    print(f"Table version: {table.version}")
```

---

## Side-by-Side Comparison

### Feature Comparison

| Feature | Spark/PySpark | Geneva |
|---------|---------------|--------|
| SQL analytics | Excellent | Via Lance/DuckDB |
| ETL pipelines | Industry standard | Via Ray Data |
| Enterprise ecosystem | Mature | Growing |
| Structured Streaming | Built-in | Via Ray |
| **UDF model caching** | Complex [1] | setup() automatic |
| **GPU scheduling** | Manual tuning [2] | num_gpus=1 |
| **Serialization overhead** | JVM/Arrow [3] | Native Python |
| **Batch job checkpointing** | Not included [4] | Fragment + batch |
| **Incremental processing** | Build yourself | Automatic |
| **Dependency isolation** | Cluster-wide [5] | Per-UDF automatic |
| **Table versioning** | Delta Lake | Built-in |

**Spark limitations for ML inference:**
1. Requires mapInPandas iterator pattern or predict_batch_udf (3.4+)
2. Must tune spark.task.resource.gpu.amount manually
3. Data crosses JVM/Python boundary via Arrow
4. Checkpointing only for Structured Streaming, not batch jobs
5. All executors must have same packages pre-installed

### Code Comparison

| Component | Spark/PySpark | Geneva |
|-----------|---------------|--------|
| UDF definitions | ~80 lines [1] | ~45 lines |
| Basic pipeline | ~40 lines | ~15 lines |
| Checkpoint manager | ~50 lines (custom) | 0 (included) |
| Incremental processor | ~80 lines (custom) | 0 (included) |
| GPU config | ~20 lines | 0 (automatic) |
| Schema definitions | ~15 lines | 0 (from types) |
| **TOTAL (production)** | **~285 lines** | **~60 lines** |

[1] Using mapInPandas iterator pattern for model caching

### Model Loading Comparison

**Spark pandas_udf (basic):**
- 100 batches × 10 partitions = model loads 1000 times
- Extremely inefficient for large models

**Spark mapInPandas (iterator):**
- 10 partitions = model loads 10 times
- Better, but still per-partition, not per-executor

**Spark predict_batch_udf (3.4+):**
- Model cached per task, but tasks >> executors
- Still inefficient, GPU memory issues

**Geneva:**
- 4 workers = model loads 4 times (once per worker)
- Efficient, predictable GPU memory usage

### Checkpoint Comparison

| Scenario | Spark Batch Job | Spark Structured Streaming | Geneva |
|----------|-----------------|---------------------------|--------|
| Job fails at row 500,000 of 1,000,000 | Restart from row 0, reprocess all 500,000 rows, no built-in checkpointing | Has checkpointing, but for streaming not batch | Restart from fragment 50 batch 15, no reprocessing |

---

## When to Use Each

### Spark is great for:
- SQL analytics on structured data
- ETL pipelines with existing Spark infrastructure
- When you need Delta Lake / Iceberg integration
- Teams with existing Spark expertise
- Streaming pipelines (Structured Streaming)

### Geneva is great for:
- GPU-heavy ML inference at scale
- Production feature engineering pipelines
- When fault tolerance for batch jobs is critical
- Teams that want to focus on ML, not infrastructure
- Multimodal data (images, video, audio)

---

## The Key Insight

**Spark was built for SQL analytics. Geneva was built for ML.**

**Spark asks**: "How do I run SQL on big data?"
**Geneva asks**: "How do I compute and maintain ML features?"

Spark's UDF system was added for extensibility.
Geneva's UDF system is the core design.

---

## Architecture

### Spark Architecture (you build the gray boxes)

**Your Code** (you build these):

| Checkpoint Manager | Incremental Processor | GPU Tuning | Iterator UDF Logic |
|--------------------|----------------------|------------|-------------------|
| (custom) | (custom) | (manual) | (complex) |

↓

**Spark**

| Catalyst Optimizer | pandas_udf API | Arrow Transfer | JVM Executors |
|--------------------|----------------|----------------|---------------|

### Geneva Architecture (everything included)

**Your Code**

| UDFs (ML logic) | backfill commands |
|-----------------|-------------------|

↓

**Geneva**

| Stateful UDFs | Fragment Checkpoint | Incremental Detection | Dependency Isolation | Table Version |
|---------------|--------------------|-----------------------|---------------------|---------------|

↓

**Ray + Lance Storage** *(No JVM, native Python execution)*
