Support TopN pushdown for text types in PostgreSQL

plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/BaseJdbcConnectorTest.java

@@ -48,6 +48,7 @@
 import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_PREDICATE_PUSHDOWN_WITH_VARCHAR_EQUALITY;
 import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_PREDICATE_PUSHDOWN_WITH_VARCHAR_INEQUALITY;
 import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_TOPN_PUSHDOWN;
+import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_TOPN_PUSHDOWN_WITH_VARCHAR;
 import static java.lang.String.format;
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -214,7 +215,7 @@ public void testCaseSensitiveTopNPushdown()
         }
 
         // topN over varchar/char columns should only be pushed down if the remote systems's sort order matches Trino
-        boolean expectTopNPushdown = hasBehavior(SUPPORTS_PREDICATE_PUSHDOWN_WITH_VARCHAR_INEQUALITY);
+        boolean expectTopNPushdown = hasBehavior(SUPPORTS_TOPN_PUSHDOWN_WITH_VARCHAR);
         PlanMatchPattern topNOverTableScan = node(TopNNode.class, anyTree(node(TableScanNode.class)));
 
         try (TestTable testTable = new TestTable(

plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlClient.java

@@ -202,6 +202,7 @@
 import static java.math.RoundingMode.UNNECESSARY;
 import static java.sql.DatabaseMetaData.columnNoNulls;
 import static java.util.Objects.requireNonNull;
+import static java.util.stream.Collectors.joining;
 
 public class PostgreSqlClient
         extends BaseJdbcClient
@@ -686,10 +687,9 @@ public boolean supportsTopN(ConnectorSession session, JdbcTableHandle handle, Li
         for (JdbcSortItem sortItem : sortOrder) {
             Type sortItemType = sortItem.getColumn().getColumnType();
             if (sortItemType instanceof CharType || sortItemType instanceof VarcharType) {
-                // PostgreSQL by default orders lowercase letters before uppercase, which is different from Trino
-                // NOTE: VarcharType also includes PostgreSQL enums
-                // TODO We could still push the sort down if we could inject a PostgreSQL-specific syntax for selecting a collation for given comparison.
-                return false;
+                if (!isCollatable(sortItem.getColumn())) {
+                    return false;
+                }
             }
         }
         return true;
@@ -698,7 +698,33 @@ public boolean supportsTopN(ConnectorSession session, JdbcTableHandle handle, Li
     @Override
     protected Optional<TopNFunction> topNFunction()
     {
-        return Optional.of(TopNFunction.sqlStandard(this::quoted));
+        return Optional.of((query, sortItems, limit) -> {
+            String orderBy = sortItems.stream()
+                    .map(sortItem -> {
+                        String ordering = sortItem.getSortOrder().isAscending() ? "ASC" : "DESC";
+                        String nullsHandling = sortItem.getSortOrder().isNullsFirst() ? "NULLS FIRST" : "NULLS LAST";
+                        String collation = "";
+                        if (isCollatable(sortItem.getColumn())) {
+                            collation = "COLLATE \"C\"";
+                        }
+                        return format("%s %s %s %s", quoted(sortItem.getColumn().getColumnName()), collation, ordering, nullsHandling);
+                    })
+                    .collect(joining(", "));
+            return format("%s ORDER BY %s LIMIT %d", query, orderBy, limit);
+        });
+    }
+
+    private boolean isCollatable(JdbcColumnHandle column)
+    {
+        if (column.getColumnType() instanceof CharType || column.getColumnType() instanceof VarcharType) {
+            String jdbcTypeName = column.getJdbcTypeHandle().getJdbcTypeName()
+                    .orElseThrow(() -> new TrinoException(JDBC_ERROR, "Type name is missing: " + column.getJdbcTypeHandle()));
+            // Only char (internally named bpchar)/varchar/text are the built-in collatable types
+            return "bpchar".equals(jdbcTypeName) || "varchar".equals(jdbcTypeName) || "text".equals(jdbcTypeName);
+        }
+
+        // non-textual types don't have the concept of collation
+        return false;
     }
 
     @Override

plugin/trino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlConnectorTest.java

@@ -82,6 +82,7 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior)
                 return false;
 
             case SUPPORTS_TOPN_PUSHDOWN:
+            case SUPPORTS_TOPN_PUSHDOWN_WITH_VARCHAR:
                 return true;
 
             case SUPPORTS_JOIN_PUSHDOWN:
@@ -763,7 +764,7 @@ public void testLimitPushdown()
         // with TopN over numeric column
         assertThat(query("SELECT * FROM (SELECT regionkey FROM nation ORDER BY nationkey ASC LIMIT 10) LIMIT 5")).isFullyPushedDown();
         // with TopN over varchar column
-        assertThat(query("SELECT * FROM (SELECT regionkey FROM nation ORDER BY name ASC LIMIT 10) LIMIT 5")).isNotFullyPushedDown(TopNNode.class);
+        assertThat(query("SELECT * FROM (SELECT regionkey FROM nation ORDER BY name ASC LIMIT 10) LIMIT 5")).isFullyPushedDown();
 
         // LIMIT with JOIN
         assertThat(query(joinPushdownEnabled(getSession()), "" +

testing/trino-testing/src/main/java/io/trino/testing/TestingConnectorBehavior.java

@@ -26,6 +26,7 @@ public enum TestingConnectorBehavior
     SUPPORTS_LIMIT_PUSHDOWN,
 
     SUPPORTS_TOPN_PUSHDOWN,
+    SUPPORTS_TOPN_PUSHDOWN_WITH_VARCHAR(fallback -> fallback.test(SUPPORTS_TOPN_PUSHDOWN) && fallback.test(SUPPORTS_PREDICATE_PUSHDOWN_WITH_VARCHAR_INEQUALITY)),
 
     SUPPORTS_AGGREGATION_PUSHDOWN,