Support multi-arg AggregateFunction types

src/codegen/type-mapper.test.ts

@@ -173,6 +173,30 @@ describe("clickhouseTypeToValidator", () => {
       );
     });
 
+    it("handles AggregateFunction with multiple state types", () => {
+      expect(clickhouseTypeToValidator("AggregateFunction(argMax, String, DateTime)")).toBe(
+        't.aggregateFunction("argMax", t.string(), t.dateTime())'
+      );
+    });
+
+    it("handles other AggregateFunction variants with multiple state types", () => {
+      expect(clickhouseTypeToValidator("AggregateFunction(argMin, Float64, DateTime)")).toBe(
+        't.aggregateFunction("argMin", t.float64(), t.dateTime())'
+      );
+      expect(clickhouseTypeToValidator("AggregateFunction(corr, Float64, Float64)")).toBe(
+        't.aggregateFunction("corr", t.float64(), t.float64())'
+      );
+      expect(clickhouseTypeToValidator("AggregateFunction(sumMap, Array(String), Array(UInt64))")).toBe(
+        't.aggregateFunction("sumMap", t.array(t.string()), t.array(t.uint64()))'
+      );
+    });
+
+    it("handles AggregateFunction names with parameters containing commas", () => {
+      expect(clickhouseTypeToValidator("AggregateFunction(quantiles(0.5, 0.9), UInt64)")).toBe(
+        't.aggregateFunction("quantiles(0.5, 0.9)", t.uint64())'
+      );
+    });
+
     it("handles AggregateFunction(count) without explicit state type", () => {
       expect(clickhouseTypeToValidator("AggregateFunction(count)")).toBe(
         't.aggregateFunction("count")'

src/codegen/type-mapper.ts

@@ -244,20 +244,17 @@ export function clickhouseTypeToValidator(chType: string): string {
     return `t.simpleAggregateFunction("${func}", ${innerType})`;
   }
 
-  // AggregateFunction(func, T)
-  const aggMatch = chType.match(/^AggregateFunction\((\w+),\s*(.+)\)$/);
+  // AggregateFunction(func, T1, T2, ...)
+  const aggMatch = chType.match(/^AggregateFunction\((.+)\)$/);
   if (aggMatch) {
-    const func = aggMatch[1];
-    const innerType = clickhouseTypeToValidator(aggMatch[2]);
-    return `t.aggregateFunction("${func}", ${innerType})`;
-  }
-
-  // AggregateFunction(count)
-  const aggNoArgMatch = chType.match(/^AggregateFunction\((\w+)\)$/);
-  if (aggNoArgMatch) {
-    const func = aggNoArgMatch[1];
-    if (func === "count") {
-      return 't.aggregateFunction("count")';
+    const args = splitTopLevelComma(aggMatch[1]);
+    if (args.length === 1) {
+      return `t.aggregateFunction(${JSON.stringify(args[0])})`;
+    }
+    if (args.length > 1) {
+      const [func, ...stateTypes] = args;
+      const innerTypes = stateTypes.map((type) => clickhouseTypeToValidator(type));
+      return `t.aggregateFunction(${JSON.stringify(func)}, ${innerTypes.join(", ")})`;
     }
     return `t.string() /* TODO: Unknown type: ${chType} */`;
   }

src/schema/types.test.ts

@@ -222,6 +222,36 @@ describe("Type Validators (t.*)", () => {
       expectTypeOf(type._type).toEqualTypeOf<string>();
     });
 
+    it("generates AggregateFunction with multiple explicit state types", () => {
+      const type = t.aggregateFunction("argMax", t.string(), t.dateTime());
+      expect(type._tinybirdType).toBe("AggregateFunction(argMax, String, DateTime)");
+      expectTypeOf(type._type).toEqualTypeOf<string>();
+    });
+
+    it("generates AggregateFunction with multiple explicit state types for other functions", () => {
+      expect(t.aggregateFunction("argMin", t.float64(), t.dateTime())._tinybirdType).toBe(
+        "AggregateFunction(argMin, Float64, DateTime)"
+      );
+      expect(t.aggregateFunction("corr", t.float64(), t.float64())._tinybirdType).toBe(
+        "AggregateFunction(corr, Float64, Float64)"
+      );
+      expect(t.aggregateFunction("sumMap", t.array(t.string()), t.array(t.uint64()))._tinybirdType).toBe(
+        "AggregateFunction(sumMap, Array(String), Array(UInt64))"
+      );
+    });
+
+    it("generates AggregateFunction with function parameters and multiple state types", () => {
+      const type = t.aggregateFunction(
+        "sequenceMatch('(?1)(?2)')",
+        t.dateTime(),
+        t.uint8(),
+        t.uint8()
+      );
+      expect(type._tinybirdType).toBe(
+        "AggregateFunction(sequenceMatch('(?1)(?2)'), DateTime, UInt8, UInt8)"
+      );
+    });
+
     it("generates AggregateFunction without an explicit state type", () => {
       const type = t.aggregateFunction("count");
       expect(type._tinybirdType).toBe("AggregateFunction(count)");

src/schema/types.ts

@@ -228,29 +228,59 @@ type AggregateFunctionValidator = {
   ): TypeValidator<unknown, `AggregateFunction(${TFunc})`, TypeModifiers>;
   <
     TFunc extends string,
-    TType extends TypeValidator<unknown, string, TypeModifiers>,
+    TTypes extends readonly [
+      TypeValidator<unknown, string, TypeModifiers>,
+      ...TypeValidator<unknown, string, TypeModifiers>[],
+    ],
   >(
     func: TFunc,
-    type: TType,
+    ...types: TTypes
   ): TypeValidator<
-    TType["_type"],
-    `AggregateFunction(${TFunc}, ${TType["_tinybirdType"]})`,
+    AggregateFunctionValue<TTypes>,
+    AggregateFunctionTinybirdType<TFunc, TTypes>,
     TypeModifiers
   >;
 };
 
+type AggregateFunctionValue<
+  TTypes extends readonly TypeValidator<unknown, string, TypeModifiers>[],
+> = TTypes extends readonly [
+  infer TFirst extends TypeValidator<unknown, string, TypeModifiers>,
+  ...TypeValidator<unknown, string, TypeModifiers>[],
+]
+  ? TFirst["_type"]
+  : unknown;
+
+type AggregateFunctionArgs<
+  TTypes extends readonly TypeValidator<unknown, string, TypeModifiers>[],
+> = TTypes extends readonly [
+  infer TFirst extends TypeValidator<unknown, string, TypeModifiers>,
+  ...infer TRest extends TypeValidator<unknown, string, TypeModifiers>[],
+]
+  ? TRest extends []
+    ? TFirst["_tinybirdType"]
+    : `${TFirst["_tinybirdType"]}, ${AggregateFunctionArgs<TRest>}`
+  : never;
+
+type AggregateFunctionTinybirdType<
+  TFunc extends string,
+  TTypes extends readonly TypeValidator<unknown, string, TypeModifiers>[],
+> = TTypes extends []
+  ? `AggregateFunction(${TFunc})`
+  : `AggregateFunction(${TFunc}, ${AggregateFunctionArgs<TTypes>})`;
+
 const aggregateFunction = ((
   func: string,
-  type?: TypeValidator<unknown, string, TypeModifiers>,
+  ...types: TypeValidator<unknown, string, TypeModifiers>[]
 ) => {
-  if (type === undefined) {
+  if (types.length === 0) {
     return createValidator<unknown, `AggregateFunction(${string})`>(
       `AggregateFunction(${func})`,
     );
   }
 
   return createValidator<unknown, `AggregateFunction(${string}, ${string})`>(
-    `AggregateFunction(${func}, ${type._tinybirdType})`,
+    `AggregateFunction(${func}, ${types.map((type) => type._tinybirdType).join(", ")})`,
   );
 }) as AggregateFunctionValidator;