Using validity bitmap as Arrow does (#14409)

Author: JaroslavTulach

std-bits/table/src/main/java/org/enso/table/data/column/builder/BoolBuilder.java

@@ -13,13 +13,13 @@
 /** A builder for boolean columns. */
 final class BoolBuilder implements BuilderForBoolean, BuilderWithRetyping {
   private final BitSet vals;
-  private final BitSet isNothing;
+  private final BitSet validityMap;
   int size = 0;
 
   // ** Creates a new builder for boolean columns. Should be built via Builder.getForBoolean. */
   BoolBuilder(int capacity) {
     vals = new BitSet(capacity);
-    isNothing = new BitSet(capacity);
+    validityMap = new BitSet(capacity);
   }
 
   @Override
@@ -31,6 +31,7 @@ public BoolBuilder append(Object o) {
         if (b) {
           vals.set(size);
         }
+        validityMap.set(size);
       } else {
         throw new ValueTypeMismatchException(getType(), o);
       }
@@ -54,13 +55,14 @@ public BoolBuilder appendBoolean(boolean value) {
     if (value) {
       vals.set(size);
     }
+    validityMap.set(size, true);
     size++;
     return this;
   }
 
   @Override
   public BoolBuilder appendNulls(int count) {
-    isNothing.set(size, size + count);
+    validityMap.set(size, size + count, false);
     size += count;
     return this;
   }
@@ -71,7 +73,7 @@ public void appendBulkStorage(ColumnStorage<?> storage) {
       // We know this is valid for a BoolStorage.
       int toCopy = (int) boolStorage.getSize();
       BitSets.copy(boolStorage.getValues(), vals, size, toCopy);
-      BitSets.copy(boolStorage.getIsNothingMap(), isNothing, size, toCopy);
+      boolStorage.getValidityMap().copyTo(validityMap, size, toCopy);
       size += toCopy;
     } else if (storage instanceof ColumnBooleanStorage columnBooleanStorage) {
       for (long i = 0; i < columnBooleanStorage.getSize(); i++) {
@@ -90,7 +92,7 @@ public void appendBulkStorage(ColumnStorage<?> storage) {
 
   @Override
   public ColumnStorage<Boolean> seal() {
-    return new BoolStorage(vals, isNothing, size, false);
+    return new BoolStorage(vals, validityMap, size, false);
   }
 
   @Override
@@ -101,7 +103,7 @@ public long getCurrentSize() {
   @Override
   public void copyDataTo(Object[] items) {
     for (int i = 0; i < size; i++) {
-      if (isNothing.get(i)) {
+      if (!validityMap.get(i)) {
         items[i] = null;
       } else {
         items[i] = vals.get(i);

std-bits/table/src/main/java/org/enso/table/data/column/builder/Builder.java

@@ -54,8 +54,12 @@ static ColumnStorage<?> fromRepeatedItem(Object item, long size) {
     // Create a single storage item based on the type of the item.
     return switch (item) {
       case null -> new NullBuilder().appendNulls(checkSize(size)).seal();
-      case Boolean booleanValue ->
-          new BoolStorage(new BitSet(), new BitSet(), checkSize(size), booleanValue);
+      case Boolean booleanValue -> {
+        var s = checkSize(size);
+        var validity = new BitSet();
+        validity.set(0, s, true);
+        yield new BoolStorage(new BitSet(), validity, s, booleanValue);
+      }
       default -> {
         var storageType = StorageType.forBoxedItem(item, PreciseTypeOptions.DEFAULT);
         var builder = Builder.getForType(storageType, size, BlackholeProblemAggregator.INSTANCE);

std-bits/table/src/main/java/org/enso/table/data/column/builder/DoubleBuilder.java

@@ -13,10 +13,10 @@
 import org.enso.table.data.column.storage.type.StorageType;
 import org.enso.table.error.ValueTypeMismatchException;
 import org.enso.table.problems.ProblemAggregator;
-import org.enso.table.util.BitSets;
 
 /** A builder for floating point columns. */
-class DoubleBuilder extends NumericBuilder implements BuilderForDouble {
+sealed class DoubleBuilder extends NumericBuilder implements BuilderForDouble
+    permits InferredDoubleBuilder {
   protected final PrecisionLossAggregator precisionLossAggregator;
   protected double[] data;
 
@@ -80,6 +80,7 @@ public DoubleBuilder append(Object o) {
     }
 
     ensureSpaceToAppend();
+    setValid(currentSize);
     data[currentSize++] = value;
     return this;
   }
@@ -91,7 +92,7 @@ public void appendBulkStorage(ColumnStorage<?> storage) {
         int n = (int) doubleStorage.getSize();
         ensureFreeSpaceFor(n);
         System.arraycopy(doubleStorage.getData(), 0, data, currentSize, n);
-        BitSets.copy(doubleStorage.getIsNothingMap(), isNothing, currentSize, n);
+        appendValidityMap(doubleStorage.getValidityMap(), n);
         currentSize += n;
       } else {
         var doubleStorage = floatType.asTypedStorage(storage);
@@ -147,8 +148,10 @@ public void appendBulkStorage(ColumnStorage<?> storage) {
    *
    * @param value the double to append
    */
+  @Override
   public DoubleBuilder appendDouble(double value) {
     ensureSpaceToAppend();
+    setValid(currentSize);
     data[currentSize++] = value;
     return this;
   }
@@ -158,14 +161,15 @@ public DoubleBuilder appendDouble(double value) {
    *
    * <p>It ensures that any loss of precision is reported.
    */
+  @Override
   public DoubleBuilder appendLong(long value) {
     appendDouble(convertLongToDouble(value));
     return this;
   }
 
   @Override
   public ColumnStorage<Double> seal() {
-    return new DoubleStorage(data, currentSize, isNothing);
+    return new DoubleStorage(data, currentSize, validityMap());
   }
 
   /**

std-bits/table/src/main/java/org/enso/table/data/column/builder/InferredDoubleBuilder.java

@@ -66,7 +66,7 @@ static InferredDoubleBuilder retypeFromLongBuilder(
   public void copyDataTo(Object[] items) {
     int rawN = rawData == null ? 0 : rawData.length;
     for (int i = 0; i < currentSize; i++) {
-      if (isNothing.get(i)) {
+      if (!isValid(i)) {
         items[i] = null;
       } else {
         if (isLongCompactedAsDouble.get(i)) {
@@ -160,7 +160,7 @@ public Builder retypeTo(StorageType<?> type) {
     if (type instanceof BigDecimalType) {
       Builder res = Builder.getForBigDecimal(data.length);
       for (int i = 0; i < currentSize; i++) {
-        if (isNothing.get(i)) {
+        if (!isValid(i)) {
           res.appendNulls(1);
         } else {
           BigDecimal bigDecimal = BigDecimal.valueOf(data[i]);

std-bits/table/src/main/java/org/enso/table/data/column/builder/LongBuilder.java

@@ -13,10 +13,10 @@
 import org.enso.table.data.column.storage.type.StorageType;
 import org.enso.table.error.ValueTypeMismatchException;
 import org.enso.table.problems.ProblemAggregator;
-import org.enso.table.util.BitSets;
 
 /** A builder for integer columns. */
-class LongBuilder extends NumericBuilder implements BuilderForLong, BuilderWithRetyping {
+sealed class LongBuilder extends NumericBuilder implements BuilderForLong, BuilderWithRetyping
+    permits BoundCheckedIntegerBuilder {
   protected final ProblemAggregator problemAggregator;
   protected long[] data;
 
@@ -49,7 +49,7 @@ protected void resize(int desiredCapacity) {
   @Override
   public void copyDataTo(Object[] items) {
     for (int i = 0; i < currentSize; i++) {
-      if (isNothing.get(i)) {
+      if (!isValid(i)) {
         items[i] = null;
       } else {
         items[i] = data[i];
@@ -96,7 +96,7 @@ public void appendBulkStorage(ColumnStorage<?> storage) {
           int n = (int) longStorage.getSize();
           ensureFreeSpaceFor(n);
           System.arraycopy(longStorage.getData(), 0, data, currentSize, n);
-          BitSets.copy(longStorage.getIsNothingMap(), isNothing, currentSize, n);
+          appendValidityMap(longStorage.getValidityMap(), n);
           currentSize += n;
         } else {
           // No conversions needed, but we need to iterate over the items.
@@ -134,6 +134,7 @@ public void appendBulkStorage(ColumnStorage<?> storage) {
    */
   public LongBuilder appendLong(long value) {
     ensureSpaceToAppend();
+    this.setValid(currentSize);
     this.data[currentSize++] = value;
     return this;
   }
@@ -143,7 +144,7 @@ public boolean isNothing(long index) {
     if (index >= currentSize) {
       throw new IndexOutOfBoundsException();
     } else {
-      return isNothing.get((int) index);
+      return !isValid((int) index);
     }
   }
 
@@ -185,6 +186,6 @@ public LongBuilder append(Object o) {
 
   @Override
   public ColumnStorage<Long> seal() {
-    return new LongStorage(data, currentSize, isNothing, getType());
+    return new LongStorage(data, currentSize, validityMap(), getType());
   }
 }

std-bits/table/src/main/java/org/enso/table/data/column/builder/NumericBuilder.java

@@ -1,28 +1,58 @@
 package org.enso.table.data.column.builder;
 
 import java.util.BitSet;
+import org.enso.table.util.ImmutableBitSet;
 
 /** A common base for numeric builders. */
-abstract class NumericBuilder implements Builder {
-  protected BitSet isNothing;
-  protected int currentSize;
+abstract sealed class NumericBuilder implements Builder permits DoubleBuilder, LongBuilder {
+  private BitSet validityMap;
+  int currentSize;
 
-  protected NumericBuilder() {
-    this.isNothing = new BitSet();
-    this.currentSize = 0;
+  protected NumericBuilder() {}
+
+  private BitSet getValidityMap() {
+    if (validityMap == null) {
+      validityMap = new BitSet();
+      validityMap.set(0, currentSize);
+    }
+    return validityMap;
   }
 
-  protected void doAppendNulls(int count) {
-    isNothing.set(currentSize, currentSize + count);
+  protected final void doAppendNulls(int count) {
+    getValidityMap().set(currentSize, currentSize + count, false);
     currentSize += count;
   }
 
+  protected final boolean isValid(int i) {
+    return validityMap == null || validityMap.get(i);
+  }
+
+  protected final void setValid(int i) {
+    if (validityMap != null) {
+      validityMap.set(i);
+    }
+  }
+
+  protected final ImmutableBitSet validityMap() {
+    if (validityMap == null) {
+      return ImmutableBitSet.allTrue(currentSize);
+    } else {
+      return new ImmutableBitSet(validityMap, currentSize);
+    }
+  }
+
+  protected final void appendValidityMap(ImmutableBitSet validity, int n) {
+    if (validity.cardinality() < n || validityMap != null) {
+      validity.copyTo(getValidityMap(), currentSize, n);
+    }
+  }
+
   @Override
   public long getCurrentSize() {
     return currentSize;
   }
 
-  protected void ensureFreeSpaceFor(int additionalSize) {
+  protected final void ensureFreeSpaceFor(int additionalSize) {
     if (currentSize + additionalSize > getDataSize()) {
       resize(currentSize + additionalSize);
     }
@@ -35,7 +65,7 @@ protected void ensureFreeSpaceFor(int additionalSize) {
    * appends. It tries to keep the invariant that after calling `grow` the array has at least one
    * free slot.
    */
-  protected void ensureSpaceToAppend() {
+  protected final void ensureSpaceToAppend() {
     int dataLength = getDataSize();
 
     // Check current size. If there is space, we don't need to grow.

std-bits/table/src/main/java/org/enso/table/data/column/operation/IsInOperation.java

@@ -14,7 +14,7 @@
 import org.enso.table.data.column.storage.ColumnBooleanStorage;
 import org.enso.table.data.column.storage.ColumnStorage;
 import org.enso.table.data.column.storage.ColumnStorageWithInferredStorage;
-import org.enso.table.data.column.storage.ColumnStorageWithNothingMap;
+import org.enso.table.data.column.storage.ColumnStorageWithValidityMap;
 import org.enso.table.data.column.storage.type.AnyObjectType;
 import org.enso.table.data.column.storage.type.BigDecimalType;
 import org.enso.table.data.column.storage.type.BigIntegerType;
@@ -28,6 +28,7 @@
 import org.enso.table.data.column.storage.type.TimeOfDayType;
 import org.enso.table.data.table.Column;
 import org.enso.table.data.table.problems.MapOperationProblemAggregator;
+import org.enso.table.util.ImmutableBitSet;
 
 /**
  * The IsInOperation class provides a way to check if a value is in a set of values. It checks if
@@ -265,8 +266,8 @@ private static ColumnStorage<?> applyBooleanIsIn(
 
     // If had both true and false, then return all true when not nothing
     if (flags.hadTrue && flags.hadFalse) {
-      var isNothing = makeIsNothingMap(boolStorage, checkedSize);
-      return new BoolStorage(new BitSet(), isNothing, checkedSize, true);
+      var validityMap = makeValidityMap(boolStorage, checkedSize);
+      return new BoolStorage(new BitSet(), validityMap, checkedSize, true);
     }
 
     // Only have one of true or false
@@ -296,34 +297,41 @@ private static ColumnStorage<?> applyBooleanIsIn(
   private static ColumnStorage<?> applyBoolStorage(
       boolean keepValue, BoolStorage boolStorage, int checkedSize) {
     BitSet values = boolStorage.getValues();
-    BitSet isNothing = boolStorage.getIsNothingMap();
+    BitSet isNothing = boolStorage.getValidityMap().cloneBitSet();
+    isNothing.flip(0, Math.toIntExact(boolStorage.getSize()));
 
     if (keepValue) {
       var newIsNothing =
-          boolStorage.isNegated() ? or(isNothing, values) : orNot(isNothing, values, checkedSize);
+          boolStorage.isNegated()
+              ? or(isNothing, values, checkedSize)
+              : orNot(isNothing, values, checkedSize);
+      newIsNothing.flip(0, checkedSize);
       return new BoolStorage(values, newIsNothing, checkedSize, boolStorage.isNegated());
     } else {
       var newIsNothing =
-          boolStorage.isNegated() ? orNot(isNothing, values, checkedSize) : or(isNothing, values);
+          boolStorage.isNegated()
+              ? orNot(isNothing, values, checkedSize)
+              : or(isNothing, values, checkedSize);
+      newIsNothing.flip(0, checkedSize);
       return new BoolStorage(values, newIsNothing, checkedSize, !boolStorage.isNegated());
     }
   }
 
-  private static BitSet makeIsNothingMap(ColumnStorage<?> storage, int size) {
-    if (storage instanceof ColumnStorageWithNothingMap withNothingMap) {
-      return withNothingMap.getIsNothingMap();
+  private static ImmutableBitSet makeValidityMap(ColumnStorage<?> storage, int size) {
+    if (storage instanceof ColumnStorageWithValidityMap withNothingMap) {
+      return withNothingMap.getValidityMap();
     }
 
-    BitSet isNothingMap = new BitSet(size);
+    BitSet validityMap = new BitSet(size);
     for (int i = 0; i < size; i++) {
-      if (storage.isNothing(i)) {
-        isNothingMap.set(i);
+      if (!storage.isNothing(i)) {
+        validityMap.set(i);
       }
     }
-    return isNothingMap;
+    return new ImmutableBitSet(validityMap, size);
   }
 
-  private static BitSet or(BitSet left, BitSet right) {
+  private static BitSet or(BitSet left, BitSet right, int sizeIsIgnored) {
     BitSet result = (BitSet) left.clone();
     result.or(right);
     return result;