featurizer

  72class Featurizer:
  73    """
  74    PostgreSQL implementation of the DFS algorithm (adapted for temporal data sets).
  75
  76    Coordinates configuration loading, feature planning, SQL rendering, and optional
  77    database execution.
  78    """
  79
  80    def __init__(
  81        self,
  82        config_file: str,
  83        *,
  84        debug: bool = False,
  85        validate: bool = True,
  86        materialize_threshold: int | None = None,
  87        connection: Any = None,
  88    ) -> None:
  89        """Initialize Featurizer from a YAML configuration file.
  90
  91        Args:
  92            config_file: Path to YAML configuration file
  93            debug: Enable debug logging with icecream. Can also be set via FEATURIZER_DEBUG env var.
  94            validate: Run enhanced validation checks (default: True)
  95            materialize_threshold: Column width above which an oversized
  96                non-target child CTE is materialized into TEMP-table shards
  97                (issue #7). Defaults to PostgreSQL's hard 1664-entry limit; lower
  98                it to materialize earlier (advanced / testing).
  99            connection: Optional psycopg connection used only to resolve a
 100                ``role: categorical`` direct variable's vocabulary from its
 101                PostgreSQL ``ENUM`` labels when no ``vocabulary`` is declared in
 102                the config. When omitted, one is opened from ``DATABASE_URL`` /
 103                ``PG*`` if any categorical actually needs it; a declared
 104                vocabulary needs no database (``query`` / ``--show-sql`` stay
 105                DB-free). The connection is never used to scan data.
 106
 107        Raises:
 108            FileNotFoundError: If config file doesn't exist
 109            ValueError: If config is missing required keys or has invalid values,
 110                or a ``role: categorical`` variable has neither a declared
 111                ``vocabulary`` nor an introspectable PostgreSQL ``ENUM``.
 112        """
 113        config = self._load_config(config_file, validate=validate)
 114        self._materialize_threshold = materialize_threshold
 115
 116        self._debug_enabled: bool = debug or self._env_debug_enabled()
 117        if self._debug_enabled:
 118            ic.configureOutput(prefix="[Featurizer] ", includeContext=True)
 119
 120        self.max_depth: int = config["max_depth"]
 121        self.intervals: List[str] = config["intervals"]
 122        # Point-in-time boundary mode: ``inclusive`` (default, ``<=``) keeps an
 123        # event dated exactly on the as_of_date knowable; ``exclusive`` (``<``)
 124        # treats it as not-yet-knowable. Validated in ConfigValidator.
 125        self.as_of_boundary: AsOfBoundary = config.get(
 126            "as_of_boundary", DEFAULT_BOUNDARY
 127        )
 128
 129        # Paired cohort (issue #10): the column of the caller's as_of_dates
 130        # table that holds the target id. Absent -> the dense cohort, every
 131        # target row under every as-of date, rendered exactly as before.
 132        self.cohort_id_column: Optional[str] = (config.get("as_of_dates") or {}).get(
 133            "id_column"
 134        )
 135
 136        self.graph: ERGraph = ERGraph(
 137            config["entities"],
 138            config["relationships"],
 139            config.get("spatial_relationships"),
 140            config.get("graph_relationships"),
 141        )
 142        self.target: Entity = self._get_entity(config["target"])
 143
 144        # Resolve fixed vocabularies for the target's role: categorical direct
 145        # variables (declared list, else introspected ENUM) before planning so
 146        # the planner stays DB-free. Fails loud if a categorical can be neither.
 147        self._resolve_categorical_vocabularies(connection)
 148
 149        # Primitive selection: config may override the active set; an absent or
 150        # null key applies the curated module defaults. An explicit empty list
 151        # suppresses that layer instead: `aggregations: []` builds zero
 152        # aggregation features, and `transformations: []` passes features
 153        # through unchanged — spelled as identity because the transform CTE
 154        # also feeds child aggregations, so a truly empty transformer set
 155        # would drop every feature column. Unknown names raise in get_* (and
 156        # are caught earlier with suggestions by the validator when
 157        # validate=True).
 158        agg_names = config.get("aggregations")
 159        if agg_names is None:
 160            agg_names = DEFAULT_AGGREGATIONS
 161        tx_names = config.get("transformations")
 162        if tx_names is None:
 163            tx_names = DEFAULT_TRANSFORMATIONS
 164        elif not tx_names:
 165            tx_names = ("identity",)
 166        self.aggregations: AggregationRegistry = get_aggregations(agg_names)
 167        self.transformations: TransformationRegistry = get_transformers(tx_names)
 168
 169        planner = FeaturePlanner(
 170            graph=self.graph,
 171            target_alias=self.target.alias,
 172            max_depth=self.max_depth,
 173            intervals=self.intervals,
 174            aggregations=self.aggregations,
 175            transformations=self.transformations,
 176            boundary=self.as_of_boundary,
 177            debug=self._debug_enabled,
 178            cohort_id_column=self.cohort_id_column,
 179        )
 180        self._plan: PlannerResult = planner.plan()
 181
 182        self.features: Dict[str, Set[Feature]] = {
 183            alias: set(features) for alias, features in self._plan.features.items()
 184        }
 185        self.ctes: List[str] = list(self._plan.ctes)
 186        self.joins: Dict[str, List[str]] = {
 187            alias: list(joins) for alias, joins in self._plan.joins.items()
 188        }
 189
 190        self._renderer: SQLRenderer = SQLRenderer()
 191        self._executor: QueryExecutor = QueryExecutor()
 192
 193    def _resolve_categorical_vocabularies(self, connection: Any) -> None:
 194        """Bake a fixed vocabulary onto each target ``role: categorical`` variable.
 195
 196        Declared vocabularies need no database. If any categorical lacks one, a
 197        connection is required to read its PostgreSQL ``ENUM`` labels: the passed
 198        ``connection`` is used, else one is opened from ``DATABASE_URL`` / ``PG*``;
 199        if neither resolves, :func:`resolve_vocabulary` raises a loud, actionable
 200        error. The data is never scanned for distinct values (split-blind).
 201        """
 202        pending = [
 203            feature
 204            for feature in self.target.features
 205            if isinstance(feature, Variable) and feature.role == ROLE_CATEGORICAL
 206        ]
 207        if not pending:
 208            return
 209
 210        needs_db = any(not var.vocabulary for var in pending)
 211        conn = connection
 212        own_connection: Any = None
 213        if needs_db and conn is None:
 214            own_connection = conn = self._maybe_env_connection()
 215        try:
 216            for var in pending:
 217                var.vocabulary = resolve_vocabulary(var, self.target, conn)
 218        finally:
 219            if own_connection is not None:
 220                own_connection.close()
 221
 222    @staticmethod
 223    def _maybe_env_connection() -> Any:
 224        """A psycopg connection from the environment, or ``None`` if unconfigured.
 225
 226        Mirrors the connection sourcing used by the Arrow/Parquet output paths
 227        (``DATABASE_URL`` / ``PG*``). Returns ``None`` rather than raising when no
 228        database is configured, so the caller can produce the precise
 229        declare-vocabulary-or-ENUM error instead of a generic connection error.
 230        """
 231        from .arrow import default_connection
 232
 233        try:
 234            return default_connection()
 235        except RuntimeError:
 236            return None
 237
 238    # ------------------------------------------------------------------ #
 239    # Public API
 240    # ------------------------------------------------------------------ #
 241
 242    @property
 243    def entities(self) -> Iterable[Entity]:
 244        """Return all entities in the graph."""
 245        return self.graph.entities.values()
 246
 247    @property
 248    def relationships(self) -> List[Any]:
 249        """Return all relationships in the graph."""
 250        return self.graph.relationships
 251
 252    @property
 253    def feature_manifest(self) -> "List[ManifestEntry]":
 254        """Map every output column to its full, untruncated intended name.
 255
 256        One row per output feature column (in output order), each carrying the
 257        rendered ``column`` name, the human-readable ``label`` (recovered even
 258        when the 63-byte identifier cap truncated the column name), a
 259        ``truncated`` flag, the ``kind`` (``one_hot`` | ``variable`` |
 260        ``derived``), the owning ``entity``, and — for one-hot columns — the
 261        ``source_column`` and ``value`` they encode. Useful for human/partner
 262        labels, plot legends, and joining readable names back onto the matrix.
 263        """
 264        from .manifest import build_feature_manifest
 265
 266        return build_feature_manifest(self._plan.target_output_features)
 267
 268    def manifest_dataframe(self) -> "pd.DataFrame":
 269        """The feature manifest as a pandas DataFrame (table / plots / joins)."""
 270        from .manifest import manifest_dataframe
 271
 272        return manifest_dataframe(self.feature_manifest)
 273
 274    def manifest_matching(
 275        self, pattern: str, *, allow_empty: bool = False
 276    ) -> "List[ManifestEntry]":
 277        """Manifest entries whose full ``label`` matches a shell-style glob.
 278
 279        Match against the **label**, never the physical column name: names
 280        longer than PostgreSQL's 63-byte identifier limit are hash-truncated
 281        from the tail, so a pattern targeting an inner fragment
 282        (``*(inspections.kw_*``) misses every truncated column. The label is
 283        the only lossless surface — ``definition`` inherits the parent's
 284        truncation and drops the interval, so windowed variants share one.
 285
 286        Args:
 287            pattern: Glob matched against ``label``. Case-sensitive. ``_`` is a
 288                literal here (unlike SQL ``LIKE`` — see
 289                :func:`~featurizer.manifest.glob_to_like` if you are querying
 290                the persisted manifest table by hand).
 291            allow_empty: Return ``[]`` instead of raising when nothing matches.
 292
 293        Returns:
 294            Matching entries in output order.
 295
 296        Raises:
 297            LookupError: When nothing matches and ``allow_empty`` is False —
 298                a silently empty selection is the failure this helper exists to
 299                prevent. The message suggests near-miss labels.
 300
 301        Example:
 302            >>> f.manifest_matching("*kw_rodent*")[0].column
 303            'CUM_SUM(facilities.MEAN(inspections.ABS(inspections.kw_rod~978dcf98'
 304        """
 305        from .manifest import filter_manifest
 306
 307        return filter_manifest(self.feature_manifest, pattern, allow_empty=allow_empty)
 308
 309    def columns_matching(self, pattern: str, *, allow_empty: bool = False) -> List[str]:
 310        """Physical column names whose full ``label`` matches a glob, in output order.
 311
 312        The projection of :meth:`manifest_matching` onto the rendered column
 313        name — what you put in a ``select`` list, a feature-group definition, or
 314        a DataFrame projection.
 315
 316        Building a triage ``feature_groups.definitions`` entry from physical
 317        names directly will abort the run on every truncated column ("matches no
 318        feature_groups.definitions glob"); resolve it here first::
 319
 320            definitions = {"cardiac": f.columns_matching("*frecuencia_cardiaca*")}
 321
 322        Args:
 323            pattern: Glob matched against ``label``. Case-sensitive.
 324            allow_empty: Return ``[]`` instead of raising when nothing matches.
 325
 326        Returns:
 327            The rendered (possibly truncated) column names, in output order.
 328
 329        Raises:
 330            LookupError: When nothing matches and ``allow_empty`` is False.
 331        """
 332        return [
 333            entry.column
 334            for entry in self.manifest_matching(pattern, allow_empty=allow_empty)
 335        ]
 336
 337    @property
 338    def query(self) -> str:
 339        """Generate the single SQL query for this featurizer configuration.
 340
 341        Raises:
 342            ValueError: If the configuration is too wide to render as one valid
 343                query — the ``<target>_transform`` CTE (or an intermediate CTE)
 344                would exceed PostgreSQL's 1664-entry target-list limit. Use
 345                :attr:`query_groups` / :meth:`to_parquet` / :meth:`to_arrow`
 346                (column-group sharding, issue #7) instead. The matrix is never
 347                silently truncated.
 348        """
 349        sharder = self._sharder()
 350        if not sharder.fits_single_group:
 351            n_groups = sharder.n_groups
 352            raise ValueError(
 353                f"Feature matrix for target '{self.target.alias}' is too wide for a "
 354                "single query: it exceeds PostgreSQL's 1664-entry target-list limit "
 355                f"and partitions into {n_groups} column groups. Use "
 356                "`.query_groups` (group_id -> SQL), `.to_parquet(dir)` (one Parquet "
 357                "per group), or `.to_arrow()` (list of tables); all groups re-join "
 358                "on (as_of_date, id). `.to_arrow`/`.to_parquet`/`.to_dataframe` also "
 359                "materialize any oversized child CTE automatically (issue #7). See "
 360                "docs/adr/0005-column-group-sharding.md."
 361            )
 362        return self._renderer.render(self._plan)
 363
 364    def _sharder(
 365        self, *, max_columns_per_group: int | None = None
 366    ) -> "ColumnGroupSharder":
 367        """A ColumnGroupSharder for this plan honouring ``materialize_threshold``.
 368
 369        ``max_columns_per_group`` overrides the default per-group column cap —
 370        used by :meth:`to_tables` to downshift the partition when the default
 371        cap would produce a heap row over PostgreSQL's ~8160-byte page limit.
 372        """
 373        from .sharding import (
 374            DEFAULT_MAX_COLUMNS_PER_GROUP,
 375            PG_MAX_TARGET_LIST,
 376            ColumnGroupSharder,
 377        )
 378
 379        threshold = (
 380            self._materialize_threshold
 381            if self._materialize_threshold is not None
 382            else PG_MAX_TARGET_LIST
 383        )
 384        return ColumnGroupSharder(
 385            self._plan,
 386            max_columns_per_group=(
 387                max_columns_per_group
 388                if max_columns_per_group is not None
 389                else DEFAULT_MAX_COLUMNS_PER_GROUP
 390            ),
 391            materialize_threshold=threshold,
 392        )
 393
 394    def _grouped(self) -> "GroupedQueries":
 395        """The grouped queries + any temp-table materialization preamble.
 396
 397        A config that fits one valid query short-circuits to a single
 398        ``group_000`` equal to :attr:`query` (no preamble). A wide or
 399        oversized-child config returns the partitioned/rewritten group queries
 400        and, when a child CTE had to be materialized, the preamble on
 401        ``GroupedQueries.materialization`` (issue #7).
 402        """
 403        from collections import OrderedDict
 404
 405        from .sharding import GroupedQueries
 406
 407        sharder = self._sharder()
 408        sharder.warn_oversized()
 409        sharder.warn_plan_size()  # pre-flight PG planner-blowup prediction
 410        if sharder.fits_single_group:
 411            return GroupedQueries(
 412                queries=OrderedDict([("group_000", self._renderer.render(self._plan))]),
 413                key_columns=sharder.key_columns,
 414                fits_single=True,
 415                materialization=None,
 416            )
 417        return sharder.build()
 418
 419    def _grouped_for_tables(
 420        self,
 421    ) -> "Tuple[GroupedQueries, OrderedDict[str, List[str]]]":
 422        """Grouped queries for :meth:`to_tables` + the matching column→group map.
 423
 424        The CTAS path differs from the SELECT/fetch paths in one way: each
 425        group becomes a heap *table*, and a heap tuple must fit one 8 KiB page
 426        (~8160 bytes). A group that is a perfectly valid query (≤1664 output
 427        columns) can still fail ``create table … as`` with ``row is too big``
 428        once ~1000+ mostly fixed-width feature columns land on a page.
 429        Pre-flight: estimate every group's row width
 430        (:func:`~featurizer.sharding.estimate_heap_row_width` — 8 bytes per
 431        column + header + null bitmap); when any group exceeds
 432        :data:`~featurizer.sharding.HEAP_ROW_BUDGET_BYTES`, re-partition with a
 433        per-group cap sized to the budget instead of letting PostgreSQL fail.
 434
 435        The returned mapping is the exact partition the queries were rendered
 436        from, so the manifest's ``feature_group`` tags always name a
 437        really-persisted table (see :meth:`_write_manifest_table`).
 438        """
 439        from collections import OrderedDict
 440
 441        from .sharding import (
 442            HEAP_ROW_BUDGET_BYTES,
 443            GroupedQueries,
 444            estimate_heap_row_width,
 445            max_heap_safe_columns,
 446        )
 447
 448        sharder = self._sharder()
 449        sharder.warn_oversized()
 450        sharder.warn_plan_size()
 451        n_keys = len(sharder.key_columns)
 452        groups = sharder.column_groups()
 453        oversized = {
 454            gid: estimate_heap_row_width(len(cols) + n_keys)
 455            for gid, cols in groups.items()
 456            if estimate_heap_row_width(len(cols) + n_keys) > HEAP_ROW_BUDGET_BYTES
 457        }
 458        if oversized:
 459            cap = max_heap_safe_columns(n_keys)
 460            worst_gid, worst_width = max(oversized.items(), key=lambda kv: kv[1])
 461            logger.warning(
 462                "to_tables heap-row pre-flight: {} of {} column group(s) "
 463                "estimate over the {}-byte page budget (worst: {} ≈ {} bytes; "
 464                "PostgreSQL rejects heap rows over ~8160 bytes with 'row is "
 465                "too big'). Re-partitioning with max {} feature columns per "
 466                "group — more, narrower tables; SELECT/fetch paths are "
 467                "unaffected.",
 468                len(oversized),
 469                len(groups),
 470                HEAP_ROW_BUDGET_BYTES,
 471                worst_gid,
 472                worst_width,
 473                cap,
 474            )
 475            sharder = self._sharder(max_columns_per_group=cap)
 476            groups = sharder.column_groups()
 477
 478        if sharder.fits_single_group and len(groups) == 1:
 479            # Same short-circuit as _grouped(): one group, the plain query.
 480            grouped = GroupedQueries(
 481                queries=OrderedDict([("group_000", self._renderer.render(self._plan))]),
 482                key_columns=sharder.key_columns,
 483                fits_single=True,
 484                materialization=None,
 485            )
 486        else:
 487            # NOTE: also taken when the config fits a single *query* but
 488            # partitions into >1 groups (e.g. >cap output columns under 1664,
 489            # or the heap downshift above) — the persisted tables must follow
 490            # the partition or the manifest's feature_group tags would name
 491            # tables that were never written.
 492            grouped = sharder.build()
 493        return grouped, groups
 494
 495    @property
 496    def query_groups(self) -> "OrderedDict[str, str]":
 497        """SQL for each column group: ``group_<NNN>`` -> self-contained query.
 498
 499        Partitions the (possibly very wide) feature matrix into ordered column
 500        groups, each a valid query whose every CTE tuple is under PostgreSQL's
 501        1664-entry limit (issue #7). A config that fits in one query returns a
 502        single ``group_000`` entry equal to :attr:`query`. Every group leads
 503        with ``(as_of_date, <target id>)`` so the groups re-join into the full
 504        matrix.
 505
 506        When an oversized non-target child CTE had to be materialized, these
 507        queries reference TEMP-table shards and **presuppose**
 508        :attr:`materialization_ddl` was executed first on the same session;
 509        :meth:`to_arrow` / :meth:`to_parquet` do that automatically.
 510        """
 511        return self._grouped().queries
 512
 513    @property
 514    def materialization_ddl(self) -> List[str]:
 515        """The ``CREATE TEMP TABLE`` preamble (issue #7) that :attr:`query_groups`
 516        presupposes, or ``[]`` when no oversized child CTE needs materializing.
 517
 518        Run these statements on the same connection/session before executing the
 519        grouped queries. :meth:`to_arrow` / :meth:`to_parquet` / :meth:`to_dataframe`
 520        run them for you.
 521        """
 522        mplan = self._grouped().materialization
 523        return list(mplan.ddl) if mplan is not None else []
 524
 525    def to_dataframe(
 526        self, *, connection: Any = None, impute: bool = False, **impute_kwargs: Any
 527    ) -> pd.DataFrame:
 528        """Execute the query and return results as a DataFrame.
 529
 530        Args:
 531            connection: An open psycopg connection to run on (required when the
 532                query references session ``TEMP`` tables — the integration
 533                harness). When ``None``, the fast single-query path uses
 534                ``records`` and the grouped/materialized path builds its own
 535                connection from the environment and closes it afterwards.
 536            impute: When True, run the opt-in imputation pass (count-like
 537                features → 0, measures left NULL unless ``measure_strategy`` is
 538                given, with ``<feature>__missing`` indicator columns). The
 539                default keeps the raw NULLs, since missingness is signal.
 540            **impute_kwargs: Forwarded to
 541                :func:`featurizer.imputation.impute_features`.
 542
 543        A config that fits one valid query uses the fast (single-query) path; a
 544        wide or oversized-child config (issue #7) runs the column-group queries —
 545        and any TEMP-table materialization preamble — on one connection and
 546        re-joins them on ``(as_of_date, target_id)`` into the same indexed frame.
 547        Passing ``connection`` forces the one-connection path (so it can see
 548        session TEMP tables) regardless of width.
 549
 550        Returns:
 551            DataFrame indexed by ['as_of_date', target_id]
 552
 553        Raises:
 554            ValueError: If target entity doesn't define a primary ID
 555        """
 556        if self.target.id is None:
 557            raise ValueError(
 558                f"Target entity '{self.target.alias}' does not define a primary id."
 559            )
 560        grouped = self._grouped()
 561        if grouped.fits_single and connection is None:
 562            df = self._executor.to_dataframe(self.query, self.target.id.name)
 563        else:
 564            df = self._executor.to_dataframe_materialized(
 565                preamble_ddl=(
 566                    grouped.materialization.ddl
 567                    if grouped.materialization is not None
 568                    else []
 569                ),
 570                group_queries=grouped.queries,
 571                target_id=self.target.id.name,
 572                connection=connection,
 573                key_columns=grouped.key_columns,
 574            )
 575        if impute:
 576            from .imputation import guard_full_matrix_fit, impute_features
 577
 578            # Engine path: this fits over the whole returned matrix, so gate the
 579            # leaky measure strategies (ADR-0001). The pure impute_features helper
 580            # stays ungated for callers that pre-split their own data.
 581            guard_full_matrix_fit(
 582                impute_kwargs.get("measure_strategy", "none"),
 583                allow_full_matrix_fit=bool(
 584                    impute_kwargs.pop("allow_full_matrix_fit", False)
 585                ),
 586                caller="to_dataframe",
 587            )
 588            df = impute_features(df, **impute_kwargs)
 589        return df
 590
 591    def to_arrow(
 592        self,
 593        *,
 594        connection: Any = None,
 595        numeric_as_float: bool = True,
 596        impute: bool = False,
 597        **impute_kwargs: Any,
 598    ) -> "Any":
 599        """Execute the query and return Arrow output, no pandas hop.
 600
 601        Streams the result out of PostgreSQL with binary ``COPY`` and decodes it
 602        column-by-column into Arrow, so SQL NULLs are preserved as Arrow nulls
 603        (never coerced to ``NaN``) and the full result set never materializes as
 604        a pandas frame. ``as_of_date`` and the target id are ordinary leading
 605        columns (no index), unlike :meth:`to_dataframe`.
 606
 607        Sharding (issue #7): when the matrix fits in one query a single
 608        :class:`pyarrow.Table` is returned. When it is too wide for a single
 609        valid query (over PostgreSQL's 1664-entry target-list limit), an
 610        ``OrderedDict[str, pyarrow.Table]`` of column groups is returned instead
 611        — ``group_<NNN>`` -> table. Every group table leads with
 612        ``(as_of_date, <target id>)`` and the groups re-join on those keys to
 613        reconstruct the full matrix.
 614
 615        Args:
 616            connection: An open psycopg connection to run ``COPY`` on. Required
 617                when the rendered query references session ``TEMP`` tables (the
 618                integration harness). When ``None``, a connection is built from
 619                ``DATABASE_URL`` / ``PG*`` and closed afterwards. A single
 620                connection is reused across all groups.
 621            numeric_as_float: Cast PostgreSQL ``numeric`` aggregate columns to
 622                ``float64`` (ML-ready, ``to_dataframe``-comparable). Set ``False``
 623                to keep exact ``decimal128``. NULLs are preserved either way.
 624            impute: When True, apply the Arrow-native imputation pass
 625                (:func:`featurizer.imputation.impute_arrow`) to each group:
 626                count-like features → 0, measures left null unless
 627                ``measure_strategy`` is given, with stable ``<feature>__missing``
 628                indicator columns. ``as_of_date`` and the target id are passed as
 629                ``key_columns`` and left untouched.
 630            **impute_kwargs: Forwarded to ``impute_arrow``. ``measure_strategy`` in
 631                ``{"mean","median"}`` additionally requires
 632                ``allow_full_matrix_fit=True`` (ADR-0001 leakage gate).
 633
 634        Returns:
 635            A ``pyarrow.Table`` for a single-group config, otherwise an
 636            ``OrderedDict[str, pyarrow.Table]`` keyed by group id.
 637
 638        Raises:
 639            ImportError: If pyarrow (the ``[parquet]`` extra) is not installed.
 640            ValueError: If the target entity does not define a primary id, or a
 641                leaky measure strategy is requested without the opt-in.
 642        """
 643        if self.target.id is None:
 644            raise ValueError(
 645                f"Target entity '{self.target.alias}' does not define a primary id."
 646            )
 647        groups = self._arrow_groups(
 648            connection=connection,
 649            numeric_as_float=numeric_as_float,
 650            impute=impute,
 651            **impute_kwargs,
 652        )
 653        if len(groups) == 1:
 654            # Single-group config: preserve the original single-Table contract.
 655            return next(iter(groups.values()))
 656        return groups
 657
 658    def to_parquet(
 659        self,
 660        path: str,
 661        *,
 662        connection: Any = None,
 663        numeric_as_float: bool = True,
 664        impute: bool = False,
 665        **impute_kwargs: Any,
 666    ) -> None:
 667        """Execute the query and write the result to Parquet.
 668
 669        Thin wrapper over the Arrow path plus ``pyarrow.parquet.write_table``;
 670        all arguments (including the imputation contract and its ADR-0001 leakage
 671        gate) behave exactly as in :meth:`to_arrow`. NULLs are written as Parquet
 672        nulls.
 673
 674        Sharding (issue #7): when the matrix fits in one query, a single Parquet
 675        file is written at ``path``. When it is too wide for one valid query,
 676        ``path`` is treated as a **directory** and one Parquet file per column
 677        group is written under it as ``group_<NNN>.parquet``. All group files
 678        re-join on ``(as_of_date, <target id>)`` to reconstruct the full matrix.
 679
 680        Args:
 681            path: Destination ``.parquet`` file (single group) or output
 682                directory (multiple groups; created if absent).
 683            connection: See :meth:`to_arrow`.
 684            numeric_as_float: See :meth:`to_arrow`.
 685            impute: See :meth:`to_arrow`.
 686            **impute_kwargs: See :meth:`to_arrow`.
 687
 688        Raises:
 689            ImportError: If pyarrow (the ``[parquet]`` extra) is not installed.
 690        """
 691        import pyarrow.parquet as pq  # pyright: ignore[reportMissingImports]
 692
 693        groups = self._arrow_groups(
 694            connection=connection,
 695            numeric_as_float=numeric_as_float,
 696            impute=impute,
 697            **impute_kwargs,
 698        )
 699        if len(groups) == 1:
 700            pq.write_table(next(iter(groups.values())), path)
 701            return
 702
 703        import os as _os
 704
 705        _os.makedirs(path, exist_ok=True)
 706        for gid, table in groups.items():
 707            pq.write_table(table, _os.path.join(path, f"{gid}.parquet"))
 708        logger.info(
 709            "Wrote {} column-group Parquet files to {} (re-join on {}).",
 710            len(groups),
 711            path,
 712            ("as_of_date", self.target.id.name) if self.target.id else "(as_of_date,)",
 713        )
 714
 715    def to_tables(
 716        self,
 717        schema: str,
 718        *,
 719        connection: Any = None,
 720        table_prefix: str | None = None,
 721        create_schema: bool = True,
 722    ) -> List["FeatureGroupTable"]:
 723        """Persist the feature matrix as triage-style feature-group tables.
 724
 725        Writes each column group as a persistent table
 726        ``"<schema>"."<stem>_group_<NNN>"`` keyed on ``(as_of_date, <target id>)``,
 727        the feature-group contract triage-pg consumes (issue #7). A config that
 728        fits one query writes a single ``<stem>_group_000``; a wide or
 729        oversized-child config writes one table per column group, all re-joinable
 730        on the keys. The issue-#7 intermediate shards stay ephemeral ``TEMP``
 731        tables — only the final groups are persisted.
 732
 733        Idempotent: each target table is ``DROP TABLE IF EXISTS`` + ``CREATE TABLE
 734        … AS`` so a re-run replaces it cleanly.
 735
 736        Alongside the group tables, the feature manifest is persisted as
 737        ``"<schema>"."<stem>_manifest"`` — one row per output column (label,
 738        lineage, generated description, and the ``feature_group`` it landed
 739        in), joinable to the group tables by column name. The returned list
 740        contains the feature-group tables only, as before.
 741
 742        Args:
 743            schema: Destination schema (created if absent unless
 744                ``create_schema=False``).
 745            connection: An open psycopg connection to write on. When supplied the
 746                caller owns the transaction (nothing is committed here — the
 747                integration harness verifies within its rolled-back transaction);
 748                when ``None`` a connection is built from the environment, committed
 749                so the tables persist, and closed.
 750            table_prefix: Table-name stem; defaults to the target alias
 751                (``stores`` -> ``stores_group_000``).
 752            create_schema: Run ``CREATE SCHEMA IF NOT EXISTS`` first.
 753
 754        Returns:
 755            The ordered manifest of created :class:`FeatureGroupTable`s.
 756
 757        Raises:
 758            ValueError: If the target entity does not define a primary id.
 759        """
 760        if self.target.id is None:
 761            raise ValueError(
 762                f"Target entity '{self.target.alias}' does not define a primary id."
 763            )
 764        from .arrow import default_connection
 765        from .sharding import FeatureGroupTable
 766
 767        grouped, column_groups = self._grouped_for_tables()
 768        preamble = (
 769            grouped.materialization.ddl if grouped.materialization is not None else []
 770        )
 771        stem = table_prefix or self.target.alias
 772        keys = list(grouped.key_columns)
 773
 774        own_connection = connection is None
 775        conn = connection if connection is not None else default_connection()
 776        tables: List["FeatureGroupTable"] = []
 777        try:
 778            # jit off from the preamble on (issue #53); a caller's value is put
 779            # back before their transaction goes on.
 780            with jit_disabled(conn):
 781                with conn.cursor() as cur:
 782                    if create_schema:
 783                        cur.execute(f'create schema if not exists "{schema}"')
 784                    for ddl in preamble:
 785                        cur.execute(ddl)
 786                analyze_as_of_dates(conn)  # planner-stats optimization (see executor)
 787                if own_connection:  # PLANNER_TUNING never reaches a caller's
 788                    apply_planner_tuning(conn)
 789                with conn.cursor() as cur:
 790                    for gid, sql in grouped.queries.items():
 791                        name = f'"{schema}"."{stem}_{gid}"'
 792                        cur.execute(f"drop table if exists {name}")
 793                        cur.execute(f"create table {name} as\n{sql}")
 794                        tables.append(
 795                            FeatureGroupTable(
 796                                name=name, group=gid, key_columns=list(keys)
 797                            )
 798                        )
 799                    self._write_manifest_table(cur, schema, stem, column_groups)
 800            if own_connection:
 801                conn.commit()
 802        finally:
 803            if own_connection:
 804                conn.close()
 805        logger.info(
 806            "Persisted {} feature-group table(s) + manifest to schema {!r} "
 807            "(re-join on {}).",
 808            len(tables),
 809            schema,
 810            tuple(keys),
 811        )
 812        return tables
 813
 814    def _write_manifest_table(
 815        self,
 816        cur: Any,
 817        schema: str,
 818        stem: str,
 819        column_groups: "OrderedDict[str, List[str]]",
 820    ) -> None:
 821        """Persist the feature manifest as ``"<schema>"."<stem>_manifest"``.
 822
 823        One row per output feature column, including which feature-group table
 824        the column landed in (``feature_group``, joinable back to the
 825        ``<stem>_group_<NNN>`` tables by column name). Same contracts as the
 826        group tables: idempotent DROP+CREATE, and the caller owns the
 827        transaction. Values are inserted parameterized — labels and definitions
 828        contain quotes and arbitrary SQL text.
 829
 830        ``column_groups`` is the exact partition the group tables were written
 831        from (:meth:`_grouped_for_tables`) — never a freshly computed one, so
 832        the ``feature_group`` tags cannot drift from the persisted tables. A
 833        manifest column absent from the partition raises rather than being
 834        silently mis-tagged: triage consumes this table for lineage, and a
 835        wrong ``feature_group`` corrupts it downstream.
 836        """
 837        column_to_group: Dict[str, str] = {}
 838        for gid, names in column_groups.items():
 839            for column_name in names:
 840                column_to_group[column_name.replace('"', "")] = gid
 841
 842        orphaned = [
 843            entry.column
 844            for entry in self.feature_manifest
 845            if entry.column not in column_to_group
 846        ]
 847        if orphaned:
 848            raise RuntimeError(
 849                f"Feature manifest column(s) map to no column group: "
 850                f"{orphaned[:5]!r}{' …' if len(orphaned) > 5 else ''} "
 851                f"(available groups: {', '.join(column_groups) or '(none)'}). "
 852                "Persisting would mis-tag feature_group in "
 853                f'"{schema}"."{stem}_manifest" and corrupt lineage for '
 854                "downstream consumers (triage joins on it). This is a "
 855                "featurizer bug — the manifest and the column partition must "
 856                "describe the same plan; please report it with your config."
 857            )
 858
 859        name = f'"{schema}"."{stem}_manifest"'
 860        cur.execute(f"drop table if exists {name}")
 861        cur.execute(f"""
 862            create table {name} (
 863                "column_name"   text not null,
 864                "label"         text not null,
 865                "truncated"     boolean not null,
 866                "kind"          text not null,
 867                "entity"        text,
 868                "source_alias"  text,
 869                "depth"         integer not null,
 870                "parents"       text[] not null,
 871                "interval"      text,
 872                "source_column" text,
 873                "value"         text,
 874                "description"   text not null,
 875                "definition"    text,
 876                "feature_group" text not null
 877            )
 878            """)
 879        rows = [
 880            (
 881                entry.column,
 882                entry.label,
 883                entry.truncated,
 884                entry.kind,
 885                entry.entity,
 886                entry.source_alias,
 887                entry.depth,
 888                entry.parents,
 889                entry.interval,
 890                entry.source_column,
 891                entry.value,
 892                entry.description,
 893                entry.definition,
 894                column_to_group[entry.column],
 895            )
 896            for entry in self.feature_manifest
 897        ]
 898        if rows:
 899            cur.executemany(
 900                f"insert into {name} values "
 901                "(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)",
 902                rows,
 903            )
 904
 905    def _arrow_groups(
 906        self,
 907        *,
 908        connection: Any = None,
 909        numeric_as_float: bool = True,
 910        impute: bool = False,
 911        **impute_kwargs: Any,
 912    ) -> "OrderedDict[str, Any]":
 913        """Run every column group through the Arrow exporter on one connection.
 914
 915        Always returns an ``OrderedDict`` (a single-group config yields one
 916        entry). Imputation, when requested, runs per group with the same
 917        ADR-0001 leakage gate as the single-query path; the gate is checked once
 918        up front so a leaky strategy is refused before any query runs.
 919        """
 920        from .arrow import ArrowExporter
 921
 922        if impute:
 923            from .imputation import guard_full_matrix_fit
 924
 925            guard_full_matrix_fit(
 926                impute_kwargs.get("measure_strategy", "none"),
 927                allow_full_matrix_fit=bool(impute_kwargs.get("allow_full_matrix_fit")),
 928                caller="to_arrow",
 929            )
 930
 931        # Grouped queries + any TEMP-table materialization preamble (issue #7).
 932        # The single-group short-circuit + oversized warning live in _grouped().
 933        grouped = self._grouped()
 934        preamble = (
 935            grouped.materialization.ddl if grouped.materialization is not None else []
 936        )
 937
 938        exporter = ArrowExporter()
 939        own_connection = connection is None
 940        conn = connection if connection is not None else exporter.open_connection()
 941        try:
 942            from collections import OrderedDict as _OrderedDict
 943
 944            # Run the CREATE TEMP TABLE preamble first, on the *same* connection,
 945            # so the group queries' shard references resolve. The connection is
 946            # non-autocommit (default_connection / the harness), so ON COMMIT DROP
 947            # shards live for the whole transaction and drop when it closes.
 948            tables: "OrderedDict[str, Any]" = _OrderedDict()
 949            # jit off from the preamble on (issue #53); a caller's value is put
 950            # back before their transaction goes on.
 951            with jit_disabled(conn):
 952                if preamble:
 953                    with conn.cursor() as cur:
 954                        for ddl in preamble:
 955                            cur.execute(ddl)
 956                analyze_as_of_dates(conn)  # planner-stats optimization (see executor)
 957                if own_connection:  # PLANNER_TUNING never reaches a caller's
 958                    apply_planner_tuning(conn)
 959
 960                for gid, sql in grouped.queries.items():
 961                    table = exporter.to_arrow(
 962                        sql, connection=conn, numeric_as_float=numeric_as_float
 963                    )
 964                    if impute:
 965                        table = self._impute_group(table, **impute_kwargs)
 966                    tables[gid] = table
 967            return tables
 968        finally:
 969            if own_connection:
 970                conn.close()
 971
 972    def _impute_group(self, table: Any, **impute_kwargs: Any) -> Any:
 973        """Apply the Arrow imputation pass to one group's table.
 974
 975        The leakage gate is already checked in :meth:`_arrow_groups`; strip the
 976        consumed ``allow_full_matrix_fit`` flag so ``impute_arrow`` does not see
 977        an unexpected keyword.
 978        """
 979        from .imputation import impute_arrow
 980
 981        kwargs = dict(impute_kwargs)
 982        kwargs.pop("allow_full_matrix_fit", None)
 983        assert self.target.id is not None  # guarded by callers
 984        return impute_arrow(
 985            table,
 986            key_columns=("as_of_date", self.target.id.name),
 987            **kwargs,
 988        )
 989
 990    # ------------------------------------------------------------------ #
 991    # Internal helpers
 992    # ------------------------------------------------------------------ #
 993
 994    def _get_entity(self, alias: str) -> Entity:
 995        """Get an entity by alias from the graph.
 996
 997        Args:
 998            alias: Entity alias to look up
 999
1000        Returns:
1001            Entity with the given alias
1002
1003        Raises:
1004            ValueError: If entity with alias doesn't exist
1005        """
1006        entity = self.graph.entities.get(alias)
1007        if entity is None:
1008            raise ValueError(f"Unknown target entity alias '{alias}'.")
1009        return entity
1010
1011    @staticmethod
1012    def _env_debug_enabled() -> bool:
1013        """Check if debug mode is enabled via environment variable."""
1014        value = os.getenv("FEATURIZER_DEBUG", "")
1015        return value.lower() in {"1", "true", "yes", "on"}
1016
1017    @staticmethod
1018    def _load_config(config_file: str, validate: bool = True) -> Dict[str, Any]:
1019        """Load and validate configuration from YAML file.
1020
1021        Args:
1022            config_file: Path to YAML configuration file
1023            validate: Run enhanced validation checks
1024
1025        Returns:
1026            Validated configuration dictionary
1027
1028        Raises:
1029            FileNotFoundError: If config file doesn't exist
1030            ValueError: If config is invalid or missing required keys
1031        """
1032        try:
1033            with open(config_file) as f:
1034                config = yaml.safe_load(f) or {}
1035        except FileNotFoundError as exc:
1036            raise FileNotFoundError(f"Config file not found: {config_file}") from exc
1037        except yaml.YAMLError as exc:
1038            raise ValueError(f"Invalid YAML in config file {config_file}") from exc
1039
1040        # Run enhanced validation if enabled
1041        if validate:
1042            validator = ConfigValidator(mode="strict")
1043            result = validator.validate(config)
1044
1045            if not result.is_valid:
1046                raise ValueError(
1047                    f"Configuration validation failed:\n{result.format_errors()}"
1048                )
1049
1050            # Log warnings
1051            for warning in result.warnings:
1052                location = f"[{warning.location}] " if warning.location else ""
1053                logger.warning(f"{location}{warning.message}")
1054
1055        # Backwards compatibility: Basic validation
1056        required_keys = {"target", "max_depth", "intervals", "entities"}
1057        missing = [key for key in required_keys if key not in config]
1058        if missing:
1059            raise ValueError(f"Config missing required keys: {', '.join(missing)}")
1060
1061        if not isinstance(config["target"], str) or not config["target"].strip():
1062            raise ValueError("'target' must be a non-empty string.")
1063
1064        if not isinstance(config["max_depth"], int) or config["max_depth"] < 1:
1065            raise ValueError("'max_depth' must be a positive integer.")
1066
1067        if not isinstance(config["entities"], list) or not config["entities"]:
1068            raise ValueError("Config must declare at least one entity in 'entities'.")
1069
1070        if not isinstance(config["intervals"], list):
1071            raise ValueError("'intervals' must be a list of interval strings.")
1072
1073        relationships = config.get("relationships")
1074        if relationships is None:
1075            logger.debug(
1076                "No relationships defined in config; defaulting to empty list."
1077            )
1078            config["relationships"] = []
1079        elif not isinstance(relationships, list):
1080            raise ValueError("'relationships' must be a list when provided.")
1081
1082        return config