GH-39435: [JS] Add Vector.nullable (#39436)

apache · Jan 4, 2024 · 5c0fa71 · 5c0fa71
1 parent 0e597ab
commit 5c0fa71
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 9 deletions.
diff --git a/js/src/table.ts b/js/src/table.ts
@@ -114,7 +114,7 @@ export class Table<T extends TypeMap = any> {
                 } else if (typeof x === 'object') {
                     const keys = Object.keys(x) as (keyof T)[];
                     const vecs = keys.map((k) => new Vector([x[k]]));
-                    const batchSchema = schema ?? new Schema(keys.map((k, i) => new Field(String(k), vecs[i].type, vecs[i].nullCount > 0)));
+                    const batchSchema = schema ?? new Schema(keys.map((k, i) => new Field(String(k), vecs[i].type, vecs[i].nullable)));
                     const [, batches] = distributeVectorsIntoRecordBatches(batchSchema, vecs);
                     return batches.length === 0 ? [new RecordBatch(x)] : batches;
                 }

diff --git a/js/src/util/chunk.ts b/js/src/util/chunk.ts
@@ -51,6 +51,11 @@ export class ChunkedIterator<T extends DataType> implements IterableIterator<T['
     }
 }
 
+/** @ignore */
+export function computeChunkNullable<T extends DataType>(chunks: ReadonlyArray<Data<T>>) {
+    return chunks.some(chunk => chunk.nullable);
+}
+
 /** @ignore */
 export function computeChunkNullCounts<T extends DataType>(chunks: ReadonlyArray<Data<T>>) {
     return chunks.reduce((nullCount, chunk) => nullCount + chunk.nullCount, 0);

diff --git a/js/src/vector.ts b/js/src/vector.ts
@@ -24,6 +24,7 @@ import { BigIntArray, TypedArray, TypedArrayDataType } from './interfaces.js';
 import {
     isChunkedValid,
     computeChunkOffsets,
+    computeChunkNullable,
     computeChunkNullCounts,
     sliceChunks,
     wrapChunkedCall1,
@@ -132,6 +133,13 @@ export class Vector<T extends DataType = any> {
         return this.data.reduce((byteLength, data) => byteLength + data.byteLength, 0);
     }
 
+    /**
+     * Whether this Vector's elements can contain null values.
+     */
+    public get nullable() {
+        return computeChunkNullable(this.data);
+    }
+
     /**
      * The number of null elements in this Vector.
      */

diff --git a/js/test/unit/table-tests.ts b/js/test/unit/table-tests.ts
@@ -139,30 +139,32 @@ describe(`Table`, () => {
             const i32 = makeVector([i32s]);
             expect(i32).toHaveLength(i32s.length);
             expect(i32.nullCount).toBe(0);
+            expect(i32.nullable).toBe(true);
 
             const table = new Table({ i32 });
             const i32Field = table.schema.fields[0];
 
             expect(i32Field.name).toBe('i32');
             expect(i32).toHaveLength(i32s.length);
-            expect(i32Field.nullable).toBe(false);
+            expect(i32Field.nullable).toBe(true);
             expect(i32.nullCount).toBe(0);
 
             expect(i32).toEqualVector(makeVector(i32s));
         });
 
-        test(`creates a new Table from a Typed Array and force nullable`, () => {
+        test(`creates a new Table from a Typed Array and force not nullable`, () => {
             const i32s = new Int32Array(arange(new Array<number>(10)));
             const i32 = makeVector([i32s]);
             expect(i32).toHaveLength(i32s.length);
             expect(i32.nullCount).toBe(0);
+            expect(i32.nullable).toBe(true);
 
-            const table = new Table(new Schema([new Field('i32', new Int32, true)]), { i32 });
+            const table = new Table(new Schema([new Field('i32', new Int32, false)]), { i32 });
             const i32Field = table.schema.fields[0];
 
             expect(i32Field.name).toBe('i32');
             expect(i32).toHaveLength(i32s.length);
-            expect(i32Field.nullable).toBe(true);
+            expect(i32Field.nullable).toBe(false);
             expect(i32.nullCount).toBe(0);
 
             expect(i32).toEqualVector(makeVector(i32s));
@@ -187,8 +189,8 @@ describe(`Table`, () => {
             expect(f32Field.name).toBe('f32');
             expect(i32).toHaveLength(i32s.length);
             expect(f32).toHaveLength(f32s.length);
-            expect(i32Field.nullable).toBe(false);
-            expect(f32Field.nullable).toBe(false);
+            expect(i32Field.nullable).toBe(true);
+            expect(f32Field.nullable).toBe(true);
             expect(i32.nullCount).toBe(0);
             expect(f32.nullCount).toBe(0);
 
@@ -222,7 +224,7 @@ describe(`Table`, () => {
 
             expect(i32Vector).toHaveLength(i32s.length);
             expect(f32Vector).toHaveLength(i32s.length); // new length should be the same as the longest sibling
-            expect(i32Field.nullable).toBe(false);
+            expect(i32Field.nullable).toBe(true);
             expect(f32Field.nullable).toBe(true); // true, with 12 additional nulls
             expect(i32Vector.nullCount).toBe(0);
             expect(f32Vector.nullCount).toBe(i32s.length - f32s.length);
@@ -264,7 +266,7 @@ describe(`Table`, () => {
             expect(f32RenamedField.name).toBe('f32Renamed');
             expect(i32Renamed).toHaveLength(i32s.length);
             expect(f32Renamed).toHaveLength(i32s.length); // new length should be the same as the longest sibling
-            expect(i32RenamedField.nullable).toBe(false);
+            expect(i32RenamedField.nullable).toBe(true);
             expect(f32RenamedField.nullable).toBe(true); // true, with 4 additional nulls
             expect(i32Renamed.nullCount).toBe(0);
             expect(f32Renamed.nullCount).toBe(i32s.length - f32s.length);