Address remaining TODOs after custom executor implementation (#14)

Author: OlegWock

vegafusion-core/src/proto/pretransform.proto

@@ -127,7 +127,7 @@ message ExportUpdate {
   string namespace = 1;
   string name = 2;
   repeated uint32 scope = 3;
-  tasks.TaskValue value = 4;
+  tasks.MaterializedTaskValue value = 4;
 }
 
 message PreTransformLogicalPlanWarning {

vegafusion-core/src/proto/tasks.proto

@@ -4,26 +4,8 @@ package tasks;
 import "expression.proto";
 import "transforms.proto";
 
-// ## Task Value
-message TaskValue {
-  oneof data {
-  /*
-   * Representation of scalar as single column, single row, record batch in Arrow IPC format
-   */
-    bytes scalar = 1;
-
-  /*
-   * Serialized Arrow record batch in Arrow IPC format
-   */
-    bytes table = 2;
-  /*
-   * Serialized DataFusion LogicalPlan in protobuf format
-   */
-   bytes plan = 3;
-  }
-}
-
 // ## Materialized Task Value
+// Represents a fully materialized (computed) task value, either a scalar or table
 message MaterializedTaskValue {
   oneof data {
   /*
@@ -128,7 +110,7 @@ message Task {
   Variable variable = 1;
   repeated uint32 scope = 2;
   oneof task_kind {
-    TaskValue value = 3;
+    MaterializedTaskValue value = 3;
     DataValuesTask data_values = 4;
     DataUrlTask data_url = 5;
     DataSourceTask data_source = 6;

vegafusion-core/src/task_graph/graph.rs

@@ -11,9 +11,9 @@ use std::collections::HashMap;
 
 use crate::task_graph::task_value::TaskValue;
 
+use crate::proto::gen::tasks::materialized_task_value::Data;
 use crate::proto::gen::tasks::task::TaskKind;
-use crate::proto::gen::tasks::task_value::Data;
-use crate::proto::gen::tasks::TaskValue as ProtoTaskValue;
+use crate::proto::gen::tasks::MaterializedTaskValue as ProtoMaterializedTaskValue;
 use std::convert::TryFrom;
 use std::hash::{BuildHasher, Hash, Hasher};
 
@@ -216,7 +216,6 @@ impl TaskGraph {
                 match value.data.as_ref().unwrap() {
                     Data::Scalar(_) => "scalar".hash(&mut hasher),
                     Data::Table(_) => "data".hash(&mut hasher),
-                    Data::Plan(_) => "plan".hash(&mut hasher),
                 }
             } else {
                 // Include id_fingerprint of parents in the hash
@@ -303,7 +302,9 @@ impl TaskGraph {
         node.task = Some(Task {
             variable: node.task().variable.clone(),
             scope: node.task().scope.clone(),
-            task_kind: Some(TaskKind::Value(ProtoTaskValue::try_from(&value)?)),
+            task_kind: Some(TaskKind::Value(ProtoMaterializedTaskValue::try_from(
+                &value,
+            )?)),
             tz_config: None,
         });
 

vegafusion-core/src/task_graph/task.rs

@@ -3,7 +3,7 @@ use crate::proto::gen::tasks::{
     task::TaskKind, DataSourceTask, DataUrlTask, DataValuesTask, NodeValueIndex, Task, TzConfig,
     Variable,
 };
-use crate::proto::gen::tasks::{SignalTask, TaskValue as ProtoTaskValue};
+use crate::proto::gen::tasks::{MaterializedTaskValue as ProtoMaterializedTaskValue, SignalTask};
 use crate::task_graph::task_value::TaskValue;
 use std::convert::TryFrom;
 
@@ -33,7 +33,9 @@ impl Task {
         Self {
             variable: Some(variable),
             scope: Vec::from(scope),
-            task_kind: Some(TaskKind::Value(ProtoTaskValue::try_from(&value).unwrap())),
+            task_kind: Some(TaskKind::Value(
+                ProtoMaterializedTaskValue::try_from(&value).unwrap(),
+            )),
             tz_config: None,
         }
     }

vegafusion-core/src/task_graph/task_value.rs

@@ -1,9 +1,7 @@
 use crate::proto::gen::tasks::materialized_task_value::Data as MaterializedTaskValueData;
-use crate::proto::gen::tasks::task_value::Data as TaskValueData;
 use crate::proto::gen::tasks::ResponseTaskValue;
 use crate::proto::gen::tasks::{
-    MaterializedTaskValue as ProtoMaterializedTaskValue, TaskGraphValueResponse,
-    TaskValue as ProtoTaskValue, Variable,
+    MaterializedTaskValue as ProtoMaterializedTaskValue, TaskGraphValueResponse, Variable,
 };
 use crate::runtime::PlanExecutor;
 use crate::task_graph::memory::{
@@ -105,30 +103,26 @@ impl MaterializedTaskValue {
     }
 }
 
-impl TryFrom<&ProtoTaskValue> for TaskValue {
+impl TryFrom<&ProtoMaterializedTaskValue> for TaskValue {
     type Error = VegaFusionError;
 
-    fn try_from(value: &ProtoTaskValue) -> std::result::Result<Self, Self::Error> {
+    fn try_from(value: &ProtoMaterializedTaskValue) -> std::result::Result<Self, Self::Error> {
         match value.data.as_ref().unwrap() {
-            TaskValueData::Table(value) => Ok(Self::Table(VegaFusionTable::from_ipc_bytes(value)?)),
-            TaskValueData::Scalar(value) => {
+            MaterializedTaskValueData::Table(value) => {
+                Ok(Self::Table(VegaFusionTable::from_ipc_bytes(value)?))
+            }
+            MaterializedTaskValueData::Scalar(value) => {
                 let scalar_table = VegaFusionTable::from_ipc_bytes(value)?;
                 let scalar_rb = scalar_table.to_record_batch()?;
                 let scalar_array = scalar_rb.column(0);
                 let scalar = ScalarValue::try_from_array(scalar_array, 0)?;
                 Ok(Self::Scalar(scalar))
             }
-            // TODO: we could use datafusion_proto::bytes::logical_plan_from_bytes here, but that
-            // requires adding datafusion_proto to vegafusion-core deps, as well as passing
-            // datafusion session (maybe empty one?) to unserialize plan
-            TaskValueData::Plan(_value) => Err(VegaFusionError::internal(
-                "Deserialization of Plan TaskValue not yet implemented",
-            )),
         }
     }
 }
 
-impl TryFrom<&TaskValue> for ProtoTaskValue {
+impl TryFrom<&TaskValue> for ProtoMaterializedTaskValue {
     type Error = VegaFusionError;
 
     fn try_from(value: &TaskValue) -> std::result::Result<Self, Self::Error> {
@@ -138,17 +132,14 @@ impl TryFrom<&TaskValue> for ProtoTaskValue {
                 let scalar_rb = RecordBatch::try_from_iter(vec![("value", scalar_array)])?;
                 let ipc_bytes = VegaFusionTable::from(scalar_rb).to_ipc_bytes()?;
                 Ok(Self {
-                    data: Some(TaskValueData::Scalar(ipc_bytes)),
+                    data: Some(MaterializedTaskValueData::Scalar(ipc_bytes)),
                 })
             }
             TaskValue::Table(table) => Ok(Self {
-                data: Some(TaskValueData::Table(table.to_ipc_bytes()?)),
+                data: Some(MaterializedTaskValueData::Table(table.to_ipc_bytes()?)),
             }),
-            // TODO: we could use datafusion_proto::bytes::logical_plan_to_bytes here, but that
-            // requires adding datafusion_proto to vegafusion-core deps, as well as passing
-            // datafusion session (maybe empty one?) to unserialize plan
             TaskValue::Plan(_) => Err(VegaFusionError::internal(
-                "Cannot convert Plan TaskValue to protobuf representation",
+                "TaskValue::Plan cannot be serialized to protobuf. Plans are intermediate values that should be materialized to tables using .to_materialized(plan_executor) before serialization.",
             )),
         }
     }
@@ -174,25 +165,6 @@ impl TryFrom<&MaterializedTaskValue> for ProtoMaterializedTaskValue {
     }
 }
 
-impl TryFrom<&ProtoMaterializedTaskValue> for TaskValue {
-    type Error = VegaFusionError;
-
-    fn try_from(value: &ProtoMaterializedTaskValue) -> std::result::Result<Self, Self::Error> {
-        match value.data.as_ref().unwrap() {
-            MaterializedTaskValueData::Table(value) => {
-                Ok(Self::Table(VegaFusionTable::from_ipc_bytes(value)?))
-            }
-            MaterializedTaskValueData::Scalar(value) => {
-                let scalar_table = VegaFusionTable::from_ipc_bytes(value)?;
-                let scalar_rb = scalar_table.to_record_batch()?;
-                let scalar_array = scalar_rb.column(0);
-                let scalar = ScalarValue::try_from_array(scalar_array, 0)?;
-                Ok(Self::Scalar(scalar))
-            }
-        }
-    }
-}
-
 impl TaskGraphValueResponse {
     pub fn deserialize(self) -> Result<Vec<(Variable, Vec<u32>, TaskValue)>> {
         self.response_values

vegafusion-python/tests/test_spark_e2e.py

@@ -40,8 +40,8 @@ def spark():
         .getOrCreate()
     )
 
-    # TODO: this is required for properly handling temporal. We need to check if we can work around different
-    # timezone or if we should require users to setup their Spark sessions to operate in UTC
+    # This is required for properly handling temporal. Due to how Spark handles dates (which doesn't match
+    # proper SQL standart), we have to set timezone to UTC. We mention this is our docs for users.
     session.sql("SET TIME ZONE 'UTC'")
 
     sales_data_df = session.read.parquet(str(SALES_DATA_PATH))

vegafusion-runtime/src/data/tasks.rs

@@ -118,7 +118,7 @@ impl TaskCall for DataUrlTask {
         let url = match self.url.as_ref().unwrap() {
             Url::String(url) => url.clone(),
             Url::Expr(expr) => {
-                let compiled = compile(expr, &config, None)?;
+                let compiled = compile(expr, &config, None).await?;
                 let url_scalar = compiled.eval_to_scalar()?;
                 url_scalar.to_scalar_string()?
             }

vegafusion-runtime/src/expression/compiler/array.rs

@@ -5,14 +5,14 @@ use vegafusion_common::datafusion_common::DFSchema;
 use vegafusion_core::error::Result;
 use vegafusion_core::proto::gen::expression::ArrayExpression;
 
-pub fn compile_array(
+pub async fn compile_array(
     node: &ArrayExpression,
     config: &CompilationConfig,
     schema: &DFSchema,
 ) -> Result<Expr> {
     let mut elements: Vec<Expr> = Vec::new();
     for el in &node.elements {
-        let phys_expr = compile(el, config, Some(schema))?;
+        let phys_expr = compile(el, config, Some(schema)).await?;
         elements.push(phys_expr);
     }
     Ok(make_array(elements))

vegafusion-runtime/src/expression/compiler/binary.rs

@@ -12,14 +12,14 @@ use vegafusion_core::arrow::datatypes::DataType;
 use vegafusion_core::error::Result;
 use vegafusion_core::proto::gen::expression::{BinaryExpression, BinaryOperator};
 
-pub fn compile_binary(
+pub async fn compile_binary(
     node: &BinaryExpression,
     config: &CompilationConfig,
     schema: &DFSchema,
 ) -> Result<Expr> {
     // First, compile argument
-    let lhs = compile(node.left(), config, Some(schema))?;
-    let rhs = compile(node.right(), config, Some(schema))?;
+    let lhs = compile(node.left(), config, Some(schema)).await?;
+    let rhs = compile(node.right(), config, Some(schema)).await?;
 
     let lhs_dtype = data_type(&lhs, schema)?;
     let rhs_dtype = data_type(&rhs, schema)?;

vegafusion-runtime/src/expression/compiler/call.rs

@@ -44,6 +44,7 @@ use vegafusion_common::data::table::VegaFusionTable;
 use vegafusion_common::datafusion_common::DFSchema;
 use vegafusion_common::datatypes::cast_to;
 use vegafusion_common::error::{Result, ResultWithContext, VegaFusionError};
+use vegafusion_core::data::dataset::VegaFusionDataset;
 use vegafusion_core::proto::gen::expression::{
     expression, literal, CallExpression, Expression, Literal,
 };
@@ -96,15 +97,15 @@ pub enum VegaFusionCallable {
     Scale,
 }
 
-pub fn compile_scalar_arguments(
+pub async fn compile_scalar_arguments(
     node: &CallExpression,
     config: &CompilationConfig,
     schema: &DFSchema,
     cast: &Option<DataType>,
 ) -> Result<Vec<Expr>> {
     let mut args: Vec<Expr> = Vec::new();
     for arg in &node.arguments {
-        let compiled_arg = compile(arg, config, Some(schema))?;
+        let compiled_arg = compile(arg, config, Some(schema)).await?;
         let arg_expr = match cast {
             None => compiled_arg,
             Some(dtype) => cast_to(compiled_arg, dtype, schema)?,
@@ -114,7 +115,7 @@ pub fn compile_scalar_arguments(
     Ok(args)
 }
 
-pub fn compile_call(
+pub async fn compile_call(
     node: &CallExpression,
     config: &CompilationConfig,
     schema: &DFSchema,
@@ -127,10 +128,10 @@ pub fn compile_call(
         VegaFusionCallable::Macro(callable) => {
             // Apply macro then recursively compile
             let new_expr = callable(&node.arguments)?;
-            compile(&new_expr, config, Some(schema))
+            compile(&new_expr, config, Some(schema)).await
         }
         VegaFusionCallable::ScalarUDF { udf, cast } => {
-            let args = compile_scalar_arguments(node, config, schema, cast)?;
+            let args = compile_scalar_arguments(node, config, schema, cast).await?;
             Ok(Expr::ScalarFunction(expr::ScalarFunction {
                 func: udf.clone(),
                 args,
@@ -144,18 +145,19 @@ pub fn compile_call(
                         ..
                     }) => {
                         if let Some(dataset) = config.data_scope.get(name) {
-                            if let Some(table) = dataset.as_table() {
-                                let tz_config = config.tz_config.with_context(|| {
-                                    "No local timezone info provided".to_string()
-                                })?;
-                                callee(table, &node.arguments[1..], schema, &tz_config)
-                            } else {
-                                // TODO: If a plan executor is available in config, try to materialize here.
-                                Err(VegaFusionError::compilation(format!(
-                                    "Dataset {} is not materialized as table",
-                                    name
-                                )))
-                            }
+                            let table = match dataset {
+                                VegaFusionDataset::Table { table, .. } => table.clone(),
+                                // Materialize plans on-demand when needed for expression evaluation.
+                                // This allows lazy evaluation of data transformations - plans are only
+                                // executed when the data is actually needed by an expression.
+                                VegaFusionDataset::Plan { plan } => {
+                                    config.plan_executor.execute_plan(plan.clone()).await?
+                                }
+                            };
+                            let tz_config = config
+                                .tz_config
+                                .with_context(|| "No local timezone info provided".to_string())?;
+                            callee(&table, &node.arguments[1..], schema, &tz_config)
                         } else {
                             Err(VegaFusionError::internal(format!(
                                 "No dataset named {}. Available: {:?}",
@@ -179,11 +181,11 @@ pub fn compile_call(
             }
         }
         VegaFusionCallable::Transform(callable) => {
-            let args = compile_scalar_arguments(node, config, schema, &None)?;
+            let args = compile_scalar_arguments(node, config, schema, &None).await?;
             callable(&args, schema)
         }
         VegaFusionCallable::UnaryTransform(callable) => {
-            let mut args = compile_scalar_arguments(node, config, schema, &None)?;
+            let mut args = compile_scalar_arguments(node, config, schema, &None).await?;
             if args.len() != 1 {
                 Err(VegaFusionError::internal(format!(
                     "The {} function requires 1 argument. Received {}",
@@ -195,14 +197,14 @@ pub fn compile_call(
             }
         }
         VegaFusionCallable::LocalTransform(callable) => {
-            let args = compile_scalar_arguments(node, config, schema, &None)?;
+            let args = compile_scalar_arguments(node, config, schema, &None).await?;
             let tz_config = config
                 .tz_config
                 .with_context(|| "No local timezone info provided".to_string())?;
             callable(&tz_config, &args, schema)
         }
         VegaFusionCallable::UtcTransform(callable) => {
-            let args = compile_scalar_arguments(node, config, schema, &None)?;
+            let args = compile_scalar_arguments(node, config, schema, &None).await?;
             let tz_config = RuntimeTzConfig {
                 local_tz: chrono_tz::UTC,
                 default_input_tz: chrono_tz::UTC,