use crate::aggregate::utils::{down_cast_any_ref, ordering_fields};
use crate::expressions::format_state_name;
use crate::{AggregateExpr, LexOrdering, PhysicalExpr, PhysicalSortExpr};
use arrow::array::ArrayRef;
use arrow::datatypes::{DataType, Field};
use arrow_array::Array;
use datafusion_common::{Result, ScalarValue};
use datafusion_expr::Accumulator;
use arrow::compute;
use arrow_array::cast::AsArray;
use datafusion_common::utils::get_row_at_idx;
use std::any::Any;
use std::sync::Arc;
#[derive(Debug)]
pub struct FirstValue {
name: String,
input_data_type: DataType,
order_by_data_types: Vec<DataType>,
expr: Arc<dyn PhysicalExpr>,
ordering_req: LexOrdering,
}
impl FirstValue {
pub fn new(
expr: Arc<dyn PhysicalExpr>,
name: impl Into<String>,
input_data_type: DataType,
ordering_req: LexOrdering,
order_by_data_types: Vec<DataType>,
) -> Self {
Self {
name: name.into(),
input_data_type,
order_by_data_types,
expr,
ordering_req,
}
}
}
impl AggregateExpr for FirstValue {
fn as_any(&self) -> &dyn Any {
self
}
fn field(&self) -> Result<Field> {
Ok(Field::new(&self.name, self.input_data_type.clone(), true))
}
fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
Ok(Box::new(FirstValueAccumulator::try_new(
&self.input_data_type,
&self.order_by_data_types,
)?))
}
fn state_fields(&self) -> Result<Vec<Field>> {
let mut fields = vec![Field::new(
format_state_name(&self.name, "first_value"),
self.input_data_type.clone(),
true,
)];
fields.extend(ordering_fields(
&self.ordering_req,
&self.order_by_data_types,
));
fields.push(Field::new(
format_state_name(&self.name, "is_set"),
DataType::Boolean,
true,
));
Ok(fields)
}
fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
vec![self.expr.clone()]
}
fn order_bys(&self) -> Option<&[PhysicalSortExpr]> {
if self.ordering_req.is_empty() {
None
} else {
Some(&self.ordering_req)
}
}
fn name(&self) -> &str {
&self.name
}
fn reverse_expr(&self) -> Option<Arc<dyn AggregateExpr>> {
let name = if self.name.starts_with("FIRST") {
format!("LAST{}", &self.name[5..])
} else {
format!("LAST_VALUE({})", self.expr)
};
Some(Arc::new(LastValue::new(
self.expr.clone(),
name,
self.input_data_type.clone(),
self.ordering_req.clone(),
self.order_by_data_types.clone(),
)))
}
fn create_sliding_accumulator(&self) -> Result<Box<dyn Accumulator>> {
Ok(Box::new(FirstValueAccumulator::try_new(
&self.input_data_type,
&self.order_by_data_types,
)?))
}
}
impl PartialEq<dyn Any> for FirstValue {
fn eq(&self, other: &dyn Any) -> bool {
down_cast_any_ref(other)
.downcast_ref::<Self>()
.map(|x| {
self.name == x.name
&& self.input_data_type == x.input_data_type
&& self.order_by_data_types == x.order_by_data_types
&& self.expr.eq(&x.expr)
})
.unwrap_or(false)
}
}
#[derive(Debug)]
struct FirstValueAccumulator {
first: ScalarValue,
is_set: bool,
orderings: Vec<ScalarValue>,
}
impl FirstValueAccumulator {
pub fn try_new(data_type: &DataType, ordering_dtypes: &[DataType]) -> Result<Self> {
let orderings = ordering_dtypes
.iter()
.map(ScalarValue::try_from)
.collect::<Result<Vec<_>>>()?;
ScalarValue::try_from(data_type).map(|value| Self {
first: value,
is_set: false,
orderings,
})
}
}
impl Accumulator for FirstValueAccumulator {
fn state(&self) -> Result<Vec<ScalarValue>> {
let mut result = vec![self.first.clone()];
result.extend(self.orderings.iter().cloned());
result.push(ScalarValue::Boolean(Some(self.is_set)));
Ok(result)
}
fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
if !values[0].is_empty() && !self.is_set {
let row = get_row_at_idx(values, 0)?;
self.first = row[0].clone();
self.orderings = row[1..].to_vec();
self.is_set = true;
}
Ok(())
}
fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
let last_idx = states.len() - 1;
let is_set_flags = &states[last_idx];
let flags = is_set_flags.as_boolean();
let mut filtered_first_vals = vec![];
for state in states.iter().take(last_idx - 1) {
filtered_first_vals.push(compute::filter(state, flags)?)
}
self.update_batch(&filtered_first_vals)
}
fn evaluate(&self) -> Result<ScalarValue> {
Ok(self.first.clone())
}
fn size(&self) -> usize {
std::mem::size_of_val(self) - std::mem::size_of_val(&self.first)
+ self.first.size()
+ ScalarValue::size_of_vec(&self.orderings)
- std::mem::size_of_val(&self.orderings)
}
}
#[derive(Debug)]
pub struct LastValue {
name: String,
input_data_type: DataType,
order_by_data_types: Vec<DataType>,
expr: Arc<dyn PhysicalExpr>,
ordering_req: LexOrdering,
}
impl LastValue {
pub fn new(
expr: Arc<dyn PhysicalExpr>,
name: impl Into<String>,
input_data_type: DataType,
ordering_req: LexOrdering,
order_by_data_types: Vec<DataType>,
) -> Self {
Self {
name: name.into(),
input_data_type,
order_by_data_types,
expr,
ordering_req,
}
}
}
impl AggregateExpr for LastValue {
fn as_any(&self) -> &dyn Any {
self
}
fn field(&self) -> Result<Field> {
Ok(Field::new(&self.name, self.input_data_type.clone(), true))
}
fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
Ok(Box::new(LastValueAccumulator::try_new(
&self.input_data_type,
&self.order_by_data_types,
)?))
}
fn state_fields(&self) -> Result<Vec<Field>> {
let mut fields = vec![Field::new(
format_state_name(&self.name, "last_value"),
self.input_data_type.clone(),
true,
)];
fields.extend(ordering_fields(
&self.ordering_req,
&self.order_by_data_types,
));
fields.push(Field::new(
format_state_name(&self.name, "is_set"),
DataType::Boolean,
true,
));
Ok(fields)
}
fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
vec![self.expr.clone()]
}
fn order_bys(&self) -> Option<&[PhysicalSortExpr]> {
if self.ordering_req.is_empty() {
None
} else {
Some(&self.ordering_req)
}
}
fn name(&self) -> &str {
&self.name
}
fn reverse_expr(&self) -> Option<Arc<dyn AggregateExpr>> {
let name = if self.name.starts_with("LAST") {
format!("FIRST{}", &self.name[4..])
} else {
format!("FIRST_VALUE({})", self.expr)
};
Some(Arc::new(FirstValue::new(
self.expr.clone(),
name,
self.input_data_type.clone(),
self.ordering_req.clone(),
self.order_by_data_types.clone(),
)))
}
fn create_sliding_accumulator(&self) -> Result<Box<dyn Accumulator>> {
Ok(Box::new(LastValueAccumulator::try_new(
&self.input_data_type,
&self.order_by_data_types,
)?))
}
}
impl PartialEq<dyn Any> for LastValue {
fn eq(&self, other: &dyn Any) -> bool {
down_cast_any_ref(other)
.downcast_ref::<Self>()
.map(|x| {
self.name == x.name
&& self.input_data_type == x.input_data_type
&& self.order_by_data_types == x.order_by_data_types
&& self.expr.eq(&x.expr)
})
.unwrap_or(false)
}
}
#[derive(Debug)]
struct LastValueAccumulator {
last: ScalarValue,
is_set: bool,
orderings: Vec<ScalarValue>,
}
impl LastValueAccumulator {
pub fn try_new(data_type: &DataType, ordering_dtypes: &[DataType]) -> Result<Self> {
let orderings = ordering_dtypes
.iter()
.map(ScalarValue::try_from)
.collect::<Result<Vec<_>>>()?;
Ok(Self {
last: ScalarValue::try_from(data_type)?,
is_set: false,
orderings,
})
}
}
impl Accumulator for LastValueAccumulator {
fn state(&self) -> Result<Vec<ScalarValue>> {
let mut result = vec![self.last.clone()];
result.extend(self.orderings.clone());
result.push(ScalarValue::Boolean(Some(self.is_set)));
Ok(result)
}
fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
if !values[0].is_empty() {
let row = get_row_at_idx(values, values[0].len() - 1)?;
self.last = row[0].clone();
self.orderings = row[1..].to_vec();
self.is_set = true;
}
Ok(())
}
fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
let last_idx = states.len() - 1;
let is_set_flags = &states[last_idx];
let flags = is_set_flags.as_boolean();
let mut filtered_first_vals = vec![];
for state in states.iter().take(last_idx - 1) {
filtered_first_vals.push(compute::filter(state, flags)?)
}
self.update_batch(&filtered_first_vals)
}
fn evaluate(&self) -> Result<ScalarValue> {
Ok(self.last.clone())
}
fn size(&self) -> usize {
std::mem::size_of_val(self) - std::mem::size_of_val(&self.last)
+ self.last.size()
+ ScalarValue::size_of_vec(&self.orderings)
- std::mem::size_of_val(&self.orderings)
}
}
#[cfg(test)]
mod tests {
use crate::aggregate::first_last::{FirstValueAccumulator, LastValueAccumulator};
use arrow_array::{ArrayRef, Int64Array};
use arrow_schema::DataType;
use datafusion_common::{Result, ScalarValue};
use datafusion_expr::Accumulator;
use std::sync::Arc;
#[test]
fn test_first_last_value_value() -> Result<()> {
let mut first_accumulator =
FirstValueAccumulator::try_new(&DataType::Int64, &[])?;
let mut last_accumulator = LastValueAccumulator::try_new(&DataType::Int64, &[])?;
let ranges: Vec<(i64, i64)> = vec![(0, 10), (1, 11), (2, 13)];
let arrs = ranges
.into_iter()
.map(|(start, end)| {
Arc::new(Int64Array::from((start..end).collect::<Vec<_>>())) as ArrayRef
})
.collect::<Vec<_>>();
for arr in arrs {
first_accumulator.update_batch(&[arr.clone()])?;
last_accumulator.update_batch(&[arr])?;
}
assert_eq!(first_accumulator.evaluate()?, ScalarValue::Int64(Some(0)));
assert_eq!(last_accumulator.evaluate()?, ScalarValue::Int64(Some(12)));
Ok(())
}
}