DRILL-8527: Hive Limit Push Down (#2997)

shfshihuafeng · web-flow · commit 78a52087bb85 · 2025-07-10T10:28:49.000-04:00
diff --git a/contrib/storage-hive/core/src/main/java/org/apache/drill/exec/store/hive/HiveScan.java b/contrib/storage-hive/core/src/main/java/org/apache/drill/exec/store/hive/HiveScan.java
@@ -65,7 +65,7 @@ public class HiveScan extends AbstractGroupScan {
   private final HiveReadEntry hiveReadEntry;
   private final HiveMetadataProvider metadataProvider;
   private final Map<String, String> confProperties;
-
+  private final int maxRecords;
   private List<List<LogicalInputSplit>> mappings;
   private List<LogicalInputSplit> inputSplits;
 
@@ -77,21 +77,23 @@ public HiveScan(@JsonProperty("userName") final String userName,
                   @JsonProperty("hiveStoragePluginConfig") final HiveStoragePluginConfig hiveStoragePluginConfig,
                   @JsonProperty("columns") final List<SchemaPath> columns,
                   @JsonProperty("confProperties") final Map<String, String> confProperties,
+                  @JsonProperty("maxRecords") final int maxRecords,
                   @JacksonInject final StoragePluginRegistry pluginRegistry) throws ExecutionSetupException {
     this(userName,
         hiveReadEntry,
         pluginRegistry.resolve(hiveStoragePluginConfig, HiveStoragePlugin.class),
         columns,
-        null, confProperties);
+        null, confProperties, maxRecords);
   }
 
   public HiveScan(final String userName, final HiveReadEntry hiveReadEntry, final HiveStoragePlugin hiveStoragePlugin,
-                  final List<SchemaPath> columns, final HiveMetadataProvider metadataProvider, final Map<String, String> confProperties) throws ExecutionSetupException {
+                  final List<SchemaPath> columns, final HiveMetadataProvider metadataProvider, final Map<String, String> confProperties, int maxRecords) throws ExecutionSetupException {
     super(userName);
     this.hiveReadEntry = hiveReadEntry;
     this.columns = columns;
     this.hiveStoragePlugin = hiveStoragePlugin;
     this.confProperties = confProperties;
+    this.maxRecords = maxRecords;
     if (metadataProvider == null) {
       this.metadataProvider = new HiveMetadataProvider(userName, hiveReadEntry, getHiveConf());
     } else {
@@ -106,10 +108,20 @@ public HiveScan(final HiveScan that) {
     this.hiveStoragePlugin = that.hiveStoragePlugin;
     this.metadataProvider = that.metadataProvider;
     this.confProperties = that.confProperties;
+    this.maxRecords = that.maxRecords;
+  }
+  public HiveScan(final HiveScan that, int maxRecords) {
+    super(that);
+    this.columns = that.columns;
+    this.hiveReadEntry = that.hiveReadEntry;
+    this.hiveStoragePlugin = that.hiveStoragePlugin;
+    this.metadataProvider = that.metadataProvider;
+    this.confProperties = that.confProperties;
+    this.maxRecords = maxRecords;
   }
 
   public HiveScan clone(final HiveReadEntry hiveReadEntry) throws ExecutionSetupException {
-    return new HiveScan(getUserName(), hiveReadEntry, hiveStoragePlugin, columns, metadataProvider, confProperties);
+    return new HiveScan(getUserName(), hiveReadEntry, hiveStoragePlugin, columns, metadataProvider, confProperties, maxRecords);
   }
 
   @JsonProperty
@@ -133,6 +145,11 @@ public Map<String, String> getConfProperties() {
     return confProperties;
   }
 
+  @JsonProperty
+  public int getMaxRecords() {
+    return maxRecords;
+  }
+
   @JsonIgnore
   public HiveStoragePlugin getStoragePlugin() {
     return hiveStoragePlugin;
@@ -167,6 +184,19 @@ public void applyAssignments(final List<CoordinationProtos.DrillbitEndpoint> end
     }
   }
 
+  @Override
+  public GroupScan applyLimit(int maxRecords) {
+    if (maxRecords == this.maxRecords){
+      return null;
+    }
+      return new HiveScan(this, maxRecords);
+  }
+
+  @Override
+  public boolean supportsLimitPushdown() {
+    return true;
+  }
+
   @Override
   public SubScan getSpecificScan(final int minorFragmentId) throws ExecutionSetupException {
     try {
@@ -189,7 +219,7 @@ public SubScan getSpecificScan(final int minorFragmentId) throws ExecutionSetupE
       }
 
       final HiveReadEntry subEntry = new HiveReadEntry(hiveReadEntry.getTableWrapper(), parts);
-      return new HiveSubScan(getUserName(), encodedInputSplits, subEntry, splitTypes, columns, hiveStoragePlugin, confProperties);
+      return new HiveSubScan(getUserName(), encodedInputSplits, subEntry, splitTypes, columns, hiveStoragePlugin, maxRecords, confProperties);
     } catch (IOException | ReflectiveOperationException e) {
       throw new ExecutionSetupException(e);
     }
@@ -271,6 +301,7 @@ public String toString() {
         + ", partitions= " + partitions
         + ", inputDirectories=" + metadataProvider.getInputDirectories(hiveReadEntry)
         + ", confProperties=" + confProperties
+        + ", maxRecords=" + maxRecords
         + "]";
   }
 
diff --git a/contrib/storage-hive/core/src/main/java/org/apache/drill/exec/store/hive/HiveStoragePlugin.java b/contrib/storage-hive/core/src/main/java/org/apache/drill/exec/store/hive/HiveStoragePlugin.java
@@ -113,7 +113,7 @@ public HiveScan getPhysicalScan(String userName, JSONOptions selection, List<Sch
         }
       }
 
-      return new HiveScan(userName, hiveReadEntry, this, columns, null, confProperties);
+      return new HiveScan(userName, hiveReadEntry, this, columns, null, confProperties, -1);
     } catch (ExecutionSetupException e) {
       throw new IOException(e);
     }
diff --git a/contrib/storage-hive/core/src/main/java/org/apache/drill/exec/store/hive/HiveSubScan.java b/contrib/storage-hive/core/src/main/java/org/apache/drill/exec/store/hive/HiveSubScan.java
@@ -58,7 +58,7 @@ public class HiveSubScan extends AbstractBase implements SubScan {
   private final List<HivePartition> partitions;
   private final List<SchemaPath> columns;
   private final Map<String, String> confProperties;
-
+  private final int maxRecords;
   @JsonCreator
   public HiveSubScan(@JacksonInject StoragePluginRegistry registry,
                      @JsonProperty("userName") String userName,
@@ -67,6 +67,7 @@ public HiveSubScan(@JacksonInject StoragePluginRegistry registry,
                      @JsonProperty("splitClasses") List<String> splitClasses,
                      @JsonProperty("columns") List<SchemaPath> columns,
                      @JsonProperty("hiveStoragePluginConfig") HiveStoragePluginConfig hiveStoragePluginConfig,
+                     @JsonProperty("maxRecords") int maxRecords,
                      @JsonProperty("confProperties") Map<String, String> confProperties)
       throws IOException, ExecutionSetupException, ReflectiveOperationException {
     this(userName,
@@ -75,6 +76,7 @@ public HiveSubScan(@JacksonInject StoragePluginRegistry registry,
         splitClasses,
         columns,
         registry.resolve(hiveStoragePluginConfig, HiveStoragePlugin.class),
+        maxRecords,
         confProperties);
   }
 
@@ -84,6 +86,7 @@ public HiveSubScan(final String userName,
                      final List<String> splitClasses,
                      final List<SchemaPath> columns,
                      final HiveStoragePlugin hiveStoragePlugin,
+                     final Integer maxRecords,
                      final Map<String, String> confProperties)
     throws IOException, ReflectiveOperationException {
     super(userName);
@@ -94,6 +97,7 @@ public HiveSubScan(final String userName,
     this.splitClasses = splitClasses;
     this.columns = columns;
     this.hiveStoragePlugin = hiveStoragePlugin;
+    this.maxRecords = maxRecords;
     this.confProperties = confProperties;
 
     for (int i = 0; i < splits.size(); i++) {
@@ -121,6 +125,10 @@ public List<SchemaPath> getColumns() {
     return columns;
   }
 
+  public int getMaxRecords() {
+    return maxRecords;
+  }
+
   @JsonProperty
   public HiveStoragePluginConfig getHiveStoragePluginConfig() {
     return hiveStoragePlugin.getConfig();
@@ -164,7 +172,7 @@ public <T, X, E extends Throwable> T accept(PhysicalVisitor<T, X, E> physicalVis
   @Override
   public PhysicalOperator getNewWithChildren(List<PhysicalOperator> children) throws ExecutionSetupException {
     try {
-      return new HiveSubScan(getUserName(), splits, hiveReadEntry, splitClasses, columns, hiveStoragePlugin, confProperties);
+      return new HiveSubScan(getUserName(), splits, hiveReadEntry, splitClasses, columns, hiveStoragePlugin, maxRecords, confProperties);
     } catch (IOException | ReflectiveOperationException e) {
       throw new ExecutionSetupException(e);
     }
diff --git a/contrib/storage-hive/core/src/main/java/org/apache/drill/exec/store/hive/readers/HiveDefaultRecordReader.java b/contrib/storage-hive/core/src/main/java/org/apache/drill/exec/store/hive/readers/HiveDefaultRecordReader.java
@@ -217,7 +217,7 @@ public class HiveDefaultRecordReader extends AbstractRecordReader {
    */
   private StructField[] selectedStructFieldRefs;
 
-
+  private final int maxRecords;
   /**
    * Readers constructor called by initializer.
    *
@@ -231,7 +231,7 @@ public class HiveDefaultRecordReader extends AbstractRecordReader {
    */
   public HiveDefaultRecordReader(HiveTableWithColumnCache table, HivePartition partition,
                                  Collection<InputSplit> inputSplits, List<SchemaPath> projectedColumns,
-                                 FragmentContext context, HiveConf hiveConf, UserGroupInformation proxyUgi) {
+                                 FragmentContext context, HiveConf hiveConf, UserGroupInformation proxyUgi, int maxRecords) {
     this.hiveTable = table;
     this.partition = partition;
     this.hiveConf = hiveConf;
@@ -243,6 +243,7 @@ public HiveDefaultRecordReader(HiveTableWithColumnCache table, HivePartition par
     this.partitionValues = new Object[0];
     setColumns(projectedColumns);
     this.fragmentContext = context;
+    this.maxRecords = maxRecords;
   }
 
   @Override
@@ -396,7 +397,8 @@ public int next() {
 
     try {
       int recordCount;
-      for (recordCount = 0; (recordCount < TARGET_RECORD_COUNT && hasNextValue(valueHolder)); recordCount++) {
+      int record = maxRecords > 0 ? maxRecords : TARGET_RECORD_COUNT;
+      for (recordCount = 0; (recordCount < record && hasNextValue(valueHolder)); recordCount++) {
         Object deserializedHiveRecord = partitionToTableSchemaConverter.convert(partitionDeserializer.deserialize((Writable) valueHolder));
         outputWriter.setPosition(recordCount);
         readHiveRecordAndInsertIntoRecordBatch(deserializedHiveRecord);
diff --git a/contrib/storage-hive/core/src/main/java/org/apache/drill/exec/store/hive/readers/HiveTextRecordReader.java b/contrib/storage-hive/core/src/main/java/org/apache/drill/exec/store/hive/readers/HiveTextRecordReader.java
@@ -58,8 +58,8 @@ public class HiveTextRecordReader extends HiveDefaultRecordReader {
    */
   public HiveTextRecordReader(HiveTableWithColumnCache table, HivePartition partition,
                               Collection<InputSplit> inputSplits, List<SchemaPath> projectedColumns,
-                              FragmentContext context, HiveConf hiveConf, UserGroupInformation proxyUgi) {
-    super(table, partition, inputSplits, projectedColumns, context, hiveConf, proxyUgi);
+                              FragmentContext context, HiveConf hiveConf, UserGroupInformation proxyUgi, int maxRecords) {
+    super(table, partition, inputSplits, projectedColumns, context, hiveConf, proxyUgi, maxRecords);
   }
 
   @Override
diff --git a/contrib/storage-hive/core/src/main/java/org/apache/drill/exec/store/hive/readers/ReadersInitializer.java b/contrib/storage-hive/core/src/main/java/org/apache/drill/exec/store/hive/readers/ReadersInitializer.java
@@ -56,10 +56,10 @@ public static List<RecordReader> init(ExecutorFragmentContext ctx, HiveSubScan c
     final UserGroupInformation proxyUgi = ImpersonationUtil.createProxyUgi(config.getUserName(), ctx.getQueryUserName());
     final List<List<InputSplit>> inputSplits = config.getInputSplits();
     final HiveConf hiveConf = config.getHiveConf();
-
+    final int maxRecords = config.getMaxRecords();
     if (inputSplits.isEmpty()) {
       return Collections.singletonList(
-          readerFactory.createReader(config.getTable(), null /*partition*/, null /*split*/, config.getColumns(), ctx, hiveConf, proxyUgi)
+          readerFactory.createReader(config.getTable(), null /*partition*/, null /*split*/, config.getColumns(), ctx, hiveConf, proxyUgi, maxRecords)
       );
     } else {
       IndexedPartitions partitions = getPartitions(config);
@@ -70,7 +70,7 @@ public static List<RecordReader> init(ExecutorFragmentContext ctx, HiveSubScan c
                   partitions.get(idx),
                   inputSplits.get(idx),
                   config.getColumns(),
-                  ctx, hiveConf, proxyUgi))
+                  ctx, hiveConf, proxyUgi, maxRecords))
           .collect(Collectors.toList());
     }
   }
@@ -109,8 +109,7 @@ private interface HiveReaderFactory {
 
     RecordReader createReader(HiveTableWithColumnCache table, HivePartition partition,
                               Collection<InputSplit> inputSplits, List<SchemaPath> projectedColumns,
-                              FragmentContext context, HiveConf hiveConf, UserGroupInformation proxyUgi);
-
+                              FragmentContext context, HiveConf hiveConf, UserGroupInformation proxyUgi, int maxRecords);
   }
 
   /**
diff --git a/contrib/storage-hive/core/src/test/java/org/apache/drill/exec/hive/TestHivePushDown.java b/contrib/storage-hive/core/src/test/java/org/apache/drill/exec/hive/TestHivePushDown.java
@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.drill.exec.hive;
+
+import org.apache.drill.categories.HiveStorageTest;
+import org.apache.drill.categories.SlowTest;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.rules.ExpectedException;
+
+import static org.junit.Assert.assertEquals;
+
+
+@Category({SlowTest.class, HiveStorageTest.class})
+public class TestHivePushDown extends HiveTestBase {
+
+  @Rule
+  public ExpectedException thrown = ExpectedException.none();
+
+  @Test
+  public void testLimitPushDown() throws Exception {
+    String query = "SELECT * FROM hive.`default`.kv LIMIT 1";
+
+    int actualRowCount = testSql(query);
+    assertEquals("Expected and actual row count should match", 1, actualRowCount);
+
+    testPlanMatchingPatterns(query, new String[]{"LIMIT"}, new String[]{"maxRecords=1"});
+  }
+}

Original file line number	Diff line number	Diff line change
`@@ -113,7 +113,7 @@ public HiveScan getPhysicalScan(String userName, JSONOptions selection, List<Sch`
`113`	`113`	`}`
`114`	`114`	`}`
`115`	`115`
`116`		`- return new HiveScan(userName, hiveReadEntry, this, columns, null, confProperties);`
	`116`	`+ return new HiveScan(userName, hiveReadEntry, this, columns, null, confProperties, -1);`
`117`	`117`	`} catch (ExecutionSetupException e) {`
`118`	`118`	`throw new IOException(e);`
`119`	`119`	`}`