From 12bc753ca8632a6e30da1f9af2f554f1b5c6c9a9 Mon Sep 17 00:00:00 2001
From: Ruben Quesada Lopez <rubenql@gmail.com>
Date: Thu, 21 Mar 2024 19:28:50 +0000
Subject: [PATCH] Refactor implementation

---
 .../calcite/rel/metadata/RelMdCollation.java  | 23 ++++++-------------
 1 file changed, 7 insertions(+), 16 deletions(-)

diff --git a/core/src/main/java/org/apache/calcite/rel/metadata/RelMdCollation.java b/core/src/main/java/org/apache/calcite/rel/metadata/RelMdCollation.java
index fd2adbd12359..9c728384f4a5 100644
--- a/core/src/main/java/org/apache/calcite/rel/metadata/RelMdCollation.java
+++ b/core/src/main/java/org/apache/calcite/rel/metadata/RelMdCollation.java
@@ -329,22 +329,13 @@ public static List<RelCollation> sort(RelCollation collation) {
         if (integers.isEmpty()) {
           continue loop; // cannot do this collation
         }
-        if (integers.size() == 1) {
-          for (List<RelFieldCollation> fieldCollations : fieldCollationsList) {
-            fieldCollations.add(ifc.withFieldIndex(integers.iterator().next()));
-          }
-        } else {
-          List<List<RelFieldCollation>> newFieldCollationsList =
-              new ArrayList<>(fieldCollationsList.size() * integers.size());
-          for (List<RelFieldCollation> fieldCollations : fieldCollationsList) {
-            for (Integer integer : integers) {
-              List<RelFieldCollation> newFieldCollations = new ArrayList<>(fieldCollations);
-              newFieldCollations.add(ifc.withFieldIndex(integer));
-              newFieldCollationsList.add(newFieldCollations);
-            }
-          }
-          fieldCollationsList = newFieldCollationsList;
-        }
+        fieldCollationsList = fieldCollationsList.stream()
+            .flatMap(fieldCollations -> integers.stream()
+                .map(integer -> {
+                  List<RelFieldCollation> newFieldCollations = new ArrayList<>(fieldCollations);
+                  newFieldCollations.add(ifc.withFieldIndex(integer));
+                  return newFieldCollations;
+                })).collect(Collectors.toList());
       }
       assert !fieldCollationsList.isEmpty();
       for (List<RelFieldCollation> fieldCollations : fieldCollationsList) {