feat: add support for self-joins with sub-query

document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/JoinExpression.java

@@ -0,0 +1,38 @@
+package org.hypertrace.core.documentstore.expression.impl;
+
+import lombok.AccessLevel;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Value;
+import org.hypertrace.core.documentstore.expression.type.FilterTypeExpression;
+import org.hypertrace.core.documentstore.expression.type.FromTypeExpression;
+import org.hypertrace.core.documentstore.parser.FilterTypeExpressionVisitor;
+import org.hypertrace.core.documentstore.parser.FromTypeExpressionVisitor;
+import org.hypertrace.core.documentstore.parser.SelectTypeExpressionVisitor;
+import org.hypertrace.core.documentstore.query.Query;
+
+/**
+ * Expression representing a condition for filtering
+ *
+ * <p>Example: <code>
+ * company IN ('Traceable', 'Harness')
+ * </code> can be constructed as <code>
+ * RelationalExpression.of( IdentifierExpression.of("company"), RelationalOperator.IN,
+ * ConstantExpression.ofStrings("Traceable", "Harness"))));
+ * </code>
+ */
+@Value
+@Builder(toBuilder = true)
+@AllArgsConstructor(access = AccessLevel.PRIVATE)
+public class JoinExpression implements FromTypeExpression {
+  FromTypeExpression left;
+  FromTypeExpression right;
+  JoinType joinType;                    // INNER, LEFT, RIGHT, etc.
+  FilterTypeExpression onCondition;     // The ON clause
+
+  @Override
+  public <T> T accept(FromTypeExpressionVisitor visitor) {
+    return visitor.visit(this);
+  }
+}
+

document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/JoinType.java

@@ -0,0 +1,9 @@
+package org.hypertrace.core.documentstore.expression.impl;
+
+public enum JoinType {
+  INNER,
+  LEFT,
+  RIGHT,
+  FULL,
+  CROSS
+}

document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/SubQueryFromExpression.java

@@ -0,0 +1,26 @@
+package org.hypertrace.core.documentstore.expression.impl;
+
+import lombok.AccessLevel;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Value;
+import org.hypertrace.core.documentstore.expression.type.FromTypeExpression;
+import org.hypertrace.core.documentstore.parser.FromTypeExpressionVisitor;
+import org.hypertrace.core.documentstore.query.Query;
+
+/**
+ * Allows embedding an entire Query in the FROM clause.
+ */
+@Value
+@Builder(toBuilder = true)
+@AllArgsConstructor(access = AccessLevel.PRIVATE)
+public class SubQueryFromExpression implements FromTypeExpression {
+  Query subQuery;
+  String alias; // e.g. "rightTable"
+
+  @Override
+  public <T> T accept(FromTypeExpressionVisitor visitor) {
+    return visitor.visit(this);
+  }
+}
+

document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/TableFromExpression.java

@@ -0,0 +1,24 @@
+package org.hypertrace.core.documentstore.expression.impl;
+
+import lombok.AccessLevel;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Value;
+import org.hypertrace.core.documentstore.expression.type.FromTypeExpression;
+import org.hypertrace.core.documentstore.parser.FromTypeExpressionVisitor;
+
+/**
+ * A simple wrapper for referencing a physical collection/table in the FROM clause.
+ */
+@Value
+@Builder(toBuilder = true)
+@AllArgsConstructor(access = AccessLevel.PRIVATE)
+public class TableFromExpression implements FromTypeExpression {
+  // Currently, this supports joins on the same table/collection only. In the future, we can also add a tableName field here to support joins across tables/collections.
+  String alias; // e.g. "leftTable"
+
+  @Override
+  public <T> T accept(FromTypeExpressionVisitor visitor) {
+    return visitor.visit(this);
+  }
+}

document-store/src/main/java/org/hypertrace/core/documentstore/parser/FromTypeExpressionVisitor.java

@@ -1,7 +1,16 @@
 package org.hypertrace.core.documentstore.parser;
 
+import org.hypertrace.core.documentstore.expression.impl.JoinExpression;
+import org.hypertrace.core.documentstore.expression.impl.SubQueryFromExpression;
+import org.hypertrace.core.documentstore.expression.impl.TableFromExpression;
 import org.hypertrace.core.documentstore.expression.impl.UnnestExpression;
 
 public interface FromTypeExpressionVisitor {
   <T> T visit(UnnestExpression unnestExpression);
+
+  <T> T visit(JoinExpression joinExpression);
+
+  <T> T visit(TableFromExpression tableFromExpression);
+
+  <T> T visit(SubQueryFromExpression subQueryFromExpression);
 }

