From 98cc72dac0aa5dc995dace1a0641db53cb19c2ca Mon Sep 17 00:00:00 2001 From: Jon Mease Date: Tue, 13 Aug 2024 19:45:02 -0400 Subject: [PATCH] Cast pivot value to type of columns --- vegafusion-runtime/src/transform/pivot.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/vegafusion-runtime/src/transform/pivot.rs b/vegafusion-runtime/src/transform/pivot.rs index 778cc209..2e608c73 100644 --- a/vegafusion-runtime/src/transform/pivot.rs +++ b/vegafusion-runtime/src/transform/pivot.rs @@ -10,7 +10,7 @@ use vegafusion_common::arrow::datatypes::DataType; use vegafusion_common::column::{flat_col, unescaped_col}; use vegafusion_common::data::scalar::ScalarValue; use vegafusion_common::data::ORDER_COL; -use vegafusion_common::datatypes::{cast_to, data_type, is_string_datatype}; +use vegafusion_common::datatypes::{cast_to, data_type, is_string_datatype, to_numeric}; use vegafusion_common::error::{Result, ResultWithContext, VegaFusionError}; use vegafusion_common::escape::unescape_field; use vegafusion_core::proto::gen::transforms::{AggregateOp, Pivot}; @@ -156,6 +156,8 @@ async fn pivot_case( return Err(VegaFusionError::internal("Unexpected empty pivot dataset")); } + let schema = dataframe.schema_df()?; + // Process aggregate operation let agg_op: AggregateOp = tx .op @@ -168,7 +170,7 @@ async fn pivot_case( for pivot_val in pivot_vec.iter() { let predicate_expr = unescaped_col(&tx.field).eq(lit(pivot_val.as_str())); - let value_expr = unescaped_col(tx.value.as_str()); + let value_expr = to_numeric(unescaped_col(tx.value.as_str()), &schema)?; let agg_col = when(predicate_expr, value_expr).otherwise(if fill_zero { // Replace null with zero for certain aggregates lit(0) @@ -176,7 +178,7 @@ async fn pivot_case( lit(ScalarValue::Null) })?; - let agg_expr = make_agg_expr_for_col_expr(agg_col, &agg_op, &dataframe.schema_df()?)?; + let agg_expr = make_agg_expr_for_col_expr(agg_col, &agg_op, &schema)?; // Compute pivot column name, replacing null placeholder with "null" let col_name = if pivot_val == NULL_PLACEHOLDER_NAME {