clickhouse: PARTITION BY expression (#3341)

I was wrong, making partitioning as inflexible as ORDER BY columns doesn't work

Most easily is that you should pretty much never partition by a datetime without rounding off most of its lower bits

Previous partitioning should be considered deprecated

Author: serprex

flow/connectors/clickhouse/normalize.go

@@ -222,7 +222,7 @@ func (c *ClickHouseConnector) generateCreateTableSQLForNormalizedTable(
 
 	if tmEngine != protos.TableEngine_CH_ENGINE_NULL {
 		hasNullableKeyFn := buildIsNullableKeyFn(tableMapping, tableSchema.Columns, tableSchema.NullableEnabled)
-		orderByColumns, hasNullableOrderKey := getOrderedOrderByColumns(
+		orderByColumns, allowNullableKey := getOrderedOrderByColumns(
 			tableMapping, colNameMap, tableSchema.PrimaryKeyColumns, hasNullableKeyFn)
 		if sourceSchemaAsDestinationColumn {
 			orderByColumns = append([]string{sourceSchemaColName}, orderByColumns...)
@@ -235,13 +235,21 @@ func (c *ClickHouseConnector) generateCreateTableSQLForNormalizedTable(
 			stmtBuilder.WriteString(" ORDER BY tuple()")
 		}
 
-		partitionByColumns, hasNullablePartitionKey := getOrderedPartitionByColumns(tableMapping, colNameMap, hasNullableKeyFn)
-		if len(partitionByColumns) > 0 {
-			partitionByStr := strings.Join(partitionByColumns, ",")
-			fmt.Fprintf(&stmtBuilder, " PARTITION BY (%s)", partitionByStr)
+		if tableMapping != nil && tableMapping.PartitionByExpr != "" {
+			allowNullableKey = true
+			fmt.Fprintf(&stmtBuilder, " PARTITION BY (%s)", tableMapping.PartitionByExpr)
+		} else {
+			partitionByColumns, hasNullablePartitionKey := getOrderedPartitionByColumns(tableMapping, colNameMap, hasNullableKeyFn)
+			if hasNullablePartitionKey {
+				allowNullableKey = true
+			}
+			if len(partitionByColumns) > 0 {
+				partitionByStr := strings.Join(partitionByColumns, ",")
+				fmt.Fprintf(&stmtBuilder, " PARTITION BY (%s)", partitionByStr)
+			}
 		}
 
-		if hasNullableOrderKey || hasNullablePartitionKey {
+		if allowNullableKey {
 			stmtBuilder.WriteString(" SETTINGS allow_nullable_key = 1")
 		}
 

flow/e2e/clickhouse/peer_flow_ch_test.go

@@ -2434,3 +2434,52 @@ func (s ClickHouseSuite) Test_PartitionBy() {
 	env.Cancel(s.t.Context())
 	e2e.RequireEnvCanceled(s.t, env)
 }
+
+func (s ClickHouseSuite) Test_PartitionByExpr() {
+	srcTableName := "test_partition_by_expr"
+	srcFullName := s.attachSchemaSuffix(srcTableName)
+	dstTableName := "test_partition_by_expr"
+
+	require.NoError(s.t, s.source.Exec(s.t.Context(),
+		fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s (id INT PRIMARY KEY, num INT, val TEXT NOT NULL)`, srcFullName)))
+
+	connectionGen := e2e.FlowConnectionGenerationConfig{
+		FlowJobName: s.attachSuffix("clickhouse_partition_by"),
+		TableMappings: []*protos.TableMapping{{
+			SourceTableIdentifier:      srcFullName,
+			DestinationTableIdentifier: dstTableName,
+			PartitionByExpr:            "num%2,val",
+			Columns: []*protos.ColumnSetting{
+				{SourceName: "id", NullableEnabled: true},
+				{SourceName: "num", NullableEnabled: true},
+				{SourceName: "val", NullableEnabled: true},
+			},
+		}},
+		Destination: s.Peer().Name,
+	}
+
+	flowConnConfig := connectionGen.GenerateFlowConnectionConfigs(s)
+	tc := e2e.NewTemporalClient(s.t)
+	env := e2e.ExecutePeerflow(s.t, tc, flowConnConfig)
+	e2e.SetupCDCFlowStatusQuery(s.t, env, flowConnConfig)
+
+	e2e.EnvWaitForEqualTablesWithNames(env, s, "table setup", srcTableName, dstTableName, "id")
+
+	var partitionKey string
+	ch, err := connclickhouse.Connect(s.t.Context(), nil, s.Peer().GetClickhouseConfig())
+	require.NoError(s.t, err)
+	var dstTableSuffix string
+	if s.cluster {
+		dstTableSuffix = "_shard"
+	}
+	require.NoError(s.t,
+		ch.QueryRow(s.t.Context(),
+			"select partition_key from system.tables where name="+clickhouse.QuoteLiteral(dstTableName+dstTableSuffix),
+		).Scan(&partitionKey),
+	)
+	require.NoError(s.t, ch.Close())
+	require.Equal(s.t, "(num % 2, val)", partitionKey)
+
+	env.Cancel(s.t.Context())
+	e2e.RequireEnvCanceled(s.t, env)
+}

protos/flow.proto

@@ -33,6 +33,7 @@ message TableMapping {
   TableEngine engine = 6;
   string sharding_key = 7;
   string policy_name = 8;
+  string partition_by_expr = 9;
 }
 
 message SetupInput {

ui/app/mirrors/create/cdc/schemabox.tsx

@@ -137,6 +137,13 @@ export default function SchemaBox({
     setRows(newRows);
   };
 
+  const updatePartitionByExpr = (source: string, partitionByExpr: string) => {
+    const newRows = [...rows];
+    const index = newRows.findIndex((row) => row.source === source);
+    newRows[index] = { ...newRows[index], partitionByExpr };
+    setRows(newRows);
+  };
+
   const addTableColumns = useCallback(
     (table: string) => {
       const [schemaName, tableName] = table.split('.');
@@ -228,6 +235,7 @@ export default function SchemaBox({
               row.partitionKey = existingRow.partitionKey;
               row.shardingKey = existingRow.shardingKey;
               row.policyName = existingRow.policyName;
+              row.partitionByExpr = existingRow.partitionByExpr;
               row.exclude = new Set(existingRow.exclude ?? []);
               row.destination = existingRow.destinationTableIdentifier;
               addTableColumns(row.source);
@@ -366,6 +374,7 @@ export default function SchemaBox({
                           rowGap: '0.5rem',
                           columnGap: '3rem',
                           display: row.selected ? 'flex' : 'none',
+                          flexWrap: 'wrap',
                         }}
                         key={row.source}
                       >
@@ -466,6 +475,27 @@ export default function SchemaBox({
                                 }
                               />
                             </div>
+                            <div style={{ width: '30%', fontSize: 12 }}>
+                              Partition By Expr:
+                              <TextField
+                                disabled={row.editingDisabled}
+                                style={{
+                                  marginTop: '0.5rem',
+                                  cursor: 'pointer',
+                                }}
+                                variant='simple'
+                                placeholder='Partition By expression (optional)'
+                                value={row.partitionByExpr}
+                                onChange={(
+                                  e: React.ChangeEvent<HTMLInputElement>
+                                ) =>
+                                  updatePartitionByExpr(
+                                    row.source,
+                                    e.target.value
+                                  )
+                                }
+                              />
+                            </div>
                           </>
                         )}
                       </div>

ui/app/mirrors/create/handlers.ts

@@ -188,6 +188,7 @@ export function reformattedTableMapping(
       engine: row.engine,
       shardingKey: row.shardingKey,
       policyName: row.policyName,
+      partitionByExpr: row.partitionByExpr,
     }));
 }
 
@@ -454,6 +455,7 @@ export async function fetchTables(
         engine: TableEngine.CH_ENGINE_REPLACING_MERGE_TREE,
         shardingKey: '',
         policyName: '',
+        partitionByExpr: '',
       });
     }
   }