Limit pushdown for Druid connector

presto-druid/src/main/java/com/facebook/presto/druid/DruidQueryGenerator.java

@@ -20,6 +20,7 @@
 import com.facebook.presto.spi.function.FunctionMetadataManager;
 import com.facebook.presto.spi.function.StandardFunctionResolution;
 import com.facebook.presto.spi.plan.FilterNode;
+import com.facebook.presto.spi.plan.LimitNode;
 import com.facebook.presto.spi.plan.PlanNode;
 import com.facebook.presto.spi.plan.PlanVisitor;
 import com.facebook.presto.spi.plan.ProjectNode;
@@ -202,6 +203,15 @@ public DruidQueryGeneratorContext visitProject(ProjectNode node, DruidQueryGener
             return context.withProject(newSelections);
         }
 
+        @Override
+        public DruidQueryGeneratorContext visitLimit(LimitNode node, DruidQueryGeneratorContext context)
+        {
+            checkArgument(!node.isPartial(), "Druid query generator cannot handle partial limit");
+            context = node.getSource().accept(this, context);
+            requireNonNull(context, "context is null");
+            return context.withLimit(node.getCount()).withOutputColumns(node.getOutputVariables());
+        }
+
         @Override
         public DruidQueryGeneratorContext visitTableScan(TableScanNode node, DruidQueryGeneratorContext contextIn)
         {

presto-druid/src/main/java/com/facebook/presto/druid/DruidQueryGeneratorContext.java

@@ -20,18 +20,20 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.OptionalLong;
 import java.util.stream.Collectors;
 
 import static com.facebook.presto.druid.DruidErrorCode.DRUID_QUERY_GENERATOR_FAILURE;
 import static com.google.common.base.MoreObjects.toStringHelper;
-import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
 import static java.util.Objects.requireNonNull;
 
 public class DruidQueryGeneratorContext
 {
     private final Map<VariableReferenceExpression, Selection> selections;
     private final Optional<String> from;
     private final Optional<String> filter;
+    private final OptionalLong limit;
 
     @Override
     public String toString()
@@ -40,6 +42,7 @@ public String toString()
                 .add("selections", selections)
                 .add("from", from)
                 .add("filter", filter)
+                .add("limit", limit)
                 .toString();
     }
 
@@ -55,34 +58,57 @@ public String toString()
         this(
                 selections,
                 Optional.ofNullable(from),
-                Optional.empty());
+                Optional.empty(),
+                OptionalLong.empty());
     }
 
     private DruidQueryGeneratorContext(
             Map<VariableReferenceExpression, Selection> selections,
             Optional<String> from,
-            Optional<String> filter)
+            Optional<String> filter,
+            OptionalLong limit)
     {
         this.selections = new LinkedHashMap<>(requireNonNull(selections, "selections can't be null"));
         this.from = requireNonNull(from, "source can't be null");
         this.filter = requireNonNull(filter, "filter is null");
+        this.limit = requireNonNull(limit, "limit is null");
     }
 
     public DruidQueryGeneratorContext withFilter(String filter)
     {
-        checkArgument(!hasFilter(), "Druid doesn't support filters at multiple levels");
+        checkState(!hasFilter(), "Druid doesn't support filters at multiple levels");
         return new DruidQueryGeneratorContext(
                 selections,
                 from,
-                Optional.of(filter));
+                Optional.of(filter),
+                limit);
     }
 
     public DruidQueryGeneratorContext withProject(Map<VariableReferenceExpression, Selection> newSelections)
     {
         return new DruidQueryGeneratorContext(
                 newSelections,
                 from,
-                filter);
+                filter,
+                limit);
+    }
+
+    public DruidQueryGeneratorContext withLimit(long limit)
+    {
+        if (limit <= 0 || limit > Long.MAX_VALUE) {
+            throw new PrestoException(DRUID_QUERY_GENERATOR_FAILURE, "Invalid limit: " + limit);
+        }
+        checkState(!hasLimit(), "Limit already exists. Druid doesn't support limit on top of another limit");
+        return new DruidQueryGeneratorContext(
+                selections,
+                from,
+                filter,
+                OptionalLong.of(limit));
+    }
+
+    private boolean hasLimit()
+    {
+        return limit.isPresent();
     }
 
     private boolean hasFilter()
@@ -113,6 +139,11 @@ public DruidQueryGenerator.GeneratedDql toQuery()
             query += " WHERE " + filter.get();
             pushdown = true;
         }
+
+        if (limit.isPresent()) {
+            query += " LIMIT " + limit.getAsLong();
+            pushdown = true;
+        }
         return new DruidQueryGenerator.GeneratedDql(tableName, query, pushdown);
     }
 
@@ -133,7 +164,7 @@ public DruidQueryGeneratorContext withOutputColumns(List<VariableReferenceExpres
         Map<VariableReferenceExpression, Selection> newSelections = new LinkedHashMap<>();
         outputColumns.forEach(o -> newSelections.put(o, requireNonNull(selections.get(o), "Cannot find the selection " + o + " in the original context " + this)));
 
-        return new DruidQueryGeneratorContext(newSelections, from, filter);
+        return new DruidQueryGeneratorContext(newSelections, from, filter, limit);
     }
 
     public enum Origin

presto-druid/src/test/java/com/facebook/presto/druid/TestDruidQueryGenerator.java

@@ -80,16 +80,16 @@ private PlanNode buildPlan(Function<PlanBuilder, PlanNode> consumer)
     @Test
     public void testSimpleSelectStar()
     {
-        testDQL(planBuilder -> tableScan(planBuilder, druidTable, regionId, city, fare, secondsSinceEpoch),
-                "SELECT regionId, city, fare, secondsSinceEpoch FROM realtimeOnly");
-        testDQL(planBuilder -> tableScan(planBuilder, druidTable, regionId, secondsSinceEpoch),
-                "SELECT regionId, secondsSinceEpoch FROM realtimeOnly");
+        testDQL(planBuilder -> limit(planBuilder, 50L, tableScan(planBuilder, druidTable, regionId, city, fare, secondsSinceEpoch)),
+                "SELECT regionId, city, fare, secondsSinceEpoch FROM realtimeOnly LIMIT 50");
+        testDQL(planBuilder -> limit(planBuilder, 10L, tableScan(planBuilder, druidTable, regionId, secondsSinceEpoch)),
+                "SELECT regionId, secondsSinceEpoch FROM realtimeOnly LIMIT 10");
     }
 
     @Test
-    public void testSimpleSelectWithFilter()
+    public void testSimpleSelectWithFilterLimit()
     {
-        testDQL(planBuilder -> project(planBuilder, filter(planBuilder, tableScan(planBuilder, druidTable, regionId, city, fare, secondsSinceEpoch), getRowExpression("secondssinceepoch > 20", defaultSessionHolder)), ImmutableList.of("city", "secondssinceepoch")),
-                "SELECT city, secondsSinceEpoch FROM realtimeOnly WHERE (secondsSinceEpoch > 20)");
+        testDQL(planBuilder -> limit(planBuilder, 30L, project(planBuilder, filter(planBuilder, tableScan(planBuilder, druidTable, regionId, city, fare, secondsSinceEpoch), getRowExpression("secondssinceepoch > 20", defaultSessionHolder)), ImmutableList.of("city", "secondssinceepoch"))),
+                "SELECT city, secondsSinceEpoch FROM realtimeOnly WHERE (secondsSinceEpoch > 20) LIMIT 30");
     }
 }