feat(hogql): tuple access

posthog/hogql/constants.py

@@ -1,8 +1,10 @@
 # HogQL -> ClickHouse allowed transformations
-from typing import Optional, Dict, Tuple
+from typing import Optional, Dict, Tuple, Literal
 
 from pydantic import BaseModel, Extra
 
+ConstantDataType = Literal["int", "float", "str", "bool", "array", "tuple", "date", "datetime", "uuid", "unknown"]
+
 CLICKHOUSE_FUNCTIONS: Dict[str, Tuple[str, int | None, int | None]] = {
     # arithmetic
     "plus": ("plus", 2, 2),
@@ -501,6 +503,8 @@
     "argMax": 2,
     "argMin": 2,
     # TODO: more aggregate functions?
+    "groupArray": 1,
+    "sumMap": 2,
 }
 ADD_TIMEZONE_TO_FUNCTIONS = ("now", "now64", "NOW", "toDateTime")
 # Keywords passed to ClickHouse without transformation

posthog/hogql/database.py

@@ -126,7 +126,7 @@ def select_from_persons_table(requested_fields: Dict[str, Any]):
         select_from=ast.JoinExpr(table=ast.Field(chain=["raw_persons"])),
         group_by=[id],
         having=ast.CompareOperation(
-            op=ast.CompareOperationType.Eq,
+            op=ast.CompareOperationOp.Eq,
             left=argmax_version(ast.Field(chain=["is_deleted"])),
             right=ast.Constant(value=0),
         ),
