faster insert wip

Author: mkeskells

roaringbitmap/src/main/java/org/roaringbitmap/art/Art.java

@@ -1,12 +1,17 @@
 package org.roaringbitmap.art;
 
 import org.roaringbitmap.ArraysShim;
+import org.roaringbitmap.Container;
 import org.roaringbitmap.longlong.LongUtils;
 
 import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.function.Supplier;
+import java.util.function.ToLongFunction;
+
+import static org.roaringbitmap.art.BranchNode.ILLEGAL_IDX;
 
 /**
  * See: https://db.in.tum.de/~leis/papers/ART.pdf a cpu cache friendly main memory data structure.
@@ -222,6 +227,124 @@ protected Toolkit removeSpecifyKey(Node node, byte[] key, int dep) {
     return null;
   }
 
+    /**
+     * Find or create a leaf node by the high part of the key.
+     * If the leaf node is not found, it will be created with the value returned by
+     * nextContainer.applyAsLong(ifNotFound).
+     *
+     * @param highPart the high part of the key
+     * @param nextContainer a function to get the next container index
+     * @param ifNotFound a supplier to provide a value if the key is not found
+     * @return the container index of the found or created leaf node
+     */
+  public <T> long findOrCreateByKey(long highPart, ToLongFunction<Supplier<T>> nextContainer, Supplier<T> ifNotFound) {
+    LeafNode result;
+
+    if (root == null) {
+      result = new LeafNode(highPart, nextContainer.applyAsLong(ifNotFound));
+      root = result;
+      return result.getContainerIdx();
+    }
+      byte depth = 0;
+      byte parentKeyInGrandParent = 0;
+      BranchNode grandParent = null;
+      Node parent = root;
+      Node originalParent;
+      // on  each cycle
+      // if gParent == null, parent will be the next root
+      // if gParent != null, parent should replace grandParent at key parentKeyInGrandParent
+      // depth is the depth of parent in the trie
+      // result is the leaf node, or null if we are just finding
+
+      //Parent can grow, so we need to keep track of the parent index in grandParent
+      // but grandParent cant.
+      while (true) {
+        //keep a copy of the original parent, so we can see if it changes, and adjust the tree
+        originalParent = parent;
+
+        //usually a branch node, so lets test that first
+        if (parent instanceof BranchNode) {
+          BranchNode parentBranch = (BranchNode) parent;
+          //is parent the real parent - does the prefix match?
+          byte prefixLength = parentBranch.prefixLength();
+          if (prefixLength > 0) {
+            byte matchLength = prefixMatchLength(LongUtils.fromArray(parentBranch.prefix), depth, prefixLength, highPart);
+            if (matchLength == prefixLength) {
+              // prefix matches, so we can just continue
+              depth += prefixLength;
+              continue;
+            } else {
+              //so we have a partial match, we need to split the branch
+              // for example, if the prefix is [1,2,3] and the highPart is 0xab99010204, and depth 2
+              // we have a match length of 2,
+              // so we create a new parent with prefix is [1,2]
+              // with 2 children - the current parent prefix shrunk to [], at key 3
+              // and add a leaf node with key 4
+              byte branchKey = parentBranch.prefix[matchLength];
+              parentBranch = parentBranch.shrinkPrefixBy(matchLength + 1);
+              result = new LeafNode(highPart, nextContainer.applyAsLong(ifNotFound));
+              parent = Node4.create(parentBranch, result, branchKey, LongUtils.getByte(highPart, depth + matchLength), highPart, depth, (byte) depth + matchLength);
+              break;
+            }
+          }
+          //OK so parent is ok, and depth is adjusted. Let try to move to the next level
+
+          byte childKey = LongUtils.getByte(highPart, depth);
+          // We optimise for the case where the childKey is already present at this level, and we are walking a tree
+          // it should be the common case, so we try to find the child at this level
+          //we could calculate the position for the access, but that means that we have to have 2 accesses on the common path,
+          // so they shoul dbe faster
+          Node childNode = parentBranch.getChildAtKey(childKey);
+          if (childNode == null) {
+            result = new LeafNode(highPart, nextContainer.applyAsLong(ifNotFound));
+            parent = parentBranch.insert(result, childKey);
+            break;
+          } else {
+            grandParent = parentBranch;
+            parentKeyInGrandParent = childKey;
+            parent = childNode;
+            depth += 1;
+
+          }
+        } else {
+          LeafNode leafNode = (LeafNode) parent;
+          long leafNodeKey = leafNode.getKey();
+          if (leafNodeKey == highPart) {
+            result = leafNode;
+          } else {
+            //we have to create a new Node4 with the two leaves
+            result = new LeafNode(highPart, nextContainer.applyAsLong(ifNotFound));
+            byte matchLength = prefixMatchLength(leafNodeKey, depth, (byte)6, highPart);
+
+            // create a new parent node4 with the current leaf node and the new leaf node
+            Node4 split = Node4.create(leafNode, result,
+                    LongUtils.getByte(leafNodeKey, depth + matchLength), LongUtils.getByte(highPart, depth + matchLength),
+                    highPart, depth, (byte) (depth + matchLength));
+            parent = split;
+          }
+          break;
+        }
+      }
+      if (grandParent == null) {
+        root = parent;
+      } else if (parent != originalParent) {
+        // if the parent has changed, we need to replace it in the grandParent
+        int pos = grandParent.getChildPos(parentKeyInGrandParent);
+        grandParent.replaceNode(pos, parent);
+      }
+    return result.getContainerIdx();
+  }
+
+  private byte prefixMatchLength(long prefix, byte depth, byte prefixLength, long highPart) {
+    for (byte i = 0; i < prefixLength; i++) {
+      if (LongUtils.getByte(prefix,i) != LongUtils.getByte(highPart, depth + i)) {
+        return i;
+      }
+    }
+    return prefixLength;
+  }
+
+
   class Toolkit {
 
     Node freshMatchedParentNode; // indicating a fresh parent node while the original

roaringbitmap/src/main/java/org/roaringbitmap/art/BranchNode.java

@@ -3,6 +3,7 @@
 import java.io.DataOutput;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.Arrays;
 
 public abstract class BranchNode extends Node {
 
@@ -233,5 +234,21 @@ protected int serializeHeaderSizeInBytes() {
         return super.serializeHeaderSizeInBytes() + prefixLength();
     }
 
-
+    /**
+     * shrink the prefix by the specified number of bytes
+     * @param bytes the number of bytes to shrink, must be greater than 0 and less than or equal to the current prefix length
+     * @return this node, or an equivenent node with a smaller prefix
+     */
+    BranchNode shrinkPrefixBy(int bytes) {
+        assert bytes > 0 && bytes <= prefixLength();
+        byte prefixLength = this.prefixLength();
+        if (bytes == prefixLength) {
+            // shrink to empty prefix
+            this.prefix = Art.EMPTY_BYTES;
+        } else {
+            prefix = Arrays.copyOfRange(prefix, bytes, prefixLength);
+        }
+        //in the future we may remove the prefix as an array, and in that case the node may be different
+        return this;
+    }
 }

roaringbitmap/src/main/java/org/roaringbitmap/art/Node4.java

@@ -16,6 +16,40 @@ public class Node4 extends BranchNode {
   public Node4(int compressedPrefixSize) {
     super(compressedPrefixSize);
   }
+  public static Node4 create(Node node1, Node node2,  byte node1Key, byte node2Key,
+                             long prefixSource, byte prefixStart, byte prefixlength) {
+    if ((char)node1Key > (char)node2Key) {
+      return createOrdered(node1, node2, node1Key, node2Key, prefixSource, prefixStart, prefixlength);
+    } else {
+      return createOrdered(node2, node1, node2Key, node1Key, prefixSource, prefixStart, prefixlength);
+    }
+  }
+  /**
+   * create a Node4 with two child nodes and their keys.
+   * It is required that node1Key < node2Key when conpared as unsigned bytes.
+   * The child nodes must have the prefixes adjusted before callingt this method.
+   * @param node1 the first node
+   * @param node2 the second node
+   * @param node1Key the key of the first node
+   * @param node2Key the key of the second node
+   * @param prefixSource the source of the prefix, in big endian order
+   * @param prefixStart the start index of the prefix in the source
+   * @param prefixLength the end index of the prefix in the source
+   * @return a Node4 instance with the two nodes as children
+   */
+  public static Node4 createOrdered(Node node1, Node node2,  byte node1Key, byte node2Key,
+                             long prefixSource, byte prefixStart, byte prefixLength) {
+    Node4 node4 = new Node4(prefixLength);
+    node4.count = 2;
+    node4.key = IntegerUtil.init2Bytes(node1Key, node2Key);
+    node4.children[0] = node1;
+    node4.children[1] = node2;
+    for (int i = prefixStart; i < prefixStart + prefixLength; i++) {
+      node4.prefix[i - prefixStart] = LongUtils.getByte(prefixSource, i);
+    }
+    return node4;
+  }
+
 
   @Override
   protected NodeType nodeType() {

roaringbitmap/src/main/java/org/roaringbitmap/longlong/HighLowContainer.java

@@ -16,6 +16,7 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.NoSuchElementException;
+import java.util.function.Supplier;
 
 public class HighLowContainer {
 
@@ -72,6 +73,12 @@ public void put(byte[] highPart, Container container) {
     art.insert(highPart, containerIdx);
   }
 
+  public ContainerWithIndex findOrCreateContainer(long highPart, Supplier<Container> ifNotFound) {
+    long containerIdx = art.findOrCreateByKey(highPart, containers::addContainer, ifNotFound);
+      Container container = containers.getContainer(containerIdx);
+      return new ContainerWithIndex(container, containerIdx);
+  }
+
   /**
    * Attempt to remove the container that corresponds to the 48 bit key.
    * @param highPart the 48 bit key
@@ -325,4 +332,5 @@ public boolean equals(Object object) {
     }
     return false;
   }
+
 }

roaringbitmap/src/main/java/org/roaringbitmap/longlong/IntegerUtil.java

@@ -41,6 +41,9 @@ public static int setByte(int v, byte bv, int pos) {
     v |= (bv & 0xFF) << i;
     return v;
   }
+  public static int init2Bytes(byte b1, byte b2) {
+    return (b1 << 24) | ((b2 & 0xFF) << 16);
+  }
 
   /**
    * shift the byte left from the specified position

roaringbitmap/src/main/java/org/roaringbitmap/longlong/LongUtils.java

@@ -114,6 +114,19 @@ public static long fromKey(byte[] key) {
             | (long) (key[4] & 0xff) << 24
             | (long) (key[5] & 0xff) << 16;
   }
+  /**
+   * get the long from the big endian representation bytes
+   *
+   * @param key the byte array. Max length 6 bytes
+   * @return the long data
+   */
+  public static long fromArray(byte[] key) {
+    long result = 0;
+    for (int i = 0; i < key.length; i++) {
+      result |= ((long) (key[i] & 0xff)) << ((5 - i) * 8);
+    }
+    return result;
+  }
 
   /**
    * initialize a long value with the given fist 32 bit

roaringbitmap/src/main/java/org/roaringbitmap/longlong/Roaring64Bitmap.java

@@ -67,6 +67,14 @@ public void addLong(long x) {
       highLowContainer.put(high, arrayContainer);
     }
   }
+  public void addLongNew(long x) {
+    long high = LongUtils.highPartOnly(x);
+    char low = LongUtils.lowPart(x);
+    ContainerWithIndex containerWithIndex = highLowContainer.searchOrCreateContainer(high, ArrayContainer::new);
+    Container container = containerWithIndex.getContainer();
+    Container freshOne = container.add(low);
+    highLowContainer.replaceContainer(containerWithIndex.getContainerIdx(), freshOne);
+  }
 
   /**
    * Returns the number of distinct integers added to the bitmap (e.g., number of bits set).

roaringbitmap/src/test/java/bench/SimpleLookup.java

@@ -0,0 +1,107 @@
+package bench;
+
+import org.roaringbitmap.longlong.Roaring64Bitmap;
+
+import java.util.Arrays;
+import java.util.Random;
+import com.sun.management.ThreadMXBean;
+import java.lang.management.ManagementFactory;
+import java.util.function.Supplier;
+
+public class SimpleLookup {
+
+  public static void main(String[] args) throws Exception {
+    SimpleLookup lookup = new SimpleLookup();
+
+//    lookup.runBenchmark(lookup::lookupPopulated, "lookupPopulated", 5);
+    lookup.runBenchmark(lookup::addSorted, "addSorted", 5);
+//    lookup.runBenchmark(lookup::addUnsorted, "addUnsorted", 5);
+
+//    lookup.runBenchmark(lookup::addSortedSmall, "addSortedSmall", 50);
+//    lookup.runBenchmark(lookup::addUnsortedSmall, "addUnsortedSmall", 50);
+  }
+
+  private void runBenchmark(Supplier<Object> fn, String name, int runs) {
+    long[] durations = new long[runs];
+    long[] allocations = new long[runs];
+    Object[] results = new Object[runs];
+    long currentThreadId = Thread.currentThread().getId();
+    ThreadMXBean threadMxBean = (ThreadMXBean)
+            ManagementFactory.getThreadMXBean();
+
+    for (int i = 0; i < 5; i++) {
+      System.out.println("Running " + name + " benchmark warmup " + i);
+      results[i] = fn.get();
+    }
+
+    for (int i = 0; i < runs; i++) {
+      System.out.println("Running " + name + " benchmark iteration " + i);
+      long allocBefore = threadMxBean.getThreadAllocatedBytes(currentThreadId);
+      long startTime = System.nanoTime();
+      Object result = fn.get();
+      long endTime = System.nanoTime();
+      long allocAfter = threadMxBean.getThreadAllocatedBytes(currentThreadId);
+      durations[i] = endTime - startTime;
+      allocations[i] = allocAfter - allocBefore;
+      results[i] = result;
+      System.out.println("Iteration " + i + ": Duration: " + durations[i] + " ns, Allocated: " + allocations[i] + " bytes, Result: " + results[i]);
+    }
+    System.out.println("Average Duration: " + Arrays.stream(durations).asDoubleStream().map(x -> x / 1000000).summaryStatistics() + " ms");
+    System.out.println("Average Allocation: " + Arrays.stream(allocations).asDoubleStream().map(x -> x / 1024 / 1024).summaryStatistics() + " MB");
+
+  }
+
+  //  private Roaring64Bitmap testPopulated = new Roaring64Bitmap();
+  private SimpleLookup() {
+    //    for (long index : indexes) {
+    //      testPopulated.addLong(index);
+    //    }
+  }
+
+  private long[] indexes = new Random(0L).longs().limit(10000000).toArray();
+  private long[] sortedIndexes = Arrays.stream(indexes).sorted().toArray();
+
+  private long[] indexesSmall = new Random(0L).longs().limit(100000).toArray();
+  private long[] sortedIndexesSmall = Arrays.stream(indexesSmall).sorted().toArray();
+
+
+//  public int lookupPopulated() {
+//    int result = 0;
+//    for (long index : indexes) {
+//      if (testPopulated.contains(index)) {;
+//        result++;
+//      }
+//    }
+//    return result;
+//  }
+
+  public Boolean addSorted() {
+    Roaring64Bitmap test = new Roaring64Bitmap();
+    for (long index : sortedIndexes) {
+      test.addLong(index);
+    }
+    return test.isEmpty();
+  }
+  public Boolean addUnsorted() {
+    Roaring64Bitmap test = new Roaring64Bitmap();
+    for (long index : indexes) {
+      test.addLong(index);
+    }
+    return test.isEmpty();
+  }
+  public Boolean addSortedSmall() {
+    Roaring64Bitmap test = new Roaring64Bitmap();
+    for (long index : sortedIndexesSmall) {
+      test.addLong(index);
+    }
+    return test.isEmpty();
+  }
+  public Boolean addUnsortedSmall() {
+    Roaring64Bitmap test = new Roaring64Bitmap();
+    for (long index : indexesSmall) {
+      test.addLong(index);
+    }
+    return test.isEmpty();
+  }
+
+}