document-store/src/test/java/org/hypertrace/core/documentstore/expression/impl/JoinExpressionTest.java

@@ -0,0 +1,107 @@
+package org.hypertrace.core.documentstore.expression.impl;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import org.hypertrace.core.documentstore.expression.operators.AggregationOperator;
+import org.hypertrace.core.documentstore.expression.operators.RelationalOperator;
+import org.hypertrace.core.documentstore.expression.type.FilterTypeExpression;
+import org.hypertrace.core.documentstore.expression.type.FromTypeExpression;
+import org.hypertrace.core.documentstore.query.FromClause;
+import org.hypertrace.core.documentstore.query.Query;
+import org.hypertrace.core.documentstore.query.SelectionSpec;
+
+class JoinExpressionTest {
+
+  /*
+This is the query we want to execute:
+SELECT sa.suite_id, sa.vulnerability_count, sa.scan_run_number
+FROM scan_analytics sa
+JOIN (
+    SELECT suite_id, MAX(scan_run_number) AS latest_scan_run_number
+    FROM scan_analytics
+    GROUP BY suite_id
+) latest
+ON sa.suite_id = latest.suite_id
+AND sa.scan_run_number = latest.latest_scan_run_number;
+   */
+
+  void exampleUsage() {
+    // The right subquery:
+// SELECT suite_id, MAX(scan_run_number) AS latest_scan_run_number
+// FROM scan_analytics
+// GROUP BY suite_id
+    Query subQuery = Query.builder()
+        // SELECT
+        .addSelection(
+            SelectionSpec.of(IdentifierExpression.of("suite_id"), "suite_id")
+        )
+        .addSelection(
+            SelectionSpec.of(
+                AggregateExpression.of(AggregationOperator.MAX, IdentifierExpression.of("scan_run_number")),
+                "latest_scan_run_number"
+        ))
+        // FROM
+        .addFromClause(
+            TableFromExpression.builder()
+                .alias("ignored")  // or skip if not needed
+                .build()
+        )
+        // GROUP BY
+        .addAggregation(
+            IdentifierExpression.of("suite_id")
+        )
+        .build();
+
+// The main FROM side: "scan_analytics sa"
+    FromTypeExpression leftTable = TableFromExpression.builder()
+        .alias("sa")
+        .build();
+
+// The subquery side: "(...subQuery...) latest"
+    FromTypeExpression rightSubQuery = SubQueryFromExpression.builder()
+        .subQuery(subQuery)
+        .alias("latest")
+        .build();
+
+// The ON condition:
+//   sa.suite_id = latest.suite_id
+//   AND sa.scan_run_number = latest.latest_scan_run_number
+    FilterTypeExpression onCondition = LogicalExpression.and(
+            RelationalExpression.of(
+                IdentifierExpression.of("sa.suite_id"),
+                RelationalOperator.EQ,
+                IdentifierExpression.of("latest.suite_id")
+                ),
+            RelationalExpression.of(
+                IdentifierExpression.of("a.scan_run_number"),
+                RelationalOperator.EQ,
+                IdentifierExpression.of("latest.latest_scan_run_number")
+            )
+    );
+
+    JoinExpression joinExpression = JoinExpression.builder()
+        .left(leftTable)
+        .right(rightSubQuery)
+        .joinType(JoinType.INNER)
+        .onCondition(onCondition)
+        .build();
+
+// Now build the top-level Query:
+// SELECT sa.suite_id, sa.vulnerability_count, sa.scan_run_number
+    Query mainQuery = Query.builder()
+        .addSelection(
+            SelectionSpec.of(IdentifierExpression.of("sa.suite_id"), "suite_id")
+        )
+        .addSelection(
+            SelectionSpec.of(IdentifierExpression.of("sa.vulnerability_count"), "vulnerability_count")
+        )
+        .addSelection(
+            SelectionSpec.of(IdentifierExpression.of("sa.scan_run_number"), "scan_run_number")
+        )
+        .addFromClause(joinExpression)
+        .build();
+
+    System.out.println(mainQuery);
+
+  }
+}