diff --git a/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlV2Test.java b/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlV2Test.java index ca0c524b4a1..ec3c36daa89 100644 --- a/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlV2Test.java +++ b/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlV2Test.java @@ -507,6 +507,46 @@ SELECT name, ROW_NUMBER() OVER (ORDER BY id) AS rn FROM catalog.employees """); } + @Test + public void testWindowRank() { + givenQuery( + """ + SELECT name, RANK() OVER (ORDER BY age) AS r FROM catalog.employees + """) + .assertPlan( + """ + LogicalProject(name=[$1], r=[RANK() OVER (ORDER BY $2 NULLS FIRST)]) + LogicalTableScan(table=[[catalog, employees]]) + """); + } + + @Test + public void testWindowDenseRank() { + givenQuery( + """ + SELECT name, DENSE_RANK() OVER (ORDER BY age) AS r FROM catalog.employees + """) + .assertPlan( + """ + LogicalProject(name=[$1], r=[DENSE_RANK() OVER (ORDER BY $2 NULLS FIRST)]) + LogicalTableScan(table=[[catalog, employees]]) + """); + } + + @Test + public void testWindowRankPartitionBy() { + givenQuery( + """ + SELECT name, RANK() OVER (PARTITION BY department ORDER BY age DESC) AS r + FROM catalog.employees + """) + .assertPlan( + """ + LogicalProject(name=[$1], r=[RANK() OVER (PARTITION BY $3 ORDER BY $2 DESC NULLS FIRST)]) + LogicalTableScan(table=[[catalog, employees]]) + """); + } + @Test public void testGroupByExpression() { givenQuery("SELECT LENGTH(name), COUNT(*) FROM catalog.employees GROUP BY LENGTH(name)") diff --git a/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerTest.java b/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerTest.java index 008121e8377..aab37d8a6c4 100644 --- a/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerTest.java +++ b/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerTest.java @@ -165,10 +165,10 @@ public void unsupportedWindowFunctionIsRethrownAsSemanticCheckException() { // CalciteRexNodeVisitor#visitWindowFunction's // orElseThrow. The throw site emits CalciteUnsupportedException so this path normalizes to a // 4xx SemanticCheckException rather than escaping as a 500. - givenInvalidQuery("source = catalog.employees | eventstats rank()") + givenInvalidQuery("source = catalog.employees | eventstats percent_rank()") .assertErrorType(SemanticCheckException.class) .assertCauseType(CalciteUnsupportedException.class) - .assertErrorMessageContains("Unexpected window function: rank"); + .assertErrorMessageContains("Unexpected window function: percent_rank"); } @Test diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java index 1bf1e217b51..b6500dd8087 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java @@ -770,9 +770,11 @@ public RexNode visitWindowFunction(WindowFunction node, CalcitePlanContext conte (arguments.isEmpty() || arguments.size() == 1) ? Collections.emptyList() : arguments.subList(1, arguments.size()); - // ROW_NUMBER takes no field/args and isn't in aggFunctionRegistry, + // These take no field/args and aren't in aggFunctionRegistry, // so skip aggregate signature validation. - if (functionName == BuiltinFunctionName.ROW_NUMBER) { + if (functionName == BuiltinFunctionName.ROW_NUMBER + || functionName == BuiltinFunctionName.RANK + || functionName == BuiltinFunctionName.DENSE_RANK) { return PlanUtils.makeOver( context, functionName, diff --git a/core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java b/core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java index f899f747421..10865d2532e 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java +++ b/core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java @@ -232,6 +232,23 @@ static RexNode makeOver( true, lowerBound, upperBound); + // Calcite rank operators disallow framing, so the ROWS/RANGE flag below is normalized away. + case RANK: + return withOver( + context.relBuilder.aggregateCall(SqlStdOperatorTable.RANK), + partitions, + orderKeys, + false, + lowerBound, + upperBound); + case DENSE_RANK: + return withOver( + context.relBuilder.aggregateCall(SqlStdOperatorTable.DENSE_RANK), + partitions, + orderKeys, + false, + lowerBound, + upperBound); case NTH_VALUE: return withOver( context.relBuilder.aggregateCall(SqlStdOperatorTable.NTH_VALUE, field, argList.get(0)), diff --git a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java index e30d723ccfc..88fca8a9aec 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java @@ -429,6 +429,8 @@ public enum BuiltinFunctionName { .put("distinct_count", BuiltinFunctionName.DISTINCT_COUNT_APPROX) .put("pattern", BuiltinFunctionName.INTERNAL_PATTERN) .put("row_number", BuiltinFunctionName.ROW_NUMBER) + .put("rank", BuiltinFunctionName.RANK) + .put("dense_rank", BuiltinFunctionName.DENSE_RANK) .build(); public static Optional of(String str) { diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/WindowFunctionIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/WindowFunctionIT.java index ad248b98d62..7fa9637760c 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/WindowFunctionIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/WindowFunctionIT.java @@ -188,4 +188,93 @@ public void testPercentilePartition() { rows("Duke Willmington", 5686), rows("Ratliff", 16418)); } + + @Test + public void testRankOverNull() { + JSONObject response = + new JSONObject( + executeQuery( + """ + SELECT lastname, RANK() OVER() FROM %s\ + """ + .formatted(TestsConstants.TEST_INDEX_BANK), + "jdbc")); + + verifyDataRows( + response, + rows("Duke Willmington", 1), + rows("Bond", 1), + rows("Bates", 1), + rows("Adams", 1), + rows("Ratliff", 1), + rows("Ayala", 1), + rows("Mcpherson", 1)); + } + + @Test + public void testRankOver() { + JSONObject response = + new JSONObject( + executeQuery( + """ + SELECT age, RANK() OVER(ORDER BY age DESC) FROM %s\ + """ + .formatted(TestsConstants.TEST_INDEX_BANK), + "jdbc")); + + verifyDataRows( + response, + rows(39, 1), + rows(36, 2), + rows(36, 2), + rows(34, 4), + rows(33, 5), + rows(32, 6), + rows(28, 7)); + } + + @Test + public void testRankPartition() { + JSONObject response = + new JSONObject( + executeQuery( + """ + SELECT lastname, RANK() OVER(PARTITION BY gender ORDER BY age DESC)\ + FROM %s\ + """ + .formatted(TestsConstants.TEST_INDEX_BANK), + "jdbc")); + + verifyDataRows( + response, + rows("Bond", 1), + rows("Ratliff", 1), + rows("Adams", 3), + rows("Duke Willmington", 4), + rows("Ayala", 1), + rows("Mcpherson", 2), + rows("Bates", 3)); + } + + @Test + public void testDenseRankOver() { + JSONObject response = + new JSONObject( + executeQuery( + """ + SELECT age, DENSE_RANK() OVER(ORDER BY age DESC) FROM %s\ + """ + .formatted(TestsConstants.TEST_INDEX_BANK), + "jdbc")); + + verifyDataRows( + response, + rows(39, 1), + rows(36, 2), + rows(36, 2), + rows(34, 3), + rows(33, 4), + rows(32, 5), + rows(28, 6)); + } }