Sparse support in arrow v1

cnuernber · cnuernber · commit 2e14d774cb92 · 2025-09-12T05:53:56.000-06:00
diff --git a/deps.edn b/deps.edn
@@ -244,6 +244,13 @@
               "--illegal-access=permit" "--add-opens=java.base/java.lang=ALL-UNNAMED"
               "--add-opens=java.base/java.util.concurrent.atomic=ALL-UNNAMED"
               "--add-opens=java.base/java.util.concurrent=ALL-UNNAMED"
+              "--add-opens=java.base/java.lang=ALL-UNNAMED"
+              "--add-opens=java.base/java.math=ALL-UNNAMED"
+              "--add-opens=java.base/jdk.internal.foreign=ALL-UNNAMED"
+              "--add-opens=java.base/sun.nio.ch=ALL-UNNAMED"
+              "--add-opens=java.base/java.io=ALL-UNNAMED"
+	      "--add-opens=java.base/java.util=ALL-UNNAMED"
+              "--add-opens=java.base/java.lang.ref=ALL-UNNAMED"
               ]
    }
 
diff --git a/java/tech/v3/dataset/ByteValidity.java b/java/tech/v3/dataset/ByteValidity.java
@@ -0,0 +1,79 @@
+package tech.v3.dataset;
+
+import clojure.lang.RT;
+import clojure.lang.IDeref;
+import java.util.function.LongConsumer;
+import ham_fisted.ArrayLists;
+import org.roaringbitmap.RoaringBitmap;
+
+
+public class ByteValidity {
+  public static int trimIndexes(int[] indexes, int nIndexes, long maxIdx) {
+    while(nIndexes > 0) {
+      if(Integer.toUnsignedLong(indexes[nIndexes-1]) >= maxIdx)
+	--nIndexes;
+      else
+	break;
+    }
+    return nIndexes;
+  }
+  public static abstract class ValidityBase implements LongConsumer, IDeref {
+    long nElems;
+    int idx;
+    int nIndexes;
+    int[] indexes;
+    public ValidityBase(long nElems, long maxIndexes) {
+      this.nElems = nElems;
+      indexes = new int[(int)maxIndexes];
+      nIndexes = 0;
+      idx = 0;
+    }
+  }
+  public static class ValidityIndexReducer extends ValidityBase {
+    public ValidityIndexReducer(long nElems, long maxIndexes) {
+      super(nElems, maxIndexes);
+    }
+    public void accept(long value) {
+      if(value != 0) {
+	int intVal = (int)value;
+	int offset = idx * 8;
+	if( (intVal & 1) == 1) indexes[nIndexes++] = offset;
+	if( (intVal & 2) == 2) indexes[nIndexes++] = offset+1;
+	if( (intVal & 4) == 4) indexes[nIndexes++] = offset+2;
+	if( (intVal & 8) == 8) indexes[nIndexes++] = offset+3;
+	if( (intVal & 16) == 16) indexes[nIndexes++] = offset+4;
+	if( (intVal & 32) == 32) indexes[nIndexes++] = offset+5;
+	if( (intVal & 64) == 64) indexes[nIndexes++] = offset+6;
+	if( (intVal & 128) == 128) indexes[nIndexes++] = offset+7;
+      }
+      ++idx;
+    }
+    public Object deref() {
+      return ArrayLists.toList(indexes).subList(0, trimIndexes(indexes, nIndexes, nElems)); }
+  }
+  public static class MissingIndexReducer extends ValidityBase {
+    public MissingIndexReducer(long nElems, long maxIndexes) {
+      super(nElems, maxIndexes);
+    }
+    public void accept(long value) {
+      if(value != -1) {
+	int intVal = (int)value;
+	int offset = idx * 8;
+	if( (intVal & 1) != 1) indexes[nIndexes++] = offset;
+	if( (intVal & 2) != 2) indexes[nIndexes++] = offset+1;
+	if( (intVal & 4) != 4) indexes[nIndexes++] = offset+2;
+	if( (intVal & 8) != 8) indexes[nIndexes++] = offset+3;
+	if( (intVal & 16) != 16) indexes[nIndexes++] = offset+4;
+	if( (intVal & 32) != 32) indexes[nIndexes++] = offset+5;
+	if( (intVal & 64) != 64) indexes[nIndexes++] = offset+6;
+	if( (intVal & 128) != 128) indexes[nIndexes++] = offset+7;
+      }
+      ++idx;
+    }
+    public Object deref() {
+      RoaringBitmap rb = new RoaringBitmap();
+      rb.addN(indexes, 0, trimIndexes(indexes, nIndexes, nElems));
+      return rb;
+    }
+  }
+}
diff --git a/src/tech/v3/dataset/base.clj b/src/tech/v3/dataset/base.clj
@@ -858,9 +858,7 @@
     (throw (Exception.
             (format "Column %s does not have :string datatype"
                     (ds-col/column-name col)))))
