pass ResultWrapper to native and offload transformData from the eventloop

KSDaemon · KSDaemon · commit 5657478fb7d8 · 2025-01-23T19:30:41.000+02:00
diff --git a/packages/cubejs-backend-native/js/ResultWrapper.ts b/packages/cubejs-backend-native/js/ResultWrapper.ts
@@ -186,6 +186,8 @@ class BaseWrapperArray extends BaseWrapper {
     return [transformDataJson, rawData, resultDataJson];
   }
 
+  // Is invoked from the native side to get
+  // an array of all raw wrapped results
   public getResults(): ResultWrapper[] {
     return this.results;
   }
diff --git a/packages/cubejs-backend-native/js/index.ts b/packages/cubejs-backend-native/js/index.ts
@@ -267,7 +267,7 @@ function wrapNativeFunctionWithStream(
       } else if (response.error) {
         writerOrChannel.reject(errorString(response));
       } else if (response.isWrapper) { // Native wrapped result
-        writerOrChannel.resolve(await response.getFinalResult());
+        writerOrChannel.resolve(response);
       } else {
         writerOrChannel.resolve(JSON.stringify(response));
       }
diff --git a/packages/cubejs-backend-native/src/channel.rs b/packages/cubejs-backend-native/src/channel.rs
@@ -4,10 +4,10 @@ use std::collections::HashMap;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 
+use crate::orchestrator::ResultWrapper;
 use crate::transport::MapCubeErrExt;
 use crate::utils::bind_method;
 use async_trait::async_trait;
-use cubeorchestrator::query_result_transform::RequestResultArray;
 use cubesql::transport::{SqlGenerator, SqlTemplates};
 use cubesql::CubeError;
 #[cfg(debug_assertions)]
@@ -197,7 +197,7 @@ where
 #[derive(Debug)]
 pub enum ValueFromJs {
     String(String),
-    RequestResultArray(RequestResultArray),
+    ResultWrapper(Vec<ResultWrapper>),
 }
 
 #[allow(clippy::type_complexity)]
diff --git a/packages/cubejs-backend-native/src/orchestrator.rs b/packages/cubejs-backend-native/src/orchestrator.rs
@@ -1,14 +1,18 @@
 use crate::node_obj_deserializer::JsValueDeserializer;
+use crate::transport::MapCubeErrExt;
 use cubeorchestrator::query_message_parser::QueryResult;
 use cubeorchestrator::query_result_transform::{
-    get_final_cubestore_result_array, RequestResultArray, RequestResultData, RequestResultDataMulti,
+    get_final_cubestore_result_array, RequestResultArray, RequestResultData,
+    RequestResultDataMulti, TransformedData,
 };
 use cubeorchestrator::transport::{JsRawData, TransformDataRequest};
+use cubesql::CubeError;
 use neon::context::{Context, FunctionContext, ModuleContext};
 use neon::handle::Handle;
 use neon::object::Object;
 use neon::prelude::{
-    JsArray, JsArrayBuffer, JsBox, JsBuffer, JsObject, JsPromise, JsResult, JsValue, NeonResult,
+    JsArray, JsArrayBuffer, JsBox, JsBuffer, JsFunction, JsObject, JsPromise, JsResult, JsValue,
+    NeonResult,
 };
 use neon::types::buffer::TypedArray;
 use serde::Deserialize;
@@ -27,6 +31,108 @@ pub fn register_module(cx: &mut ModuleContext) -> NeonResult<()> {
     Ok(())
 }
 
