Address comments

qianheng-aws · qianheng-aws · commit c8f2a7f083db · 2025-12-24T14:47:28.000+08:00
Signed-off-by: Heng Qian &lt;qianheng@amazon.com&gt;
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/planner/rules/DedupPushdownRule.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/planner/rules/DedupPushdownRule.java
@@ -125,10 +125,8 @@ protected void apply(
     PlanUtils.addIgnoreNullBucketHintToAggregate(relBuilder);
     // peek the aggregate after hint being added
     LogicalAggregate aggregate = (LogicalAggregate) relBuilder.build();
-    // assert aggregate.getInput().getRowType().getFieldCount() == groupByList.size() :
-    // String.format("The input's field size should be trimmed to equal to group list size %d, but
-    // got %d", groupByList.size(), aggregate.getInput().getRowType().getFieldCount());
-    assert aggregate.getGroupSet().asList().equals(newGroupByList);
+    assert aggregate.getGroupSet().asList().equals(newGroupByList)
+        : "The group set of aggregate should be exactly the same as the generated group list";
 
     CalciteLogicalIndexScan newScan =
         (CalciteLogicalIndexScan) scan.pushDownAggregate(aggregate, targetChildProject);
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java
@@ -9,7 +9,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Optional;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import lombok.Getter;
@@ -403,19 +402,15 @@ public AbstractRelNode pushDownLimit(LogicalSort sort, Integer limit, Integer of
         // Since the AggPushDownAction is shared among different PushDownContext, its size() may be
         // inaccurate(<= the actual size).
         // So take the previous limit into account to decide whether it can update the context.
-        boolean canReduceEstimatedRowsCount = true;
-        if (pushDownContext.isLimitPushed()) {
-          Optional<Integer> previousRowCount =
-              pushDownContext.getQueue().reversed().stream()
-                  .takeWhile(operation -> operation.type() != PushDownType.AGGREGATION)
-                  .filter(operation -> operation.type() == PushDownType.LIMIT)
-                  .findFirst()
-                  .map(operation -> (LimitDigest) operation.digest())
-                  .map(limitDigest -> limitDigest.offset() + limitDigest.limit());
-          if (previousRowCount.isPresent()) {
-            canReduceEstimatedRowsCount = totalSize < previousRowCount.get();
-          }
-        }
+        boolean canReduceEstimatedRowsCount =
+            !pushDownContext.isLimitPushed()
+                || pushDownContext.getQueue().reversed().stream()
+                    .takeWhile(op -> op.type() != PushDownType.AGGREGATION)
+                    .filter(op -> op.type() == PushDownType.LIMIT)
+                    .findFirst()
+                    .map(op -> (LimitDigest) op.digest())
+                    .map(d -> totalSize < d.offset() + d.limit())
+                    .orElse(true);
 
         // Push down the limit into the aggregation bucket in advance to detect whether the limit
         // can update the aggregation builder
