gortiz commented on code in PR #15135:
URL: https://github.com/apache/pinot/pull/15135#discussion_r1973274698
##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SpoolIntegrationTest.java:
##########
@@ -117,13 +121,146 @@ public void intermediateSpool()
JsonNode stats = jsonNode.get("stageStats");
assertNoError(jsonNode);
DocumentContext parsed = JsonPath.parse(stats.toString());
- List<Map<String, Object>> stage4On3 = parsed.read("$..[?(@.stage ==
3)]..[?(@.stage == 4)]");
- Assert.assertEquals(stage4On3.size(), 1, "Stage 4 should be descended from
stage 3 exactly once");
- List<Map<String, Object>> stage4On7 = parsed.read("$..[?(@.stage ==
7)]..[?(@.stage == 4)]");
- Assert.assertEquals(stage4On3.size(), 1, "Stage 4 should be descended from
stage 7 exactly once");
+ checkSpoolTimes(parsed, 4, 3, 1);
+ checkSpoolTimes(parsed, 4, 7, 1);
+ checkSpoolSame(parsed, 4, 3, 7);
+ }
+
+ /**
+ * Test a complex with nested spools. Don't try to understand it, just check
that the spools are correct.
+ * The test name corresponds to the PR that fixed the issue.
+ */
+ @Test
+ public void test15135()
+ throws Exception {
+ JsonNode jsonNode = postQuery("SET useSpools = true;\n"
+ + "\n"
+ + "WITH\n"
+ + " q1 AS (\n"
+ + " SELECT ArrTimeBlk as userUUID,\n"
+ + " Dest as deviceOS,\n"
+ + " SUM(ArrTime) AS totalTrips\n"
+ + " FROM mytable\n"
+ + " GROUP BY ArrTimeBlk, Dest\n"
+ + " ),\n"
+ + " q2 AS (\n"
+ + " SELECT userUUID,\n"
+ + " deviceOS,\n"
+ + " SUM(totalTrips) AS totalTrips,\n"
+ + " COUNT(DISTINCT userUUID) AS reach\n"
+ + " FROM q1\n"
+ + " GROUP BY userUUID,\n"
+ + " deviceOS\n"
+ + " ),\n"
+ + " q3 AS (\n"
+ + " SELECT userUUID,\n"
+ + " (totalTrips / reach) AS frequency\n"
+ + " FROM q2\n"
+ + " ),\n"
+ + " q4 AS (\n"
+ + " SELECT rd.userUUID,\n"
+ + " rd.deviceOS,\n"
+ + " rd.totalTrips as totalTrips,\n"
+ + " rd.reach AS reach\n"
+ + " FROM q2 rd\n"
+ + " ),\n"
+ + " q5 AS (\n"
+ + " SELECT userUUID,\n"
+ + " SUM(totalTrips) AS totalTrips\n"
+ + " FROM q4\n"
+ + " GROUP BY userUUID\n"
+ + " ),\n"
+ + " q6 AS (\n"
+ + " SELECT s.userUUID,\n"
+ + " s.totalTrips,\n"
+ + " (s.totalTrips / o.frequency) AS reach,\n"
+ + " 'Traditional TV + OTT' AS deviceOS\n"
+ + " FROM q5 s\n"
+ + " JOIN q3 o ON s.userUUID = o.userUUID\n"
+ + " ),\n"
+ + " q7 AS (\n"
+ + " SELECT rd.userUUID,\n"
+ + " rd.totalTrips,\n"
+ + " rd.reach,\n"
+ + " rd.deviceOS\n"
+ + " FROM q4 rd\n"
+ + " UNION ALL\n"
+ + " SELECT f.userUUID,\n"
+ + " f.totalTrips,\n"
+ + " f.reach,\n"
+ + " f.deviceOS\n"
+ + " FROM q6 f\n"
+ + " ),\n"
+ + " q8 AS (\n"
+ + " SELECT sd.*\n"
+ + " FROM q7 sd\n"
+ + " JOIN (\n"
+ + " SELECT deviceOS,\n"
+ + " PERCENTILETDigest(totalTrips, 20) AS p20\n"
+ + " FROM q7\n"
+ + " GROUP BY deviceOS\n"
+ + " ) q ON sd.deviceOS = q.deviceOS\n"
+ + " )\n"
+ + "SELECT *\n"
+ + "FROM q8");
+ JsonNode stats = jsonNode.get("stageStats");
+ assertNoError(jsonNode);
+ DocumentContext parsed = JsonPath.parse(stats.toString());
+
+ checkSpoolTimes(parsed, 6, 5, 1);
+ checkSpoolTimes(parsed, 6, 14, 1);
+ checkSpoolSame(parsed, 6, 5, 14);
+
+ checkSpoolTimes(parsed, 7, 6, 2);
+
+ checkSpoolTimes(parsed, 4, 3, 1);
+ checkSpoolTimes(parsed, 4, 7, 2); // because there are 2 copies of 7 as
well
+ checkSpoolTimes(parsed, 4, 9, 1);
+ checkSpoolTimes(parsed, 4, 12, 1);
+ checkSpoolTimes(parsed, 4, 18, 1);
+ checkSpoolSame(parsed, 4, 3, 7, 9, 12, 18);
+ }
+
+ private List<Map<String, Object>> findDescendantById(DocumentContext stats,
int parent, int descendant) {
+ return stats.read(parentDescendantJsonPathExpression(parent, descendant));
+ }
+
+ private void checkSpoolTimes(DocumentContext stats, int spoolStageId, int
parent, int times) {
+ List<Map<String, Object>> descendants = findDescendantById(stats, parent,
spoolStageId);
+ Assert.assertEquals(descendants.size(), times, "Stage " + spoolStageId + "
should be descended from stage "
+ + parent + " exactly " + times + " times");
+ Map<String, Object> firstSpool = descendants.get(0);
+ for (int i = 1; i < descendants.size(); i++) {
+ Assert.assertEquals(descendants.get(i), firstSpool, "Stage " +
spoolStageId + " should be the same in "
+ + "all " + times + " descendants");
+ }
+ }
+
+ private void checkSpoolSame(DocumentContext stats, int spoolStageId, int...
parents) {
+ List<Pair<Integer, List<Map<String, Object>>>> spools =
Arrays.stream(parents)
+ .mapToObj(parent -> Pair.of(parent, findDescendantById(stats, parent,
spoolStageId)))
+ .collect(Collectors.toList());
+ Pair<Integer, List<Map<String, Object>>> notEmpty = spools.stream()
+ .filter(s -> !s.getValue().isEmpty())
+ .findFirst()
+ .orElse(null);
+ if (notEmpty == null) {
+ Assert.fail("None of the parent nodes " + Arrays.toString(parents) + "
have a descendant with id "
+ + spoolStageId);
+ }
+ List<Pair<Integer, List<Map<String, Object>>>> allNotEqual =
spools.stream()
+ .filter(s -> !s.getValue().get(0).equals(notEmpty.getValue().get(0)))
+ .collect(Collectors.toList());
+ if (!allNotEqual.isEmpty()) {
+ Assert.fail("The descendant with id " + spoolStageId + " is not the same
in all parent nodes "
+ + spools);
+ }
+ }
- Assert.assertEquals(stage4On3, stage4On7, "Stage 4 should be the same in
both stage 3 and stage 7");
+ @Language("jsonpath")
+ private String parentDescendantJsonPathExpression(int parent, int child) {
Review Comment:
it just returns the same JSON path expression we used before, but with the
ability to change parent and child. The expression says: give me all objects
whose stage id is `child,` and one of its descendants (not just direct parent,
any ancestor) is `parent`.
It isn't super efficient. On my machine, the new test takes 6 seconds to
run, but with this method, the tests are easier to read than before.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
