#!/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: phrase_stats_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 'phrase_stats/category=comedy/date=2020-01-01' into components.""" match = re.match(r'phrase_stats/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) configs = [] # Process each partition reference for partition_ref in args: 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: phrase models and categorized reviews for the category and date phrase_models_ref = f"phrase_models/category={category}/date={date_str}" categorized_reviews_ref = f"categorized_reviews/category={category}/date={date_str}" configs.append({ "outputs": [{"str": partition_ref}], "inputs": [ {"dep_type": 1, "partition_ref": {"str": phrase_models_ref}}, {"dep_type": 1, "partition_ref": {"str": categorized_reviews_ref}} ], "args": [category, date_str], "env": { "PARTITION_REF": partition_ref, "TARGET_CATEGORY": category, "TARGET_DATE": date_str } }) config = {"configs": configs} print(json.dumps(config)) def handle_exec(args): import time, random, os; time.sleep(float(os.getenv('EXEC_SLEEP', '0')) * random.random()) 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'phrase_stats/category={target_category}/date={target_date}') # Input paths phrase_models_file = f"/tmp/databuild_test/examples/podcast_reviews/phrase_models/category={target_category}/date={target_date}/phrase_models.parquet" categorized_reviews_file = f"/tmp/databuild_test/examples/podcast_reviews/categorized_reviews/category={target_category}/date={target_date}/categorized_reviews.parquet" # Check input files exist if not os.path.exists(phrase_models_file): print(f"Phrase models file not found: {phrase_models_file}", file=sys.stderr) sys.exit(1) if not os.path.exists(categorized_reviews_file): print(f"Categorized reviews file not found: {categorized_reviews_file}", file=sys.stderr) sys.exit(1) # Output path output_dir = Path(f"/tmp/databuild_test/examples/podcast_reviews/phrase_stats/category={target_category}/date={target_date}") output_dir.mkdir(parents=True, exist_ok=True) output_file = output_dir / "phrase_stats.parquet" try: # Calculate phrase statistics per podcast calculate_phrase_stats_for_category_date( phrase_models_file, categorized_reviews_file, target_category, target_date, str(output_file) ) print(f"Successfully calculated phrase stats for category {target_category} on {target_date}") print(f"Output written to: {output_file}") # Create manifest manifest = { "outputs": [{"str": partition_ref}], "inputs": [ {"str": f"phrase_models/category={target_category}/date={target_date}"}, {"str": f"categorized_reviews/category={target_category}/date={target_date}"} ], "start_time": datetime.now().isoformat(), "end_time": datetime.now().isoformat(), "task": { "job": {"label": "//examples/podcast_reviews:phrase_stats_job"}, "config": { "outputs": [{"str": partition_ref}], "inputs": [ {"dep_type": 1, "partition_ref": {"str": f"phrase_models/category={target_category}/date={target_date}"}}, {"dep_type": 1, "partition_ref": {"str": f"categorized_reviews/category={target_category}/date={target_date}"}} ], "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 calculating phrase stats: {e}", file=sys.stderr) sys.exit(1) def calculate_phrase_stats_for_category_date( phrase_models_file: str, categorized_reviews_file: str, target_category: str, target_date: str, output_file: str ): """Calculate phrase statistics per podcast by joining phrase models with reviews.""" # 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") # Check if we have phrase models phrase_count = duckdb_conn.execute(f"SELECT COUNT(*) FROM parquet_scan('{phrase_models_file}')").fetchone()[0] if phrase_count == 0: print(f"No phrase models found, creating empty phrase stats") create_empty_phrase_stats(target_category, target_date, output_file, duckdb_conn) return # Query to calculate phrase statistics per podcast query = f""" WITH phrase_matches AS ( SELECT r.podcast_id, r.podcast_title, r.rating, r.content, p.ngram, p.score as phrase_score FROM parquet_scan('{categorized_reviews_file}') r JOIN parquet_scan('{phrase_models_file}') p ON lower(r.content) LIKE '%' || lower(p.ngram) || '%' WHERE r.content IS NOT NULL AND r.content != '' AND p.ngram IS NOT NULL AND p.ngram != '' ), podcast_phrase_stats AS ( SELECT '{target_date}' as date, '{target_category}' as category, podcast_id, podcast_title, ngram, COUNT(*) as count, AVG(rating::FLOAT) as avg_rating, MIN(rating) as min_rating, MAX(rating) as max_rating, AVG(phrase_score) as avg_phrase_score, COUNT(*) * AVG(phrase_score) * AVG(rating::FLOAT) / 5.0 as weighted_score FROM phrase_matches GROUP BY podcast_id, podcast_title, ngram HAVING COUNT(*) >= 2 -- Only include phrases that appear at least twice per podcast ) SELECT date, category, podcast_id, podcast_title, ngram, count, avg_rating, min_rating, max_rating, avg_phrase_score, weighted_score FROM podcast_phrase_stats ORDER BY weighted_score DESC, count DESC """ # 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"Calculated phrase statistics for {row_count} podcast-phrase combinations") if row_count == 0: print(f"Warning: No phrase matches found for category '{target_category}' on date '{target_date}'") create_empty_phrase_stats(target_category, target_date, output_file, duckdb_conn) finally: duckdb_conn.close() def create_empty_phrase_stats(category: str, date: str, output_file: str, duckdb_conn): """Create empty phrase stats parquet file with correct schema.""" duckdb_conn.execute("DROP TABLE IF EXISTS empty_phrase_stats") duckdb_conn.execute(""" CREATE TABLE empty_phrase_stats ( date VARCHAR, category VARCHAR, podcast_id VARCHAR, podcast_title VARCHAR, ngram VARCHAR, count BIGINT, avg_rating DOUBLE, min_rating INTEGER, max_rating INTEGER, avg_phrase_score DOUBLE, weighted_score DOUBLE ) """) duckdb_conn.execute(f"COPY empty_phrase_stats TO '{output_file}' (FORMAT PARQUET)") print("Created empty phrase stats file") if __name__ == "__main__": main()