perspective_js/
typed_array.rs1use 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#[derive(Default, serde::Deserialize, TS)]
36pub struct TypedArrayWindow {
37 #[serde(flatten)]
38 pub view_window: ViewWindow,
39
40 #[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
77pub(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 let mut f32_storage: Vec<Box<[f32]>> = Vec::new();
118 let mut f64_storage: Vec<Box<[f64]>> = Vec::new();
119
120 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 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 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 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 drop(f32_storage);
274 drop(f64_storage);
275
276 Ok(())
277}