Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 19 additions & 13 deletions packages/cubejs-backend-native/js/ResultWrapper.ts
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,24 @@ export function rowsToColumnar(rawData: any): JsRawColumnarData {
return { members, columns };
}

/**
* Pivot to columnar before serializing: the row-oriented form repeats
* every column name on every row, which inflates JSON size and forces
* the Rust side to allocate a per-row map before transposing back to
* its native columnar `QueryResult` representation.
*
* Serialize to a Buffer so the Rust side can decode via
* serde_json::from_slice instead of walking a JsValue through the
* Neon bridge with JsValueDeserializer. On 5 MB of AoO rows
* (~21k rows × 8 fields) the JsValue walk costs ~80 ms locally;
* Buffer + serde_json is ~7× faster (M3 MAX) and tracks V8's JSON.parse
* (~11 ms on the same payload). On a real server it should be 3-6× slower,
* so avoiding the JsValue walk matters even more there.
*/
export function rowsToColumnarBuffer(rawData: any): Buffer {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Low — rowsToColumnar derives the schema from rows[0] only. Object.keys(rows[0]) at line 39 fixes the column set from row 0. If later rows in the stream buffer ever include additional keys, those columns are silently dropped before crossing the bridge. Postgres cursors return uniform row shape today, so this is likely safe — but this invariant now matters per streaming chunk (not just per full result), so worth a short doc note here or in rowsToColumnar making it explicit.

return Buffer.from(JSON.stringify(rowsToColumnar(rawData)));
}

class BaseWrapper {
public readonly isWrapper: boolean = true;
}
Expand Down Expand Up @@ -191,19 +209,7 @@ export class ResultWrapper extends BaseWrapper implements DataResult {
return [this[NATIVE_REFERENCE]];
}

// Pivot to columnar before serializing: the row-oriented form repeats
// every column name on every row, which inflates JSON size and forces
// the Rust side to allocate a per-row map before transposing back to
// its native columnar `QueryResult` representation.
//
// Serialize to a Buffer so the Rust side can decode via
// serde_json::from_slice instead of walking a JsValue through the
// Neon bridge with JsValueDeserializer. On 5 MB of AoO rows
// (~21k rows × 8 fields) the JsValue walk costs ~80 ms locally;
// Buffer + serde_json is ~7× faster (M3 MAX) and tracks V8's JSON.parse
// (~11 ms on the same payload). On a real server it should be 3-6× slower,
// so avoiding the JsValue walk matters even more there.
return [Buffer.from(JSON.stringify(rowsToColumnar(this.jsResult)))];
return [rowsToColumnarBuffer(this.jsResult)];
}

public setTransformData(td: any) {
Expand Down
6 changes: 3 additions & 3 deletions packages/cubejs-backend-native/js/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import path from 'path';
import { Writable } from 'stream';
import type { Request as ExpressRequest } from 'express';
import { CacheMode } from '@cubejs-backend/shared';
import { NativeQueryResultRef, ResultWrapper } from './ResultWrapper';
import { NativeQueryResultRef, ResultWrapper, rowsToColumnarBuffer } from './ResultWrapper';

export * from './ResultWrapper';

Expand Down Expand Up @@ -294,7 +294,7 @@ function wrapNativeFunctionWithStream(
if (chunkBuffer.length < chunkLength) {
callback(null);
} else {
Comment thread
claude[bot] marked this conversation as resolved.
const toSend = chunkBuffer;
const toSend = rowsToColumnarBuffer(chunkBuffer);
chunkBuffer = [];
writerOrChannel.chunk(toSend, callback);
}
Expand All @@ -308,7 +308,7 @@ function wrapNativeFunctionWithStream(
}
};
if (chunkBuffer.length > 0) {
const toSend = chunkBuffer;
const toSend = rowsToColumnarBuffer(chunkBuffer);
chunkBuffer = [];
writerOrChannel.chunk(toSend, end);
} else {
Expand Down
88 changes: 2 additions & 86 deletions packages/cubejs-backend-native/src/orchestrator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,10 @@ use crate::node_obj_deserializer::JsValueDeserializer;
use crate::transport::MapCubeErrExt;
use cubeorchestrator::query_message_parser::QueryResult;
use cubeorchestrator::query_result_transform::{
DBResponsePrimitive, InternedKeyLookup, RequestResultData, RequestResultDataMulti,
TransformedData,
DBResponsePrimitive, RequestResultData, RequestResultDataMulti, TransformedData,
};
use cubeorchestrator::transport::{JsRawColumnarData, TransformDataRequest};
use cubesql::compile::engine::df::scan::{ColumnarValueObject, FieldValue, ValueObject};
use cubesql::compile::engine::df::scan::{ColumnarValueObject, FieldValue};
use cubesql::CubeError;
use neon::context::{Context, FunctionContext, ModuleContext};
use neon::handle::Handle;
Expand Down Expand Up @@ -144,89 +143,6 @@ fn db_primitive_to_field_value(value: &DBResponsePrimitive) -> FieldValue<'_> {
}
}

