Support limit pushdown in Loki connector.

plugin/trino-loki/pom.xml

@@ -43,6 +43,11 @@
             <artifactId>json</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>io.airlift</groupId>
+            <artifactId>log</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>io.airlift</groupId>
             <artifactId>units</artifactId>
@@ -51,7 +56,7 @@
         <dependency>
             <groupId>io.github.jeschkies</groupId>
             <artifactId>loki-client</artifactId>
-            <version>0.0.4</version>
+            <version>0.0.5</version>
         </dependency>
 
         <dependency>
@@ -106,12 +111,6 @@
             <scope>runtime</scope>
         </dependency>
 
-        <dependency>
-            <groupId>io.airlift</groupId>
-            <artifactId>log</artifactId>
-            <scope>runtime</scope>
-        </dependency>
-
         <dependency>
             <groupId>io.airlift</groupId>
             <artifactId>log-manager</artifactId>

plugin/trino-loki/src/main/java/io/trino/plugin/loki/LokiMetadata.java

@@ -26,6 +26,7 @@
 import io.trino.spi.connector.ConnectorTableHandle;
 import io.trino.spi.connector.ConnectorTableMetadata;
 import io.trino.spi.connector.ConnectorTableVersion;
+import io.trino.spi.connector.LimitApplicationResult;
 import io.trino.spi.connector.SchemaTableName;
 import io.trino.spi.connector.TableFunctionApplicationResult;
 import io.trino.spi.function.table.ConnectorTableFunctionHandle;
@@ -53,6 +54,8 @@ public class LokiMetadata
 
     private final LokiClient lokiClient;
 
+    private Data.ResultType expectedResultType;
+
     @Inject
     public LokiMetadata(LokiClient lokiClient, TypeManager typeManager)
     {
@@ -108,7 +111,8 @@ public List<ColumnHandle> getColumnHandles(String query)
         columnsBuilder.add(new LokiColumnHandle("timestamp", TIMESTAMP_TZ_MILLIS, 1));
 
         try {
-            Type valueType = lokiClient.getExpectedResultType(query) == Data.ResultType.Matrix ? DOUBLE : VARCHAR;
+            this.expectedResultType = lokiClient.getExpectedResultType(query);
+            Type valueType = this.expectedResultType == Data.ResultType.Matrix ? DOUBLE : VARCHAR;
             columnsBuilder.add(new LokiColumnHandle("value", valueType, 2));
         }
         catch (LokiClientException e) {
@@ -117,4 +121,23 @@ public List<ColumnHandle> getColumnHandles(String query)
 
         return columnsBuilder.build();
     }
+
+    @Override
+    public Optional<LimitApplicationResult<ConnectorTableHandle>> applyLimit(ConnectorSession session, ConnectorTableHandle tableHandle, long limit)
+    {
+        LokiTableHandle lokiTableHandle = (LokiTableHandle) tableHandle;
+
+        // Metric queries do not support setting a limit
+        if (this.expectedResultType != null && this.expectedResultType == Data.ResultType.Matrix) {
+            return Optional.empty();
+        }
+
+        // This change has no effect
+        if (lokiTableHandle.limit().isPresent() && limit == lokiTableHandle.limit().getAsLong()) {
+            return Optional.empty();
+        }
+
+        lokiTableHandle = lokiTableHandle.withLimit(limit);
+        return Optional.of(new LimitApplicationResult<>(lokiTableHandle, true, false));
+    }
 }

plugin/trino-loki/src/main/java/io/trino/plugin/loki/LokiRecordSet.java

@@ -14,6 +14,7 @@
 package io.trino.plugin.loki;
 
 import com.google.common.collect.ImmutableList;
+import io.airlift.log.Logger;
 import io.github.jeschkies.loki.client.LokiClient;
 import io.github.jeschkies.loki.client.LokiClientException;
 import io.github.jeschkies.loki.client.model.Matrix;
@@ -28,11 +29,14 @@
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.trino.plugin.loki.LokiErrorCode.LOKI_CLIENT_ERROR;
+import static java.lang.Math.toIntExact;
 import static java.util.Objects.requireNonNull;
 
 public class LokiRecordSet
         implements RecordSet
 {
+    private static final Logger log = Logger.get(LokiRecordSet.class);
+
     private final List<LokiColumnHandle> columnHandles;
     private final List<Type> columnTypes;
 
@@ -49,7 +53,8 @@ public LokiRecordSet(LokiClient lokiClient, LokiSplit split, List<LokiColumnHand
 
         // Execute the query
         try {
-            this.result = lokiClient.rangeQuery(split.query(), split.start(), split.end(), split.step());
+            log.info("querying %s with limit %d", split.query(), split.limit());
+            this.result = lokiClient.rangeQuery(split.query(), split.start(), split.end(), split.step(), toIntExact(split.limit()));
         }
         catch (LokiClientException e) {
             throw new TrinoException(LOKI_CLIENT_ERROR, e);

plugin/trino-loki/src/main/java/io/trino/plugin/loki/LokiSplit.java

@@ -16,10 +16,11 @@
 import io.trino.spi.connector.ConnectorSplit;
 
 import java.time.Instant;
+import java.util.OptionalLong;
 
 import static java.util.Objects.requireNonNull;
 
-public record LokiSplit(String query, Instant start, Instant end, int step)
+public record LokiSplit(String query, Instant start, Instant end, int step, long limit)
         implements ConnectorSplit
 {
     public LokiSplit

plugin/trino-loki/src/main/java/io/trino/plugin/loki/LokiSplitManager.java

@@ -39,7 +39,9 @@ public ConnectorSplitSource getSplits(
     {
         final LokiTableHandle table = (LokiTableHandle) connectorTableHandle;
 
-        List<ConnectorSplit> splits = ImmutableList.of(new LokiSplit(table.query(), table.start(), table.end(), table.step()));
+        // Set the limit to 0 as the default. Since the client ignores 0.
+        long limit = table.limit().orElse(0);
+        List<ConnectorSplit> splits = ImmutableList.of(new LokiSplit(table.query(), table.start(), table.end(), table.step(), limit));
         return new FixedSplitSource(splits);
     }
 }

plugin/trino-loki/src/main/java/io/trino/plugin/loki/LokiTableHandle.java

@@ -19,17 +19,24 @@
 
 import java.time.Instant;
 import java.util.List;
+import java.util.OptionalLong;
 
 import static java.util.Objects.requireNonNull;
 
-public record LokiTableHandle(String query, Instant start, Instant end, int step, List<ColumnHandle> columnHandles)
+public record LokiTableHandle(String query, Instant start, Instant end, int step, List<ColumnHandle> columnHandles, OptionalLong limit)
         implements ConnectorTableHandle
 {
     public LokiTableHandle
     {
         requireNonNull(query, "query is null");
         requireNonNull(start, "start is null");
         requireNonNull(end, "end is null");
+        requireNonNull(limit, "limit is null");
         columnHandles = ImmutableList.copyOf(columnHandles);
     }
+
+    public LokiTableHandle withLimit(long limit)
+    {
+        return new LokiTableHandle(query, start, end, step, columnHandles, OptionalLong.of(limit));
+    }
 }

plugin/trino-loki/src/main/java/io/trino/plugin/loki/QueryRangeTableFunction.java

@@ -36,6 +36,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.OptionalLong;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
@@ -130,7 +131,8 @@ public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransact
                 start,
                 end,
                 step.intValue(),
-                columnHandles);
+                columnHandles,
+                OptionalLong.empty());
 
         return TableFunctionAnalysis.builder()
                 .returnedType(returnedType)

plugin/trino-loki/src/test/java/io/trino/plugin/loki/TestLokiIntegration.java

@@ -25,6 +25,7 @@
 import java.time.ZoneOffset;
 import java.time.format.DateTimeFormatter;
 import java.time.temporal.ChronoUnit;
+import java.util.concurrent.TimeUnit;
 
 import static java.lang.String.format;
 
@@ -81,6 +82,42 @@
                 "VALUES ('line 1')");
     }
 
