Python Data Manipulation with Pandas Questions

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

# sample data
df = pd.DataFrame({
    "grp": np.repeat(np.arange(1000), 100),
    "value": np.random.randn(1000*100),
    "w": np.random.rand(1000*100)
})

# prepare arrays: group ids (sorted) and contiguous slices
order = np.argsort(df["grp"].values)
grp_sorted = df["grp"].values[order]
val_sorted = df["value"].values[order]
w_sorted = df["w"].values[order]

# compute group boundaries
@njit
def group_boundaries(grp):
    n = grp.size
    boundaries = []
    if n == 0:
        return np.empty((0,2), dtype=np.int64)
    start = 0
    for i in range(1, n):
        if grp[i] != grp[i-1]:
            boundaries.append((start, i))  # [start, end)
            start = i
    boundaries.append((start, n))
    out = np.empty((len(boundaries), 2), dtype=np.int64)
    for i in range(len(boundaries)):
        out[i, 0] = boundaries[i][0]
        out[i, 1] = boundaries[i][1]
    return out

@njit
def weighted_trimmed_mean(vals, weights, start, end, trim_frac=0.05):
    # copy slice into local arrays
    m = end - start
    if m == 0:
        return 0.0
    vs = np.empty(m)
    ws = np.empty(m)
    for i in range(m):
        vs[i] = vals[start + i]
        ws[i] = weights[start + i]
    # sort by value (simple argsort)
    idx = np.argsort(vs)
    cum_w = 0.0
    total_w = 0.0
    for i in range(m):
        total_w += ws[i]
    # trim by cumulative weight at both ends
    low_cut = trim_frac * total_w
    high_cut = (1 - trim_frac) * total_w
    s = 0.0
    weighted_sum = 0.0
    kept_w = 0.0
    for j in range(m):
        i = idx[j]
        s += ws[i]
        if s <= low_cut:
            continue
        if s > high_cut:
            break
        weighted_sum += vs[i] * ws[i]
        kept_w += ws[i]
    if kept_w == 0:
        return 0.0
    return weighted_sum / kept_w

@njit
def compute_per_group(vals, weights, boundaries):
    G = boundaries.shape[0]
    out = np.empty(G)
    for gi in range(G):
        s = boundaries[gi, 0]
        e = boundaries[gi, 1]
        out[gi] = weighted_trimmed_mean(vals, weights, s, e, 0.05)
    return out

# run
bounds = group_boundaries(grp_sorted)
res = compute_per_group(val_sorted, w_sorted, bounds)

