featurizer

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