Add in tests for Colocated joins with table hints by deepthi912 · Pull Request #17846 · apache/pinot

deepthi912 · 2026-03-10T00:29:36Z

In reference to #16603 there have been changes to the colocated joins and the query hints used there. This PR adds in integration tests that can help us test the configs through which the colocated joins can be identified and be enabled for the user so the hash join and regular data shuffles are eliminated as the data and worker operations is colocated on a single server for the same partitions of the tables that would be joined in the query.

codecov-commenter · 2026-03-10T01:15:26Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 63.25%. Comparing base (097a89f) to head (fa1d65d).
⚠️ Report is 10 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #17846      +/-   ##
============================================
- Coverage     63.27%   63.25%   -0.02%     
- Complexity     1466     1481      +15     
============================================
  Files          3190     3190              
  Lines        192101   192287     +186     
  Branches      29433    29471      +38     
============================================
+ Hits         121547   121627      +80     
- Misses        61040    61131      +91     
- Partials       9514     9529      +15

Flag	Coverage Δ
custom-integration1	`100.00% <ø> (ø)`
integration	`100.00% <ø> (ø)`
integration1	`100.00% <ø> (ø)`
integration2	`0.00% <ø> (ø)`
java-11	`63.21% <ø> (-0.01%)`	⬇️
java-21	`63.22% <ø> (-0.03%)`	⬇️
temurin	`63.25% <ø> (-0.02%)`	⬇️
unittests	`63.24% <ø> (-0.02%)`	⬇️
unittests1	`55.58% <ø> (-0.01%)`	⬇️
unittests2	`34.25% <ø> (-0.02%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

...ion-tests/src/test/java/org/apache/pinot/integration/tests/ColocatedJoinIntegrationTest.java

reverse merge

yashmayya

I've added a number of suggestions here, please take a look.

yashmayya · 2026-03-17T14:03:42Z

...test/java/org/apache/pinot/integration/tests/ColocatedJoinMultiPartitionIntegrationTest.java

+  /** Three-table join: result correctness with colocated hint. */
+  @Test
+  public void testThreeTableJoinResultCorrectness()
+      throws Exception {
+    String tableOpt = getTableOptPerTableHint();
+    String sql = "SELECT COUNT(*) FROM userAttributes /*+ " + tableOpt + " */ ua "
+        + "JOIN userGroups /*+ " + tableOpt + " */ ug ON ua.userUUID = ug.userUUID "
+        + "JOIN userFactEvents /*+ " + tableOpt + " */ ue ON ua.userUUID = ue.userUUID";
+    JsonNode result = postQuery(sql);
+    assertNoExceptions(result);
+    long count = getCountFromResult(result);
+    assertTrue(count >= 1, "Three-table join should return at least one row");
+  }


How is this test validating join result correctness?

IMO the result should be manually validated and exact counts should be asserted instead.

Alternatively, we can also validate against results using non colocated join queries.

yashmayya · 2026-03-17T14:37:27Z

...tests/src/test/java/org/apache/pinot/integration/tests/ColocatedJoinIntegrationTestBase.java

+    startZk();
+    startController();
+    startBrokers(1);
+    startServers(1);


How are we even testing colocated joins if there's a single server setup?

yashmayya · 2026-03-17T14:38:07Z

...tests/src/test/java/org/apache/pinot/integration/tests/ColocatedJoinIntegrationTestBase.java

+    }
+  }
+
+  /** Two-table join with mismatched partition_size in hints: planner rejects with partition size mismatch error


Maybe add one for mismatched partition function too?

yashmayya · 2026-03-17T14:38:36Z

...tests/src/test/java/org/apache/pinot/integration/tests/ColocatedJoinIntegrationTestBase.java

+/**
+ * Base integration tests for colocated joins tests(multi-stage engine).
+ */
+public abstract class ColocatedJoinIntegrationTestBase extends BaseClusterIntegrationTestSet {


I think we don't have any coverage for composite join keys (multi column)?

Another case is when table is partitioned on one column but we're using a different column as the join key. We need positive / negative tests for this.

yashmayya · 2026-03-17T14:39:43Z

...tests/src/test/java/org/apache/pinot/integration/tests/ColocatedJoinIntegrationTestBase.java

+/**
+ * Base integration tests for colocated joins tests(multi-stage engine).
+ */
+public abstract class ColocatedJoinIntegrationTestBase extends BaseClusterIntegrationTestSet {


What happens if there are equi + non-equi join conditions with colocated join hints?

yashmayya · 2026-03-17T14:40:36Z

...tests/src/test/java/org/apache/pinot/integration/tests/ColocatedJoinIntegrationTestBase.java

+
+  /** Two-table join with joinOptions colocated but only one table has tableOptions hint (partial config). */
+  @Test
+  public void testTwoTableJoinWithPartialTableOptionsResultCorrectness()


Why / how does this work?

cc - @gortiz

yashmayya · 2026-03-17T14:41:00Z

...tests/src/test/java/org/apache/pinot/integration/tests/ColocatedJoinIntegrationTestBase.java

+  /** Two-table join with joinOptions(is_colocated_by_join_keys='null'): explicit 'null' disables colocated exchange
+   * (plan does not use [PARTITIONED]); query still succeeds with correct result. */
+  @Test
+  public void testTwoTableJoinWithJoinOptionsColocatedKeyExplicitNullResultCorrectness()


What is the point of this test?

yashmayya · 2026-03-17T14:42:25Z

...test/java/org/apache/pinot/integration/tests/ColocatedJoinMultiPartitionIntegrationTest.java

+ * assertion {@code [PARTITIONED]}). Adds third table (userFactEvents) and three-table colocated
+ * join tests.
+ */
+public class ColocatedJoinMultiPartitionIntegrationTest extends ColocatedJoinIntegrationTestBase {


Can we make these tests also support realtime or hybrid tables? The partitioning info is calculated differently in those cases

yashmayya · 2026-03-17T14:42:42Z

...test/java/org/apache/pinot/integration/tests/ColocatedJoinMultiPartitionIntegrationTest.java

+ * assertion {@code [PARTITIONED]}). Adds third table (userFactEvents) and three-table colocated
+ * join tests.
+ */
+public class ColocatedJoinMultiPartitionIntegrationTest extends ColocatedJoinIntegrationTestBase {


Can we also have some tests for partition_parallelism?

This one supports multiple workers per partition

yashmayya · 2026-03-17T14:44:53Z

...tests/src/test/java/org/apache/pinot/integration/tests/ColocatedJoinIntegrationTestBase.java

+    assertNoExceptions(result);
+    String plan = extractImplementationPlan(result);
+    assertNotNull(plan, "implementation plan should be present");
+    assertTrue(plan.contains(PLAN_PARTITIONED_MARKER), message + plan);


This is a very weak assertion. We should check it's on the correct stage / exchange, no unnecessary shuffles elsewhere in the plan, number of PARTITIONED markers etc. Also ideally we want to assert on the actual query stage stats too (validating that raw messages is 0 / there are only in-memory messages across the colocated stages) since this PARTITIONED marker is a weak assertion (you can trace the code path that adds this to the explain plan output).

Add in Integration test cases for Colocated joins with table hints

ffc304a

yashmayya reviewed Mar 10, 2026

View reviewed changes

...ion-tests/src/test/java/org/apache/pinot/integration/tests/ColocatedJoinIntegrationTest.java Outdated Show resolved Hide resolved

...ion-tests/src/test/java/org/apache/pinot/integration/tests/ColocatedJoinIntegrationTest.java Outdated Show resolved Hide resolved

deepthi912 and others added 4 commits March 12, 2026 14:53

Merge pull request #152 from apache/master

e719e92

reverse merge

Fix checkstyle Exceptions and add in some more tests

d67038e

Remove some unused comments

cf466da

Change the class name

fa1d65d

yashmayya requested changes Mar 17, 2026

View reviewed changes

Conversation

deepthi912 commented Mar 10, 2026

Uh oh!

codecov-commenter commented Mar 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

yashmayya left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov-commenter commented Mar 10, 2026 •

edited

Loading