+// #[derive(Debug, Clone)]
+// pub enum ResultWrapperData {
+//     JsObj(RecordBatch),
+//     Native()
+// }
+
+#[derive(Debug, Clone)]
+pub struct ResultWrapper {
+    transform_data: TransformDataRequest,
+    data: Arc<QueryResult>,
+}
+
+impl ResultWrapper {
+    pub fn from_js_result_wrapper(
+        cx: &mut FunctionContext<'_>,
+        js_result_wrapper_val: Handle<JsValue>,
+    ) -> Result<Self, CubeError> {
+        let js_result_wrapper = js_result_wrapper_val
+            .downcast::<JsObject, _>(cx)
+            .map_cube_err("Can't downcast JS ResultWrapper to object")?;
+
+        let get_transform_data_js_method: Handle<JsFunction> = js_result_wrapper
+            .get(cx, "getTransformData")
+            .map_cube_err("Can't get getTransformData() method from JS ResultWrapper object")?;
+
+        let transform_data_js_arr = get_transform_data_js_method
+            .call(cx, js_result_wrapper.upcast::<JsValue>(), [])
+            .map_cube_err("Error calling getTransformData() method of ResultWrapper object")?
+            .downcast::<JsArray, _>(cx)
+            .map_cube_err("Can't downcast JS transformData to array")?
+            .to_vec(cx)
+            .map_cube_err("Can't convert JS transformData to array")?;
+
+        let transform_data_js = transform_data_js_arr.first().unwrap();
+
+        let deserializer = JsValueDeserializer::new(cx, *transform_data_js);
+        let transform_request: TransformDataRequest = match Deserialize::deserialize(deserializer) {
+            Ok(data) => data,
+            Err(_) => {
+                return Err(CubeError::internal(
+                    "Can't deserialize transformData from JS ResultWrapper object".to_string(),
+                ))
+            }
+        };
+
+        let get_raw_data_js_method: Handle<JsFunction> = js_result_wrapper
+            .get(cx, "getRawData")
+            .map_cube_err("Can't get getRawData() method from JS ResultWrapper object")?;
+
+        let raw_data_js_arr = get_raw_data_js_method
+            .call(cx, js_result_wrapper.upcast::<JsValue>(), [])
+            .map_cube_err("Error calling getRawData() method of ResultWrapper object")?
+            .downcast::<JsArray, _>(cx)
+            .map_cube_err("Can't downcast JS rawData to array")?
+            .to_vec(cx)
+            .map_cube_err("Can't convert JS rawData to array")?;
+
+        let raw_data_js = raw_data_js_arr.first().unwrap();
+
+        let query_result =
+            if let Ok(js_box) = raw_data_js.downcast::<JsBox<Arc<QueryResult>>, _>(cx) {
+                Arc::clone(&js_box)
+            } else if let Ok(js_array) = raw_data_js.downcast::<JsArray, _>(cx) {
+                let deserializer = JsValueDeserializer::new(cx, js_array.upcast());
+                let js_raw_data: JsRawData = match Deserialize::deserialize(deserializer) {
+                    Ok(data) => data,
+                    Err(_) => {
+                        return Err(CubeError::internal(
+                            "Can't deserialize results raw data from JS ResultWrapper object".to_string(),
+                        ));
+                    }
+                };
+
+                QueryResult::from_js_raw_data(js_raw_data)
+                    .map(Arc::new)
+                    .map_cube_err("Can't build results data from JS rawData")?
+            } else {
+                return Err(CubeError::internal(
+                    "Can't deserialize results raw data from JS ResultWrapper object".to_string(),
+                ));
+            };
+
+        Ok(Self {
+            transform_data: transform_request,
+            data: query_result,
+        })
+    }
+
+    pub fn transform_result(&self) -> Result<TransformedData, CubeError> {
+        let transformed = TransformedData::transform(&self.transform_data, &self.data)
+            .map_cube_err("Can't prepare transformed data")?;
+
+        Ok(transformed)
+    }
+}
+
+pub type JsResultDataVectors = (
+    Vec<TransformDataRequest>,
+    Vec<Arc<QueryResult>>,
+    Vec<RequestResultData>,
+);
+
 fn json_to_array_buffer<'a, C>(
     mut cx: C,
     json_data: Result<String, anyhow::Error>,
@@ -144,12 +250,6 @@ pub fn final_query_result(mut cx: FunctionContext) -> JsResult<JsPromise> {
     Ok(promise)
 }
 
