Python for Data Analysis Questions

import numpy as np
import pandas as pd
from numba import njit

@njit
def topk_per_group(vals, k):
    n = vals.shape[0]
    if k <= 0 or n == 0:
        return np.empty(0, dtype=vals.dtype)
    # keep an unsorted buffer of up to k largest seen so far
    buf = np.empty(min(k, n), dtype=vals.dtype)
    m = 0
    for i in range(n):
        v = vals[i]
        if m < k:
            buf[m] = v
            m += 1
        else:
            # find min index in buffer
            min_idx = 0
            min_val = buf[0]
            for j in range(1, m):
                if buf[j] < min_val:
                    min_val = buf[j]
                    min_idx = j
            if v > min_val:
                buf[min_idx] = v
    # sort descending before returning
    # simple in-place sort (selection sort for small k)
    for i in range(m):
        max_idx = i
        for j in range(i+1, m):
            if buf[j] > buf[max_idx]:
                max_idx = j
        if max_idx != i:
            tmp = buf[i]; buf[i] = buf[max_idx]; buf[max_idx] = tmp
    return buf[:m]

# pandas integration
def group_topk(df, group_col, value_col, k):
    groups = []
    for name, grp in df.groupby(group_col):
        arr = grp[value_col].to_numpy()  # plain NumPy array
        topk = topk_per_group(arr, k)    # call Numba function
        # store result as rows: one row per returned top value (or pack)
        for rank, v in enumerate(topk, start=1):
            groups.append((name, rank, v))
    return pd.DataFrame(groups, columns=[group_col, "rank", value_col])

# small demo
df = pd.DataFrame({
    "g": np.repeat(np.arange(3), 6),
    "val": np.random.randn(18)
})
print(group_topk(df, "g", "val", k=3))

import numpy as np, pandas as pd, json, hashlib
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
seed = 42
np.random.seed(seed)

df = pd.read_csv("data/raw_v1.csv")            # read exact data version (from DVC/Delta)
# schema validation (e.g., pandera) here

# deterministic shuffle + split
train_idx, val_idx = train_test_split(
    np.arange(len(df)), test_size=0.2, random_state=seed, shuffle=True
)

scaler = StandardScaler().fit(df.loc[train_idx, ["num_col"]])
# serialize artifacts
with open("preproc_meta.json","w") as f:
    json.dump({
      "seed": seed,
      "git_sha": "<git_sha>",
      "dvc_data_hash": "<dvc_hash>",
      "split": {"train": train_idx.tolist()[:10], "val": val_idx.tolist()[:10]},
      "scaler_mean": scaler.mean_.tolist(),
      "requirements_hash": "<requirements_hash>"
    }, f)

# assumes smallish left, large right on disk-backed format
left = left.set_index('user_id')
right = right.set_index('user_id')
res = left.join(right, how='inner')

cats, codes = pd.factorize(df['user_id'])
df['user_id_code'] = codes
# join on user_id_code

def join_chunks(left_iter, right_path):
    for left_chunk in pd.read_csv(left_path, chunksize=1_000_000):
        # load right filtered by keys in chunk (if indexed or pre-partitioned)
        keys = left_chunk['user_id'].unique()
        right_subset = load_right_for_keys(keys)  # from a DB or pre-partitioned files
        out = left_chunk.merge(right_subset, on='user_id', how='inner')
        out.to_parquet('out.parquet', append=True)

import dask.dataframe as dd
L = dd.read_parquet('left/*.parquet')
R = dd.read_parquet('right/*.parquet')
res = L.merge(R, on='user_id', how='inner')  # triggers shuffle
res.to_parquet('joined/')

import duckdb
duckdb.query("CREATE TABLE left AS SELECT * FROM 'left.parquet'")
duckdb.query("CREATE TABLE right AS SELECT * FROM 'right.parquet'")
duckdb.query("CREATE TABLE joined AS SELECT * FROM left JOIN right USING(user_id)")

import numpy as np
from numba import njit, prange

@njit(parallel=True, fastmath=False)
def bayesian_smoothed_rate(successes, trials, alpha, beta):
    """
    successes, trials: 1D numpy arrays (same length)
    alpha, beta: scalar prior params (floats, >0)
    returns: 1D array of posterior means (float64)
    """
    n = successes.shape[0]
    out = np.empty(n, dtype=np.float64)
    prior_mean = alpha / (alpha + beta)
    denom_base = alpha + beta  # float64
    for i in prange(n):
        s = successes[i]
        t = trials[i]
        # guard: NaNs or negative counts -> yield NaN to signal bad input
        if t != t or s != s or s < 0 or t < 0:
            out[i] = np.nan
            continue
        if t == 0:
            out[i] = prior_mean
            continue
        out[i] = (denom_base + s) / (denom_base + t)
    return out