Guarantee limit for broker queries

Since broker query has one split, limit is guaranteed.
trinodb · Jun 14, 2021 · 0d5ae90 · 0d5ae90
1 parent 92ec228
commit 0d5ae90
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/PinotMetadata.java b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/PinotMetadata.java
@@ -248,7 +248,8 @@ public Optional<LimitApplicationResult<ConnectorTableHandle>> applyLimit(Connect
                 handle.getConstraint(),
                 OptionalLong.of(limit),
                 dynamicTable);
-        return Optional.of(new LimitApplicationResult<>(handle, false, false));
+        boolean singleSplit = dynamicTable.isPresent();
+        return Optional.of(new LimitApplicationResult<>(handle, singleSplit, false));
     }
 
     @Override

diff --git a/plugin/trino-pinot/src/test/java/io/trino/plugin/pinot/TestPinotIntegrationSmokeTest.java b/plugin/trino-pinot/src/test/java/io/trino/plugin/pinot/TestPinotIntegrationSmokeTest.java
@@ -20,6 +20,7 @@
 import io.confluent.kafka.serializers.KafkaAvroSerializer;
 import io.trino.plugin.pinot.client.PinotHostMapper;
 import io.trino.sql.planner.plan.FilterNode;
+import io.trino.sql.planner.plan.LimitNode;
 import io.trino.sql.planner.plan.ProjectNode;
 import io.trino.testing.AbstractTestQueryFramework;
 import io.trino.testing.MaterializedResult;
@@ -785,4 +786,14 @@ public void testArrayFilter()
         // Array filters are not pushed down, as there are no array literals in pinot
         assertThat(query("SELECT price, vendor FROM " + JSON_TABLE + " WHERE prices = ARRAY[3.5, 5.5]")).isNotFullyPushedDown(FilterNode.class);
     }
+
+    @Test
+    public void testLimitPushdown()
+    {
+        assertThat(query("SELECT string_col, long_col FROM " + "\"SELECT string_col, long_col, bool_col FROM " + ALL_TYPES_TABLE + " WHERE int_col > 0\" " +
+                "  WHERE bool_col = 'false' LIMIT " + MAX_ROWS_PER_SPLIT_FOR_SEGMENT_QUERIES))
+                .isFullyPushedDown();
+        assertThat(query("SELECT string_col, long_col FROM " + ALL_TYPES_TABLE + "  WHERE int_col >0 AND bool_col = 'false' LIMIT " + MAX_ROWS_PER_SPLIT_FOR_SEGMENT_QUERIES))
+                .isNotFullyPushedDown(LimitNode.class);
+    }
 }