Cast pivot value to type of columns

vega · Aug 13, 2024 · 98cc72d · 98cc72d
1 parent a9e08ad
commit 98cc72d
Showing 1 changed file with 5 additions and 3 deletions.
diff --git 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,15 +170,15 @@ 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)
         } else {
             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 {