impl ValueObject for ResultWrapper {
fn len(&mut self) -> Result<usize, CubeError> {
if self.transformed_data.is_none() {
self.transform_result()?;
}

let data = self.transformed_data.as_ref().unwrap();

match data {
TransformedData::Compact {
members: _members,
dataset,
} => Ok(dataset.len()),
TransformedData::Columnar {
members: _members,
columns,
} => Ok(columns.first().map(|c| c.len()).unwrap_or(0)),
TransformedData::Vanilla(dataset) => Ok(dataset.len()),
}
}

fn get(&mut self, index: usize, field_name: &str) -> Result<FieldValue<'_>, CubeError> {
if self.transformed_data.is_none() {
self.transform_result()?;
}

let data = self.transformed_data.as_ref().unwrap();

let value = match data {
TransformedData::Compact { members, dataset } => {
let Some(row) = dataset.get(index) else {
return Err(CubeError::internal(format!(
"Unexpected response from Cube, can't get {} row",
index
)));
};

let Some(member_index) = members.iter().position(|m| m == field_name) else {
// Missing field → NULL, matching `Vanilla` semantics below.
return Ok(FieldValue::Null);
};

row.get(member_index).unwrap_or(&DBResponsePrimitive::Null)
}
TransformedData::Columnar { members, columns } => {
let Some(member_index) = members.iter().position(|m| m == field_name) else {
// Missing field → NULL, matching `Vanilla` semantics below.
return Ok(FieldValue::Null);
};

let Some(column) = columns.get(member_index) else {
return Err(CubeError::internal(format!(
"Unexpected response from Cube, missing column for '{}'",
field_name
)));
};

let Some(value) = column.get(index) else {
return Err(CubeError::user(format!(
"Unexpected response from Cube, can't get {} row",
index
)));
};

value
}
TransformedData::Vanilla(dataset) => {
let Some(row) = dataset.get(index) else {
return Err(CubeError::internal(format!(
"Unexpected response from Cube, can't get {} row",
index
)));
};

row.get(&InternedKeyLookup::new(field_name))
.unwrap_or(&DBResponsePrimitive::Null)
}
};

Ok(db_primitive_to_field_value(value))
}
}

impl ColumnarValueObject for ResultWrapper {
fn len(&mut self) -> Result<usize, CubeError> {
if self.transformed_data.is_none() {
Expand Down
109 changes: 14 additions & 95 deletions packages/cubejs-backend-native/src/stream.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
use cubesql::compile::engine::df::scan::{
transform_response, FieldValue, MemberField, RecordBatch, SchemaRef, ValueObject,
transform_response, JsonColumnarValueObject, MemberField, RecordBatch, SchemaRef,
};
use std::borrow::Cow;

use std::cell::RefCell;
use std::future::Future;
Expand All @@ -15,13 +14,12 @@ use crate::channel::call_js_fn;
use cubesql::CubeError;

use neon::prelude::*;
use neon::types::buffer::TypedArray;
use tokio::sync::{oneshot, Semaphore};

#[cfg(feature = "neon-debug")]
use log::trace;

use neon::types::JsDate;

use crate::utils::bind_method;

use tokio::sync::mpsc::{channel as mpsc_channel, Receiver, Sender};
Expand Down Expand Up @@ -256,110 +254,31 @@ fn wait_for_future_and_execute_callback(
});
}

pub struct JsValueObject<'a> {
pub cx: FunctionContext<'a>,
pub handle: Handle<'a, JsArray>,
}

fn js_value_to_json_string<'a, C: Context<'a>>(
cx: &mut C,
value: Handle<'a, JsValue>,
) -> Result<String, CubeError> {
let global = cx.global_object();
let json = global
.get::<JsObject, _, _>(cx, "JSON")
.map_err(|e| CubeError::internal(format!("Can't get JSON global: {}", e)))?;
let stringify = json
.get::<JsFunction, _, _>(cx, "stringify")
.map_err(|e| CubeError::internal(format!("Can't get JSON.stringify: {}", e)))?;
let undefined = cx.undefined().upcast::<JsValue>();
let result = stringify
.call(cx, undefined, [value])
.map_err(|e| CubeError::internal(format!("JSON.stringify failed: {}", e)))?;
let s = result.downcast::<JsString, _>(cx).map_err(|e| {
CubeError::internal(format!("JSON.stringify did not return a string: {}", e))
})?;
Ok(s.value(cx))
}

impl ValueObject for JsValueObject<'_> {
fn len(&mut self) -> Result<usize, CubeError> {
Ok(self.handle.len(&mut self.cx) as usize)
}

fn get(&mut self, index: usize, field_name: &str) -> Result<FieldValue<'_>, CubeError> {
let value = self
.handle
.get::<JsObject, _, _>(&mut self.cx, index as u32)
.map_err(|e| {
CubeError::internal(format!("Can't get object at array index {}: {}", index, e))
})?
.get::<JsValue, _, _>(&mut self.cx, field_name)
.map_err(|e| {
CubeError::internal(format!("Can't get '{}' field value: {}", field_name, e))
})?;
if let Ok(s) = value.downcast::<JsString, _>(&mut self.cx) {
Ok(FieldValue::String(Cow::Owned(s.value(&mut self.cx))))
} else if let Ok(n) = value.downcast::<JsNumber, _>(&mut self.cx) {
Ok(FieldValue::Number(n.value(&mut self.cx)))
} else if let Ok(b) = value.downcast::<JsBoolean, _>(&mut self.cx) {
Ok(FieldValue::Bool(b.value(&mut self.cx)))
} else if value.downcast::<JsUndefined, _>(&mut self.cx).is_ok()
|| value.downcast::<JsNull, _>(&mut self.cx).is_ok()
{
Ok(FieldValue::Null)
} else if value.is_a::<JsArray, _>(&mut self.cx) {
Ok(FieldValue::String(Cow::Owned(js_value_to_json_string(
&mut self.cx,
value,
)?)))
} else if let Ok(b) = value.downcast::<JsDate, _>(&mut self.cx) {
// TODO: Support it?
Err(CubeError::internal(format!(
"Expected primitive value but found JsDate({:?})",
b
)))
} else if value.is_a::<JsObject, _>(&mut self.cx) {
Ok(FieldValue::String(Cow::Owned(js_value_to_json_string(
&mut self.cx,
value,
)?)))
} else {
Err(CubeError::internal(format!(
"Expected primitive value but found: {:?}",
value
)))
}
}
}

fn js_stream_push_chunk(mut cx: FunctionContext) -> JsResult<JsUndefined> {
#[cfg(feature = "neon-debug")]
trace!("JsWriteStream.push_chunk");

let this = cx
.this::<JsValue>()?
.downcast_or_throw::<JsBox<JsWriteStream>, _>(&mut cx)?;
let chunk_array = cx.argument::<JsArray>(0)?;
let callback = cx.argument::<JsFunction>(1)?.root(&mut cx);
let mut value_object = JsValueObject {
cx,
handle: chunk_array,
};

let chunk_buffer = cx.argument::<JsBuffer>(0)?;
let mut value_object =
match serde_json::from_slice::<JsonColumnarValueObject>(chunk_buffer.as_slice(&cx)) {
Ok(v) => v,
Err(e) => return cx.throw_error(format!("Can't parse columnar chunk JSON: {}", e)),
};
let value =
match transform_response(&mut value_object, this.schema.clone(), &this.member_fields) {
Ok(value) => value,
Err(e) => return value_object.cx.throw_error(e.message),
Err(e) => return cx.throw_error(e.message),
};

let future = this.push_chunk(value);
wait_for_future_and_execute_callback(
this.tokio_handle.clone(),
value_object.cx.channel(),
callback,
future,
);

Ok(value_object.cx.undefined())
wait_for_future_and_execute_callback(this.tokio_handle.clone(), cx.channel(), callback, future);

Ok(cx.undefined())
}

fn js_stream_start(mut cx: FunctionContext) -> JsResult<JsUndefined> {
Expand Down
16 changes: 4 additions & 12 deletions packages/cubejs-backend-native/src/transport.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ use crate::{
use async_trait::async_trait;
use cubeorchestrator::query_result_transform::RequestResultData;
use cubesql::compile::engine::df::scan::{
build_response_schema, convert_transport_response_columnar, transform_columnar_response,
CacheMode, MemberField, RecordBatch, SchemaRef,
build_response_schema, convert_transport_response, transform_response, CacheMode, MemberField,
RecordBatch, SchemaRef,
};
use cubesql::compile::engine::df::wrapper::SqlQuery;
use cubesql::transport::{
Expand Down Expand Up @@ -524,11 +524,7 @@ impl TransportService for NodeBridgeTransport {
}
};

break convert_transport_response_columnar(
response,
schema.clone(),
member_fields,
);
break convert_transport_response(response, schema.clone(), member_fields);
}
ValueFromJs::ResultWrapper(result_wrappers) => {
break result_wrappers
Expand All @@ -540,11 +536,7 @@ impl TransportService for NodeBridgeTransport {
wrapper.external,
);

transform_columnar_response(
&mut wrapper,
updated_schema,
&member_fields,
)
transform_response(&mut wrapper, updated_schema, &member_fields)
})
.collect::<Result<Vec<_>, _>>();
}
Expand Down
Loading
Loading