crdc-demo/src/hooks/useDrawDistribution.js

import { useEffect, useState } from 'react'
import { getDb } from '../utils/duckdbClient.js'

const HF_BASE = 'https://huggingface.co/datasets/civilytics/crdc-school-arrest-rates/resolve/main/parquet'

// Module-level cache: one registered duckdb-wasm file buffer per
// (model, year, state) shard, shared across every component instance and
// district navigated to in this browser session. See Global Constraints —
// in-memory only, no persistence across page loads.
const shardCache = new Map()

function shardKey(model, year, state) {
  return `${model}__${year}__${state}`
}

function ensureShardRegistered(db, model, year, state) {
  const key = shardKey(model, year, state)
  if (!shardCache.has(key)) {
    shardCache.set(
      key,
      (async () => {
        try {
          const url = `${HF_BASE}/model_id=${model}/YEAR=${year}/LEA_STATE=${state}/data_0.parquet`
          const res = await fetch(url)
          if (!res.ok) throw new Error(`Failed to fetch draw shard: HTTP ${res.status}`)
          const buffer = new Uint8Array(await res.arrayBuffer())
          const fileName = `${key}.parquet`
          await db.registerFileBuffer(fileName, buffer)
          return fileName
        } catch (err) {
          // Don't let a transient failure (network blip, HF outage) poison the
          // cache forever — remove the rejected entry so the next caller for
          // this shard gets a fresh attempt instead of the same dead promise.
          shardCache.delete(key)
          throw err
        }
      })(),
    )
  }
  return shardCache.get(key)
}

/**
 * Fetches real posterior draws for one district/model/year from the Hugging
 * Face parquet dataset via duckdb-wasm, converts predicted counts to
 * rate-per-1,000 using each group's stu_enroll (not present in the draws
 * table itself — joined here from data this app already has), and returns
 * them keyed by "RACE_SEX".
 *
 * @param {{leaid: string, state: string, model: string, year: string,
 *   groups: Array<{race: string, sex: string, stuEnroll: number}>}} params
 * @returns {{status: 'loading'|'ready'|'error', drawsByGroup: Record<string, number[]> | null}}
 */
export function useDrawDistribution({ leaid, state, model, year, groups }) {
  const [status, setStatus] = useState('loading')
  const [drawsByGroup, setDrawsByGroup] = useState(null)

  // groups is typically a fresh array literal every render; derive a stable
  // primitive so the effect only re-runs when its actual content changes.
  const groupsSignature = (groups || []).map((g) => `${g.race}:${g.sex}:${g.stuEnroll}`).join(',')

  useEffect(() => {
    if (!leaid || !state || !model || !year || !groups?.length) return
    let cancelled = false
    setStatus('loading')

    async function run() {
      let conn
      try {
        const db = await getDb()
        const fileName = await ensureShardRegistered(db, model, year, state)
        conn = await db.connect()
        const stmt = await conn.prepare(`SELECT RACE, SEX, pred FROM read_parquet('${fileName}') WHERE LEAID = ?`)
        const table = await stmt.query(leaid)
        await stmt.close()
        const rows = table.toArray().map((r) => r.toJSON())

        const enrollByGroup = {}
        for (const g of groups) enrollByGroup[`${g.race}_${g.sex}`] = g.stuEnroll || 0

        const byGroup = {}
        for (const row of rows) {
          const key = `${row.RACE}_${row.SEX}`
          const enroll = enrollByGroup[key]
          if (!enroll) continue
          const rate = (Number(row.pred) / enroll) * 1000
          ;(byGroup[key] ??= []).push(rate)
        }

        // A successful query with zero matching rows (this district isn't in
        // the draws shard, or none of its rows matched a known group) is not
        // "ready" — there's no real data to show, so treat it like a failure
        // and let the caller fall back, rather than silently claiming real
        // draws while every group actually uses the fitted approximation.
        const hasDraws = Object.values(byGroup).some((draws) => draws.length > 0)

        if (!cancelled) {
          if (hasDraws) {
            setDrawsByGroup(byGroup)
            setStatus('ready')
          } else {
            console.warn('useDrawDistribution: query succeeded but returned no usable draws for', { leaid, model, year, state })
            setDrawsByGroup(null)
            setStatus('error')
          }
        }
      } catch (err) {
        console.error('useDrawDistribution failed:', err)
        if (!cancelled) setStatus('error')
      } finally {
        if (conn) await conn.close()
      }
    }

    run()
    return () => {
      cancelled = true
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [leaid, state, model, year, groupsSignature])

  return { status, drawsByGroup }
}