Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion benchmarks/random-access-bench/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -229,7 +229,11 @@ fn measurement_name(
format: Format,
remote: bool,
) -> String {
let fmt = format.ext();
let fmt = if format == Format::VortexCompact {
format.name()
} else {
format.ext()
};
let suffix = source_suffix(remote);
match pattern {
Some(p) => format!("random-access/{}/{}/{}-{}", dataset, p.name(), fmt, suffix),
Expand Down
1 change: 1 addition & 0 deletions vortex-bench/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,7 @@ pub const ALLOWED_FORMATS: &[Format] = &[
Format::ArrowIpc,
Format::Parquet,
Format::OnDiskVortex,
Format::VortexCompact,
Format::Lance,
];

Expand Down
36 changes: 26 additions & 10 deletions vortex-bench/src/random_access/take.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,12 @@ use tokio::fs::File;
use vortex::array::Canonical;
use vortex::array::IntoArray;
use vortex::array::VortexSessionExecute;
use vortex::array::stream::ArrayStreamAdapter;
use vortex::array::stream::ArrayStreamExt;
use vortex::buffer::Buffer;
use vortex::file::OpenOptionsSessionExt;
use vortex::file::VortexFile;
use vortex::layout::scan::v2::ScanBuilder;
use vortex::scan::strict_sorted_buffer::StrictSortedBuffer;
use vortex::utils::aliases::hash_map::HashMap;

Expand Down Expand Up @@ -128,11 +130,14 @@ impl RandomAccessor for ArrowIpcRandomAccessor {

/// Random accessor for Vortex format files.
///
/// The file handle is opened at construction time and reused across `take()` calls.
/// The file handle is opened at construction time and reused across `take()` calls. Set
/// `VORTEX_BENCH_SCAN=v1` to compare with V1 using the same open options and materialization.
/// V2 requires files without list layouts.
pub struct VortexRandomAccessor {
name: String,
format: Format,
file: VortexFile,
v2: bool,
}

impl VortexRandomAccessor {
Expand All @@ -150,6 +155,7 @@ impl VortexRandomAccessor {
Ok(Self {
name: name.into(),
format,
v2: !std::env::var("VORTEX_BENCH_SCAN").is_ok_and(|scan| scan == "v1"),
file,
})
}
Expand All @@ -169,6 +175,7 @@ impl VortexRandomAccessor {
Ok(Self {
name: name.into(),
format,
v2: !std::env::var("VORTEX_BENCH_SCAN").is_ok_and(|scan| scan == "v1"),
file,
})
}
Expand All @@ -185,16 +192,25 @@ impl RandomAccessor for VortexRandomAccessor {
}

async fn take(&self, indices: &[u64]) -> anyhow::Result<RandomAccessorRet> {
let indices_buf: Buffer<u64> = Buffer::from(indices.to_vec());
let array = self
.file
.scan()?
.with_row_indices(StrictSortedBuffer::try_new(indices_buf)?)
.into_array_stream()?
.read_all()
.await?;
let indices = StrictSortedBuffer::try_new(Buffer::from(indices.to_vec()))?;
let array = if self.v2 {
let stream = ScanBuilder::new(SESSION.clone(), self.file.scan_file())
.with_row_indices(indices)
.into_stream()?;
ArrayStreamAdapter::new(self.file.dtype().clone(), stream)
.read_all()
.await?
} else {
self.file
.scan()?
.with_row_indices(indices)
.into_array_stream()?
.read_all()
.await?
};

// We canonicalize / decompress for equivalence to Arrow's `RecordBatch`es.
// Keep the historical benchmark contract: canonicalize the outer array, leaving nested
// values lazy. Both scan implementations must perform the same materialization.
let mut ctx = SESSION.create_execution_ctx();
let canonical = array.execute::<Canonical>(&mut ctx)?.into_array();
Ok(RandomAccessorRet::ArrayRef(canonical))
Expand Down
1 change: 1 addition & 0 deletions vortex-layout/src/plan/physical/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ mod optimize;
pub mod optimizer;
pub mod pipeline;
mod plans;
pub(crate) mod selection;
mod typed;
mod vtable;

Expand Down
88 changes: 88 additions & 0 deletions vortex-layout/src/plan/physical/pipeline/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,11 @@ use crate::layouts::chunked::ChunkedLayout;
use crate::layouts::dict::DictLayout;
use crate::layouts::flat::FlatLayout;
use crate::layouts::struct_::StructLayout;
use crate::plan::physical::ConcatPlan;
use crate::plan::physical::EvalPlan;
use crate::plan::physical::Filter;
use crate::plan::physical::PackPlan;
use crate::plan::physical::PlanChildren;
use crate::plan::physical::PlanRef;
use crate::plan::physical::SegmentScan;
use crate::plan::physical::Take;
Expand All @@ -60,6 +62,7 @@ use crate::plan::physical::operator::SourceReady;
use crate::plan::physical::operator::Step;
use crate::plan::physical::operators::PackSource;
use crate::plan::physical::optimize;
use crate::plan::physical::selection::selected;
use crate::segments::SegmentId;
use crate::test::SESSION;

Expand Down Expand Up @@ -834,3 +837,88 @@ fn pack_readiness_tracks_changes_and_consumption(
}
Ok(())
}

#[rstest]
#[case::sparse(2..19, vec![0, 5, 14, 16])]
#[case::contiguous(2..19, (5..10).collect())]
#[case::empty(2..19, vec![])]
fn bound_selection_matches_projection(
#[case] rows: Range<u64>,
#[case] indices: Vec<usize>,
) -> VortexResult<()> {
let mut store = Store::default();
let (plan, expected) = fixture(&mut store)?;
let mask = Mask::from_indices((rows.end - rows.start) as usize, indices);
let plan = selected(plan, rows.clone(), mask.clone());
let result = run(
&store,
&plan,
rows.clone(),
Mask::new_true(mask.len()),
delivery(Delivery::Lifo),
)?;
assert_view(&expected, &rows, &mask, result.arrays)
}

#[test]
fn bound_selection_never_initializes_skipped_chunks() -> VortexResult<()> {
let mut store = Store::default();
let array = PrimitiveArray::from_iter(0..20i32).into_array();
let first = lower(&store.flat(&array.slice(0..5)?)?)?;
let last = lower(&store.flat(&array.slice(15..20)?)?)?;
let children = PlanChildren::lazy(4, move |index| match index {
0 => Ok(first.clone()),
3 => Ok(last.clone()),
_ => Err(vortex_err!("unselected chunk {index} was initialized")),
});
// SAFETY: All four logical chunks have five rows of the same primitive dtype. The two
// unselected initializers deliberately fail to detect discovery as well as compilation.
let plan = unsafe {
ConcatPlan::from_children_unchecked(
array.dtype().clone(),
20,
vec![0, 5, 10, 15].into(),
children,
)
}
.into_plan();
let mask = Mask::from_indices(20, [2, 17]);
let plan = selected(plan, 0..20, mask.clone());
let result = run(
&store,
&plan,
0..20,
Mask::new_true(20),
delivery(Delivery::Fifo),
)?;
assert_eq!(reads(&result.events), 2);
assert_view(&array, &(0..20), &mask, result.arrays)
}

#[test]
fn bound_selection_keeps_whole_dictionary_values() -> VortexResult<()> {
let mut store = Store::default();
let values = PrimitiveArray::from_iter(100..120i32).into_array();
let codes = PrimitiveArray::from_iter((0..20u32).rev()).into_array();
let layout = DictLayout::new(
store.chunked(&values, &[5, 5, 5, 5])?,
store.chunked(&codes, &[5, 5, 5, 5])?,
)
.into_layout();
let mask = Mask::from_indices(20, [1, 18]);
let plan = selected(lower(&layout)?, 0..20, mask);
let result = run(
&store,
&plan,
0..20,
Mask::new_true(20),
delivery(Delivery::Lifo),
)?;
assert_eq!(reads(&result.events), 6, "all values and two code chunks");
assert_arrays_eq!(
join(values.dtype(), result.arrays)?,
PrimitiveArray::from_iter([118, 101]).into_array(),
&mut SESSION.create_execution_ctx()
);
Ok(())
}
25 changes: 16 additions & 9 deletions vortex-layout/src/plan/physical/plans/concat.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ use crate::plan::physical::pipeline::Chain;
use crate::plan::physical::pipeline::Compiler;
use crate::plan::physical::pipeline::Reach;
use crate::plan::physical::pipeline::overlapping;
use crate::plan::physical::selection::selected_slice;
use crate::segments::SegmentId;

/// Concatenates its children row-wise.
Expand Down Expand Up @@ -157,20 +158,26 @@ impl PlanVTable for Concat {
mask: &Mask,
compiler: &mut Compiler<'_>,
) -> VortexResult<Option<Chain>> {
let mut chains = Vec::with_capacity(overlapping(plan, &rows).size_hint().1.unwrap_or(0));
let capacity = if mask.all_true() {
overlapping(plan, &rows).size_hint().1.unwrap_or(0)
} else {
0
};
let mut chains = Vec::with_capacity(capacity);
for (index, start, local) in overlapping(plan, &rows) {
let local_mask = mask.slice(
let Some((span, local_mask)) = selected_slice(
mask,
usize::try_from(local.start - rows.start)?
..usize::try_from(local.end - rows.start)?,
);
// A chunk the mask selects nothing of is never built, so never read.
if local_mask.all_false() {
) else {
continue;
}
};
let chunk = plan.child_required(index)?;
if let Some(chain) =
compiler.compile(&chunk, local.start - start..local.end - start, &local_mask)?
{
if let Some(chain) = compiler.compile(
&chunk,
rows.start + span.start as u64 - start..rows.start + span.end as u64 - start,
&local_mask,
)? {
chains.push(chain);
}
}
Expand Down
6 changes: 3 additions & 3 deletions vortex-layout/src/plan/physical/plans/filter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -103,13 +103,13 @@ impl PlanVTable for Filter {
mask: &Mask,
compiler: &mut Compiler<'_>,
) -> VortexResult<Option<Chain>> {
if mask.all_false() {
return Ok(None);
}
let child = plan.child_plan()?;
if let Some(scan) = child.as_opt::<SegmentScan>() {
return compiler.scan(scan, rows, Some(mask.clone()));
}
if mask.all_false() {
return Ok(None);
}
let len = usize::try_from(rows.end - rows.start)?;
let predicate = plan.dtype().is_boolean();
let chain = compiler.compile(&child, rows, &Mask::new_true(len))?;
Expand Down
Loading
Loading