use std::fmt;
use mz_lowertest::MzReflect;
use mz_repr::{ColumnType, Datum, Row, RowArena, ScalarType};
use proptest_derive::Arbitrary;
use serde::{Deserialize, Serialize};
use crate::scalar::func::{stringify_datum, LazyUnaryFunc};
use crate::{EvalError, MirScalarExpr};
#[derive(
Arbitrary, Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash, MzReflect,
)]
pub struct CastListToString {
pub ty: ScalarType,
}
impl LazyUnaryFunc for CastListToString {
fn eval<'a>(
&'a self,
datums: &[Datum<'a>],
temp_storage: &'a RowArena,
a: &'a MirScalarExpr,
) -> Result<Datum<'a>, EvalError> {
let a = a.eval(datums, temp_storage)?;
if a.is_null() {
return Ok(Datum::Null);
}
let mut buf = String::new();
stringify_datum(&mut buf, a, &self.ty)?;
Ok(Datum::String(temp_storage.push_string(buf)))
}
fn output_type(&self, input_type: ColumnType) -> ColumnType {
ScalarType::String.nullable(input_type.nullable)
}
fn propagates_nulls(&self) -> bool {
true
}
fn introduces_nulls(&self) -> bool {
false
}
fn preserves_uniqueness(&self) -> bool {
true
}
fn inverse(&self) -> Option<crate::UnaryFunc> {
None
}
fn is_monotone(&self) -> bool {
false
}
}
impl fmt::Display for CastListToString {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.write_str("listtostr")
}
}
#[derive(
Arbitrary, Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash, MzReflect,
)]
pub struct CastListToJsonb {
pub cast_element: Box<MirScalarExpr>,
}
impl LazyUnaryFunc for CastListToJsonb {
fn eval<'a>(
&'a self,
datums: &[Datum<'a>],
temp_storage: &'a RowArena,
a: &'a MirScalarExpr,
) -> Result<Datum<'a>, EvalError> {
let a = a.eval(datums, temp_storage)?;
if a.is_null() {
return Ok(Datum::Null);
}
let mut row = Row::default();
row.packer().push_list_with(|packer| {
for elem in a.unwrap_list().iter() {
let elem = match self.cast_element.eval(&[elem], temp_storage)? {
Datum::Null => Datum::JsonNull,
d => d,
};
packer.push(elem);
}
Ok::<_, EvalError>(())
})?;
Ok(temp_storage.push_unary_row(row))
}
fn output_type(&self, input_type: ColumnType) -> ColumnType {
ScalarType::Jsonb.nullable(input_type.nullable)
}
fn propagates_nulls(&self) -> bool {
true
}
fn introduces_nulls(&self) -> bool {
false
}
fn preserves_uniqueness(&self) -> bool {
true
}
fn inverse(&self) -> Option<crate::UnaryFunc> {
None
}
fn is_monotone(&self) -> bool {
false
}
}
impl fmt::Display for CastListToJsonb {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.write_str("listtojsonb")
}
}
#[derive(Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash, MzReflect)]
pub struct CastList1ToList2 {
pub return_ty: ScalarType,
pub cast_expr: Box<MirScalarExpr>,
}
impl LazyUnaryFunc for CastList1ToList2 {
fn eval<'a>(
&'a self,
datums: &[Datum<'a>],
temp_storage: &'a RowArena,
a: &'a MirScalarExpr,
) -> Result<Datum<'a>, EvalError> {
let a = a.eval(datums, temp_storage)?;
if a.is_null() {
return Ok(Datum::Null);
}
let mut cast_datums = Vec::new();
for el in a.unwrap_list().iter() {
cast_datums.push(self.cast_expr.eval(&[el], temp_storage)?);
}
Ok(temp_storage.make_datum(|packer| packer.push_list(cast_datums)))
}
fn output_type(&self, input_type: ColumnType) -> ColumnType {
self.return_ty
.without_modifiers()
.nullable(input_type.nullable)
}
fn propagates_nulls(&self) -> bool {
true
}
fn introduces_nulls(&self) -> bool {
false
}
fn preserves_uniqueness(&self) -> bool {
false
}
fn inverse(&self) -> Option<crate::UnaryFunc> {
None
}
fn is_monotone(&self) -> bool {
false
}
}
impl fmt::Display for CastList1ToList2 {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.write_str("list1tolist2")
}
}
#[derive(
Arbitrary, Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash, MzReflect,
)]
pub struct ListLength;
impl LazyUnaryFunc for ListLength {
fn eval<'a>(
&'a self,
datums: &[Datum<'a>],
temp_storage: &'a RowArena,
a: &'a MirScalarExpr,
) -> Result<Datum<'a>, EvalError> {
let a = a.eval(datums, temp_storage)?;
if a.is_null() {
return Ok(Datum::Null);
}
let count = a.unwrap_list().iter().count();
match count.try_into() {
Ok(c) => Ok(Datum::Int32(c)),
Err(_) => Err(EvalError::Int32OutOfRange(count.to_string())),
}
}
fn output_type(&self, input_type: ColumnType) -> ColumnType {
ScalarType::Int32.nullable(input_type.nullable)
}
fn propagates_nulls(&self) -> bool {
true
}
fn introduces_nulls(&self) -> bool {
false
}
fn preserves_uniqueness(&self) -> bool {
false
}
fn inverse(&self) -> Option<crate::UnaryFunc> {
None
}
fn is_monotone(&self) -> bool {
false
}
}
impl fmt::Display for ListLength {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.write_str("list_length")
}
}