Add review fixes for Top-N pushdown support

Author: codluca

plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoMetadata.java

@@ -137,6 +137,9 @@ public class MongoMetadata
 
     private static final int MAX_QUALIFIED_IDENTIFIER_BYTE_LENGTH = 120;
 
+    public static final int MONGO_SORT_ASC = 1;
+    private static final int MONGO_SORT_DESC = -1;
+
     private final MongoSession mongoSession;
 
     private final AtomicReference<Runnable> rollbackAction = new AtomicReference<>();
@@ -634,7 +637,12 @@ public Optional<LimitApplicationResult<ConnectorTableHandle>> applyLimit(Connect
     }
 
     @Override
-    public Optional<TopNApplicationResult<ConnectorTableHandle>> applyTopN(ConnectorSession session, ConnectorTableHandle table, long topNCount, List<SortItem> sortItems, Map<String, ColumnHandle> assignments)
+    public Optional<TopNApplicationResult<ConnectorTableHandle>> applyTopN(
+            ConnectorSession session,
+            ConnectorTableHandle table,
+            long topNCount,
+            List<SortItem> sortItems,
+            Map<String, ColumnHandle> assignments)
     {
         MongoTableHandle handle = (MongoTableHandle) table;
 
@@ -651,15 +659,15 @@ public Optional<TopNApplicationResult<ConnectorTableHandle>> applyTopN(Connector
         Document sortNullFieldsDocument = null;
         for (SortItem sortItem : sortItems) {
             String columnName = sortItem.getName();
-            int direction = (sortItem.getSortOrder() == SortOrder.ASC_NULLS_FIRST || sortItem.getSortOrder() == SortOrder.ASC_NULLS_LAST) ? 1 : -1;
+            int direction = (sortItem.getSortOrder() == SortOrder.ASC_NULLS_FIRST || sortItem.getSortOrder() == SortOrder.ASC_NULLS_LAST) ? MONGO_SORT_ASC : MONGO_SORT_DESC;
 
             // MongoDB considers null values to be less than any other value.
             // When we have sort items with SortOrder.ASC_NULLS_LAST or SortOrder.DESC_NULLS_FIRST,
             // we need to add computed fields to sort correctly.
             if (sortItem.getSortOrder() == SortOrder.ASC_NULLS_LAST || sortItem.getSortOrder() == SortOrder.DESC_NULLS_FIRST) {
                 String sortColumnName = "_sortNulls_" + columnName;
                 Document condition = new Document();
-                condition.append("$cond", List.of(new Document("$eq", new ArrayList<>(Arrays.asList("$" + columnName, null))), 1, 0));
+                condition.append("$cond", ImmutableList.of(new Document("$eq", Arrays.asList("$" + columnName, null)), 1, 0));
                 if (sortNullFieldsDocument == null) {
                     sortNullFieldsDocument = new Document();
                 }
@@ -670,7 +678,7 @@ public Optional<TopNApplicationResult<ConnectorTableHandle>> applyTopN(Connector
             sortDocument.append(columnName, direction);
         }
         List<MongoTableSort> tableSortList = handle.sort().orElse(new ArrayList<>());
-        MongoTableSort tableSort = new MongoTableSort(sortDocument, sortNullFieldsDocument == null ? Optional.empty() : Optional.of(sortNullFieldsDocument), toIntExact(topNCount));
+        MongoTableSort tableSort = new MongoTableSort(sortDocument, Optional.ofNullable(sortNullFieldsDocument), toIntExact(topNCount));
         tableSortList.add(tableSort);
 
         return Optional.of(new TopNApplicationResult<>(

plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java

@@ -558,7 +558,7 @@ public MongoCursor<Document> execute(MongoTableHandle tableHandle, List<MongoCol
                 for (String sortField : tableSort.sort().keySet()) {
                     // Sorting on the field does not work unless we add it to the projection
                     if (!projection.containsKey(sortField)) {
-                        projection.append(sortField, 1);
+                        projection.append(sortField, MongoMetadata.MONGO_SORT_ASC);
                     }
                 }
             }

plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoTableSort.java

@@ -13,6 +13,7 @@
  */
 package io.trino.plugin.mongodb;
 
+import com.google.common.collect.ImmutableList;
 import org.bson.Document;
 
 import java.util.List;
@@ -38,8 +39,7 @@ public String toString()
                 .toList();
         builder.append(", orderBy = ").append(sortEntries);
         if (sortNullFields.isPresent()) {
-            List<String> nullFields = sortNullFields.get().keySet().stream()
-                    .toList();
+            List<String> nullFields = ImmutableList.copyOf(sortNullFields.get().keySet());
             builder.append(", nullFields=").append(nullFields);
         }
         return builder.toString();