25756190722

Committed 12 May 2026 07:06PM UTC coverage: 75.119% (+0.2%) from 74.932%

Build # 25756190722

Build Type

push

github

Committed by

web-flow

Commit Message

feat(arrow) Support RecordBatch stream to ArrowGPUTable (#2611)

Coverage Stats

5973 of 8932 branches covered (66.87%)

Branch coverage included in aggregate %.

271 of 333 new or added lines in 9 files covered. (81.38%)

3 existing lines in 2 files now uncovered.

13032 of 16368 relevant lines covered (79.62%)

831.06 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

78.31

/modules/arrow/src/arrow/arrow-gpu-vector.ts

// luma.gl
// SPDX-License-Identifier: MIT
// Copyright (c) vis.gl contributors

import {
  Buffer,
  Device,
  type BufferAttributeLayout,
  type BufferLayout,
  type BufferProps,
  type BigTypedArray
} from '@luma.gl/core';
import * as arrow from 'apache-arrow';
import type {AttributeArrowType, NumericArrowType} from './arrow-types';
import {getArrowVectorBufferSource} from './arrow-fixed-size-list';

/** Buffer creation props forwarded when uploading Arrow vector memory to the GPU. */
export type ArrowGPUVectorBufferProps = Omit<BufferProps, 'byteLength' | 'data'>;

/** @deprecated Use {@link ArrowGPUVectorBufferProps}. */
export type ArrowGPUVectorProps = ArrowGPUVectorBufferProps;

/** Constructor props that upload an Arrow vector into a new GPU buffer. */
export type ArrowGPUVectorFromArrowProps<T extends AttributeArrowType = AttributeArrowType> = {
  /** Discriminator for Arrow-vector upload construction. */
  type: 'arrow';
  /** Name used when this vector is added to an {@link ArrowGPUTable}. */
  name: string;
  /** Device that creates the GPU buffer. */
  device: Device;
  /** Arrow vector whose value memory is uploaded. */
  vector: arrow.Vector<T>;
  /** Buffer creation props forwarded to the GPU buffer. */
  bufferProps?: ArrowGPUVectorBufferProps;
};

/** Constructor props that wrap an existing typed GPU buffer. */
export type ArrowGPUVectorFromBufferProps<T extends AttributeArrowType = AttributeArrowType> = {
  /** Discriminator for existing-buffer construction. */
  type: 'buffer';
  /** Name used when this vector is added to an {@link ArrowGPUTable}. */
  name: string;
  /** Existing GPU buffer. */
  buffer: Buffer;
  /** Arrow type that describes the values in the buffer. */
  arrowType: T;
  /** Number of logical rows in the buffer. */
  length: number;
  /** Byte offset of the first logical row. */
  byteOffset?: number;
  /** Bytes between adjacent logical rows. Defaults to the byte width of `arrowType`. */
  byteStride?: number;
  /**
   * Whether this vector should destroy the buffer.
   *
   * Defaults to `false` for wrapped buffers because ownership remains with the caller unless
   * explicitly transferred or opted in.
   */
  ownsBuffer?: boolean;
};

/** Constructor props that wrap one interleaved GPU buffer as opaque Arrow binary rows. */
export type ArrowGPUVectorFromInterleavedProps = {
  /** Discriminator for interleaved-buffer construction. */
  type: 'interleaved';
  /** Name used when this vector is added to an {@link ArrowGPUTable}. */
  name: string;
  /** Existing interleaved GPU buffer. */
  buffer: Buffer;
  /** Number of logical rows in the buffer. */
  length: number;
  /** Byte offset of the first logical row. */
  byteOffset?: number;
  /** Bytes between adjacent logical rows. */
  byteStride: number;
  /** Attribute views stored in each interleaved row. */
  attributes: BufferAttributeLayout[];
  /**
   * Whether this vector should destroy the buffer.
   *
   * Defaults to `false` for wrapped buffers because ownership remains with the caller unless
   * explicitly transferred or opted in.
   */
  ownsBuffer?: boolean;
};

/** Discriminated constructor props for {@link ArrowGPUVector}. */
export type ArrowGPUVectorCreateProps<T extends arrow.DataType = AttributeArrowType> =
  | (T extends AttributeArrowType
      ? ArrowGPUVectorFromArrowProps<T> | ArrowGPUVectorFromBufferProps<T>
      : never)
  | ArrowGPUVectorFromInterleavedProps;

type ArrowGPUVectorReadableBuffer = Pick<Buffer, 'readAsync'>;

type ArrowGPUVectorReadProps<T extends AttributeArrowType> = {
  type: T;
  buffer: ArrowGPUVectorReadableBuffer;
  length: number;
  byteOffset: number;
  byteStride: number;
};

type NumericTypedArrayConstructor = {
  readonly BYTES_PER_ELEMENT: number;
  new (buffer: ArrayBufferLike, byteOffset?: number, length?: number): BigTypedArray;
};

const makeNumericData = arrow.makeData as <T extends NumericArrowType>(props: {
  type: T;
  length: number;
  data: T['TArray'];
}) => arrow.Data<T>;

const makeFixedSizeListData = arrow.makeData as <T extends NumericArrowType>(props: {
  type: arrow.FixedSizeList<T>;
  length: number;
  nullCount: number;
  nullBitmap: null;
  child: arrow.Data<T>;
}) => arrow.Data<arrow.FixedSizeList<T>>;

/** Read GPU bytes and reconstruct one non-null Arrow vector with the supplied Arrow type. */
export async function readArrowGPUVectorAsync<T extends AttributeArrowType>(
  props: ArrowGPUVectorReadProps<T>
): Promise<arrow.Vector<T>> {
  const {buffer, type, length, byteOffset, byteStride} = props;
  const rowByteWidth = getArrowTypeByteStride(type);
  if (byteStride < rowByteWidth) {
    throw new Error(
      `ArrowGPUVector.readAsync() byteStride ${byteStride} is smaller than row byte width ${rowByteWidth}`
    );
  }

  const readByteLength = length === 0 ? 0 : (length - 1) * byteStride + rowByteWidth;
  const bytes =
    readByteLength === 0 ? new Uint8Array(0) : await buffer.readAsync(byteOffset, readByteLength);
  const packedBytes =
    byteStride === rowByteWidth
      ? bytes
      : compactStridedRows(bytes, length, byteStride, rowByteWidth);

  return makeArrowVectorFromPackedBytes(type, length, packedBytes);
}

/**
 * GPU memory and Arrow type metadata derived from one Arrow vector.
 *
 * The Arrow vector is a construction input only. ArrowGPUVector does not retain
 * the source vector; it keeps a GPU buffer plus the type, length, and stride that
 * describe the uploaded memory.
 *
 * Ownership is tracked separately from the buffer reference. Vectors constructed
 * from Arrow data allocate and own their buffers. Vectors wrapping an existing
 * buffer default to non-owning unless `ownsBuffer` is supplied. In-place
 * operations can use {@link transferBufferOwnership} to consume one logical
 * vector and return another view that becomes responsible for destroying the
 * shared buffer.
 */
export class ArrowGPUVector<T extends arrow.DataType = AttributeArrowType> {
  /** Name used when this vector is added to an {@link ArrowGPUTable}. */
  readonly name: string;
  /** GPU buffer containing the Arrow vector's attribute-compatible value memory. */
  readonly buffer: Buffer;
  /** Arrow type that describes the uploaded vector memory. */
  readonly type: T;
  /** Number of logical Arrow vector rows uploaded into the GPU buffer. */
  readonly length: number;
  /** Number of scalar values per logical vector row. */
  readonly stride: number;
  /** Byte offset of the first logical row in {@link buffer}. */
  readonly byteOffset: number;
  /** Bytes between adjacent logical rows in {@link buffer}. */
  readonly byteStride: number;
  /** Optional GPU buffer layout described by this vector. */
  readonly bufferLayout?: BufferLayout;
  /** Whether this vector is responsible for destroying {@link buffer}. */
  private _ownsBuffer: boolean;

  /** Creates a GPU representation from an Arrow vector without retaining the source vector. */
  constructor(
    device: Device,
    vector: arrow.Vector<T & AttributeArrowType>,
    props?: ArrowGPUVectorBufferProps
  );
  /** Creates a GPU representation using discriminated construction props. */
  constructor(props: ArrowGPUVectorCreateProps<T>);
  constructor(
    deviceOrProps: Device | ArrowGPUVectorCreateProps<any>,
    vector?: arrow.Vector<T & AttributeArrowType>,
    props: ArrowGPUVectorBufferProps = {}
  ) {
    const constructionProps =
      deviceOrProps instanceof Device
        ? ({
            type: 'arrow',
            name: 'vector',
            device: deviceOrProps,
            vector: vector!,
            bufferProps: props
          } satisfies ArrowGPUVectorFromArrowProps)
        : deviceOrProps;

    switch (constructionProps.type) {
      case 'arrow': {
        const {name, device, vector: arrowVector, bufferProps = {}} = constructionProps;
        this.name = name;
        this.type = arrowVector.type;
        this.length = arrowVector.length;
        this.stride = getArrowVectorStride(arrowVector);
        this.byteOffset = 0;
        this.byteStride = getArrowTypeByteStride(arrowVector.type);
        this.buffer = device.createBuffer({
          usage: Buffer.VERTEX | Buffer.STORAGE | Buffer.COPY_DST | Buffer.COPY_SRC,
          ...bufferProps,
          data: getArrowVectorBufferSource(arrowVector as any)
        });
        this._ownsBuffer = true;
        return;
      }

      case 'buffer': {
        const {
          name,
          buffer,
          arrowType,
          length,
          byteOffset = 0,
          byteStride = getArrowTypeByteStride(arrowType),
          ownsBuffer = false
        } = constructionProps;
        this.name = name;
        this.buffer = buffer;
        this.type = arrowType;
        this.length = length;
        this.stride = getArrowTypeStride(arrowType);
        this.byteOffset = byteOffset;
        this.byteStride = byteStride;
        this._ownsBuffer = ownsBuffer;
        return;
      }

      case 'interleaved': {
        const {
          name,
          buffer,
          length,
          byteOffset = 0,
          byteStride,
          attributes,
          ownsBuffer = false
        } = constructionProps;
        this.name = name;
        this.buffer = buffer;
        this.type = new arrow.Binary() as T;
        this.length = length;
        this.stride = byteStride;
        this.byteOffset = byteOffset;
        this.byteStride = byteStride;
        this.bufferLayout = {name, byteStride, attributes};
        this._ownsBuffer = ownsBuffer;
        return;
      }
    }
  }

  /**
   * Whether this vector is responsible for destroying {@link buffer}.
   *
   * `destroy()` only releases the buffer when this is `true`. This value can
   * change when ownership is transferred to another same-buffer view.
   */
  get ownsBuffer(): boolean {
    return this._ownsBuffer;
  }

  /**
   * Transfers buffer ownership to another vector that views the same GPU buffer.
   *
   * This is intended for in-place operations that consume one logical vector and
   * return a new logical interpretation of the same bytes. After transfer,
   * destroying this vector will not destroy the buffer; destroying `target` will
   * destroy it if this vector previously owned it.
   */
  transferBufferOwnership(target: ArrowGPUVector): void {
    if (target.buffer !== this.buffer) {
      throw new Error('ArrowGPUVector ownership can only be transferred to the same buffer');
    }
    target._ownsBuffer = this._ownsBuffer;
    this._ownsBuffer = false;
  }

  /** Reads the GPU buffer contents back into a single non-null Arrow vector. */
  async readAsync(): Promise<arrow.Vector<T>> {
    if (this.bufferLayout) {
      throw new Error('ArrowGPUVector.readAsync() does not support interleaved vectors');
    }

    return readArrowGPUVectorAsync({
      type: this.type as unknown as AttributeArrowType,
      buffer: this.buffer,
      length: this.length,
      byteOffset: this.byteOffset,
      byteStride: this.byteStride
    }) as unknown as Promise<arrow.Vector<T>>;
  }

  destroy(): void {
    if (this._ownsBuffer) {
      this.buffer.destroy();
      this._ownsBuffer = false;
    }
  }
}

function getArrowVectorStride(vector: arrow.Vector<AttributeArrowType>): number {
  return arrow.DataType.isFixedSizeList(vector.type) ? vector.type.listSize : 1;
}

function getArrowTypeStride(type: arrow.DataType): number {
  return arrow.DataType.isFixedSizeList(type) ? type.listSize : 1;
}

function getArrowTypeByteStride(type: arrow.DataType): number {
  if (arrow.DataType.isFixedSizeList(type)) {
    return type.listSize * getArrowTypeByteStride(type.children[0].type);
  }
  if (arrow.DataType.isInt(type)) {
    return type.bitWidth / 8;
  }
  if (arrow.DataType.isFloat(type)) {
    switch (type.precision) {
      case arrow.Precision.HALF:
        return 2;
      case arrow.Precision.SINGLE:
        return 4;
      case arrow.Precision.DOUBLE:
        return 8;
    }
  }
  throw new Error(`Cannot determine byte stride for Arrow type ${type}`);
}

function compactStridedRows(
  bytes: Uint8Array,
  length: number,
  byteStride: number,
  rowByteWidth: number
): Uint8Array {
  const packedBytes = new Uint8Array(length * rowByteWidth);
  for (let rowIndex = 0; rowIndex < length; rowIndex++) {
    const sourceOffset = rowIndex * byteStride;
    const targetOffset = rowIndex * rowByteWidth;
    packedBytes.set(bytes.subarray(sourceOffset, sourceOffset + rowByteWidth), targetOffset);
  }
  return packedBytes;
}

function makeArrowVectorFromPackedBytes<T extends AttributeArrowType>(
  type: T,
  length: number,
  bytes: Uint8Array
): arrow.Vector<T> {
  if (arrow.DataType.isFixedSizeList(type)) {
    const childType = type.children[0].type as NumericArrowType;
    const values = makeNumericTypedArray(childType, bytes, length * type.listSize);
    const childData = makeNumericData({
      type: childType,
      length: length * type.listSize,
      data: values as NumericArrowType['TArray']
    });
    const listData = makeFixedSizeListData({
      type: type as arrow.FixedSizeList<NumericArrowType>,
      length,
      nullCount: 0,
      nullBitmap: null,
      child: childData
    });
    return arrow.makeVector(listData) as arrow.Vector<T>;
  }

  const numericType = type as NumericArrowType;
  const values = makeNumericTypedArray(numericType, bytes, length);
  const data = makeNumericData({
    type: numericType,
    length,
    data: values as NumericArrowType['TArray']
  });
  return arrow.makeVector(data) as arrow.Vector<T>;
}

function makeNumericTypedArray(
  type: NumericArrowType,
  bytes: Uint8Array,
  length: number
): BigTypedArray {
  if (arrow.DataType.isInt(type)) {
    if (type.isSigned) {
      switch (type.bitWidth) {
        case 8:
          return makeTypedArrayView(Int8Array, bytes, length);
        case 16:
          return makeTypedArrayView(Int16Array, bytes, length);
        case 32:
          return makeTypedArrayView(Int32Array, bytes, length);
        case 64:
          return makeTypedArrayView(BigInt64Array, bytes, length);
      }
    }

    switch (type.bitWidth) {
      case 8:
        return makeTypedArrayView(Uint8Array, bytes, length);
      case 16:
        return makeTypedArrayView(Uint16Array, bytes, length);
      case 32:
        return makeTypedArrayView(Uint32Array, bytes, length);
      case 64:
        return makeTypedArrayView(BigUint64Array, bytes, length);
    }
  }

  if (arrow.DataType.isFloat(type)) {
    switch (type.precision) {
      case arrow.Precision.HALF:
        return makeTypedArrayView(Uint16Array, bytes, length);
      case arrow.Precision.SINGLE:
        return makeTypedArrayView(Float32Array, bytes, length);
      case arrow.Precision.DOUBLE:
        return makeTypedArrayView(Float64Array, bytes, length);
    }
  }

  throw new Error(`ArrowGPUVector.readAsync() does not support Arrow type ${type}`);
}

function makeTypedArrayView<T extends BigTypedArray>(
  TypedArrayConstructor: NumericTypedArrayConstructor,
  bytes: Uint8Array,
  length: number
): T {
  const byteLength = length * TypedArrayConstructor.BYTES_PER_ELEMENT;
  if (bytes.byteOffset % TypedArrayConstructor.BYTES_PER_ELEMENT === 0) {
    return new TypedArrayConstructor(bytes.buffer, bytes.byteOffset, length) as T;
  }

  const alignedBytes = new Uint8Array(byteLength);
  alignedBytes.set(bytes.subarray(0, byteLength));
  return new TypedArrayConstructor(alignedBytes.buffer, 0, length) as T;
}

1	// luma.gl
2	// SPDX-License-Identifier: MIT
3	// Copyright (c) vis.gl contributors
4
5	import {
6	Buffer,
7	Device,
8	type BufferAttributeLayout,
9	type BufferLayout,
10	type BufferProps,
11	type BigTypedArray
12	} from '@luma.gl/core';
13	import * as arrow from 'apache-arrow';
14	import type {AttributeArrowType, NumericArrowType} from './arrow-types';
15	import {getArrowVectorBufferSource} from './arrow-fixed-size-list';
16
17	/** Buffer creation props forwarded when uploading Arrow vector memory to the GPU. */
18	export type ArrowGPUVectorBufferProps = Omit<BufferProps, 'byteLength' \| 'data'>;
19
20	/** @deprecated Use {@link ArrowGPUVectorBufferProps}. */
21	export type ArrowGPUVectorProps = ArrowGPUVectorBufferProps;
22
23	/** Constructor props that upload an Arrow vector into a new GPU buffer. */
24	export type ArrowGPUVectorFromArrowProps<T extends AttributeArrowType = AttributeArrowType> = {
25	/** Discriminator for Arrow-vector upload construction. */
26	type: 'arrow';
27	/** Name used when this vector is added to an {@link ArrowGPUTable}. */
28	name: string;
29	/** Device that creates the GPU buffer. */
30	device: Device;
31	/** Arrow vector whose value memory is uploaded. */
32	vector: arrow.Vector<T>;
33	/** Buffer creation props forwarded to the GPU buffer. */
34	bufferProps?: ArrowGPUVectorBufferProps;
35	};
36
37	/** Constructor props that wrap an existing typed GPU buffer. */
38	export type ArrowGPUVectorFromBufferProps<T extends AttributeArrowType = AttributeArrowType> = {
39	/** Discriminator for existing-buffer construction. */
40	type: 'buffer';
41	/** Name used when this vector is added to an {@link ArrowGPUTable}. */
42	name: string;
43	/** Existing GPU buffer. */
44	buffer: Buffer;
45	/** Arrow type that describes the values in the buffer. */
46	arrowType: T;
47	/** Number of logical rows in the buffer. */
48	length: number;
49	/** Byte offset of the first logical row. */
50	byteOffset?: number;
51	/** Bytes between adjacent logical rows. Defaults to the byte width of `arrowType`. */
52	byteStride?: number;
53	/**
54	* Whether this vector should destroy the buffer.
55	*
56	* Defaults to `false` for wrapped buffers because ownership remains with the caller unless
57	* explicitly transferred or opted in.
58	*/
59	ownsBuffer?: boolean;
60	};
61
62	/** Constructor props that wrap one interleaved GPU buffer as opaque Arrow binary rows. */
63	export type ArrowGPUVectorFromInterleavedProps = {
64	/** Discriminator for interleaved-buffer construction. */
65	type: 'interleaved';
66	/** Name used when this vector is added to an {@link ArrowGPUTable}. */
67	name: string;
68	/** Existing interleaved GPU buffer. */
69	buffer: Buffer;
70	/** Number of logical rows in the buffer. */
71	length: number;
72	/** Byte offset of the first logical row. */
73	byteOffset?: number;
74	/** Bytes between adjacent logical rows. */
75	byteStride: number;
76	/** Attribute views stored in each interleaved row. */
77	attributes: BufferAttributeLayout[];
78	/**
79	* Whether this vector should destroy the buffer.
80	*
81	* Defaults to `false` for wrapped buffers because ownership remains with the caller unless
82	* explicitly transferred or opted in.
83	*/
84	ownsBuffer?: boolean;
85	};
86
87	/** Discriminated constructor props for {@link ArrowGPUVector}. */
88	export type ArrowGPUVectorCreateProps<T extends arrow.DataType = AttributeArrowType> =
89	\| (T extends AttributeArrowType
90	? ArrowGPUVectorFromArrowProps<T> \| ArrowGPUVectorFromBufferProps<T>
91	: never)
92	\| ArrowGPUVectorFromInterleavedProps;
93
94	type ArrowGPUVectorReadableBuffer = Pick<Buffer, 'readAsync'>;
95
96	type ArrowGPUVectorReadProps<T extends AttributeArrowType> = {
97	type: T;
98	buffer: ArrowGPUVectorReadableBuffer;
99	length: number;
100	byteOffset: number;
101	byteStride: number;
102	};
103
104	type NumericTypedArrayConstructor = {
105	readonly BYTES_PER_ELEMENT: number;
106	new (buffer: ArrayBufferLike, byteOffset?: number, length?: number): BigTypedArray;
107	};
108
109	const makeNumericData = arrow.makeData as <T extends NumericArrowType>(props: {	9✔
110	type: T;
111	length: number;
112	data: T['TArray'];
113	}) => arrow.Data<T>;
114
115	const makeFixedSizeListData = arrow.makeData as <T extends NumericArrowType>(props: {	9✔
116	type: arrow.FixedSizeList<T>;
117	length: number;
118	nullCount: number;
119	nullBitmap: null;
120	child: arrow.Data<T>;
121	}) => arrow.Data<arrow.FixedSizeList<T>>;
122
123	/** Read GPU bytes and reconstruct one non-null Arrow vector with the supplied Arrow type. */
124	export async function readArrowGPUVectorAsync<T extends AttributeArrowType>(
125	props: ArrowGPUVectorReadProps<T>
126	): Promise<arrow.Vector<T>> {
127	const {buffer, type, length, byteOffset, byteStride} = props;	5✔
128	const rowByteWidth = getArrowTypeByteStride(type);	5✔
129	if (byteStride < rowByteWidth) {	5!
NEW 130	throw new Error(	×
131	`ArrowGPUVector.readAsync() byteStride ${byteStride} is smaller than row byte width ${rowByteWidth}`
132	);
133	}
134
135	const readByteLength = length === 0 ? 0 : (length - 1) * byteStride + rowByteWidth;	5✔
136	const bytes =
137	readByteLength === 0 ? new Uint8Array(0) : await buffer.readAsync(byteOffset, readByteLength);	5✔
138	const packedBytes =
139	byteStride === rowByteWidth	4✔
140	? bytes
141	: compactStridedRows(bytes, length, byteStride, rowByteWidth);
142
143	return makeArrowVectorFromPackedBytes(type, length, packedBytes);	5✔
144	}
145
146	/**
147	* GPU memory and Arrow type metadata derived from one Arrow vector.
148	*
149	* The Arrow vector is a construction input only. ArrowGPUVector does not retain
150	* the source vector; it keeps a GPU buffer plus the type, length, and stride that
151	* describe the uploaded memory.
152	*
153	* Ownership is tracked separately from the buffer reference. Vectors constructed
154	* from Arrow data allocate and own their buffers. Vectors wrapping an existing
155	* buffer default to non-owning unless `ownsBuffer` is supplied. In-place
156	* operations can use {@link transferBufferOwnership} to consume one logical
157	* vector and return another view that becomes responsible for destroying the
158	* shared buffer.
159	*/
160	export class ArrowGPUVector<T extends arrow.DataType = AttributeArrowType> {
161	/** Name used when this vector is added to an {@link ArrowGPUTable}. */
162	readonly name: string;
163	/** GPU buffer containing the Arrow vector's attribute-compatible value memory. */
164	readonly buffer: Buffer;
165	/** Arrow type that describes the uploaded vector memory. */
166	readonly type: T;
167	/** Number of logical Arrow vector rows uploaded into the GPU buffer. */
168	readonly length: number;
169	/** Number of scalar values per logical vector row. */
170	readonly stride: number;
171	/** Byte offset of the first logical row in {@link buffer}. */
172	readonly byteOffset: number;
173	/** Bytes between adjacent logical rows in {@link buffer}. */
174	readonly byteStride: number;
175	/** Optional GPU buffer layout described by this vector. */
176	readonly bufferLayout?: BufferLayout;
177	/** Whether this vector is responsible for destroying {@link buffer}. */
178	private _ownsBuffer: boolean;
179
180	/** Creates a GPU representation from an Arrow vector without retaining the source vector. */
181	constructor(
182	device: Device,
183	vector: arrow.Vector<T & AttributeArrowType>,
184	props?: ArrowGPUVectorBufferProps
185	);
186	/** Creates a GPU representation using discriminated construction props. */
187	constructor(props: ArrowGPUVectorCreateProps<T>);
188	constructor(
189	deviceOrProps: Device \| ArrowGPUVectorCreateProps<any>,
190	vector?: arrow.Vector<T & AttributeArrowType>,
191	props: ArrowGPUVectorBufferProps = {}	31✔
192	) {
193	const constructionProps =
194	deviceOrProps instanceof Device	31✔
195	? ({
196	type: 'arrow',
197	name: 'vector',
198	device: deviceOrProps,
199	vector: vector!,
200	bufferProps: props
201	} satisfies ArrowGPUVectorFromArrowProps)
202	: deviceOrProps;
203
204	switch (constructionProps.type) {	31✔
205	case 'arrow': {
206	const {name, device, vector: arrowVector, bufferProps = {}} = constructionProps;	22✔
207	this.name = name;	22✔
208	this.type = arrowVector.type;	22✔
209	this.length = arrowVector.length;	22✔
210	this.stride = getArrowVectorStride(arrowVector);	22✔
211	this.byteOffset = 0;	22✔
212	this.byteStride = getArrowTypeByteStride(arrowVector.type);	22✔
213	this.buffer = device.createBuffer({	22✔
214	usage: Buffer.VERTEX \| Buffer.STORAGE \| Buffer.COPY_DST \| Buffer.COPY_SRC,
215	...bufferProps,
216	data: getArrowVectorBufferSource(arrowVector as any)
217	});
218	this._ownsBuffer = true;	22✔
219	return;	22✔
220	}
221
222	case 'buffer': {
223	const {
224	name,
225	buffer,
226	arrowType,
227	length,
228	byteOffset = 0,	6✔
229	byteStride = getArrowTypeByteStride(arrowType),	6✔
230	ownsBuffer = false	6✔
231	} = constructionProps;	6✔
232	this.name = name;	6✔
233	this.buffer = buffer;	6✔
234	this.type = arrowType;	6✔
235	this.length = length;	6✔
236	this.stride = getArrowTypeStride(arrowType);	6✔
237	this.byteOffset = byteOffset;	6✔
238	this.byteStride = byteStride;	6✔
239	this._ownsBuffer = ownsBuffer;	6✔
240	return;	6✔
241	}
242
243	case 'interleaved': {
244	const {
245	name,
246	buffer,
247	length,
248	byteOffset = 0,	3✔
249	byteStride,
250	attributes,
251	ownsBuffer = false	3✔
252	} = constructionProps;	3✔
253	this.name = name;	3✔
254	this.buffer = buffer;	3✔
255	this.type = new arrow.Binary() as T;	3✔
256	this.length = length;	3✔
257	this.stride = byteStride;	3✔
258	this.byteOffset = byteOffset;	3✔
259	this.byteStride = byteStride;	3✔
260	this.bufferLayout = {name, byteStride, attributes};	3✔
261	this._ownsBuffer = ownsBuffer;	3✔
262	return;	3✔
263	}
264	}
265	}
266
267	/**
268	* Whether this vector is responsible for destroying {@link buffer}.
269	*
270	* `destroy()` only releases the buffer when this is `true`. This value can
271	* change when ownership is transferred to another same-buffer view.
272	*/
273	get ownsBuffer(): boolean {
274	return this._ownsBuffer;	3✔
275	}
276
277	/**
278	* Transfers buffer ownership to another vector that views the same GPU buffer.
279	*
280	* This is intended for in-place operations that consume one logical vector and
281	* return a new logical interpretation of the same bytes. After transfer,
282	* destroying this vector will not destroy the buffer; destroying `target` will
283	* destroy it if this vector previously owned it.
284	*/
285	transferBufferOwnership(target: ArrowGPUVector): void {
286	if (target.buffer !== this.buffer) {	1!
287	throw new Error('ArrowGPUVector ownership can only be transferred to the same buffer');	×
288	}
289	target._ownsBuffer = this._ownsBuffer;	1✔
290	this._ownsBuffer = false;	1✔
291	}
292
293	/** Reads the GPU buffer contents back into a single non-null Arrow vector. */
294	async readAsync(): Promise<arrow.Vector<T>> {
295	if (this.bufferLayout) {	4✔
296	throw new Error('ArrowGPUVector.readAsync() does not support interleaved vectors');	1✔
297	}
298
299	return readArrowGPUVectorAsync({	3✔
300	type: this.type as unknown as AttributeArrowType,
301	buffer: this.buffer,
302	length: this.length,
303	byteOffset: this.byteOffset,
304	byteStride: this.byteStride
305	}) as unknown as Promise<arrow.Vector<T>>;
306	}
307
308	destroy(): void {
309	if (this._ownsBuffer) {	31✔
310	this.buffer.destroy();	28✔
311	this._ownsBuffer = false;	28✔
312	}
313	}
314	}
315
316	function getArrowVectorStride(vector: arrow.Vector<AttributeArrowType>): number {
317	return arrow.DataType.isFixedSizeList(vector.type) ? vector.type.listSize : 1;	22✔
318	}
319
320	function getArrowTypeStride(type: arrow.DataType): number {
321	return arrow.DataType.isFixedSizeList(type) ? type.listSize : 1;	6✔
322	}
323
324	function getArrowTypeByteStride(type: arrow.DataType): number {
325	if (arrow.DataType.isFixedSizeList(type)) {	54✔
326	return type.listSize * getArrowTypeByteStride(type.children[0].type);	22✔
327	}
328	if (arrow.DataType.isInt(type)) {	32✔
329	return type.bitWidth / 8;	12✔
330	}
331	if (arrow.DataType.isFloat(type)) {	20!
332	switch (type.precision) {	20!
333	case arrow.Precision.HALF:
334	return 2;	×
335	case arrow.Precision.SINGLE:
336	return 4;	20✔
337	case arrow.Precision.DOUBLE:
338	return 8;	×
339	}
340	}
341	throw new Error(`Cannot determine byte stride for Arrow type ${type}`);	×
342	}
343
344	function compactStridedRows(
345	bytes: Uint8Array,
346	length: number,
347	byteStride: number,
348	rowByteWidth: number
349	): Uint8Array {
350	const packedBytes = new Uint8Array(length * rowByteWidth);	1✔
351	for (let rowIndex = 0; rowIndex < length; rowIndex++) {	1✔
352	const sourceOffset = rowIndex * byteStride;	2✔
353	const targetOffset = rowIndex * rowByteWidth;	2✔
354	packedBytes.set(bytes.subarray(sourceOffset, sourceOffset + rowByteWidth), targetOffset);	2✔
355	}
356	return packedBytes;	1✔
357	}
358
359	function makeArrowVectorFromPackedBytes<T extends AttributeArrowType>(
360	type: T,
361	length: number,
362	bytes: Uint8Array
363	): arrow.Vector<T> {
364	if (arrow.DataType.isFixedSizeList(type)) {	5✔
365	const childType = type.children[0].type as NumericArrowType;	1✔
366	const values = makeNumericTypedArray(childType, bytes, length * type.listSize);	1✔
367	const childData = makeNumericData({	1✔
368	type: childType,
369	length: length * type.listSize,
370	data: values as NumericArrowType['TArray']
371	});
372	const listData = makeFixedSizeListData({	1✔
373	type: type as arrow.FixedSizeList<NumericArrowType>,
374	length,
375	nullCount: 0,
376	nullBitmap: null,
377	child: childData
378	});
379	return arrow.makeVector(listData) as arrow.Vector<T>;	1✔
380	}
381
382	const numericType = type as NumericArrowType;	4✔
383	const values = makeNumericTypedArray(numericType, bytes, length);	4✔
384	const data = makeNumericData({	4✔
385	type: numericType,
386	length,
387	data: values as NumericArrowType['TArray']
388	});
389	return arrow.makeVector(data) as arrow.Vector<T>;	4✔
390	}
391
392	function makeNumericTypedArray(
393	type: NumericArrowType,
394	bytes: Uint8Array,
395	length: number
396	): BigTypedArray {
397	if (arrow.DataType.isInt(type)) {	5✔
398	if (type.isSigned) {	1!
399	switch (type.bitWidth) {	1!
400	case 8:
NEW 401	return makeTypedArrayView(Int8Array, bytes, length);	×
402	case 16:
NEW 403	return makeTypedArrayView(Int16Array, bytes, length);	×
404	case 32:
405	return makeTypedArrayView(Int32Array, bytes, length);	1✔
406	case 64:
NEW 407	return makeTypedArrayView(BigInt64Array, bytes, length);	×
408	}
409	}
410
NEW 411	switch (type.bitWidth) {	×
412	case 8:
NEW 413	return makeTypedArrayView(Uint8Array, bytes, length);	×
414	case 16:
NEW 415	return makeTypedArrayView(Uint16Array, bytes, length);	×
416	case 32:
NEW 417	return makeTypedArrayView(Uint32Array, bytes, length);	×
418	case 64:
NEW 419	return makeTypedArrayView(BigUint64Array, bytes, length);	×
420	}
421	}
422
423	if (arrow.DataType.isFloat(type)) {	4!
424	switch (type.precision) {	4!
425	case arrow.Precision.HALF:
NEW 426	return makeTypedArrayView(Uint16Array, bytes, length);	×
427	case arrow.Precision.SINGLE:
428	return makeTypedArrayView(Float32Array, bytes, length);	4✔
429	case arrow.Precision.DOUBLE:
NEW 430	return makeTypedArrayView(Float64Array, bytes, length);	×
431	}
432	}
433
NEW 434	throw new Error(`ArrowGPUVector.readAsync() does not support Arrow type ${type}`);	×
435	}
436
437	function makeTypedArrayView<T extends BigTypedArray>(
438	TypedArrayConstructor: NumericTypedArrayConstructor,
439	bytes: Uint8Array,
440	length: number
441	): T {
442	const byteLength = length * TypedArrayConstructor.BYTES_PER_ELEMENT;	5✔
443	if (bytes.byteOffset % TypedArrayConstructor.BYTES_PER_ELEMENT === 0) {	5!
444	return new TypedArrayConstructor(bytes.buffer, bytes.byteOffset, length) as T;	5✔
445	}
446
NEW 447	const alignedBytes = new Uint8Array(byteLength);	×
NEW 448	alignedBytes.set(bytes.subarray(0, byteLength));	×
NEW 449	return new TypedArrayConstructor(alignedBytes.buffer, 0, length) as T;	×
450	}

visgl / luma.gl / 25756190722

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous