databuild/examples/podcast_reviews/categorize_reviews_job.py

#!/usr/bin/env python3

import sys
import json
import os
import duckdb
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Any
import re

def main():
    if len(sys.argv) < 2:
        print("Usage: categorize_reviews_job.py {config|exec} [args...]", file=sys.stderr)
        sys.exit(1)
    
    command = sys.argv[1]
    
    if command == "config":
        handle_config(sys.argv[2:])
    elif command == "exec":
        handle_exec(sys.argv[2:])
    else:
        print(f"Unknown command: {command}", file=sys.stderr)
        sys.exit(1)

def parse_partition_ref(partition_ref: str) -> Dict[str, str]:
    """Parse partition ref like 'categorized_reviews/category=comedy/date=2020-01-01' into components."""
    match = re.match(r'categorized_reviews/category=([^/]+)/date=(\d{4}-\d{2}-\d{2})', partition_ref)
    if not match:
        raise ValueError(f"Invalid partition ref format: {partition_ref}")
    return {"category": match.group(1), "date": match.group(2)}

def handle_config(args):
    if len(args) < 1:
        print("Config mode requires partition ref", file=sys.stderr)
        sys.exit(1)
    
    partition_ref = args[0]
    
    try:
        parsed = parse_partition_ref(partition_ref)
        category = parsed["category"]
        date_str = parsed["date"]
    except ValueError as e:
        print(f"Error parsing partition ref: {e}", file=sys.stderr)
        sys.exit(1)
    
    # Dependencies: reviews for the date and podcast metadata
    reviews_ref = f"reviews/date={date_str}"
    podcasts_ref = "podcasts/all"
    
    config = {
        "configs": [{
            "outputs": [{"str": partition_ref}],
            "inputs": [
                {"dep_type": 1, "partition_ref": {"str": reviews_ref}},
                {"dep_type": 1, "partition_ref": {"str": podcasts_ref}}
            ],
            "args": [category, date_str],
            "env": {
                "PARTITION_REF": partition_ref,
                "TARGET_CATEGORY": category,
                "TARGET_DATE": date_str
            }
        }]
    }
    
    print(json.dumps(config))

def handle_exec(args):
    if len(args) < 2:
        print("Exec mode requires category and date arguments", file=sys.stderr)
        sys.exit(1)
    
    target_category = args[0]
    target_date = args[1]
    partition_ref = os.getenv('PARTITION_REF', f'categorized_reviews/category={target_category}/date={target_date}')
    
    # Input paths
    reviews_file = f"/tmp/databuild_test/examples/podcast_reviews/reviews/date={target_date}/reviews.parquet"
    podcasts_file = "/tmp/databuild_test/examples/podcast_reviews/podcasts/podcasts.parquet"
    
    # Check input files exist
    if not os.path.exists(reviews_file):
        print(f"Reviews file not found: {reviews_file}", file=sys.stderr)
        sys.exit(1)
    
    if not os.path.exists(podcasts_file):
        print(f"Podcasts file not found: {podcasts_file}", file=sys.stderr)
        sys.exit(1)
    
    # Output path
    output_dir = Path(f"/tmp/databuild_test/examples/podcast_reviews/categorized_reviews/category={target_category}/date={target_date}")
    output_dir.mkdir(parents=True, exist_ok=True)
    output_file = output_dir / "categorized_reviews.parquet"
    
    try:
        # Categorize reviews by joining with podcast metadata
        categorize_reviews_for_category_date(reviews_file, podcasts_file, target_category, str(output_file))
        
        print(f"Successfully categorized reviews for category {target_category} on {target_date}")
        print(f"Output written to: {output_file}")
        
        # Create manifest
        manifest = {
            "outputs": [{"str": partition_ref}],
            "inputs": [
                {"str": f"reviews/date={target_date}"},
                {"str": "podcasts/all"}
            ],
            "start_time": datetime.now().isoformat(),
            "end_time": datetime.now().isoformat(),
            "task": {
                "job": {"label": "//examples/podcast_reviews:categorize_reviews_job"},
                "config": {
                    "outputs": [{"str": partition_ref}],
                    "inputs": [
                        {"dep_type": 1, "partition_ref": {"str": f"reviews/date={target_date}"}},
                        {"dep_type": 1, "partition_ref": {"str": "podcasts/all"}}
                    ],
                    "args": [target_category, target_date],
                    "env": {"PARTITION_REF": partition_ref, "TARGET_CATEGORY": target_category, "TARGET_DATE": target_date}
                }
            }
        }
        
        manifest_file = output_dir / "manifest.json"
        with open(manifest_file, 'w') as f:
            json.dump(manifest, f, indent=2)
            
    except Exception as e:
        print(f"Error categorizing reviews: {e}", file=sys.stderr)
        sys.exit(1)

def categorize_reviews_for_category_date(reviews_file: str, podcasts_file: str, target_category: str, output_file: str):
    """Join reviews with podcast categories and filter for target category."""
    
    # Connect to DuckDB for processing
    duckdb_conn = duckdb.connect()
    
    try:
        # Try to install and load parquet extension, but don't fail if it's already installed
        try:
            duckdb_conn.execute("INSTALL parquet")
        except Exception:
            pass  # Extension might already be installed
        
        duckdb_conn.execute("LOAD parquet")
        
        # Query to join reviews with podcasts and filter by category
        query = f"""
        SELECT 
            r.podcast_id,
            r.review_title,
            r.content,
            r.rating,
            r.author_id,
            r.created_at,
            r.review_date,
            p.title as podcast_title,
            p.primary_category,
            p.all_categories,
            '{target_category}' as target_category
        FROM parquet_scan('{reviews_file}') r
        JOIN parquet_scan('{podcasts_file}') p ON r.podcast_id = p.podcast_id
        WHERE p.primary_category = '{target_category}' 
           OR p.all_categories LIKE '%{target_category}%'
        ORDER BY r.created_at
        """
        
        # Execute query and save to parquet
        duckdb_conn.execute(f"COPY ({query}) TO '{output_file}' (FORMAT PARQUET)")
        
        # Get row count for logging
        count_result = duckdb_conn.execute(f"SELECT COUNT(*) FROM ({query})").fetchone()
        
        row_count = count_result[0] if count_result else 0
        print(f"Categorized {row_count} reviews for category '{target_category}'")
        
        if row_count == 0:
            print(f"Warning: No reviews found for category '{target_category}' on date '{reviews_file.split('date=')[1].split('/')[0]}'")
        
    finally:
        duckdb_conn.close()

if __name__ == "__main__":
    main()