Switch MemtableAverageRowSize to use trie directly and expand test

Author: blambov

src/java/org/apache/cassandra/db/IDataSize.java

@@ -0,0 +1,25 @@
+/*
+ * 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.cassandra.db;
+
+/// Shared interface for providing data size information
+public interface IDataSize
+{
+    int dataSize();
+}

src/java/org/apache/cassandra/db/memtable/MemtableAverageRowSize.java

@@ -18,11 +18,15 @@
 
 package org.apache.cassandra.db.memtable;
 
+import java.util.function.Consumer;
+
 import org.apache.cassandra.db.DataRange;
+import org.apache.cassandra.db.IDataSize;
 import org.apache.cassandra.db.filter.ColumnFilter;
 import org.apache.cassandra.db.rows.Row;
 import org.apache.cassandra.db.rows.Unfiltered;
 import org.apache.cassandra.db.rows.UnfilteredRowIterator;
+import org.apache.cassandra.db.tries.Trie;
 
 class MemtableAverageRowSize
 {
@@ -31,6 +35,35 @@ class MemtableAverageRowSize
     public final long rowSize;
     public final long operations;
 
+    public MemtableAverageRowSize(Memtable memtable, Trie<?> trie)
+    {
+        // If this is a trie-based memtable, get the row sizes from the trie elements. This achieves two things:
+        // - makes sure the size used is the size reflected in the memtable's dataSize
+        //   (which e.g. excludes clustering keys)
+        // - avoids the conversion to Row, which has non-trivial cost
+
+        class SizeCalculator implements Trie.ValueConsumer<Object>
+        {
+            long totalSize = 0;
+            long count = 0;
+
+            @Override
+            public void accept(Object o)
+            {
+                if (o instanceof IDataSize)
+                {
+                    totalSize += ((IDataSize) o).dataSize();
+                    ++count;
+                }
+            }
+        }
+
+        SizeCalculator sizeCalculator = new SizeCalculator();
+        trie.forEachValue(sizeCalculator);
+
+        this.rowSize = sizeCalculator.count > 0 ? sizeCalculator.totalSize / sizeCalculator.count : 0;
+        this.operations = memtable.getOperations();
+    }
 
     public MemtableAverageRowSize(Memtable memtable)
     {

src/java/org/apache/cassandra/db/memtable/TrieMemtable.java

@@ -299,7 +299,7 @@ public int getShardCount()
     public long getEstimatedAverageRowSize()
     {
         if (estimatedAverageRowSize == null || currentOperations.get() > estimatedAverageRowSize.operations * 1.5)
-            estimatedAverageRowSize = new MemtableAverageRowSize(this);
+            estimatedAverageRowSize = new MemtableAverageRowSize(this, mergedTrie.contentOnlyTrie());
         return estimatedAverageRowSize.rowSize;
     }
 

src/java/org/apache/cassandra/db/memtable/TrieMemtableStage2.java

@@ -307,7 +307,7 @@ public int getShardCount()
     public long getEstimatedAverageRowSize()
     {
         if (estimatedAverageRowSize == null || currentOperations.get() > estimatedAverageRowSize.operations * 1.5)
-            estimatedAverageRowSize = new MemtableAverageRowSize(this);
+            estimatedAverageRowSize = new MemtableAverageRowSize(this, mergedTrie);
         return estimatedAverageRowSize.rowSize;
     }
 

src/java/org/apache/cassandra/db/partitions/TrieBackedPartition.java

@@ -30,6 +30,7 @@
 import org.apache.cassandra.db.ClusteringPrefix;
 import org.apache.cassandra.db.DecoratedKey;
 import org.apache.cassandra.db.DeletionTime;
+import org.apache.cassandra.db.IDataSize;
 import org.apache.cassandra.db.LivenessInfo;
 import org.apache.cassandra.db.RegularAndStaticColumns;
 import org.apache.cassandra.db.Slice;
@@ -113,7 +114,7 @@ public static boolean isPartitionBoundary(Object content)
     /// The representation of a row stored at the leaf of a trie. Does not contain the row key.
     ///
     /// The method [#toRow] combines this with a clustering for the represented [Row].
-    public static class RowData
+    public static class RowData implements IDataSize
     {
         final Object[] columnsBTree;
         final LivenessInfo livenessInfo;

src/java/org/apache/cassandra/db/partitions/TrieBackedPartitionStage2.java

@@ -31,6 +31,7 @@
 import org.apache.cassandra.db.DecoratedKey;
 import org.apache.cassandra.db.DeletionInfo;
 import org.apache.cassandra.db.DeletionTime;
+import org.apache.cassandra.db.IDataSize;
 import org.apache.cassandra.db.LivenessInfo;
 import org.apache.cassandra.db.MutableDeletionInfo;
 import org.apache.cassandra.db.RangeTombstone;
@@ -98,7 +99,7 @@ public class TrieBackedPartitionStage2 implements Partition
      *
      * The methods toRow and copyToOnHeapRow combine this with a clustering for the represented Row.
      */
