feat: add partial index (WHERE clause) support to Alembic BM25 ops and autogenerate comparator

Author: ankitml

paradedb/sqlalchemy/alembic.py

@@ -33,12 +33,14 @@ def __init__(
         key_field: str,
         *,
         table_schema: str | None = None,
+        where: str | None = None,
     ) -> None:
         self.index_name = index_name
         self.table_name = table_name
         self.expressions = expressions
         self.key_field = key_field
         self.table_schema = table_schema
+        self.where = where
 
     @classmethod
     def create_bm25_index(
@@ -50,6 +52,7 @@ def create_bm25_index(
         *,
         key_field: str,
         table_schema: str | None = None,
+        where: str | None = None,
     ) -> MigrateOperation:
         return operations.invoke(
             cls(
@@ -58,6 +61,7 @@ def create_bm25_index(
                 expressions,
                 key_field,
                 table_schema=table_schema,
+                where=where,
             )
         )
 
@@ -70,6 +74,8 @@ def _create_bm25_index_impl(operations: Operations, operation: CreateBM25IndexOp
         f"ON {_quote_qualified(operation.table_schema, operation.table_name)} "
         f"USING bm25 ({expressions_sql}) WITH (key_field={_quote_literal(operation.key_field)})"
     )
+    if operation.where is not None:
+        sql += f" WHERE {operation.where}"
     operations.execute(sql)
 
 
@@ -83,6 +89,8 @@ def _render_create_bm25_index_op(autogen_context, op: CreateBM25IndexOp) -> str:
     ]
     if op.table_schema is not None:
         parts.append(f"table_schema={op.table_schema!r}")
+    if op.where is not None:
+        parts.append(f"where={op.where!r}")
     return f"op.create_bm25_index({', '.join(parts)})"
 
 
@@ -164,8 +172,8 @@ def _autogen_bm25_meta_indexes(
 
 
 def _autogen_bm25_db_indexes(conn, effective_schemas: set[str]) -> dict[tuple[str, str], dict]:
-    """Return {(schema, index_name): {table_name, expressions, key_field}} from pg_indexes."""
-    from .indexing import _extract_bm25_field_list, _extract_key_field
+    """Return {(schema, index_name): {table_name, expressions, key_field, where}} from pg_indexes."""
+    from .indexing import _extract_bm25_field_list, _extract_key_field, _extract_where_clause
 
     result: dict[tuple[str, str], dict] = {}
     for schema in effective_schemas:
@@ -187,6 +195,7 @@ def _autogen_bm25_db_indexes(conn, effective_schemas: set[str]) -> dict[tuple[st
                 "table_name": row.tablename,
                 "expressions": raw_fields,
                 "key_field": _extract_key_field(row.indexdef) or "",
+                "where": _extract_where_clause(row.indexdef),
             }
     return result
 
@@ -264,6 +273,35 @@ def _normalized_expression_list(expressions: list[str]) -> list[str]:
     return [_normalize_bm25_expression(expr) for expr in expressions]
 
 
+def _normalize_where(clause: str | None) -> str | None:
+    """Normalize a WHERE clause string for comparison.
+
+    Strips whitespace, removes double quotes, and lowercases to reduce
+    false-positive drift detection between PostgreSQL's normalized form
+    and the SQLAlchemy-compiled form.
+    """
+    if clause is None:
+        return None
+    normalized = " ".join(clause.split())
+    normalized = normalized.replace('"', "")
+    return normalized.lower()
+
+
+def _render_where_from_index(index) -> str | None:
+    """Compile the ``postgresql_where`` clause from a SQLAlchemy Index to SQL text."""
+    where_clause = index.dialect_options["postgresql"].get("where")
+    if where_clause is None:
+        return None
+    if isinstance(where_clause, ClauseElement):
+        return str(
+            where_clause.compile(
+                dialect=postgresql.dialect(),
+                compile_kwargs={"literal_binds": True},
+            )
+        )
+    return str(where_clause)
+
+
 def _suppress_standard_bm25_ops(upgrade_ops, bm25_names: set[str]) -> None:
     """Remove any standard Alembic CreateIndexOp/DropIndexOp for BM25 indexes."""
     from alembic.operations.ops import CreateIndexOp, DropIndexOp, ModifyTableOps
@@ -317,38 +355,36 @@ def _compare_bm25_indexes(autogen_context, upgrade_ops, schemas) -> PriorityDisp
             upgrade_ops.ops.append(DropBM25IndexOp(index_name=key[1], if_exists=True, schema=key[0]))
 
     # Emit create ops for indexes present in MetaData but absent from DB.
-    # Also re-create indexes whose expression list or key_field differs from the DB.
+    # Also re-create indexes whose expression list, key_field, or WHERE clause differs from the DB.
     for key, index in meta_bm25.items():
         with_opts = index.dialect_options["postgresql"].get("with") or {}
         key_field = with_opts.get("key_field", "")
         expressions = [
             _strip_relation_qualifiers(_render_bm25_expression(expr), index.table.name)
             for expr in index.expressions
         ]
+        meta_where = _render_where_from_index(index)
+
+        def _make_create_op():
+            return CreateBM25IndexOp(
+                index_name=index.name,
+                table_name=index.table.name,
+                expressions=expressions,
+                key_field=key_field,
+                table_schema=key[0],
+                where=meta_where,
+            )
 
         if key not in db_bm25:
-            upgrade_ops.ops.append(
-                CreateBM25IndexOp(
-                    index_name=index.name,
-                    table_name=index.table.name,
-                    expressions=expressions,
-                    key_field=key_field,
-                    table_schema=key[0],
-                )
-            )
+            upgrade_ops.ops.append(_make_create_op())
         else:
             db = db_bm25[key]
-            if _normalized_expression_list(db["expressions"]) != _normalized_expression_list(expressions) or db["key_field"] != key_field:
+            expressions_changed = _normalized_expression_list(db["expressions"]) != _normalized_expression_list(expressions)
+            key_field_changed = db["key_field"] != key_field
+            where_changed = _normalize_where(db.get("where")) != _normalize_where(meta_where)
+            if expressions_changed or key_field_changed or where_changed:
                 # Index configuration changed: drop the old one, create the new one.
                 upgrade_ops.ops.append(DropBM25IndexOp(index_name=key[1], if_exists=True, schema=key[0]))
-                upgrade_ops.ops.append(
-                    CreateBM25IndexOp(
-                        index_name=index.name,
-                        table_name=index.table.name,
-                        expressions=expressions,
-                        key_field=key_field,
-                        table_schema=key[0],
-                    )
-                )
+                upgrade_ops.ops.append(_make_create_op())
 
     return PriorityDispatchResult.CONTINUE

paradedb/sqlalchemy/indexing.py

@@ -578,6 +578,21 @@ def _extract_key_field(indexdef: str) -> str | None:
     return None
 
 
+_WHERE_CLAUSE_RE = re.compile(r"\bWHERE\s*\((.+)\)\s*$", re.IGNORECASE)
+
+
+def _extract_where_clause(indexdef: str) -> str | None:
+    """Extract the WHERE predicate from a ``pg_indexes.indexdef`` string.
+
+    PostgreSQL normalizes partial index predicates as ``WHERE (<condition>)``.
+    Returns the inner condition text or ``None`` if not present.
+    """
+    match = _WHERE_CLAUSE_RE.search(indexdef)
+    if match:
+        return match.group(1).strip()
+    return None
+
+
 def _extract_alias(index_expr: str) -> str | None:
     match = _ALIAS_RE.search(index_expr)
     if match: