feat(hogql): support join USING

hogql_parser/CONTRIBUTING.md

@@ -68,3 +68,9 @@ The three pages below are must-reads though. They're key to writing production-r
     ```bash
     pytest posthog/hogql/
     ```
+
+## How to install dependencies on Ubuntu
+
+Antlr runtime provided in Ubuntu packages might be of an older version, which results in compilation errors.
+
+In that case run commands from [this step](https://github.com/PostHog/posthog/blob/4fba6a63e351131fdb27b85e7ba436446fdb3093/.github/actions/run-backend-tests/action.yml#L100).

hogql_parser/parser.cpp

@@ -857,9 +857,6 @@ class HogQLParseTreeConverter : public HogQLParserBaseVisitor {
   VISIT_UNSUPPORTED(JoinOpCross)
 
   VISIT(JoinConstraintClause) {
-    if (ctx->USING()) {
-      throw NotImplementedError("Unsupported: JOIN ... USING");
-    }
     PyObject* column_expr_list = visitAsPyObject(ctx->columnExprList());
     Py_ssize_t column_expr_list_size = PyList_Size(column_expr_list);
     if (column_expr_list_size == -1) {
@@ -872,7 +869,7 @@ class HogQLParseTreeConverter : public HogQLParserBaseVisitor {
     }
     PyObject* expr = Py_NewRef(PyList_GET_ITEM(column_expr_list, 0));
     Py_DECREF(column_expr_list);
-    RETURN_NEW_AST_NODE("JoinConstraint", "{s:N}", "expr", expr);
+    RETURN_NEW_AST_NODE("JoinConstraint", "{s:N,s:s}", "expr", expr, "constraint_type", ctx->USING() ? "USING" : "ON");
   }
 
   VISIT(SampleClause) {

mypy-baseline.txt

@@ -217,8 +217,6 @@ posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression
 posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Resolver" has incompatible type "JoinExpr | None"; expected "Expr"  [arg-type]
 posthog/hogql/resolver.py:0: error: Statement is unreachable  [unreachable]
 posthog/hogql/resolver.py:0: error: Item "None" of "JoinExpr | None" has no attribute "join_type"  [union-attr]
-posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "JoinConstraint | None")  [assignment]
-posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Resolver" has incompatible type "JoinConstraint | None"; expected "Expr"  [arg-type]
 posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "SampleExpr | None")  [assignment]
 posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Resolver" has incompatible type "SampleExpr | None"; expected "Expr"  [arg-type]
 posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Visitor" has incompatible type "SelectQuery | SelectUnionQuery | Field | None"; expected "AST"  [arg-type]
@@ -230,8 +228,6 @@ posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression
 posthog/hogql/resolver.py:0: error: Argument 1 to "append" of "list" has incompatible type "BaseTableType | SelectUnionQueryType | SelectQueryType | SelectQueryAliasType | SelectViewType | None"; expected "SelectQueryType | SelectUnionQueryType"  [arg-type]
 posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "JoinExpr | None")  [assignment]
 posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Resolver" has incompatible type "JoinExpr | None"; expected "Expr"  [arg-type]
-posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "JoinConstraint | None")  [assignment]
-posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Resolver" has incompatible type "JoinConstraint | None"; expected "Expr"  [arg-type]
 posthog/hogql/resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "SampleExpr | None")  [assignment]
 posthog/hogql/resolver.py:0: error: Argument 1 to "visit" of "Resolver" has incompatible type "SampleExpr | None"; expected "Expr"  [arg-type]
 posthog/hogql/resolver.py:0: error: Argument 2 to "convert_hogqlx_tag" has incompatible type "int | None"; expected "int"  [arg-type]
@@ -514,6 +510,13 @@ posthog/hogql/test/test_resolver.py:0: error: Incompatible types in assignment (
 posthog/hogql/test/test_resolver.py:0: error: "TestResolver" has no attribute "snapshot"  [attr-defined]
 posthog/hogql/test/test_resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "SelectQuery")  [assignment]
 posthog/hogql/test/test_resolver.py:0: error: "TestResolver" has no attribute "snapshot"  [attr-defined]
+posthog/hogql/test/test_resolver.py:0: error: Item "None" of "JoinExpr | None" has no attribute "next_join"  [union-attr]
+posthog/hogql/test/test_resolver.py:0: error: Item "None" of "JoinExpr | Any | None" has no attribute "constraint"  [union-attr]
+posthog/hogql/test/test_resolver.py:0: error: Item "None" of "JoinConstraint | Any | None" has no attribute "constraint_type"  [union-attr]
+posthog/hogql/test/test_resolver.py:0: error: Item "None" of "JoinConstraint | Any | None" has no attribute "expr"  [union-attr]
+posthog/hogql/test/test_resolver.py:0: error: Item "None" of "JoinExpr | None" has no attribute "next_join"  [union-attr]
+posthog/hogql/test/test_resolver.py:0: error: Item "None" of "JoinExpr | Any | None" has no attribute "constraint"  [union-attr]
+posthog/hogql/test/test_resolver.py:0: error: Item "None" of "JoinConstraint | Any | None" has no attribute "constraint_type"  [union-attr]
 posthog/hogql/test/test_resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "SelectQuery")  [assignment]
 posthog/hogql/test/test_resolver.py:0: error: "TestResolver" has no attribute "snapshot"  [attr-defined]
 posthog/hogql/test/test_resolver.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "SelectQuery")  [assignment]

posthog/hogql/ast.py

@@ -571,6 +571,7 @@ class Call(Expr):
 @dataclass(kw_only=True)
 class JoinConstraint(Expr):
     expr: Expr
+    constraint_type: Literal["ON", "USING"]
 
 
 @dataclass(kw_only=True)

posthog/hogql/database/schema/event_sessions.py

@@ -205,7 +205,8 @@ def join_with_events_table_session_duration(
             op=ast.CompareOperationOp.Eq,
             left=ast.Field(chain=[from_table, "$session_id"]),
             right=ast.Field(chain=[to_table, "id"]),
-        )
+        ),
+        constraint_type="ON",
     )
 
     return join_expr

posthog/hogql/database/schema/groups.py

@@ -61,7 +61,8 @@ def join_with_group_table(
                 op=ast.CompareOperationOp.Eq,
                 left=ast.Field(chain=[from_table, f"$group_{group_index}"]),
                 right=ast.Field(chain=[to_table, "key"]),
-            )
+            ),
+            constraint_type="ON",
         )
 
         return join_expr

posthog/hogql/database/schema/person_distinct_id_overrides.py

@@ -58,7 +58,8 @@ def join_with_person_distinct_id_overrides_table(
             op=ast.CompareOperationOp.Eq,
             left=ast.Field(chain=[from_table, "distinct_id"]),
             right=ast.Field(chain=[to_table, "distinct_id"]),
-        )
+        ),
+        constraint_type="ON",
     )
     return join_expr
 

posthog/hogql/database/schema/person_distinct_ids.py

@@ -58,7 +58,8 @@ def join_with_person_distinct_ids_table(
             op=ast.CompareOperationOp.Eq,
             left=ast.Field(chain=[from_table, "distinct_id"]),
             right=ast.Field(chain=[to_table, "distinct_id"]),
-        )
+        ),
+        constraint_type="ON",
     )
     return join_expr
 

posthog/hogql/database/schema/person_overrides.py

@@ -53,7 +53,8 @@ def join_with_person_overrides_table(
             op=ast.CompareOperationOp.Eq,
             left=ast.Field(chain=[from_table, "event_person_id"]),
             right=ast.Field(chain=[to_table, "old_person_id"]),
-        )
+        ),
+        constraint_type="ON",
     )
     return join_expr
 

posthog/hogql/database/schema/persons.py

@@ -100,7 +100,8 @@ def join_with_persons_table(
             op=ast.CompareOperationOp.Eq,
             left=ast.Field(chain=[from_table, "person_id"]),
             right=ast.Field(chain=[to_table, "id"]),
-        )
+        ),
+        constraint_type="ON",
     )
     return join_expr
 

posthog/hogql/database/schema/persons_pdi.py

@@ -47,7 +47,8 @@ def persons_pdi_join(
             op=ast.CompareOperationOp.Eq,
             left=ast.Field(chain=[from_table, "id"]),
             right=ast.Field(chain=[to_table, "person_id"]),
-        )
+        ),
+        constraint_type="ON",
     )
     return join_expr
 

posthog/hogql/database/schema/session_replay_events.py

@@ -51,7 +51,8 @@ def join_with_events_table(
             op=ast.CompareOperationOp.Eq,
             left=ast.Field(chain=[from_table, "session_id"]),
             right=ast.Field(chain=[to_table, "$session_id"]),
-        )
+        ),
+        constraint_type="ON",
     )
 
     return join_expr
@@ -103,7 +104,8 @@ def join_with_console_logs_log_entries_table(
             op=ast.CompareOperationOp.Eq,
             left=ast.Field(chain=[from_table, "session_id"]),
             right=ast.Field(chain=[to_table, "log_source_id"]),
-        )
+        ),
+        constraint_type="ON",
     )
 
     return join_expr

posthog/hogql/database/schema/sessions.py

@@ -265,7 +265,8 @@ def join_events_table_to_sessions_table(
             op=ast.CompareOperationOp.Eq,
             left=ast.Field(chain=[from_table, "$session_id"]),
             right=ast.Field(chain=[to_table, "session_id"]),
-        )
+        ),
+        constraint_type="ON",
     )
     return join_expr
 

posthog/hogql/parser.py

@@ -366,12 +366,10 @@ def visitJoinOpCross(self, ctx: HogQLParser.JoinOpCrossContext):
         raise NotImplementedError(f"Unsupported node: JoinOpCross")
 
     def visitJoinConstraintClause(self, ctx: HogQLParser.JoinConstraintClauseContext):
-        if ctx.USING():
-            raise NotImplementedError(f"Unsupported: JOIN ... USING")
         column_expr_list = self.visit(ctx.columnExprList())
         if len(column_expr_list) != 1:
             raise NotImplementedError(f"Unsupported: JOIN ... ON with multiple expressions")
-        return ast.JoinConstraint(expr=column_expr_list[0])
+        return ast.JoinConstraint(expr=column_expr_list[0], constraint_type="USING" if ctx.USING() else "ON")
 
     def visitSampleClause(self, ctx: HogQLParser.SampleClauseContext):
         ratio_expressions = ctx.ratioExpr()

posthog/hogql/printer.py

@@ -452,7 +452,7 @@ def visit_join_expr(self, node: ast.JoinExpr) -> JoinExprResponse:
                 join_strings.append(sample_clause)
 
         if node.constraint is not None:
-            join_strings.append(f"ON {self.visit(node.constraint)}")
+            join_strings.append(f"{node.constraint.constraint_type} {self.visit(node.constraint)}")
 
         return JoinExprResponse(printed_sql=" ".join(join_strings), where=extra_where)
 

posthog/hogql/resolver.py

@@ -275,12 +275,17 @@
                 return response
 
         if isinstance(node.table, ast.Field):
+            node = cast(ast.JoinExpr, clone_expr(node))
+            if node.constraint and node.constraint.constraint_type == "USING":
+                # visit USING constraint before adding the table to avoid ambiguous names
+                node.constraint = self.visit_join_constraint(node.constraint)
+
             table_name = node.table.chain[0]
             table_alias = node.alias or table_name
             if table_alias in scope.tables:
                 raise QueryError(f'Already have joined a table called "{table_alias}". Can\'t redefine.')
 
             database_table = self.database.get_table(table_name)
 
             if isinstance(database_table, SavedQuery):
                 self.current_view_depth += 1
@@ -293,7 +298,7 @@
                 if isinstance(node.table, ast.SelectQuery):
                     node.table.view_name = database_table.name
 
                 node.alias = table_alias or database_table.name
                 node = self.visit(node)
 
                 self.current_view_depth -= 1
@@ -307,13 +312,13 @@
             # Always add an alias for function call tables. This way `select table.* from table` is replaced with
             # `select table.* from something() as table`, and not with `select something().* from something()`.
             if table_alias != table_name or isinstance(database_table, FunctionCallTable):
                 node_type = ast.TableAliasType(alias=table_alias, table_type=node_table_type)
             else:
                 node_type = node_table_type
+
             scope.tables[table_alias] = node_type
 
             # :TRICKY: Make sure to clone and visit _all_ JoinExpr fields/nodes.
-            node = cast(ast.JoinExpr, clone_expr(node))
             node.type = node_type
             node.table = cast(ast.Field, clone_expr(node.table))
             node.table.type = node_table_type
@@ -330,7 +335,8 @@
             if is_global:
                 node.next_join.join_type = "GLOBAL JOIN"
 
-            node.constraint = self.visit(node.constraint)
+            if node.constraint and node.constraint.constraint_type == "ON":
+                node.constraint = self.visit_join_constraint(node.constraint)
             node.sample = self.visit(node.sample)
 
             # In case we had a function call table, and had to add an alias where none was present, mark it here
@@ -341,6 +347,9 @@
 
         elif isinstance(node.table, ast.SelectQuery) or isinstance(node.table, ast.SelectUnionQuery):
             node = cast(ast.JoinExpr, clone_expr(node))
+            if node.constraint and node.constraint.constraint_type == "USING":
+                # visit USING constraint before adding the table to avoid ambiguous names
+                node.constraint = self.visit_join_constraint(node.constraint)
 
             node.table = super().visit(node.table)
             if isinstance(node.table, ast.SelectQuery) and node.table.view_name is not None and node.alias is not None:
@@ -365,7 +374,8 @@
 
             # :TRICKY: Make sure to clone and visit _all_ JoinExpr fields/nodes.
             node.next_join = self.visit(node.next_join)
-            node.constraint = self.visit(node.constraint)
+            if node.constraint and node.constraint.constraint_type == "ON":
+                node.constraint = self.visit_join_constraint(node.constraint)
             node.sample = self.visit(node.sample)
 
             return node

posthog/hogql/test/_test_parser.py

@@ -790,7 +790,7 @@ def test_select_from_join(self):
                         next_join=ast.JoinExpr(
                             join_type="JOIN",
                             table=ast.Field(chain=["events2"]),
-                            constraint=ast.JoinConstraint(expr=ast.Constant(value=1)),
+                            constraint=ast.JoinConstraint(expr=ast.Constant(value=1), constraint_type="ON"),
                         ),
                     ),
                 ),
@@ -804,7 +804,7 @@ def test_select_from_join(self):
                         next_join=ast.JoinExpr(
                             join_type="LEFT OUTER JOIN",
                             table=ast.Field(chain=["events2"]),
-                            constraint=ast.JoinConstraint(expr=ast.Constant(value=1)),
+                            constraint=ast.JoinConstraint(expr=ast.Constant(value=1), constraint_type="ON"),
                         ),
                     ),
                 ),
@@ -818,16 +818,30 @@ def test_select_from_join(self):
                         next_join=ast.JoinExpr(
                             join_type="LEFT OUTER JOIN",
                             table=ast.Field(chain=["events2"]),
-                            constraint=ast.JoinConstraint(expr=ast.Constant(value=1)),
+                            constraint=ast.JoinConstraint(expr=ast.Constant(value=1), constraint_type="ON"),
                             next_join=ast.JoinExpr(
                                 join_type="RIGHT ANY JOIN",
                                 table=ast.Field(chain=["events3"]),
-                                constraint=ast.JoinConstraint(expr=ast.Constant(value=2)),
+                                constraint=ast.JoinConstraint(expr=ast.Constant(value=2), constraint_type="ON"),
                             ),
                         ),
                     ),
                 ),
             )
+            self.assertEqual(
+                self._select("select 1 from events JOIN events2 USING 1"),
+                ast.SelectQuery(
+                    select=[ast.Constant(value=1)],
+                    select_from=ast.JoinExpr(
+                        table=ast.Field(chain=["events"]),
+                        next_join=ast.JoinExpr(
+                            join_type="JOIN",
+                            table=ast.Field(chain=["events2"]),
+                            constraint=ast.JoinConstraint(expr=ast.Constant(value=1), constraint_type="USING"),
+                        ),
+                    ),
+                ),
+            )
 
         def test_select_from_join_multiple(self):
             node = self._select(
@@ -863,7 +877,8 @@ def test_select_from_join_multiple(self):
                                     op=ast.CompareOperationOp.Eq,
                                     left=ast.Field(chain=["pdi", "distinct_id"]),
                                     right=ast.Field(chain=["e", "distinct_id"]),
-                                )
+                                ),
+                                constraint_type="ON",
                             ),
                             next_join=ast.JoinExpr(
                                 join_type="LEFT JOIN",
@@ -874,7 +889,8 @@ def test_select_from_join_multiple(self):
                                         op=ast.CompareOperationOp.Eq,
                                         left=ast.Field(chain=["p", "id"]),
                                         right=ast.Field(chain=["pdi", "person_id"]),
-                                    )
+                                    ),
+                                    constraint_type="ON",
                                 ),
                             ),
                         ),

