apacheGH-40851: [JS] Fix nullcount and make vectors created from type…

…d arrays not nullable
domoritz · Mar 27, 2024 · 30b984f · 30b984f
1 parent f710ac5
commit 30b984f
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 10 deletions.
diff --git a/js/src/data.ts b/js/src/data.ts
@@ -109,7 +109,10 @@ export class Data<T extends DataType = DataType> {
         let nullCount = this._nullCount;
         let nullBitmap: Uint8Array | undefined;
         if (nullCount <= kUnknownNullCount && (nullBitmap = this.nullBitmap)) {
-            this._nullCount = nullCount = this.length - popcnt_bit_range(nullBitmap, this.offset, this.offset + this.length);
+            this._nullCount = nullCount = nullBitmap.length === 0 ?
+                // no null bitmap, so all values are valid
+                0 :
+                this.length - popcnt_bit_range(nullBitmap, this.offset, this.offset + this.length);
         }
         return nullCount;
     }
@@ -177,16 +180,16 @@ export class Data<T extends DataType = DataType> {
                 // if we have a nullBitmap, truncate + slice and set it over the pre-filled 1s
                 if (this.nullCount > 0) {
                     nullBitmap.set(truncateBitmap(offset, length, this.nullBitmap), 0);
+                    Object.assign(this, { nullBitmap });
+                } else {
+                    Object.assign(this, { nullBitmap, _nullCount: 0 });
                 }
-                Object.assign(this, { nullBitmap, _nullCount: -1 });
             }
 
             const byte = nullBitmap[byteOffset];
 
             prev = (byte & mask) !== 0;
-            value ?
-                (nullBitmap[byteOffset] = byte | mask) :
-                (nullBitmap[byteOffset] = byte & ~mask);
+            nullBitmap[byteOffset] = value ? (byte | mask) : (byte & ~mask);
         }
 
         if (prev !== !!value) {

diff --git a/js/src/vector.ts b/js/src/vector.ts
@@ -445,7 +445,7 @@ export function makeVector(init: any) {
             if (init instanceof DataView) {
                 init = new Uint8Array(init.buffer);
             }
-            const props = { offset: 0, length: init.length, nullCount: 0, data: init };
+            const props = { offset: 0, length: init.length, nullCount: -1, data: init };
             if (init instanceof Int8Array) { return new Vector([makeData({ ...props, type: new dtypes.Int8 })]); }
             if (init instanceof Int16Array) { return new Vector([makeData({ ...props, type: new dtypes.Int16 })]); }
             if (init instanceof Int32Array) { return new Vector([makeData({ ...props, type: new dtypes.Int32 })]); }

diff --git a/js/test/unit/vector/vector-tests.ts b/js/test/unit/vector/vector-tests.ts
@@ -16,7 +16,7 @@
 // under the License.
 
 import {
-    Bool, DateDay, DateMillisecond, Dictionary, Float64, Int32, List, makeVector, Struct, Timestamp, TimeUnit, Utf8, LargeUtf8, util, Vector, vectorFromArray
+    Bool, DateDay, DateMillisecond, Dictionary, Float64, Int32, List, makeVector, Struct, Timestamp, TimeUnit, Utf8, LargeUtf8, util, Vector, vectorFromArray, makeData
 } from 'apache-arrow';
 
 describe(`makeVectorFromArray`, () => {
@@ -33,6 +33,47 @@ describe(`makeVectorFromArray`, () => {
     });
 });
 
+describe(`basic vector methods`, () => {
+    test(`not nullable`, () => {
+        const vector = makeVector([makeData({ data: new Int32Array([1, 2, 3]), nullCount: -1, type: new Int32() })]);
+        expect(vector.nullable).toBe(false);
+        expect(vector.nullCount).toBe(0);
+    });
+
+    test(`nullable`, () => {
+        const vector = makeVector([makeData({ data: new Int32Array([1, 2, 3]), nullCount: 0, type: new Int32() })]);
+        expect(vector.nullable).toBe(true);
+        expect(vector.nullCount).toBe(0);
+        expect(vector.isValid(0)).toBe(true);
+
+        // set a value to null
+        vector.set(0, null);
+        expect(vector.nullable).toBe(true);
+        expect(vector.nullCount).toBe(1);
+        expect(vector.isValid(0)).toBe(false);
+
+        // set the same value to null which should not change anything
+        vector.set(0, null);
+        expect(vector.nullable).toBe(true);
+        expect(vector.nullCount).toBe(1);
+
+        // set a different value to null
+        vector.set(1, null);
+        expect(vector.nullable).toBe(true);
+        expect(vector.nullCount).toBe(2);
+
+        // set first value to non-null
+        vector.set(0, 1);
+        expect(vector.nullable).toBe(true);
+        expect(vector.nullCount).toBe(1);
+
+        // set last null to non-null
+        vector.set(1, 2);
+        expect(vector.nullable).toBe(true);
+        expect(vector.nullCount).toBe(0);
+    });
+});
+
 describe(`StructVector`, () => {
     test(`makeVectorFromArray`, () => {
         const values: { a?: number; b?: string | null; c?: boolean | null }[] = [
@@ -108,7 +149,6 @@ describe(`DateVector`, () => {
 });
 
 describe(`DictionaryVector`, () => {
-
     const dictionary = ['foo', 'bar', 'baz'];
     const extras = ['abc', '123']; // values to search for that should NOT be found
     const dictionary_vec = vectorFromArray(dictionary, new Utf8).memoize();
@@ -117,7 +157,6 @@ describe(`DictionaryVector`, () => {
     const validity = Array.from({ length: indices.length }, () => Math.random() > 0.2);
 
     describe(`index with nullCount == 0`, () => {
-
         const values = indices.map((d) => dictionary[d]);
         const vector = makeVector({
             data: indices,
@@ -133,7 +172,6 @@ describe(`DictionaryVector`, () => {
     });
 
     describe(`index with nullCount > 0`, () => {
-
         const nullBitmap = util.packBools(validity);
         const nullCount = validity.reduce((acc, d) => acc + (d ? 0 : 1), 0);
         const values = indices.map((d, i) => validity[i] ? dictionary[d] : null);