use std::cmp::Ordering;
use std::collections::HashSet;
use rudb_common::{Error, LogicalType, Result, Value};
use rudb_vector::{Buffer, Data, Live, Validity, Vector, interleave};
use crate::aggregate::{Accumulator, NOWHERE, finish_run, update_runs, update_scattered};
use crate::compare::order;
use crate::datetime;
use crate::number::integral;
const NULL_PICK: &str = "NULLs are not allowed as list elements in the second input parameter.";
pub(crate) fn before_nulls(
name: &str,
args: &[Value],
returns: &LogicalType,
) -> Option<Result<Value>> {
Some(match (name, args) {
("list_position", [Value::Null, _]) => Ok(Value::Null),
("list_position", [Value::List { values, .. }, needle]) => position(values, needle),
("list_resize" | "list_intersect", [Value::Null, ..]) => Ok(Value::Null),
("list_intersect", [Value::List { .. }, Value::Null]) => listed(Vec::new(), returns),
("list_resize", [Value::List { values, .. }, size, filler @ ..]) => {
resize(values, size, filler.first().unwrap_or(&Value::Null), returns)
}
_ => return None,
})
}
pub(crate) fn value(name: &str, args: &[Value], returns: &LogicalType) -> Option<Result<Value>> {
let list = |values| listed(values, returns);
Some(match (name, args) {
("list_contains", [Value::List { values, .. }, needle]) => {
found(values, needle).map(|at| Value::Boolean(at.is_some()))
}
("list_has_any", [Value::List { values, .. }, Value::List { values: wanted, .. }]) => {
has_any(values, wanted).map(Value::Boolean)
}
("list_has_all", [Value::List { values, .. }, Value::List { values: wanted, .. }]) => {
has_any_missing(values, wanted).map(|missing| Value::Boolean(!missing))
}
("list_distinct", [Value::List { values, .. }]) => distinct(values).and_then(&list),
("list_unique", [Value::List { values, .. }]) => {
distinct(values).map(|kept| Value::UBigInt(kept.len() as u64))
}
("list_intersect", [Value::List { values, .. }, Value::List { values: other, .. }]) => {
intersect(values, other).and_then(&list)
}
("list_where", [Value::List { values, .. }, Value::List { values: mask, .. }]) => {
masked(values, mask).and_then(&list)
}
("list_select", [Value::List { values, .. }, Value::List { values: indexes, .. }]) => {
selected(values, indexes).and_then(&list)
}
("list_sort", [Value::List { values, .. }, spelled @ ..]) => {
let order = spelled.first().map(spelled_order).transpose();
let nulls = spelled.get(1).map(spelled_nulls).transpose();
match (order, nulls) {
(Ok(order), Ok(nulls)) => {
sort(values, order.unwrap_or(false), nulls.unwrap_or(false)).and_then(list)
}
(Err(error), _) | (_, Err(error)) => Err(error),
}
}
("range" | "generate_series", _) => ranged(name == "generate_series", args).and_then(list),
("list_grade_up", [Value::List { values, .. }, spelled @ ..]) => {
let order = spelled.first().map(spelled_order).transpose();
let nulls = spelled.get(1).map(spelled_nulls).transpose();
match (order, nulls) {
(Ok(order), Ok(nulls)) => {
graded(values, order.unwrap_or(false), nulls.unwrap_or(false)).and_then(list)
}
(Err(error), _) | (_, Err(error)) => Err(error),
}
}
("list_reverse_sort", [Value::List { values, .. }, spelled @ ..]) => {
match spelled.first().map(spelled_nulls).transpose() {
Ok(nulls) => sort(values, true, nulls.unwrap_or(false)).and_then(list),
Err(error) => Err(error),
}
}
("list_reverse", [Value::List { values, .. }]) => {
list(values.iter().rev().cloned().collect())
}
("flatten", [Value::List { values, .. }]) => {
let mut flat = Vec::new();
for inner in values {
if let Value::List { values: held, .. } = inner {
flat.extend(held.iter().cloned());
}
}
list(flat)
}
_ => return None,
})
}
fn listed(values: Vec<Value>, returns: &LogicalType) -> Result<Value> {
let LogicalType::List(element) = returns else {
return Err(Error::internal(format!("a list function returning {returns}")));
};
Ok(Value::List { element: (**element).clone(), values })
}
fn ranged(inclusive: bool, args: &[Value]) -> Result<Vec<Value>> {
if let [start, stop, Value::Interval { months, days, micros }] = args {
return stepped(inclusive, start, stop, (*months, *days, *micros));
}
let whole = |value: &Value| {
integral(value)
.and_then(|held| i64::try_from(held).ok())
.ok_or_else(|| Error::internal(format!("a range over a {}", value.logical_type())))
};
let (start, stop, step) = match args {
[stop] => (0, whole(stop)?, 1),
[start, stop] => (whole(start)?, whole(stop)?, 1),
[start, stop, step] => (whole(start)?, whole(stop)?, whole(step)?),
_ => return Err(Error::internal(format!("a range over {} arguments", args.len()))),
};
let count = series_length(start, stop, step, inclusive)?;
let mut at = start;
let mut values = Vec::with_capacity(count);
for _ in 0..count {
values.push(Value::BigInt(at));
at = at.wrapping_add(step);
}
Ok(values)
}
fn series_length(start: i64, stop: i64, step: i64, inclusive: bool) -> Result<usize> {
if step == 0 || (start > stop && step > 0) || (start < stop && step < 0) {
return Ok(0);
}
let apart = stop.abs_diff(start);
let by = step.unsigned_abs();
let (whole, over) = if by == 1 { (apart, false) } else { (apart / by, apart % by != 0) };
let count = u128::from(whole) + u128::from(inclusive || over);
usize::try_from(count).ok().filter(|&count| count <= MAX_SERIES).ok_or_else(too_long)
}
fn stepped(
inclusive: bool,
start: &Value,
stop: &Value,
interval: (i32, i32, i64),
) -> Result<Vec<Value>> {
let moment = |value: &Value| match value {
Value::Timestamp(stamp) | Value::TimestampTz(stamp) => Ok(*stamp),
other => Err(Error::internal(format!("a range from a {}", other.logical_type()))),
};
let mut stamps = Vec::new();
step_stamps(inclusive, moment(start)?, moment(stop)?, interval, &mut stamps)?;
let zoned = matches!(start, Value::TimestampTz(_));
Ok(stamps
.into_iter()
.map(|stamp| if zoned { Value::TimestampTz(stamp) } else { Value::Timestamp(stamp) })
.collect())
}
fn step_stamps(
inclusive: bool,
start: i64,
end: i64,
(months, days, micros): (i32, i32, i64),
out: &mut Vec<i64>,
) -> Result<usize> {
let forward = months > 0 || days > 0 || micros > 0;
let backward = months < 0 || days < 0 || micros < 0;
if forward && backward {
return Err(Error::invalid_input(
"Interval with mix of negative/positive entries not supported",
));
}
if [start, end].iter().any(|&stamp| stamp == i64::MAX || stamp == -i64::MAX) {
return Err(Error::invalid_input("Interval infinite bounds not supported"));
}
let from = out.len();
if !forward && !backward {
return Ok(0);
}
let whole = i128::from(days) * i128::from(datetime::MICROS_PER_DAY) + i128::from(micros);
if let (0, Ok(step)) = (months, i64::try_from(whole)) {
let count = series_length(start, end, step, inclusive)?;
let steps = i64::try_from(count).map_err(|_| too_long())?;
out.reserve(count);
out.extend((0..steps).map(|at| start.wrapping_add(at.wrapping_mul(step))));
return Ok(count);
}
let (months, days, micros) = (i64::from(months), i64::from(days), i128::from(micros));
let mut at = start;
loop {
let past = if forward { at > end } else { at < end };
if past || (at == end && !inclusive) {
return Ok(out.len() - from);
}
if out.len() - from == MAX_SERIES {
return Err(too_long());
}
let next = datetime::shifted_stamp(at, months, days, micros)?;
out.push(at);
at = next;
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Stepping {
Even {
start: i64,
step: i64,
count: usize,
},
Listed(Vec<i64>),
}
impl Stepping {
#[must_use]
pub fn len(&self) -> usize {
match self {
Self::Even { count, .. } => *count,
Self::Listed(stamps) => stamps.len(),
}
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
#[must_use]
pub fn at(&self, position: usize) -> i64 {
match self {
Self::Even { start, step, .. } => {
let steps = i64::try_from(position).unwrap_or(i64::MAX);
start.wrapping_add(steps.wrapping_mul(*step))
}
Self::Listed(stamps) => stamps[position],
}
}
}
pub fn moment_steps(
inclusive: bool,
start: i64,
end: i64,
interval: (i32, i32, i64),
) -> Result<Stepping> {
let (months, days, micros) = interval;
let whole = i128::from(days) * i128::from(datetime::MICROS_PER_DAY) + i128::from(micros);
if let (0, Ok(step)) = (months, i64::try_from(whole)) {
return Ok(Stepping::Even {
start,
step,
count: series_length(start, end, step, inclusive)?,
});
}
let mut stamps = Vec::new();
step_stamps(inclusive, start, end, interval, &mut stamps)?;
Ok(Stepping::Listed(stamps))
}
const MAX_SERIES: usize = u32::MAX as usize;
fn too_long() -> Error {
Error::invalid_input("Lists larger than 2^32 elements are not supported")
}
fn same(left: &Value, right: &Value) -> Result<bool> {
Ok(match (left.is_null(), right.is_null()) {
(true, true) => true,
(true, false) | (false, true) => false,
(false, false) => order(left, right)? == Ordering::Equal,
})
}
fn found(values: &[Value], needle: &Value) -> Result<Option<usize>> {
for (at, value) in values.iter().enumerate() {
if same(value, needle)? {
return Ok(Some(at));
}
}
Ok(None)
}
fn position(values: &[Value], needle: &Value) -> Result<Value> {
Ok(match found(values, needle)? {
Some(at) => Value::Integer(i32::try_from(at + 1).map_err(|_| {
Error::out_of_range(format!("a list position of {} does not fit in INTEGER", at + 1))
})?),
None => Value::Null,
})
}
fn sorted(values: &[Value]) -> Result<Vec<&Value>> {
let mut held: Vec<&Value> = values.iter().filter(|value| !value.is_null()).collect();
let mut failed = None;
held.sort_by(|left, right| {
order(left, right).unwrap_or_else(|error| {
failed.get_or_insert(error);
Ordering::Equal
})
});
match failed {
Some(error) => Err(error),
None => Ok(held),
}
}
fn contains(haystack: &[&Value], needle: &Value) -> Result<bool> {
let mut failed = None;
let hit = haystack
.binary_search_by(|probe| {
order(probe, needle).unwrap_or_else(|error| {
failed.get_or_insert(error);
Ordering::Equal
})
})
.is_ok();
match failed {
Some(error) => Err(error),
None => Ok(hit),
}
}
fn has_any(values: &[Value], wanted: &[Value]) -> Result<bool> {
let haystack = sorted(values)?;
for value in wanted.iter().filter(|value| !value.is_null()) {
if contains(&haystack, value)? {
return Ok(true);
}
}
Ok(false)
}
fn has_any_missing(values: &[Value], wanted: &[Value]) -> Result<bool> {
let haystack = sorted(values)?;
for value in wanted.iter().filter(|value| !value.is_null()) {
if !contains(&haystack, value)? {
return Ok(true);
}
}
Ok(false)
}
fn distinct(values: &[Value]) -> Result<Vec<Value>> {
let mut at: Vec<usize> = (0..values.len()).filter(|&at| !values[at].is_null()).collect();
let mut failed = None;
at.sort_by(|&left, &right| match order(&values[left], &values[right]) {
Ok(Ordering::Equal) => left.cmp(&right),
Ok(ordering) => ordering,
Err(error) => {
failed.get_or_insert(error);
Ordering::Equal
}
});
if let Some(error) = failed {
return Err(error);
}
let mut kept = Vec::with_capacity(at.len());
for (index, &here) in at.iter().enumerate() {
if index == 0 || order(&values[at[index - 1]], &values[here])? != Ordering::Equal {
kept.push(here);
}
}
kept.sort_unstable();
Ok(kept.into_iter().map(|at| values[at].clone()).collect())
}
fn intersect(values: &[Value], other: &[Value]) -> Result<Vec<Value>> {
let haystack = sorted(other)?;
let mut kept = Vec::new();
for value in distinct(values)? {
if contains(&haystack, &value)? {
kept.push(value);
}
}
Ok(kept)
}
fn masked(values: &[Value], mask: &[Value]) -> Result<Vec<Value>> {
let mut kept = Vec::new();
for (at, flag) in mask.iter().enumerate() {
match flag {
Value::Boolean(true) => kept.push(values.get(at).cloned().unwrap_or(Value::Null)),
Value::Boolean(false) => {}
Value::Null => return Err(Error::invalid_input(NULL_PICK)),
other => {
return Err(Error::internal(format!(
"list_where with a {} mask",
other.logical_type()
)));
}
}
}
Ok(kept)
}
fn selected(values: &[Value], indexes: &[Value]) -> Result<Vec<Value>> {
let mut kept = Vec::with_capacity(indexes.len());
for index in indexes {
if index.is_null() {
return Err(Error::invalid_input(NULL_PICK));
}
let picked = index
.as_i64()
.and_then(|index| usize::try_from(index).ok())
.and_then(|index| index.checked_sub(1))
.and_then(|at| values.get(at));
kept.push(picked.cloned().unwrap_or(Value::Null));
}
Ok(kept)
}
fn resize(values: &[Value], size: &Value, filler: &Value, returns: &LogicalType) -> Result<Value> {
let size = match size {
Value::Null => 0,
Value::UBigInt(size) => usize::try_from(*size).map_err(|_| {
Error::out_of_range(format!("a list of {size} elements is too long to build"))
})?,
other => {
return Err(Error::internal(format!("list_resize to a {}", other.logical_type())));
}
};
let element = match returns {
LogicalType::List(element) => element,
_ => &LogicalType::Null,
};
let width = element.physical().size().max(1);
if (size as u128) * (width as u128) > MAX_VECTOR_BYTES {
return Err(Error::out_of_range(format!(
"Cannot resize vector to {size} rows: maximum allowed vector size is 128.0 GiB"
)));
}
let mut kept: Vec<Value> = values.iter().take(size).cloned().collect();
kept.resize(size, filler.clone());
listed(kept, returns)
}
const MAX_VECTOR_BYTES: u128 = 1 << 37;
fn spelled_order(spelled: &Value) -> Result<bool> {
let spelled = spelled.to_string().to_uppercase();
match spelled.as_str() {
"ASC" | "ASCENDING" | "DEFAULT" | "ORDER_DEFAULT" => Ok(false),
"DESC" | "DESCENDING" => Ok(true),
_ => Err(unrecognized(&spelled, "OrderType")),
}
}
fn spelled_nulls(spelled: &Value) -> Result<bool> {
let spelled = spelled.to_string().to_uppercase();
match spelled.as_str() {
"NULLS FIRST" | "NULLS_FIRST" => Ok(true),
"NULLS LAST" | "NULLS_LAST" | "DEFAULT" | "ORDER_DEFAULT" => Ok(false),
_ => Err(unrecognized(&spelled, "OrderByNullType")),
}
}
fn unrecognized(spelled: &str, kind: &str) -> Error {
Error::not_implemented(format!(
"Enum value: unrecognized value \"{spelled}\" for enum \"{kind}\""
))
}
fn sort(values: &[Value], descending: bool, nulls_first: bool) -> Result<Vec<Value>> {
Ok(grade(values, descending, nulls_first)?.into_iter().map(|at| values[at].clone()).collect())
}
fn graded(values: &[Value], descending: bool, nulls_first: bool) -> Result<Vec<Value>> {
grade(values, descending, nulls_first)?
.into_iter()
.map(|at| {
Ok(Value::BigInt(
i64::try_from(at + 1).map_err(|error| Error::internal(error.to_string()))?,
))
})
.collect()
}
fn grade(values: &[Value], descending: bool, nulls_first: bool) -> Result<Vec<usize>> {
let (mut held, nulls): (Vec<usize>, Vec<usize>) =
(0..values.len()).partition(|&at| !values[at].is_null());
let mut failed = None;
held.sort_by(|&left, &right| {
let ordering = order(&values[left], &values[right]).unwrap_or_else(|error| {
failed.get_or_insert(error);
Ordering::Equal
});
if descending { ordering.reverse() } else { ordering }
});
if let Some(error) = failed {
return Err(error);
}
Ok(if nulls_first { [nulls, held].concat() } else { [held, nulls].concat() })
}
pub(crate) fn vectorized<V: AsRef<Vector>>(
name: &str,
args: &[V],
returns: &LogicalType,
rows: usize,
) -> Result<Option<Vector>> {
match (name, args) {
("list_value", [_, ..]) => built(args, returns, rows),
("range" | "generate_series", [_, ..]) => {
series(name == "generate_series", args, returns, rows)
}
("list_aggr", [list, aggregate]) => aggregated(list.as_ref(), aggregate.as_ref(), returns),
("list_reverse", [list]) => reversed(list.as_ref()),
("length" | "array_length", [list]) => counted(list.as_ref()),
("list_distinct", [list]) => deduplicated(false, list.as_ref()),
("list_unique", [list]) => deduplicated(true, list.as_ref()),
("list_contains" | "list_position", [list, needle]) => {
searched(name == "list_position", list.as_ref(), needle.as_ref())
}
("list_sort" | "list_grade_up", [list, spelled @ ..]) => ordered(
list.as_ref(),
spelled.first().map(AsRef::as_ref),
spelled.get(1).map(AsRef::as_ref),
false,
name == "list_grade_up",
),
("list_reverse_sort", [list, spelled @ ..]) => {
ordered(list.as_ref(), None, spelled.first().map(AsRef::as_ref), true, false)
}
_ => Ok(None),
}
}
fn built<V: AsRef<Vector>>(
args: &[V],
returns: &LogicalType,
rows: usize,
) -> Result<Option<Vector>> {
let LogicalType::List(element) = returns else {
return Ok(None);
};
if nested_or_null(element) || args.iter().any(|arg| arg.as_ref().logical_type() != &**element) {
return Ok(None);
}
let pieces = args.iter().map(|arg| arg.as_ref().flatten()).collect::<Result<Vec<_>>>()?;
let width = args.len();
let order: Vec<usize> =
(0..rows).flat_map(|row| (0..width).map(move |at| at * rows + row)).collect();
let child = interleave(element, &pieces, &order)?;
let count = entry(width)?;
let entries = (0..rows).map(|row| Ok((entry(row * width)?, count))).collect::<Result<_>>()?;
Vector::list(entries, child).map(Some)
}
fn series<V: AsRef<Vector>>(
inclusive: bool,
args: &[V],
returns: &LogicalType,
rows: usize,
) -> Result<Option<Vector>> {
if let [start, stop, step] = args {
if step.as_ref().logical_type() == &LogicalType::Interval {
return timed(inclusive, [start.as_ref(), stop.as_ref(), step.as_ref()], returns, rows);
}
}
if args.iter().any(|arg| arg.as_ref().logical_type() != &LogicalType::BigInt) {
return Ok(None);
}
let flat: Vec<Vector> = args.iter().map(|arg| arg.as_ref().flatten()).collect::<Result<_>>()?;
let mut columns = Vec::with_capacity(flat.len());
for vector in &flat {
let Some(Data::Int64(values)) = vector.data() else {
return Ok(None);
};
columns.push((values.as_slice(), vector.validity().live()));
}
let mut runs = Vec::with_capacity(rows);
let mut live = vec![true; rows];
let mut total = 0_usize;
for (row, live) in live.iter_mut().enumerate() {
let mut held = [0_i64; 3];
for (at, (values, valid)) in columns.iter().enumerate() {
match values.get(row) {
Some(&value) if valid.at(row) => held[at] = value,
_ => *live = false,
}
}
let (start, stop, step) = match columns.len() {
1 => (0, held[0], 1),
2 => (held[0], held[1], 1),
_ => (held[0], held[1], held[2]),
};
let count = if *live { series_length(start, stop, step, inclusive)? } else { 0 };
runs.push((start, step, count));
total += count;
}
let mut entries = Vec::with_capacity(rows);
let mut child = Vec::with_capacity(total);
for &(start, step, count) in &runs {
entries.push((entry(child.len())?, entry(count)?));
let steps = i64::try_from(count).map_err(|_| too_long())?;
child.extend((0..steps).map(|at| start.wrapping_add(at.wrapping_mul(step))));
}
let child = Vector::flat(LogicalType::BigInt, Data::Int64(Buffer::from(child)))?;
let validity = Validity::from_iter(rows, |row| live[row]).normalize(rows);
Ok(Some(Vector::list(entries, child)?.with_validity(validity)))
}
fn timed(
inclusive: bool,
args: [&Vector; 3],
returns: &LogicalType,
rows: usize,
) -> Result<Option<Vector>> {
let LogicalType::List(element) = returns else {
return Ok(None);
};
if args[..2].iter().any(|arg| arg.logical_type() != &**element) {
return Ok(None);
}
let [start, stop, step] = args.map(Vector::flatten);
let (start, stop, step) = (start?, stop?, step?);
let (Some(Data::Int64(starts)), Some(Data::Int64(stops)), Some(Data::Interval(steps))) =
(start.data(), stop.data(), step.data())
else {
return Ok(None);
};
let (starts, stops, steps) = (starts.as_slice(), stops.as_slice(), steps.as_slice());
let lives = [start.validity().live(), stop.validity().live(), step.validity().live()];
let mut entries = Vec::with_capacity(rows);
let mut child = Vec::new();
let mut live = vec![true; rows];
for row in 0..rows {
let at = entry(child.len())?;
let held = (starts.get(row), stops.get(row), steps.get(row));
let (Some(&from), Some(&to), Some(&interval)) = held else {
return Ok(None);
};
if lives.iter().any(|live| !live.at(row)) {
entries.push((at, 0));
live[row] = false;
continue;
}
let count = step_stamps(inclusive, from, to, interval, &mut child)?;
entries.push((at, entry(count)?));
}
let child = Vector::flat((**element).clone(), Data::Int64(Buffer::from(child)))?;
let validity = Validity::from_iter(rows, |row| live[row]).normalize(rows);
Ok(Some(Vector::list(entries, child)?.with_validity(validity)))
}
fn counted(list: &Vector) -> Result<Option<Vector>> {
let Some((entries, _)) = list.list_parts() else {
return Ok(None);
};
if !matches!(list.logical_type(), LogicalType::List(_)) {
return Ok(None);
}
let lengths: Vec<i64> = entries.iter().map(|&(_, len)| i64::from(len)).collect();
let answer = Vector::flat(LogicalType::BigInt, Data::Int64(Buffer::from(lengths)))?;
Ok(Some(answer.with_validity(list.validity().clone())))
}
fn aggregated(list: &Vector, aggregate: &Vector, returns: &LogicalType) -> Result<Option<Vector>> {
let (Some((entries, child)), Some(Value::Varchar(aggregate))) =
(list.list_parts(), aggregate.constant_value())
else {
return Ok(None);
};
let live = list.validity().live();
let mut slots = vec![NOWHERE; child.len()];
for (row, &(start, len)) in entries.iter().enumerate() {
if !live.at(row) {
continue;
}
let (start, len) = (start as usize, len as usize);
let Some(held) = slots.get_mut(start..start + len) else {
return Err(Error::internal("a list entry past the end of its child"));
};
if held.iter().any(|&slot| slot != NOWHERE) {
return Ok(None);
}
held.fill(row);
}
let rows = entries.len();
let mut states = vec![Accumulator::new(aggregate, returns)?; rows];
let mut runs: Vec<(usize, usize)> = Vec::new();
for (at, &slot) in slots.iter().enumerate() {
match runs.last_mut() {
Some((held, end)) if *held == slot => *end = at + 1,
_ => runs.push((slot, at + 1)),
}
}
if !update_runs(&mut states, &runs, 1, 0, Some(child), child.len())? {
update_scattered(&mut states, &slots, 1, 0, Some(child), child.len())?;
}
let every: Vec<usize> = (0..rows).collect();
let answer = match finish_run(&states, &every, 1, 0, returns)? {
Some(answer) => answer,
None => {
let values = states.iter().map(Accumulator::finish).collect::<Result<Vec<_>>>()?;
Vector::from_values(returns.clone(), &values)?
}
};
let validity = answer.validity().and(list.validity(), rows);
Ok(Some(answer.with_validity(validity)))
}
fn reversed(list: &Vector) -> Result<Option<Vector>> {
let Some((entries, child)) = list.list_parts() else {
return Ok(None);
};
if !matches!(list.logical_type(), LogicalType::List(_)) {
return Ok(None);
}
let live = list.validity().live();
let mut indices = Vec::with_capacity(child.len());
let mut placed = Vec::with_capacity(entries.len());
for (row, &(start, len)) in entries.iter().enumerate() {
let at = entry(indices.len())?;
if live.at(row) {
indices.extend((start..start + len).rev());
placed.push((at, len));
} else {
placed.push((at, 0));
}
}
let child = child.gather(&indices)?;
Ok(Some(Vector::list(placed, child)?.with_validity(list.validity().clone())))
}
fn searched(position: bool, list: &Vector, needle: &Vector) -> Result<Option<Vector>> {
let (Some((entries, child)), Some(wanted)) = (list.list_parts(), needle.constant_value())
else {
return Ok(None);
};
let plain = plain(child.logical_type());
let Some(wanted) =
integral(wanted).filter(|_| plain && needle.logical_type() == child.logical_type())
else {
return Ok(None);
};
let elements = child.validity().live();
macro_rules! scan {
($($variant:ident),+) => {
match child.data() {
$(Some(Data::$variant(values)) => {
first_places(entries, values.as_slice(), elements, wanted)
})+
_ => return Ok(None),
}
};
}
let found = scan!(Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64);
let rows = list.validity().live();
if position {
let validity = Validity::from_iter(entries.len(), |row| rows.at(row) && found[row] != 0);
let data =
Data::Int32(Buffer::from(found.iter().map(|&place| place as i32).collect::<Vec<_>>()));
let answer = Vector::flat(LogicalType::Integer, data)?;
return Ok(Some(answer.with_validity(validity.normalize(entries.len()))));
}
let data = Data::Bool(Buffer::from(found.iter().map(|&place| place != 0).collect::<Vec<_>>()));
let answer = Vector::flat(LogicalType::Boolean, data)?;
Ok(Some(answer.with_validity(list.validity().clone())))
}
fn ordered(
list: &Vector,
order: Option<&Vector>,
nulls: Option<&Vector>,
reverse: bool,
grade: bool,
) -> Result<Option<Vector>> {
let Some((entries, child)) = list.list_parts() else {
return Ok(None);
};
if !plain(child.logical_type()) || list.validity().count_valid(list.len()) == 0 {
return Ok(None);
}
let spelled = |arg: Option<&Vector>| match arg.map(Vector::constant_value) {
None => Some(None),
Some(Some(value @ Value::Varchar(_))) => Some(Some(value.clone())),
Some(_) => None,
};
let (Some(order), Some(nulls)) = (spelled(order), spelled(nulls)) else {
return Ok(None);
};
let descending = reverse || order.as_ref().map(spelled_order).transpose()?.unwrap_or(false);
let nulls_first = nulls.as_ref().map(spelled_nulls).transpose()?.unwrap_or(false);
let rows = list.validity().live();
let elements = child.validity().live();
macro_rules! permute {
($($variant:ident),+) => {
match child.data() {
$(Some(Data::$variant(values)) => {
let values = values.as_slice();
permutation(entries, rows, elements, nulls_first, |left, right| {
let ordering = values[left as usize].cmp(&values[right as usize]);
if descending { ordering.reverse() } else { ordering }
})?
})+
_ => return Ok(None),
}
};
}
let (placed, indices) = permute!(Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64);
let child = if grade {
let mut places = Vec::with_capacity(indices.len());
for (&(at, len), &(start, _)) in placed.iter().zip(entries) {
let run = &indices[at as usize..(at + len) as usize];
places.extend(run.iter().map(|&index| i64::from(index - start) + 1));
}
Vector::flat(LogicalType::BigInt, Data::Int64(Buffer::from(places)))?
} else {
child.gather(&indices)?
};
Ok(Some(Vector::list(placed, child)?.with_validity(list.validity().clone())))
}
fn deduplicated(unique: bool, list: &Vector) -> Result<Option<Vector>> {
let Some((entries, child)) = list.list_parts() else {
return Ok(None);
};
if !plain(child.logical_type()) {
return Ok(None);
}
let rows = list.validity().live();
let elements = child.validity().live();
macro_rules! keep {
($($variant:ident),+) => {
match child.data() {
$(Some(Data::$variant(values)) => {
let values = values.as_slice();
firsts(entries, rows, elements, |at| i128::from(values[at as usize]))?
})+
_ => return Ok(None),
}
};
}
let (placed, indices) = keep!(Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64);
if unique {
let counts: Vec<u64> = placed.iter().map(|&(_, len)| u64::from(len)).collect();
let answer = Vector::flat(LogicalType::UBigInt, Data::UInt64(Buffer::from(counts)))?;
return Ok(Some(answer.with_validity(list.validity().clone())));
}
let child = child.gather(&indices)?;
Ok(Some(Vector::list(placed, child)?.with_validity(list.validity().clone())))
}
fn firsts(
entries: &[(u32, u32)],
rows: Live<'_>,
elements: Live<'_>,
key: impl Fn(u32) -> i128,
) -> Result<Permuted> {
const SHORT: u32 = 32;
let mut indices = Vec::new();
let mut placed = Vec::with_capacity(entries.len());
let mut kept: Vec<i128> = Vec::new();
let mut seen: HashSet<i128> = HashSet::new();
for (row, &(start, len)) in entries.iter().enumerate() {
let at = entry(indices.len())?;
if !rows.at(row) {
placed.push((at, 0));
continue;
}
let from = indices.len();
kept.clear();
seen.clear();
for index in start..start + len {
if !elements.at(index as usize) {
continue;
}
let value = key(index);
let fresh = if len <= SHORT {
let fresh = !kept.contains(&value);
if fresh {
kept.push(value);
}
fresh
} else {
seen.insert(value)
};
if fresh {
indices.push(index);
}
}
placed.push((at, entry(indices.len() - from)?));
}
Ok((placed, indices))
}
type Permuted = (Vec<(u32, u32)>, Vec<u32>);
fn permutation(
entries: &[(u32, u32)],
rows: Live<'_>,
elements: Live<'_>,
nulls_first: bool,
compare: impl Fn(u32, u32) -> Ordering,
) -> Result<Permuted> {
let mut indices = Vec::new();
let mut placed = Vec::with_capacity(entries.len());
let mut nulls = Vec::new();
for (row, &(start, len)) in entries.iter().enumerate() {
let at = entry(indices.len())?;
if !rows.at(row) {
placed.push((at, 0));
continue;
}
nulls.clear();
let from = indices.len();
for index in start..start + len {
if elements.at(index as usize) {
indices.push(index);
} else {
nulls.push(index);
}
}
indices[from..].sort_by(|&left, &right| compare(left, right));
if nulls_first {
indices.splice(from..from, nulls.iter().copied());
} else {
indices.extend_from_slice(&nulls);
}
placed.push((at, len));
}
Ok((placed, indices))
}
fn plain(ty: &LogicalType) -> bool {
matches!(
ty,
LogicalType::TinyInt
| LogicalType::SmallInt
| LogicalType::Integer
| LogicalType::BigInt
| LogicalType::UTinyInt
| LogicalType::USmallInt
| LogicalType::UInteger
| LogicalType::UBigInt
)
}
fn first_places<T: Copy + PartialEq + TryFrom<i128>>(
entries: &[(u32, u32)],
values: &[T],
elements: Live<'_>,
wanted: i128,
) -> Vec<u32> {
let Ok(wanted) = T::try_from(wanted) else {
return vec![0; entries.len()];
};
let place = |at: Option<usize>| at.map_or(0, |at| at as u32 + 1);
entries
.iter()
.map(|&(start, len)| {
let start = start as usize;
let run = &values[start..start + len as usize];
match elements {
Live::All => place(run.iter().position(|&value| value == wanted)),
_ => place(
run.iter()
.enumerate()
.position(|(at, &value)| value == wanted && elements.at(start + at)),
),
}
})
.collect()
}
fn nested_or_null(element: &LogicalType) -> bool {
matches!(
element,
LogicalType::Null
| LogicalType::List(_)
| LogicalType::Array(..)
| LogicalType::Struct(_)
| LogicalType::Map(..)
| LogicalType::Union(_)
)
}
fn entry(offset: usize) -> Result<u32> {
u32::try_from(offset)
.map_err(|_| Error::out_of_range(format!("a list child of {offset} elements")))
}