feat: Add weighted mean for retention insights

cypress/e2e/insights-navigation-open-sql-insight-first.cy.ts

@@ -46,7 +46,7 @@ describe('Insights', () => {
             it('can open a new retention insight', () => {
                 insight.clickTab('RETENTION')
                 cy.get('.RetentionContainer canvas').should('exist')
-                cy.get('.RetentionTable__Tab').should('have.length', 66)
+                cy.get('.RetentionTable__Tab').should('have.length', 77)
             })
 
             it('can open a new paths insight', () => {

cypress/e2e/insights-navigation.cy.ts

@@ -43,7 +43,7 @@ describe('Insights', () => {
             cy.get('[data-attr="insight-save-button"]').click()
 
             cy.get('.RetentionContainer canvas').should('exist')
-            cy.get('.RetentionTable__Tab').should('have.length', 66)
+            cy.get('.RetentionTable__Tab').should('have.length', 77)
         })
 
         it('can open a new SQL insight and navigate to a different one, then back to SQL, and back again', () => {

frontend/src/lib/constants.tsx

@@ -256,6 +256,8 @@ export enum FunnelLayout {
 
 export const BIN_COUNT_AUTO = 'auto' as const
 
+export const RETENTION_MEAN_NONE = 'none' as const
+
 // Cohort types
 export enum CohortTypeEnum {
     Static = 'static',

frontend/src/queries/nodes/InsightQuery/defaults.ts

@@ -54,6 +54,7 @@ const retentionQueryDefault: RetentionQuery = {
             type: 'events',
         },
         retentionType: 'retention_first_time',
+        showMean: 'simple',
     },
 }
 

frontend/src/queries/nodes/InsightQuery/utils/filtersToQueryNode.test.ts

@@ -669,7 +669,7 @@ describe('filtersToQueryNode', () => {
                 returning_entity: { id: '1' },
                 target_entity: { id: '1' },
                 period: RetentionPeriod.Day,
-                show_mean: true,
+                show_mean: 'simple',
             }
 
             const result = filtersToQueryNode(filters)
@@ -683,7 +683,7 @@ describe('filtersToQueryNode', () => {
                     returningEntity: { id: '1' },
                     targetEntity: { id: '1' },
                     period: RetentionPeriod.Day,
-                    showMean: true,
+                    showMean: 'simple',
                 },
             }
             expect(result).toEqual(query)

frontend/src/queries/nodes/InsightViz/InsightDisplayConfig.tsx

@@ -22,7 +22,7 @@ import { ShowMultipleYAxesFilter } from 'scenes/insights/EditorFilters/ShowMulti
 import { ValueOnSeriesFilter } from 'scenes/insights/EditorFilters/ValueOnSeriesFilter'
 import { InsightDateFilter } from 'scenes/insights/filters/InsightDateFilter'
 import { RetentionCumulativeCheckbox } from 'scenes/insights/filters/RetentionCumulativeCheckbox'
-import { RetentionMeanCheckbox } from 'scenes/insights/filters/RetentionMeanCheckbox'
+import { RetentionMeanDropdown } from 'scenes/insights/filters/RetentionMeanDropdown'
 import { RetentionReferencePicker } from 'scenes/insights/filters/RetentionReferencePicker'
 import { insightLogic } from 'scenes/insights/insightLogic'
 import { insightVizDataLogic } from 'scenes/insights/insightVizDataLogic'
@@ -178,7 +178,7 @@ export function InsightDisplayConfig(): JSX.Element {
                     <ConfigFilter>
                         <RetentionDatePicker />
                         <RetentionReferencePicker />
-                        <RetentionMeanCheckbox />
+                        <RetentionMeanDropdown />
                         <RetentionCumulativeCheckbox />
                     </ConfigFilter>
                 )}

frontend/src/queries/schema.json

@@ -1258,7 +1258,8 @@
                 },
                 "showMean": {
                     "description": "Whether an additional series should be shown, showing the mean conversion for each period across cohorts.",
-                    "type": "boolean"
+                    "enum": ["simple", "weighted", "none"],
+                    "type": "string"
                 },
                 "targetEntity": {
                     "$ref": "#/definitions/RetentionEntity",
@@ -12479,7 +12480,8 @@
                     "$ref": "#/definitions/RetentionEntity"
                 },
                 "showMean": {
-                    "type": "boolean"
+                    "enum": ["simple", "weighted", "none"],
+                    "type": "string"
                 },
                 "targetEntity": {
                     "$ref": "#/definitions/RetentionEntity"
@@ -12513,7 +12515,8 @@
                     "$ref": "#/definitions/RetentionEntity"
                 },
                 "show_mean": {
-                    "type": "boolean"
+                    "enum": ["simple", "weighted", "none"],
+                    "type": "string"
                 },
                 "target_entity": {
                     "$ref": "#/definitions/RetentionEntity"

frontend/src/scenes/insights/filters/RetentionMeanDropdown.tsx

@@ -0,0 +1,52 @@
+import { LemonSelect } from '@posthog/lemon-ui'
+import { useActions, useValues } from 'kea'
+import { RETENTION_MEAN_NONE } from 'lib/constants'
+import { insightLogic } from 'scenes/insights/insightLogic'
+import { insightVizDataLogic } from 'scenes/insights/insightVizDataLogic'
+
+export type RetentionMeanType = 'simple' | 'weighted' | typeof RETENTION_MEAN_NONE
+
+export function RetentionMeanDropdown(): JSX.Element | null {
+    const { insightProps, canEditInsight } = useValues(insightLogic)
+
+    const { retentionFilter } = useValues(insightVizDataLogic(insightProps))
+    const { updateInsightFilter } = useActions(insightVizDataLogic(insightProps))
+
+    const showMean = retentionFilter?.showMean || RETENTION_MEAN_NONE
+
+    if (!canEditInsight) {
+        return null
+    }
+
+    return (
+        <LemonSelect
+            className="w-44"
+            size="small"
+            value={showMean}
+            onChange={(showMean) => {
+                updateInsightFilter({ showMean })
+            }}
+            options={[
+                {
+                    value: RETENTION_MEAN_NONE,
+                    labelInMenu: 'No mean calculation',
+                    label: 'No mean calculation',
+                },
+                {
+                    value: 'simple',
+                    labelInMenu: 'Simple mean',
+                    label: 'Simple mean',
+                    tooltip:
+                        'Calculates the average retention rate across all cohorts by giving equal weight to each cohort, regardless of its size.',
+                },
+                {
+                    value: 'weighted',
+                    labelInMenu: 'Weighted mean',
+                    label: 'Weighted mean',
+                    tooltip:
+                        'Calculates the average retention rate by giving more weight to larger cohorts, accounting for different cohort sizes in the final mean.',
+                },
+            ]}
+        />
+    )
+}

frontend/src/scenes/insights/utils.tsx

@@ -570,28 +570,35 @@ export function isQueryTooLarge(query: Node<Record<string, any>>): boolean {
     return queryLength > 1024 * 1024
 }
 
-export function parseDraftQueryFromLocalStorage(
-    query: string
-): { query: Node<Record<string, any>>; timestamp: number } | null {
+function parseAndMigrateQuery<T>(query: string): T | null {
     try {
-        return JSON.parse(query)
+        const parsedQuery = JSON.parse(query)
+        // We made a database migration to convert showMean from a boolean to a string,
+        // to allow for weighted and simple mean in retention tables. This ensures older URLs
+        // are parsed correctly.
+        const retentionFilter = parsedQuery?.source?.retentionFilter
+        if (retentionFilter && 'showMean' in retentionFilter && typeof retentionFilter.showMean === 'boolean') {
+            retentionFilter.showMean = retentionFilter.showMean ? 'simple' : null
+        }
+        return parsedQuery
     } catch (e) {
         console.error('Error parsing query', e)
         return null
     }
 }
 
+export function parseDraftQueryFromLocalStorage(
+    query: string
+): { query: Node<Record<string, any>>; timestamp: number } | null {
+    return parseAndMigrateQuery(query)
+}
+
 export function crushDraftQueryForLocalStorage(query: Node<Record<string, any>>, timestamp: number): string {
     return JSON.stringify({ query, timestamp })
 }
 
 export function parseDraftQueryFromURL(query: string): Node<Record<string, any>> | null {
-    try {
-        return JSON.parse(query)
-    } catch (e) {
-        console.error('Error parsing query', e)
-        return null
-    }
+    return parseAndMigrateQuery(query)
 }
 
 export function crushDraftQueryForURL(query: Node<Record<string, any>>): string {

frontend/src/scenes/insights/utils/cleanFilters.ts

@@ -5,6 +5,7 @@ import {
     NON_VALUES_ON_SERIES_DISPLAY_TYPES,
     PERCENT_STACK_VIEW_DISPLAY_TYPE,
     RETENTION_FIRST_TIME,
+    RETENTION_MEAN_NONE,
     ShownAsValue,
 } from 'lib/constants'
 import { clamp } from 'lib/utils'
@@ -307,7 +308,7 @@ export function cleanFilters(
             breakdowns: filters.breakdowns,
             breakdown_type: filters.breakdown_type,
             retention_reference: filters.retention_reference,
-            show_mean: filters.show_mean,
+            ...(filters.show_mean && filters.show_mean !== RETENTION_MEAN_NONE ? { show_mean: filters.show_mean } : {}),
             cumulative: filters.cumulative,
             total_intervals: Math.min(Math.max(filters.total_intervals ?? 11, 0), 100),
             ...(filters.aggregation_group_type_index != undefined

frontend/src/scenes/retention/RetentionTable.tsx

@@ -1,7 +1,7 @@
 import './RetentionTable.scss'
 
 import clsx from 'clsx'
-import { mean } from 'd3'
+import { mean, sum } from 'd3'
 import { useActions, useValues } from 'kea'
 import { Tooltip } from 'lib/lemon-ui/Tooltip'
 import { gradateColor, range } from 'lib/utils'
@@ -17,7 +17,7 @@ export function RetentionTable({ inSharedMode = false }: { inSharedMode?: boolea
     const { openModal } = useActions(retentionModalLogic(insightProps))
     const backgroundColor = theme?.['preset-1'] || '#000000' // Default to black if no color found
     const backgroundColorMean = theme?.['preset-2'] || '#000000' // Default to black if no color found
-    const showMean = retentionFilter?.showMean || false
+    const showMean = retentionFilter?.showMean || null
 
     return (
         <table
@@ -37,7 +37,51 @@ export function RetentionTable({ inSharedMode = false }: { inSharedMode?: boolea
                     ))}
                 </tr>
 
-                {showMean && tableRows.length > 0 ? (
+                {showMean === 'weighted' && tableRows.length > 0 ? (
+                    <tr className="border-b" key={-2}>
+                        {range(0, tableRows[0].length).map((columnIndex) => (
+                            <td key={columnIndex} className="pb-2">
+                                {columnIndex <= (hideSizeColumn ? 0 : 1) ? (
+                                    columnIndex == 0 ? (
+                                        <span className="RetentionTable__TextTab">Weighted Mean</span>
+                                    ) : null
+                                ) : (
+                                    <CohortDay
+                                        percentage={
+                                            (() => {
+                                                const validRows = tableRows.filter((row) => {
+                                                    return !(
+                                                        (columnIndex >= row.length - 1 && isLatestPeriod) ||
+                                                        !row[columnIndex] ||
+                                                        row[columnIndex].count <= 0
+                                                    )
+                                                })
+                                                if (validRows.length === 0) {
+                                                    return 0
+                                                }
+                                                const weights = validRows.map((row) =>
+                                                    parseInt(row[1]?.toString() || '0')
+                                                )
+                                                const weightedSum = sum(
+                                                    validRows.map(
+                                                        (row, i) => (row[columnIndex]?.percentage || 0) * weights[i]
+                                                    )
+                                                )
+                                                const totalWeight = sum(weights)
+                                                return totalWeight > 0 ? weightedSum / totalWeight : 0
+                                            })() || 0
+                                        }
+                                        latest={isLatestPeriod && columnIndex == tableRows[0].length - 1}
+                                        clickable={false}
+                                        backgroundColor={backgroundColorMean}
+                                    />
+                                )}
+                            </td>
+                        ))}
+                    </tr>
+                ) : undefined}
+
+                {showMean === 'simple' && tableRows.length > 0 ? (
                     <tr className="border-b" key={-1}>
                         {range(0, tableRows[0].length).map((columnIndex) => (
                             <td key={columnIndex} className="pb-2">
@@ -50,10 +94,12 @@ export function RetentionTable({ inSharedMode = false }: { inSharedMode?: boolea
                                         percentage={
                                             mean(
                                                 tableRows.map((row) => {
-                                                    // Stop before the last item in a row, which is an incomplete time period
+                                                    // Don't include the last item in a row, which is an incomplete time period
+                                                    // Also don't include the percentage if the cohort size (count) is 0 or less
                                                     if (
                                                         (columnIndex >= row.length - 1 && isLatestPeriod) ||
-                                                        !row[columnIndex]
+                                                        !row[columnIndex] ||
+                                                        row[columnIndex].count <= 0
                                                     ) {
                                                         return null
                                                     }

frontend/src/types.ts

@@ -16,6 +16,7 @@ import {
     PluginsAccessLevel,
     PROPERTY_MATCH_TYPE,
     RETENTION_FIRST_TIME,
+    RETENTION_MEAN_NONE,
     RETENTION_RECURRING,
     ShownAsValue,
     TeamMembershipLevel,
@@ -2435,7 +2436,7 @@ export interface RetentionFilterType extends FilterType {
     cumulative?: boolean
 
     //frontend only
-    show_mean?: boolean
+    show_mean?: 'simple' | 'weighted' | typeof RETENTION_MEAN_NONE
 }
 export interface LifecycleFilterType extends FilterType {
     /** @deprecated */

posthog/hogql_queries/legacy_compatibility/test/test_filter_to_query.py

@@ -1504,7 +1504,7 @@ def test_retention_filter(self):
             },
             "target_entity": {"id": "$pageview", "name": "$pageview", "type": "events"},
             "period": "Week",
-            "show_mean": True,
+            "show_mean": "simple",
             "cumulative": True,
         }
 
@@ -1531,7 +1531,7 @@ def test_retention_filter(self):
                     "custom_name": None,
                     "order": None,
                 },
-                showMean=True,
+                showMean="simple",
                 cumulative=True,
             ),
         )