data_loader.py

# MovieMindAI - IMDb Data Loader
# Efficiently loads large TSV files using chunked reading

import pandas as pd
from pathlib import Path
from typing import Optional, Generator
import config

# TSV reading configuration
TSV_PARAMS = {
    "sep": "\t",
    "na_values": "\\N",
    "quoting": 3,  # QUOTE_NONE - IMDb files have no quoting
    "low_memory": False,
}


def load_title_basics(
    filter_types: Optional[list] = None,
    min_year: Optional[int] = None
) -> pd.DataFrame:
    """
    Load title.basics.tsv with optional filtering.
    
    Columns: tconst, titleType, primaryTitle, originalTitle, 
             isAdult, startYear, endYear, runtimeMinutes, genres
    """
    print("Loading title.basics.tsv...")
    
    df = pd.read_csv(
        config.TITLE_BASICS_PATH,
        dtype={
            "tconst": str,
            "titleType": str,
            "primaryTitle": str,
            "originalTitle": str,
            "isAdult": str,
            "startYear": str,
            "endYear": str,
            "runtimeMinutes": str,
            "genres": str,
        },
        **TSV_PARAMS
    )
    
    # Convert numeric columns
    df["startYear"] = pd.to_numeric(df["startYear"], errors="coerce")
    df["endYear"] = pd.to_numeric(df["endYear"], errors="coerce")
    df["runtimeMinutes"] = pd.to_numeric(df["runtimeMinutes"], errors="coerce")
    df["isAdult"] = df["isAdult"] == "1"
    
    # Filter out adult content
    df = df[~df["isAdult"]]
    
    # Apply title type filter
    if filter_types:
        df = df[df["titleType"].isin(filter_types)]
    
    # Apply year filter
    if min_year:
        df = df[df["startYear"] >= min_year]
    
    print(f"  Loaded {len(df):,} titles")
    return df


def load_title_ratings() -> pd.DataFrame:
    """
    Load title.ratings.tsv.
    
    Columns: tconst, averageRating, numVotes
    """
    print("Loading title.ratings.tsv...")
    
    df = pd.read_csv(
        config.TITLE_RATINGS_PATH,
        dtype={
            "tconst": str,
            "averageRating": float,
            "numVotes": int,
        },
        **TSV_PARAMS
    )
    
    print(f"  Loaded {len(df):,} ratings")
    return df


def load_title_crew() -> pd.DataFrame:
    """
    Load title.crew.tsv.
    
    Columns: tconst, directors, writers
    (directors and writers are comma-separated nconst values)
    """
    print("Loading title.crew.tsv...")
    
    df = pd.read_csv(
        config.TITLE_CREW_PATH,
        dtype={
            "tconst": str,
            "directors": str,
            "writers": str,
        },
        **TSV_PARAMS
    )
    
    print(f"  Loaded {len(df):,} crew records")
    return df


def load_name_basics() -> pd.DataFrame:
    """
    Load name.basics.tsv.
    
    Columns: nconst, primaryName, birthYear, deathYear, 
             primaryProfession, knownForTitles
    """
    print("Loading name.basics.tsv...")
    
    df = pd.read_csv(
        config.NAME_BASICS_PATH,
        dtype={
            "nconst": str,
            "primaryName": str,
            "birthYear": str,
            "deathYear": str,
            "primaryProfession": str,
            "knownForTitles": str,
        },
        **TSV_PARAMS
    )
    
    # We only need nconst and primaryName for lookups
    df = df[["nconst", "primaryName"]]
    
    print(f"  Loaded {len(df):,} names")
    return df


def load_title_principals_filtered(title_ids: set) -> pd.DataFrame:
    """
    Load title.principals.tsv, filtering to only specified titles.
    Uses chunked reading due to file size.
    
    Columns: tconst, ordering, nconst, category, job, characters
    """
    print("Loading title.principals.tsv (chunked)...")
    
    chunks = []
    chunk_size = 1_000_000
    
    for chunk in pd.read_csv(
        config.TITLE_PRINCIPALS_PATH,
        dtype={
            "tconst": str,
            "ordering": int,
            "nconst": str,
            "category": str,
            "job": str,
            "characters": str,
        },
        chunksize=chunk_size,
        **TSV_PARAMS
    ):
        # Filter to our target titles
        filtered = chunk[chunk["tconst"].isin(title_ids)]
        if len(filtered) > 0:
            chunks.append(filtered)
    
    df = pd.concat(chunks, ignore_index=True) if chunks else pd.DataFrame()
    
    print(f"  Loaded {len(df):,} principal records")
    return df


def load_title_akas_filtered(title_ids: set) -> pd.DataFrame:
    """
    Load title.akas.tsv, filtering to only specified titles.
    Uses chunked reading due to file size.
    
    Columns: titleId, ordering, title, region, language, 
             types, attributes, isOriginalTitle
    """
    print("Loading title.akas.tsv (chunked)...")
    
    chunks = []
    chunk_size = 1_000_000
    
    for chunk in pd.read_csv(
        config.TITLE_AKAS_PATH,
        dtype={
            "titleId": str,
            "ordering": int,
            "title": str,
            "region": str,
            "language": str,
            "types": str,
            "attributes": str,
            "isOriginalTitle": str,
        },
        chunksize=chunk_size,
        **TSV_PARAMS
    ):
        # Filter to our target titles
        filtered = chunk[chunk["titleId"].isin(title_ids)]
        if len(filtered) > 0:
            chunks.append(filtered)
    
    df = pd.concat(chunks, ignore_index=True) if chunks else pd.DataFrame()
    
    # Convert isOriginalTitle to boolean
    df["isOriginalTitle"] = df["isOriginalTitle"] == "1"
    
    print(f"  Loaded {len(df):,} aka records")
    return df


if __name__ == "__main__":
    # Test loading
    print("Testing data loader...")
    
    # Load basics with filters
    basics = load_title_basics(
        filter_types=config.VALID_TITLE_TYPES,
        min_year=config.MIN_YEAR
    )
    print(f"\nFiltered basics: {len(basics):,} titles")
    print(f"Title types: {basics['titleType'].value_counts().to_dict()}")