posthog/hogql/test/test_resolver.py

@@ -270,6 +270,19 @@ def test_ctes_with_union_all(self):
                         """),
         )
 
+    def test_join_using(self):
+        node = self._select(
+            "WITH my_table AS (SELECT 1 AS a) SELECT q1.a FROM my_table AS q1 INNER JOIN my_table AS q2 USING a"
+        )
+        node = cast(ast.SelectQuery, resolve_types(node, self.context, dialect="clickhouse"))
+        constraint = cast(ast.SelectQuery, node).select_from.next_join.constraint
+        assert constraint.constraint_type == "USING"
+        assert cast(ast.Field, cast(ast.Alias, constraint.expr).expr).chain == ["a"]
+
+        node = self._select("SELECT q1.event FROM events AS q1 INNER JOIN events AS q2 USING event")
+        node = cast(ast.SelectQuery, resolve_types(node, self.context, dialect="clickhouse"))
+        assert cast(ast.SelectQuery, node).select_from.next_join.constraint.constraint_type == "USING"
+
     @override_settings(PERSON_ON_EVENTS_OVERRIDE=False, PERSON_ON_EVENTS_V2_OVERRIDE=False)
     @pytest.mark.usefixtures("unittest_snapshot")
     def test_asterisk_expander_table(self):

posthog/hogql/test/test_visitor.py

@@ -85,7 +85,8 @@ def test_everything_visitor(self):
                                     op=ast.CompareOperationOp.Eq,
                                     left=ast.Field(chain=["d"]),
                                     right=ast.Field(chain=["e"]),
-                                )
+                                ),
+                                constraint_type="ON",
                             ),
                         ),
                         sample=ast.SampleExpr(

posthog/hogql/transforms/in_cohort.py

@@ -237,7 +237,8 @@ def get_dynamic_cohort_clause():
                         op=ast.CompareOperationOp.Eq,
                         left=ast.Constant(value=1),
                         right=ast.Constant(value=1),
-                    )
+                    ),
+                    constraint_type="ON",
                 ),
             )
 
@@ -377,7 +378,8 @@ def _add_join_for_cohort(
                         op=ast.CompareOperationOp.Eq,
                         left=ast.Constant(value=1),
                         right=ast.Constant(value=1),
-                    )
+                    ),
+                    constraint_type="ON",
                 ),
             )
             new_join = cast(

posthog/hogql/visitor.py

@@ -528,8 +528,8 @@ def visit_window_frame_expr(self, node: ast.WindowFrameExpr):
             frame_value=node.frame_value,
         )
 
-    def visit_join_constraint(self, node: ast.JoinConstraint):
-        return ast.JoinConstraint(expr=self.visit(node.expr))
+    def visit_join_constraint(self, node: ast.JoinConstraint) -> ast.JoinConstraint:
+        return ast.JoinConstraint(expr=self.visit(node.expr), constraint_type=node.constraint_type)
 
     def visit_hogqlx_tag(self, node: ast.HogQLXTag):
         return ast.HogQLXTag(kind=node.kind, attributes=[self.visit(a) for a in node.attributes])

posthog/hogql_queries/actors_query_runner.py

@@ -156,7 +156,8 @@ def source_table_join(self) -> ast.JoinExpr:
                         op=ast.CompareOperationOp.Eq,
                         left=ast.Field(chain=[self.strategy.origin, self.strategy.origin_id]),
                         right=ast.Field(chain=[source_alias, *source_id_chain]),
-                    )
+                    ),
+                    constraint_type="ON",
                 ),
             ),
         )