I have the following code in databricks under serverless and i want to know how to improve it to make it more efficient and run faster without having the results change (standard errors change slightly when i try to improve it):

# Databricks Serverless - Taylor Row Percent

# ============================================================

# Runs on Databricks Serverless (pandas on Spark / pandas UDF pattern).

# The core estimation logic is kept in pandas because Taylor

# linearisation and survey-variance math has no native Spark equivalent.

# Usage

# -----

# Input : pass a pandas DataFrame directly, OR a Spark DataFrame

# (it will be converted automatically via .toPandas()).

# Output : pandas DataFrame (call display() or wrap in spark.createDataFrame()

# if you need a Spark / Delta result).

# ============================================================

from __future__ import annotations

import numpy as np

import pandas as pd

from itertools import product

from typing import Dict, List, Optional, Union

# Databricks / PySpark imports (available on every Serverless cluster)

from pyspark.sql import DataFrame as SparkDataFrame

# ------------------------------------------------------------------

# Helper: accept either a Spark or pandas DataFrame

# ------------------------------------------------------------------

def _to_pandas(df: Union[pd.DataFrame, SparkDataFrame]) -> pd.DataFrame:

"""Convert Spark DataFrame → pandas if needed."""

if isinstance(df, SparkDataFrame):

return df.toPandas()

return df.copy()

# ------------------------------------------------------------------

# Main function

# ------------------------------------------------------------------

def taylor_row_percent(

df: Union[pd.DataFrame, SparkDataFrame],

row_domains: List[str],

measure: str,

exclusions: Dict[str, List],

weight_col: Optional[str] = None,

strata_col: Optional[str] = None,

fpc: Optional[Union[Dict[object, float], str]] = None,

p_01_adjust: Optional[float] = 1.0,

zcritical: float = 1.96,

output_spark: bool = False, # set True to get a Spark DataFrame back

output_delta_table: Optional[str] = None, # e.g. "catalog.schema.table"

) -> Union[pd.DataFrame, SparkDataFrame]:

"""

Taylor-linearisation row-percent estimator.

Parameters

----------

df : pandas or Spark DataFrame

row_domains : columns that define the row grouping

measure : column whose value distribution is estimated

exclusions : {col: [values_to_exclude]} – NaN supported

weight_col : survey weight column (None → unweighted)

strata_col : stratification column (None → single stratum)

fpc : finite-population correction – dict {stratum: N} or column name

p_01_adjust : Agresti-Coull scaling factor for p=0/1 cells

zcritical : z-value for confidence interval (default 1.96)

output_spark : if True, return a Spark DataFrame instead of pandas

output_delta_table : if set, write result to this Delta table and return it

"""

# ------------------------------------------------------------------

# 0. Normalise input

# ------------------------------------------------------------------

df = _to_pandas(df)

# ------------------------------------------------------------------

# 1. Weight handling

# ------------------------------------------------------------------

if weight_col is None:

df["_w"] = 1.0

weight_col = "_w"

# ------------------------------------------------------------------

# 2. Strata handling

# ------------------------------------------------------------------

if strata_col is None:

df["_strata"] = "all"

strata_col = "_strata"

fpc = None

# ------------------------------------------------------------------

# 3. FPC handling

# ------------------------------------------------------------------

if fpc is None:

df["_N_h"] = np.inf

fpc_col = "_N_h"

elif isinstance(fpc, dict😞

df["_N_h"] = df[strata_col].map(fpc)

fpc_col = "_N_h"

else:

fpc_col = fpc

# ------------------------------------------------------------------

# 4. Scope flag (mirrors SAS exclusion logic)

# ------------------------------------------------------------------

df["scope_fg"] = 1

for var, excl_vals in exclusions.items(...

New board topics in Databricks Community

Message History

Delivery option(s)