Hi @Jake3,

Your Taylor-linearisation row-percent estimator is well structured. The main performance bottleneck is the Python-level loop over every domain/measure combination, with a full DataFrame copy (df.copy()) happening inside each iteration. Here are concrete changes that should give you a significant speedup without altering the statistical results.

OPTIMIZATION 1: REMOVE THE df.copy() INSIDE THE LOOP

The line df2 = df.copy() creates an entire copy of your DataFrame for every single domain/measure combination. That is the single most expensive operation in your code. Instead, compute the "u" column directly on sliced views and avoid the copy entirely:

# Replace the df2 block with direct computation:
w_arr = df.loc[mask_row, weight_col].values
c_arr = mask_cell[mask_row].astype(float).values
r_arr = np.ones(n_row)
u_arr = (w_arr / W_row) * (c_arr - p_hat * r_arr)

Then in the strata variance loop, group by strata within the mask_row subset rather than copying the entire DataFrame.

OPTIMIZATION 2: USE groupby INSTEAD OF LOOPING OVER COMBINATIONS

Instead of building all_rows via itertools.product and looping, use pandas groupby to iterate only over combinations that actually exist in your data:

group_cols = row_domains + [measure]
grouped = df.loc[df["scope_fg"] == 1].groupby(group_cols, observed=True)

This automatically skips empty combinations and avoids repeated Boolean mask construction for each combo. For the domain totals (W_row), pre-compute them with a separate groupby on just the row_domains:

domain_totals = (
  df.loc[df["scope_fg"] == 1]
    .groupby(row_domains, observed=True)[weight_col]
    .sum()
)

Then look up the domain total for each group with a simple index lookup.

OPTIMIZATION 3: VECTORIZE THE STRATA VARIANCE CALCULATION

The inner loop over strata can be vectorized using groupby + transform:

# Pre-compute per-stratum stats in one pass
strata_groups = sub_df.groupby(strata_col)
n_h_map = strata_groups["u"].transform("count")
mean_u = strata_groups["u"].transform("mean")

# Variance contribution per record
sub_df["u_dev_sq"] = (sub_df["u"] - mean_u) ** 2

strata_stats = sub_df.groupby(strata_col).agg(
  n_h=("u", "count"),
  S2=("u_dev_sq", lambda x: x.sum() / (len(x) - 1) if len(x) > 1 else 0),
  N_h=(fpc_col, "first")
)
strata_stats["f_h"] = np.where(
  np.isfinite(strata_stats["N_h"]),
  strata_stats["n_h"] / strata_stats["N_h"],
  0.0
)
var_p = ((1 - strata_stats["f_h"]) * strata_stats["n_h"] * strata_stats["S2"]).sum()

OPTIMIZATION 4: PRE-COMPUTE MASKS WITH CATEGORICAL COLUMNS

Convert your domain and measure columns to pandas Categorical dtype before processing. This makes groupby and equality comparisons faster, especially with string columns:

for c in row_domains + [measure]:
  df[c] = df[c].astype("category")

OPTIMIZATION 5: AVOID .toPandas() ON LARGE DATASETS

If your Spark DataFrame is very large, the .toPandas() call collects everything to the driver. Consider whether you can filter or aggregate in Spark first. For example, if you only need in-scope records:

# Filter in Spark before collecting
spark_df = spark_df.filter(~col("some_col").isin(exclusion_values))
df = spark_df.toPandas()

You can also enable Apache Arrow for faster Spark-to-pandas conversion:

spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")

This is typically enabled by default on serverless, but worth confirming on other compute types.

PUTTING IT ALL TOGETHER

Here is a refactored version of the main estimation loop that incorporates optimizations 1-3:

# Pre-compute domain-level weighted totals
inscope = df[df["scope_fg"] == 1].copy()
for c in row_domains + [measure]:
  inscope[c] = inscope[c].astype("category")

domain_W = inscope.groupby(row_domains, observed=True)[weight_col].sum()
domain_N = inscope.groupby(row_domains, observed=True)[weight_col].count()

group_cols = row_domains + [measure]
cell_agg = inscope.groupby(group_cols, observed=True).agg(
  W_cell=(weight_col, "sum"),
  n_cell=(weight_col, "count")
).reset_index()

# Merge domain totals
cell_agg = cell_agg.merge(
  domain_W.rename("W_row").reset_index(),
  on=row_domains, how="left"
)
cell_agg = cell_agg.merge(
  domain_N.rename("n_row").reset_index(),
  on=row_domains, how="left"
)
cell_agg["p_hat"] = cell_agg["W_cell"] / cell_agg["W_row"]

# Then loop only for Taylor variance (which needs record-level data)
# but now you skip empty combos and avoid df.copy()

This approach computes the point estimates (p_hat) in one vectorized pass. You only need the record-level loop for the Taylor variance computation, which still needs individual "u" values per stratum.

A NOTE ON NUMERICAL PRECISION

You mentioned that standard errors change slightly when you try to optimize. This is almost always caused by floating-point summation order differences. If you switch from a Python loop to vectorized operations, the order in which values are summed may change slightly. To verify correctness:

1. Compare results on a small test dataset where you can verify by hand.
2. Check that differences are within floating-point tolerance (typically 1e-10 or smaller).
3. If you need exact reproducibility with SAS output, keep the loop-based variance computation but apply the other optimizations (removing df.copy(), pre-computing domain totals).

REFERENCE DOCUMENTATION

- Apache Arrow optimization for pandas conversion: https://docs.databricks.com/en/pandas/pandas-function-apis.html
- Performance tuning on Databricks: https://docs.databricks.com/en/optimizations/index.html
- pandas best practices for performance: https://pandas.pydata.org/docs/user_guide/enhancingperf.html

* This reply used an agent system I built to research and draft this response based on the wide set of documentation I have available and previous memory. I personally review the draft for any obvious issues and for monitoring system reliability and update it when I detect any drift, but there is still a small chance that something is inaccurate, especially if you are experimenting with brand new features.

If this answer resolves your question, could you mark it as "Accept as Solution"? That helps other users quickly find the correct fix.