use super::columnar::{
Column, ColumnBatch, Number, TypeTag, frame_bytes_column, gather_fixed_column,
};
use super::zone_map::ColumnStats;
use alloc::vec::Vec;
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum PredicateSupport {
Unsupported,
PruneOnly,
Exact,
}
#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Hash)]
pub enum PredicateApply {
#[default]
Filter,
Prune,
}
#[derive(Debug, Clone)]
pub struct ColumnRangePredicate {
pub column_id: u16,
pub lower: Option<Vec<u8>>,
pub upper: Option<Vec<u8>>,
pub apply: PredicateApply,
}
impl ColumnRangePredicate {
#[must_use]
pub fn support(&self, type_tag: Option<TypeTag>) -> PredicateSupport {
match type_tag {
Some(TypeTag::Bytes | TypeTag::Number(_)) => match self.apply {
PredicateApply::Filter => PredicateSupport::Exact,
PredicateApply::Prune => PredicateSupport::PruneOnly,
},
Some(TypeTag::Fixed(_)) | None => PredicateSupport::Unsupported,
}
}
#[must_use]
pub fn support_in(&self, batch: &ColumnBatch) -> PredicateSupport {
self.support(
batch
.columns
.iter()
.find(|c| c.column_id == self.column_id)
.map(|c| c.type_tag),
)
}
pub(crate) fn bounds(&self, type_tag: TypeTag) -> Option<super::columnar::Bounds<'_>> {
use super::columnar::Bounds;
match type_tag {
TypeTag::Bytes => Some(Bounds::Bytes {
lower: self.lower.as_deref(),
upper: self.upper.as_deref(),
}),
TypeTag::Number(number) => Some(match self.ordinal_span(number) {
Some((lo, hi)) => Bounds::Ordinals { number, lo, hi },
None => Bounds::Nothing,
}),
TypeTag::Fixed(_) => None,
}
}
pub(crate) fn ordinal_span(&self, number: Number) -> Option<(u128, u128)> {
let width = usize::from(number.width());
let max = number.max_ordinal();
let be = |bytes: &[u8]| {
bytes
.iter()
.fold(0u128, |acc, &b| (acc << 8) | u128::from(b))
};
let lo = match self.lower.as_deref() {
None | Some([]) => 0,
Some(l) if l.len() <= width => be(l) << (8 * (width - l.len())),
Some(l) => be(l.get(..width)?).checked_add(1).filter(|&v| v <= max)?,
};
let hi = match self.upper.as_deref() {
None => max,
Some(u) if u.len() >= width => be(u.get(..width)?),
Some(u) => {
let below = be(u).checked_sub(1)?;
let fill = 8 * (width - u.len());
(below << fill) | ((1u128 << fill) - 1)
}
};
(lo <= hi).then_some((lo, hi))
}
#[must_use]
pub fn can_skip_block(&self, stats: &[ColumnStats]) -> bool {
let target = u32::from(self.column_id);
let Some(s) = stats.iter().find(|s| s.column_id == target) else {
return false;
};
if let Some(lo) = &self.lower
&& s.max.as_slice() < lo.as_slice()
{
return true;
}
if let Some(hi) = &self.upper
&& s.min.as_slice() > hi.as_slice()
{
return true;
}
false
}
#[must_use]
pub fn matching_rows(&self, batch: &ColumnBatch) -> Vec<bool> {
let matcher = self.matcher(batch);
(0..batch.row_count)
.map(|row| matcher.matches(row))
.collect()
}
pub(crate) fn matcher<'a>(&'a self, batch: &'a ColumnBatch) -> RowMatcher<'a> {
let rows = batch.row_count as usize;
let Some(col) = batch.columns.iter().find(|c| c.column_id == self.column_id) else {
return RowMatcher::Every;
};
match col.type_tag {
TypeTag::Bytes => RowMatcher::Bytes {
pred: self,
col,
rows,
},
TypeTag::Number(number) => RowMatcher::Number {
col,
number,
span: self.ordinal_span(number),
},
TypeTag::Fixed(_) => RowMatcher::Every,
}
}
}
pub(crate) enum RowMatcher<'a> {
Every,
Bytes {
pred: &'a ColumnRangePredicate,
col: &'a Column,
rows: usize,
},
Number {
col: &'a Column,
number: Number,
span: Option<(u128, u128)>,
},
}
impl RowMatcher<'_> {
pub(crate) fn matches(&self, row: u32) -> bool {
let cell = match *self {
Self::Every => return true,
Self::Bytes { col, rows, .. } => bytes_row(&col.data, rows, row as usize),
Self::Number { col, number, .. } => {
let width = usize::from(number.width());
(row as usize)
.checked_mul(width)
.and_then(|start| col.data.get(start..start.checked_add(width)?))
}
};
cell.is_some_and(|cell| self.matches_cell(row, cell))
}
pub(crate) fn matches_cell(&self, row: u32, cell: &[u8]) -> bool {
let row = row as usize;
match *self {
Self::Every => true,
Self::Bytes { pred, col, .. } => {
row_valid(col, row)
&& pred.lower.as_deref().is_none_or(|lo| cell >= lo)
&& pred.upper.as_deref().is_none_or(|hi| cell <= hi)
}
Self::Number { col, number, span } => {
let Some((lo, hi)) = span else {
return false;
};
if cell.len() != usize::from(number.width()) {
return false;
}
let value = number.ordinal(cell);
lo <= value && value <= hi && row_valid(col, row)
}
}
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Selection {
words: Vec<u64>,
len: u32,
}
impl Selection {
#[must_use]
pub fn none(len: u32) -> Self {
Self {
words: alloc::vec![0; (len as usize).div_ceil(64)],
len,
}
}
#[must_use]
pub fn all(len: u32) -> Self {
let mut selection = Self::none(len);
selection.insert_range(0..len);
selection
}
#[must_use]
pub fn from_bitmap(bitmap: &[u8], len: u32) -> Self {
let mut selection = Self::none(len);
for (word, chunk) in selection.words.iter_mut().zip(bitmap.chunks(8)) {
let mut bytes = [0u8; 8];
for (dst, &src) in bytes.iter_mut().zip(chunk) {
*dst = src;
}
*word = u64::from_le_bytes(bytes);
}
let tail = len % 64;
if tail != 0
&& let Some(last) = selection.words.last_mut()
{
*last &= (1u64 << tail) - 1;
}
selection
}
#[must_use]
pub const fn len(&self) -> u32 {
self.len
}
#[must_use]
pub const fn is_empty(&self) -> bool {
self.len == 0
}
pub fn insert(&mut self, row: u32) {
if row < self.len
&& let Some(word) = self.words.get_mut((row / 64) as usize)
{
*word |= 1 << (row % 64);
}
}
pub fn remove(&mut self, row: u32) {
if let Some(word) = self.words.get_mut((row / 64) as usize) {
*word &= !(1 << (row % 64));
}
}
pub fn insert_range(&mut self, rows: core::ops::Range<u32>) {
let end = rows.end.min(self.len);
let mut row = rows.start;
while row < end {
let bit = row % 64;
let count = (64 - bit).min(end - row);
let mask = if count == 64 {
u64::MAX
} else {
((1u64 << count) - 1) << bit
};
if let Some(word) = self.words.get_mut((row / 64) as usize) {
*word |= mask;
}
row += count;
}
}
#[must_use]
pub fn contains(&self, row: u32) -> bool {
self.words
.get((row / 64) as usize)
.is_some_and(|word| word >> (row % 64) & 1 == 1)
}
pub fn intersect(&mut self, other: &Self) {
let mut theirs = other.words.iter();
for word in &mut self.words {
*word &= theirs.next().copied().unwrap_or(0);
}
}
#[must_use]
pub fn count(&self) -> u32 {
self.words.iter().map(|w| w.count_ones()).sum()
}
pub fn rows(&self) -> impl Iterator<Item = u32> + '_ {
(0u32..).zip(&self.words).flat_map(|(index, &word)| {
let mut rest = word;
core::iter::from_fn(move || {
if rest == 0 {
return None;
}
let bit = rest.trailing_zeros();
rest &= rest - 1;
Some(index * 64 + bit)
})
})
}
#[must_use]
pub fn to_mask(&self) -> Vec<bool> {
(0..self.len).map(|row| self.contains(row)).collect()
}
}
pub fn filter_batch(batch: &ColumnBatch, mask: &[bool]) -> crate::Result<ColumnBatch> {
let mut kept = Vec::with_capacity(mask.len());
for (row, &keep) in mask.iter().enumerate() {
if keep {
kept.push(u32::try_from(row).unwrap_or(u32::MAX));
}
}
take_rows(batch, &kept)
}
pub(crate) fn take_rows(batch: &ColumnBatch, indices: &[u32]) -> crate::Result<ColumnBatch> {
let rows = batch.row_count as usize;
let columns = batch
.columns
.iter()
.map(|c| take_column(c, rows, indices))
.collect::<crate::Result<_>>()?;
let row_count = u32::try_from(indices.len()).unwrap_or(u32::MAX);
Ok(ColumnBatch { row_count, columns })
}
pub(crate) fn take_column(col: &Column, rows: usize, indices: &[u32]) -> crate::Result<Column> {
let data = match col.type_tag.fixed_width() {
Some(width) => {
let width = width as usize;
gather_fixed_column(
width,
indices.len(),
indices.iter().map(|&i| {
let start = (i as usize).checked_mul(width)?;
col.data.get(start..start.checked_add(width)?)
}),
)
}
None => frame_bytes_column(indices.len(), || {
indices
.iter()
.map(|&i| bytes_row(&col.data, rows, i as usize).unwrap_or(&[]))
})?,
};
let validity = col
.validity
.as_ref()
.map(|bits| take_validity(bits, indices));
Ok(Column {
column_id: col.column_id,
type_tag: col.type_tag,
validity,
data,
})
}
pub(crate) fn take_validity(bits: &[u8], indices: &[u32]) -> Vec<u8> {
let mut out = alloc::vec![0u8; indices.len().div_ceil(8)];
for (o, &i) in indices.iter().enumerate() {
let i = i as usize;
let valid = bits.get(i / 8).is_some_and(|b| b & (1u8 << (i % 8)) != 0);
if valid && let Some(byte) = out.get_mut(o / 8) {
*byte |= 1u8 << (o % 8);
}
}
out
}
fn row_valid(col: &Column, row: usize) -> bool {
match &col.validity {
None => true,
Some(bits) => bits
.get(row / 8)
.is_some_and(|byte| byte & (1u8 << (row % 8)) != 0),
}
}
fn bytes_row(data: &[u8], rows: usize, row: usize) -> Option<&[u8]> {
let table_len = rows.checked_add(1)?.checked_mul(4)?;
let off_i = read_u32_le(data.get(row.checked_mul(4)?..)?)? as usize;
let off_next = read_u32_le(data.get(row.checked_add(1)?.checked_mul(4)?..)?)? as usize;
data.get(table_len.checked_add(off_i)?..table_len.checked_add(off_next)?)
}
fn read_u32_le(bytes: &[u8]) -> Option<u32> {
let head = bytes.get(..4)?;
let mut arr = [0u8; 4];
for (dst, &src) in arr.iter_mut().zip(head) {
*dst = src;
}
Some(u32::from_le_bytes(arr))
}
#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
cpufeatures::new!(cpu_avx2_byte_eq, "avx2");
#[must_use]
pub fn byte_eq_mask(batch: &ColumnBatch, column_id: u16, value: u8) -> Vec<bool> {
let rows = batch.row_count as usize;
let Some(col) = batch.columns.iter().find(|c| c.column_id == column_id) else {
return alloc::vec![true; rows];
};
if col.type_tag.fixed_width() != Some(1) {
return alloc::vec![true; rows];
}
byte_eq_dispatch(&col.data, value)
}
fn byte_eq_scalar(data: &[u8], value: u8) -> Vec<bool> {
data.iter().map(|&b| b == value).collect()
}
fn byte_eq_dispatch(data: &[u8], value: u8) -> Vec<bool> {
#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
{
if cpu_avx2_byte_eq::get() {
return unsafe { byte_eq_avx2(data, value) };
}
byte_eq_scalar(data, value)
}
#[cfg(target_arch = "aarch64")]
{
unsafe { byte_eq_neon(data, value) }
}
#[cfg(not(any(target_arch = "x86_64", target_arch = "x86", target_arch = "aarch64")))]
{
byte_eq_scalar(data, value)
}
}
#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
#[target_feature(enable = "avx2")]
unsafe fn byte_eq_avx2(data: &[u8], value: u8) -> Vec<bool> {
#[cfg(target_arch = "x86")]
use core::arch::x86::{
_mm256_cmpeq_epi8, _mm256_loadu_si256, _mm256_movemask_epi8, _mm256_set1_epi8,
};
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::{
_mm256_cmpeq_epi8, _mm256_loadu_si256, _mm256_movemask_epi8, _mm256_set1_epi8,
};
let (chunks, tail) = data.as_chunks::<32>();
let mut out = Vec::with_capacity(data.len());
unsafe {
let needle = _mm256_set1_epi8(i8::from_ne_bytes([value]));
for chunk in chunks {
let v = _mm256_loadu_si256(chunk.as_ptr().cast());
let cmp = _mm256_movemask_epi8(_mm256_cmpeq_epi8(v, needle));
let mask = u32::from_ne_bytes(cmp.to_ne_bytes());
for i in 0..32u32 {
out.push((mask >> i) & 1 == 1);
}
}
}
out.extend(byte_eq_scalar(tail, value));
out
}
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
unsafe fn byte_eq_neon(data: &[u8], value: u8) -> Vec<bool> {
use core::arch::aarch64::{vceqq_u8, vdupq_n_u8, vld1q_u8, vst1q_u8};
let (chunks, tail) = data.as_chunks::<16>();
let mut out = Vec::with_capacity(data.len());
unsafe {
let needle = vdupq_n_u8(value);
for chunk in chunks {
let eq = vceqq_u8(vld1q_u8(chunk.as_ptr()), needle); let mut lanes = [0u8; 16];
vst1q_u8(lanes.as_mut_ptr(), eq);
for &lane in &lanes {
out.push(lane != 0);
}
}
}
out.extend(byte_eq_scalar(tail, value));
out
}
#[cfg(test)]
#[expect(clippy::expect_used, clippy::indexing_slicing, reason = "test code")]
mod tests;