feat(arrow): integrate color conversion

Author: ibgreen

modules/arrow-layers/src/layers/arrow-layer-input.ts

@@ -3,10 +3,16 @@
 // Copyright (c) vis.gl contributors
 
 import {
+  canConvertColors,
+  convertArrowColors,
+  convertColors,
   getArrowPaths,
   getArrowRecordBatchAsyncIterator,
+  readArrowGPUVectorAsync,
+  type ArrowColorType,
   type ArrowRecordBatchSource
 } from '@luma.gl/arrow';
+import type {Device} from '@luma.gl/core';
 import {GPUVector, type GPUVectorFormat} from '@luma.gl/tables';
 import {RecordBatch, Table, type DataType, type Vector} from 'apache-arrow';
 
@@ -56,6 +62,59 @@ export function assertArrowLayerGPUVector(
   }
 }
 
+/** Validates a caller-owned color vector that is normalized already or convertible by luma.gl. */
+export function assertArrowLayerColorGPUVector(
+  ownerName: string,
+  vector: GPUVector,
+  expectedLength?: number
+): void {
+  if (vector.format === 'vertex-list<unorm8x4>') {
+    assertArrowLayerGPUVector(ownerName, vector, ['vertex-list<unorm8x4>'], expectedLength);
+    return;
+  }
+  if (!canConvertColors(vector)) {
+    throw new Error(
+      `${ownerName} GPUVector.format "${vector.format ?? 'undefined'}" must be a convertible fixed-width RGB/RGBA format or vertex-list<unorm8x4>`
+    );
+  }
+  assertArrowLayerGPUVector(ownerName, vector, [vector.format!], expectedLength);
+}
+
+/** Returns normalized RGBA8 colors, preserving caller ownership when conversion is unnecessary. */
+export async function convertArrowLayerColorGPUVector(
+  device: Device,
+  vector: GPUVector,
+  name: string
+): Promise<{
+  vector: GPUVector<'unorm8x4' | 'vertex-list<unorm8x4>'>;
+  converted: boolean;
+}> {
+  if (vector.format === 'unorm8x4' || vector.format === 'vertex-list<unorm8x4>') {
+    return {
+      vector: vector as GPUVector<'unorm8x4' | 'vertex-list<unorm8x4>'>,
+      converted: false
+    };
+  }
+  return {vector: await convertColors(device, vector, {name}), converted: true};
+}
+
+/** Normalizes one supported fixed-width Arrow RGB/RGBA vector to Uint8 RGBA rows. */
+export async function convertArrowLayerColorVector(
+  device: Device,
+  vector: Vector,
+  name: string
+): Promise<Vector> {
+  if (!canConvertColors(vector)) {
+    throw new Error(`Arrow color vector type ${vector.type} is not convertible to Uint8 RGBA`);
+  }
+  const converted = await convertArrowColors(device, vector as Vector<ArrowColorType>, {name});
+  try {
+    return await readArrowGPUVectorAsync(converted);
+  } finally {
+    converted.destroy();
+  }
+}
+
 function hasNullRows(nullBitmap: Uint8Array, length: number): boolean {
   for (let rowIndex = 0; rowIndex < length; rowIndex++) {
     if ((nullBitmap[rowIndex >> 3] & (1 << (rowIndex & 7))) === 0) {

modules/arrow-layers/src/layers/arrow-path-layer.ts

@@ -17,6 +17,7 @@ import {
   getArrowRecordBatchAsyncIterator,
   readArrowGPUVectorAsync,
   resolveArrowPathSourceVectors,
+  type ArrowColorType,
   type ArrowPathSourceVectors,
   type ArrowPathSourceVectorSelectors,
   type ArrowRecordBatchSource,
@@ -35,7 +36,7 @@ import {
   GPUVector,
   PathAttributeModel,
   PathStorageModel,
-  type VertexList
+  type GPUVectorFormat
 } from '@luma.gl/tables';
 import {
   DataType,
@@ -57,7 +58,10 @@ import {
   type ArrowLayerPickingInfo
 } from './arrow-layer-types';
 import {
+  assertArrowLayerColorGPUVector,
   assertArrowLayerGPUVector,
+  convertArrowLayerColorGPUVector,
+  convertArrowLayerColorVector,
   getArrowLayerInputNullValue,
   getArrowLayerInputSource,
   hasArrowLayerColumn,
@@ -69,11 +73,10 @@ import {
 } from './arrow-layer-input';
 
 type ArrowPathColor = [number, number, number, number];
-type ArrowPathRowColorType = FixedSizeList<Uint8>;
-type ArrowPathVertexColorType = List<ArrowPathRowColorType>;
+type ArrowPathRowColorType = ArrowColorType;
+type ArrowPathVertexColorType = List<FixedSizeList<Uint8>>;
 type ArrowPathColorSource = ArrowLayerColumnSource<
-  ArrowPathRowColorType | ArrowPathVertexColorType,
-  'unorm8x4' | VertexList<'unorm8x4'>
+  ArrowPathRowColorType | ArrowPathVertexColorType
 >;
 type ArrowPathWidthSource = ArrowLayerColumnSource<Float32, 'float32'>;
 
@@ -424,6 +427,7 @@ type ArrowPathLayerBatch = {
   rowIndexOffset: number;
   rowCount: number;
   temporalBuffer: Buffer | null;
+  convertedColorVector: GPUVector | null;
 };
 
 /** Deck-facing props for an Arrow-backed path layer. */
@@ -640,10 +644,7 @@ export class ArrowPathLayer extends Layer<ArrowPathLayerProps> {
       widthNullValue
     );
     if (isArrowLayerGPUVector(colorSource)) {
-      assertArrowLayerGPUVector('ArrowPathLayer color', colorSource, [
-        'unorm8x4',
-        'vertex-list<unorm8x4>'
-      ]);
+      assertArrowLayerColorGPUVector('ArrowPathLayer color', colorSource);
     }
     if (isArrowLayerGPUVector(widthSource)) {
       assertArrowLayerGPUVector('ArrowPathLayer width', widthSource, ['float32']);
@@ -656,18 +657,28 @@ export class ArrowPathLayer extends Layer<ArrowPathLayerProps> {
       model === 'storage' && isArrowLayerGPUVector(widthSource)
         ? undefined
         : await this.resolvePathStyleSource(widthSource, rowIndexOffset, recordBatch?.numRows);
-    const sourceVectors = fillNullablePathStyleVectors(
+    let sourceVectors = fillNullablePathStyleVectors(
       resolveArrowPathSourceVectors(model === 'storage' ? PathStorageModel : PathAttributeModel, {
         data: recordBatch,
         selectors: {
           paths: props.paths,
-          colors: colorSelector ?? null,
+          colors: (colorSelector ?? null) as ArrowPathSourceVectorSelectors['colors'],
           widths: widthSelector ?? null
         }
       }),
       colorNullValue,
       widthNullValue
     );
+    if (sourceVectors.colors && !DataType.isList(sourceVectors.colors.type)) {
+      sourceVectors = {
+        ...sourceVectors,
+        colors: (await convertArrowLayerColorVector(
+          this.context.device,
+          sourceVectors.colors,
+          `${props.id}-batch-${batchIndex}-arrow-colors`
+        )) as NonNullable<ArrowPathSourceVectors['colors']>
+      };
+    }
     const prepared = await ArrowPathRenderer.convertToGPUVectors(
       this.context.device,
       sourceVectors,
@@ -682,10 +693,32 @@ export class ArrowPathLayer extends Layer<ArrowPathLayerProps> {
       return;
     }
 
-    const preparedColorColumn =
-      model === 'storage' && isArrowLayerGPUVector(colorSource)
-        ? getPathStorageGPUVectorBatch(colorSource, batchIndex, sourceVectors.paths.length)
-        : prepared.colors;
+    let convertedColorVector: GPUVector | null = null;
+    let preparedColorColumn = prepared.colors;
+    if (model === 'storage' && isArrowLayerGPUVector(colorSource)) {
+      try {
+        const sourceColorBatch = getPathStorageGPUVectorBatch(
+          colorSource,
+          batchIndex,
+          sourceVectors.paths.length
+        );
+        const normalized = await convertArrowLayerColorGPUVector(
+          this.context.device,
+          sourceColorBatch,
+          `${props.id}-batch-${batchIndex}-colors`
+        );
+        preparedColorColumn = normalized.vector;
+        convertedColorVector = normalized.converted ? normalized.vector : null;
+      } catch (error) {
+        prepared.destroy();
+        throw error;
+      }
+      if (!this.isActiveLoad(loadVersion)) {
+        convertedColorVector?.destroy();
+        prepared.destroy();
+        return;
+      }
+    }
     const preparedWidthColumn =
       model === 'storage' && isArrowLayerGPUVector(widthSource)
         ? getPathStorageGPUVectorBatch(widthSource, batchIndex, sourceVectors.paths.length)
@@ -748,12 +781,14 @@ export class ArrowPathLayer extends Layer<ArrowPathLayerProps> {
       }
     } catch (error) {
       destroyBuffers(constantStyle.buffers);
+      convertedColorVector?.destroy();
       prepared.destroy();
       throw error;
     }
     if (!this.isActiveLoad(loadVersion)) {
       renderModel.destroy();
       destroyBuffers(constantStyle.buffers);
+      convertedColorVector?.destroy();
       prepared.destroy();
       return;
     }
@@ -766,7 +801,8 @@ export class ArrowPathLayer extends Layer<ArrowPathLayerProps> {
       batchIndex,
       rowIndexOffset,
       rowCount,
-      temporalBuffer: constantStyle.temporalBuffer
+      temporalBuffer: constantStyle.temporalBuffer,
+      convertedColorVector
     };
     const batches = [...this.getLayerState().batches, batch];
     this.setState({batches, loadVersion, sourceInitialized: true});
@@ -807,7 +843,20 @@ export class ArrowPathLayer extends Layer<ArrowPathLayerProps> {
     const cache = this.getLayerState().gpuVectorSourceCache;
     let arrowVectorPromise = cache.get(source);
     if (!arrowVectorPromise) {
-      arrowVectorPromise = readArrowGPUVectorAsync(source);
+      arrowVectorPromise = (async () => {
+        const normalized = await convertArrowLayerColorGPUVector(
+          this.context.device,
+          source,
+          `${this.id}-colors`
+        );
+        try {
+          return await readArrowGPUVectorAsync(normalized.vector);
+        } finally {
+          if (normalized.converted) {
+            normalized.vector.destroy();
+          }
+        }
+      })();
       cache.set(source, arrowVectorPromise);
     }
     const arrowVector = await arrowVectorPromise;
@@ -905,9 +954,11 @@ function hasArrowDataNulls(data: Data): boolean {
   return data.nullCount > 0 || data.children.some(hasArrowDataNulls);
 }
 
-function getPathStorageGPUVectorBatch<
-  Format extends 'unorm8x4' | VertexList<'unorm8x4'> | 'float32'
->(vector: GPUVector<Format>, batchIndex: number, expectedRowCount: number): GPUVector<Format> {
+function getPathStorageGPUVectorBatch<Format extends GPUVectorFormat>(
+  vector: GPUVector<Format>,
+  batchIndex: number,
+  expectedRowCount: number
+): GPUVector<Format> {
   const useWholeVector =
     vector.data.length === 1 && batchIndex === 0 && vector.length === expectedRowCount;
   const data = useWholeVector ? vector.data[0] : vector.data[batchIndex];
@@ -1085,6 +1136,7 @@ function destroyPathBatches(batches: ArrowPathLayerBatch[]): void {
   for (const batch of batches) {
     batch.model.destroy();
     destroyBuffers(batch.constantBuffers);
+    batch.convertedColorVector?.destroy();
     batch.prepared.destroy();
   }
 }

modules/arrow-layers/src/layers/arrow-polygon-layer.ts

@@ -15,6 +15,7 @@ import {
 import {
   ArrowPolygonRenderer,
   readArrowGPUVectorAsync,
+  type ArrowColorType,
   type ArrowPolygonRendererDataBatchUpdate,
   type ArrowPolygonRendererProps
 } from '@luma.gl/arrow';
@@ -31,12 +32,13 @@ import type {ArrowLayerPickingInfo} from './arrow-layer-types';
 import {
   GPUVector,
   POLYGON_STORAGE_SHADER_LAYOUT,
-  type GPURecordBatchSourceInfo,
-  type VertexList
+  type GPURecordBatchSourceInfo
 } from '@luma.gl/tables';
-import type {Vector} from 'apache-arrow';
+import {Vector} from 'apache-arrow';
 import {
-  assertArrowLayerGPUVector,
+  assertArrowLayerColorGPUVector,
+  convertArrowLayerColorGPUVector,
+  convertArrowLayerColorVector,
   getArrowLayerInputNullValue,
   getArrowLayerInputSource,
   inspectArrowLayerColumn,
@@ -48,7 +50,8 @@ import {
 type ArrowPolygonColor = [number, number, number, number];
 type ArrowPolygonColorSource =
   | Exclude<ArrowPolygonRendererProps['colors'], null | undefined>
-  | GPUVector<'unorm8x4' | VertexList<'unorm8x4'>>;
+  | Vector<ArrowColorType>
+  | GPUVector;
 type ArrowPolygonResolvedColorSource = {
   value: Exclude<ArrowPolygonColorSource, GPUVector> | null;
 };
@@ -299,13 +302,12 @@ export class ArrowPolygonLayer extends Layer<ArrowPolygonLayerProps> {
       ...this.getRendererModelProps(props),
       data: props.data,
       polygons: props.polygons,
-      colors:
-        resolvedColorSource?.value ??
+      colors: (resolvedColorSource?.value ??
         (resolvedColorSource
           ? null
           : isArrowLayerGPUVector(colorSource)
             ? null
-            : (colorSource ?? null)),
+            : (colorSource ?? null))) as ArrowPolygonRendererProps['colors'],
       tessellated: props.tessellated,
       color: getArrowLayerInputNullValue(props.color, DEFAULT_POLYGON_COLOR, isArrowLayerColor),
       center: props.center,
@@ -330,6 +332,29 @@ export class ArrowPolygonLayer extends Layer<ArrowPolygonLayerProps> {
     const state = this.getLayerState();
     const colorResolveVersion = state.colorResolveVersion + 1;
     state.colorResolveVersion = colorResolveVersion;
+    if (colorSource instanceof Vector) {
+      renderer.setProps(this.getRendererProps(props, {value: null}));
+      void convertArrowLayerColorVector(this.context.device, colorSource, `${this.id}-colors`)
+        .then(colorVector => {
+          if (
+            this.getRendererOrNull() === renderer &&
+            this.getLayerState().colorResolveVersion === colorResolveVersion
+          ) {
+            renderer.setProps(
+              this.getRendererProps(props, {value: colorVector as Vector as never})
+            );
+            this.watchRendererPipeline(renderer);
+            this.setNeedsRedraw();
+          }
+        })
+        .catch(error => {
+          if (this.getLayerState().colorResolveVersion === colorResolveVersion) {
+            if (props.onDataError) props.onDataError(error);
+            else this.raiseError(toError(error), `resolving Arrow polygon color in ${this}`);
+          }
+        });
+      return;
+    }
     if (!isArrowLayerGPUVector(colorSource)) {
       if (typeof colorSource === 'string') {
         if (!props.data) {
@@ -380,14 +405,22 @@ export class ArrowPolygonLayer extends Layer<ArrowPolygonLayerProps> {
     }
     const polygonRowCount =
       props.polygons && typeof props.polygons !== 'string' ? props.polygons.length : undefined;
-    assertArrowLayerGPUVector(
-      'ArrowPolygonLayer color',
-      colorSource,
-      ['unorm8x4', 'vertex-list<unorm8x4>'],
-      polygonRowCount
-    );
+    assertArrowLayerColorGPUVector('ArrowPolygonLayer color', colorSource, polygonRowCount);
     renderer.setProps({data: null});
-    void readArrowGPUVectorAsync(colorSource)
+    void (async () => {
+      const normalized = await convertArrowLayerColorGPUVector(
+        this.context.device,
+        colorSource,
+        `${this.id}-colors`
+      );
+      try {
+        return await readArrowGPUVectorAsync(normalized.vector);
+      } finally {
+        if (normalized.converted) {
+          normalized.vector.destroy();
+        }
+      }
+    })()
       .then(colorVector => {
         if (
           this.getRendererOrNull() === renderer &&