vegafusion_runtime/expression/compiler/builtin_functions/array/
span.rs

1use datafusion_common::DFSchema;
2use datafusion_expr::{lit, when, Expr, ExprSchemable};
3use datafusion_functions_nested::expr_fn::array_element;
4use datafusion_functions_nested::length::array_length;
5use std::ops::Sub;
6use vegafusion_common::arrow::datatypes::DataType;
7use vegafusion_common::error::{ResultWithContext, VegaFusionError};
8
9pub fn span_transform(args: &[Expr], schema: &DFSchema) -> vegafusion_common::error::Result<Expr> {
10    if args.len() == 1 {
11        let arg = args[0].clone();
12        let dtype = arg
13            .get_type(schema)
14            .with_context(|| format!("Failed to infer type of expression: {arg:?}"))?;
15
16        match dtype {
17            DataType::List(field)
18            | DataType::LargeList(field)
19            | DataType::FixedSizeList(field, _) => {
20                if field.data_type().is_numeric() {
21                    let len = array_length(arg.clone()).cast_to(&DataType::Int32, schema)?;
22                    let first_el = array_element(arg.clone(), lit(1));
23                    let last_el = array_element(arg.clone(), len.clone());
24                    Ok(when(len.eq(lit(0)), lit(0.0)).otherwise(last_el.sub(first_el))?)
25                } else {
26                    Ok(lit(0.0))
27                }
28            }
29            _ => {
30                // Span of non-array is zero
31                Ok(lit(0.0))
32            }
33        }
34    } else {
35        Err(VegaFusionError::parse(format!(
36            "span requires a single argument. Received {} arguments",
37            args.len()
38        )))
39    }
40}