aetherdialect 0.1.0__py3-none-any.whl

text2sql/qsim_sample.py

@@ -0,0 +1,609 @@
+"""Value sampling and instantiation for the question-generation simulator.
+
+Implements operator-aware value selection across categorical, numeric, boolean, and temporal domains, with coordinated range sampling for decomposed BETWEEN pairs (lower from [0.15, 0.35], upper from [0.65, 0.85]) and deterministic HAVING value generation keyed by variant index.
+Skips instantiation for IS NULL filters and column-to-column comparisons.
+Populates QSimIntent param_values via index-based keys, computes per-intent variance scores, and performs single-pass proportional variant allocation across the full intent set.
+"""
+
+from __future__ import annotations
+
+import random
+from dataclasses import replace
+from datetime import datetime, timedelta
+
+from .config import (
+    AGG_PATTERN,
+    HAVING_COUNT_VALUES,
+    HAVING_MIN_MAX_VALUES,
+    HAVING_SUM_AVG_VALUES,
+    QSimConfig,
+)
+from .contracts_base import SchemaGraph, ValueDomain
+from .contracts_core import QSimFilter, QSimHaving, QSimIntent
+from .core_utils import debug
+from .qsim_struct import decompose_between_filter
+
+
+def _is_integer_type(data_type: str | None) -> bool:
+    """Determine if data type represents an integer."""
+    if not data_type:
+        return False
+    dtype_lower = data_type.lower()
+    if dtype_lower in (
+        "integer",
+        "int",
+        "bigint",
+        "smallint",
+        "tinyint",
+        "long",
+        "short",
+    ):
+        return True
+    if "int" in dtype_lower or dtype_lower in ("long", "short"):
+        if "interval" not in dtype_lower:
+            return True
+    return False
+
+
+def _parse_date(val: str) -> datetime | None:
+    """Parse date string to datetime object."""
+    if "T" in val:
+        val = val.split("T")[0]
+    elif " " in val:
+        val = val.split(" ")[0]
+
+    for fmt in ("%Y-%m-%d", "%Y/%m/%d", "%d-%m-%Y", "%d/%m/%Y"):
+        try:
+            return datetime.strptime(val, fmt)
+        except ValueError:
+            continue
+    return None
+
+
+def _format_date(dt: datetime) -> str:
+    """Format datetime to date-only string."""
+    return dt.strftime("%Y-%m-%d")
+
+
+def _extract_date_part(val: str) -> str:
+    """Extract date part from datetime string."""
+    if "T" in val:
+        return val.split("T")[0]
+    if " " in val:
+        return val.split(" ")[0]
+    return val
+
+
+def _sample_categorical(domain: ValueDomain, variant_idx: int) -> str | None:
+    """Sample categorical value."""
+    values_list = domain.values
+    if values_list:
+        idx = variant_idx % len(values_list)
+        return values_list[idx]
+    if domain.min_val is not None and domain.max_val is not None:
+        try:
+            min_v = int(float(domain.min_val))
+            max_v = int(float(domain.max_val))
+            range_size = max(1, max_v - min_v + 1)
+            value = min_v + (variant_idx % range_size)
+            return str(value)
+        except (ValueError, TypeError):
+            return str(domain.min_val)
+    return None
+
+
+def _sample_boolean(domain: ValueDomain, variant_idx: int) -> str | None:
+    """Sample boolean value."""
+    values_list = domain.values
+    if values_list:
+        normalized = []
+        for v in values_list:
+            if isinstance(v, bool):
+                normalized.append("true" if v else "false")
+            elif isinstance(v, str):
+                normalized.append(v.lower() if v.lower() in ("true", "false") else v)
+            else:
+                normalized.append(str(v))
+        idx = variant_idx % len(normalized)
+        return normalized[idx]
+    default_bools = ["true", "false"]
+    idx = variant_idx % len(default_bools)
+    return default_bools[idx]
+
+
+def _sample_numeric_categorical(domain: ValueDomain, variant_idx: int) -> str | None:
+    """Sample numeric categorical value from discrete set."""
+    values_list = domain.values
+    if values_list:
+        idx = variant_idx % len(values_list)
+        val = values_list[idx]
+        return str(int(float(val))) if isinstance(val, int | float) else str(val)
+    if domain.min_val is not None and domain.max_val is not None:
+        try:
+            min_v = int(float(domain.min_val))
+            max_v = int(float(domain.max_val))
+            range_size = max(1, max_v - min_v + 1)
+            value = min_v + (variant_idx % range_size)
+            return str(value)
+        except (ValueError, TypeError):
+            return str(int(float(domain.min_val)))
+    return None
+
+
+def _sample_numeric(domain: ValueDomain, op: str, variant_idx: int) -> str | None:
+    """Sample numeric value with operator awareness."""
+    if domain.min_val is not None and domain.max_val is not None:
+        try:
+            min_v = float(domain.min_val)
+            max_v = float(domain.max_val)
+            range_size = max_v - min_v
+            is_integer = _is_integer_type(domain.data_type)
+
+            if op == "=":
+                if is_integer:
+                    int_range = max(1, int(range_size + 1))
+                    value = int(min_v + (variant_idx % int_range))
+                else:
+                    segment = (variant_idx % 10) / 10.0
+                    value = min_v + segment * range_size
+                    value = round(value, 2) if abs(value) >= 1 else round(value, 4)
+            elif op in (">", ">="):
+                lower_bound = min_v + range_size * 0.2
+                upper_bound = min_v + range_size * 0.5
+                value = lower_bound + (variant_idx % 5) * (upper_bound - lower_bound) / 5
+                value = int(round(value)) if is_integer else (round(value, 2) if abs(value) >= 1 else round(value, 4))
+            elif op in ("<", "<="):
+                lower_bound = min_v + range_size * 0.5
+                upper_bound = min_v + range_size * 0.8
+                value = lower_bound + (variant_idx % 5) * (upper_bound - lower_bound) / 5
+                value = int(round(value)) if is_integer else (round(value, 2) if abs(value) >= 1 else round(value, 4))
+            else:
+                if range_size > 0:
+                    segment = (variant_idx % 10) / 10.0
+                    value = min_v + segment * range_size
+                else:
+                    value = min_v
+                value = int(round(value)) if is_integer else (round(value, 2) if abs(value) >= 1 else round(value, 4))
+
+            return str(value)
+        except (ValueError, TypeError):
+            pass
+
+    values_list = domain.values
+    if values_list:
+        idx = variant_idx % len(values_list)
+        return values_list[idx]
+
+    return None
+
+
+def _sample_temporal(domain: ValueDomain, op: str, variant_idx: int) -> str | None:
+    """Sample temporal value with date interpolation."""
+    if domain.min_val is not None and domain.max_val is not None:
+        try:
+            min_dt = _parse_date(str(domain.min_val))
+            max_dt = _parse_date(str(domain.max_val))
+
+            if min_dt is None or max_dt is None:
+                return _extract_date_part(str(domain.min_val))
+
+            total_days = (max_dt - min_dt).days
+            if total_days <= 0:
+                return _format_date(min_dt)
+
+            if op in (">", ">="):
+                segment = 0.2 + ((variant_idx % 5) / 5.0) * 0.15
+            elif op in ("<", "<="):
+                segment = 0.65 + ((variant_idx % 5) / 5.0) * 0.15
+            else:
+                segment = (variant_idx % 10) / 10.0
+
+            offset_days = int(total_days * segment)
+            result_dt = min_dt + timedelta(days=offset_days)
+            return _format_date(result_dt)
+        except (ValueError, TypeError):
+            pass
+
+    values_list = domain.values
+    if values_list:
+        idx = variant_idx % len(values_list)
+        return _extract_date_part(values_list[idx])
+
+    return None
+
+
+def _sample_in_values(domain: ValueDomain, value_type: str, variant_idx: int) -> str | None:
+    """Sample multiple values for IN/NOT IN operators."""
+    if value_type == "categorical":
+        values_list = domain.values
+        if values_list:
+            n_values = min(3 + (variant_idx % 3), len(values_list))
+            start_idx = variant_idx % max(1, len(values_list) - n_values + 1)
+            values = values_list[start_idx : start_idx + n_values]
+            return "'" + "','".join(values) + "'"
+
+    elif value_type == "numeric_categorical":
+        values_list = domain.values
+        if values_list:
+            n_values = min(3 + (variant_idx % 3), len(values_list))
+            start_idx = variant_idx % max(1, len(values_list) - n_values + 1)
+            values = values_list[start_idx : start_idx + n_values]
+            int_values = [str(int(float(v))) if isinstance(v, int | float) else str(v) for v in values]
+            return ",".join(int_values)
+        if domain.min_val is not None and domain.max_val is not None:
+            try:
+                min_v = int(float(domain.min_val))
+                max_v = int(float(domain.max_val))
+                range_size = max(1, max_v - min_v + 1)
+                n_values = min(3 + (variant_idx % 3), range_size)
+                values = []
+                for i in range(n_values):
+                    value = min_v + ((variant_idx + i) % range_size)
+                    values.append(str(value))
+                return ",".join(values)
+            except (ValueError, TypeError):
+                pass
+
+    elif value_type == "boolean":
+        values_list = domain.values
+        if values_list:
+            normalized = []
+            for v in values_list:
+                if isinstance(v, bool):
+                    normalized.append("true" if v else "false")
+                elif isinstance(v, str):
+                    normalized.append(v.lower() if v.lower() in ("true", "false") else v)
+                else:
+                    normalized.append(str(v))
+            return ",".join(normalized)
+        return "true,false"
+
+    elif value_type in ("numeric", "temporal"):
+        if domain.min_val is not None and domain.max_val is not None:
+            try:
+                min_v = float(domain.min_val)
+                max_v = float(domain.max_val)
+                range_size = max_v - min_v
+                is_integer = _is_integer_type(domain.data_type)
+                n_values = 2 + (variant_idx % 3)
+                values = []
+                for i in range(n_values):
+                    segment = ((variant_idx + i) % 10) / 10.0
+                    val = min_v + segment * range_size
+                    val = int(round(val)) if is_integer else (round(val, 2) if abs(val) >= 1 else round(val, 4))
+                    values.append(str(val))
+                return ",".join(values)
+            except (ValueError, TypeError):
+                pass
+
+    return None
+
+
+def sample_value_from_domain(domain: ValueDomain, value_type: str, op: str = "=", variant_idx: int = 0) -> str | None:
+    """Sample a concrete value from a column's domain with operator awareness."""
+    if value_type == "null" or op in ("is null", "is not null"):
+        return None
+
+    if op in ("in", "not in"):
+        return _sample_in_values(domain, value_type, variant_idx)
+
+    if value_type == "categorical":
+        return _sample_categorical(domain, variant_idx)
+
+    if value_type == "numeric_categorical":
+        return _sample_numeric_categorical(domain, variant_idx)
+
+    if value_type == "numeric":
+        return _sample_numeric(domain, op, variant_idx)
+
+    if value_type == "temporal":
+        return _sample_temporal(domain, op, variant_idx)
+
+    if value_type == "boolean":
+        return _sample_boolean(domain, variant_idx)
+
+    return None
+
+
+def _identify_range_pairs(filters: list[QSimFilter]) -> dict[str, dict[str, int]]:
+    """Identify columns with paired range filters (>= and </<= on same column)."""
+    column_ops: dict[str, dict[str, int]] = {}
+    for idx, f in enumerate(filters):
+        if f.is_expr_comparison:
+            continue
+        if f.op in (">", ">="):
+            column_ops.setdefault(f.column, {})["lower_idx"] = idx
+        elif f.op in ("<", "<="):
+            column_ops.setdefault(f.column, {})["upper_idx"] = idx
+    return {col: ops for col, ops in column_ops.items() if "lower_idx" in ops and "upper_idx" in ops}
+
+
+def _sample_numeric_range(domain: ValueDomain, variant_idx: int) -> tuple[str | None, str | None]:
+    """Sample coordinated numeric range values."""
+    if domain.min_val is None or domain.max_val is None:
+        return None, None
+
+    try:
+        min_v = float(domain.min_val)
+        max_v = float(domain.max_val)
+        range_size = max_v - min_v
+        if range_size <= 0:
+            return None, None
+
+        is_integer = _is_integer_type(domain.data_type)
+
+        lower_segment = 0.15 + ((variant_idx % 5) / 5.0) * 0.2
+        upper_segment = 0.65 + ((variant_idx % 5) / 5.0) * 0.2
+
+        lower_val = min_v + lower_segment * range_size
+        upper_val = min_v + upper_segment * range_size
+
+        if is_integer:
+            lower_val = int(round(lower_val))
+            upper_val = int(round(upper_val))
+            if lower_val >= upper_val:
+                upper_val = min(lower_val + 1, int(max_v))
+        else:
+            lower_val = round(lower_val, 2) if abs(lower_val) >= 1 else round(lower_val, 4)
+            upper_val = round(upper_val, 2) if abs(upper_val) >= 1 else round(upper_val, 4)
+
+        return str(lower_val), str(upper_val)
+    except (ValueError, TypeError):
+        return None, None
+
+
+def _sample_temporal_range(domain: ValueDomain, variant_idx: int) -> tuple[str | None, str | None]:
+    """Sample coordinated temporal range values with date interpolation."""
+    if domain.min_val is None or domain.max_val is None:
+        return None, None
+
+    try:
+        min_dt = _parse_date(str(domain.min_val))
+        max_dt = _parse_date(str(domain.max_val))
+
+        if min_dt is None or max_dt is None:
+            lower_val = _extract_date_part(str(domain.min_val))
+            upper_val = _extract_date_part(str(domain.max_val))
+            return lower_val, upper_val
+
+        total_days = (max_dt - min_dt).days
+        if total_days <= 0:
+            return _format_date(min_dt), _format_date(max_dt)
+
+        lower_segment = 0.15 + ((variant_idx % 5) / 5.0) * 0.2
+        upper_segment = 0.65 + ((variant_idx % 5) / 5.0) * 0.2
+
+        lower_days = int(total_days * lower_segment)
+        upper_days = int(total_days * upper_segment)
+
+        lower_dt = min_dt + timedelta(days=lower_days)
+        upper_dt = min_dt + timedelta(days=upper_days)
+
+        return _format_date(lower_dt), _format_date(upper_dt)
+    except (ValueError, TypeError):
+        return None, None
+
+
+def sample_coordinated_range(domain: ValueDomain, value_type: str, variant_idx: int) -> tuple[str | None, str | None]:
+    """Sample coordinated lower and upper values for range pairs."""
+    if value_type not in ("numeric", "temporal"):
+        return None, None
+
+    if value_type == "numeric":
+        return _sample_numeric_range(domain, variant_idx)
+
+    if value_type == "temporal":
+        return _sample_temporal_range(domain, variant_idx)
+
+    return None, None
+
+
+def deterministic_having_value(agg_func: str, variant_idx: int, having_idx: int = 0) -> str:
+    """Generate deterministic HAVING value from predefined pools."""
+    offset = variant_idx * 3 + having_idx
+
+    if agg_func == "count":
+        value = HAVING_COUNT_VALUES[offset % len(HAVING_COUNT_VALUES)]
+        return str(value)
+
+    if agg_func in {"sum", "avg"}:
+        value = HAVING_SUM_AVG_VALUES[offset % len(HAVING_SUM_AVG_VALUES)]
+        return str(value)
+
+    if agg_func in {"min", "max"}:
+        value = HAVING_MIN_MAX_VALUES[offset % len(HAVING_MIN_MAX_VALUES)]
+        return str(value)
+
+    idx = offset % len(HAVING_COUNT_VALUES)
+    return str(HAVING_COUNT_VALUES[idx])
+
+
+def _compute_intent_variance(intent: QSimIntent, value_domains: dict[str, ValueDomain]) -> int:
+    """Compute variance score for intent instantiation potential."""
+    variance_score = 0
+
+    for f in intent.filters_param:
+        if f.is_expr_comparison:
+            continue
+        col_key = f.column
+        domain = value_domains.get(col_key)
+        if domain:
+            if domain.values:
+                variance_score += len(domain.values)
+            elif domain.min_val is not None and domain.max_val is not None:
+                variance_score += 10
+
+    if intent.filters_param:
+        variance_score += 10 * len(intent.having_param)
+    else:
+        variance_score += 5 * len(intent.having_param)
+
+    return variance_score
+
+
+def _instantiate_intent(
+    intent: QSimIntent, value_domains: dict[str, ValueDomain], variant_idx: int = 0
+) -> QSimIntent | None:
+    """Populate QSimIntent filter/having values via index-based param_values."""
+
+    decomposed_filters: list[QSimFilter] = []
+    for f in intent.filters_param:
+        decomposed_filters.extend(decompose_between_filter(f))
+
+    range_pairs = _identify_range_pairs(decomposed_filters)
+    range_values: dict[str, tuple[str, str]] = {}
+
+    for col_key, pair_indices in range_pairs.items():
+        domain = value_domains.get(col_key)
+        if domain is None:
+            continue
+        lower_idx = pair_indices["lower_idx"]
+        value_type = decomposed_filters[lower_idx].value_type
+        lower_val, upper_val = sample_coordinated_range(domain, value_type, variant_idx)
+        if lower_val is not None and upper_val is not None:
+            range_values[col_key] = (lower_val, upper_val)
+
+    new_filters: list[QSimFilter] = []
+    new_param_values: dict[str, any] = {}
+
+    for filter_idx, f in enumerate(decomposed_filters):
+        param_key = f"f{filter_idx}"
+
+        if f.is_expr_comparison:
+            new_filters.append(f)
+            debug(f"[qsim_sample.instantiate_intent] expr_comparison: {f.column} {f.op} {f.right_column}")
+            continue
+
+        col_key = f.column
+        value_type = f.value_type
+        op = f.op
+
+        if value_type == "null" or op in ("is null", "is not null"):
+            new_filters.append(replace(f, value_type="null"))
+            debug(f"[qsim_sample.instantiate_intent] null_filter: {col_key} {op}")
+            continue
+
+        domain = value_domains.get(col_key)
+
+        if domain is None:
+            debug(f"[qsim_sample.instantiate_intent] no_domain: {col_key}")
+            new_filters.append(f)
+            continue
+
+        if col_key in range_values:
+            lower_val, upper_val = range_values[col_key]
+            if f.op in (">", ">="):
+                value = lower_val
+            elif f.op in ("<", "<="):
+                value = upper_val
+            else:
+                combined_idx = variant_idx * len(decomposed_filters) + filter_idx
+                value = sample_value_from_domain(domain, value_type, f.op, combined_idx)
+        else:
+            combined_idx = variant_idx * len(decomposed_filters) + filter_idx
+            value = sample_value_from_domain(domain, value_type, f.op, combined_idx)
+
+        if value is not None:
+            new_param_values[param_key] = value
+
+        new_filters.append(f)
+
+    new_having: list[QSimHaving] = []
+    for having_idx, h in enumerate(intent.having_param):
+        param_key = f"h{having_idx}"
+
+        agg_match = AGG_PATTERN.match(h.expression)
+        agg_func = agg_match.group(1).lower() if agg_match else "count"
+        value = deterministic_having_value(agg_func, variant_idx, having_idx)
+        new_param_values[param_key] = value
+
+        new_having.append(h)
+
+    return QSimIntent(
+        intent_id=intent.intent_id,
+        tables=intent.tables,
+        grain=intent.grain,
+        select_cols=intent.select_cols,
+        group_by_cols=intent.group_by_cols,
+        order_by_cols=intent.order_by_cols,
+        filters_param=new_filters,
+        having_param=new_having,
+        param_values=new_param_values,
+        question="",
+        variant_idx=variant_idx,
+        limit=intent.limit,
+        distinct=intent.distinct,
+    )
+
+
+def instantiate_all(intents: list[QSimIntent], schema: SchemaGraph, num_questions: int = None) -> list[QSimIntent]:
+    """Generate QSimIntent instances with populated filter/having values via single-pass proportional allocation."""
+    if num_questions is None:
+        num_questions = QSimConfig.QUESTIONS_COUNT
+
+    random.seed(QSimConfig.RANDOM_SEED)
+
+    avg_variants = num_questions / len(intents) if intents else 0
+    if avg_variants < QSimConfig.MIN_AVG_VARIANTS_PER_INTENT:
+        debug(
+            f"[qsim_sample.instantiate_all] WARNING: avg_variants={avg_variants:.2f} below MIN={QSimConfig.MIN_AVG_VARIANTS_PER_INTENT}"
+        )
+    if avg_variants > QSimConfig.MAX_AVG_VARIANTS_PER_INTENT:
+        raise ValueError(
+            f"Intent/variant ratio unrealistic: {len(intents)} intents cannot generate {num_questions} diverse questions (avg={avg_variants:.1f} > max={QSimConfig.MAX_AVG_VARIANTS_PER_INTENT})"
+        )
+
+    value_domains: dict[str, ValueDomain] = {}
+    for table_name, table_meta in schema.tables.items():
+        for col_name, col_meta in table_meta.columns.items():
+            col_key = f"{table_name}.{col_name}"
+            value_domains[col_key] = ValueDomain(
+                values=col_meta.top_k_values or [],
+                min_val=col_meta.min_val,
+                max_val=col_meta.max_val,
+                data_type=col_meta.data_type,
+            )
+    debug(f"[qsim_sample.instantiate_all] value_domains: {len(value_domains)} columns")
+
+    variances: dict[str, float] = {}
+    for intent in intents:
+        variances[intent.intent_id] = _compute_intent_variance(intent, value_domains)
+
+    total_variance = sum(v for v in variances.values() if v > 0)
+
+    allocations: dict[str, int] = {}
+    if total_variance == 0:
+        for intent in intents:
+            allocations[intent.intent_id] = 1
+    else:
+        for intent in intents:
+            v = variances[intent.intent_id]
+            if v == 0:
+                allocations[intent.intent_id] = 1
+            else:
+                share = v / total_variance
+                allocations[intent.intent_id] = max(1, round(num_questions * share))
+
+    debug(
+        f"[qsim_sample.instantiate_all] total_variance={total_variance:.2f}, allocations_sum={sum(allocations.values())}"
+    )
+
+    instantiated: list[QSimIntent] = []
+
+    for intent in intents:
+        max_variants = allocations[intent.intent_id]
+        for variant_idx in range(max_variants):
+            result = _instantiate_intent(intent, value_domains, variant_idx)
+            if result is not None:
+                instantiated.append(result)
+
+    if len(instantiated) > num_questions:
+        random.shuffle(instantiated)
+        instantiated = instantiated[:num_questions]
+        debug(f"[qsim_sample.instantiate_all] truncated: {len(instantiated)}/{num_questions}")
+    elif len(instantiated) < num_questions:
+        debug(f"[qsim_sample.instantiate_all] limit_reached: {len(instantiated)}/{num_questions}")
+    else:
+        debug(f"[qsim_sample.instantiate_all] created: {len(instantiated)} intents")
+
+    return instantiated