[Backport 2.19-dev] Pushdown case function in aggregations as range queries (#4400)

core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java

@@ -1032,19 +1032,15 @@ void populate() {
           XOR,
           SqlStdOperatorTable.NOT_EQUALS,
           PPLTypeChecker.family(SqlTypeFamily.BOOLEAN, SqlTypeFamily.BOOLEAN));
-      // SqlStdOperatorTable.CASE.getOperandTypeChecker is null. We manually create a
-      // type checker
-      // for it. The second and third operands are required to be of the same type. If
-      // not,
-      // it will throw an IllegalArgumentException with information Can't find
-      // leastRestrictive type
+      // SqlStdOperatorTable.CASE.getOperandTypeChecker is null. We manually create a type checker
+      // for it. The second and third operands are required to be of the same type. If not,  it will
+      // throw an IllegalArgumentException with information Can't find leastRestrictive type
       registerOperator(
           IF,
           SqlStdOperatorTable.CASE,
           PPLTypeChecker.family(SqlTypeFamily.BOOLEAN, SqlTypeFamily.ANY, SqlTypeFamily.ANY));
       // Re-define the type checker for is not null, is present, and is null since
-      // their original
-      // type checker ANY isn't compatible with struct types.
+      // their original type checker ANY isn't compatible with struct types.
       registerOperator(
           IS_NOT_NULL,
           SqlStdOperatorTable.IS_NOT_NULL,

docs/user/ppl/functions/condition.rst

@@ -207,6 +207,14 @@ Argument type: all the supported data type, (NOTE : there is no comma before "el
 
 Return type: any
 
+Limitations
+>>>>>>>>>>>
+
+When each condition is a field comparison with a numeric literal and each result expression is a string literal, the query will be optimized as `range aggregations <https://docs.opensearch.org/latest/aggregations/bucket/range>`_ if pushdown optimization is enabled. However, this optimization has the following limitations:
+
+- Null values will not be grouped into any bucket of a range aggregation and will be ignored
+- The default ELSE clause will use the string literal ``"null"`` instead of actual NULL values
+
 Example::
 
     os> source=accounts | eval result = case(age > 35, firstname, age < 30, lastname else employer) | fields result, firstname, lastname, age, employer

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java

@@ -10,6 +10,7 @@
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_LOGS;
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_NESTED_SIMPLE;
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_STRINGS;
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_TIME_DATA;
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_WEBLOGS;
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_WORKER;
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_WORK_INFORMATION;
@@ -18,6 +19,7 @@
 
 import java.io.IOException;
 import java.util.Locale;
+import org.junit.Assume;
 import org.junit.Ignore;
 import org.junit.Test;
 import org.opensearch.sql.ppl.ExplainIT;
@@ -512,22 +514,6 @@ public void testExplainStatsWithSubAggregation() throws IOException {
                 + " @timestamp, region"));
   }
 
-  @Test
-  public void bucketNullableNotSupportSubAggregation() throws IOException {
-    // TODO: Don't throw exception after addressing
-    // https://github.com/opensearch-project/sql/issues/4317
-    // When bucketNullable is true, sub aggregation is not supported. Hence we cannot pushdown the
-    // aggregation in this query. Caused by issue
-    // https://github.com/opensearch-project/sql/issues/4317,
-    // bin aggregation on timestamp field won't work if not been push down.
-    enabledOnlyWhenPushdownIsEnabled();
-    assertThrows(
-        Exception.class,
-        () ->
-            explainQueryToString(
-                "source=events | bin @timestamp bins=3 | stats count() by @timestamp, region"));
-  }
-
   @Test
   public void testExplainBinWithSpan() throws IOException {
     String expected = loadExpectedPlan("explain_bin_span.yaml");
@@ -1169,4 +1155,118 @@ public void testReplaceCommandExplain() throws IOException {
                 "source=%s | replace 'IL' WITH 'Illinois' IN state | fields state",
                 TEST_INDEX_ACCOUNT)));
   }