-  (if (not (instance? StringTable (.data ^Column col)))
-    (str-table/string-table-from-strings col)
-    (.data ^Column col)))
+  (str-table/->string-table (ds-proto/column-data col)))
 
 
 (defn ensure-dataset-string-tables
diff --git a/src/tech/v3/dataset/impl/column.clj b/src/tech/v3/dataset/impl/column.clj
@@ -30,7 +30,7 @@
 (set! *unchecked-math* :warn-on-boxed)
 
 
-(declare new-column)
+(declare new-column construct-column)
 
 
 (defn ->persistent-map
@@ -381,6 +381,8 @@
   (row-count [this] (dtype/ecount data))
   ds-proto/PMissing
   (missing [this] missing)
+  (num-missing [this] (dtype/ecount missing))
+  (any-missing? [this] (boolean (not (.isEmpty missing))))
   ds-proto/PValidRows
   (valid-rows [this] (.xor (->bitmap (hamf/range (dtype/ecount this))) missing))
   ds-proto/PSelectRows
@@ -400,6 +402,8 @@
   (is-column? [_this] true)
   (column-buffer [_this] data)
   (empty-column? [this] (== (dtype/ecount this) (dtype/ecount missing)))
+  (column-data [_this] data)
+  (with-column-data [this new-data] (construct-column missing new-data metadata))
   ds-proto/PColumnName
   (column-name [this] (get metadata :name))
   IMutList
diff --git a/src/tech/v3/dataset/impl/dataset.clj b/src/tech/v3/dataset/impl/dataset.clj
@@ -254,6 +254,10 @@
   ds-proto/PMissing
   (missing [this]
     (RoaringBitmap/or (.iterator ^Iterable (lznc/map ds-proto/missing (.values this)))))
+  (num-missing [this]
+    (dtype/ecount (ds-proto/missing this)))
+  (any-missing? [this]
+    (boolean (some ds-proto/any-missing? (.values this))))
 
   ds-proto/PValidRows
   (valid-rows [this]
diff --git a/src/tech/v3/dataset/impl/sparse_column.clj b/src/tech/v3/dataset/impl/sparse_column.clj
@@ -7,6 +7,7 @@
             [tech.v3.datatype.pprint :as dtype-pp]
             [tech.v3.dataset.protocols :as ds-proto]
             [tech.v3.dataset.impl.column :as col-impl]
+            [tech.v3.dataset.string-table :as str-t]
             [ham-fisted.set :as set]
             [ham-fisted.api :as hamf]
             [ham-fisted.reduce :as hamf-rf]
@@ -18,7 +19,7 @@
 
 (set! *warn-on-reflection* true)
 (set! *unchecked-math* :warn-on-boxed)
-(declare scol->reader reduce-scol kv-reduce-scol)
+(declare scol->reader reduce-scol kv-reduce-scol construct-sparse-col)
 
 (defn- index-bin-search
   ^clojure.lang.IFn$LL [^IMutList indexes]
@@ -87,12 +88,16 @@
         (let [miss-data (set/difference (bitmap/->bitmap (hamf/range rc)) (bitmap/->bitmap indexes))]
           (set! (.-msng this) miss-data)
           miss-data))))
+  (num-missing [this] (- rc (.size indexes)))
+  (any-missing? [this] (not (== rc (.size indexes))))
   ds-proto/PValidRows
   (valid-rows [this] indexes)
   ds-proto/PColumn
   (is-column? [this] true)
   (column-buffer [this] [indexes data])
   (empty-column? [this] (.isEmpty indexes))
+  (column-data [this] data)
+  (with-column-data [this new-data] (construct-sparse-col indexes new-data rc metadata))
   ds-proto/PColumnName
   (column-name [this] (get metadata :name))
   ds-proto/PRowCount
@@ -280,9 +285,6 @@
               nil valid-indexes)
       (SparseCol. valid-indexes data rc (meta col) nil nil))))
 
-(defn is-scol? [sc] (instance? SparseCol sc))
-(defn as-scol ^SparseCol [sc] (when (instance? SparseCol sc) sc))
-
 (defn scol->reader
   (^Buffer [^SparseCol scol] (scol->reader scol 0 (.-rc scol)))
   (^Buffer [^SparseCol scol ^long sidx ^long eidx]
@@ -356,14 +358,18 @@
          (ChunkedList/sublistCheck ssidx seidx rc)
          (sparse-rows ds (+ sidx ssidx) (+ sidx seidx)))))))
 
