use alloc::vec::Vec;
use super::projection::{MISTYPED, ProjectedField, conform, conform_lenient};
use super::rows::SourceCursor;
use super::{ColumnarScan, Segment, SegmentCursor, drop_columns, key_in_bounds};
use super::{PredicateTiming, TombstoneSweep, is_operand, value_types};
use super::{Resolved, UNREADABLE_BY_ID};
use crate::table::columnar::{
COL_SEQNO, COL_USER_KEY, COL_VALUE, COL_VALUE_TYPE, Column, ColumnBatch, TypeTag,
bytes_column_row, bytes_column_span, fixed_u64_row, frame_bytes_column,
};
use crate::table::columnar_predicate::{PredicateApply, PredicateSupport, RowMatcher};
use crate::{Error, SeqNo, UserKey};
const TARGET_ROWS: usize = 4_096;
pub(super) const COL_WHOLE_VALUE: u16 = u16::MAX;
struct MergeSource {
cursor: SourceCursor,
whole: bool,
batch: Option<ColumnBatch>,
types: FieldTypes,
stored: StoredCells,
row: u32,
key_col: usize,
seqno_col: usize,
key: core::ops::Range<usize>,
eff: SeqNo,
global: SeqNo,
all_visible: bool,
threshold: SeqNo,
rank: usize,
referenced: bool,
late: LatePayload,
}
#[derive(Default)]
struct LatePayload {
columns: Vec<(u16, u16)>,
eager_ids: Vec<u16>,
lean_ids: Vec<u16>,
eager: bool,
batch_late: bool,
filled: Vec<u16>,
batch_at: Option<crate::table::columnar_cursor::PageAt>,
batch_hit: bool,
page_bytes: u64,
page_useful: u64,
pages: u32,
recent_hits: u8,
}
const DENSITY_WINDOW: u32 = 4;
const DENSITY_MASK: u8 = (1 << DENSITY_WINDOW) - 1;
const JUDGE_COLUMNS: [u16; 2] = [
crate::blob_tree::field_row::CELL_REFS_COLUMN,
COL_WHOLE_VALUE,
];
impl LatePayload {
fn incidental(&mut self) -> u64 {
debug_assert!(self.page_useful <= self.page_bytes);
let incidental = self.page_bytes - self.page_useful;
self.page_bytes = 0;
self.page_useful = 0;
incidental
}
fn spent(&mut self) -> Option<&[u16]> {
if self.columns.is_empty() {
return None;
}
self.recent_hits = (self.recent_hits << 1 | u8::from(self.batch_hit)) & DENSITY_MASK;
self.batch_hit = false;
if self.pages < DENSITY_WINDOW {
self.pages += 1;
if self.pages < DENSITY_WINDOW {
return None;
}
}
let hit = self.recent_hits.count_ones();
if !self.eager && hit * 4 >= DENSITY_WINDOW * 3 {
self.eager = true;
return Some(&self.eager_ids);
}
if self.eager && hit * 4 <= DENSITY_WINDOW {
self.eager = false;
return Some(&self.lean_ids);
}
None
}
}
enum StoredCells {
None,
All,
Some(Vec<u8>),
}
impl StoredCells {
fn of(column: Option<&Column>) -> Self {
match column.map(|c| &c.validity) {
None => Self::None,
Some(None) => Self::All,
Some(Some(bits)) => Self::Some(bits.clone()),
}
}
fn holds(&self, row: u32) -> bool {
match self {
Self::None => false,
Self::All => true,
Self::Some(bits) => bits
.get(row as usize / 8)
.is_some_and(|byte| byte >> (row % 8) & 1 == 1),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum FieldTypes {
AsDeclared,
Mistyped,
MistypedUnderPredicate,
}
impl MergeSource {
fn held_bytes(&self) -> u64 {
self.batch.as_ref().map_or(0, |b| b.data_size() as u64) + self.cursor.held_bytes()
}
}
enum Position {
At,
MustGather,
Exhausted,
}
#[derive(Clone, Copy)]
struct Pick {
source: usize,
row: u32,
eff: SeqNo,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Late {
No,
Values,
Cells,
}
pub(super) struct MergeStream {
sources: Vec<MergeSource>,
dropped: Vec<u16>,
fields: Vec<ProjectedField>,
loose: Vec<u16>,
loose_predicate: Option<u16>,
late: Late,
raw: RawValue,
predicate_column: Option<u16>,
deletes: bool,
rts: TombstoneSweep,
last_key: Option<Vec<u8>>,
pending: Vec<Pick>,
heap: Vec<usize>,
started: bool,
parked: Option<usize>,
}
impl MergeStream {
pub(super) fn open(
scan: &ColumnarScan,
segments: &[Segment],
rts: Vec<(UserKey, UserKey, SeqNo)>,
) -> crate::Result<Self> {
let deletes = scan.resolver.is_some()
|| scan.cells.is_some()
|| segments.iter().any(Segment::records_deletions);
let mut needed = alloc::vec![COL_USER_KEY, COL_SEQNO];
if let Some(pred) = &scan.predicate {
needed.push(pred.column_id);
}
if deletes {
needed.push(COL_VALUE_TYPE);
}
let (augmented, dropped) = scan.augment(&needed);
let fields = scan
.fields
.iter()
.cloned()
.chain(
dropped
.iter()
.copied()
.map(ProjectedField::by_id)
.filter(|field| field.type_tag().is_some()),
)
.collect();
let mut loose: Vec<u16> = scan
.fields
.iter()
.filter(|f| f.type_tag().is_none())
.map(ProjectedField::column_id)
.collect();
let loose_predicate = scan.predicate.as_ref().map(|p| p.column_id).filter(|&id| {
ProjectedField::by_id(id).type_tag().is_none()
&& scan
.fields
.iter()
.all(|f| f.column_id() != id || f.type_tag().is_none())
});
if let Some(id) = loose_predicate
&& !loose.contains(&id)
{
loose.push(id);
}
let share = scan.budget / (segments.len() as u64).max(1);
let sources = segments
.iter()
.map(|seg| {
let SegmentCursor {
mut cursor,
whole,
ids,
} = scan.segment_cursor(seg, &augmented, None, share)?;
let columns = scan.late_payload(seg);
let late = if columns.is_empty() {
LatePayload::default()
} else {
let lean_ids: Vec<u16> = ids
.iter()
.copied()
.filter(|id| columns.iter().all(|&(stored, _)| stored != *id))
.collect();
cursor.set_projection(&lean_ids);
LatePayload {
columns,
eager_ids: ids,
lean_ids,
..LatePayload::default()
}
};
Ok(MergeSource {
cursor,
whole,
late,
batch: None,
types: FieldTypes::AsDeclared,
stored: StoredCells::None,
row: 0,
key_col: 0,
seqno_col: 0,
key: 0..0,
eff: 0,
global: seg.global,
all_visible: seg.visibility == crate::table::SeqnoVisibility::All,
threshold: scan.seqno.saturating_sub(seg.global),
rank: seg.recency_rank,
referenced: false,
})
})
.collect::<crate::Result<Vec<_>>>()?;
let late = if scan.cells.is_some() {
Late::Cells
} else if segments.iter().any(|seg| scan.reads_late(seg)) {
Late::Values
} else {
Late::No
};
Ok(Self {
sources,
dropped,
fields,
loose,
loose_predicate,
late,
raw: if scan.raw_value_read() {
RawValue::Kept
} else {
RawValue::Moved
},
predicate_column: scan.predicate.as_ref().map(|p| p.column_id),
deletes,
rts: TombstoneSweep::new(rts, scan.comparator.as_ref()),
last_key: None,
pending: Vec::new(),
heap: Vec::with_capacity(segments.len()),
started: false,
parked: None,
})
}
pub(super) fn next_batch(
&mut self,
scan: &ColumnarScan,
support: &mut PredicateSupport,
) -> crate::Result<Option<ColumnBatch>> {
let cmp = scan.comparator.clone();
let cmp = cmp.as_ref();
if !self.started {
self.started = true;
for i in 0..self.sources.len() {
self.enter(i, scan, cmp)?;
}
}
if let Some(i) = self.parked.take() {
self.enter(i, scan, cmp)?;
}
loop {
let Some(&winner) = self.heap.first() else {
return self.gather(scan, support);
};
let source = self.sources.get(winner).ok_or(NO_SOURCE)?;
let row = source.row;
let (key, eff) = current(source)?;
let shadowed = self
.last_key
.as_deref()
.is_some_and(|last| cmp.compare(last, key).is_eq());
let keep = if shadowed {
false
} else {
let keep = decide(scan, source, self.deletes, &mut self.rts, key, eff)?;
let last = self.last_key.get_or_insert_with(Vec::new);
last.clear();
last.extend_from_slice(key);
keep
};
let source = self.sources.get_mut(winner).ok_or(NO_SOURCE)?;
source.row += 1;
if keep {
source.referenced = true;
self.pending.push(Pick {
source: winner,
row,
eff,
});
}
match self.position(winner, scan, cmp)? {
Position::At => sift_down(&mut self.heap, 0, &self.sources, cmp)?,
Position::Exhausted => self.leave_top(cmp)?,
Position::MustGather => {
self.leave_top(cmp)?;
self.parked = Some(winner);
let out = self.gather(scan, support)?;
if out.is_some() {
return Ok(out);
}
self.parked = None;
self.enter(winner, scan, cmp)?;
}
}
if self.pending.len() >= TARGET_ROWS
&& let Some(out) = self.gather(scan, support)?
{
return Ok(Some(out));
}
}
}
fn enter(
&mut self,
i: usize,
scan: &ColumnarScan,
cmp: &dyn crate::comparator::UserComparator,
) -> crate::Result<()> {
match self.position(i, scan, cmp)? {
Position::At => {
let at = self.heap.len();
self.heap.push(i);
sift_up(&mut self.heap, at, &self.sources, cmp)
}
Position::Exhausted => Ok(()),
Position::MustGather => Err(Error::InvalidHeader(
"columnar_scan: a merge source entered holding chosen rows",
)),
}
}
fn leave_top(&mut self, cmp: &dyn crate::comparator::UserComparator) -> crate::Result<()> {
self.heap.swap_remove(0);
sift_down(&mut self.heap, 0, &self.sources, cmp)
}
pub(super) fn held_bytes(&self) -> u64 {
self.sources.iter().map(MergeSource::held_bytes).sum()
}
pub(super) fn take_oversized(&mut self) -> u64 {
self.sources
.iter_mut()
.map(|s| s.cursor.take_oversized())
.sum()
}
fn position(
&mut self,
i: usize,
scan: &ColumnarScan,
cmp: &dyn crate::comparator::UserComparator,
) -> crate::Result<Position> {
let mut incidental = 0u64;
let positioned = self.position_source(i, cmp, &mut incidental);
scan.record_payload(0, incidental);
let (position, peak) = positioned?;
if let Some(peak) = peak {
let own = self.sources.get(i).map_or(0, MergeSource::held_bytes);
scan.observe_payload(self.held_bytes() - own + peak);
}
Ok(position)
}
fn position_source(
&mut self,
i: usize,
cmp: &dyn crate::comparator::UserComparator,
incidental: &mut u64,
) -> crate::Result<(Position, Option<u64>)> {
let last_key = self.last_key.as_deref();
let (fields, late, raw, predicate_column) =
(&self.fields, self.late, self.raw, self.predicate_column);
let Some(source) = self.sources.get_mut(i) else {
return Ok((Position::Exhausted, None));
};
let mut loaded: Option<u64> = None;
loop {
if let Some(batch) = &source.batch {
let keys = batch.columns.get(source.key_col).ok_or(MISSING_COLUMN)?;
let seqnos = batch.columns.get(source.seqno_col).ok_or(MISSING_COLUMN)?;
while source.row < batch.row_count {
let local = fixed_u64_row(&seqnos.data, source.row)?;
if source.all_visible || local < source.threshold {
let span = bytes_column_span(&keys.data, batch.row_count, source.row)?;
let shadowed = last_key.is_some_and(|last| {
keys.data
.get(span.clone())
.is_some_and(|key| cmp.compare(last, key).is_eq())
});
if !shadowed {
source.eff =
local
.checked_add(source.global)
.ok_or(Error::InvalidHeader(
"columnar_scan: effective seqno overflow",
))?;
source.key = span;
return Ok((Position::At, loaded));
}
}
source.row += 1;
}
if source.referenced {
return Ok((Position::MustGather, loaded));
}
}
if source.batch.take().is_some() {
*incidental += source.late.incidental();
if let Some(ids) = source.late.spent() {
source.cursor.set_projection(ids);
}
}
match source.cursor.next_located() {
None => return Ok((Position::Exhausted, loaded)),
Some(located) => {
let (batch, at) = located?;
source.late.batch_late = at.is_some()
&& !source.late.columns.is_empty()
&& source.late.columns.iter().all(|&(stored, _)| {
batch.columns.iter().all(|c| c.column_id != stored)
});
source.late.batch_at = at;
source.late.filled.clear();
let batch = carry_whole_value(batch, source.whole, raw);
source.stored = StoredCells::of(
predicate_column
.and_then(|id| batch.columns.iter().find(|c| c.column_id == id)),
);
let (batch, mistyped) = conform_lenient(batch, fields, predicate_column)?;
let batch = if late == Late::No {
batch
} else {
last_whole_value(batch, late == Late::Cells)?
};
if batch
.columns
.iter()
.any(|c| c.column_id == COL_USER_KEY && c.type_tag != TypeTag::Bytes)
{
return Err(Error::InvalidHeader(
"columnar_scan: key column is not a bytes column",
));
}
for column in &batch.columns {
column.validate(batch.row_count)?;
}
let place = |id| {
batch
.columns
.iter()
.position(|c| c.column_id == id)
.ok_or(MISSING_COLUMN)
};
source.key_col = place(COL_USER_KEY)?;
source.seqno_col = place(COL_SEQNO)?;
source.batch = Some(batch);
source.types = if mistyped.is_empty() {
FieldTypes::AsDeclared
} else if predicate_column.is_some_and(|id| mistyped.contains(&id)) {
FieldTypes::MistypedUnderPredicate
} else {
FieldTypes::Mistyped
};
source.row = 0;
let held = source.held_bytes();
loaded = Some(loaded.map_or(held, |peak| peak.max(held)));
}
}
}
}
fn gather(
&mut self,
scan: &ColumnarScan,
support: &mut PredicateSupport,
) -> crate::Result<Option<ColumnBatch>> {
let pending = core::mem::take(&mut self.pending);
let built = if pending.is_empty() {
None
} else {
self.judge_and_build(scan, pending, support)?
};
for source in &mut self.sources {
source.referenced = false;
}
let Some(Built {
merged,
pending,
deferred,
judged,
early,
left_out,
}) = built
else {
return Ok(None);
};
scan.record_gather(&merged);
let (merged, resolved) = if self.late == Late::No {
(merged, Resolved::default())
} else {
let (merged, resolved) = scan.read_late(merged, early)?;
let (merged, mistyped) = conform_lenient(merged, &scan.fields, self.predicate_column)?;
if !mistyped.is_empty() {
return Err(MISTYPED);
}
(merged, resolved)
};
if pending
.iter()
.any(|pick| self.source_types(pick) == Some(FieldTypes::MistypedUnderPredicate))
{
let returned = self.returned_picks(&merged, &pending)?;
let keys = merged
.columns
.iter()
.find(|c| c.column_id == COL_USER_KEY)
.ok_or(MISSING_COLUMN)?;
for (row, pick) in (0..merged.row_count).zip(&returned) {
if self.source_types(pick) == Some(FieldTypes::MistypedUnderPredicate)
&& !resolved.holds(bytes_column_row(&keys.data, merged.row_count, row)?)
{
return Err(MISTYPED);
}
}
}
let judged = judged && !(resolved.any() && self.loose_predicate.is_some());
let mut merged = if early {
merged
} else if judged {
scan.filter_after_dedup(merged, scan.predicate.as_ref(), support)?
} else {
*support = (*support).min(PredicateSupport::Unsupported);
merged
};
let mut waiting: Vec<usize> = pending
.iter()
.zip(&deferred)
.filter(|&(pick, &defer)| {
defer
&& self
.sources
.get(pick.source)
.is_some_and(|s| !s.late.batch_hit)
})
.map(|(pick, _)| pick.source)
.collect();
if !waiting.is_empty() {
waiting.sort_unstable();
waiting.dedup();
let mut hit = Vec::with_capacity(waiting.len());
self.walk_returned(&merged, &pending, |index| {
if deferred.get(index).copied().unwrap_or(false)
&& let Some(pick) = pending.get(index)
&& let Ok(at) = waiting.binary_search(&pick.source)
{
waiting.remove(at);
hit.push(pick.source);
}
!waiting.is_empty()
})?;
for source in hit {
if let Some(source) = self.sources.get_mut(source) {
source.late.batch_hit = true;
}
}
}
if merged.row_count == 0 {
return Ok(None);
}
if pending.iter().any(|pick| self.taken_from_mistyped(pick)) || resolved.any() {
let returned = self.returned_picks(&merged, &pending)?;
let keys = merged
.columns
.iter()
.find(|c| c.column_id == COL_USER_KEY)
.ok_or(MISSING_COLUMN)?;
for (row, pick) in (0..merged.row_count).zip(&returned) {
let key = bytes_column_row(&keys.data, merged.row_count, row)?;
if resolved.unreadable(key) {
return Err(UNREADABLE_BY_ID);
}
if self.taken_from_mistyped(pick) && !resolved.holds(key) {
return Err(MISTYPED);
}
}
}
if !left_out.is_empty() {
self.fill_left_out(&mut merged, &pending, &left_out, &scan.fields)?;
}
drop_columns(&mut merged, &self.dropped);
let merged = conform(merged, &scan.fields)?;
Ok(Some(merged))
}
fn returned_picks<'p>(
&self,
merged: &ColumnBatch,
pending: &'p [Pick],
) -> crate::Result<Vec<&'p Pick>> {
let mut returned = Vec::with_capacity(merged.row_count as usize);
self.walk_returned(merged, pending, |index| {
returned.extend(pending.get(index));
true
})?;
Ok(returned)
}
fn walk_returned(
&self,
merged: &ColumnBatch,
pending: &[Pick],
mut visit: impl FnMut(usize) -> bool,
) -> crate::Result<()> {
let keys = merged
.columns
.iter()
.find(|c| c.column_id == COL_USER_KEY)
.ok_or(MISSING_COLUMN)?;
let mut picks = pending.iter().enumerate();
for row in 0..merged.row_count {
let key = bytes_column_row(&keys.data, merged.row_count, row)?;
let (index, _) = picks
.find(|(_, pick)| {
self.sources
.get(pick.source)
.and_then(|s| {
let batch = s.batch.as_ref()?;
let column = batch.columns.get(s.key_col)?;
bytes_column_row(&column.data, batch.row_count, pick.row).ok()
})
.is_some_and(|chosen| chosen == key)
})
.ok_or(Error::InvalidHeader(
"columnar_scan: a returned row is not among the rows chosen",
))?;
if !visit(index) {
break;
}
}
Ok(())
}
fn taken_from_mistyped(&self, pick: &Pick) -> bool {
self.source_types(pick)
.is_some_and(|types| types != FieldTypes::AsDeclared)
}
fn source_types(&self, pick: &Pick) -> Option<FieldTypes> {
self.sources.get(pick.source).map(|s| s.types)
}
fn fill_left_out(
&self,
merged: &mut ColumnBatch,
pending: &[Pick],
left_out: &[u16],
fields: &[ProjectedField],
) -> crate::Result<()> {
let returned = self.returned_picks(merged, pending)?;
let rows = returned.len();
for &id in left_out {
if !fields.iter().any(|f| f.column_id() == id) {
continue;
}
let mut located: Vec<(&Column, u32, u32)> = Vec::with_capacity(rows);
for pick in &returned {
let Some(batch) = self.sources.get(pick.source).and_then(|s| s.batch.as_ref())
else {
break;
};
let Some(column) = batch.columns.iter().find(|c| c.column_id == id) else {
break;
};
located.push((column, batch.row_count, pick.row));
}
let Some(&(first, _, _)) = located.first() else {
continue;
};
let type_tag = first.type_tag;
if located.len() != rows || located.iter().any(|(c, _, _)| c.type_tag != type_tag) {
continue;
}
let data = if let Some(width) = type_tag.fixed_width() {
let width = usize::from(width);
let mut out = Vec::with_capacity(rows * width);
for &(column, _, row) in &located {
let start = row as usize * width;
out.extend_from_slice(
column
.data
.get(start..start + width)
.ok_or(MISSING_COLUMN)?,
);
}
crate::Slice::from(out)
} else {
let mut cells: Vec<&[u8]> = Vec::with_capacity(rows);
for &(column, count, row) in &located {
cells.push(bytes_column_row(&column.data, count, row)?);
}
frame_bytes_column(rows, || cells.iter().copied())?
};
let validity = located
.iter()
.any(|(c, _, _)| c.validity.is_some())
.then(|| {
let mut bits = alloc::vec![0u8; rows.div_ceil(8)];
for (at, &(column, _, row)) in located.iter().enumerate() {
if column.is_valid(row)
&& let Some(byte) = bits.get_mut(at / 8)
{
*byte |= 1 << (at % 8);
}
}
bits
});
merged.columns.push(Column {
column_id: id,
type_tag,
validity,
data,
});
}
Ok(())
}
fn judge_and_build(
&mut self,
scan: &ColumnarScan,
pending: Vec<Pick>,
support: &mut PredicateSupport,
) -> crate::Result<Option<Built>> {
let layout = self.layout(&pending)?;
let types: Vec<Option<&[u8]>> = if scan.resolver.is_some() {
self.sources
.iter()
.map(|s| s.batch.as_ref().and_then(value_types))
.collect()
} else {
Vec::new()
};
let timing = if layout.judged {
scan.predicate_timing(|| pending.iter().any(|pick| holds_operand(&types, pick)))
} else {
PredicateTiming::AfterValues
};
let mut early = timing == PredicateTiming::BeforeValues;
let (pending, open) = match timing {
PredicateTiming::BeforeValues => {
let pending = self.judge(scan, pending, None, &layout, support)?;
let open = alloc::vec![false; pending.len()];
(pending, open)
}
PredicateTiming::BeforeValuesExceptOperands => {
let pending = self.judge(scan, pending, Some(&types), &layout, support)?;
let open = alloc::vec![false; pending.len()];
(pending, open)
}
PredicateTiming::AfterValues if layout.judged => {
let (pending, open) = self.prejudge(scan, pending);
let (pending, open) = if open.contains(&true) {
self.judge_on_objects(scan, pending, open, &layout)?
} else {
(pending, open)
};
if !open.contains(&true)
&& let Some(pred) = scan.predicate.as_ref()
{
*support = (*support).min(pred.support(layout.type_of(pred.column_id)));
early = true;
}
(pending, open)
}
PredicateTiming::AfterValues => {
let open = alloc::vec![false; pending.len()];
(pending, open)
}
};
if pending.is_empty() {
return Ok(None);
}
let filters_after = scan
.predicate
.as_ref()
.is_some_and(|p| p.apply == PredicateApply::Filter);
let mut deferred = Vec::with_capacity(pending.len());
for (pick, open) in pending.iter().zip(open) {
let Some(source) = self.sources.get_mut(pick.source) else {
deferred.push(false);
continue;
};
let defer =
open && filters_after && !source.late.columns.is_empty() && !source.late.batch_late;
if !defer {
source.late.batch_hit = true;
}
deferred.push(defer);
}
self.fill_late(scan, &pending)?;
scan.record_payload(self.late_useful(&pending)?, 0);
let reread = self.late != Late::No
|| !layout.left_out.is_empty()
|| (!early && scan.predicate.is_some())
|| self
.sources
.iter()
.any(|s| s.referenced && s.types != FieldTypes::AsDeclared);
let skip: &[u16] = if reread { &[] } else { &self.dropped };
let merged = self.build(&pending, &layout, skip)?;
Ok(Some(Built {
merged,
pending,
deferred,
judged: layout.judged,
early,
left_out: layout.left_out,
}))
}
fn prejudge(&self, scan: &ColumnarScan, pending: Vec<Pick>) -> (Vec<Pick>, Vec<bool>) {
let Some(pred) = scan
.predicate
.as_ref()
.filter(|p| p.apply == PredicateApply::Filter)
else {
let open = alloc::vec![true; pending.len()];
return (pending, open);
};
if scan.resolver.is_some() {
let open = alloc::vec![true; pending.len()];
return (pending, open);
}
let matchers: Vec<Option<(RowMatcher<'_>, &StoredCells)>> = self
.sources
.iter()
.map(|s| {
if s.whole || !s.referenced || s.types != FieldTypes::AsDeclared {
return None;
}
let batch = s.batch.as_ref()?;
Some((pred.matcher(batch), &s.stored))
})
.collect();
let mut kept = Vec::with_capacity(pending.len());
let mut open = Vec::with_capacity(pending.len());
for pick in pending {
match matchers.get(pick.source).and_then(Option::as_ref) {
Some((matcher, stored)) if stored.holds(pick.row) => {
if matcher.matches(pick.row) {
kept.push(pick);
open.push(false);
}
}
_ => {
kept.push(pick);
open.push(true);
}
}
}
(kept, open)
}
fn judge_on_objects(
&mut self,
scan: &ColumnarScan,
pending: Vec<Pick>,
open: Vec<bool>,
layout: &Layout,
) -> crate::Result<(Vec<Pick>, Vec<bool>)> {
let Some(cells) = scan.cells.as_ref() else {
return Ok((pending, open));
};
let judging = |source: &MergeSource| {
source.late.batch_late
&& !source.whole
&& source.types == FieldTypes::AsDeclared
&& source
.late
.columns
.iter()
.any(|&(_, carried)| !JUDGE_COLUMNS.contains(&carried))
};
let ask: Vec<bool> = pending
.iter()
.zip(&open)
.map(|(pick, &open)| open && self.sources.get(pick.source).is_some_and(judging))
.collect();
let asked: Vec<Pick> = pending
.iter()
.zip(&ask)
.filter(|&(_, &ask)| ask)
.map(|(pick, _)| *pick)
.collect();
if asked.is_empty() {
return Ok((pending, open));
}
self.fill_late_columns(scan, &asked, Some(&JUDGE_COLUMNS))?;
let predicate = scan.predicate.as_ref().map(|p| p.column_id);
let skip: Vec<u16> = layout
.heads
.iter()
.map(|&(id, _)| id)
.filter(|&id| {
id != COL_USER_KEY
&& id != COL_VALUE_TYPE
&& Some(id) != predicate
&& !JUDGE_COLUMNS.contains(&id)
})
.collect();
let verdicts = super::cells::judge(scan, cells, self.build(&asked, layout, &skip)?)?;
debug_assert_eq!(verdicts.len(), asked.len());
let mut verdicts = verdicts.into_iter();
let mut kept = Vec::with_capacity(pending.len());
let mut still_open = Vec::with_capacity(pending.len());
for ((pick, open), ask) in pending.into_iter().zip(open).zip(ask) {
if ask {
if verdicts.next().unwrap_or(true) {
kept.push(pick);
still_open.push(false);
}
} else {
kept.push(pick);
still_open.push(open);
}
}
Ok((kept, still_open))
}
fn late_useful(&mut self, pending: &[Pick]) -> crate::Result<u64> {
let mut useful = 0u64;
for pick in pending {
let Some(source) = self.sources.get_mut(pick.source) else {
continue;
};
let MergeSource { batch, late, .. } = source;
let Some(batch) = batch.as_ref().filter(|_| late.page_bytes > 0) else {
continue;
};
let mut taken = 0u64;
for &(_, carried) in &late.columns {
let Some(column) = batch
.columns
.iter()
.find(|c| c.column_id == carried && c.is_valid(pick.row))
else {
continue;
};
taken += match column.type_tag.fixed_width() {
Some(width) => u64::from(width),
None => bytes_column_row(&column.data, batch.row_count, pick.row)?.len() as u64,
};
}
late.page_useful += taken;
useful += taken;
}
Ok(useful)
}
fn fill_late(&mut self, scan: &ColumnarScan, pending: &[Pick]) -> crate::Result<()> {
self.fill_late_columns(scan, pending, None)
}
fn fill_late_columns(
&mut self,
scan: &ColumnarScan,
pending: &[Pick],
only: Option<&[u16]>,
) -> crate::Result<()> {
let mut copied = 0usize;
let result = (|| {
for (index, source) in self.sources.iter_mut().enumerate() {
let MergeSource {
cursor,
batch,
late,
types,
..
} = source;
if !late.batch_late || !pending.iter().any(|pick| pick.source == index) {
continue;
}
let (Some(at), Some(table), Some(batch)) =
(late.batch_at.as_ref(), cursor.table(), batch.as_mut())
else {
continue;
};
let stored: Vec<u16> = late
.columns
.iter()
.filter(|&&(_, carried)| {
!late.filled.contains(&carried)
&& only.is_none_or(|only| only.contains(&carried))
})
.map(|&(id, _)| id)
.collect();
if only.is_some() {
late.filled.extend(
late.columns
.iter()
.filter_map(|&(id, carried)| stored.contains(&id).then_some(carried)),
);
} else {
late.batch_late = false;
late.filled.clear();
}
if stored.is_empty() {
continue;
}
for column in table.columns_at(at, &stored, &mut copied)? {
let Some(&(_, carried)) =
late.columns.iter().find(|&&(id, _)| id == column.column_id)
else {
continue;
};
column.validate(batch.row_count)?;
let Some(slot) = batch.columns.iter_mut().find(|c| c.column_id == carried)
else {
continue;
};
if slot.type_tag != column.type_tag {
if *types == FieldTypes::AsDeclared {
*types = FieldTypes::Mistyped;
}
continue;
}
late.page_bytes += column.data.len() as u64;
*slot = Column {
column_id: carried,
..column
};
}
}
Ok(())
})();
scan.record_copied(copied);
result
}
fn judge(
&self,
scan: &ColumnarScan,
pending: Vec<Pick>,
operands: Option<&[Option<&[u8]>]>,
layout: &Layout,
support: &mut PredicateSupport,
) -> crate::Result<Vec<Pick>> {
let Some(pred) = scan.predicate.as_ref() else {
return Ok(pending);
};
if operands.is_none() {
*support = (*support).min(pred.support(layout.type_of(pred.column_id)));
}
if pred.apply != PredicateApply::Filter {
return Ok(pending);
}
let matchers: Vec<Option<RowMatcher<'_>>> = self
.sources
.iter()
.map(|s| {
s.batch
.as_ref()
.filter(|_| s.referenced)
.map(|b| pred.matcher(b))
})
.collect();
let seqno = pred.column_id == COL_SEQNO;
let mut kept = Vec::with_capacity(pending.len());
for pick in pending {
let matcher = matchers
.get(pick.source)
.and_then(Option::as_ref)
.ok_or(NOT_KEPT)?;
let keep = if seqno {
matcher.matches_cell(pick.row, &pick.eff.to_le_bytes())
} else {
matcher.matches(pick.row)
};
if keep || operands.is_some_and(|types| holds_operand(types, &pick)) {
kept.push(pick);
}
}
Ok(kept)
}
fn layout(&self, pending: &[Pick]) -> crate::Result<Layout> {
let batch_of = |pick: &Pick| self.sources.get(pick.source).and_then(|s| s.batch.as_ref());
let template = pending.first().and_then(batch_of).ok_or(NOT_KEPT)?;
let type_of = |batch: &ColumnBatch, id: u16| {
batch
.columns
.iter()
.find(|c| c.column_id == id)
.map(|c| c.type_tag)
};
let mut left_out: Vec<u16> = Vec::new();
for &id in &self.loose {
let expected = type_of(template, id);
for source in self.sources.iter().filter(|s| s.referenced) {
let batch = source.batch.as_ref().ok_or(NOT_KEPT)?;
if type_of(batch, id) != expected {
left_out.push(id);
break;
}
}
}
let judged = self
.loose_predicate
.is_none_or(|id| !left_out.contains(&id));
let uniform = left_out.is_empty();
let kept = |id: u16| !left_out.contains(&id);
let heads: Vec<(u16, TypeTag)> = template
.columns
.iter()
.filter(|c| kept(c.column_id))
.map(|c| (c.column_id, c.type_tag))
.collect();
let mut places: Option<Vec<Vec<usize>>> =
(!uniform).then(|| alloc::vec![Vec::new(); self.sources.len()]);
for (i, source) in self
.sources
.iter()
.enumerate()
.filter(|(_, s)| s.referenced)
{
let batch = source.batch.as_ref().ok_or(NOT_KEPT)?;
let agree = if let Some(places) = places.as_mut() {
let at: Vec<usize> = batch
.columns
.iter()
.enumerate()
.filter(|(_, c)| kept(c.column_id))
.map(|(at, _)| at)
.collect();
let agree = at.len() == heads.len()
&& at.iter().zip(&heads).all(|(&at, &(id, type_tag))| {
batch
.columns
.get(at)
.is_some_and(|c| c.column_id == id && c.type_tag == type_tag)
});
if let Some(slot) = places.get_mut(i) {
*slot = at;
}
agree
} else {
batch.columns.len() == heads.len()
&& batch
.columns
.iter()
.zip(&heads)
.all(|(c, &(id, type_tag))| c.column_id == id && c.type_tag == type_tag)
};
if !agree {
return Err(Error::InvalidHeader(
"columnar_scan: merged segments disagree on their columns",
));
}
}
Ok(Layout {
heads,
places,
left_out,
judged,
})
}
fn build(&self, pending: &[Pick], layout: &Layout, skip: &[u16]) -> crate::Result<ColumnBatch> {
use crate::table::columnar::{Column, frame_bytes_column, gather_fixed_column};
let Layout { heads, places, .. } = layout;
let batch_of = |pick: &Pick| self.sources.get(pick.source).and_then(|s| s.batch.as_ref());
let count = pending.len();
let row_count = u32::try_from(count).map_err(|_| {
Error::InvalidHeader("columnar_scan: merged batch row count exceeds u32")
})?;
let cell = |pick: &Pick, index: usize| {
let at = match &places {
None => index,
Some(places) => *places.get(pick.source)?.get(index)?,
};
batch_of(pick).and_then(|b| Some((b.columns.get(at)?, b.row_count)))
};
let mut columns = Vec::with_capacity(heads.len());
for (index, &(column_id, type_tag)) in heads.iter().enumerate() {
if skip.contains(&column_id) {
continue;
}
let data = if column_id == COL_SEQNO {
let mut out = Vec::with_capacity(count * 8);
for pick in pending {
out.extend_from_slice(&pick.eff.to_le_bytes());
}
crate::Slice::from(out)
} else if let Some(width) = type_tag.fixed_width() {
let width = usize::from(width);
gather_fixed_column(
width,
count,
pending.iter().map(|pick| {
let (column, _) = cell(pick, index)?;
let start = pick.row as usize * width;
column.data.get(start..start + width)
}),
)
} else {
frame_bytes_column(count, || {
pending.iter().map(move |pick| {
cell(pick, index)
.and_then(|(column, rows)| {
bytes_column_row(&column.data, rows, pick.row).ok()
})
.unwrap_or_default()
})
})?
};
let validity = pending
.iter()
.any(|pick| cell(pick, index).is_some_and(|(c, _)| c.validity.is_some()))
.then(|| {
let mut bits = alloc::vec![0u8; count.div_ceil(8)];
for (at, pick) in pending.iter().enumerate() {
let present = cell(pick, index).is_some_and(|(c, _)| {
c.validity.as_deref().is_none_or(|v| {
v.get(pick.row as usize / 8)
.is_some_and(|byte| byte >> (pick.row % 8) & 1 == 1)
})
});
if present && let Some(byte) = bits.get_mut(at / 8) {
*byte |= 1 << (at % 8);
}
}
bits
});
columns.push(Column {
column_id,
type_tag,
validity,
data,
});
}
Ok(ColumnBatch { row_count, columns })
}
}
fn holds_operand(types: &[Option<&[u8]>], pick: &Pick) -> bool {
types
.get(pick.source)
.copied()
.flatten()
.and_then(|cells| cells.get(pick.row as usize))
.is_some_and(|&byte| is_operand(byte))
}
const NOT_KEPT: Error = Error::InvalidHeader("columnar_scan: a chosen row's batch was not kept");
struct Layout {
heads: Vec<(u16, TypeTag)>,
places: Option<Vec<Vec<usize>>>,
left_out: Vec<u16>,
judged: bool,
}
impl Layout {
fn type_of(&self, id: u16) -> Option<TypeTag> {
self.heads
.iter()
.find(|&&(column_id, _)| column_id == id)
.map(|&(_, type_tag)| type_tag)
}
}
struct Built {
merged: ColumnBatch,
pending: Vec<Pick>,
deferred: Vec<bool>,
judged: bool,
early: bool,
left_out: Vec<u16>,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum RawValue {
Moved,
Kept,
}
fn carry_whole_value(mut batch: ColumnBatch, whole: bool, raw: RawValue) -> ColumnBatch {
if !whole {
return batch;
}
if let Some(at) = batch.columns.iter().position(|c| c.column_id == COL_VALUE) {
if raw == RawValue::Kept {
if let Some(value) = batch.columns.get(at).cloned() {
batch.columns.push(Column {
column_id: COL_WHOLE_VALUE,
..value
});
}
} else if let Some(value) = batch.columns.get_mut(at) {
value.column_id = COL_WHOLE_VALUE;
}
}
batch
}
fn last_whole_value(mut batch: ColumnBatch, cells: bool) -> crate::Result<ColumnBatch> {
let carried: &[u16] = if cells {
&[
crate::blob_tree::field_row::CELL_REFS_COLUMN,
COL_WHOLE_VALUE,
]
} else {
&[COL_WHOLE_VALUE]
};
for &id in carried {
let at = batch.columns.iter().position(|c| c.column_id == id);
let column = if let Some(at) = at {
batch.columns.remove(at)
} else {
let rows = batch.row_count as usize;
Column {
column_id: id,
type_tag: TypeTag::Bytes,
validity: Some(alloc::vec![0u8; rows.div_ceil(8)]),
data: frame_bytes_column(rows, || core::iter::repeat_n(&[][..], rows))?,
}
};
batch.columns.push(column);
}
Ok(batch)
}
const NO_SOURCE: Error = Error::InvalidHeader("columnar_scan: merge source out of range");
fn before(
sources: &[MergeSource],
a: usize,
b: usize,
cmp: &dyn crate::comparator::UserComparator,
) -> crate::Result<bool> {
let (sa, sb) = (
sources.get(a).ok_or(NO_SOURCE)?,
sources.get(b).ok_or(NO_SOURCE)?,
);
let ((ka, ea), (kb, eb)) = (current(sa)?, current(sb)?);
Ok(cmp
.compare(ka, kb)
.then_with(|| eb.cmp(&ea))
.then_with(|| sa.rank.cmp(&sb.rank))
.is_lt())
}
fn sift_up(
heap: &mut [usize],
mut at: usize,
sources: &[MergeSource],
cmp: &dyn crate::comparator::UserComparator,
) -> crate::Result<()> {
while at > 0 {
let parent = (at - 1) / 2;
let (&child, &up) = (
heap.get(at).ok_or(NO_SOURCE)?,
heap.get(parent).ok_or(NO_SOURCE)?,
);
if !before(sources, child, up, cmp)? {
break;
}
heap.swap(at, parent);
at = parent;
}
Ok(())
}
fn sift_down(
heap: &mut [usize],
mut at: usize,
sources: &[MergeSource],
cmp: &dyn crate::comparator::UserComparator,
) -> crate::Result<()> {
loop {
let left = 2 * at + 1;
let Some(&l) = heap.get(left) else {
return Ok(());
};
let (mut child, mut least) = (left, l);
if let Some(&r) = heap.get(left + 1)
&& before(sources, r, l, cmp)?
{
(child, least) = (left + 1, r);
}
let &parent = heap.get(at).ok_or(NO_SOURCE)?;
if !before(sources, least, parent, cmp)? {
return Ok(());
}
heap.swap(at, child);
at = child;
}
}
fn current(source: &MergeSource) -> crate::Result<(&[u8], SeqNo)> {
let key = source
.batch
.as_ref()
.and_then(|b| b.columns.get(source.key_col))
.and_then(|c| c.data.get(source.key.clone()))
.ok_or(Error::InvalidHeader(
"columnar_scan: a positioned merge source has no row",
))?;
Ok((key, source.eff))
}
const MISSING_COLUMN: Error = Error::InvalidHeader(
"columnar_scan: a merged group batch is missing a column the merge decoded",
);
fn decide(
scan: &ColumnarScan,
source: &MergeSource,
deletes: bool,
rts: &mut TombstoneSweep,
key: &[u8],
eff: SeqNo,
) -> crate::Result<bool> {
let cmp = scan.comparator.as_ref();
if !scan.range_is_full() && !key_in_bounds(key, &scan.lo, &scan.hi, cmp) {
return Ok(false);
}
if deletes {
let batch = source.batch.as_ref().ok_or(Error::InvalidHeader(
"columnar_scan: a positioned merge source has no batch",
))?;
let byte = *column(batch, COL_VALUE_TYPE)?
.get(source.row as usize)
.ok_or(Error::InvalidHeader(
"columnar_scan: value-type column shorter than the row count",
))?;
let value_type = crate::ValueType::try_from(byte)
.map_err(|()| Error::InvalidTag(("ValueType", byte)))?;
if value_type.is_tombstone() {
return Ok(false);
}
}
Ok(rts.is_empty() || !rts.covers(key, eff, cmp))
}
fn column(batch: &ColumnBatch, column_id: u16) -> crate::Result<&[u8]> {
batch
.columns
.iter()
.find(|c| c.column_id == column_id)
.map(|c| &*c.data)
.ok_or(MISSING_COLUMN)
}