mostlyai-engine 1.6.0__tar.gz → 1.7.0__tar.gz

{mostlyai_engine-1.6.0 → mostlyai_engine-1.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.6.0
+Version: 1.7.0
 Summary: Synthetic Data Engine
 Project-URL: homepage, https://github.com/mostly-ai/mostlyai-engine
 Project-URL: repository, https://github.com/mostly-ai/mostlyai-engine

{mostlyai_engine-1.6.0 → mostlyai_engine-1.7.0}/mostlyai/engine/__init__.py

@@ -22,7 +22,7 @@ from mostlyai.engine.splitting import split
 from mostlyai.engine.training import train
 
 __all__ = ["split", "analyze", "encode", "train", "generate", "init_logging", "set_random_state"]
-__version__ = "1.6.0"
+__version__ = "1.7.0"
 
 # suppress specific warning related to os.fork() in multi-threaded processes
 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")

{mostlyai_engine-1.6.0 → mostlyai_engine-1.7.0}/mostlyai/engine/_common.py

@@ -618,10 +618,23 @@ class FixedSizeSampleBuffer:
         self.n_clears += 1
 
 
-def _get_log_histogram_edges(idx: int, bins: int = 64) -> tuple[float, float]:
+def _get_log_histogram_bin_bounds(idx: int, bins: int = 64) -> tuple[float, float]:
     """
+    Compute the lower and upper boundaries for a logarithmically-spaced histogram bin.
+
+    Creates symmetric logarithmic bins around zero that efficiently represent values
+    across many orders of magnitude. With bins=64, creates 128 total bins covering
+    negative powers of 2, the range [-1, 1], and positive powers of 2.
+
     Modified from OpenDP's SmartNoise SDK (MIT License)
     Source: https://github.com/opendp/smartnoise-sdk/blob/main/sql/snsql/sql/_mechanisms/approx_bounds.py
+
+    Args:
+        idx: The bin index (0 to bins*2-1)
+        bins: Number of bins per side (default 64, creating 128 total bins)
+
+    Returns:
+        Tuple of (lower_edge, upper_edge) for the bin
     """
     if idx == bins:
         return (0.0, 1.0)
@@ -635,50 +648,59 @@ def _get_log_histogram_edges(idx: int, bins: int = 64) -> tuple[float, float]:
 
 def compute_log_histogram(values: np.ndarray, bins: int = 64) -> list[int]:
     """
+    Compute a histogram using logarithmically-spaced bins for efficient distribution analysis.
+
+    This creates a histogram that can efficiently represent values spanning many orders of
+    magnitude (e.g., 0.001 to 1,000,000) using relatively few bins. The bins are symmetric
+    around zero with exponentially increasing widths away from zero.
+
     Modified from OpenDP's SmartNoise SDK (MIT License)
     Source: https://github.com/opendp/smartnoise-sdk/blob/main/sql/snsql/sql/_mechanisms/approx_bounds.py
-    """
-    hist = [0.0] * bins * 2
 
+    Args:
+        values: Array of numeric values to histogram
+        bins: Number of bins per side (default 64, creating 128 total bins)
+
+    Returns:
+        List of counts for each bin. Invalid values (NaN, inf) are filtered out.
+    """
+    # filter out invalid values
     values = np.array(values, dtype=np.float64)
-    values = values[values != np.inf]
-    values = values[values != -np.inf]
-    values = values[~np.isnan(values)]
-    edge_list = [_get_log_histogram_edges(idx) for idx in range(len(hist))]
-    min_val = min([lower for lower, _ in edge_list])
-    max_val = max([upper for _, upper in edge_list]) - 1
+    values = values[~np.isinf(values) & ~np.isnan(values)]
+
+    # generate all bin edges efficiently
+    edge_list = [_get_log_histogram_bin_bounds(idx, bins) for idx in range(bins * 2)]
+    bin_edges = np.array([lower for lower, _ in edge_list] + [edge_list[-1][1]])
+
+    # clip values to be within the bin edges to ensure all values are counted
+    min_val = bin_edges[0]
+    max_val = bin_edges[-1]
     values = np.clip(values, min_val, max_val)
 
-    # compute histograms
-    for v in values:
-        bin = None
-        for idx, (lower, upper) in enumerate(edge_list):
-            if lower <= v < upper:
-                bin = idx
-                break
-        if bin is None:
-            bin = idx
-        hist[bin] += 1
+    # use numpy's histogram for efficient binning (O(n log bins) vs O(n * bins))
+    hist, _ = np.histogram(values, bins=bin_edges)
 
-        # for testing
-        lower, upper = _get_log_histogram_edges(bin)
-    return hist
+    return hist.tolist()
 
 
 def dp_approx_bounds(hist: list[int], epsilon: float) -> tuple[float | None, float | None]:
     """
+    Estimate the minimum and maximum values using a differentially private histogram.
+
+    Uses Laplace noise on histogram bin counts, then finds the lowest and highest bins
+    that exceed a threshold (based on failure probability). Returns None if insufficient
+    data or privacy budget makes reliable estimation impossible.
+
+    Reference: https://desfontain.es/thesis/Usability.html#usability-u-ding-
     Modified from OpenDP's SmartNoise SDK (MIT License)
     Source: https://github.com/opendp/smartnoise-sdk/blob/main/sql/snsql/sql/_mechanisms/approx_bounds.py
 
-    Estimate the minimium and maximum values of a list of values.
-    from: https://desfontain.es/thesis/Usability.html#usability-u-ding-
-
     Args:
-        hist (list[int]): A list of log histogram counts.
-        epsilon (float): The privacy budget to spend estimating the bounds.
+        hist: A list of log histogram counts (typically from compute_log_histogram).
+        epsilon: The privacy budget to spend estimating the bounds.
 
     Returns:
-        tuple[float | None, float | None]: A tuple of the estimated minimum and maximum values.
+        Tuple of (min, max) estimates, or (None, None) if bounds cannot be reliably estimated.
     """
 
     n_bins = len(hist)
@@ -700,8 +722,8 @@ def dp_approx_bounds(hist: list[int], epsilon: float) -> tuple[float | None, flo
         return (None, None)
 
     lower_bin, upper_bin = min(exceeds), max(exceeds)
-    lower, _ = _get_log_histogram_edges(lower_bin)
-    _, upper = _get_log_histogram_edges(upper_bin)
+    lower, _ = _get_log_histogram_bin_bounds(lower_bin)
+    _, upper = _get_log_histogram_bin_bounds(upper_bin)
     return (float(lower), float(upper))
 
 
@@ -709,18 +731,23 @@ def _dp_bounded_quantiles(
     values: np.ndarray, quantiles: list[float], epsilon: float, lower: float, upper: float
 ) -> list[float]:
     """
-    Estimate the quantile.
-    from: http://cs-people.bu.edu/ads22/pubs/2011/stoc194-smith.pdf
+    Estimate quantiles with differential privacy using the Smith (2011) smooth sensitivity method.
+
+    Assumes values are bounded within [lower, upper]. Uses exponential mechanism to sample
+    quantile estimates with noise proportional to local sensitivity. Privacy budget is split
+    evenly across all requested quantiles. Results are post-processed to ensure monotonicity.
+
+    Reference: http://cs-people.bu.edu/ads22/pubs/2011/stoc194-smith.pdf
 
     Args:
-        values (np.ndarray): A 1D array of numeric values.
-        quantiles (list[float]): List of probabilities of the quantiles to estimate.
-        epsilon (float): Privacy budget.
-        lower (float): A bounding parameter. The quantile will be estimated only for values greater than or equal to this bound.
-        upper (float): A bounding parameter. The quantile will be estimated only for values less than or equal to this bound.
+        values: A 1D array of numeric values.
+        quantiles: List of quantile probabilities to estimate (e.g., [0.05, 0.5, 0.95]).
+        epsilon: Privacy budget (split evenly across quantiles).
+        lower: Lower bound for clipping values.
+        upper: Upper bound for clipping values.
 
     Returns:
-        list[float]: The estimated quantile.
+        List of differentially private quantile estimates (monotonically ordered).
     """
 
     _LOG.info(f"compute DP bounded quantiles within [{lower}, {upper}]")
@@ -747,89 +774,23 @@ def _dp_bounded_quantiles(
     return results
 
 
-# NOTE: the unbounded method is not used in the current implementation
-# def _dp_unbounded_quantiles(
-#     values: np.ndarray, quantiles: list[float], epsilon: float, beta: float = 1.01
-# ) -> tuple[list[float], float]:
-#     """
-#     Fully unbounded differentially private quantile estimation using two AboveThreshold calls
-#     with Exponential noise (one-sided Laplace).
-
-#     Implements Algorithm 4 from Durfee (2023):
-#       1) AboveThreshold on positives: T1 = q*n, f_i = |{x_j + 1 < beta^i}|
-#       2) AboveThreshold on negatives: T2 = (1-q)*n, f_i = |{x_j - 1 > -beta^i}|
-#       3) If first halts at k>0: return  beta^k - 1
-#       4) If second halts at k>0: return -beta^k + 1
-#       5) Otherwise return 0
-
-#     Args:
-#         values (np.ndarray): A 1D array of numeric values.
-#         quantiles (list[float]): List of probabilities of the quantiles to estimate.
-#         epsilon (float): Privacy budget.
-#         beta (float): Multiplicative step size (default 1.01). Section 6.4 from Durfee (2023) suggests the range [1.01, 1.001] and 1.01 for general use, especially for more significant
-#         decreases in epsilon or in the data size.
-
-#     Returns:
-#         list[float]: Differentially private estimates of the quantiles.
-#     """
-
-#     def above_threshold(
-#         values: np.ndarray, q: float, eps: float, beta: float, is_positive_side: bool
-#     ) -> tuple[int, float]:
-#         n = len(values)
-#         eps1 = eps2 = eps / 2.0
-#         T = q * n if is_positive_side else (1 - q) * n
-#         noisy_T = T + np.random.exponential(scale=1 / eps1)
-#         i = 0
-#         while True:
-#             candidate = beta**i - 1 if is_positive_side else -(beta**i - 1)
-#             f_i = (values < candidate).sum() if is_positive_side else (values > candidate).sum()
-#             noisy_f_i = f_i + np.random.exponential(scale=1 / eps2)
-#             if noisy_f_i >= noisy_T:
-#                 return i, candidate
-#             i += 1
-
-#     _LOG.info("compute DP unbounded quantiles")
-#     # Split epsilon across quantiles and the two AboveThreshold calls per quantile
-#     eps_pass = epsilon / len(quantiles) / 2.0
-
-#     results = []
-#     for q in quantiles:
-#         # 1) Positive-side AboveThreshold
-#         k, candidate = above_threshold(values, q, eps_pass, beta, is_positive_side=True)
-#         if k > 0:
-#             results.append(candidate)
-#         else:
-#             # 2) Continue with negative-side AboveThreshold only if the first one did not halt at k > 0
-#             k, candidate = above_threshold(values, q, eps_pass, beta, is_positive_side=False)
-#             if k > 0:
-#                 results.append(candidate)
-#             else:
-#                 # 3) Return 0 if both AboveThreshold calls did not halt at k > 0
-#                 results.append(0.0)
-
-#     # ensure monotonicity of results with respect to quantiles
-#     sorted_indices = [t[0] for t in sorted(enumerate(quantiles), key=lambda x: x[1])]
-#     sorted_results = sorted(results)
-#     results = [sorted_results[sorted_indices.index(i)] for i in range(len(quantiles))]
-
-#     # NOTE: consider returning the actual epsilon spent in the future, so that the unused budget can be used for training later
-#     return results
-
-
 def dp_quantiles(values: list | np.ndarray, quantiles: list[float], epsilon: float) -> list[float]:
     """
-    Differentially private quantile estimation.
-    First, estimate the bounds of the values, then use the bounds to estimate the quantiles.
-    If the bounds are not available, estimate the quantiles using the unbounded method.
+    Estimate quantiles with differential privacy using a two-phase approach.
+
+    Phase 1: Estimate data bounds using dp_approx_bounds on a log histogram.
+    Phase 2: Estimate quantiles within those bounds using _dp_bounded_quantiles.
+
+    Privacy budget is split as epsilon/(m+1) for bounds and m*epsilon/(m+1) for m quantiles.
+    Returns None values if bounds cannot be reliably estimated (insufficient data/privacy budget).
 
     Args:
-        values (list | np.ndarray): A list of numeric values.
-        quantiles (list[float]): List of probabilities of the quantiles to estimate.
-        epsilon (float): Privacy budget.
+        values: A list or array of numeric values.
+        quantiles: List of quantile probabilities to estimate (e.g., [0.05, 0.95]).
+        epsilon: Total privacy budget to allocate.
 
     Returns:
-        list[float]: The estimated quantiles.
+        List of differentially private quantile estimates, or list of None if estimation fails.
     """
     values = np.array(values)
 
@@ -850,17 +811,21 @@ def dp_quantiles(values: list | np.ndarray, quantiles: list[float], epsilon: flo
 
 def dp_non_rare(value_counts: dict[str, int], epsilon: float, threshold: int = 5) -> tuple[list[str], float]:
     """
-    Differentially private selection of all categories whose true count >= threshold,
-    via the Laplace vector mechanism + post-processing.
+    Select non-rare categories (count >= threshold) with differential privacy.
+
+    Uses the Laplace vector mechanism: adds independent Laplace(1/ε) noise to each count,
+    then selects categories where noisy_count >= threshold. Also computes the non-rare ratio
+    (fraction of total counts in selected categories).
+
+    Provides ε-differential privacy via the Laplace mechanism with L1 sensitivity = 1.
 
     Args:
-        value_counts (dict): Mapping from category to its count.
-        epsilon (float): Privacy budget.
-        threshold (int): Threshold for non-rare values.
+        value_counts: Mapping from category name to its count.
+        epsilon: Privacy budget.
+        threshold: Minimum count threshold for non-rare categories (default: 5).
 
     Returns:
-        list[str]: Categories whose noisy counts are above the threshold (DP guarantee: ε-DP).
-        float: Non-rare ratio (DP guarantee: ε-DP).
+        Tuple of (selected_categories, non_rare_ratio), both with ε-DP guarantees.
     """
 
     # 1. Add independent Laplace(1/ε) noise to each count (vector Laplace mechanism)
@@ -883,6 +848,20 @@ def dp_non_rare(value_counts: dict[str, int], epsilon: float, threshold: int = 5
 
 
 def get_stochastic_rare_threshold(min_threshold: int = 5, noise_multiplier: float = 3) -> int:
+    """
+    Generate a randomized threshold for rare category detection.
+
+    Adds uniform random noise to the base threshold to prevent adversaries from
+    exploiting knowledge of exact threshold values. The threshold is sampled from
+    [min_threshold, min_threshold + noise_multiplier).
+
+    Args:
+        min_threshold: Base threshold value (default: 5).
+        noise_multiplier: Maximum noise to add (default: 3).
+
+    Returns:
+        Integer threshold in range [min_threshold, min_threshold + noise_multiplier).
+    """
     return min_threshold + int(noise_multiplier * np.random.uniform())
 
 

{mostlyai_engine-1.6.0 → mostlyai_engine-1.7.0}/mostlyai/engine/_tabular/argn.py

@@ -983,6 +983,7 @@ class FlatModel(nn.Module):
         top_p: float | None = None,
         return_probs: list[str] | None = None,
         fairness_transforms: dict[str, Any] | None = None,
+        column_order: list[str] | None = None,
     ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
         fixed_probs = fixed_probs or {}
         fixed_values = fixed_values or {}
@@ -990,6 +991,7 @@ class FlatModel(nn.Module):
         outputs = {}
         probs = {}
         fairness_transforms = fairness_transforms or {}
+        effective_column_order = column_order or self.column_order
 
         if mode == "trn":
             # forward pass through context compressor
@@ -1007,7 +1009,7 @@ class FlatModel(nn.Module):
             col_mask = _make_permutation_mask(
                 col_embedding_dims=self.column_embedders.dims,
                 columns=self.tgt_columns,
-                column_order=self.column_order,
+                column_order=effective_column_order,
                 is_sequential=False,
                 device=self.device,
             )
@@ -1046,7 +1048,7 @@ class FlatModel(nn.Module):
             col_embeddings = torch.cat(list(tgt_col_embeds.values()), dim=-1)
 
             # take sub columns in the specified generation order
-            column_order = self.column_order or self.tgt_columns
+            column_order = effective_column_order or self.tgt_columns
             sub_column_order = [sub_col for col in column_order for sub_col in self.tgt_column_sub_columns[col]]
 
             for sub_col in sub_column_order:
@@ -1267,12 +1269,15 @@ class SequentialModel(nn.Module):
         history=None,
         history_state=None,
         context=None,
+        column_order: list[str] | None = None,
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor, torch.Tensor]:
         fixed_probs = fixed_probs or {}
         fixed_values = fixed_values or {}
         if context is None:
             context = self.context_compressor(x)
 
+        effective_column_order = column_order or self.column_order
+
         has_ridx = any(sub_col.startswith(RIDX_SUB_COLUMN_PREFIX) for sub_col in self.tgt_cardinalities)
 
         # SLEN and RIDX are masked for history
@@ -1318,7 +1323,7 @@ class SequentialModel(nn.Module):
             col_mask = _make_permutation_mask(
                 col_embedding_dims=self.column_embedders.dims,
                 columns=self.tgt_columns,
-                column_order=self.column_order,
+                column_order=effective_column_order,
                 is_sequential=True,
                 device=self.device,
             )
@@ -1387,7 +1392,7 @@ class SequentialModel(nn.Module):
             col_embeddings = torch.cat(list(tgt_col_embeds.values()), dim=-1)
 
             # take sub columns in the specified generation order
-            column_order = self.column_order or self.tgt_columns
+            column_order = effective_column_order or self.tgt_columns
             sub_column_order = [sub_col for col in column_order for sub_col in self.tgt_column_sub_columns[col]]
 
             for sub_col in sub_column_order:

{mostlyai_engine-1.6.0 → mostlyai_engine-1.7.0}/mostlyai/engine/_tabular/generation.py

@@ -16,6 +16,7 @@ import logging
 import random
 import time
 import uuid
+from collections.abc import Callable
 from functools import partial
 from pathlib import Path
 from typing import Literal
@@ -231,6 +232,113 @@ def _regroup_partial_sequences_by_length(
     return ctx_data, len(new_batches)
 
 
+def _flat_null_pattern(group: pd.DataFrame, relevant_columns: list[str]) -> tuple:
+    """returns tuple of bools: True if column is fully NULL"""
+    return tuple(group[col].isna().all() for col in relevant_columns)
+
+
+def _trailing_null_pattern(group: pd.DataFrame, relevant_columns: list[str]) -> tuple:
+    """returns tuple of bools: True if column has trailing NULLs"""
+    pattern = []
+    for col in relevant_columns:
+        col_values = group[col].reset_index(drop=True)
+        non_null_mask = col_values.notna()
+        if non_null_mask.any():
+            last_non_null_idx = non_null_mask[::-1].idxmax()
+            has_trailing_nulls = last_non_null_idx < len(col_values) - 1
+        else:
+            has_trailing_nulls = True
+        pattern.append(has_trailing_nulls)
+    return tuple(pattern)
+
+
+def _regroup_by_pattern(
+    ctx_data: pd.DataFrame,
+    seed_data: pd.DataFrame,
+    ctx_primary_key: str,
+    imputation_columns: list[str],
+    pattern_fn: Callable[[pd.DataFrame, list[str]], tuple],
+    *,
+    groupby_key: str | None = None,
+    use_vectorized_regroup: bool = True,
+) -> tuple[pd.DataFrame, int]:
+    """regroup batches so that rows/sequences with the same NULL pattern are together
+
+    Args:
+        pattern_fn: function that computes NULL pattern for a grouped dataframe
+        groupby_key: key to group seed_data by (defaults to ctx_primary_key)
+        use_vectorized_regroup: if True, use fast categorical factorization;
+                                if False, use nested loop approach
+    """
+    # only consider columns that are BOTH in imputation_columns AND in seed_data
+    relevant_columns = [col for col in imputation_columns if col in seed_data.columns]
+
+    # early exit: no relevant columns
+    if not relevant_columns or seed_data[relevant_columns].isna().all().all():
+        return ctx_data, ctx_data["__BATCH"].nunique()
+
+    # compute NULL pattern for each group
+    groupby_key = groupby_key or ctx_primary_key
+    seed_data_grouped = seed_data.groupby(groupby_key, sort=False)
+    null_patterns = seed_data_grouped.apply(
+        lambda group: pattern_fn(group, relevant_columns),
+        include_groups=False,
+    ).rename("__NULL_PATTERN")
+
+    # early exit: all NULL patterns are the same
+    if null_patterns.nunique() == 1:
+        return ctx_data, ctx_data["__BATCH"].nunique()
+
+    # add __NULL_PATTERN to ctx_data
+    ctx_data = ctx_data.assign(__NULL_PATTERN=ctx_data[ctx_primary_key].map(null_patterns))
+
+    # regroup batches
+    if use_vectorized_regroup:
+        # vectorized approach for flat data
+        ctx_data = ctx_data.assign(
+            __COMPOSITE_KEY=ctx_data["__BATCH"].astype(str) + "_" + ctx_data["__NULL_PATTERN"].astype(str)
+        )
+        composite_cat = pd.Categorical(ctx_data["__COMPOSITE_KEY"], categories=ctx_data["__COMPOSITE_KEY"].unique())
+        ctx_data = ctx_data.assign(__BATCH=pd.factorize(composite_cat)[0] + 1)
+        num_batches = ctx_data["__BATCH"].max()
+        ctx_data = ctx_data.drop(columns=["__NULL_PATTERN", "__COMPOSITE_KEY"]).reset_index(drop=True)
+    else:
+        # nested loop approach for sequential data
+        new_batches = []
+        for _, old_batch_df in ctx_data.groupby("__BATCH", sort=False):
+            for _, new_batch_df in old_batch_df.groupby("__NULL_PATTERN", sort=False):
+                new_batch_df = new_batch_df.assign(__BATCH=len(new_batches) + 1)
+                new_batches.append(new_batch_df)
+        ctx_data = pd.concat(new_batches, axis=0).drop(columns=["__NULL_PATTERN"]).reset_index(drop=True)
+        num_batches = len(new_batches)
+
+    return ctx_data, num_batches
+
+
+def _regroup_by_null_pattern(
+    ctx_data: pd.DataFrame,
+    seed_data: pd.DataFrame,
+    ctx_primary_key: str,
+    tgt_context_key: str,
+    imputation_columns: list[str],
+    is_sequential: bool,
+) -> tuple[pd.DataFrame, int]:
+    """regroup batches by NULL pattern (flat) or trailing NULL pattern (sequential)"""
+    pattern_fn = _trailing_null_pattern if is_sequential else _flat_null_pattern
+    groupby_key = tgt_context_key if is_sequential else None
+    use_vectorized_regroup = not is_sequential
+
+    return _regroup_by_pattern(
+        ctx_data,
+        seed_data,
+        ctx_primary_key,
+        imputation_columns,
+        pattern_fn=pattern_fn,
+        groupby_key=groupby_key,
+        use_vectorized_regroup=use_vectorized_regroup,
+    )
+
+
 def _reshape_pt_to_pandas(
     data: list[torch.Tensor], sub_cols: list[str], keys: list[pd.Series], key_name: str
 ) -> pd.DataFrame:
@@ -262,18 +370,34 @@ def _reshape_pt_to_pandas(
     return pd.concat([keys, df], axis=1)
 
 
+def _drop_fully_null_imputed_columns(seed_batch: pd.DataFrame, imputation_columns: list[str]) -> pd.DataFrame:
+    """drop columns from seed_batch that are in imputation_columns and are fully NULL
+
+    this allows the model to freely generate these columns rather than conditioning on NULL values.
+    """
+    if not imputation_columns or seed_batch.empty:
+        return seed_batch
+
+    fully_null_cols = [col for col in imputation_columns if col in seed_batch.columns and seed_batch[col].isna().all()]
+    return seed_batch.drop(columns=fully_null_cols) if fully_null_cols else seed_batch
+
+
 def _post_process_decoding(
     syn: pd.DataFrame,
     tgt_primary_key: str | None = None,
 ) -> pd.DataFrame:
-    # remove dummy context key (if exists)
+    # sort by dummy context key to restore original order (if exists)
     if DUMMY_CONTEXT_KEY in syn:
+        syn = syn.sort_values(DUMMY_CONTEXT_KEY).reset_index(drop=True)
         syn = syn.drop(columns=DUMMY_CONTEXT_KEY)
 
     # generate primary keys, if they are not present
     if tgt_primary_key and tgt_primary_key not in syn:
         syn[tgt_primary_key] = _generate_primary_keys(len(syn), type="uuid")
 
+    # reset index to ensure sequential indices for consistent test assertions
+    syn = syn.reset_index(drop=True)
+
     return syn
 
 
@@ -580,6 +704,7 @@ def decode_buffered_samples(
     tgt_primary_key: str,
     tgt_context_key: str,
     decode_prev_steps: dict | None = None,
+    impute_columns: list[str] | None = None,
 ) -> pd.DataFrame:
     is_sequential = tgt_stats["is_sequential"]
     seq_len_stats = get_sequence_length_stats(tgt_stats)
@@ -614,10 +739,12 @@ def decode_buffered_samples(
 
     # preserve all seed values
     df_seed = pd.concat(seed_data, axis=0).reset_index(drop=True) if seed_data else pd.DataFrame()
+
     if not df_seed.empty:
         seed_columns = [col for col in df_seed.columns]
         if is_sequential:
             # overwrite first steps of each sequence in synthetic data with values from seed data
+            impute_columns = impute_columns or []
             df_syn["__SEQ_IDX"] = df_syn.groupby(tgt_context_key).cumcount()
             df_seed["__SEQ_IDX"] = df_seed.groupby(tgt_context_key).cumcount()
             # df_overwrite is a dataframe with the same shape as df_syn, but with the seed values for the first steps of each sequence
@@ -630,12 +757,47 @@ def decode_buffered_samples(
             )
             # project df_overwrite onto df_syn
             seed_rows = df_overwrite["__INDICATOR"] == "both"
-            df_syn.loc[seed_rows, seed_columns] = df_overwrite.loc[seed_rows, seed_columns]
+            # overwrite columns based on imputation logic
+            for col in seed_columns:
+                if col in [tgt_context_key, "__SEQ_IDX"]:
+                    continue  # skip the key columns
+                if col not in impute_columns:
+                    # non-impute columns: override all values
+                    df_syn.loc[seed_rows, col] = df_overwrite.loc[seed_rows, col]
+                else:
+                    # impute columns: override only non-NULL seed values
+                    mask = seed_rows & df_overwrite[col].notna()
+                    df_syn.loc[mask, col] = df_overwrite.loc[mask, col]
             df_syn.drop(columns=["__SEQ_IDX"], inplace=True)
             df_seed.drop(columns=["__SEQ_IDX"], inplace=True)
         else:
-            # for flat data, just overwrite all seed columns
-            df_syn[seed_columns] = df_seed[seed_columns].copy()
+            # for flat data, overwrite seed columns using merge to handle reordered rows
+            # for non-impute columns: override all values
+            # for impute columns: override only non-NULL seed values (let model impute NULL values)
+            impute_columns = impute_columns or []
+
+            # use merge on context key to properly align seed values with synthetic data
+            df_overwrite = pd.merge(
+                df_syn[[tgt_context_key]].copy(),
+                df_seed,
+                on=tgt_context_key,
+                how="left",
+                suffixes=("", "_seed"),
+            )
+
+            # overwrite columns based on imputation logic
+            for col in seed_columns:
+                if col == tgt_context_key:
+                    continue  # skip the key column itself
+                seed_col_name = col if col in df_overwrite.columns else f"{col}_seed"
+                if seed_col_name in df_overwrite.columns:
+                    if col not in impute_columns:
+                        # non-impute columns: override all values
+                        df_syn[col] = df_overwrite[seed_col_name]
+                    else:
+                        # impute columns: override only non-NULL seed values
+                        mask = df_overwrite[seed_col_name].notna()
+                        df_syn.loc[mask, col] = df_overwrite.loc[mask, seed_col_name]
 
     # postprocess generated data
     _LOG.info(f"post-process generated data {df_syn.shape}")
@@ -740,10 +902,11 @@ def generate(
             fairness=fairness,
         )
         _LOG.info(f"{gen_column_order=}")
+        trn_column_order = get_columns_from_cardinalities(tgt_cardinalities)
+        _LOG.info(f"{trn_column_order=}")
+
         if not enable_flexible_generation:
             # check if resolved column order is the same as the one from training
-            trn_column_order = get_columns_from_cardinalities(tgt_cardinalities)
-            _LOG.info(f"{trn_column_order=}")
             if gen_column_order != trn_column_order:
                 raise ValueError(
                     "The column order for generation does not match the column order from training, due to seed, rebalancing, fairness or imputation configs. "
@@ -925,6 +1088,12 @@ def generate(
                 ctx_data, seed_data, ctx_primary_key, tgt_context_key
             )
 
+        # regroup by NULL pattern if imputation is enabled
+        if imputation and seed_data is not None and len(seed_data) > 0:
+            ctx_data, no_of_batches = _regroup_by_null_pattern(
+                ctx_data, seed_data, ctx_primary_key, tgt_context_key, imputation.columns, is_sequential
+            )
+
         # keep at most 500k samples in memory before decoding and writing to disk
         buffer = FixedSizeSampleBuffer(capacity=500_000)
 
@@ -937,6 +1106,9 @@ def generate(
             batch_size = len(ctx_batch)
 
             seed_batch = seed_data[seed_data[tgt_context_key].isin(ctx_batch[ctx_primary_key])]
+            # drop fully-NULL imputation columns from seed_batch to allow conditional generation
+            if imputation:
+                seed_batch = _drop_fully_null_imputed_columns(seed_batch, imputation.columns)
             seed_batch = apply_encoding_type_dtypes(seed_batch, seed_encoding_types)
 
             if ctx_primary_key not in ctx_batch.columns:
@@ -1014,9 +1186,18 @@ def generate(
 
                     # get seed data for current step
                     seed_step = seed_batch_grouped.nth(seq_step) if seq_step < n_seed_steps else pd.DataFrame()
-                    seed_step_encoded = (
-                        seed_batch_encoded_grouped.nth(seq_step) if seq_step < n_seed_steps else pd.DataFrame()
-                    )
+
+                    # drop NULL imputation columns from seed_step to allow conditional generation
+                    if imputation and len(seed_step) > 0:
+                        seed_step = _drop_fully_null_imputed_columns(seed_step, imputation.columns)
+
+                    # encode seed_step (after dropping NULL imputation columns)
+                    if len(seed_step) > 0:
+                        seed_step_encoded, _, _ = encode_df(
+                            df=seed_step, stats=tgt_stats, tgt_context_key=tgt_context_key
+                        )
+                    else:
+                        seed_step_encoded = pd.DataFrame()
 
                     # fix SIDX by incrementing ourselves instead of sampling
                     sidx = pd.Series([seq_step] * step_size)
@@ -1079,11 +1260,19 @@ def generate(
                                 torch.as_tensor(seed_step_encoded[col].to_numpy(), device=model.device).type(torch.int),
                                 dim=-1,
                             )
-                            for col in seed_batch_encoded.columns
+                            for col in seed_step_encoded.columns
                             if col in tgt_sub_columns
                         }
 
                     fixed_values = sidx_vals | slen_vals | ridx_vals | sdec_vals | seed_vals
+                    column_order = _resolve_gen_column_order(
+                        column_stats=tgt_stats["columns"],
+                        cardinalities=tgt_cardinalities,
+                        rebalancing=rebalancing,
+                        imputation=imputation,
+                        seed_data=seed_step,
+                        fairness=fairness,
+                    )
                     out_dct, history, history_state = model(
                         x=None,  # not used in generation forward pass
                         mode="gen",
@@ -1095,6 +1284,7 @@ def generate(
                         history=history,
                         history_state=history_state,
                         context=context,
+                        column_order=column_order,
                     )
 
                     # transform output dict to tensor for memory efficiency
@@ -1147,6 +1337,9 @@ def generate(
                         ]
                         history = history[include_mask, ...]
                         history_state = tuple(h[:, include_mask, ...] for h in history_state)
+                    # filter seed_step to match step_ctx_keys (always, not just when filtering above)
+                    if len(seed_step) > 0:
+                        seed_step = seed_step[seed_step[tgt_context_key].isin(step_ctx_keys)].reset_index(drop=True)
                     # accumulate outputs in memory
                     buffer.add((out_pt, step_ctx_keys, seed_step))
                     # increment progress by 1 for each step
@@ -1160,6 +1353,7 @@ def generate(
                             tgt_primary_key=tgt_primary_key,
                             tgt_context_key=tgt_context_key,
                             decode_prev_steps=decode_prev_steps,
+                            impute_columns=imputation.columns if imputation else None,
                         )
                         persist_data_part(syn, output_path, f"{buffer.n_clears:06}.{0:06}")
                         buffer.clear()
@@ -1190,6 +1384,14 @@ def generate(
                     if col in tgt_sub_columns
                 }
 
+                column_order = _resolve_gen_column_order(
+                    column_stats=tgt_stats["columns"],
+                    cardinalities=tgt_cardinalities,
+                    rebalancing=rebalancing,
+                    imputation=imputation,
+                    seed_data=seed_batch,
+                    fairness=fairness,
+                )
                 out_dct, _ = model(
                     x,
                     mode="gen",
@@ -1199,6 +1401,7 @@ def generate(
                     temperature=sampling_temperature,
                     top_p=sampling_top_p,
                     fairness_transforms=fairness_transforms,
+                    column_order=column_order,
                 )
 
                 syn = pd.concat(
@@ -1224,6 +1427,7 @@ def generate(
                     tgt_primary_key=tgt_primary_key,
                     tgt_context_key=tgt_context_key,
                     decode_prev_steps=decode_prev_steps,
+                    impute_columns=imputation.columns if imputation else None,
                 )
                 persist_data_part(syn, output_path, f"{buffer.n_clears:06}.{0:06}")
                 buffer.clear()
@@ -1239,6 +1443,7 @@ def generate(
                 tgt_primary_key=tgt_primary_key,
                 tgt_context_key=tgt_context_key,
                 decode_prev_steps=decode_prev_steps,
+                impute_columns=imputation.columns if imputation else None,
             )
             persist_data_part(syn, output_path, f"{buffer.n_clears:06}.{0:06}")
             buffer.clear()

{mostlyai_engine-1.6.0 → mostlyai_engine-1.7.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mostlyai-engine"
-version = "1.6.0"
+version = "1.7.0"
 description = "Synthetic Data Engine"
 authors = [{ name = "MOSTLY AI", email = "dev@mostly.ai" }]
 requires-python = ">=3.10"
@@ -55,6 +55,7 @@ gpu = [
 [dependency-groups]
 dev = [
     "pytest>=8.0",
+    "pytest-rerunfailures>=15.0",
     "ruff>=0.11",  # sync'ed with .pre-commit-config
     "pre-commit>=4.0",
     "twine>=6.1",