-(defn ds->sparse-ds
-  [ds]
-  (let [rc (long (ds-proto/row-count ds))]
-    (reduce (fn [ds col]
-              (if (> (dt/ecount (ds/missing col)) (quot rc 2))
-                (assoc ds (ds-proto/column-name col) (->scol col))
-                ds))
-            ds (.values ^java.util.Map ds))))
+(defn ->sparse-ds
+  ([ds] (->sparse-ds ds 0.5))
+  ([ds missing-frac]
+   (let [rc (long (ds-proto/row-count ds))
+         missing-cutoff (long (* rc (double missing-frac)))]
+     (reduce (fn [ds col]
+               (if (>= (dt/ecount (ds-proto/missing col)) missing-cutoff)
+                 (assoc ds (ds-proto/column-name col) (->scol col))
+                 ds))
+             ds (.values ^java.util.Map ds)))))
+
+(defn is-sparse? [col] (instance? SparseCol col))
 
 (comment
   (require '[tech.v3.dataset :as ds])
diff --git a/src/tech/v3/dataset/protocols.clj b/src/tech/v3/dataset/protocols.clj
@@ -14,7 +14,12 @@
   (^long column-count [this]))
 
 (defprotocol PMissing
-  (missing [this]))
+  (missing [this])
+  (num-missing [this])
+  (any-missing? [this]))
+
+(defn missing-count
+  ^long [this] (num-missing this))
 
 (defprotocol PValidRows
   (valid-rows [this]))
@@ -30,9 +35,13 @@
 
 (defprotocol PColumn
   (is-column? [col])
-  (column-buffer [col])
-  (empty-column? [col]))
-
+  (column-buffer [col]
+    "Return either the backing dataset for dense columns or tuple of [indexes,data] for sparse columns")
+  (empty-column? [col])
+  (column-data [col]
+    "Return the backing data store.  Note for sparse columns this may be much shorter than the column")
+  (with-column-data [col new-data]
+    "new data must be same length as old data"))
 
 (defprotocol PDataset
   (is-dataset? [item])
diff --git a/src/tech/v3/dataset/string_table.clj b/src/tech/v3/dataset/string_table.clj
@@ -189,24 +189,28 @@
     (StringTable. (hamf/vec int->str) (.clone str->int)
                   (compress-indexes indexes (long (.size str->int))))))
 
-(defn fast-string-container []
-  (let [str->int (hamf/java-hashmap)
-        int->str (ArrayList.)
-        _ (do (.put str->int "" 0)
-              (.add int->str ""))]
-    (FastStringContainer. (hamf/long-array-list)
-                          int->str str->int)))
+(defn fast-string-container
+  ([str->int int->str]
+   (FastStringContainer. (hamf/long-array-list)
+                         int->str str->int))
+  ([]
+   (let [str->int (hamf/java-hashmap)
+         int->str (ArrayList.)]
+     (.put str->int "" 0)
+     (.add int->str "")
+     (fast-string-container str->int int->str))))
 
 (defn string-table-from-strings
-  [str-data]
-  (hamf-rf/reduce-reducer (hamf-rf/consumer-reducer fast-string-container) str-data))
+  ([str-data] (string-table-from-strings (fast-string-container) str-data))
+  ([fast-string-container str-data]
+   (hamf-rf/reduce-reducer (hamf-rf/consumer-reducer (constantly fast-string-container)) str-data)))
 
 
 (defn ->string-table
   ^StringTable [str-t]
-  (errors/when-not-errorf (instance? StringTable str-t)
-                          "string table is wrong type: %s" str-t)
-  str-t)
+  (if (instance? StringTable str-t)
+    str-t
+    (string-table-from-strings str-t)))
 
 
 (defn indices
diff --git a/src/tech/v3/libs/arrow.clj b/src/tech/v3/libs/arrow.clj
diff --git a/test/tech/v3/libs/arrow_test.clj b/test/tech/v3/libs/arrow_test.clj

Original file line number	Diff line number	Diff line change
`@@ -244,6 +244,13 @@`
`244`	`244`	`"--illegal-access=permit" "--add-opens=java.base/java.lang=ALL-UNNAMED"`
`245`	`245`	`"--add-opens=java.base/java.util.concurrent.atomic=ALL-UNNAMED"`
`246`	`246`	`"--add-opens=java.base/java.util.concurrent=ALL-UNNAMED"`
	`247`	`+ "--add-opens=java.base/java.lang=ALL-UNNAMED"`
	`248`	`+ "--add-opens=java.base/java.math=ALL-UNNAMED"`
	`249`	`+ "--add-opens=java.base/jdk.internal.foreign=ALL-UNNAMED"`
	`250`	`+ "--add-opens=java.base/sun.nio.ch=ALL-UNNAMED"`
	`251`	`+ "--add-opens=java.base/java.io=ALL-UNNAMED"`
	`252`	`+ "--add-opens=java.base/java.util=ALL-UNNAMED"`
	`253`	`+ "--add-opens=java.base/java.lang.ref=ALL-UNNAMED"`
`247`	`254`	`]`
`248`	`255`	`}`
`249`	`256`