Address comments

Author: iwanttobepowerful

core/src/main/java/org/apache/calcite/sql2rel/RelDecorrelator.java

@@ -575,7 +575,30 @@ protected RexNode removeCorrelationExpr(
     }
 
     if (isCorVarDefined && (rel.fetch != null || rel.offset != null)) {
-      return decorrelateFetchOneSort(rel, frame);
+      if (rel.offset == null
+          && rel.fetch != null
+          && RexLiteral.longValue(rel.fetch) == 0) {
+        return null;
+      }
+
+      //
+      // Rewrite logic:
+      //
+      // For correlated Sort with LIMIT/OFFSET:
+      // Special case: if OFFSET is null and FETCH = 1,
+      // we may rewrite as an Aggregate using MIN/MAX.
+      Frame aggFrame = decorrelateSortAsAggregate(rel, frame);
+      if (aggFrame != null) {
+        return aggFrame;
+      }
+
+      // General case: rewrite as
+      //   Project(original_fields..., corVars..., rn)
+      //     where rn = ROW_NUMBER() OVER (PARTITION BY corVars ORDER BY sortExprs
+      //                                   ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)
+      //   Filter(rn > offset, rn <= offset + fetch)
+      // This preserves per-corVar LIMIT/OFFSET semantics.
+      return decorrelateSortWithRowNumber(rel, frame);
     }
 
     final RelNode newInput = frame.r;
@@ -1036,32 +1059,7 @@ private static void shiftMapping(Map<Integer, Integer> mapping, int startIndex,
     return null;
   }
 
-  protected @Nullable Frame decorrelateFetchOneSort(Sort sort, final Frame frame) {
-    if (sort.offset == null
-        && sort.fetch != null
-        && RexLiteral.intValue(sort.fetch) == 0) {
-      return null;
-    }
-
-    //
-    // Rewrite logic:
-    //
-    // For correlated Sort with LIMIT/OFFSET:
-    //   1) Special case: if OFFSET is null and FETCH = 1, we may rewrite as an Aggregate
-    //      using MIN/MAX (see decorrelateSortAsAggregate).
-    //   2) General case: rewrite as
-    //        Project(original_fields..., corVars..., rn)
-    //          where rn = ROW_NUMBER() OVER (PARTITION BY corVars ORDER BY sortExprs
-    //                                        ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)
-    //        Filter(rn > offset, rn <= offset + fetch)
-    //      This preserves per-corVar LIMIT/OFFSET semantics.
-    //
-
-    Frame aggFrame = decorrelateSortAsAggregate(sort, frame);
-    if (aggFrame != null) {
-      return aggFrame;
-    }
-
+  protected @Nullable Frame decorrelateSortWithRowNumber(Sort sort, final Frame frame) {
     final Map<Integer, Integer> mapOldToNewOutputs = new HashMap<>();
     final NavigableMap<CorDef, Integer> corDefOutputs = new TreeMap<>();
 
@@ -1143,7 +1141,7 @@ private static void shiftMapping(Map<Integer, Integer> mapping, int startIndex,
     if (sort.offset != null) {
       return null;
     }
-    if (sort.fetch == null || RexLiteral.intValue(sort.fetch) != 1) {
+    if (sort.fetch == null || RexLiteral.longValue(sort.fetch) != 1) {
       return null;
     }
 
@@ -1376,7 +1374,7 @@ private static void shiftMapping(Map<Integer, Integer> mapping, int startIndex,
       final int newOutput =
           newLocalOutputs.indexOf(newLocalOutput)
               + requireNonNull(mapNewInputToNewOffset.get(newInput),
-                  () -> "mapNewInputToNewOffset.get(" + newInput + ")")
+              () -> "mapNewInputToNewOffset.get(" + newInput + ")")
               + valueGenFieldOffset;
 
       corDefOutputs.put(corRef.def(), newOutput);
@@ -3339,7 +3337,7 @@ private CorelMap(Multimap<RelNode, CorRef> mapRefRelToCorRef,
           && mapRefRelToCorRef.equals(((CorelMap) obj).mapRefRelToCorRef)
           && mapCorToCorRel.equals(((CorelMap) obj).mapCorToCorRel)
           && mapFieldAccessToCorRef.equals(
-              ((CorelMap) obj).mapFieldAccessToCorRef);
+          ((CorelMap) obj).mapFieldAccessToCorRef);
     }
 
     @Override public int hashCode() {