Write the shard worker (`process_shard`)

import asyncio
import os
from pathlib import Path
from typing import Any

import pandas as pd
from render_sdk import Workflows

from sharding import NUM_SHARDS, get_shard_id
from enrichment import enrich_profile

app = Workflows()

DATA_DIR = Path(os.getenv("DATA_DIR", "../sample_data"))


def sanitize(obj: Any) -> Any:
    if pd.isna(obj):
        return None
    if hasattr(obj, "item"):
        return obj.item()
    if isinstance(obj, dict):
        return {k: sanitize(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [sanitize(v) for v in obj]
    return obj


def load_csv(filename: str) -> list[dict]:
    df = pd.read_csv(DATA_DIR / filename)
    return [sanitize(r) for r in df.to_dict("records")]


@app.task
def process_shard(shard_id: int) -> dict:
    sources = {
        name: load_csv(f"{name}.csv")
        for name in ("crm", "billing", "product", "support")
    }

    # Filter each source down to this shard's customers
    indexed: dict[str, dict[str, dict]] = {}
    for source_name, rows in sources.items():
        for row in rows:
            cid = str(row["customer_id"])
            if get_shard_id(cid) != shard_id:
                continue
            indexed.setdefault(cid, {})[source_name] = row

    # Merge and enrich each customer
    profiles = []
    for customer_id, by_source in indexed.items():
        merged = {"customer_id": customer_id}
        for source_name in ("crm", "billing", "product", "support"):
            if source_name in by_source:
                merged.update(by_source[source_name])
        profiles.append(enrich_profile(merged))

    return {
        "shard_id": shard_id,
        "profiles": profiles,
        "count": len(profiles),
    }


if __name__ == "__main__":
    app.start()

import { readFileSync } from "node:fs";
import { join } from "node:path";
import Papa from "papaparse";
import { task } from "@renderinc/sdk/workflows";
import { NUM_SHARDS, getShardId } from "./sharding.js";
import { enrichProfile, type EnrichedProfile } from "./enrichment.js";

const DATA_DIR = process.env.DATA_DIR ?? "../sample_data";

type Row = Record<string, unknown>;

function loadCsv(filename: string): Row[] {
  const content = readFileSync(join(DATA_DIR, filename), "utf-8");
  const result = Papa.parse<Row>(content, {
    header: true,
    skipEmptyLines: true,
    dynamicTyping: true,
  });
  return result.data;
}

export const processShard = task(
  { name: "process_shard" },
  function processShard(shardId: number) {
    const sources = {
      crm:     loadCsv("crm.csv"),
      billing: loadCsv("billing.csv"),
      product: loadCsv("product.csv"),
      support: loadCsv("support.csv"),
    };

    // Filter each source down to this shard's customers
    const indexed = new Map<string, Record<string, Row>>();
    for (const [sourceName, rows] of Object.entries(sources)) {
      for (const row of rows) {
        const cid = String(row.customer_id);
        if (getShardId(cid) !== shardId) continue;
        if (!indexed.has(cid)) indexed.set(cid, {});
        indexed.get(cid)![sourceName] = row;
      }
    }

    // Merge and enrich each customer
    const profiles: EnrichedProfile[] = [];
    for (const [customerId, bySource] of indexed) {
      const merged: Record<string, unknown> = { customer_id: customerId };
      for (const sourceName of ["crm", "billing", "product", "support"]) {
        if (bySource[sourceName]) Object.assign(merged, bySource[sourceName]);
      }
      profiles.push(enrichProfile(merged));
    }

    return {
      shard_id: shardId,
      profiles,
      count: profiles.length,
    };
  },
);