feat: Support property group migrations on distributed and sharded tables

posthog/clickhouse/migrations/0074_add_custom_and_features_property_groups.py

@@ -1,10 +1,10 @@
 from posthog.clickhouse.client.migration_tools import run_sql_with_exceptions
-from posthog.clickhouse.property_groups import sharded_events_property_groups
+from posthog.clickhouse.property_groups import property_groups
 
 operations = [
     run_sql_with_exceptions(statement)
     for statement in [
-        *sharded_events_property_groups.get_alter_create_statements("custom"),
-        *sharded_events_property_groups.get_alter_create_statements("feature_flags"),
+        *property_groups.get_alter_create_statements("sharded_events", "properties", "custom"),
+        *property_groups.get_alter_create_statements("sharded_events", "properties", "feature_flags"),
     ]
 ]

posthog/clickhouse/migrations/0075_add_custom_and_features_property_groups_events.py

@@ -0,0 +1,10 @@
+from posthog.clickhouse.client.migration_tools import run_sql_with_exceptions
+from posthog.clickhouse.property_groups import property_groups
+
+operations = [
+    run_sql_with_exceptions(statement)
+    for statement in [
+        *property_groups.get_alter_create_statements("events", "properties", "custom"),
+        *property_groups.get_alter_create_statements("events", "properties", "feature_flags"),
+    ]
+]

posthog/clickhouse/property_groups.py

@@ -1,40 +1,75 @@
-from collections.abc import Iterable, MutableMapping
+import dataclasses
+from collections.abc import Iterable, Mapping
 from dataclasses import dataclass
 
 from posthog import settings
 
 
 @dataclass
 class PropertyGroupDefinition:
-    filter_expression: str
+    key_filter_expression: str
     codec: str = "ZSTD(1)"
+    is_materialized: bool = True
+
+
+TableName = str
+ColumnName = str
+PropertyGroupName = str
 
 
 class PropertyGroupManager:
-    def __init__(self, cluster: str, table: str, source_column: str) -> None:
+    def __init__(
+        self,
+        cluster: str,
+        groups: Mapping[TableName, Mapping[ColumnName, Mapping[PropertyGroupName, PropertyGroupDefinition]]],
+    ) -> None:
         self.__cluster = cluster
-        self.__table = table
-        self.__source_column = source_column
-        self.__groups: MutableMapping[str, PropertyGroupDefinition] = {}
+        self.__groups = groups
+
+    def __get_map_column_name(self, column: ColumnName, group_name: PropertyGroupName) -> str:
+        return f"{column}_group_{group_name}"
 
-    def register(self, name: str, definition: PropertyGroupDefinition) -> None:
-        assert name not in self.__groups, "property group names can only be used once"
-        self.__groups[name] = definition
+    def __get_column_definition(self, table: TableName, column: ColumnName, group_name: PropertyGroupName) -> str:
+        group_definition = self.__groups[table][column][group_name]
+        map_column_name = self.__get_map_column_name(column, group_name)
+        column_definition = f"{map_column_name} Map(String, String)"
+        if not group_definition.is_materialized:
+            return column_definition
+        else:
+            return f"""\
+                {column_definition}
+                MATERIALIZED mapSort(
+                    mapFilter((key, _) -> {group_definition.key_filter_expression},
+                    CAST(JSONExtractKeysAndValues({column}, 'String'), 'Map(String, String)'))
+                )
+                CODEC({group_definition.codec})
+            """
 
-    def __get_map_expression(self, definition: PropertyGroupDefinition) -> str:
-        return f"mapSort(mapFilter((key, _) -> {definition.filter_expression}, CAST(JSONExtractKeysAndValues({self.__source_column}, 'String'), 'Map(String, String)')))"
+    def __get_index_definitions(
+        self, table: TableName, column: ColumnName, group_name: PropertyGroupName
+    ) -> Iterable[str]:
+        group_definition = self.__groups[table][column][group_name]
+        if not group_definition.is_materialized:
+            return
 
-    def get_alter_create_statements(self, name: str) -> Iterable[str]:
-        column_name = f"{self.__source_column}_group_{name}"
-        definition = self.__groups[name]
-        return [
-            f"ALTER TABLE {self.__table} ON CLUSTER {self.__cluster} ADD COLUMN {column_name} Map(String, String) MATERIALIZED {self.__get_map_expression(definition)} CODEC({definition.codec})",
-            f"ALTER TABLE {self.__table} ON CLUSTER {self.__cluster} ADD INDEX {column_name}_keys_bf mapKeys({column_name}) TYPE bloom_filter",
-            f"ALTER TABLE {self.__table} ON CLUSTER {self.__cluster} ADD INDEX {column_name}_values_bf mapValues({column_name}) TYPE bloom_filter",
-        ]
+        map_column_name = self.__get_map_column_name(column, group_name)
+        yield f"{map_column_name}_keys_bf mapKeys({map_column_name}) TYPE bloom_filter"
+        yield f"{map_column_name}_values_bf mapValues({map_column_name}) TYPE bloom_filter"
 