+
+  @Test
+  public void testCasePushdownAsRangeQueryExplain() throws IOException {
+    // CASE 1: Range - Metric
+    // 1.1 Range - Metric
+    assertYamlEqualsIgnoreId(
+        loadExpectedPlan("agg_range_metric_push.yaml"),
+        explainQueryYaml(
+            String.format(
+                "source=%s | eval age_range = case(age < 30, 'u30', age < 40, 'u40' else 'u100') |"
+                    + " stats avg(age) as avg_age by age_range",
+                TEST_INDEX_BANK)));
+
+    // 1.2 Range - Metric (COUNT)
+    assertYamlEqualsIgnoreId(
+        loadExpectedPlan("agg_range_count_push.yaml"),
+        explainQueryYaml(
+            String.format(
+                "source=%s | eval age_range = case(age < 30, 'u30', age >= 30 and age < 40, 'u40'"
+                    + " else 'u100') | stats avg(age) by age_range",
+                TEST_INDEX_BANK)));
+
+    // 1.3 Range - Range - Metric
+    assertYamlEqualsIgnoreId(
+        loadExpectedPlan("agg_range_range_metric_push.yaml"),
+        explainQueryYaml(
+            String.format(
+                "source=%s | eval age_range = case(age < 30, 'u30', age < 40, 'u40' else 'u100'),"
+                    + " balance_range = case(balance < 20000, 'medium' else 'high') | stats"
+                    + " avg(balance) as avg_balance by age_range, balance_range",
+                TEST_INDEX_BANK)));
+
+    // 1.5 Should not be pushed because the range is not closed-open
+    assertYamlEqualsIgnoreId(
+        loadExpectedPlan("agg_case_cannot_push.yaml"),
+        explainQueryYaml(
+            String.format(
+                "source=%s | eval age_range = case(age < 30, 'u30', age >= 30 and age <= 40, 'u40'"
+                    + " else 'u100') | stats avg(age) as avg_age by age_range",
+                TEST_INDEX_BANK)));
+
+    // 1.6 Should not be pushed as range query because the result expression is not a string
+    // literal.
+    // Range aggregation keys must be strings
+    assertYamlEqualsIgnoreId(
+        loadExpectedPlan("agg_case_num_res_cannot_push.yaml"),
+        explainQueryYaml(
+            String.format(
+                "source=%s | eval age_range = case(age < 30, 30 else 100) | stats count() by"
+                    + " age_range",
+                TEST_INDEX_BANK)));
+
+    // CASE 2: Composite - Range - Metric
+    // 2.1 Composite (term) - Range - Metric
+    assertYamlEqualsIgnoreId(
+        loadExpectedPlan("agg_composite_range_metric_push.yaml"),
+        explainQueryYaml(
+            String.format(
+                "source=%s | eval age_range = case(age < 30, 'u30' else 'a30') | stats avg(balance)"
+                    + " by state, age_range",
+                TEST_INDEX_BANK)));
+
+    // 2.2 Composite (date histogram) - Range - Metric
+    assertYamlEqualsIgnoreId(
+        loadExpectedPlan("agg_composite_date_range_push.yaml"),
+        explainQueryYaml(
+            "source=opensearch-sql_test_index_time_data | eval value_range = case(value < 7000,"
+                + " 'small' else 'large') | stats avg(value) by value_range, span(@timestamp,"
+                + " 1h)"));
+
+    // 2.3 Composite(2 fields) - Range - Metric (with count)
+    assertYamlEqualsIgnoreId(
+        loadExpectedPlan("agg_composite2_range_count_push.yaml"),
+        explainQueryYaml(
+            String.format(
+                "source=%s | eval age_range = case(age < 30, 'u30' else 'a30') | stats"
+                    + " avg(balance), count() by age_range, state, gender",
+                TEST_INDEX_BANK)));
+
+    // 2.4 Composite (2 fields) - Range - Range - Metric (with count)
+    assertYamlEqualsIgnoreId(
+        loadExpectedPlan("agg_composite2_range_range_count_push.yaml"),
+        explainQueryYaml(
+            String.format(
+                "source=%s | eval age_range = case(age < 35, 'u35' else 'a35'), balance_range ="
+                    + " case(balance < 20000, 'medium' else 'high') | stats avg(balance) as"
+                    + " avg_balance by age_range, balance_range, state",
+                TEST_INDEX_BANK)));
+
+    // 2.5 Should not be pushed down as range query because case result expression is not constant
+    assertYamlEqualsIgnoreId(
+        loadExpectedPlan("agg_case_composite_cannot_push.yaml"),
+        explainQueryYaml(
+            String.format(
+                "source=%s | eval age_range = case(age < 35, 'u35' else email) | stats avg(balance)"
+                    + " as avg_balance by age_range, state",
+                TEST_INDEX_BANK)));
+  }
+
+  @Test
+  public void testNestedAggregationsExplain() throws IOException {
+    // TODO: Remove after resolving: https://github.com/opensearch-project/sql/issues/4578
+    Assume.assumeFalse(
+        "The query runs into error when pushdown is disabled due to bin's implementation",
+        isPushdownDisabled());
+    assertYamlEqualsIgnoreId(
+        loadExpectedPlan("agg_composite_autodate_range_metric_push.yaml"),
+        explainQueryYaml(
+            String.format(
+                "source=%s | bin timestamp bins=3 | eval value_range = case(value < 7000, 'small'"
+                    + " else 'great') | stats bucket_nullable=false avg(value), count() by"
+                    + " timestamp, value_range, category",
+                TEST_INDEX_TIME_DATA)));
+  }
 }