+    @Test
+    void testLimitedLogsQuery()
+            throws Exception
+    {
+        Instant start = Instant.now().minus(Duration.ofHours(3));
+        Instant end = start.plus(Duration.ofHours(2));
+
+        // Loki has a default of 100. Setting it to 120 verifies that the limit is propagated.
+        long limit = 120;
+        long numberOfRows = 150;
+
+        for (int i = 0; i < numberOfRows; i++) {
+            client.pushLogLine("line " + i, end.minus(Duration.ofSeconds(i)), ImmutableMap.of("test", "limited_logs_query"));
+        }
+        client.flush();
+
+        // We expect the last 120 lines
+        StringBuilder expected = new StringBuilder("VALUES ");
+        for (long i = numberOfRows - limit; i < numberOfRows; i++) {
+            expected.append("('line ").append(i).append("')");
+            if (i < numberOfRows - 1) {
+                expected.append(", ");
+            }
+        }
+        assertQueryEventually(getSession(), format("""
+                        SELECT value FROM
+                        TABLE(system.query_range(
+                         '{test="limited_logs_query"}',
+                         TIMESTAMP '%s',
+                         TIMESTAMP '%s'
+                        ))
+                        LIMIT %d
+                        """, timestampFormatter.format(start), timestampFormatter.format(end), limit),
+                expected.toString(), new io.airlift.units.Duration(30, TimeUnit.SECONDS));
+    }
+
     @Test
     void testMetricsQuery()
             throws Exception
@@ -92,7 +129,7 @@
         client.pushLogLine("line 2", end.minus(Duration.ofMinutes(2)), ImmutableMap.of("test", "metrics_query"));
         client.pushLogLine("line 3", end.minus(Duration.ofMinutes(1)), ImmutableMap.of("test", "metrics_query"));
         client.flush();
         assertQuery(format("""
                         SELECT value FROM
                         TABLE(system.query_range(
                          'count_over_time({test="metrics_query"}[5m])',
@@ -115,7 +152,7 @@
         client.pushLogLine("line 2", end.minus(Duration.ofMinutes(2)), ImmutableMap.of("test", "labels"));
         client.pushLogLine("line 3", end.minus(Duration.ofMinutes(1)), ImmutableMap.of("test", "labels"));
         client.flush();
         assertQuery(format("""
                         SELECT labels['test'] FROM
                         TABLE(system.query_range(
                          'count_over_time({test="labels"}[5m])',
@@ -189,7 +226,7 @@
         this.client.pushLogLine("line 2", start.plus(Duration.ofHours(2)), ImmutableMap.of("test", "timestamp_metrics_query"));
         this.client.pushLogLine("line 3", start.plus(Duration.ofHours(3)), ImmutableMap.of("test", "timestamp_metrics_query"));
         this.client.flush();
         assertQuery(format("""
                         SELECT to_iso8601(timestamp), value FROM
                         TABLE(system.query_range(
                          'count_over_time({test="timestamp_metrics_query"}[5m])',
@@ -211,29 +248,27 @@
     @Test
     void testQueryRangeInvalidArguments()
     {
-        assertQueryFails(
-                """
-                SELECT to_iso8601(timestamp), value FROM
-                TABLE(system.query_range(
-                 'count_over_time({test="timestamp_metrics_query"}[5m])',
-                 TIMESTAMP '2012-08-08',
-                 TIMESTAMP '2012-08-09',
-                 -300
-                ))
-                LIMIT 1
-                """,
+        assertQueryFails("""
+                        SELECT to_iso8601(timestamp), value FROM
+                        TABLE(system.query_range(
+                         'count_over_time({test="timestamp_metrics_query"}[5m])',
+                         TIMESTAMP '2012-08-08',
+                         TIMESTAMP '2012-08-09',
+                         -300
+                        ))
+                        LIMIT 1
+                        """,
                 "step must be positive");
-        assertQueryFails(
-                """
-                SELECT to_iso8601(timestamp), value FROM
-                TABLE(system.query_range(
-                 'count_over_time({test="timestamp_metrics_query"}[5m])',
-                 TIMESTAMP '2012-08-08',
-                 TIMESTAMP '2012-08-09',
-                 NULL
-                ))
-                LIMIT 1
-                """,
+        assertQueryFails("""
+                        SELECT to_iso8601(timestamp), value FROM
+                        TABLE(system.query_range(
+                         'count_over_time({test="timestamp_metrics_query"}[5m])',
+                         TIMESTAMP '2012-08-08',
+                         TIMESTAMP '2012-08-09',
+                         NULL
+                        ))
+                        LIMIT 1
+                        """,
                 "step must be positive");
     }
 }