tracks total bytes, not average bytes

Author: dnhatn

core/src/main/java/org/elasticsearch/index/shard/DocsStats.java

@@ -25,35 +25,31 @@
 import org.elasticsearch.common.io.stream.Streamable;
 import org.elasticsearch.common.xcontent.ToXContentFragment;
 import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.index.store.StoreStats;
 
 import java.io.IOException;
 
 public class DocsStats implements Streamable, ToXContentFragment {
 
     long count = 0;
     long deleted = 0;
-    long averageSizeInBytes = 0;
+    long totalSizeInBytes = 0;
 
     public DocsStats() {
 
     }
 
-    public DocsStats(long count, long deleted, long averageSizeInBytes) {
+    public DocsStats(long count, long deleted, long totalSizeInBytes) {
         this.count = count;
         this.deleted = deleted;
-        this.averageSizeInBytes = averageSizeInBytes;
+        this.totalSizeInBytes = totalSizeInBytes;
     }
 
     public void add(DocsStats that) {
         if (that == null) {
             return;
         }
-        long totalBytes = this.averageSizeInBytes * (this.count + this.deleted)
-                        + that.averageSizeInBytes * (that.count + that.deleted);
-        long totalDocs = this.count + this.deleted + that.count + that.deleted;
-        if (totalDocs > 0) {
-            this.averageSizeInBytes = totalBytes / totalDocs;
-        }
+        this.totalSizeInBytes += that.totalSizeInBytes;
         this.count += that.count;
         this.deleted += that.deleted;
     }
@@ -66,16 +62,28 @@ public long getDeleted() {
         return this.deleted;
     }
 
+    /**
+     * Returns the total size in bytes of all documents in this stats.
+     * This value may be more reliable than {@link StoreStats#getSizeInBytes()} in estimating the index size.
+     */
+    public long getTotalSizeInBytes() {
+        return totalSizeInBytes;
+    }
+
+    /**
+     * Returns the average size in bytes of all documents in this stats.
+     */
     public long getAverageSizeInBytes() {
-        return averageSizeInBytes;
+        long totalDocs = count + deleted;
+        return totalDocs == 0 ? 0 : totalSizeInBytes / totalDocs;
     }
 
     @Override
     public void readFrom(StreamInput in) throws IOException {
         count = in.readVLong();
         deleted = in.readVLong();
         if (in.getVersion().onOrAfter(Version.V_7_0_0_alpha1)) {
-            averageSizeInBytes = in.readVLong();
+            totalSizeInBytes = in.readVLong();
         }
     }
 
@@ -84,7 +92,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeVLong(count);
         out.writeVLong(deleted);
         if (out.getVersion().onOrAfter(Version.V_7_0_0_alpha1)) {
-            out.writeVLong(averageSizeInBytes);
+            out.writeVLong(totalSizeInBytes);
         }
     }
 

core/src/main/java/org/elasticsearch/index/shard/IndexShard.java

@@ -880,16 +880,18 @@ public FlushStats flushStats() {
     }
 
     public DocsStats docStats() {
-        long totalDocsInSegments = 0L;
-        long totalBytesInSegments = 0L;
-        for (Segment segment : segments(false)) {
-            totalDocsInSegments += segment.getNumDocs();
-            totalBytesInSegments += segment.getSizeInBytes();
-        }
-        final long avgDocSize = totalDocsInSegments > 0 ? totalBytesInSegments / totalDocsInSegments : 0;
-        try (Engine.Searcher searcher = acquireSearcher("doc_stats")) {
-            return new DocsStats(searcher.reader().numDocs(), searcher.reader().numDeletedDocs(), avgDocSize);
+        long numDocs = 0;
+        long numDeletedDocs = 0;
+        long sizeInBytes = 0;
+        List<Segment> segments = segments(false);
+        for (Segment segment : segments) {
+            if (segment.search) {
+                numDocs += segment.getNumDocs();
+                numDeletedDocs += segment.getDeletedDocs();
+                sizeInBytes += segment.getSizeInBytes();
+            }
         }
+        return new DocsStats(numDocs, numDeletedDocs, sizeInBytes);
     }
 
     /**

core/src/test/java/org/elasticsearch/index/shard/DocsStatsTests.java

@@ -29,29 +29,17 @@
 public class DocsStatsTests extends ESTestCase {
 
     public void testCalculateAverageDocSize() throws Exception {
-        DocsStats stats = new DocsStats(10, 2, 10);
+        DocsStats stats = new DocsStats(10, 2, 120);
         assertThat(stats.getAverageSizeInBytes(), equalTo(10L));
 
-        stats.add(new DocsStats(0, 0, randomNonNegativeLong()));
+        stats.add(new DocsStats(0, 0, 0));
         assertThat(stats.getAverageSizeInBytes(), equalTo(10L));
 
-        // (38*900 + 12*10) / 50 = 686L
-        stats.add(new DocsStats(8, 30, 900));
+        stats.add(new DocsStats(8, 30, 480));
         assertThat(stats.getCount(), equalTo(18L));
         assertThat(stats.getDeleted(), equalTo(32L));
-        assertThat(stats.getAverageSizeInBytes(), equalTo(686L));
-
-        // (50*686 + 40*120) / 90 = 434L
-        stats.add(new DocsStats(0, 40, 120));
-        assertThat(stats.getCount(), equalTo(18L));
-        assertThat(stats.getDeleted(), equalTo(72L));
-        assertThat(stats.getAverageSizeInBytes(), equalTo(434L));
-
-        // (90*434 + 35*99) / 125 = 340L
-        stats.add(new DocsStats(35, 0, 99));
-        assertThat(stats.getCount(), equalTo(53L));
-        assertThat(stats.getDeleted(), equalTo(72L));
-        assertThat(stats.getAverageSizeInBytes(), equalTo(340L));
+        assertThat(stats.getTotalSizeInBytes(), equalTo(600L));
+        assertThat(stats.getAverageSizeInBytes(), equalTo(12L));
     }
 
     public void testSerialize() throws Exception {

core/src/test/java/org/elasticsearch/index/shard/IndexShardTests.java

@@ -88,6 +88,7 @@
 import org.elasticsearch.index.seqno.SequenceNumbers;
 import org.elasticsearch.index.snapshots.IndexShardSnapshotStatus;
 import org.elasticsearch.index.store.Store;
+import org.elasticsearch.index.store.StoreStats;
 import org.elasticsearch.index.translog.Translog;
 import org.elasticsearch.index.translog.TranslogTests;
 import org.elasticsearch.indices.IndicesQueryCache;
@@ -2273,54 +2274,57 @@ public void testDocStats() throws IOException {
         }
     }
 
-    public void testEstimateAverageDocSize() throws Exception {
+    public void testEstimateTotalDocSize() throws Exception {
         IndexShard indexShard = null;
         try {
             indexShard = newStartedShard(true);
-            int smallDocNum = randomIntBetween(5, 100);
-            for (int i = 0; i < smallDocNum; i++) {
-                indexDoc(indexShard, "test", "small-" + i);
-            }
-            // Average document size is estimated by sampling segments, thus it should be zero without flushing.
-            DocsStats withoutFlush = indexShard.docStats();
-            assertThat(withoutFlush.averageSizeInBytes, equalTo(0L));
 
-            indexShard.flush(new FlushRequest());
-            indexShard.refresh("test");
-            DocsStats smallStats = indexShard.docStats();
-            assertThat(smallStats.averageSizeInBytes, greaterThan(10L));
-
-            long storedAvgSize = indexShard.storeStats().sizeInBytes() / smallDocNum;
-            assertThat("Estimated average document size is too small compared with the average stored size",
-                smallStats.averageSizeInBytes, greaterThanOrEqualTo(storedAvgSize * 80/100));
-            assertThat("Estimated average document size is too large compared with the average stored size",
-                smallStats.averageSizeInBytes, lessThanOrEqualTo(storedAvgSize * 120/100));
-
-            // Indexing large documents should increase the average document size.
-            int largeDocNum = randomIntBetween(100, 200);
-            for (int i = 0; i < largeDocNum; i++) {
+            int numDoc = randomIntBetween(100, 200);
+            for (int i = 0; i < numDoc; i++) {
                 String doc = XContentFactory.jsonBuilder()
                     .startObject()
                         .field("count", randomInt())
                         .field("point", randomFloat())
                         .field("description", randomUnicodeOfCodepointLength(100))
                     .endObject().string();
-                indexDoc(indexShard, "test", "large-" + i, doc);
+                indexDoc(indexShard, "doc", Integer.toString(i), doc);
             }
+
+            assertThat("Without flushing, segment sizes should be zero",
+                indexShard.docStats().getTotalSizeInBytes(), equalTo(0L));
+
             indexShard.flush(new FlushRequest());
             indexShard.refresh("test");
-            DocsStats largeStats = indexShard.docStats();
-            assertThat(largeStats.averageSizeInBytes, greaterThan(100L));
-            assertThat(largeStats.averageSizeInBytes, greaterThan(smallStats.averageSizeInBytes));
+            {
+                final DocsStats docsStats = indexShard.docStats();
+                final StoreStats storeStats = indexShard.storeStats();
+                assertThat(storeStats.sizeInBytes(), greaterThan(numDoc * 100L)); // A doc should be more than 100 bytes.
 
-            int deleteDocs = randomIntBetween(1, smallDocNum / 2);
-            for (int i = 0; i < deleteDocs; i++) {
-                deleteDoc(indexShard, "test", "small-" + i);
+                assertThat("Estimated total document size is too small compared with the stored size",
+                    docsStats.getTotalSizeInBytes(), greaterThanOrEqualTo(storeStats.sizeInBytes() * 80/100));
+                assertThat("Estimated total document size is too large compared with the stored size",
+                    docsStats.getTotalSizeInBytes(), lessThanOrEqualTo(storeStats.sizeInBytes() * 120/100));
             }
+
+            // Do some updates and deletes, then recheck the correlation again.
+            for (int i = 0; i < numDoc / 2; i++) {
+                if (randomBoolean()) {
+                    deleteDoc(indexShard, "doc", Integer.toString(i));
+                } else {
+                    indexDoc(indexShard, "doc", Integer.toString(i), "{\"foo\": \"bar\"}");
+                }
+            }
+
             indexShard.flush(new FlushRequest());
             indexShard.refresh("test");
-            DocsStats withDeletedStats = indexShard.docStats();
-            assertThat(withDeletedStats.averageSizeInBytes, greaterThan(largeStats.averageSizeInBytes));
+            {
+                final DocsStats docsStats = indexShard.docStats();
+                final StoreStats storeStats = indexShard.storeStats();
+                assertThat("Estimated total document size is too small compared with the stored size",
+                    docsStats.getTotalSizeInBytes(), greaterThanOrEqualTo(storeStats.sizeInBytes() * 80/100));
+                assertThat("Estimated total document size is too large compared with the stored size",
+                    docsStats.getTotalSizeInBytes(), lessThanOrEqualTo(storeStats.sizeInBytes() * 120/100));
+            }
 
         } finally {
             closeShards(indexShard);