Merge remote-tracking branch 'origin/main' into addEarliest

Author: xinyual

core/src/main/java/org/opensearch/sql/analysis/ExpressionAnalyzer.java

@@ -292,15 +292,38 @@ public Expression visitIn(In node, AnalysisContext context) {
 
   private Expression visitIn(
       UnresolvedExpression field, List<UnresolvedExpression> valueList, AnalysisContext context) {
-    if (valueList.size() == 1) {
-      return visitCompare(new Compare("=", field, valueList.get(0)), context);
-    } else if (valueList.size() > 1) {
-      return DSL.or(
-          visitCompare(new Compare("=", field, valueList.get(0)), context),
-          visitIn(field, valueList.subList(1, valueList.size()), context));
-    } else {
+    if (valueList.isEmpty()) {
       throw new SemanticCheckException("Values in In clause should not be empty");
     }
+
+    Expression[] expressions = new Expression[valueList.size()];
+
+    for (int i = 0; i < expressions.length; i++) {
+      expressions[i] = visitCompare(new Compare("=", field, valueList.get(i)), context);
+    }
+
+    return buildOrTree(expressions, 0, expressions.length);
+  }
+
+  /**
+   * `DSL.or` can only take two arguments. To represent large lists without massive recursion, we
+   * want to represent the expression as a balanced tree. This builds that tree from a node list.
+   *
+   * @param children The list of expressions to merge.
+   * @param start The starting position (inclusive) for the current combination step.
+   * @param end The ending position (exclusive) for the current combination step. If <= start,
+   *     children[start] is returned.
+   * @return The final `DSL.or` expression.
+   */
+  private Expression buildOrTree(Expression[] children, int start, int end) {
+    if (end - start <= 1) {
+      return children[start];
+    }
+    if (end - start == 2) {
+      return DSL.or(children[start], children[end - 1]);
+    }
+    int split = start + (end - start) / 2;
+    return DSL.or(buildOrTree(children, start, split), buildOrTree(children, split, end));
   }
 
   @Override

core/src/main/java/org/opensearch/sql/calcite/udf/udaf/PercentileApproxFunction.java

@@ -5,7 +5,7 @@
 
 package org.opensearch.sql.calcite.udf.udaf;
 
-import com.tdunning.math.stats.AVLTreeDigest;
+import com.tdunning.math.stats.MergingDigest;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
@@ -85,7 +85,7 @@ public void evaluate(double value) {
     public Object value(Object... argList) {
       double percent = (double) argList[1];
       double compression = (double) argList[0];
-      AVLTreeDigest tree = new AVLTreeDigest(compression);
+      MergingDigest tree = new MergingDigest(compression);
       for (Number num : candidate) {
         tree.add(num.doubleValue());
       }

core/src/main/java/org/opensearch/sql/expression/aggregation/PercentileApproximateAggregator.java

@@ -8,7 +8,7 @@
 import static org.opensearch.sql.data.model.ExprValueUtils.doubleValue;
 import static org.opensearch.sql.utils.ExpressionUtils.format;
 
-import com.tdunning.math.stats.AVLTreeDigest;
+import com.tdunning.math.stats.MergingDigest;
 import java.util.List;
 import org.opensearch.sql.common.utils.StringUtils;
 import org.opensearch.sql.data.model.ExprNullValue;
@@ -56,11 +56,11 @@ public String toString() {
   }
 
   /**
-   * PercentileApproximateState is used to store the AVLTreeDigest state for percentile estimation.
+   * PercentileApproximateState is used to store the MergingDigest state for percentile estimation.
    */
-  protected static class PercentileApproximateState extends AVLTreeDigest
+  protected static class PercentileApproximateState extends MergingDigest
       implements AggregationState {
-    // The compression level for the AVLTreeDigest, keep the same default value as OpenSearch core.
+    // The compression level for the MergingDigest, keep the same default value as OpenSearch core.
     public static final double DEFAULT_COMPRESSION = 100.0;
     private final double percent;
 

core/src/test/java/org/opensearch/sql/analysis/ExpressionAnalyzerTest.java

@@ -27,6 +27,7 @@
 import static org.opensearch.sql.expression.DSL.ref;
 
 import com.google.common.collect.ImmutableMap;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -401,6 +402,17 @@ void visit_in() {
         () -> analyze(AstDSL.in(field("integer_value"), Collections.emptyList())));
   }
 
+  @Test
+  void visit_in_large_list() {
+    List<UnresolvedExpression> ints = new ArrayList<>();
+    for (int i = 0; i < 10000; i++) {
+      ints.add(intLiteral(i));
+    }
+
+    // Shouldn't crash
+    analyze(AstDSL.in(field("integer_value"), ints));
+  }
+
   @Test
   void multi_match_expression() {
     assertAnalyzeEqual(

core/src/test/java/org/opensearch/sql/expression/aggregation/PercentileApproxAggregatorTest.java

@@ -68,21 +68,21 @@ public void test_percentile_field_expression_with_user_defined_compression() {
         aggregation(
             DSL.percentile(DSL.ref("integer_value", INTEGER), DSL.literal(50), DSL.literal(0.1)),
             tuples);
-    assertEquals(2.5, result.value());
+    assertEquals(3.0, result.value());
     result =
         aggregation(
             DSL.percentile(DSL.ref("long_value", LONG), DSL.literal(50), DSL.literal(0.1)), tuples);
-    assertEquals(2.5, result.value());
+    assertEquals(3.0, result.value());
     result =
         aggregation(
             DSL.percentile(DSL.ref("double_value", DOUBLE), DSL.literal(50), DSL.literal(0.1)),
             tuples);
-    assertEquals(2.5, result.value());
+    assertEquals(3.0, result.value());
     result =
         aggregation(
             DSL.percentile(DSL.ref("float_value", FLOAT), DSL.literal(50), DSL.literal(0.1)),
             tuples);
-    assertEquals(2.5, result.value());
+    assertEquals(3.0, result.value());
   }
 
   @Test

docs/user/dql/aggregations.rst

@@ -378,6 +378,8 @@ Description
 
 Usage: PERCENTILE(expr, percent) or PERCENTILE_APPROX(expr, percent). Returns the approximate percentile value of `expr` at the specified percentage. `percent` must be a constant between 0 and 100.
 
+Note: From 3.1.0, the percentile implementation is switched to MergingDigest from AVLTreeDigest. Ref `issue link <https://github.com/opensearch-project/OpenSearch/issues/18122>`_.
+
 Example::
 
     os> SELECT gender, percentile(age, 90) as p90 FROM accounts GROUP BY gender;

docs/user/ppl/cmd/stats.rst

@@ -309,6 +309,8 @@ Usage: PERCENTILE(expr, percent) or PERCENTILE_APPROX(expr, percent). Return the
 
 * percent: The number must be a constant between 0 and 100.
 
+Note: From 3.1.0, the percentile implementation is switched to MergingDigest from AVLTreeDigest. Ref `issue link <https://github.com/opensearch-project/OpenSearch/issues/18122>`_.
+
 Example::
 
     os> source=accounts | stats percentile(age, 90) by gender;

integ-test/src/test/java/org/opensearch/sql/ppl/ExplainIT.java

@@ -95,7 +95,6 @@ public void testSortPushDownExplain() throws Exception {
 
   @Test
   public void testLimitPushDownExplain() throws Exception {
-    // TODO fix after https://github.com/opensearch-project/sql/issues/3381
     String expected =
         isCalciteEnabled()
             ? loadFromFile("expectedOutput/calcite/explain_limit_push.json")
@@ -110,6 +109,104 @@ public void testLimitPushDownExplain() throws Exception {
                 + "| fields ageMinus"));
   }
 
+  @Test
+  public void testLimitWithFilterPushdownExplain() throws Exception {
+    String expectedFilterThenLimit =
+        isCalciteEnabled()
+            ? loadFromFile("expectedOutput/calcite/explain_filter_then_limit_push.json")
+            : loadFromFile("expectedOutput/ppl/explain_filter_then_limit_push.json");
+    assertJsonEqualsIgnoreId(
+        expectedFilterThenLimit,
+        explainQueryToString(
+            "source=opensearch-sql_test_index_account"
+                + "| where age > 30 "
+                + "| head 5 "
+                + "| fields age"));
+
+    // The filter in limit-then-filter queries should not be pushed since the current DSL will
+    // execute it as filter-then-limit
+    String expectedLimitThenFilter =
+        isCalciteEnabled()
+            ? loadFromFile("expectedOutput/calcite/explain_limit_then_filter_push.json")
+            : loadFromFile("expectedOutput/ppl/explain_limit_then_filter_push.json");
+    assertJsonEqualsIgnoreId(
+        expectedLimitThenFilter,
+        explainQueryToString(
+            "source=opensearch-sql_test_index_account"
+                + "| head 5 "
+                + "| where age > 30 "
+                + "| fields age"));
+  }
+
+  @Test
+  public void testMultipleLimitExplain() throws Exception {
+    String expected5Then10 =
+        isCalciteEnabled()
+            ? loadFromFile("expectedOutput/calcite/explain_limit_5_10_push.json")
+            : loadFromFile("expectedOutput/ppl/explain_limit_5_10_push.json");
+    assertJsonEqualsIgnoreId(
+        expected5Then10,
+        explainQueryToString(
+            "source=opensearch-sql_test_index_account"
+                + "| head 5 "
+                + "| head 10 "
+                + "| fields age"));
+
+    String expected10Then5 =
+        isCalciteEnabled()
+            ? loadFromFile("expectedOutput/calcite/explain_limit_10_5_push.json")
+            : loadFromFile("expectedOutput/ppl/explain_limit_10_5_push.json");
+    assertJsonEqualsIgnoreId(
+        expected10Then5,
+        explainQueryToString(
+            "source=opensearch-sql_test_index_account"
+                + "| head 10 "
+                + "| head 5 "
+                + "| fields age"));
+
+    String expected10from1then10from2 =
+        isCalciteEnabled()
+            ? loadFromFile("expectedOutput/calcite/explain_limit_10from1_10from2_push.json")
+            : loadFromFile("expectedOutput/ppl/explain_limit_10from1_10from2_push.json");
+    assertJsonEqualsIgnoreId(
+        expected10from1then10from2,
+        explainQueryToString(
+            "source=opensearch-sql_test_index_account"
+                + "| head 10 from 1 "
+                + "| head 10 from 2 "
+                + "| fields age"));
+
+    // The second limit should not be pushed down for limit-filter-limit queries
+    String expected10ThenFilterThen5 =
+        isCalciteEnabled()
+            ? loadFromFile("expectedOutput/calcite/explain_limit_10_filter_5_push.json")
+            : loadFromFile("expectedOutput/ppl/explain_limit_10_filter_5_push.json");
+    assertJsonEqualsIgnoreId(
+        expected10ThenFilterThen5,
+        explainQueryToString(
+            "source=opensearch-sql_test_index_account"
+                + "| head 10 "
+                + "| where age > 30 "
+                + "| head 5 "
+                + "| fields age"));
+  }
+
+  @Test
+  public void testLimitWithMultipleOffsetPushdownExplain() throws Exception {
+    String expected =
+        isCalciteEnabled()
+            ? loadFromFile("expectedOutput/calcite/explain_limit_offsets_push.json")
+            : loadFromFile("expectedOutput/ppl/explain_limit_offsets_push.json");
+
+    assertJsonEqualsIgnoreId(
+        expected,
+        explainQueryToString(
+            "source=opensearch-sql_test_index_account"
+                + "| head 10 from 1 "
+                + "| head 5 from 2 "
+                + "| fields age"));
+  }
+
   @Test
   public void testFillNullPushDownExplain() throws Exception {
     String expected =

integ-test/src/test/java/org/opensearch/sql/ppl/StatsCommandIT.java

@@ -436,9 +436,17 @@ public void testStatsPercentileWithNull() throws IOException {
   public void testStatsPercentileWithCompression() throws IOException {
     JSONObject response =
         executeQuery(
-            String.format("source=%s | stats percentile(balance, 50, 1)", TEST_INDEX_BANK));
-    verifySchema(response, schema("percentile(balance, 50, 1)", null, "bigint"));
+            String.format("source=%s | stats percentile(balance, 50, 20)", TEST_INDEX_BANK));
+    verifySchema(response, schema("percentile(balance, 50, 20)", null, "bigint"));
     verifyDataRows(response, rows(32838));
+
+    // disable pushdown by adding a eval command
+    JSONObject responsePushdownDisabled =
+        executeQuery(
+            String.format(
+                "source=%s | eval a = 1 | stats percentile(balance, 50, 20)", TEST_INDEX_BANK));
+    verifySchema(responsePushdownDisabled, schema("percentile(balance, 50, 20)", null, "bigint"));
+    verifyDataRows(responsePushdownDisabled, rows(32838));
   }
 
   @Test

integ-test/src/test/resources/expectedOutput/calcite/explain_filter_then_limit_push.json

@@ -0,0 +1,6 @@
+{
+  "calcite": {
+    "logical": "LogicalProject(age=[$8])\n  LogicalSort(fetch=[5])\n    LogicalFilter(condition=[>($8, 30)])\n      CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
+    "physical": "CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[FILTER->>($8, 30), LIMIT->5, PROJECT->[age]], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"size\":5,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"_source\":{\"includes\":[\"age\"],\"excludes\":[]},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}]}, requestedTotalSize=5, pageSize=null, startFrom=0)])\n"
+  }
+}