-    public static class RowData
+    public static class RowData implements IDataSize
     {
         final Object[] columnsBTree;
         final LivenessInfo livenessInfo;

src/java/org/apache/cassandra/db/rows/Row.java

@@ -51,7 +51,7 @@
  * it's own data. For instance, a {@code Row} cannot contains a cell that is deleted by its own
  * row deletion.
  */
-public interface Row extends Unfiltered, Iterable<ColumnData>
+public interface Row extends Unfiltered, Iterable<ColumnData>, IDataSize
 {
     /**
      * The clustering values for this row.

test/unit/org/apache/cassandra/index/sai/plan/SingleRestrictionEstimatedRowCountTest.java

@@ -25,6 +25,8 @@
 import java.util.Map;
 
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
 
 import org.apache.cassandra.Util;
 import org.apache.cassandra.cql3.CQL3Type;
@@ -43,8 +45,18 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.fail;
 
+@RunWith(Parameterized.class)
 public class SingleRestrictionEstimatedRowCountTest extends SAITester
 {
+    @Parameterized.Parameter
+    public String memtableClass;
+
+    @Parameterized.Parameters(name = "{0}")
+    public static String[] memtableClasses()
+    {
+        return new String[]{ "TrieMemtable", "TrieMemtableStage2", "TrieMemtableStage1" };
+    }
+
     static protected Map<Map.Entry<Version, CQL3Type.Native>, ColumnFamilyStore> tables = new HashMap<>();
     static Version[] versions = new Version[]{ Version.DB, Version.EB };
     static CQL3Type.Native[] types = new CQL3Type.Native[]{ INT, DECIMAL, VARINT };
@@ -75,12 +87,12 @@ public void testMemtablesSAI()
         createTables();
 
         RowCountTest test = new RowCountTest(Operator.NEQ, 25);
-        test.doTest(Version.DB, INT, 97.0);
-        test.doTest(Version.EB, INT, 97.0);
+        test.doTest(Version.DB, INT, 100);
+        test.doTest(Version.EB, INT, 100);
         // Truncated numeric types planned differently
-        test.doTest(Version.DB, DECIMAL, 97.0);
-        test.doTest(Version.EB, DECIMAL, 97.0);
-        test.doTest(Version.EB, VARINT, 97.0);
+        test.doTest(Version.DB, DECIMAL, 100);
+        test.doTest(Version.EB, DECIMAL, 100);
+        test.doTest(Version.EB, VARINT, 100);
 
         test = new RowCountTest(Operator.LT, 50);
         test.doTest(Version.DB, INT, 48);
@@ -89,10 +101,10 @@ public void testMemtablesSAI()
         test.doTest(Version.EB, DECIMAL, 48);
 
         test = new RowCountTest(Operator.LT, 150);
-        test.doTest(Version.DB, INT, 97);
-        test.doTest(Version.EB, INT, 97);
-        test.doTest(Version.DB, DECIMAL, 97);
-        test.doTest(Version.EB, DECIMAL, 97);
+        test.doTest(Version.DB, INT, 100);
+        test.doTest(Version.EB, INT, 100);
+        test.doTest(Version.DB, DECIMAL, 100);
+        test.doTest(Version.EB, DECIMAL, 100);
 
         test = new RowCountTest(Operator.EQ, 31);
         test.doTest(Version.DB, INT, 15);
@@ -109,7 +121,8 @@ void createTables()
             SAIUtil.setLatestVersion(version);
             for (CQL3Type.Native type : types)
             {
-                createTable("CREATE TABLE %s (pk text PRIMARY KEY, age " + type + ')');
+                createTable("CREATE TABLE %s (pk text PRIMARY KEY, age " + type + ") " +
+                            "WITH memtable = {'class': '" + memtableClass + "'}");
                 createIndex("CREATE CUSTOM INDEX ON %s(age) USING 'StorageAttachedIndex'");
                 tables.put(tablesEntryKey(version, type), getCurrentColumnFamilyStore());
             }
@@ -159,7 +172,7 @@ void doTest(Version version, CQL3Type.Native type, double expectedRows)
 
             long totalRows = controller.planFactory.tableMetrics.rows;
             assertEquals(0, cfs.metrics().liveSSTableCount.getValue().intValue());
-            assertEquals(97, totalRows);
+            assertEquals(100, totalRows);
 
             Plan plan = controller.buildPlan();
             assert plan instanceof Plan.RowsIteration;