comments and hogql

Author: oliverb123

posthog/clickhouse/migrations/0182_shard_document_embeddings.py

@@ -4,16 +4,21 @@
 from products.error_tracking.backend.embedding import (
     DISTRIBUTED_DOCUMENT_EMBEDDINGS_TABLE_SQL,
     DOCUMENT_EMBEDDING_WRITABLE,
-    DOCUMENT_EMBEDDINGS_DATA_TABLE_SQL,
     DOCUMENT_EMBEDDINGS_MV,
     DOCUMENT_EMBEDDINGS_MV_SQL,
+    DOCUMENT_EMBEDDINGS_TABLE_SQL,
     DOCUMENT_EMBEDDINGS_WRITABLE_TABLE_SQL,
     KAFKA_DOCUMENT_EMBEDDINGS,
     KAFKA_DOCUMENT_EMBEDDINGS_TABLE_SQL,
 )
 
+# The only tricky part of this migration is that the behaviour of `DOCUMENT_EMBEDDINGS_TABLE_SQL` has changed -
+# it now creates sharded tables rather than replicated tables. We don't drop the replicated table, as in production
+# we want to keep the historical data around, at least until we get around to migrating it, but this does mean
+# all future "rebuild the world" runs of the migration set will never create that old table, only the new sharded one.
+
 operations = [
-    # 1. Drop MV
+    # 1. Drop MV to stop processing messages from kafka
     run_sql_with_exceptions(
         f"DROP TABLE IF EXISTS {DOCUMENT_EMBEDDINGS_MV}",
         node_roles=[NodeRole.INGESTION_SMALL],
@@ -23,14 +28,14 @@
         f"DROP TABLE IF EXISTS {KAFKA_DOCUMENT_EMBEDDINGS}",
         node_roles=[NodeRole.INGESTION_SMALL],
     ),
-    # 3. Drop old writable table
+    # 3. Drop old writable table (but not the old "main" table, since we want to keep the data around)
     run_sql_with_exceptions(
         f"DROP TABLE IF EXISTS {DOCUMENT_EMBEDDING_WRITABLE}",
         node_roles=[NodeRole.INGESTION_SMALL],
     ),
-    # 4. Create new sharded data table
+    # 4. Create new sharded data tables (this function used to create "posthog_document_embeddings" directly, but now creates the sharded_ version)
     run_sql_with_exceptions(
-        DOCUMENT_EMBEDDINGS_DATA_TABLE_SQL(),
+        DOCUMENT_EMBEDDINGS_TABLE_SQL(),
         node_roles=[NodeRole.DATA],
         sharded=True,
     ),
@@ -39,7 +44,7 @@
         DISTRIBUTED_DOCUMENT_EMBEDDINGS_TABLE_SQL(),
         node_roles=[NodeRole.DATA, NodeRole.COORDINATOR],
     ),
-    # 6. Create new writable distributed table pointing to sharded table
+    # 6. Create new writable distributed table pointing to sharded tables
     run_sql_with_exceptions(
         DOCUMENT_EMBEDDINGS_WRITABLE_TABLE_SQL(),
         node_roles=[NodeRole.INGESTION_SMALL],
@@ -49,7 +54,7 @@
         KAFKA_DOCUMENT_EMBEDDINGS_TABLE_SQL(),
         node_roles=[NodeRole.INGESTION_SMALL],
     ),
-    # 8. Recreate MV writing to writable table (which now writes to sharded table)
+    # 8. Recreate MV writing to writable table (which now writes to sharded tables)
     run_sql_with_exceptions(
         DOCUMENT_EMBEDDINGS_MV_SQL(),
         node_roles=[NodeRole.INGESTION_SMALL],

posthog/clickhouse/schema.py

@@ -172,8 +172,8 @@
 
 from products.error_tracking.backend.embedding import (
     DISTRIBUTED_DOCUMENT_EMBEDDINGS_TABLE_SQL,
-    DOCUMENT_EMBEDDINGS_DATA_TABLE_SQL,
     DOCUMENT_EMBEDDINGS_MV_SQL,
+    DOCUMENT_EMBEDDINGS_TABLE_SQL,
     DOCUMENT_EMBEDDINGS_WRITABLE_TABLE_SQL,
     KAFKA_DOCUMENT_EMBEDDINGS_TABLE_SQL,
 )
@@ -202,7 +202,7 @@
     PERSONS_DISTINCT_ID_TABLE_SQL,
     PERSON_DISTINCT_ID2_TABLE_SQL,
     PERSON_DISTINCT_ID_OVERRIDES_TABLE_SQL,
-    DOCUMENT_EMBEDDINGS_DATA_TABLE_SQL,
+    DOCUMENT_EMBEDDINGS_TABLE_SQL,
     ERROR_TRACKING_ISSUE_FINGERPRINT_OVERRIDES_TABLE_SQL,
     ERROR_TRACKING_FINGERPRINT_EMBEDDINGS_TABLE_SQL,
     PLUGIN_LOG_ENTRIES_TABLE_SQL,

posthog/hogql/database/schema/document_embeddings.py

@@ -13,7 +13,7 @@
     Table,
 )
 
