Merge remote-tracking branch 'upstream/main' into merge_back

Signed-off-by: Lantao Jin <ltjin@amazon.com>

Author: LantaoJin

core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java

@@ -27,9 +27,12 @@
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexLambdaRef;
+import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.SqlIntervalQualifier;
+import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.type.ArraySqlType;
 import org.apache.calcite.sql.type.SqlTypeName;
@@ -190,8 +193,15 @@ public RexNode visitXor(Xor node, CalcitePlanContext context) {
 
   @Override
   public RexNode visitNot(Not node, CalcitePlanContext context) {
-    final RexNode expr = analyze(node.getExpression(), context);
-    return context.relBuilder.not(expr);
+    // Special handling for NOT(boolean_field = true/false) - see boolean comparison helpers below
+    UnresolvedExpression inner = node.getExpression();
+    if (inner instanceof Compare compare && "=".equals(compare.getOperator())) {
+      RexNode result = tryMakeBooleanNotEquals(compare, context);
+      if (result != null) {
+        return result;
+      }
+    }
+    return context.relBuilder.not(analyze(node.getExpression(), context));
   }
 
   @Override
@@ -221,7 +231,15 @@ public RexNode visitIn(In node, CalcitePlanContext context) {
   public RexNode visitCompare(Compare node, CalcitePlanContext context) {
     RexNode left = analyze(node.getLeft(), context);
     RexNode right = analyze(node.getRight(), context);
-    return PPLFuncImpTable.INSTANCE.resolve(context.rexBuilder, node.getOperator(), left, right);
+    String op = node.getOperator();
+    // Handle boolean_field != literal -> IS_NOT_TRUE/IS_NOT_FALSE
+    if ("!=".equals(op) || "<>".equals(op)) {
+      RexNode result = tryMakeBooleanNotEquals(left, right, context);
+      if (result != null) {
+        return result;
+      }
+    }
+    return PPLFuncImpTable.INSTANCE.resolve(context.rexBuilder, op, left, right);
   }
 
   @Override
@@ -251,6 +269,65 @@ public RexNode visitEqualTo(EqualTo node, CalcitePlanContext context) {
     return context.rexBuilder.equals(left, right);
   }
 
+  // ==================== Boolean NOT comparison helpers ====================
+  // Calcite's RexSimplify transforms:
+  // - "field = true" -> "field" (handled by PredicateAnalyzer detecting boolean field)
+  // - "field = false" -> "NOT(field)" (handled by PredicateAnalyzer.prefix())
+  // - "NOT(field = true)" -> "NOT(field)" -> would generate term{false}, have conflicted semantics
+  // - "NOT(field = false)" -> "NOT(NOT(field))" -> "field" -> would generate term{true}, have
+  // conflicted semantics
+  // We intercept NOT(field = true/false) at AST level before Calcite optimization:
+  // - "NOT(field = true)" -> IS_NOT_TRUE(field): matches false, null, missing
+  // - "NOT(field = false)" -> IS_NOT_FALSE(field): matches true, null, missing
+
+  /**
+   * Try to convert boolean_field != literal or NOT(boolean_field = literal) to
+   * IS_NOT_TRUE/IS_NOT_FALSE. This preserves correct null-handling semantics.
+   */
+  private RexNode tryMakeBooleanNotEquals(RexNode left, RexNode right, CalcitePlanContext context) {
+    BooleanFieldComparison cmp = extractBooleanFieldComparison(left, right);
+    if (cmp == null) {
+      return null;
+    }
+    SqlOperator op =
+        Boolean.FALSE.equals(cmp.literalValue)
+            ? SqlStdOperatorTable.IS_NOT_FALSE
+            : SqlStdOperatorTable.IS_NOT_TRUE;
+    return context.rexBuilder.makeCall(op, cmp.field);
+  }
+
+  /** Overload for NOT(Compare) AST pattern. */
+  private RexNode tryMakeBooleanNotEquals(Compare compare, CalcitePlanContext context) {
+    return tryMakeBooleanNotEquals(
+        analyze(compare.getLeft(), context), analyze(compare.getRight(), context), context);
+  }
+
+  /** Represents a comparison between a boolean field and a boolean literal. */
+  private record BooleanFieldComparison(RexNode field, Boolean literalValue) {}
+
+  /**
+   * Extract boolean field and literal value from a comparison, normalizing operand order. Returns
+   * null if the comparison is not between a boolean field and a boolean literal.
+   */
+  private BooleanFieldComparison extractBooleanFieldComparison(RexNode left, RexNode right) {
+    if (isBooleanField(left) && isBooleanLiteral(right)) {
+      return new BooleanFieldComparison(left, ((RexLiteral) right).getValueAs(Boolean.class));
+    }
+    if (isBooleanField(right) && isBooleanLiteral(left)) {
+      return new BooleanFieldComparison(right, ((RexLiteral) left).getValueAs(Boolean.class));
+    }
+    return null;
+  }
+
+  private boolean isBooleanField(RexNode node) {
+    // Only match actual field references, not arbitrary boolean expressions like CASE
+    return node instanceof RexInputRef && node.getType().getSqlTypeName() == SqlTypeName.BOOLEAN;
+  }
+
+  private boolean isBooleanLiteral(RexNode node) {
+    return node instanceof RexLiteral && node.getType().getSqlTypeName() == SqlTypeName.BOOLEAN;
+  }
+
   /** Resolve qualified name. Note, the name should be case-sensitive. */
   @Override
   public RexNode visitQualifiedName(QualifiedName node, CalcitePlanContext context) {

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

@@ -5,6 +5,7 @@
 
 package org.opensearch.sql.calcite.remote;
 
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_ALIAS;
 import static org.opensearch.sql.util.MatcherUtils.rows;
 import static org.opensearch.sql.util.MatcherUtils.schema;
 import static org.opensearch.sql.util.MatcherUtils.verifyDataRows;
@@ -25,13 +26,14 @@
  */
 public class CalciteAliasFieldAggregationIT extends PPLIntegTestCase {
 
-  private static final String TEST_INDEX_ALIAS = "test_alias_bug";
+  private static final String TEST_ALIAS_BUG = "test_alias_bug";
 
   @Override
   public void init() throws Exception {
     super.init();
     enableCalcite();
     createTestIndexWithAliasFields();
+    loadIndex(Index.DATA_TYPE_ALIAS);
   }
 
   /**
@@ -41,14 +43,14 @@ public void init() throws Exception {
   private void createTestIndexWithAliasFields() throws IOException {
     // Delete the index if it exists (for test isolation)
     try {
-      Request deleteIndex = new Request("DELETE", "/" + TEST_INDEX_ALIAS);
+      Request deleteIndex = new Request("DELETE", "/" + TEST_ALIAS_BUG);
       client().performRequest(deleteIndex);
     } catch (ResponseException e) {
       // Index doesn't exist, which is fine
     }
 
     // Create index with alias fields
-    Request createIndex = new Request("PUT", "/" + TEST_INDEX_ALIAS);
+    Request createIndex = new Request("PUT", "/" + TEST_ALIAS_BUG);
     createIndex.setJsonEntity(
         "{\n"
             + "  \"mappings\": {\n"
@@ -63,7 +65,7 @@ private void createTestIndexWithAliasFields() throws IOException {
     client().performRequest(createIndex);
 
     // Insert test documents
-    Request bulkRequest = new Request("POST", "/" + TEST_INDEX_ALIAS + "/_bulk?refresh=true");
+    Request bulkRequest = new Request("POST", "/" + TEST_ALIAS_BUG + "/_bulk?refresh=true");
     bulkRequest.setJsonEntity(
         "{\"index\":{}}\n"
             + "{\"created_at\": \"2024-01-01T10:00:00Z\", \"value\": 100}\n"
@@ -77,15 +79,15 @@ private void createTestIndexWithAliasFields() throws IOException {
   @Test
   public void testMinWithDateAliasField() throws IOException {
     JSONObject actual =
-        executeQuery(String.format("source=%s | stats MIN(@timestamp)", TEST_INDEX_ALIAS));
+        executeQuery(String.format("source=%s | stats MIN(@timestamp)", TEST_ALIAS_BUG));
     verifySchema(actual, schema("MIN(@timestamp)", "timestamp"));
     verifyDataRows(actual, rows("2024-01-01 10:00:00"));
   }
 
   @Test
   public void testMaxWithDateAliasField() throws IOException {
     JSONObject actual =
-        executeQuery(String.format("source=%s | stats MAX(@timestamp)", TEST_INDEX_ALIAS));
+        executeQuery(String.format("source=%s | stats MAX(@timestamp)", TEST_ALIAS_BUG));
     verifySchema(actual, schema("MAX(@timestamp)", "timestamp"));
     verifyDataRows(actual, rows("2024-01-03 10:00:00"));
   }
@@ -94,8 +96,7 @@ public void testMaxWithDateAliasField() throws IOException {
   public void testMinMaxWithNumericAliasField() throws IOException {
     JSONObject actual =
         executeQuery(
-            String.format(
-                "source=%s | stats MIN(value_alias), MAX(value_alias)", TEST_INDEX_ALIAS));
+            String.format("source=%s | stats MIN(value_alias), MAX(value_alias)", TEST_ALIAS_BUG));
     verifySchemaInOrder(
         actual, schema("MIN(value_alias)", "int"), schema("MAX(value_alias)", "int"));
     verifyDataRows(actual, rows(100, 300));
@@ -105,8 +106,7 @@ public void testMinMaxWithNumericAliasField() throws IOException {
   public void testFirstWithAliasField() throws IOException {
     JSONObject actual =
         executeQuery(
-            String.format(
-                "source=%s | sort @timestamp | stats FIRST(@timestamp)", TEST_INDEX_ALIAS));
+            String.format("source=%s | sort @timestamp | stats FIRST(@timestamp)", TEST_ALIAS_BUG));
     verifySchema(actual, schema("FIRST(@timestamp)", "timestamp"));
     verifyDataRows(actual, rows("2024-01-01 10:00:00"));
   }
@@ -115,8 +115,7 @@ public void testFirstWithAliasField() throws IOException {
   public void testLastWithAliasField() throws IOException {
     JSONObject actual =
         executeQuery(
-            String.format(
-                "source=%s | sort @timestamp | stats LAST(@timestamp)", TEST_INDEX_ALIAS));
+            String.format("source=%s | sort @timestamp | stats LAST(@timestamp)", TEST_ALIAS_BUG));
     verifySchema(actual, schema("LAST(@timestamp)", "timestamp"));
     verifyDataRows(actual, rows("2024-01-03 10:00:00"));
   }
@@ -126,7 +125,7 @@ public void testTakeWithAliasField() throws IOException {
     JSONObject actual =
         executeQuery(
             String.format(
-                "source=%s | sort @timestamp | stats TAKE(@timestamp, 2)", TEST_INDEX_ALIAS));
+                "source=%s | sort @timestamp | stats TAKE(@timestamp, 2)", TEST_ALIAS_BUG));
     verifySchema(actual, schema("TAKE(@timestamp, 2)", "array"));
     verifyDataRows(actual, rows(List.of("2024-01-01T10:00:00.000Z", "2024-01-02T10:00:00.000Z")));
   }
@@ -135,7 +134,7 @@ public void testTakeWithAliasField() throws IOException {
   public void testAggregationsWithOriginalFieldsStillWork() throws IOException {
     JSONObject actual =
         executeQuery(
-            String.format("source=%s | stats MIN(created_at), MAX(value)", TEST_INDEX_ALIAS));
+            String.format("source=%s | stats MIN(created_at), MAX(value)", TEST_ALIAS_BUG));
     verifySchemaInOrder(
         actual, schema("MIN(created_at)", "timestamp"), schema("MAX(value)", "int"));
     verifyDataRows(actual, rows("2024-01-01 10:00:00", 300));
@@ -147,12 +146,50 @@ public void testUnaffectedAggregationsWithAliasFields() throws IOException {
         executeQuery(
             String.format(
                 "source=%s | stats SUM(value_alias), AVG(value_alias), COUNT(value_alias)",
-                TEST_INDEX_ALIAS));
+                TEST_ALIAS_BUG));
     verifySchemaInOrder(
         actual,
         schema("SUM(value_alias)", "bigint"),
         schema("AVG(value_alias)", "double"),
         schema("COUNT(value_alias)", "bigint"));
     verifyDataRows(actual, rows(600, 200.0, 3));
   }
+
+  @Test
+  public void testAliasTypeWithLastFirstTakeLatestEarliestAggregation() throws IOException {
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | stats take(original_text, 2), last(original_text),"
+                    + " first(original_text), take(alias_text, 2), last(alias_text),"
+                    + " first(alias_text), take(original_col, 2), last(original_col),"
+                    + " first(original_col), take(alias_col, 2), last(alias_col), first(alias_col),"
+                    + " latest(original_col), earliest(original_col), latest(alias_col),"
+                    + " earliest(alias_col),latest(original_text), earliest(original_text),"
+                    + " latest(alias_text), earliest(alias_text)",
+                TEST_INDEX_ALIAS));
+    verifyDataRows(
+        actual,
+        rows(
+            List.of("a b c", "d e f"),
+            "x y z",
+            "a b c",
+            List.of("a b c", "d e f"),
+            "x y z",
+            "a b c",
+            List.of(1, 2),
+            3,
+            1,
+            List.of(1, 2),
+            3,
+            1,
+            3,
+            1,
+            3,
+            1,
+            "x y z",
+            "a b c",
+            "x y z",
+            "a b c"));
+  }
 }