+    def get_create_table_pieces(self, table: TableName) -> Iterable[str]:
+        for column, groups in self.__groups[table].items():
+            for group_name in groups:
+                yield self.__get_column_definition(table, column, group_name)
+                for index_definition in self.__get_index_definitions(table, column, group_name):
+                    yield f"INDEX {index_definition}"
+
+    def get_alter_create_statements(
+        self, table: TableName, column: ColumnName, group_name: PropertyGroupName
+    ) -> Iterable[str]:
+        yield f"ALTER TABLE {table} ON CLUSTER {self.__cluster} ADD COLUMN IF NOT EXISTS {self.__get_column_definition(table, column, group_name)}"
+        for index_definition in self.__get_index_definitions(table, column, group_name):
+            yield f"ALTER TABLE {table} ON CLUSTER {self.__cluster} ADD INDEX IF NOT EXISTS {index_definition}"
 
-sharded_events_property_groups = PropertyGroupManager(settings.CLICKHOUSE_CLUSTER, "sharded_events", "properties")
 
 ignore_custom_properties = [
     # `token` & `distinct_id` properties are sent with ~50% of events and by
@@ -65,11 +100,25 @@ def get_alter_create_statements(self, name: str) -> Iterable[str]:
     "rdt_cid",  # reddit
 ]
 
-sharded_events_property_groups.register(
-    "custom",
-    PropertyGroupDefinition(
-        f"key NOT LIKE '$%' AND key NOT IN (" + f", ".join(f"'{name}'" for name in ignore_custom_properties) + f")"
-    ),
-)
+event_property_group_definitions = {
+    "properties": {
+        "custom": PropertyGroupDefinition(
+            f"key NOT LIKE '$%' AND key NOT IN (" + f", ".join(f"'{name}'" for name in ignore_custom_properties) + f")"
+        ),
+        "feature_flags": PropertyGroupDefinition("key like '$feature/%'"),
+    }
+}
 
-sharded_events_property_groups.register("feature_flags", PropertyGroupDefinition("key like '$feature/%'"))
+property_groups = PropertyGroupManager(
+    settings.CLICKHOUSE_CLUSTER,
+    {
+        "sharded_events": event_property_group_definitions,
+        "events": {
+            column_name: {
+                group_name: dataclasses.replace(group_definition, is_materialized=False)
+                for group_name, group_definition in column_group_definitions.items()
+            }
+            for column_name, column_group_definitions in event_property_group_definitions.items()
+        },
+    },
+)

posthog/clickhouse/test/__snapshots__/test_schema.ambr

@@ -599,6 +599,7 @@
       , elements_chain_texts Array(String) COMMENT 'column_materializer::elements_chain::texts'
       , elements_chain_ids Array(String) COMMENT 'column_materializer::elements_chain::ids'
       , elements_chain_elements Array(Enum('a', 'button', 'form', 'input', 'select', 'textarea', 'label')) COMMENT 'column_materializer::elements_chain::elements'
+      , properties_group_custom Map(String, String), properties_group_feature_flags Map(String, String)
 
 
   , _timestamp DateTime
@@ -2188,6 +2189,19 @@
       , INDEX `minmax_$group_4` `$group_4` TYPE minmax GRANULARITY 1
       , INDEX `minmax_$window_id` `$window_id` TYPE minmax GRANULARITY 1
       , INDEX `minmax_$session_id` `$session_id` TYPE minmax GRANULARITY 1