-from products.error_tracking.backend.embedding import DOCUMENT_EMBEDDINGS
+from products.error_tracking.backend.embedding import DISTRIBUTED_DOCUMENT_EMBEDDINGS
 
 DOCUMENT_EMBEDDINGS_FIELDS: dict[str, FieldOrTable] = {
     "team_id": IntegerDatabaseField(name="team_id", nullable=False),
@@ -49,7 +49,7 @@ class RawDocumentEmbeddingsTable(Table):
     fields: dict[str, FieldOrTable] = DOCUMENT_EMBEDDINGS_FIELDS
 
     def to_printed_clickhouse(self, context):
-        return DOCUMENT_EMBEDDINGS
+        return DISTRIBUTED_DOCUMENT_EMBEDDINGS
 
     def to_printed_hogql(self):
         return f"raw_document_embeddings"
@@ -67,7 +67,7 @@ def lazy_select(
         return select_from_embeddings_table(table_to_add.fields_accessed)
 
     def to_printed_clickhouse(self, context):
-        return DOCUMENT_EMBEDDINGS
+        return DISTRIBUTED_DOCUMENT_EMBEDDINGS
 
     def to_printed_hogql(self):
         return "document_embeddings"

posthog/hogql_queries/test/test_document_embeddings_query_runner.py

@@ -77,7 +77,7 @@ class DocumentEmbeddingRow:
         inserted_at: datetime
 
     def _seed_document_embeddings(self) -> list[DocumentEmbeddingRow]:
-        sync_execute("TRUNCATE TABLE posthog_document_embeddings", flush=False, team_id=self.team.pk)
+        sync_execute("TRUNCATE TABLE distributed_posthog_document_embeddings", flush=False, team_id=self.team.pk)
 
         fixtures: list[TestDocumentEmbeddingsQueryRunner.DocumentEmbeddingRow] = []
         rows: list[tuple] = []
@@ -128,7 +128,7 @@ def _seed_document_embeddings(self) -> list[DocumentEmbeddingRow]:
         if rows:
             sync_execute(
                 """
-                INSERT INTO posthog_document_embeddings (
+                INSERT INTO distributed_posthog_document_embeddings (
                     team_id,
                     product,
                     document_type,

products/error_tracking/backend/embedding.py

@@ -30,13 +30,14 @@
 """
 
 
-def DOCUMENT_EMBEDDINGS_TABLE_ENGINE():
-    return ReplacingMergeTree(
-        SHARDED_DOCUMENT_EMBEDDINGS, ver="inserted_at", replication_scheme=ReplicationScheme.SHARDED
-    )
+# The flow of this table set, as per other sharded tables, is:
+# - Kafka table exposes messages from Kafka topic
+# - Materialized view reads from Kafka table, writes to writable table, moving the kafka offset
+# - Writable table distributes writes to sharded tables
+# - Distributed table distributes reads to sharded tables
 
 
-def DOCUMENT_EMBEDDINGS_DATA_TABLE_SQL():
+def DOCUMENT_EMBEDDINGS_TABLE_SQL():
     return (
         DOCUMENT_EMBEDDINGS_TABLE_BASE_SQL
         + """
@@ -49,7 +50,9 @@ def DOCUMENT_EMBEDDINGS_DATA_TABLE_SQL():
     """
     ).format(
         table_name=SHARDED_DOCUMENT_EMBEDDINGS,
-        engine=DOCUMENT_EMBEDDINGS_TABLE_ENGINE(),
+        engine=ReplacingMergeTree(
+            SHARDED_DOCUMENT_EMBEDDINGS, ver="inserted_at", replication_scheme=ReplicationScheme.SHARDED
+        ),
         default_clause=" DEFAULT ''",
         extra_fields=f"""
     {KAFKA_COLUMNS_WITH_PARTITION}
@@ -58,6 +61,8 @@ def DOCUMENT_EMBEDDINGS_DATA_TABLE_SQL():
     )
 
 
+# The sharding keys of this and the table below are chosen mostly at random - as far as I could tell,
+# there isn't much to be gained from trying to get clever here, and it's best just to keep spread even
 def DISTRIBUTED_DOCUMENT_EMBEDDINGS_TABLE_SQL():
     return DOCUMENT_EMBEDDINGS_TABLE_BASE_SQL.format(
         table_name=DISTRIBUTED_DOCUMENT_EMBEDDINGS,
@@ -122,7 +127,3 @@ def DOCUMENT_EMBEDDINGS_MV_SQL(
 
 def TRUNCATE_DOCUMENT_EMBEDDINGS_TABLE_SQL():
     return f"TRUNCATE TABLE IF EXISTS {SHARDED_DOCUMENT_EMBEDDINGS} ON CLUSTER '{settings.CLICKHOUSE_CLUSTER}'"
-
-
-# Backwards compatibility alias for old migrations (0155, 0174)
-DOCUMENT_EMBEDDINGS_TABLE_SQL = DOCUMENT_EMBEDDINGS_DATA_TABLE_SQL