# map back to group ids
group_ids = np.empty(bounds.shape[0], dtype=grp_sorted.dtype)
for i in range(bounds.shape
Approach summary:
- Replace Python-level per-group loops with vectorized NumPy/numba to remove Python/GIL overhead. Extract minimal numeric arrays from pandas, compute contiguous group slices, run JIT-compiled functions operating on raw arrays, then map results back to a DataFrame. This is low-effort and gives large speedups vs groupby.apply.

Example: compute a weighted trimmed mean per group (drop bottom/top 5% by cumulative weight, then weighted mean). We sort once by group to get contiguous slices, then run numba on arrays with no pandas inside.

Why this helps:
- Numba operates on raw NumPy arrays in native code — eliminates per-group Python callbacks.
- Sorting once and computing contiguous boundaries minimizes overhead.
- Minimal pandas interaction: only array extraction and final mapping.

Complexity and trade-offs:
- Time: O(N log N) due to initial sort; per-group work depends on group sizes (sorting within group is O(k log k)). If groups are small, overall is fast.
- Memory: uses extra arrays (slices, boundaries) but modest compared to full DataFrame.
- Alternatives: express fully vectorized using pandas/NumPy (if possible), use Cython for more control, or leverage C-accelerated libraries (numexpr, dask, polars) or groupby-aggregation in C (polars or SQL engines).

Edge cases & tips:
- If groups are already contiguous, skip global sort.
- For very large datasets, chunk processing or out-of-core frameworks (Dask, Polars) may be needed.
- Validate numeric stability for zero weights and empty groups.
- Benchmark vs original apply to ensure speedup; warm up numba compile time before measuring.

import pyarrow as pa
import pyarrow.parquet as pq
import pyarrow.dataset as ds
import pandas as pd
import boto3
import hashlib
from io import BytesIO
import s3fs
s3 = s3fs.S3FileSystem()

DATA_URI = "s3://my-bucket/path/to/partitioned/"
INTERMEDIATE_PREFIX = "s3://my-bucket/tmp/dedupe_buckets/"
FINAL_PREFIX = "s3://my-bucket/final/deduped/"
N_BUCKETS = 256  # tune based on key cardinality & memory

def bucket_of(key):
    h = hashlib.md5(str(key).encode()).digest()
    return h[0] % N_BUCKETS

def map_phase(filter_expr=None):
    dataset = ds.dataset(DATA_URI, format="parquet", filesystem=s3)
    scanner = dataset.scanner(filter=filter_expr, batch_size=64_000)  # controlled memory
    for batch in scanner.to_batches():
        table = pa.Table.from_batches([batch])
        df = table.to_pandas()  # batch sized to fit memory
        # compute "best" per key in this batch (e.g., latest ts)
        df.sort_values(['key', 'timestamp'], ascending=[True, False], inplace=True)
        deduped = df.drop_duplicates(subset=['key'], keep='first')
        # write per-bucket parquet to S3 (append style: one file per batch per bucket)
        for b, sub in deduped.groupby(deduped['key'].apply(bucket_of)):
            buf = BytesIO()
            pq.write_table(pa.Table.from_pandas(sub), buf)
            path = f"{INTERMEDIATE_PREFIX}bucket={b}/part-{hashlib.md5(buf.getvalue()).hexdigest()}.parquet"
            with s3.open(path, "wb") as f:
                f.write(buf.getvalue())

def reduce_phase():
    for b in range(N_BUCKETS):
        bucket_path = f"{INTERMEDIATE_PREFIX}bucket={b}/"
        try:
            files = s3.glob(bucket_path + "*.parquet")
        except FileNotFoundError:
            continue
        # read all intermediate files for this bucket into a single dataframe (should fit in memory)
        parts = []
        for fpath in files:
            with s3.open(fpath, "rb") as f:
                tbl = pq.read_table(f)
            parts.append(tbl.to_pandas())
        if not parts:
            continue
        df = pd.concat(parts, ignore_index=True)
        # final dedupe per key (same rule)
        df.sort_values(['key', 'timestamp'], ascending=[True, False], inplace=True)
        final = df.drop_duplicates(subset=['key'], keep='first')
        out_tbl = pa.Table.from_pandas(final)
        out_path = f"{FINAL_PREFIX}bucket={b}.parquet"
        with s3.open(out_path, "wb") as f:
            pq.write_table(out_tbl, f)

import pandas as pd

# load small dimension entirely
dim = pd.read_csv("dim.csv")  # ~10k rows
# reduce memory: set appropriate types
dim['dim_key'] = dim['dim_key'].astype('int32')
for col in ['country', 'segment']:
    dim[col] = dim[col].astype('category')

# If join is one-to-one or many-to-one and you only need a few dim columns, use map (faster, lower mem)
map_series = dim.set_index('dim_key')['segment']  # Series for mapping

# Process fact in chunks, add mapped column, then optionally merge remaining columns
chunksize = 2_000_000
reader = pd.read_csv("fact.csv", chunksize=chunksize, dtype={'fact_key':'int32'})

out_path = "fact_enriched.csv"
first = True
for chunk in reader:
    # reduce memory
    chunk['fact_key'] = chunk['fact_key'].astype('int32')
    # prefer map if only adding single column (many-to-one)
    chunk['segment'] = chunk['fact_key'].map(map_series).astype('category')
    # if you need multiple columns or complex joins, do an indexed merge on small dim:
    # chunk = chunk.merge(dim, left_on='fact_key', right_on='dim_key', how='left')
    # write out
    chunk.to_csv(out_path, mode='w' if first else 'a', index=False, header=first)
    first = False

import pandas as pd

# sample data
users = pd.DataFrame({
    "user_id": [1,2,3,4],
    "name": ["Alice","Bob","Cara","Dan"],
    "signup_date": pd.to_datetime(["2021-01-01","2021-02-10","2021-03-05","2021-04-01"])
})
events = pd.DataFrame({
    "event_id": [10,11,12],
    "user_id": [1,2,1],
    "event_type": ["click","view","purchase"],
    "ts": pd.to_datetime(["2021-05-01","2021-05-02","2021-05-03"])
})

# Inner join (users with matching events)
inner = users.merge(events, on="user_id", how="inner", suffixes=("_user","_event"), validate="m:*>0")
# Left join (all users, events where present)
left = users.merge(events, on="user_id", how="left", suffixes=("_user","_event"), indicator=True, validate="1:m")

# Anti-join: users with NO events (left-only rows)
anti = users.merge(events[["user_id"]].drop_duplicates(), on="user_id", how="left", indicator=True)
anti_no_events = anti[anti["_merge"] == "left_only"].drop(columns=["_merge"])

# Example inspect
print("Inner:\n", inner)
print("Left with indicator:\n", left)
print("Users without events (anti-join):\n", anti_no_events)

import pandas as pd

def dedupe_latest(df: pd.DataFrame, keys: list, ts_col: str = "updated_at") -> pd.DataFrame:
    """
    Return rows deduplicated by keys, keeping the row with the max ts_col.
    Works for a full DataFrame.
    """
    # Ensure timestamp is datetime for correct comparisons
    df = df.copy()
    if not pd.api.types.is_datetime64_any_dtype(df[ts_col]):
        df[ts_col] = pd.to_datetime(df[ts_col], errors="coerce")
    # Use idxmax to pick the index of the latest row per key
    idx = df.groupby(keys)[ts_col].idxmax()
    return df.loc[idx].reset_index(drop=True)

def update_state_with_chunk(state: pd.DataFrame, chunk: pd.DataFrame, keys: list, ts_col: str = "updated_at") -> pd.DataFrame:
    """
    Given current state (latest rows per key) and a new chunk, return updated state.
    Keep only the latest row per key across state + chunk.
    """
    if state is None or state.empty:
        return dedupe_latest(chunk, keys, ts_col)
    combined = pd.concat([state, chunk], ignore_index=True, sort=False)
    return dedupe_latest(combined, keys, ts_col)