6

apache · Oct 26, 2023 · 82ea9ca · 82ea9ca
1 parent 574b582
commit 82ea9ca
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 10 deletions.
diff --git a/...e-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java b/...e-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java
@@ -623,11 +623,12 @@ public PlanFragment visitPhysicalOlapScan(PhysicalOlapScan olapScan, PlanTransla
         TableRef ref = new TableRef(tableName, null, null);
         BaseTableRef tableRef = new BaseTableRef(ref, olapTable, tableName);
         tupleDescriptor.setRef(tableRef);
-        olapScanNode.setSampleTabletIds(olapScan.getSelectedTabletIds());
         olapScanNode.setSelectedPartitionIds(olapScan.getSelectedPartitionIds());
+        olapScanNode.setSampleTabletIds(olapScan.getSelectedTabletIds()); // TODO
         if (olapScan.getTableSample().isPresent()) {
             olapScanNode.setTableSample(new TableSample(olapScan.getTableSample().get().isPercent,
                     olapScan.getTableSample().get().sampleValue, olapScan.getTableSample().get().seek));
+            olapScanNode.computeSampleTabletIds();
         }
 
         // TODO:  remove this switch?

diff --git a/fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java
@@ -59,6 +59,7 @@
 import org.apache.doris.common.Config;
 import org.apache.doris.common.ErrorCode;
 import org.apache.doris.common.ErrorReport;
+import org.apache.doris.common.FeConstants;
 import org.apache.doris.common.UserException;
 import org.apache.doris.common.util.Util;
 import org.apache.doris.nereids.glue.translator.PlanTranslatorContext;
@@ -547,6 +548,7 @@ public void init(Analyzer analyzer) throws UserException {
             computePartitionInfo();
         }
         computeTupleState(analyzer);
+        computeSampleTabletIds();
 
         /**
          * Compute InAccurate cardinality before mv selector and tablet pruning.
@@ -944,6 +946,9 @@ public void computeSampleTabletIds() {
         long selectedRows = 0;
         long totalSampleRows = 0;
         List<Long> selectedPartitionList = new ArrayList<>();
+        if (FeConstants.runningUnitTest && selectedIndexId == -1) {
+            selectedIndexId = olapTable.getBaseIndexId();
+        }
         for (Long partitionId : selectedPartitionIds) {
             final Partition partition = olapTable.getPartition(partitionId);
             final MaterializedIndex selectedTable = partition.getIndex(selectedIndexId);
@@ -954,14 +959,12 @@ public void computeSampleTabletIds() {
 
         // 2.Sampling is not required in some cases, will not take effect after clear sampleTabletIds.
         if (tableSample.isPercent()) {
-            if (tableSample.getSampleValue() == 100) {
-                sampleTabletIds.clear();
+            if (tableSample.getSampleValue() >= 100) {
                 return;
             }
             totalSampleRows = (long) Math.max(selectedRows * (tableSample.getSampleValue() / 100.0), 1);
         } else {
             if (tableSample.getSampleValue() > selectedRows) {
-                sampleTabletIds.clear();
                 return;
             }
             totalSampleRows = tableSample.getSampleValue();
@@ -994,12 +997,18 @@ public void computeSampleTabletIds() {
                     ? tableSample.getSeek() : (long) (new SecureRandom().nextDouble() * tablets.size());
             for (int j = 0; j < tablets.size(); j++) {
                 int seekTid = (int) ((j + tabletSeek) % tablets.size());
-                if (tablets.get(seekTid).getRowCount(true) == 0) {
+                long tabletRowCount;
+                if (!FeConstants.runningUnitTest) {
+                    tabletRowCount = tablets.get(seekTid).getRowCount(true);
+                } else {
+                    tabletRowCount = selectedTable.getRowCount() / tablets.size();
+                }
+                if (tabletRowCount == 0) {
                     continue;
                 }
                 sampleTabletIds.add(tablets.get(seekTid).getId());
-                sampleRows -= tablets.get(seekTid).getRowCount(true);
-                hitRows += tablets.get(seekTid).getRowCount(true);
+                sampleRows -= tabletRowCount;
+                hitRows += tabletRowCount;
                 if (sampleRows <= 0) {
                     break;
                 }
@@ -1034,7 +1043,6 @@ private void computeTabletInfo() throws UserException {
          */
         Preconditions.checkState(scanBackendIds.size() == 0);
         Preconditions.checkState(scanTabletIds.size() == 0);
-        computeSampleTabletIds();
         for (Long partitionId : selectedPartitionIds) {
             final Partition partition = olapTable.getPartition(partitionId);
             final MaterializedIndex selectedTable = partition.getIndex(selectedIndexId);

diff --git a/fe/fe-core/src/test/java/org/apache/doris/analysis/SelectStmtTest.java b/fe/fe-core/src/test/java/org/apache/doris/analysis/SelectStmtTest.java
@@ -26,6 +26,7 @@
 import org.apache.doris.catalog.Tablet;
 import org.apache.doris.common.AnalysisException;
 import org.apache.doris.common.Config;
+import org.apache.doris.common.FeConstants;
 import org.apache.doris.common.util.Util;
 import org.apache.doris.planner.OlapScanNode;
 import org.apache.doris.planner.OriginalPlanner;
@@ -67,6 +68,7 @@ public static void setUp() throws Exception {
         Config.enable_batch_delete_by_default = true;
         Config.enable_http_server_v2 = false;
         UtFrameUtils.createDorisCluster(runningDir);
+        FeConstants.runningUnitTest = true;
         String createTblStmtStr = "create table db1.tbl1(k1 varchar(32),"
                 + " k2 varchar(32), k3 varchar(32), k4 int, k5 largeint) "
                 + "AGGREGATE KEY(k1, k2,k3,k4,k5) distributed by hash(k1) buckets 3"
@@ -890,7 +892,7 @@ public void testSelectSampleHashBucketTable() throws Exception {
         String sql4 = "SELECT * FROM db1.table1 TABLESAMPLE(9500 ROWS)";
         OriginalPlanner planner4 = (OriginalPlanner) dorisAssert.query(sql4).internalExecuteOneAndGetPlan();
         Set<Long> sampleTabletIds4 = ((OlapScanNode) planner4.getScanNodes().get(0)).getSampleTabletIds();
-        Assert.assertEquals(0, sampleTabletIds4.size()); // no sample, all tablet
+        Assert.assertEquals(10, sampleTabletIds4.size());
 
         String sql5 = "SELECT * FROM db1.table1 TABLESAMPLE(11000 ROWS)";
         OriginalPlanner planner5 = (OriginalPlanner) dorisAssert.query(sql5).internalExecuteOneAndGetPlan();
@@ -995,7 +997,7 @@ public void testSelectSampleRandomBucketTable() throws Exception {
         String sql4 = "SELECT * FROM db1.table3 TABLESAMPLE(9500 ROWS)";
         OriginalPlanner planner4 = (OriginalPlanner) dorisAssert.query(sql4).internalExecuteOneAndGetPlan();
         Set<Long> sampleTabletIds4 = ((OlapScanNode) planner4.getScanNodes().get(0)).getSampleTabletIds();
-        Assert.assertEquals(0, sampleTabletIds4.size()); // no sample, all tablet
+        Assert.assertEquals(10, sampleTabletIds4.size());
 
         String sql5 = "SELECT * FROM db1.table3 TABLESAMPLE(11000 ROWS)";
         OriginalPlanner planner5 = (OriginalPlanner) dorisAssert.query(sql5).internalExecuteOneAndGetPlan();