Skip to main content

perspective_js/
typed_array.rs

1// ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
2// ┃ ██████ ██████ ██████       █      █      █      █      █ █▄  ▀███ █       ┃
3// ┃ ▄▄▄▄▄█ █▄▄▄▄▄ ▄▄▄▄▄█  ▀▀▀▀▀█▀▀▀▀▀ █ ▀▀▀▀▀█ ████████▌▐███ ███▄  ▀█ █ ▀▀▀▀▀ ┃
4// ┃ █▀▀▀▀▀ █▀▀▀▀▀ █▀██▀▀ ▄▄▄▄▄ █ ▄▄▄▄▄█ ▄▄▄▄▄█ ████████▌▐███ █████▄   █ ▄▄▄▄▄ ┃
5// ┃ █      ██████ █  ▀█▄       █ ██████      █      ███▌▐███ ███████▄ █       ┃
6// ┣━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┫
7// ┃ Copyright (c) 2017, the Perspective Authors.                              ┃
8// ┃ ╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌ ┃
9// ┃ This file is part of the Perspective library, distributed under the terms ┃
10// ┃ of the [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0). ┃
11// ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛
12
13use std::io::Cursor;
14
15use arrow_array::cast::AsArray;
16use arrow_array::types::*;
17use arrow_array::{Array as _, ArrowPrimitiveType, PrimitiveArray};
18use arrow_ipc::reader::StreamReader;
19use arrow_schema::{DataType, TimeUnit};
20use js_sys::{Array, Function, JsString, Uint8Array};
21use perspective_client::ViewWindow;
22use ts_rs::TS;
23use wasm_bindgen::JsCast;
24use wasm_bindgen::prelude::*;
25
26#[wasm_bindgen]
27unsafe extern "C" {
28    #[wasm_bindgen(typescript_type = "TypedArrayWindow")]
29    #[derive(Clone)]
30    pub type JsTypedArrayWindow;
31}
32
33/// Options for `with_typed_arrays`, extending `ViewWindow` with
34/// typed-array-specific options.
35#[derive(Default, serde::Deserialize, TS)]
36pub struct TypedArrayWindow {
37    #[serde(flatten)]
38    pub view_window: ViewWindow,
39
40    /// When `true`, Float64/Date32/Timestamp columns are output as
41    /// `Float32Array` instead of `Float64Array`.
42    #[serde(default)]
43    pub float32: bool,
44}
45
46impl From<TypedArrayWindow> for ViewWindow {
47    fn from(w: TypedArrayWindow) -> Self {
48        w.view_window
49    }
50}
51
52fn zero_invalid_slots<T: ArrowPrimitiveType>(arr: &PrimitiveArray<T>) {
53    let Some(nulls) = arr.nulls() else { return };
54    let ptr = arr.values().as_ptr() as *mut T::Native;
55    let chunks = nulls.inner().bit_chunks();
56    let mut base = 0;
57    for chunk in chunks.iter() {
58        if chunk != u64::MAX {
59            for bit in 0..64 {
60                if chunk & (1 << bit) == 0 {
61                    unsafe { ptr.add(base + bit).write(T::default_value()) };
62                }
63            }
64        }
65
66        base += 64;
67    }
68
69    let rem = chunks.remainder_bits();
70    for bit in 0..chunks.remainder_len() {
71        if rem & (1 << bit) == 0 {
72            unsafe { ptr.add(base + bit).write(T::default_value()) };
73        }
74    }
75}
76
77/// Decode an Arrow IPC batch and call `callback` once with all columns.
78///
79/// Callback signature:
80/// ```js
81/// callback(names: string[], values: TypedArray[], validities: (Uint8Array|null)[], dictionaries: (string[]|null)[]) => void | Promise<void>
82/// ```
83///
84/// If the callback returns a `Promise`, it is awaited before the Arrow
85/// batch (and therefore the zero-copy typed-array views into it) is
86/// dropped. A synchronous callback returning `undefined` is supported
87/// with no promise-handling overhead.
88pub(crate) async fn decode_and_call(
89    arrow: &[u8],
90    float32: bool,
91    callback: &Function,
92) -> Result<(), JsValue> {
93    let cursor = Cursor::new(arrow);
94    let reader = StreamReader::try_new(cursor, None)
95        .map_err(|e| JsValue::from_str(&format!("Arrow decode error: {e}")))?;
96
97    let batch = reader
98        .into_iter()
99        .next()
100        .ok_or_else(|| JsValue::from_str("Arrow IPC contained no record batches"))?
101        .map_err(|e| JsValue::from_str(&format!("Arrow batch error: {e}")))?;
102
103    let schema = batch.schema();
104    let num_cols = batch.num_columns();
105
106    let js_names = Array::new_with_length(num_cols as u32);
107    let js_values = Array::new_with_length(num_cols as u32);
108    let js_validities = Array::new_with_length(num_cols as u32);
109    let js_dicts = Array::new_with_length(num_cols as u32);
110
111    // Storage for type-conversion buffers (Int64/Date32/Timestamp and
112    // `float32` narrowing). These MUST outlive the callback because
113    // `js_sys::*Array::view()` creates zero-copy views into their heap
114    // memory. Using `Box<[T]>` (rather than `Vec<T>`) yields a stable
115    // data pointer that won't move when the outer Vec grows, so a view
116    // created before the push stays valid.
117    let mut f32_storage: Vec<Box<[f32]>> = Vec::new();
118    let mut f64_storage: Vec<Box<[f64]>> = Vec::new();
119
120    // The bytes under a NULL slot in the source Arrow are UNDEFINED —
121    // the perspective engine zero-fills them but e.g. DuckDB's Arrow
122    // output leaves NaN (which, unlike the garbage a consumer merely
123    // *displays* wrong, poisons any consumer that aggregates, e.g. the
124    // treemap's bottom-up value sums). `zero_invalid_slots` normalizes
125    // every column in place — forcing invalid slots to 0 so all backends
126    // present the same value contract — without giving up the zero-copy
127    // view.
128    for col_idx in 0..num_cols {
129        let field = schema.field(col_idx);
130        let col = batch.column(col_idx);
131        let validity = col.nulls().map(|nulls| nulls.validity());
132
133        js_names.set(col_idx as u32, JsString::from(field.name().as_str()).into());
134
135        match col.data_type() {
136            DataType::UInt32 => {
137                let typed = col.as_primitive::<UInt32Type>();
138                zero_invalid_slots(typed);
139                let arr = unsafe { js_sys::Uint32Array::view(typed.values().as_ref()) };
140                js_values.set(col_idx as u32, arr.into());
141                js_dicts.set(col_idx as u32, JsValue::NULL);
142            },
143            DataType::Int32 => {
144                let typed = col.as_primitive::<Int32Type>();
145                zero_invalid_slots(typed);
146                let arr = unsafe { js_sys::Int32Array::view(typed.values().as_ref()) };
147                js_values.set(col_idx as u32, arr.into());
148                js_dicts.set(col_idx as u32, JsValue::NULL);
149            },
150            DataType::Float32 => {
151                let typed = col.as_primitive::<Float32Type>();
152                zero_invalid_slots(typed);
153                let arr = unsafe { js_sys::Float32Array::view(typed.values().as_ref()) };
154                js_values.set(col_idx as u32, arr.into());
155                js_dicts.set(col_idx as u32, JsValue::NULL);
156            },
157            DataType::Float64 => {
158                let typed = col.as_primitive::<Float64Type>();
159                zero_invalid_slots(typed);
160                if float32 {
161                    let vals: Box<[f32]> = typed.values().iter().map(|&v| v as f32).collect();
162
163                    let arr = unsafe { js_sys::Float32Array::view(&vals) };
164                    f32_storage.push(vals);
165                    js_values.set(col_idx as u32, arr.into());
166                } else {
167                    let arr = unsafe { js_sys::Float64Array::view(typed.values().as_ref()) };
168
169                    js_values.set(col_idx as u32, arr.into());
170                }
171
172                js_dicts.set(col_idx as u32, JsValue::NULL);
173            },
174            DataType::Date32 => {
175                // Datetime values are always emitted as Float64 — narrowing
176                // epoch-ms to f32 collapses ~256 ms of resolution at modern
177                // timestamps, so the `float32` flag is intentionally ignored
178                // for date/timestamp columns.
179                let typed = col.as_primitive::<Date32Type>();
180                zero_invalid_slots(typed);
181                let vals: Box<[f64]> = typed
182                    .values()
183                    .iter()
184                    .map(|&v| v as f64 * 86_400_000.0)
185                    .collect();
186
187                let arr = unsafe { js_sys::Float64Array::view(&vals) };
188                f64_storage.push(vals);
189                js_values.set(col_idx as u32, arr.into());
190                js_dicts.set(col_idx as u32, JsValue::NULL);
191            },
192            DataType::Timestamp(TimeUnit::Millisecond, _) => {
193                let typed = col.as_primitive::<TimestampMillisecondType>();
194                zero_invalid_slots(typed);
195                let vals: Box<[f64]> = typed.values().iter().map(|&v| v as f64).collect();
196
197                let arr = unsafe { js_sys::Float64Array::view(&vals) };
198                f64_storage.push(vals);
199                js_values.set(col_idx as u32, arr.into());
200                js_dicts.set(col_idx as u32, JsValue::NULL);
201            },
202            DataType::Int64 => {
203                let typed = col.as_primitive::<Int64Type>();
204                zero_invalid_slots(typed);
205                if float32 {
206                    let vals: Box<[f32]> = typed.values().iter().map(|&v| v as f32).collect();
207
208                    let arr = unsafe { js_sys::Float32Array::view(&vals) };
209                    f32_storage.push(vals);
210                    js_values.set(col_idx as u32, arr.into());
211                } else {
212                    let vals: Box<[f64]> = typed.values().iter().map(|&v| v as f64).collect();
213
214                    let arr = unsafe { js_sys::Float64Array::view(&vals) };
215                    f64_storage.push(vals);
216                    js_values.set(col_idx as u32, arr.into());
217                }
218
219                js_dicts.set(col_idx as u32, JsValue::NULL);
220            },
221            DataType::Dictionary(..) => {
222                let dict = col.as_dictionary::<Int32Type>();
223                let keys = dict.keys();
224                zero_invalid_slots(keys);
225                let arr = unsafe { js_sys::Int32Array::view(keys.values().as_ref()) };
226                js_values.set(col_idx as u32, arr.into());
227
228                let values = dict.values().as_string::<i32>();
229                let js_dict = Array::new_with_length(values.len() as u32);
230                for i in 0..values.len() {
231                    js_dict.set(i as u32, JsValue::from_str(values.value(i)));
232                }
233                js_dicts.set(col_idx as u32, js_dict.into());
234            },
235            dt => {
236                return Err(JsValue::from_str(&format!(
237                    "Unsupported column type for typed array: {dt}"
238                )));
239            },
240        }
241
242        // SAFETY: Validity bitmap is owned by `batch` which outlives the
243        // callback — safe to view zero-copy.
244        let js_validity = validity.map(|v| unsafe { Uint8Array::view(v) });
245        js_validities.set(
246            col_idx as u32,
247            js_validity
248                .as_ref()
249                .map(JsValue::from)
250                .unwrap_or(JsValue::NULL),
251        );
252    }
253
254    let ret = callback.call4(
255        &JsValue::UNDEFINED,
256        &js_names.into(),
257        &js_values.into(),
258        &js_validities.into(),
259        &js_dicts.into(),
260    )?;
261
262    // If the callback returned a Promise, await it before releasing the
263    // batch — zero-copy TypedArray views into `batch`/`f32_storage`/
264    // `f64_storage` must remain valid for the full lifetime of the
265    // awaited work.
266    if ret.is_instance_of::<js_sys::Promise>() {
267        let promise: js_sys::Promise = ret.unchecked_into();
268        wasm_bindgen_futures::JsFuture::from(promise).await?;
269    }
270
271    // Keep storage alive until after the callback (and its awaited
272    // promise, if any) returns.
273    drop(f32_storage);
274    drop(f64_storage);
275
276    Ok(())
277}