-pub type JsResultDataVectors = (
-    Vec<TransformDataRequest>,
-    Vec<Arc<QueryResult>>,
-    Vec<RequestResultData>,
-);
-
 pub fn convert_final_query_result_array_from_js(
     cx: &mut FunctionContext<'_>,
     transform_data_array: Handle<JsValue>,
diff --git a/packages/cubejs-backend-native/src/transport.rs b/packages/cubejs-backend-native/src/transport.rs
@@ -6,15 +6,12 @@ use std::fmt::Display;
 use crate::auth::NativeAuthContext;
 use crate::channel::{call_raw_js_with_channel_as_callback, NodeSqlGenerator, ValueFromJs};
 use crate::node_obj_serializer::NodeObjSerializer;
-use crate::orchestrator::convert_final_query_result_array_from_js;
+use crate::orchestrator::ResultWrapper;
 use crate::{
     auth::TransportRequest, channel::call_js_with_channel_as_callback,
     stream::call_js_with_stream_as_callback,
 };
 use async_trait::async_trait;
-use cubeorchestrator::query_result_transform::{
-    get_final_cubestore_result_array, RequestResultArray,
-};
 use cubesql::compile::engine::df::scan::{MemberField, SchemaRef};
 use cubesql::compile::engine::df::wrapper::SqlQuery;
 use cubesql::transport::{
@@ -337,6 +334,7 @@ impl TransportService for NodeBridgeTransport {
         sql_query: Option<SqlQuery>,
         ctx: AuthContextRef,
         meta: LoadRequestMeta,
+    // ) -> Result<Vec<RecordBatch>, CubeError> {
     ) -> Result<TransportLoadResponse, CubeError> {
         trace!("[transport] Request ->");
 
@@ -378,52 +376,43 @@ impl TransportService for NodeBridgeTransport {
                 extra,
                 Box::new(|cx, v| Ok(cx.string(v).as_value(cx))),
                 Box::new(move |cx, v| {
-                    // It's too heavy/slow to get instance of ResultArrayWrapper from JS
-                    // and then call/await the .getFinalResult() method which needs be
-                    // executed again on JS side to get the actual needed date,
-                    // instead we pass it directly from JS side.
-                    // In case of wrapped result it's actually a tuple of
-                    // (transformDataJson[], rawData[], resultDataJson[])
-                    if let Ok(result_wrapped) = v.downcast::<JsArray, _>(cx) {
-                        let res_wrapped_vec = result_wrapped.to_vec(cx).map_cube_err("Can't convert JS result to array")?;
-
-                        if res_wrapped_vec.len() != 3 {
-                            return Err(CubeError::internal("Expected a tuple with 3 elements: transformDataJson[], rawData[], resultDataJson[]".to_string()));
-                        }
+                    if let Ok(js_result_wrapped) = v.downcast::<JsObject, _>(cx) {
+                        let get_results_js_method: Handle<JsFunction> = js_result_wrapped
+                            .get(cx, "getResults")
+                            .map_cube_err(
+                                "Can't get getResults() method from JS ResultWrapper object",
+                            )?;
 
-                        let transform_data_array = res_wrapped_vec.first().unwrap();
-                        let data_array = res_wrapped_vec.get(1).unwrap()
-                            .downcast_or_throw::<JsArray, _>(cx).map_cube_err("Can't downcast js data to array")?;
-                        let results_data_array = res_wrapped_vec.get(2).unwrap();
-
-                        match convert_final_query_result_array_from_js(
-                            cx,
-                            *transform_data_array,
-                            data_array,
-                            *results_data_array,
-                        ) {
-                            Ok((transform_requests, cube_store_results, mut request_results)) => {
-                                get_final_cubestore_result_array(
-                                    &transform_requests,
-                                    &cube_store_results,
-                                    &mut request_results,
-                                ).map_cube_err("Can't build result array")?;
-
-                                Ok(ValueFromJs::RequestResultArray(RequestResultArray {
-                                    results: request_results,
-                                }))
-                            }
-                            Err(err) => {
-                                Err(CubeError::internal(format!("Error converting result data: {:?}", err.to_string())))
-                            }
-                        }
+                        let results = get_results_js_method
+                            .call(cx, js_result_wrapped.upcast::<JsValue>(), [])
+                            .map_cube_err(
+                                "Error calling getResults() method of ResultWrapper object",
+                            )?;
 
+                        let js_res_wrapped_vec = results
+                            .downcast::<JsArray, _>(cx)
+                            .map_cube_err("Can't downcast JS result to array")?
+                            .to_vec(cx)
+                            .map_cube_err("Can't convert JS result to array")?;
+
+                        let native_wrapped_results = js_res_wrapped_vec
+                            .iter()
+                            .map(|r| ResultWrapper::from_js_result_wrapper(cx, *r))
+                            .collect::<Result<Vec<_>, _>>()
+                            .map_cube_err(
+                                "Can't construct result wrapper from JS ResultWrapper object",
+                            )?;
+
+                        Ok(ValueFromJs::ResultWrapper(native_wrapped_results))
                     } else if let Ok(str) = v.downcast::<JsString, _>(cx) {
                         Ok(ValueFromJs::String(str.value(cx)))
                     } else {
-                        Err(CubeError::internal("Can't downcast callback argument to string or resultWrapper object".to_string()))
+                        Err(CubeError::internal(
+                            "Can't downcast callback argument to string or resultWrapper object"
+                                .to_string(),
+                        ))
                     }
-                })
+                }),
             )
             .await;
 
@@ -475,12 +464,15 @@ impl TransportService for NodeBridgeTransport {
                     break serde_json::from_value::<TransportLoadResponse>(response)
                         .map_err(|err| CubeError::user(err.to_string()));
                 }
-                ValueFromJs::RequestResultArray(result) => {
+                ValueFromJs::ResultWrapper(result_wrappers) => {
                     let response = TransportLoadResponse {
                         pivot_query: None,
                         slow_query: None,
                         query_type: None,
-                        results: result.results.into_iter().map(|v| v.into()).collect(),
+                        results: result_wrappers
+                            .into_iter()
+                            .map(|v| v.transform_result().unwrap().into())
+                            .collect(),
                     };
                     break Ok(response);
                 }
diff --git a/rust/cubesql/cubeclient/src/models/v1_load_result.rs b/rust/cubesql/cubeclient/src/models/v1_load_result.rs
@@ -7,7 +7,7 @@
  *
  * Generated by: https://openapi-generator.tech
  */
-use cubeorchestrator::query_result_transform::{RequestResultData, TransformedData};
+use cubeorchestrator::query_result_transform::TransformedData;
 
 #[derive(Clone, Debug, PartialEq, Default, Serialize, Deserialize)]
 pub struct V1LoadResult {
@@ -35,55 +35,32 @@ impl V1LoadResult {
     }
 }
 
-impl From<RequestResultData> for V1LoadResult {
-    fn from(rrd: RequestResultData) -> Self {
-        let data = match rrd.data {
-            None => {
+impl From<TransformedData> for V1LoadResult {
+    fn from(d: TransformedData) -> Self {
+        let data = match d {
+            TransformedData::Compact {
+                members: _members,
+                dataset: _dataset,
+            } => {
+                // TODO V1LoadResult expects only Vanilla results?
                 vec![]
             }
-            Some(d) => {
-                match d {
-                    TransformedData::Compact {
-                        members: _members,
-                        dataset: _dataset,
-                    } => {
-                        // TODO V1LoadResult expects only Vanilla results?
-                        vec![]
-                    }
-                    TransformedData::Vanilla(v) => v
-                        .iter()
-                        .map(|item| serde_json::to_value(item).unwrap_or(serde_json::Value::Null))
-                        .collect(),
-                }
-            }
+            TransformedData::Vanilla(v) => v
+                .iter()
+                .map(|item| serde_json::to_value(item).unwrap_or(serde_json::Value::Null))
+                .collect(),
         };
         let annotation = crate::models::V1LoadResultAnnotation {
-            measures: rrd
-                .annotation
-                .get("measures")
-                .map(|v| serde_json::to_value(v).unwrap_or(serde_json::Value::Null))
-                .unwrap_or(serde_json::Value::Null),
-            dimensions: rrd
-                .annotation
-                .get("dimensions")
-                .map(|v| serde_json::to_value(v).unwrap_or(serde_json::Value::Null))
-                .unwrap_or(serde_json::Value::Null),
-            segments: rrd
-                .annotation
-                .get("segments")
-                .map(|v| serde_json::to_value(v).unwrap_or(serde_json::Value::Null))
-                .unwrap_or(serde_json::Value::Null),
-            time_dimensions: rrd
-                .annotation
-                .get("timeDimensions")
-                .map(|v| serde_json::to_value(v).unwrap_or(serde_json::Value::Null))
-                .unwrap_or(serde_json::Value::Null),
+            measures: serde_json::Value::Null,
+            dimensions: serde_json::Value::Null,
+            segments: serde_json::Value::Null,
+            time_dimensions: serde_json::Value::Null,
         };
         V1LoadResult {
-            data_source: rrd.data_source,
+            data_source: None,
             annotation: Box::new(annotation),
             data,
-            refresh_key_values: rrd.refresh_key_values,
+            refresh_key_values: None,
         }
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -186,6 +186,8 @@ class BaseWrapperArray extends BaseWrapper {`
`186`	`186`	`return [transformDataJson, rawData, resultDataJson];`
`187`	`187`	`}`
`188`	`188`
	`189`	`+ // Is invoked from the native side to get`
	`190`	`+ // an array of all raw wrapped results`
`189`	`191`	`public getResults(): ResultWrapper[] {`
`190`	`192`	`return this.results;`
`191`	`193`	`}`
Original file line number	Diff line number	Diff line change
`@@ -267,7 +267,7 @@ function wrapNativeFunctionWithStream(`
`267`	`267`	`} else if (response.error) {`
`268`	`268`	`writerOrChannel.reject(errorString(response));`
`269`	`269`	`} else if (response.isWrapper) { // Native wrapped result`
`270`		`- writerOrChannel.resolve(await response.getFinalResult());`
	`270`	`+ writerOrChannel.resolve(response);`
`271`	`271`	`} else {`
`272`	`272`	`writerOrChannel.resolve(JSON.stringify(response));`
`273`	`273`	`}`