+      ,                 properties_group_custom Map(String, String)
+                  MATERIALIZED mapSort(
+                      mapFilter((key, _) -> key NOT LIKE '$%' AND key NOT IN ('token', 'distinct_id', 'utm_source', 'utm_medium', 'utm_campaign', 'utm_content', 'utm_term', 'gclid', 'gad_source', 'gclsrc', 'dclid', 'gbraid', 'wbraid', 'fbclid', 'msclkid', 'twclid', 'li_fat_id', 'mc_cid', 'igshid', 'ttclid', 'rdt_cid'),
+                      CAST(JSONExtractKeysAndValues(properties, 'String'), 'Map(String, String)'))
+                  )
+                  CODEC(ZSTD(1))
+              , INDEX properties_group_custom_keys_bf mapKeys(properties_group_custom) TYPE bloom_filter, INDEX properties_group_custom_values_bf mapValues(properties_group_custom) TYPE bloom_filter,                 properties_group_feature_flags Map(String, String)
+                  MATERIALIZED mapSort(
+                      mapFilter((key, _) -> key like '$feature/%',
+                      CAST(JSONExtractKeysAndValues(properties, 'String'), 'Map(String, String)'))
+                  )
+                  CODEC(ZSTD(1))
+              , INDEX properties_group_feature_flags_keys_bf mapKeys(properties_group_feature_flags) TYPE bloom_filter, INDEX properties_group_feature_flags_values_bf mapValues(properties_group_feature_flags) TYPE bloom_filter
 
 
   , _timestamp DateTime
@@ -3282,6 +3296,19 @@
       , INDEX `minmax_$group_4` `$group_4` TYPE minmax GRANULARITY 1
       , INDEX `minmax_$window_id` `$window_id` TYPE minmax GRANULARITY 1
       , INDEX `minmax_$session_id` `$session_id` TYPE minmax GRANULARITY 1
+      ,                 properties_group_custom Map(String, String)
+                  MATERIALIZED mapSort(
+                      mapFilter((key, _) -> key NOT LIKE '$%' AND key NOT IN ('token', 'distinct_id', 'utm_source', 'utm_medium', 'utm_campaign', 'utm_content', 'utm_term', 'gclid', 'gad_source', 'gclsrc', 'dclid', 'gbraid', 'wbraid', 'fbclid', 'msclkid', 'twclid', 'li_fat_id', 'mc_cid', 'igshid', 'ttclid', 'rdt_cid'),
+                      CAST(JSONExtractKeysAndValues(properties, 'String'), 'Map(String, String)'))
+                  )
+                  CODEC(ZSTD(1))
+              , INDEX properties_group_custom_keys_bf mapKeys(properties_group_custom) TYPE bloom_filter, INDEX properties_group_custom_values_bf mapValues(properties_group_custom) TYPE bloom_filter,                 properties_group_feature_flags Map(String, String)
+                  MATERIALIZED mapSort(
+                      mapFilter((key, _) -> key like '$feature/%',
+                      CAST(JSONExtractKeysAndValues(properties, 'String'), 'Map(String, String)'))
+                  )
+                  CODEC(ZSTD(1))
+              , INDEX properties_group_feature_flags_keys_bf mapKeys(properties_group_feature_flags) TYPE bloom_filter, INDEX properties_group_feature_flags_values_bf mapValues(properties_group_feature_flags) TYPE bloom_filter
 
 
   , _timestamp DateTime

posthog/models/event/sql.py

@@ -8,6 +8,7 @@
     kafka_engine,
     trim_quotes_expr,
 )
+from posthog.clickhouse.property_groups import property_groups
 from posthog.clickhouse.table_engines import (
     Distributed,
     ReplacingMergeTree,
@@ -76,9 +77,10 @@
     , INDEX `minmax_$group_4` `$group_4` TYPE minmax GRANULARITY 1
     , INDEX `minmax_$window_id` `$window_id` TYPE minmax GRANULARITY 1
     , INDEX `minmax_$session_id` `$session_id` TYPE minmax GRANULARITY 1
+    , {", ".join(property_groups.get_create_table_pieces("sharded_events"))}
 """
 
-EVENTS_TABLE_PROXY_MATERIALIZED_COLUMNS = """
+EVENTS_TABLE_PROXY_MATERIALIZED_COLUMNS = f"""
     , $group_0 VARCHAR COMMENT 'column_materializer::$group_0'
     , $group_1 VARCHAR COMMENT 'column_materializer::$group_1'
     , $group_2 VARCHAR COMMENT 'column_materializer::$group_2'
@@ -90,6 +92,7 @@
     , elements_chain_texts Array(String) COMMENT 'column_materializer::elements_chain::texts'
     , elements_chain_ids Array(String) COMMENT 'column_materializer::elements_chain::ids'
     , elements_chain_elements Array(Enum('a', 'button', 'form', 'input', 'select', 'textarea', 'label')) COMMENT 'column_materializer::elements_chain::elements'
+    , {", ".join(property_groups.get_create_table_pieces("events"))}
 """
 
 EVENTS_DATA_TABLE_ENGINE = lambda: ReplacingMergeTree(