sync-indexes: create expression statistics alongside identifier indexes

ingestify/infra/store/dataset/sqlalchemy/repository.py

@@ -165,6 +165,15 @@ def create_identifier_indexes(self, index_configs: list[dict]):
         The WHERE clause limits the index to a single (provider, dataset_type) pair,
         making it smaller and ensuring neither column is a post-scan filter.
 
+        Alongside each index it also creates table-level *expression statistics* on the
+        same JSONB expression(s) (PostgreSQL 14+). The index lets the planner FIND rows,
+        but without statistics on a JSONB expression it mis-estimates the expression's
+        selectivity (it defaults to a poor n_distinct). For batch identifier lookups that
+        makes it prefer a full-partition hash join over a per-key index nested loop, so a
+        large table stays slow despite the index. The expression statistics give the
+        planner the real cardinality so it chooses the index plan on its own. Statistics
+        only take effect after ANALYZE, which this runs once at the end.
+
         Call this explicitly (e.g. via `ingestify sync-indexes`) when datasets
         have high-cardinality identifiers that are queried frequently.
         """
@@ -174,6 +183,10 @@ def create_identifier_indexes(self, index_configs: list[dict]):
 
         table_name = f"{self.table_prefix}dataset"
         with self.engine.connect() as conn:
+            # Expression statistics (CREATE STATISTICS ... ON (<expr>)) need PostgreSQL 14+.
+            create_stats = (
+                int(conn.execute(text("SHOW server_version_num")).scalar()) >= 140000
+            )
             for config in index_configs:
                 name = config["name"]
                 provider = config["provider"]
@@ -194,8 +207,30 @@ def create_identifier_indexes(self, index_configs: list[dict]):
                     )
                 )
                 logger.info("Created index %s on keys: %s", index_name, keys)
+
+                if create_stats:
+                    stats_name = f"{self.table_prefix}stat_dataset_identifier_{name}"
+                    conn.execute(
+                        text(
+                            f"CREATE STATISTICS IF NOT EXISTS {stats_name} "
+                            f"ON {expressions} FROM {table_name}"
+                        )
+                    )
+                    logger.info("Created statistics %s", stats_name)
             conn.commit()
 
+        # Statistics are inert until ANALYZE populates them (and ANALYZE cannot run inside
+        # the transaction above). Run it once so the new statistics take effect now instead
+        # of waiting for autovacuum.
+        if create_stats and index_configs:
+            with self.engine.connect().execution_options(
+                isolation_level="AUTOCOMMIT"
+            ) as conn:
+                logger.info(
+                    "Analyzing %s to populate identifier statistics", table_name
+                )
+                conn.execute(text(f"ANALYZE {table_name}"))
+
     def drop_all_tables(self):
         """Drop all tables in the database. Useful for test cleanup."""
         if hasattr(self, "metadata") and hasattr(self, "engine"):

ingestify/tests/test_identifier_indexes.py

@@ -35,7 +35,7 @@ def repository(ingestify_test_database_url, db_cleanup):
     provider = SqlAlchemySessionProvider(ingestify_test_database_url)
     repo = SqlAlchemyDatasetRepository(provider)
     yield repo
-    # Drop test indexes so they don't leak between runs
+    # Drop test indexes + statistics so they don't leak between runs
     if provider.engine.dialect.name == "postgresql":
         with provider.engine.connect() as conn:
             for config in INDEX_CONFIGS:
@@ -44,6 +44,11 @@ def repository(ingestify_test_database_url, db_cleanup):
                         f"DROP INDEX IF EXISTS idx_dataset_identifier_{config['name']}"
                     )
                 )
+                conn.execute(
+                    sqlalchemy.text(
+                        f"DROP STATISTICS IF EXISTS stat_dataset_identifier_{config['name']}"
+                    )
+                )
             conn.commit()
     provider.drop_all_tables()
 
@@ -73,6 +78,33 @@ def test_create_identifier_indexes_creates_indexes(repository):
     assert "idx_dataset_identifier_test_keyword_set" in index_names
 
 
+def test_create_identifier_indexes_creates_statistics(repository):
+    """On PostgreSQL 14+, expression statistics are created alongside each index so the
+    planner estimates the JSONB-expression selectivity correctly (otherwise it mis-costs a
+    batch identifier lookup and full-scans the partition instead of using the index)."""
+    provider = repository.session_provider
+    if provider.engine.dialect.name != "postgresql":
+        pytest.skip("Expression statistics require PostgreSQL")
+    with provider.engine.connect() as conn:
+        version = int(conn.execute(sqlalchemy.text("SHOW server_version_num")).scalar())
+    if version < 140000:
+        pytest.skip("Expression statistics require PostgreSQL 14+")
+
+    repository.create_identifier_indexes(INDEX_CONFIGS)
+
+    with provider.engine.connect() as conn:
+        result = conn.execute(
+            sqlalchemy.text(
+                "SELECT stxname FROM pg_statistic_ext "
+                "WHERE stxname LIKE 'stat_dataset_identifier_%'"
+            )
+        )
+        stat_names = {row[0] for row in result}
+
+    assert "stat_dataset_identifier_test_keyword_metrics" in stat_names
+    assert "stat_dataset_identifier_test_keyword_set" in stat_names
+
+
 def test_create_identifier_indexes_idempotent(repository):
     """Running twice does not raise (IF NOT EXISTS)."""
     repository.create_identifier_indexes(INDEX_CONFIGS)