@@ -142,7 +142,7 @@ def join_with_persons_table(from_table: str, to_table: str, requested_fields: Di
     join_expr.join_type = "INNER JOIN"
     join_expr.alias = to_table
     join_expr.constraint = ast.CompareOperation(
-        op=ast.CompareOperationType.Eq,
+        op=ast.CompareOperationOp.Eq,
         left=ast.Field(chain=[from_table, "person_id"]),
         right=ast.Field(chain=[to_table, "id"]),
     )
@@ -203,7 +203,7 @@ def select_from_person_distinct_ids_table(requested_fields: Dict[str, Any]):
         select_from=ast.JoinExpr(table=ast.Field(chain=["raw_person_distinct_ids"])),
         group_by=[distinct_id],
         having=ast.CompareOperation(
-            op=ast.CompareOperationType.Eq,
+            op=ast.CompareOperationOp.Eq,
             left=argmax_version(ast.Field(chain=["is_deleted"])),
             right=ast.Constant(value=0),
         ),
@@ -219,7 +219,7 @@ def join_with_person_distinct_ids_table(from_table: str, to_table: str, requeste
     join_expr.join_type = "INNER JOIN"
     join_expr.alias = to_table
     join_expr.constraint = ast.CompareOperation(
-        op=ast.CompareOperationType.Eq,
+        op=ast.CompareOperationOp.Eq,
         left=ast.Field(chain=[from_table, "distinct_id"]),
         right=ast.Field(chain=[to_table, "distinct_id"]),
     )

posthog/hogql/hogql.py

@@ -18,10 +18,10 @@ def translate_hogql(query: str, context: HogQLContext, dialect: Literal["hogql",
     try:
         # Create a fake query that selects from "events" to have fields to select from.
         context.database = context.database or create_hogql_database(context.team_id)
-        select_query_ref = ast.SelectQueryRef(tables={"events": ast.TableRef(table=context.database.events)})
+        select_query_type = ast.SelectQueryType(tables={"events": ast.TableType(table=context.database.events)})
         node = parse_expr(query, no_placeholders=True)
         select_query = ast.SelectQuery(
-            select=[node], select_from=ast.JoinExpr(table=ast.Field(chain=["events"])), ref=select_query_ref
+            select=[node], select_from=ast.JoinExpr(table=ast.Field(chain=["events"])), type=select_query_type
         )
         prepared_ast: ast.SelectQuery = cast(
             ast.SelectQuery,

posthog/hogql/parse_string.py

@@ -22,7 +22,7 @@ def parse_string(text: str) -> str:
         text = text.replace("{{", "{")
         text = text.replace("\\{", "{")
     else:
-        raise HogQLException(f"Invalid string literal, must start and end with the same quote ref: {text}")
+        raise HogQLException(f"Invalid string literal, must start and end with the same quote type: {text}")
 
     # copied from clickhouse_driver/util/escape.py
     text = text.replace("\\b", "\b")

posthog/hogql/parser.py

@@ -59,7 +59,7 @@ def get_parser(query: str) -> HogQLParser:
 
 
 class HogQLErrorListener(ErrorListener):
-    def syntaxError(self, recognizer, offendingRef, line, column, msg, e):
+    def syntaxError(self, recognizer, offendingType, line, column, msg, e):
         raise SyntaxException(msg, line=line, column=column)
 
 
@@ -372,7 +372,7 @@ def visitColumnExprExtract(self, ctx: HogQLParser.ColumnExprExtractContext):
 
     def visitColumnExprNegate(self, ctx: HogQLParser.ColumnExprNegateContext):
         return ast.BinaryOperation(
-            op=ast.BinaryOperationType.Sub, left=ast.Constant(value=0), right=self.visit(ctx.columnExpr())
+            op=ast.BinaryOperationOp.Sub, left=ast.Constant(value=0), right=self.visit(ctx.columnExpr())
         )
 
     def visitColumnExprSubquery(self, ctx: HogQLParser.ColumnExprSubqueryContext):
@@ -392,11 +392,11 @@ def visitColumnExprCast(self, ctx: HogQLParser.ColumnExprCastContext):
 
     def visitColumnExprPrecedence1(self, ctx: HogQLParser.ColumnExprPrecedence1Context):
         if ctx.SLASH():
-            op = ast.BinaryOperationType.Div
+            op = ast.BinaryOperationOp.Div
         elif ctx.ASTERISK():
-            op = ast.BinaryOperationType.Mult
+            op = ast.BinaryOperationOp.Mult
         elif ctx.PERCENT():
-            op = ast.BinaryOperationType.Mod
+            op = ast.BinaryOperationOp.Mod
         else:
             raise NotImplementedException(f"Unsupported ColumnExprPrecedence1: {ctx.operator.text}")
         left = self.visit(ctx.left)
@@ -405,9 +405,9 @@ def visitColumnExprPrecedence1(self, ctx: HogQLParser.ColumnExprPrecedence1Conte
 
     def visitColumnExprPrecedence2(self, ctx: HogQLParser.ColumnExprPrecedence2Context):
         if ctx.PLUS():
-            op = ast.BinaryOperationType.Add
+            op = ast.BinaryOperationOp.Add
         elif ctx.DASH():
-            op = ast.BinaryOperationType.Sub
+            op = ast.BinaryOperationOp.Sub
         elif ctx.CONCAT():
             raise NotImplementedException(f"Yet unsupported text concat operation: {ctx.operator.text}")
         else:
@@ -418,34 +418,34 @@ def visitColumnExprPrecedence2(self, ctx: HogQLParser.ColumnExprPrecedence2Conte
 
     def visitColumnExprPrecedence3(self, ctx: HogQLParser.ColumnExprPrecedence3Context):
         if ctx.EQ_SINGLE() or ctx.EQ_DOUBLE():
-            op = ast.CompareOperationType.Eq
+            op = ast.CompareOperationOp.Eq
         elif ctx.NOT_EQ():
-            op = ast.CompareOperationType.NotEq
+            op = ast.CompareOperationOp.NotEq
         elif ctx.LT():
-            op = ast.CompareOperationType.Lt
+            op = ast.CompareOperationOp.Lt
         elif ctx.LE():
-            op = ast.CompareOperationType.LtE
+            op = ast.CompareOperationOp.LtE
         elif ctx.GT():
-            op = ast.CompareOperationType.Gt
+            op = ast.CompareOperationOp.Gt
         elif ctx.GE():
-            op = ast.CompareOperationType.GtE
+            op = ast.CompareOperationOp.GtE
         elif ctx.LIKE():
             if ctx.NOT():
-                op = ast.CompareOperationType.NotLike
+                op = ast.CompareOperationOp.NotLike
             else:
-                op = ast.CompareOperationType.Like
+                op = ast.CompareOperationOp.Like
         elif ctx.ILIKE():
             if ctx.NOT():
-                op = ast.CompareOperationType.NotILike
+                op = ast.CompareOperationOp.NotILike
             else:
-                op = ast.CompareOperationType.ILike
+                op = ast.CompareOperationOp.ILike
         elif ctx.IN():
             if ctx.GLOBAL():
                 raise NotImplementedException(f"Unsupported node: IN GLOBAL")
             if ctx.NOT():
-                op = ast.CompareOperationType.NotIn
+                op = ast.CompareOperationOp.NotIn
             else:
-                op = ast.CompareOperationType.In
+                op = ast.CompareOperationOp.In
         else:
             raise NotImplementedException(f"Unsupported ColumnExprPrecedence3: {ctx.getText()}")
         return ast.CompareOperation(left=self.visit(ctx.left), right=self.visit(ctx.right), op=op)
@@ -476,7 +476,7 @@ def visitColumnExprIsNull(self, ctx: HogQLParser.ColumnExprIsNullContext):
         return ast.CompareOperation(
             left=self.visit(ctx.columnExpr()),
             right=ast.Constant(value=None),
-            op=ast.CompareOperationType.NotEq if ctx.NOT() else ast.CompareOperationType.Eq,
+            op=ast.CompareOperationOp.NotEq if ctx.NOT() else ast.CompareOperationOp.Eq,
         )
 
     def visitColumnExprWinFunctionTarget(self, ctx: HogQLParser.ColumnExprWinFunctionTargetContext):
@@ -536,7 +536,7 @@ def visitColumnExprOr(self, ctx: HogQLParser.ColumnExprOrContext):
         return ast.Or(exprs=left_array + right_array)
 
     def visitColumnExprTupleAccess(self, ctx: HogQLParser.ColumnExprTupleAccessContext):
-        raise NotImplementedException(f"Unsupported node: ColumnExprTupleAccess")
+        return ast.TupleAccess(tuple=self.visit(ctx.columnExpr()), index=int(ctx.DECIMAL_LITERAL().getText()))
 
     def visitColumnExprCase(self, ctx: HogQLParser.ColumnExprCaseContext):
         raise NotImplementedException(f"Unsupported node: ColumnExprCase")
@@ -588,12 +588,12 @@ def visitWithExprList(self, ctx: HogQLParser.WithExprListContext):
     def visitWithExprSubquery(self, ctx: HogQLParser.WithExprSubqueryContext):
         subquery = self.visit(ctx.selectUnionStmt())
         name = self.visit(ctx.identifier())
-        return ast.Macro(name=name, expr=subquery, type="subquery")
+        return ast.Macro(name=name, expr=subquery, macro_format="subquery")
 
     def visitWithExprColumn(self, ctx: HogQLParser.WithExprColumnContext):
         expr = self.visit(ctx.columnExpr())
         name = self.visit(ctx.identifier())
-        return ast.Macro(name=name, expr=expr, type="column")
+        return ast.Macro(name=name, expr=expr, macro_format="column")
 
     def visitColumnIdentifier(self, ctx: HogQLParser.ColumnIdentifierContext):
         if ctx.PLACEHOLDER():