diff --git a/.github/workflows/pco-attribution.yml b/.github/workflows/pco-attribution.yml new file mode 100644 index 00000000000..7e1e165439f --- /dev/null +++ b/.github/workflows/pco-attribution.yml @@ -0,0 +1,81 @@ +name: PCO attribution + +on: + pull_request: + paths: + - ".github/workflows/pco-attribution.yml" + workflow_dispatch: { } + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: false + +jobs: + profile: + runs-on: >- + ${{ github.repository == 'vortex-data/vortex' + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=pco-profile', github.run_id) + || 'ubuntu-latest' }} + timeout-minutes: 25 + steps: + - uses: runs-on/action@v2 + if: github.repository == 'vortex-data/vortex' + with: + sccache: s3 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 + with: + fetch-depth: 0 + - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} + - name: Build frozen adverse candidate and baseline + run: | + git worktree add --detach "$RUNNER_TEMP/pco-profile-candidate" 9f468abe16dd070d2d8eed5df4f45709b50ae375 + git worktree add --detach "$RUNNER_TEMP/pco-profile-base" 910f9933e8e2318c592c1eddcc2dfee8d24c6cd7 + git worktree add --detach "$RUNNER_TEMP/pco-profile-latest" HEAD + mkdir -p "$RUNNER_TEMP/pco-profile-base/vortex-file/examples" + cp vortex-file/examples/pco_access_profile.rs "$RUNNER_TEMP/pco-profile-candidate/vortex-file/examples/" + cp vortex-file/examples/pco_access_profile.rs "$RUNNER_TEMP/pco-profile-base/vortex-file/examples/" + for engine in candidate base latest; do + cp vortex-file/tests/pco_logical_pages_compat.rs "$RUNNER_TEMP/pco-profile-$engine/vortex-file/tests/" + cargo build --manifest-path "$RUNNER_TEMP/pco-profile-$engine/Cargo.toml" --profile bench --locked -p vortex-file --example pco_access_profile --message-format=json-render-diagnostics > "$RUNNER_TEMP/$engine-build.jsonl" + done + - name: Bounded software attribution, separate from timings + shell: bash + run: | + mkdir -p "$RUNNER_TEMP/attribution" + sudo apt-get update -qq + sudo apt-get install -y -qq linux-tools-generic + perf=$(find /usr/lib -path "*linux-tools*/perf" -type f | sort | tail -1) + test -n "$perf" + for engine in base candidate latest; do + executable=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/$engine-build.jsonl") + sha256sum "$executable" > "$RUNNER_TEMP/attribution/$engine-executable.sha256" + git -C "$RUNNER_TEMP/pco-profile-$engine" rev-parse HEAD > "$RUNNER_TEMP/attribution/$engine-source.txt" + nm "$executable" | awk '/vortex_pco/ && /cached_scalar|5probe9scalar_at|scalar_at_once|scalar_at_uncached/ {print $NF}' > "$RUNNER_TEMP/attribution/$engine-symbols.txt" + while IFS= read -r symbol; do + timeout --kill-after=5s 20s objdump --disassemble="$symbol" --no-show-raw-insn "$executable" >> "$RUNNER_TEMP/attribution/$engine-scalar-assembly.txt" + done < "$RUNNER_TEMP/attribution/$engine-symbols.txt" + for request in probe:dense probe:nullable probe:all-null; do + IFS=: read -r mode pattern <<< "$request" + repeats=100000 + stem="$RUNNER_TEMP/attribution/$engine-$mode-$pattern" + echo "Record $engine $mode $pattern" + timeout --signal=INT --kill-after=5s 25s sudo env DEBUGINFOD_URLS= "$perf" record -e cpu-clock -F 499 -o "$stem.data" -- "$executable" "$mode" "$pattern" 8388608 "$repeats" > "$RUNNER_TEMP/$engine-$mode-$pattern.jsonl" || echo "Record exit $?" > "$stem-status.txt" + echo "Flat report $engine $mode $pattern" + timeout --kill-after=5s 20s sudo env DEBUGINFOD_URLS= "$perf" report -i "$stem.data" --stdio --no-children --call-graph none --percent-limit 0.5 > "$stem-flat.txt" || echo "Flat report exit $?" >> "$stem-status.txt" + sudo rm -f "$stem.data" + done + done + printf '%s\n' 'Frozen 9f adverse candidate, develop910f, and current CI merge source. Software flat samples and exact linked symbol assembly only; profiled times are not uninstrumented latency controls. Startup/fixture compression is included. Public stock/metadata byte attribution is retained in run40119098145; it is not repeated here.' > "$RUNNER_TEMP/attribution/README.txt" + - uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 + if: always() + with: + name: pco-attribution + path: | + ${{ runner.temp }}/attribution + !${{ runner.temp }}/attribution/*.data + if-no-files-found: ignore diff --git a/.github/workflows/pco-compat.yml b/.github/workflows/pco-compat.yml new file mode 100644 index 00000000000..1d9a16e84d8 --- /dev/null +++ b/.github/workflows/pco-compat.yml @@ -0,0 +1,100 @@ +name: PCO selective access compatibility + +on: + pull_request: + paths: + - "encodings/pco/**" + - "vortex-file/src/lib.rs" + - "vortex-file/tests/pco_logical_pages_compat.rs" + - ".github/workflows/pco-compat.yml" + workflow_dispatch: { } + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + pco-compat: + runs-on: >- + ${{ github.repository == 'vortex-data/vortex' + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=pco-compat', github.run_id) + || 'ubuntu-latest' }} + timeout-minutes: 45 + steps: + - uses: runs-on/action@v2 + if: github.repository == 'vortex-data/vortex' + with: + sccache: s3 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 + with: + fetch-depth: 0 + - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} + - name: PCO public semantics + run: cargo test --profile ci --locked -p vortex-pco + - name: Generate new files with current public writer + env: + PCO_COMPAT_MODE: write + PCO_COMPAT_DIR: ${{ runner.temp }}/pco-new-files + run: cargo test --profile ci --locked -p vortex-file --test pco_logical_pages_compat + - name: Read new files using unchanged pre-change reader; generate legacy files + env: + PCO_COMPAT_DIR: ${{ runner.temp }}/pco-new-files + run: | + git worktree add --detach "$RUNNER_TEMP/pco-old-reader" 910f9933e8e2318c592c1eddcc2dfee8d24c6cd7 + # Only the public test harness is added. The old reader/writer/library sources remain + # exactly the pinned pre-change revision and do not know the new metadata field. + cp vortex-file/tests/pco_logical_pages_compat.rs "$RUNNER_TEMP/pco-old-reader/vortex-file/tests/" + PCO_COMPAT_MODE=read cargo test --manifest-path "$RUNNER_TEMP/pco-old-reader/Cargo.toml" --profile ci --locked -p vortex-file --test pco_logical_pages_compat + PCO_COMPAT_MODE=write PCO_COMPAT_DIR="$RUNNER_TEMP/pco-old-files" cargo test --manifest-path "$RUNNER_TEMP/pco-old-reader/Cargo.toml" --profile ci --locked -p vortex-file --test pco_logical_pages_compat + - name: Read legacy files using current reader + env: + PCO_COMPAT_MODE: read + PCO_COMPAT_DIR: ${{ runner.temp }}/pco-old-files + run: cargo test --profile ci --locked -p vortex-file --test pco_logical_pages_compat + - name: Record actual public file sizes + run: | + mkdir -p "$RUNNER_TEMP/pco-resource-results" + wc -c "$RUNNER_TEMP"/pco-{new,old}-files/*.vortex > "$RUNNER_TEMP/pco-resource-results/public-file-sizes.txt" + - name: Build matching optimized public resource harnesses + run: | + mkdir -p "$RUNNER_TEMP/pco-old-reader/vortex-file/examples" + cp vortex-file/examples/pco_access_profile.rs "$RUNNER_TEMP/pco-old-reader/vortex-file/examples/" + cargo build --profile bench --locked -p vortex-file --example pco_access_profile --message-format=json-render-diagnostics > "$RUNNER_TEMP/pco-candidate-build.jsonl" + cargo build --manifest-path "$RUNNER_TEMP/pco-old-reader/Cargo.toml" --profile bench --locked -p vortex-file --example pco_access_profile --message-format=json-render-diagnostics > "$RUNNER_TEMP/pco-baseline-build.jsonl" + - name: Matched first and warmed resource controls + shell: bash + run: | + mkdir -p "$RUNNER_TEMP/pco-resource-results" + printf '%s\n' 'Wall timings exclude accounting and post-timing scalar inspection.' 'open-only/open-first include public array-plugin metadata decode/validation/index preparation with compressed buffers and validity children already available; no file IO.' 'probe-first starts with prepared metadata/index; probe retains touched pages across iterations.' 'CPU ticks are OS process user+system ticks; sub-tick readings are resolution-limited, not zero CPU.' 'RSS snapshots may miss transient peaks. GNU time maximum RSS is whole-process peak including fixtures/compression and post-timing inspection, not operation-only peak.' > "$RUNNER_TEMP/pco-resource-results/README.txt" + getconf CLK_TCK > "$RUNNER_TEMP/pco-resource-results/clock-ticks-per-second.txt" + git rev-parse HEAD > "$RUNNER_TEMP/pco-resource-results/candidate-source.txt" + git -C "$RUNNER_TEMP/pco-old-reader" rev-parse HEAD > "$RUNNER_TEMP/pco-resource-results/baseline-source.txt" + candidate=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/pco-candidate-build.jsonl") + baseline=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/pco-baseline-build.jsonl") + sha256sum "$candidate" "$baseline" > "$RUNNER_TEMP/pco-resource-results/executables.sha256" + for mode in write full slice slice-near-full take take-large take-repeated take-mixed take-narrow take-spans take-tiny take-contiguous take-range ranges open-only open-first probe-first probe; do + for pattern in dense nullable all-null long-null compressed nullable-legacy; do + # ABBA process order; every process retains all eight iterations including cold. + order=0 + for engine in baseline candidate candidate baseline; do + order=$((order + 1)) + executable="$candidate" + if [ "$engine" = baseline ]; then + executable="$baseline" + fi + stem="$RUNNER_TEMP/pco-resource-results/$mode-$pattern-$order-$engine" + /usr/bin/time -v "$executable" "$mode" "$pattern" 8388608 8 > "$stem.jsonl" 2> "$stem.time" + done + done + done + - uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 + if: always() + with: + name: pco-resource-results + path: ${{ runner.temp }}/pco-resource-results + if-no-files-found: ignore diff --git a/encodings/pco/src/access.rs b/encodings/pco/src/access.rs new file mode 100644 index 00000000000..bf3a488d172 --- /dev/null +++ b/encodings/pco/src/access.rs @@ -0,0 +1,459 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Operation-owned selected-page values and validity. No decoded state lives in array metadata. + +use std::any::Any; +use std::collections::BTreeMap; +use std::ops::Range; +use std::sync::OnceLock; + +use pco::data_types::NumberType; +use pco::match_number_enum; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::half; +use vortex_array::match_each_native_ptype; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_mask::Mask; + +use crate::Pco; +use crate::PcoArrayExt; +use crate::array::number_type_from_ptype; +use crate::page_index::PageIndex; + +const RANK_STRIDE: usize = 512; + +#[derive(Default)] +pub(crate) struct AccessState { + legacy: Option, + validated_all_null: bool, + decoder: Option>, + pages: BTreeMap, + last_page: Option<(usize, CachedPage)>, +} + +struct CachedPage { + rows: Option>, + values: Range, + mask: Option, + decoded: Option, +} + +struct RankMask { + mask: Mask, + ranks: OnceLock>, +} + +impl RankMask { + fn new(mask: Mask) -> Self { + Self { + mask, + ranks: OnceLock::new(), + } + } + + fn boundary_rank(&self, row: usize) -> usize { + match &self.mask { + Mask::AllTrue(_) => row, + Mask::AllFalse(_) => 0, + Mask::Values(values) => { + let bits = values.bit_buffer(); + if row <= bits.len() - row { + bits.count_range(0, row) + } else { + self.mask.true_count() - bits.count_range(row, bits.len()) + } + } + } + } + + fn rank(&self, row: usize) -> usize { + match &self.mask { + Mask::AllTrue(_) => row, + Mask::AllFalse(_) => 0, + Mask::Values(values) => { + if row == self.mask.len() { + return self.mask.true_count(); + } + let ranks = self.ranks.get_or_init(|| { + let bits = values.bit_buffer(); + let mut ranks = Vec::with_capacity(bits.len().div_ceil(RANK_STRIDE)); + let mut count = 0; + for start in (0..bits.len()).step_by(RANK_STRIDE) { + ranks.push(count); + count += bits.count_range(start, (start + RANK_STRIDE).min(bits.len())); + } + ranks + }); + let block = row / RANK_STRIDE; + ranks[block] + values.bit_buffer().count_range(block * RANK_STRIDE, row) + } + } + } + + fn value_index(&self, row: usize) -> Option { + match &self.mask { + Mask::AllTrue(_) => Some(row), + Mask::AllFalse(_) => None, + Mask::Values(values) => { + let bits = values.bit_buffer(); + bits.value(row).then(|| { + if INDEXED { + self.rank(row) + } else { + self.boundary_rank(row) + } + }) + } + } + } +} + +impl AccessState { + /// Read the last fully validated decoded page directly. A miss preserves all + /// operation-owned cache state and follows the ordinary fallible resolver. + #[inline] + pub fn cached_scalar(&self, array: ArrayView<'_, Pco>, row: usize) -> Option { + if row >= array.unsliced_n_rows() { + return None; + } + let (_, cached) = self.last_page.as_ref()?; + let decoded = cached.decoded.as_ref()?; + let value = if let Some(rows) = &cached.rows { + if !rows.contains(&row) { + return None; + } + cached.mask.as_ref()?.value_index::(row - rows.start) + } else { + let value = self.legacy.as_ref()?.value_index::(row); + if let Some(value) = value { + if !cached.values.contains(&value) { + return None; + } + Some(value - cached.values.start) + } else { + None + } + }; + Some(match value { + None => Scalar::null(array.dtype().clone()), + Some(value) => match_each_native_ptype!(decoded.ptype(), |T| { + Scalar::primitive(decoded.as_slice::()[value], array.dtype().nullability()) + }), + }) + } + + pub fn validated_all_null(&self) -> bool { + self.validated_all_null + } + + fn cached(&self, page: usize) -> Option<&CachedPage> { + if let Some((last, cached)) = &self.last_page + && *last == page + { + Some(cached) + } else { + self.pages.get(&page) + } + } + + fn cached_mut(&mut self, page: usize) -> Option<&mut CachedPage> { + if self.last_page.as_ref().map(|(last, _)| *last) != Some(page) { + let cached = self.pages.remove(&page)?; + self.insert(page, cached); + } + self.last_page.as_mut().map(|(_, cached)| cached) + } + + fn insert(&mut self, page: usize, cached: CachedPage) { + if let Some((previous, cached)) = self.last_page.replace((page, cached)) + && RETAIN_PAGES + { + self.pages.insert(previous, cached); + } + } + + fn logical_page(&self, index: &PageIndex, row: usize) -> usize { + if let Some((last, _)) = &self.last_page + && index + .pages + .get(*last) + .and_then(|page| page.rows.as_ref()) + .is_some_and(|range| range.contains(&row)) + { + return *last; + } + index + .pages + .partition_point(|page| page.rows.as_ref().is_some_and(|range| range.end <= row)) + } + + fn compact_page(&self, index: &PageIndex, value: usize) -> usize { + if let Some((last, _)) = &self.last_page + && index + .pages + .get(*last) + .is_some_and(|page| page.values.contains(&value)) + { + return *last; + } + index.pages.partition_point(|page| page.values.end <= value) + } + + /// Resolve an unsliced logical row to a physical page and compact within-page index. + pub fn resolve( + &mut self, + array: ArrayView<'_, Pco>, + row: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + self.resolve_impl::(array, row, ctx) + } + + fn resolve_impl( + &mut self, + array: ArrayView<'_, Pco>, + row: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + vortex_ensure!( + row < array.unsliced_n_rows(), + "Pco row exceeds array bounds" + ); + let index = array.page_index()?; + if index.logical && !index.pages.is_empty() { + let page_idx = self.logical_page(index, row); + let page = index + .pages + .get(page_idx) + .ok_or_else(|| vortex_err!("Pco logical page index does not cover row"))?; + let rows = page + .rows + .as_ref() + .ok_or_else(|| vortex_err!("Pco logical page range is missing"))?; + if self.cached(page_idx).is_none() { + let mask = array + .unsliced_validity() + .slice(rows.clone())? + .execute_mask(rows.len(), ctx)?; + vortex_ensure_eq!( + mask.true_count(), + page.values.len(), + "Pco logical page validity count does not match stored values" + ); + self.insert( + page_idx, + CachedPage { + rows: Some(rows.clone()), + values: page.values.clone(), + mask: Some(RankMask::new(mask)), + decoded: None, + }, + ); + } + return Ok(self + .cached_mut(page_idx) + .ok_or_else(|| vortex_err!("Pco selected page cache is missing"))? + .mask + .as_ref() + .and_then(|mask| mask.value_index::(row - rows.start)) + .map(|value| (page_idx, value))); + } + + // Legacy or unknown optional metadata: execute/rank the full validity once per operation, + // never once per requested index. Empty physical arrays also take this checked path. + if self.legacy.is_none() { + let mask = array + .unsliced_validity() + .execute_mask(array.unsliced_n_rows(), ctx)?; + self.validated_all_null = index.pages.is_empty() && mask.all_false(); + self.legacy = Some(RankMask::new(mask)); + } + let Some(value) = self + .legacy + .as_ref() + .and_then(|mask| mask.value_index::(row)) + else { + return Ok(None); + }; + let page_idx = self.compact_page(index, value); + let page = index.pages.get(page_idx).ok_or_else(|| { + vortex_err!("Pco validity contains more valid rows than stored values") + })?; + if self.cached(page_idx).is_none() { + self.insert( + page_idx, + CachedPage { + rows: None, + values: page.values.clone(), + mask: None, + decoded: None, + }, + ); + } + Ok(Some((page_idx, value - page.values.start))) + } + + /// Select a contiguous logical span within one indexed page. Validity is sliced only after + /// the complete page has been checked, so copies use one typed contiguous compact-value run. + pub fn span( + &mut self, + array: ArrayView<'_, Pco>, + row: usize, + stop: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let index = array.page_index()?; + if !index.logical || index.pages.is_empty() { + if self.legacy.is_none() { + let mask = array + .unsliced_validity() + .execute_mask(array.unsliced_n_rows(), ctx)?; + let total = index.pages.last().map_or(0, |page| page.values.end); + vortex_ensure_eq!( + mask.true_count(), + total, + "Pco validity count does not match stored values" + ); + self.validated_all_null = index.pages.is_empty() && mask.all_false(); + self.legacy = Some(RankMask::new(mask)); + } + let mask = self + .legacy + .as_ref() + .ok_or_else(|| vortex_err!("Pco legacy validity is missing"))?; + let start_value = mask.rank(row); + let stop_value = mask.rank(stop); + if start_value == stop_value { + return Ok(Some(PageSpan { + page: 0, + stop, + values: 0..0, + mask: mask.mask.slice(row..stop), + })); + } + let page_idx = self.compact_page(index, start_value); + let page = index + .pages + .get(page_idx) + .ok_or_else(|| vortex_err!("Pco page does not cover compact value"))?; + let value_stop = stop_value.min(page.values.end); + let end = if value_stop == stop_value { + stop + } else { + let offset = match &mask.mask { + Mask::AllTrue(_) => value_stop - start_value - 1, + Mask::Values(values) => values + .bit_buffer() + .slice(row..stop) + .select(value_stop - start_value - 1) + .ok_or_else(|| vortex_err!("Pco legacy page rank is inconsistent"))?, + Mask::AllFalse(_) => { + return Err(vortex_err!("Pco all-null validity has compact values")); + } + }; + row.checked_add(offset) + .and_then(|end| end.checked_add(1)) + .ok_or_else(|| vortex_err!("Pco legacy logical page end overflows usize"))? + }; + let selected_mask = mask.mask.slice(row..end); + if self.cached(page_idx).is_none() { + self.insert( + page_idx, + CachedPage { + rows: None, + values: page.values.clone(), + mask: None, + decoded: None, + }, + ); + } + return Ok(Some(PageSpan { + page: page_idx, + stop: end, + values: (start_value - page.values.start)..(value_stop - page.values.start), + mask: selected_mask, + })); + } + let page_idx = self.logical_page(index, row); + let page = index + .pages + .get(page_idx) + .ok_or_else(|| vortex_err!("Pco logical page index does not cover span"))?; + let rows = page + .rows + .as_ref() + .ok_or_else(|| vortex_err!("Pco logical page range is missing"))?; + // Resolve performs whole-page validity validation, even when the first row is null. + self.resolve_impl::(array, row, ctx)?; + let cached = self + .cached(page_idx) + .and_then(|page| page.mask.as_ref()) + .ok_or_else(|| vortex_err!("Pco page validity is missing"))?; + let end = stop.min(rows.end); + let start_in_page = row - rows.start; + let end_in_page = end - rows.start; + Ok(Some(PageSpan { + page: page_idx, + stop: end, + values: cached.boundary_rank(start_in_page)..cached.boundary_rank(end_in_page), + mask: cached.mask.slice(start_in_page..end_in_page), + })) + } + + pub fn values( + &mut self, + array: ArrayView<'_, Pco>, + page_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult<&PrimitiveArray> { + let needs_decode = self + .cached(page_idx) + .ok_or_else(|| vortex_err!("Pco page must be resolved before decoding"))? + .decoded + .is_none(); + if needs_decode { + let index = array.page_index()?; + let page = index + .pages + .get(page_idx) + .ok_or_else(|| vortex_err!("Pco page is missing"))?; + let buffer = array + .pages + .get(page_idx) + .ok_or_else(|| vortex_err!("Pco compressed page is missing"))?; + let decoded = match_number_enum!( + number_type_from_ptype(array.dtype().as_ptype()), + NumberType => { + if self.decoder.is_none() { + self.decoder = Some(Box::new(crate::decoder::PageDecoder::::new(array.metadata.header.as_slice())?)); + } + let decoder = self.decoder.as_mut().and_then(|decoder| decoder.downcast_mut::>()) + .ok_or_else(|| vortex_err!("Pco operation decoder primitive type changed"))?; + decoder.read(page.chunk, array.chunk_metas[page.chunk].as_slice(), page.values.len(), buffer.as_slice(), ctx)? + } + ); + self.cached_mut(page_idx) + .ok_or_else(|| vortex_err!("Pco selected page cache is missing"))? + .decoded = Some(decoded); + } + self.cached_mut(page_idx) + .ok_or_else(|| vortex_err!("Pco selected page cache is missing"))? + .decoded + .as_ref() + .ok_or_else(|| vortex_err!("Pco page decoding did not produce values")) + } +} + +pub(crate) struct PageSpan { + pub page: usize, + pub stop: usize, + pub values: Range, + pub mask: Mask, +} diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index c1a4347cd94..380a559f9b6 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -7,6 +7,8 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; +use std::sync::Arc; +use std::sync::OnceLock; use pco::ChunkConfig; use pco::PagingSpec; @@ -49,6 +51,7 @@ use vortex_array::vtable::validity_to_child; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_error::VortexError; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -58,8 +61,10 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::PcoChunkInfo; +use crate::PcoLogicalPages; use crate::PcoMetadata; use crate::PcoPageInfo; +use crate::page_index::PageIndex; // Overall approach here: // Chunk the array into Pco chunks (currently using the default recommended size @@ -89,11 +94,22 @@ impl ArrayHash for PcoData { self.unsliced_n_rows.hash(state); self.slice_start.hash(state); self.slice_stop.hash(state); + self.metadata.header.hash(state); + for chunk in &self.metadata.chunks { + chunk.pages.len().hash(state); + for page in &chunk.pages { + page.n_values.hash(state); + } + } + if let Some(index) = &self.metadata.logical_pages { + index.version.hash(state); + index.ends.hash(state); + } // Hash chunk_metas and pages using pointer-based hashing - for chunk_meta in &self.chunk_metas { + for chunk_meta in self.chunk_metas.iter() { chunk_meta.array_hash(state, accuracy); } - for page in &self.pages { + for page in self.pages.iter() { page.array_hash(state, accuracy); } } @@ -106,15 +122,16 @@ impl ArrayEq for PcoData { || self.slice_stop != other.slice_stop || self.chunk_metas.len() != other.chunk_metas.len() || self.pages.len() != other.pages.len() + || self.metadata != other.metadata { return false; } - for (a, b) in self.chunk_metas.iter().zip(&other.chunk_metas) { + for (a, b) in self.chunk_metas.iter().zip(other.chunk_metas.iter()) { if !a.array_eq(b, accuracy) { return false; } } - for (a, b) in self.pages.iter().zip(&other.pages) { + for (a, b) in self.pages.iter().zip(other.pages.iter()) { if !a.array_eq(b, accuracy) { return false; } @@ -180,11 +197,13 @@ impl VTable for Pco { data.chunk_metas = buffers[..chunk_metas_len] .iter() .map(|buffer| buffer.clone().try_to_host_sync()) - .collect::>>()?; + .collect::>>()? + .into(); data.pages = buffers[chunk_metas_len..] .iter() .map(|buffer| buffer.clone().try_to_host_sync()) - .collect::>>()?; + .collect::>>()? + .into(); Ok(ArrayParts::new( self.clone(), array.dtype().clone(), @@ -198,7 +217,11 @@ impl VTable for Pco { array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - Ok(Some(array.metadata.clone().encode_to_vec())) + vortex_ensure!( + array.slice_start == 0 && array.slice_stop == array.unsliced_n_rows, + "Sliced Pco data must be executed to a primitive array (and optionally recompressed) before serialization" + ); + Ok(Some(array.metadata.encode_to_vec())) } fn deserialize( @@ -259,6 +282,11 @@ impl VTable for Pco { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + if array.slice_start() != 0 || array.slice_stop() != array.unsliced_n_rows() { + return Ok(ExecutionResult::done( + crate::take::slice(array.as_view(), ctx)?.into_array(), + )); + } let unsliced_validity = array.unsliced_validity(); Ok(ExecutionResult::done( array @@ -299,16 +327,16 @@ pub(crate) fn number_type_from_ptype(ptype: PType) -> NumberType { fn collect_valid( parray: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, -) -> VortexResult { +) -> VortexResult<(PrimitiveArray, vortex_mask::Mask)> { let mask = parray .array() .validity()? .execute_mask(parray.array().len(), ctx)?; let result = parray .array() - .filter(mask)? + .filter(mask.clone())? .execute::(ctx)?; - Ok(result) + Ok((result, mask)) } pub(crate) fn vortex_err_from_pco(err: PcoError) -> VortexError { @@ -384,9 +412,10 @@ impl> PcoArrayExt for T {} #[derive(Clone, Debug)] /// Encoding-specific data for a [`PcoArray`]. pub struct PcoData { - pub(crate) chunk_metas: Vec, - pub(crate) pages: Vec, - pub(crate) metadata: PcoMetadata, + pub(crate) chunk_metas: Arc<[ByteBuffer]>, + pub(crate) pages: Arc<[ByteBuffer]>, + pub(crate) metadata: Arc, + page_index: Arc>, ptype: PType, unsliced_n_rows: usize, slice_start: usize, @@ -466,6 +495,7 @@ impl PcoData { } else if validity.definitely_all_null() { vortex_ensure_eq!(n_values, 0, "Pco contains values for an all-null array"); } + self.page_index()?; Ok(()) } @@ -483,9 +513,10 @@ impl PcoData { len: usize, ) -> Self { Self { - chunk_metas, - pages, - metadata, + chunk_metas: chunk_metas.into(), + pages: pages.into(), + metadata: Arc::new(metadata), + page_index: Arc::new(OnceLock::new()), ptype, unsliced_n_rows: len, slice_start: 0, @@ -528,7 +559,7 @@ impl PcoData { .with_compression_level(level) .with_paging_spec(PagingSpec::EqualPagesUpTo(values_per_page)); - let values = collect_valid(parray, ctx)?; + let (values, mask) = collect_valid(parray, ctx)?; let n_values = values.len(); let fc = FileCompressor::default(); @@ -569,9 +600,47 @@ impl PcoData { chunk_infos.push(PcoChunkInfo { pages: page_infos }) } + let mut ends = Vec::with_capacity(page_buffers.len()); + match &mask { + vortex_mask::Mask::AllFalse(_) => {} + vortex_mask::Mask::AllTrue(_) => { + let mut end = 0usize; + for page in chunk_infos.iter().flat_map(|chunk| &chunk.pages) { + end = end + .checked_add(page.n_values as usize) + .ok_or_else(|| vortex_err!("Pco logical page prefix overflows usize"))?; + ends.push(u64::try_from(end)?); + } + } + vortex_mask::Mask::Values(mask) => { + let bits = mask.bit_buffer(); + let mut start = 0usize; + for page in chunk_infos.iter().flat_map(|chunk| &chunk.pages) { + if ends.len() + 1 == page_buffers.len() { + // No search through trailing nulls is necessary for the final page. + ends.push(u64::try_from(parray.len())?); + break; + } + // Existing bitmap select scans words only up to the next page boundary. + // Advancing the suffix avoids rescanning prefixes or visiting every null row. + let offset = bits + .slice(start..) + .select(page.n_values as usize - 1) + .ok_or_else(|| { + vortex_err!("Pco compression validity has too few valid rows") + })?; + start = start + .checked_add(offset) + .and_then(|end| end.checked_add(1)) + .ok_or_else(|| vortex_err!("Pco logical page end overflows usize"))?; + ends.push(u64::try_from(start)?); + } + } + } let metadata = PcoMetadata { header, chunks: chunk_infos, + logical_pages: Some(PcoLogicalPages { version: 1, ends }), }; // SAFETY: The compressor produced matching chunk metadata and pages from `parray`, with // the same ptype, logical length, and valid-value count. @@ -637,9 +706,21 @@ impl PcoData { ctx: &mut ExecutionCtx, ) -> VortexResult { // To start, we figure out what range of values we need to decompress. - let slice_value_indices = unsliced_validity - .execute_mask(self.unsliced_n_rows, ctx)? - .valid_counts_for_indices(&[self.slice_start, self.slice_stop]); + let mask = unsliced_validity.execute_mask(self.unsliced_n_rows, ctx)?; + let index = self.page_index()?; + if index.logical { + for page in &index.pages { + if let Some(rows) = &page.rows { + vortex_ensure_eq!( + mask.slice(rows.clone()).true_count(), + page.values.len(), + "Pco logical page validity count does not match stored values" + ); + } + } + } + let slice_value_indices = + mask.valid_counts_for_indices(&[self.slice_start, self.slice_stop]); let slice_value_start = slice_value_indices[0]; let slice_value_stop = slice_value_indices[1]; let slice_n_values = slice_value_stop - slice_value_start; @@ -655,7 +736,7 @@ impl PcoData { let mut page_idx = 0; let mut page_value_start = 0usize; let mut n_skipped_values = 0; - for (chunk_info, chunk_meta) in self.metadata.chunks.iter().zip(&self.chunk_metas) { + for (chunk_info, chunk_meta) in self.metadata.chunks.iter().zip(self.chunk_metas.iter()) { // lazily initialize chunk decompressor let mut chunk_decompressor: Option> = None; for page_info in &chunk_info.pages { @@ -722,6 +803,19 @@ impl PcoData { .into_byte_buffer()) } + pub(crate) fn page_index(&self) -> VortexResult<&PageIndex> { + if self.page_index.get().is_none() { + let index = PageIndex::try_new(&self.metadata, self.unsliced_n_rows)?; + // Concurrent first readers may prepare equivalent immutable descriptors. Only one + // index is retained, and no decoded values or validity are cached in the array. + drop(self.page_index.set(index)); + } + Ok(self + .page_index + .get() + .vortex_expect("page index was initialized")) + } + pub(crate) fn _slice(&self, start: usize, stop: usize) -> Self { PcoData { slice_start: self.slice_start + start, @@ -795,7 +889,6 @@ mod tests { use vortex_error::VortexResult; use vortex_session::VortexSession; - use super::VALUES_PER_CHUNK; use crate::Pco; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -846,19 +939,4 @@ mod tests { assert!(result.is_err()); Ok(()) } - - #[test] - fn test_validate_rejects_oversized_chunk() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let values = PrimitiveArray::from_iter([42u32]); - let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx)?; - let mut data = pco.data().clone(); - data.metadata.chunks[0].pages[0].n_values = u32::try_from(VALUES_PER_CHUNK + 1)?; - - assert!( - data.validate(pco.dtype(), pco.len(), &Validity::NonNullable) - .is_err() - ); - Ok(()) - } } diff --git a/encodings/pco/src/compute/cast.rs b/encodings/pco/src/compute/cast.rs index af1d003d7f7..17d3835d21e 100644 --- a/encodings/pco/src/compute/cast.rs +++ b/encodings/pco/src/compute/cast.rs @@ -10,7 +10,6 @@ use vortex_error::VortexResult; use crate::Pco; use crate::PcoArrayExt; -use crate::PcoData; impl CastReduce for Pco { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { @@ -35,18 +34,8 @@ impl CastReduce for Pco { return Ok(None); }; - // SAFETY: The buffers and metadata come from a validated Pco array, the primitive type is - // unchanged, and `Pco::try_new` validates the adjusted nullability and slice below. - let data = unsafe { - PcoData::new_unchecked( - array.chunk_metas.clone(), - array.pages.clone(), - dtype.as_ptype(), - array.metadata.clone(), - array.unsliced_n_rows(), - ) - } - ._slice(array.slice_start(), array.slice_stop()); + // Only nullability changes; retain the same immutable page index and compressed owners. + let data = array.data().clone(); Ok(Some( Pco::try_new(dtype.clone(), data, new_validity)?.into_array(), diff --git a/encodings/pco/src/decoder.rs b/encodings/pco/src/decoder.rs new file mode 100644 index 00000000000..620cb9c366f --- /dev/null +++ b/encodings/pco/src/decoder.rs @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! File and chunk decoding setup amortized over all selected pages in an operation. + +use std::collections::BTreeMap; +use std::collections::btree_map::Entry; + +use pco::data_types::Number; +use pco::wrapped::ChunkDecompressor; +use pco::wrapped::FileDecompressor; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::validity::Validity; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; + +use crate::array::vortex_err_from_pco; + +pub(crate) struct PageDecoder { + file: FileDecompressor, + chunks: BTreeMap>, +} + +impl PageDecoder { + pub fn new(header: &[u8]) -> VortexResult { + let (file, _) = FileDecompressor::new(header).map_err(vortex_err_from_pco)?; + Ok(Self { + file, + chunks: BTreeMap::new(), + }) + } + + pub fn read( + &mut self, + chunk: usize, + chunk_meta: &[u8], + n_values: usize, + buffer: &[u8], + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let mut values = BufferMut::::zeroed_in(n_values, ctx.allocator().clone()); + self.read_into(chunk, chunk_meta, n_values, buffer, &mut values)?; + Ok(PrimitiveArray::new(values.freeze(), Validity::NonNullable)) + } + /// Decode exactly one complete page into caller-owned initialized storage. + /// A corrupt/partial page never publishes the caller's unfinished output array. + pub fn read_into( + &mut self, + chunk: usize, + chunk_meta: &[u8], + n_values: usize, + buffer: &[u8], + output: &mut [T], + ) -> VortexResult<()> { + vortex_ensure_eq!(output.len(), n_values, "Pco page output length mismatch"); + let chunk = match self.chunks.entry(chunk) { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => { + let (decoder, _) = self + .file + .chunk_decompressor::(chunk_meta) + .map_err(vortex_err_from_pco)?; + entry.insert(decoder) + } + }; + let mut decoder = chunk + .page_decompressor(buffer, n_values) + .map_err(vortex_err_from_pco)?; + let progress = decoder.read(output).map_err(vortex_err_from_pco)?; + vortex_ensure_eq!( + progress.n_processed, + n_values, + "Pco page did not fill output" + ); + vortex_ensure!(progress.finished, "Pco page decode is incomplete"); + Ok(()) + } +} diff --git a/encodings/pco/src/lib.rs b/encodings/pco/src/lib.rs index df6581ec1d1..179f12ed078 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -12,22 +12,25 @@ //! To deserialize arrays manually, register the encoding in the array session: //! //! ```rust -//! use vortex_array::session::ArraySessionExt; //! use vortex_session::VortexSession; //! //! let session = VortexSession::empty(); -//! session.arrays().register(vortex_pco::Pco); +//! vortex_pco::initialize(&session); //! ``` +mod access; mod array; mod compute; +mod decoder; +mod page_index; mod probe; mod rules; mod slice; +mod take; pub use array::*; -#[derive(Clone, prost::Message)] +#[derive(Clone, PartialEq, Eq, prost::Message)] /// Metadata for one Pco page. pub struct PcoPageInfo { // Since pco limits to 2^24 values per chunk, u32 is sufficient for the @@ -39,7 +42,7 @@ pub struct PcoPageInfo { // We're calling this Info instead of Metadata because ChunkMeta refers to a specific // component of a Pco file. -#[derive(Clone, prost::Message)] +#[derive(Clone, PartialEq, Eq, prost::Message)] /// Metadata for one Pco chunk. pub struct PcoChunkInfo { /// Pages contained in this chunk. @@ -47,7 +50,7 @@ pub struct PcoChunkInfo { pub pages: Vec, } -#[derive(Clone, prost::Message)] +#[derive(Clone, PartialEq, Eq, prost::Message)] /// Serialized metadata for a [`PcoArray`]. pub struct PcoMetadata { // would be nice to reuse one header per vortex file, but it's really only 1 byte, so @@ -58,7 +61,41 @@ pub struct PcoMetadata { /// Metadata for each compressed chunk. #[prost(message, repeated, tag = "2")] pub chunks: Vec, + /// Optional logical-row boundaries for selective decoding. Older readers ignore this field. + #[prost(message, optional, tag = "3")] + pub logical_pages: Option, } +/// Versioned optional logical-row page index. Unknown versions retain legacy decoding. +#[derive(Clone, PartialEq, Eq, prost::Message)] +pub struct PcoLogicalPages { + /// Index interpretation version; currently one. + #[prost(uint32, tag = "1")] + pub version: u32, + /// Exclusive unsliced logical-row ends in physical page order. + /// + /// The final page includes trailing null rows. An empty index has no valid values. + #[prost(uint64, repeated, tag = "2")] + pub ends: Vec, +} + +/// Register PCO encoding and its selective take execution kernel in a session. +/// +/// Registering only `Pco` in the array registry is sufficient for legacy reads, but does not +/// enable the batched take kernel. +pub fn initialize(session: &vortex_session::VortexSession) { + use vortex_array::ArrayVTable; + use vortex_array::arrays::Dict; + use vortex_array::arrays::dict::TakeExecuteAdaptor; + use vortex_array::optimizer::kernels::ArrayKernelsExt; + use vortex_array::session::ArraySessionExt; + session.arrays().register(Pco); + session + .kernels() + .register_execute_parent_kernel(Dict.id(), Pco, TakeExecuteAdaptor(Pco)); +} + +#[cfg(test)] +mod select_tests; #[cfg(test)] mod tests; diff --git a/encodings/pco/src/page_index.rs b/encodings/pco/src/page_index.rs new file mode 100644 index 00000000000..8fe0faab72f --- /dev/null +++ b/encodings/pco/src/page_index.rs @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Immutable physical-page descriptors shared by sliced owners and access operations. + +use std::ops::Range; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; + +use crate::PcoMetadata; + +#[derive(Debug)] +pub(crate) struct Page { + pub chunk: usize, + pub values: Range, + pub rows: Option>, +} + +#[derive(Debug)] +pub(crate) struct PageIndex { + pub pages: Vec, + pub logical: bool, +} + +impl PageIndex { + pub fn try_new(metadata: &PcoMetadata, rows: usize) -> VortexResult { + let logical = metadata + .logical_pages + .as_ref() + .filter(|index| index.version == 1); + let page_count = metadata.chunks.iter().try_fold(0usize, |count, chunk| { + count + .checked_add(chunk.pages.len()) + .ok_or_else(|| vortex_err!("Pco page count overflows usize")) + })?; + if let Some(index) = logical { + vortex_ensure_eq!( + index.ends.len(), + page_count, + "Pco logical page count does not match physical pages" + ); + } + let mut pages = Vec::with_capacity(page_count); + let mut value_start = 0usize; + let mut row_start = 0usize; + for (chunk, info) in metadata.chunks.iter().enumerate() { + for page in &info.pages { + let n = page.n_values as usize; + vortex_ensure!(n != 0, "Pco contains an empty physical page"); + let end = value_start + .checked_add(n) + .ok_or_else(|| vortex_err!("Pco value prefix overflows usize"))?; + vortex_ensure!(end <= rows, "Pco value prefix exceeds logical row count"); + let logical_rows = if let Some(index) = logical { + let stop = usize::try_from(index.ends[pages.len()])?; + vortex_ensure!( + stop > row_start && stop <= rows, + "Pco logical page ends are not strictly increasing within array bounds" + ); + vortex_ensure!( + stop - row_start >= n, + "Pco page has fewer logical rows than values" + ); + let range = row_start..stop; + row_start = stop; + Some(range) + } else { + None + }; + pages.push(Page { + chunk, + values: value_start..end, + rows: logical_rows, + }); + value_start = end; + } + } + if logical.is_some() && page_count != 0 { + vortex_ensure_eq!( + row_start, + rows, + "Pco logical page ends do not cover all rows" + ); + } + Ok(Self { + pages, + logical: logical.is_some(), + }) + } +} diff --git a/encodings/pco/src/probe.rs b/encodings/pco/src/probe.rs index 00aa801254a..8df0537c2d8 100644 --- a/encodings/pco/src/probe.rs +++ b/encodings/pco/src/probe.rs @@ -3,51 +3,23 @@ //! PCO probes retain every decoded page and index the compacted, non-null value positions. -use std::ops::Range; - -use pco::data_types::Number; -use pco::data_types::NumberType; -use pco::match_number_enum; -use pco::wrapped::FileDecompressor; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::dtype::NativePType; -use vortex_array::dtype::half; use vortex_array::match_each_native_ptype; use vortex_array::scalar::Scalar; -use vortex_array::validity::Validity; -use vortex_buffer::BufferMut; -use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_mask::Mask; +use vortex_error::vortex_ensure; use crate::Pco; -use crate::PcoArrayExt; -use crate::array::number_type_from_ptype; -use crate::array::vortex_err_from_pco; - -const RANK_STRIDE: usize = 512; +use crate::access::AccessState; /// State retained by repeated PCO probes. Construction performs no allocation. /// -/// The mask and rank index use unsliced logical rows; page boundaries use compacted value -/// positions. Every page decoded so far is retained, so the state grows with the pages touched -/// and never beyond the decompressed array. +/// Decoded pages and their local validity are retained as touched, up to the full array. +/// Legacy files additionally retain a full validity mask and its sparse rank index. #[derive(Default)] pub struct PcoProbeState { - validity: Option, - rank: Vec, - pages: Vec, - /// The page the previous probe landed in, `None` until the first probe resolves one. - last_page: Option, -} - -struct Page { - values: Range, - chunk: usize, - decoded: Option, + access: AccessState, } pub(crate) fn scalar_at( @@ -57,110 +29,47 @@ pub(crate) fn scalar_at( ctx: &mut ExecutionCtx, ) -> VortexResult { let Some(state) = state else { - let validity = array.unsliced_validity(); - return array - ._slice(index, index + 1) - .decompress(&validity, ctx)? - .into_array() - .execute_scalar(0, ctx); - }; - let mask = match &mut state.validity { - Some(mask) => mask, - slot @ None => slot.insert( - array - .unsliced_validity() - .execute_mask(array.unsliced_n_rows(), ctx)?, - ), - }; - let logical_index = array.slice_start() + index; - - let value_index = match mask { - Mask::AllTrue(_) => logical_index, - Mask::AllFalse(_) => return Ok(Scalar::null(array.dtype().clone())), - Mask::Values(values) => { - let bits = values.bit_buffer(); - if !bits.value(logical_index) { - return Ok(Scalar::null(array.dtype().clone())); - } - if state.rank.is_empty() { - state.rank.reserve(bits.len().div_ceil(RANK_STRIDE)); - let mut count = 0; - for start in (0..bits.len()).step_by(RANK_STRIDE) { - state.rank.push(count); - count += bits.count_range(start, (start + RANK_STRIDE).min(bits.len())); - } - } - let block = logical_index / RANK_STRIDE; - state.rank[block] + bits.count_range(block * RANK_STRIDE, logical_index) - } + return scalar_at_once(array, index, ctx); }; - - if state.pages.is_empty() { - let mut start = 0; - for (chunk, metadata) in array.metadata.chunks.iter().enumerate() { - for page in &metadata.pages { - let end = start + page.n_values as usize; - state.pages.push(Page { - values: start..end, - chunk, - decoded: None, - }); - start = end; - } - } + if state.access.validated_all_null() { + vortex_ensure!( + index < array.len(), + "Pco scalar index exceeds visible array bounds" + ); + return Ok(Scalar::null(array.dtype().clone())); } - let page_index = match state.last_page { - Some(last) if state.pages[last].values.contains(&value_index) => last, - _ => state - .pages - .partition_point(|page| page.values.end <= value_index), - }; - state.last_page = Some(page_index); - let page = state - .pages - .get_mut(page_index) - .vortex_expect("PCO pages cover every valid value index"); - let values = match &mut page.decoded { - Some(decoded) => decoded, - slot @ None => { - let decoded = match_number_enum!( - number_type_from_ptype(array.dtype().as_ptype()), - NumberType => { - decode_page::(array, page.chunk, page.values.len(), array.pages[page_index].as_slice(), ctx)? - } - ); - slot.insert(decoded) - } - }; - Ok(match_each_native_ptype!(values.ptype(), |T| { - Scalar::primitive( - values.as_slice::()[value_index - page.values.start], - array.dtype().nullability(), - ) - })) + let row = array.slice_start() + index; + if let Some(value) = state.access.cached_scalar(array, row) { + return Ok(value); + } + scalar_at_uncached(array, row, state, ctx) } -fn decode_page( +// Keep the large temporary cache out of the retained-probe hit frame. Stateless +// probes and cache misses are ordinary workloads, so neither helper is cold. +#[inline(never)] +fn scalar_at_once( array: ArrayView<'_, Pco>, - chunk: usize, - n_values: usize, - buffer: &[u8], + index: usize, ctx: &mut ExecutionCtx, -) -> VortexResult { - let (file, _) = - FileDecompressor::new(array.metadata.header.as_slice()).map_err(vortex_err_from_pco)?; - let (mut chunk, _) = file - .chunk_decompressor::(array.chunk_metas[chunk].as_ref()) - .map_err(vortex_err_from_pco)?; - let mut decoder = chunk - .page_decompressor(buffer, n_values) - .map_err(vortex_err_from_pco)?; - let mut values = BufferMut::::with_capacity_in(n_values, ctx.allocator().clone()); - // SAFETY: the buffer reserves `n_values` elements, and the page decompressor was built for - // exactly that count, so `read` writes every element before anything observes it. - unsafe { values.set_len(n_values) }; - decoder.read(&mut values).map_err(vortex_err_from_pco)?; - Ok(PrimitiveArray::new(values.freeze(), Validity::NonNullable)) +) -> VortexResult { + scalar_at(array, index, Some(&mut PcoProbeState::default()), ctx) +} + +#[inline(never)] +fn scalar_at_uncached( + array: ArrayView<'_, Pco>, + row: usize, + state: &mut PcoProbeState, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let Some((page, value)) = state.access.resolve(array, row, ctx)? else { + return Ok(Scalar::null(array.dtype().clone())); + }; + let values = state.access.values(array, page, ctx)?; + Ok(match_each_native_ptype!(values.ptype(), |T| { + Scalar::primitive(values.as_slice::()[value], array.dtype().nullability()) + })) } #[cfg(test)] diff --git a/encodings/pco/src/select_tests.rs b/encodings/pco/src/select_tests.rs new file mode 100644 index 00000000000..e9330ff978c --- /dev/null +++ b/encodings/pco/src/select_tests.rs @@ -0,0 +1,543 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::LazyLock; + +use prost::Message; +use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PiecewiseSequenceArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::buffer::BufferHandle; +use vortex_array::serde::SerializeOptions; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::Pco; +use crate::PcoLogicalPages; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session +}); + +/// Build a wire fixture through the public serializer and deserialize its mutated protobuf +/// metadata through the public plugin boundary. No private in-memory encoding fields are used. +fn with_wire_metadata( + array: &ArrayRef, + mutate: impl FnOnce(&mut crate::PcoMetadata), +) -> VortexResult { + let mut wire = SESSION + .array_serialize(array)? + .vortex_expect("PCO is serializable"); + let mut metadata = crate::PcoMetadata::decode(wire.metadata.as_slice())?; + mutate(&mut metadata); + wire.metadata = metadata.encode_to_vec(); + let buffers = wire + .buffers + .into_iter() + .map(BufferHandle::new_host) + .collect::>(); + ArrayPlugin::deserialize( + &Pco, + ArrayDeserialization::new( + wire.serialized_id, + array.dtype(), + array.len(), + &wire.metadata, + &buffers, + &wire.children, + ), + &SESSION, + ) +} + +fn legacy(array: &ArrayRef, unknown: bool) -> VortexResult { + with_wire_metadata(array, |metadata| { + metadata.logical_pages = unknown.then_some(PcoLogicalPages { + version: 999, + ends: vec![u64::MAX], + }); + }) +} + +#[rstest] +fn reordered_nullable_take(#[values(0, 1, 2)] format: usize) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = + PrimitiveArray::from_option_iter((0..4096i32).map(|i| (i % 7 != 0).then_some(i * 19))); + let encoded = Pco::from_primitive(input.as_view(), 3, 31, &mut ctx)?.into_array(); + let encoded = match format { + 0 => encoded, + 1 => legacy(&encoded, false)?, + _ => legacy(&encoded, true)?, + }; + let source = encoded.slice(17..4001)?.slice(13..3900)?; + let input = input.slice(30..3917)?; + let mut probe = source.repeated_probe(); + for index in [3000, 0, 127, 1, 3000, 128, 1000, 0] { + assert_eq!( + probe.execute_scalar(index, &mut ctx)?, + input.execute_scalar(index, &mut ctx)? + ); + } + let indices = PrimitiveArray::from_option_iter([ + Some(3000u32), + Some(0), + None, + Some(127), + Some(1), + Some(3000), + Some(128), + None, + Some(1000), + ]) + .into_array(); + assert_arrays_eq!( + source.take(indices.clone())?, + input.take(indices)?, + &mut ctx + ); + Ok(()) +} + +#[rstest] +#[case(vec![0u64, 40, 0, 70], vec![15u64, 20, 3, 0], vec![1u64, 1, 0, 1])] +#[case(vec![100u64, 0, 20], vec![20u64, 20, 20], vec![1u64, 1, 2])] +#[case(vec![99u64, 0], vec![0u64, 0], vec![u64::MAX, 0])] +fn piecewise_ranges( + #[case] starts: Vec, + #[case] lengths: Vec, + #[case] strides: Vec, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = + PrimitiveArray::from_option_iter((0..256i64).map(|i| (i % 5 != 0).then_some(i * i))); + let encoded = Pco::from_primitive(input.as_view(), 3, 7, &mut ctx)?.into_array(); + let mut expected_indices = Vec::new(); + for ((&start, &length), &stride) in starts.iter().zip(&lengths).zip(&strides) { + expected_indices.extend((0..length).map(|i| start + i * stride)); + } + let indices = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter(starts).into_array(), + PrimitiveArray::from_iter(lengths).into_array(), + PrimitiveArray::from_iter(strides).into_array(), + expected_indices.len(), + )? + .into_array(); + assert_arrays_eq!( + encoded.take(indices)?, + input.take(PrimitiveArray::from_iter(expected_indices).into_array())?, + &mut ctx + ); + Ok(()) +} + +#[test] +fn all_null_and_long_null_runs() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + for input in [ + PrimitiveArray::new(vec![0i32; 1024], Validity::AllInvalid), + PrimitiveArray::from_option_iter( + (0..100_000).map(|i| [1, 90_000, 99_990].contains(&i).then_some(i)), + ), + ] { + let encoded = Pco::from_primitive(input.as_view(), 3, 1, &mut ctx)?.into_array(); + let indices = PrimitiveArray::from_iter([0u32, 1, 100, 1000, 2, 1]).into_array(); + assert_arrays_eq!( + encoded.take(indices.clone())?, + input.take(indices)?, + &mut ctx + ); + assert_arrays_eq!(encoded.slice(1..1000)?, input.slice(1..1000)?, &mut ctx); + } + Ok(()) +} + +#[test] +fn rejects_invalid_indices_and_descriptors() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_iter(0..64i32); + let encoded = Pco::from_primitive(input.as_view(), 3, 4, &mut ctx)?.into_array(); + for indices in [ + PrimitiveArray::from_iter([-1i64]).into_array(), + PrimitiveArray::from_iter([64i64]).into_array(), + ] { + assert!( + encoded + .take(indices) + .and_then(|array| array.execute::(&mut ctx)) + .is_err() + ); + } + for (start, length, stride, declared) in [ + (u64::MAX, 2u64, 1u64, 2usize), + (1, 2, u64::MAX, 2), + (1, 2, 1, 3), + ] { + let indices = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter([start]).into_array(), + PrimitiveArray::from_iter([length]).into_array(), + PrimitiveArray::from_iter([stride]).into_array(), + declared, + )? + .into_array(); + assert!( + encoded + .take(indices) + .and_then(|array| array.execute::(&mut ctx)) + .is_err() + ); + } + Ok(()) +} + +#[test] +fn raw_sliced_serde_requires_explicit_normalization() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_option_iter((0..128i32).map(|i| (i % 3 != 0).then_some(i))); + let encoded = Pco::from_primitive(input.as_view(), 3, 7, &mut ctx)?.into_array(); + let sliced = encoded.slice(11..77)?; + assert!( + sliced + .serialize( + &ArrayContext::empty(), + &SESSION, + &SerializeOptions::default() + ) + .is_err() + ); + let normalized = sliced.execute::(&mut ctx)?; + let recompressed = Pco::from_primitive(normalized.as_view(), 3, 7, &mut ctx)?.into_array(); + assert!( + !recompressed + .serialize( + &ArrayContext::empty(), + &SESSION, + &SerializeOptions::default() + )? + .is_empty() + ); + assert_arrays_eq!(recompressed, input.slice(11..77)?, &mut ctx); + Ok(()) +} + +#[test] +fn registry_only_session_retains_legacy_take_fallback() -> VortexResult<()> { + let session = VortexSession::empty(); + session.arrays().register(Pco); + let mut ctx = session.create_execution_ctx(); + let input = PrimitiveArray::from_option_iter([Some(5i32), None, Some(9), Some(2)]); + let encoded = Pco::from_primitive(input.as_view(), 3, 1, &mut ctx)?.into_array(); + let indices: ArrayRef = PrimitiveArray::from_iter([3u32, 1, 0, 3]).into_array(); + assert_arrays_eq!( + encoded.take(indices.clone())?, + input.take(indices)?, + &mut ctx + ); + Ok(()) +} + +#[test] +fn checked_metadata_and_touched_page_validity() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_option_iter((0..128i32).map(|i| (i % 2 == 0).then_some(i))); + let encoded = Pco::from_primitive(input.as_view(), 3, 2, &mut ctx)?.into_array(); + for ends in [vec![128u64], vec![0; 32], vec![u64::MAX; 32]] { + assert!( + with_wire_metadata(&encoded, |metadata| { + metadata.logical_pages = Some(PcoLogicalPages { version: 1, ends }); + }) + .is_err() + ); + } + assert!( + with_wire_metadata(&encoded, |metadata| { + metadata.chunks[0].pages[0].n_values = u32::try_from(pco::DEFAULT_MAX_PAGE_N + 1) + .vortex_expect("invalid fixture count fits u32"); + }) + .is_err() + ); + let malformed = with_wire_metadata(&encoded, |metadata| { + // Structurally valid but the first logical page has three valid rows, not two. + metadata + .logical_pages + .as_mut() + .vortex_expect("writer emits logical ends") + .ends[0] = 5; + })?; + assert!( + malformed + .take(PrimitiveArray::from_iter([0u32]).into_array())? + .execute::(&mut ctx) + .is_err() + ); + Ok(()) +} + +#[test] +fn large_public_compressor_selections() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let len = pco::DEFAULT_MAX_PAGE_N + 257; + let input = PrimitiveArray::from_iter( + (0..len).map(|i| u32::try_from(i).vortex_expect("fixture fits u32")), + ); + let encoded = Pco::from_primitive(input.as_view(), 0, 4096, &mut ctx)?.into_array(); + let indices = PrimitiveArray::from_iter([ + 0u32, + 900, + 260, + 3000, + u32::try_from(len - 1).vortex_expect("fixture fits u32"), + 4095, + 900, + ]) + .into_array(); + assert_arrays_eq!( + encoded.take(indices.clone())?, + input.take(indices)?, + &mut ctx + ); + assert_arrays_eq!( + encoded.slice(len - 1000..len - 1)?, + input.slice(len - 1000..len - 1)?, + &mut ctx + ); + Ok(()) +} + +#[rstest] +fn repeated_scattered_requests_preserve_order_and_nulls( + #[values(false, true)] legacy_format: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..32768i64) + .map(|row| (row % 7 != 0).then_some(row * 13)) + .collect::>(); + let input = PrimitiveArray::from_option_iter(values.iter().copied()); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + let encoded = if legacy_format { + legacy(&encoded, false)? + } else { + encoded + }; + let requests = (0..1024usize) + .map(|i| (i % 11 != 0).then_some((i * 97) % 16)) + .collect::>(); + let indices = PrimitiveArray::from_option_iter( + requests + .iter() + .map(|row| row.map(|row| u64::try_from(row).vortex_expect("small index"))), + ) + .into_array(); + let expected = PrimitiveArray::from_option_iter( + requests.iter().map(|row| row.and_then(|row| values[row])), + ); + assert_arrays_eq!(encoded.take(indices)?, expected, &mut ctx); + let invalid = PrimitiveArray::from_iter([0i64, -1, 32768]).into_array(); + assert!( + encoded + .take(invalid)? + .execute::(&mut ctx) + .is_err() + ); + Ok(()) +} + +#[test] +fn mixed_fragments_narrow_slices_and_compact_repeats() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_option_iter( + (0..32768i64).map(|row| (row % 5 != 0).then_some(row * 11)), + ); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + let mixed = PrimitiveArray::from_iter( + (0..4096u64).map(|i| if i < 2 { i } else { (i * 104729) % 32768 }), + ) + .into_array(); + assert_arrays_eq!(encoded.take(mixed.clone())?, input.take(mixed)?, &mut ctx); + let repeated = + PrimitiveArray::from_iter((0..100000u64).map(|i| if i % 2 == 0 { 0 } else { 127 })) + .into_array(); + assert_arrays_eq!( + encoded.slice(30000..30128)?.take(repeated.clone())?, + input.slice(30000..30128)?.take(repeated)?, + &mut ctx + ); + let compact = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter([1u64, 5]).into_array(), + PrimitiveArray::from_iter([500000u64, 500000]).into_array(), + PrimitiveArray::from_iter([0u64, 0]).into_array(), + 1000000, + )? + .into_array(); + let expected = PrimitiveArray::from_option_iter( + std::iter::repeat_n(Some(11i64), 500000).chain(std::iter::repeat_n(None, 500000)), + ); + assert_arrays_eq!(encoded.take(compact)?, expected, &mut ctx); + let nulls = PrimitiveArray::new(vec![0i64; 100000], Validity::AllInvalid); + let encoded_nulls = Pco::from_primitive(nulls.as_view(), 3, 1024, &mut ctx)?.into_array(); + let two = PrimitiveArray::from_iter([0u64, 99999]).into_array(); + assert_arrays_eq!(encoded_nulls.take(two.clone())?, nulls.take(two)?, &mut ctx); + assert!( + encoded_nulls + .take(PrimitiveArray::from_iter([100000u64]).into_array())? + .execute::(&mut ctx) + .is_err() + ); + Ok(()) +} + +#[test] +fn empty_take_dtype_and_mostly_null_request_storage() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_iter(0..32768i64); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + for source in [encoded.clone(), encoded.slice(0..0)?] { + for indices in [ + PrimitiveArray::from_iter(std::iter::empty::()).into_array(), + PrimitiveArray::from_option_iter(std::iter::empty::>()).into_array(), + ] { + let result = source + .take(indices.clone())? + .execute::(&mut ctx)?; + let expected = input + .slice(0..0)? + .take(indices)? + .execute::(&mut ctx)?; + assert_eq!(result.dtype(), expected.dtype()); + assert_arrays_eq!(result, expected, &mut ctx); + } + } + let indices = + PrimitiveArray::from_option_iter((0..200000).map(|i| (i == 12345).then_some(1u64))) + .into_array(); + let result = encoded.take(indices)?.execute::(&mut ctx)?; + let expected = + PrimitiveArray::from_option_iter((0..200000).map(|i| (i == 12345).then_some(1i64))); + assert_eq!(result.dtype(), expected.dtype()); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) +} + +#[rstest] +fn streamed_slices_preserve_float_bits_and_validity( + #[values(false, true)] dense: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let bits = [ + 0u64, + 1u64 << 63, + 0x7ff8000000000042, + 0x7ff0000000000000, + 0xfff0000000000000, + 0x3ff4000000000000, + ]; + let values = (0..2048) + .map(|row| f64::from_bits(bits[row % bits.len()])) + .collect::>(); + let validity = if dense { + Validity::NonNullable + } else { + Validity::from_iter((0..2048).map(|row| row % 7 != 0)) + }; + let input = PrimitiveArray::new(values.clone(), validity); + let encoded = Pco::from_primitive(input.as_view(), 3, 31, &mut ctx)?.into_array(); + for range in [1..2047, 89..129, 1000..1001] { + let actual = encoded + .slice(range.clone())? + .execute::(&mut ctx)?; + assert_eq!(actual.dtype(), input.dtype()); + let mask = actual.validity()?.execute_mask(actual.len(), &mut ctx)?; + for (offset, &value) in actual.as_slice::().iter().enumerate() { + let row = range.start + offset; + assert_eq!(mask.value(offset), dense || row % 7 != 0); + if mask.value(offset) { + assert_eq!(value.to_bits(), values[row].to_bits()); + } + } + } + Ok(()) +} + +#[test] +fn logical_null_indices_ignore_invalid_physical_values() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_iter(0..128i64); + let encoded = Pco::from_primitive(input.as_view(), 3, 7, &mut ctx)?.into_array(); + let indices = PrimitiveArray::new( + vec![i64::MIN, 4, i64::MAX], + Validity::from_iter([false, true, false]), + ) + .into_array(); + let expected = PrimitiveArray::from_option_iter([None, Some(4i64), None]); + assert_arrays_eq!(encoded.take(indices.clone())?, expected, &mut ctx); + // Reusing the immutable index array also exercises cached native extrema. + assert_arrays_eq!( + encoded.take(indices)?, + PrimitiveArray::from_option_iter([None, Some(4i64), None]), + &mut ctx + ); + assert!( + encoded + .take(PrimitiveArray::from_iter([u64::MAX]).into_array())? + .execute::(&mut ctx) + .is_err() + ); + Ok(()) +} + +#[test] +fn contiguous_take_preserves_nullable_index_dtype() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_iter(0..32768i64); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + let indices = PrimitiveArray::from_option_iter((1..32767u64).map(Some)).into_array(); + let expected = input + .take(indices.clone())? + .execute::(&mut ctx)?; + let result = encoded.take(indices)?.execute::(&mut ctx)?; + assert_eq!(result.dtype(), expected.dtype()); + assert_arrays_eq!(result, expected, &mut ctx); + let range = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter([1u64]).into_array(), + PrimitiveArray::from_iter([32766u64]).into_array(), + PrimitiveArray::from_iter([1u64]).into_array(), + 32766, + )? + .into_array(); + assert_arrays_eq!(encoded.take(range)?, input.slice(1..32767)?, &mut ctx); + Ok(()) +} + +#[test] +fn retained_all_null_probes_preserve_slice_bounds() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::new(vec![0i64; 100000], Validity::AllInvalid); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + let sliced = encoded.slice(17..145)?; + let mut probe = sliced.repeated_probe(); + for index in [0, 127, 64, 0, 127] { + assert_eq!( + probe.execute_scalar(index, &mut ctx)?, + input.execute_scalar(17 + index, &mut ctx)? + ); + } + assert!(probe.execute_scalar(128, &mut ctx).is_err()); + assert!(probe.execute_scalar(usize::MAX, &mut ctx).is_err()); + assert_eq!( + probe.execute_scalar(0, &mut ctx)?, + input.execute_scalar(17, &mut ctx)? + ); + Ok(()) +} diff --git a/encodings/pco/src/slice.rs b/encodings/pco/src/slice.rs index 316ea835aa1..80e33972153 100644 --- a/encodings/pco/src/slice.rs +++ b/encodings/pco/src/slice.rs @@ -8,12 +8,17 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::arrays::slice::SliceReduce; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use crate::Pco; use crate::PcoArrayExt; impl SliceReduce for Pco { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { + vortex_ensure!( + range.start <= range.end && range.end <= array.len(), + "Pco slice range exceeds array bounds" + ); let unsliced_validity = array.unsliced_validity(); Ok(Some( Pco::try_new( diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs new file mode 100644 index 00000000000..678798d1ed2 --- /dev/null +++ b/encodings/pco/src/take.rs @@ -0,0 +1,750 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Batched selection using compact request descriptors and one operation-owned page cache. + +use std::collections::BTreeSet; +use std::ops::Range; + +use pco::data_types::Number; +use pco::data_types::NumberType; +use pco::match_number_enum; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; +use vortex_array::aggregate_fn::fns::min_max::min_max; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PiecewiseSequence; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::dict::TakeExecute; +use vortex_array::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::half; +use vortex_array::match_each_integer_ptype; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_mask::Mask; + +use crate::Pco; +use crate::PcoArrayExt; +use crate::access::AccessState; +use crate::array::number_type_from_ptype; +use crate::decoder::PageDecoder; + +/// Compact descriptor; null indices have no start. Zero-length runs consume no source rows. +struct Run { + start: Option, + length: usize, + stride: usize, +} + +impl TakeExecute for Pco { + fn take( + array: ArrayView<'_, Pco>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if indices.is_empty() { + let dtype = array + .dtype() + .clone() + .union_nullability(indices.dtype().nullability()); + return Ok(Some(Canonical::empty(&dtype).into_array())); + } + let compact = indices.as_opt::().is_some(); + let mut grouped = None; + let runs = if let Some(sequence) = indices.as_opt::() { + let starts = sequence.starts().clone().execute::(ctx)?; + let lengths = sequence.lengths().clone().execute::(ctx)?; + let multipliers = sequence + .multipliers() + .clone() + .execute::(ctx)?; + let starts = primitive_usizes(&starts)?; + let lengths = primitive_usizes(&lengths)?; + let multipliers = primitive_usizes(&multipliers)?; + let mut runs = Vec::with_capacity(starts.len()); + for ((start, length), stride) in starts.into_iter().zip(lengths).zip(multipliers) { + runs.push(Run { + start: Some(start), + length, + stride, + }); + } + runs + } else { + let indices = indices.clone().execute::(ctx)?; + let mask = indices.validity()?.execute_mask(indices.len(), ctx)?; + // Reuse native logical-null-aware extrema and their immutable-array stats cache. + // Numeric extrema prove every defined index is representable and in bounds. + let extrema = min_max(indices.as_array(), ctx, NumericalAggregateOpts::default())?; + let summary = if let Some(extrema) = extrema { + match_each_integer_ptype!(indices.ptype(), |I| { + validate_extrema(indices.as_slice::(), &mask, &extrema, array.len())? + }) + } else { + PlainSummary { + defined: 0, + contiguous: false, + } + }; + if summary.defined == 0 { + return Ok(Some( + ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len()) + .into_array(), + )); + } + if let Some(result) = zero_page_take(array, indices.len(), ctx)? { + return Ok(Some(result)); + } + grouped = (!summary.contiguous).then_some(summary.defined); + if grouped.is_some() && descriptor_bytes(summary.defined) >= canonical_bytes(array) { + return Ok(None); + } + if grouped.is_some() + && match_each_integer_ptype!(indices.ptype(), |I| { + prefer_canonical_primitive( + array, + indices.as_slice::(), + &mask, + summary.defined, + )? + }) + { + return Ok(None); + } + match_each_integer_ptype!(indices.ptype(), |I| { + primitive_runs(indices.as_slice::(), &mask)? + }) + }; + validate_runs(&runs, array.len(), indices.len())?; + if compact { + if let Some(result) = zero_page_take(array, indices.len(), ctx)? { + return Ok(Some(result)); + } + // Compact ranges/repeats remain compact. Only genuine strided requests carry + // per-position selection work; output cardinality alone is not setup work. + let random = runs + .iter() + .filter(|run| run.stride > 1) + .map(|run| run.length) + .sum(); + if descriptor_bytes(random) >= canonical_bytes(array) && random != 0 { + return Ok(None); + } + } + let nullable = array.dtype().nullability() | indices.dtype().nullability(); + let result = match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType => { + gather::(array, &runs, indices.len(), nullable, grouped, ctx)? + }); + Ok(Some(result.into_array())) + } +} + +struct PlainSummary { + defined: usize, + contiguous: bool, +} + +fn validate_extrema( + indices: &[I], + mask: &Mask, + extrema: &MinMaxResult, + source_len: usize, +) -> VortexResult +where + usize: TryFrom, +{ + let min = extrema + .min + .as_primitive() + .try_typed_value::()? + .ok_or_else(|| vortex_err!("Pco index minimum is null"))?; + let max = extrema + .max + .as_primitive() + .try_typed_value::()? + .ok_or_else(|| vortex_err!("Pco index maximum is null"))?; + checked_index(min)?; + let max = checked_index(max)?; + vortex_ensure!( + max < source_len, + "Pco take index {max} exceeds array length {source_len}" + ); + // Bounds/error category is guaranteed above. No contract chooses the first invalid row. + // Contiguous detection exits at the first mismatch; null physical values are never inspected. + let contiguous = match mask { + Mask::AllTrue(_) => indices.windows(2).all(|pair| { + checked_index(pair[0]) + .ok() + .and_then(|row| row.checked_add(1)) + == checked_index(pair[1]).ok() + }), + _ => false, + }; + Ok(PlainSummary { + defined: mask.true_count(), + contiguous, + }) +} + +fn zero_page_take( + array: ArrayView<'_, Pco>, + len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + if !array.page_index()?.pages.is_empty() { + return Ok(None); + } + let mask = array + .unsliced_validity() + .execute_mask(array.unsliced_n_rows(), ctx)?; + vortex_ensure_eq!( + mask.true_count(), + 0, + "Pco zero-page source has defined rows" + ); + Ok(Some( + ConstantArray::new(Scalar::null(array.dtype().as_nullable()), len).into_array(), + )) +} + +fn canonical_bytes(array: ArrayView<'_, Pco>) -> usize { + // Fallback canonicalizes the visible array, not its potentially much larger owner. + array + .len() + .saturating_mul(array.dtype().as_ptype().byte_width()) +} + +fn descriptor_bytes(requests: usize) -> usize { + // Actual prospective sparse-path storage: coalesced Run upper bound, sorting pair, + // and resolved (page, compact index, repeated-request range). Output storage is common. + requests.saturating_mul( + size_of::() + size_of::<(usize, usize)>() + size_of::<(usize, usize, Range)>(), + ) +} + +/// Borrow validated primitive inputs during planning. Broad requests stop once selected +/// physical values plus prospective descriptors exceed the visible canonical values buffer. +/// Only the chosen sparse strategy allocates Runs/sorting descriptors; no source-sized bitmap. +fn prefer_canonical_primitive( + array: ArrayView<'_, Pco>, + indices: &[I], + mask: &Mask, + requests: usize, +) -> VortexResult +where + usize: TryFrom, +{ + let index = array.page_index()?; + if !index.logical { + return Ok(requests >= index.pages.len()); + } + let full_bytes = canonical_bytes(array); + let setup = descriptor_bytes(requests); + let width = array.dtype().as_ptype().byte_width(); + let mut touched = BTreeSet::new(); + let mut selected = 0usize; + for (&row, valid) in indices.iter().zip(mask.iter()) { + if !valid { + continue; + } + let row = array.slice_start() + checked_index(row)?; + let page = index + .pages + .partition_point(|page| page.rows.as_ref().is_some_and(|rows| rows.end <= row)); + if touched.insert(page) { + let page = index + .pages + .get(page) + .ok_or_else(|| vortex_err!("Pco selected row has no logical page"))?; + selected = selected.saturating_add(page.values.len()); + if setup.saturating_add(selected.saturating_mul(width)) >= full_bytes { + return Ok(true); + } + } + } + Ok(false) +} + +fn checked_index(value: I) -> VortexResult +where + usize: TryFrom, +{ + usize::try_from(value).map_err(|_| vortex_err!("Pco take index is negative or exceeds usize")) +} + +fn primitive_runs(indices: &[I], mask: &Mask) -> VortexResult> +where + usize: TryFrom, +{ + let mut runs: Vec = Vec::with_capacity(indices.len().min(64)); + for (&value, valid) in indices.iter().zip(mask.iter()) { + let start = if valid { + Some(checked_index(value)?) + } else { + None + }; + if let Some(previous) = runs.last_mut() { + let contiguous = match (previous.start, start) { + (Some(first), Some(next)) => { + previous.stride == 1 && first.checked_add(previous.length) == Some(next) + } + (None, None) => true, + _ => false, + }; + if contiguous { + previous.length += 1; + continue; + } + } + runs.push(Run { + start, + length: 1, + stride: 1, + }); + } + Ok(runs) +} + +fn primitive_usizes(array: &PrimitiveArray) -> VortexResult> { + match_each_integer_ptype!(array.ptype(), |I| { + array + .as_slice::() + .iter() + .map(|&value| checked_index(value)) + .collect() + }) +} + +fn validate_runs(runs: &[Run], source_len: usize, output_len: usize) -> VortexResult<()> { + let mut total = 0usize; + for run in runs { + total = total + .checked_add(run.length) + .ok_or_else(|| vortex_err!("Pco take output length overflows usize"))?; + if let Some(start) = run.start + && run.length != 0 + { + let last = (run.length - 1) + .checked_mul(run.stride) + .and_then(|delta| start.checked_add(delta)) + .ok_or_else(|| vortex_err!("Pco take range overflows usize"))?; + vortex_ensure!( + last < source_len, + "Pco take index {last} exceeds array length {source_len}" + ); + } + } + vortex_ensure_eq!( + total, + output_len, + "Pco take expanded length does not match declared length" + ); + Ok(()) +} + +fn gather( + array: ArrayView<'_, Pco>, + runs: &[Run], + output_len: usize, + nullability: Nullability, + grouped: Option, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if let [ + Run { + start: Some(start), + length, + stride: 1, + }, + ] = runs + { + return stream_range::(array, *start..*start + *length, nullability, ctx); + } + if array.page_index()?.pages.is_empty() { + let mask = array + .unsliced_validity() + .execute_mask(array.unsliced_n_rows(), ctx)?; + vortex_ensure_eq!( + mask.true_count(), + 0, + "Pco zero-page source has defined rows" + ); + return Ok(PrimitiveArray::new( + BufferMut::::zeroed_in(output_len, ctx.allocator().clone()).freeze(), + Validity::AllInvalid, + )); + } + if nullability.is_nullable() { + if let Some(defined) = grouped { + gather_grouped::(array, runs, output_len, defined, ctx) + } else { + gather_typed::(array, runs, output_len, ctx) + } + } else if let Some(defined) = grouped { + gather_grouped::(array, runs, output_len, defined, ctx) + } else { + gather_typed::(array, runs, output_len, ctx) + } +} + +/// Sort requests once, resolve each repeated logical position once, then invoke the native +/// primitive gather once per physical page. Only the final scatter visits output positions; +/// it does no page-index search, validity rank, decoder lookup, or cache lookup. +fn gather_grouped( + array: ArrayView<'_, Pco>, + runs: &[Run], + output_len: usize, + defined: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut requests = Vec::with_capacity(defined); + let mut output = 0usize; + for run in runs { + if let Some(start) = run.start { + for offset in 0..run.length { + requests.push(( + array.slice_start() + start + offset * run.stride, + output + offset, + )); + } + } + output += run.length; + } + requests.sort_unstable(); + let mut state = AccessState::::default(); + let mut resolved = Vec::with_capacity(requests.len().min(64)); + let mut cursor = 0; + while cursor < requests.len() { + let row = requests[cursor].0; + let end = cursor + requests[cursor..].partition_point(|request| request.0 == row); + if let Some((page, value)) = state.resolve(array, row, ctx)? { + resolved.push((page, value, cursor..end)); + } + cursor = end; + } + let mut values = BufferMut::::zeroed_in(output_len, ctx.allocator().clone()); + let mut valid = BitBufferMut::new_unset(if NULLABLE { output_len } else { 0 }); + let mut cursor = 0; + while cursor < resolved.len() { + let page = resolved[cursor].0; + let end = cursor + resolved[cursor..].partition_point(|request| request.0 == page); + let compact = PrimitiveArray::from_iter( + resolved[cursor..end] + .iter() + .map(|request| { + u64::try_from(request.1) + .map_err(|_| vortex_err!("Pco compact index exceeds u64")) + }) + .collect::>>()?, + ) + .into_array(); + let decoded = state.values(array, page, ctx)?.clone(); + let selected = decoded.take(compact)?.execute::(ctx)?; + for (request, &value) in resolved[cursor..end].iter().zip(selected.as_slice::()) { + for &(_, output) in &requests[request.2.clone()] { + values[output] = value; + if NULLABLE { + valid.set(output); + } + } + } + cursor = end; + } + let validity = if NULLABLE { + Validity::from(valid.freeze()) + } else { + Validity::NonNullable + }; + Ok(PrimitiveArray::new(values.freeze(), validity)) +} + +/// One initialized final output buffer. Nullable outputs leave null positions at the +/// allocator's safe zero initialization instead of eagerly touching every null gap. +struct Output { + values: BufferMut, + len: usize, +} + +impl Output { + fn new(len: usize, ctx: &ExecutionCtx) -> Self { + let values = if NULLABLE { + BufferMut::zeroed_in(len, ctx.allocator().clone()) + } else { + BufferMut::with_capacity_in(len, ctx.allocator().clone()) + }; + Self { values, len: 0 } + } + + fn push(&mut self, value: T) { + if NULLABLE { + self.values[self.len] = value; + } else { + self.values.push(value); + } + self.len += 1; + } + + fn extend_from_slice(&mut self, values: &[T]) { + if NULLABLE { + self.values[self.len..self.len + values.len()].copy_from_slice(values); + } else { + self.values.extend_from_slice(values); + } + self.len += values.len(); + } + + fn repeat(&mut self, value: T, len: usize) { + if NULLABLE { + self.values[self.len..self.len + len].fill(value); + } else { + self.values.extend(std::iter::repeat_n(value, len)); + } + self.len += len; + } + + fn nulls(&mut self, len: usize) { + if NULLABLE { + self.len += len; + } else { + self.repeat(T::default(), len); + } + } +} + +fn gather_typed( + array: ArrayView<'_, Pco>, + runs: &[Run], + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut state = AccessState::::default(); + let mut values = Output::::new(output_len, ctx); + // The non-nullable monomorphization has no bitmap allocation or append work. + let mut valid = BitBufferMut::with_capacity(if NULLABLE { output_len } else { 0 }); + for run in runs { + let Some(start) = run.start else { + if NULLABLE { + valid.append_n(false, run.length); + } + values.nulls(run.length); + continue; + }; + if run.stride == 0 && run.length != 0 { + if let Some((page, value)) = state.resolve(array, array.slice_start() + start, ctx)? { + let value = state.values(array, page, ctx)?.as_slice::()[value]; + values.repeat(value, run.length); + if NULLABLE { + valid.append_n(true, run.length); + } + } else { + if NULLABLE { + valid.append_n(false, run.length); + } + values.nulls(run.length); + } + continue; + } + if run.stride == 1 && run.length != 0 { + let mut row = array.slice_start() + start; + let stop = row + run.length; + while row < stop { + if let Some(span) = state.span(array, row, stop, ctx)? { + match &span.mask { + Mask::AllTrue(n) => { + values.extend_from_slice( + &state.values(array, span.page, ctx)?.as_slice::()[span.values], + ); + if NULLABLE { + valid.append_n(true, *n); + } + } + Mask::AllFalse(n) => { + values.nulls(*n); + if NULLABLE { + valid.append_n(false, *n); + } + } + Mask::Values(mask) => { + let source = state.values(array, span.page, ctx)?.as_slice::(); + let source = &source[span.values]; + let mut logical = 0usize; + let mut compact = 0usize; + // Copy defined runs straight into the final output. Null gaps are + // initialized once; no compact staging output is expanded afterward. + for (start, stop) in mask.bit_buffer().set_slices() { + values.nulls(start - logical); + let count = stop - start; + values.extend_from_slice(&source[compact..compact + count]); + compact += count; + logical = stop; + } + values.nulls(mask.len() - logical); + if NULLABLE { + valid.append_buffer(mask.bit_buffer()); + } + } + } + row = span.stop; + } else { + if let Some((page, value)) = state.resolve(array, row, ctx)? { + values.push(state.values(array, page, ctx)?.as_slice::()[value]); + if NULLABLE { + valid.append_true(); + } + } else { + values.nulls(1); + if NULLABLE { + valid.append_false(); + } + } + row += 1; + } + } + continue; + } + for offset in 0..run.length { + // The complete descriptor, including its output length and last element, was + // checked before allocating output. The source slice itself is already validated. + let row = array.slice_start() + start + offset * run.stride; + if let Some((page, value)) = state.resolve(array, row, ctx)? { + values.push(state.values(array, page, ctx)?.as_slice::()[value]); + if NULLABLE { + valid.append_true(); + } + } else { + values.nulls(1); + if NULLABLE { + valid.append_false(); + } + } + } + } + let validity = if NULLABLE { + Validity::from(valid.freeze()) + } else { + Validity::NonNullable + }; + vortex_ensure_eq!( + values.len, + output_len, + "Pco selected output length mismatch" + ); + Ok(PrimitiveArray::new(values.values.freeze(), validity)) +} + +/// Contiguous selections stream pages into the final initialized values buffer and retain +/// exact sliced source validity. No output validity copy or retained decoded-page cache. +fn stream_range( + array: ArrayView<'_, Pco>, + range: Range, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + vortex_ensure!( + range.start <= range.end && range.end <= array.len(), + "Pco contiguous range exceeds visible bounds" + ); + let global_start = array.slice_start() + range.start; + let mut values = BufferMut::::zeroed_in(range.len(), ctx.allocator().clone()); + let mut scratch = BufferMut::::with_capacity_in(0, ctx.allocator().clone()); + let mut state = AccessState::::default(); + let mut decoder = PageDecoder::::new(array.metadata.header.as_slice())?; + let stop = array.slice_start() + range.end; + let mut row = global_start; + while row < stop { + let span = state + .span(array, row, stop, ctx)? + .ok_or_else(|| vortex_err!("Pco contiguous page span is missing"))?; + vortex_ensure!( + span.stop > row && span.stop <= stop, + "Pco contiguous span does not advance within bounds" + ); + vortex_ensure_eq!( + span.mask.len(), + span.stop - row, + "Pco contiguous mask length mismatch" + ); + vortex_ensure_eq!( + span.values.len(), + span.mask.true_count(), + "Pco contiguous value count mismatch" + ); + let start = row - global_start; + let output = values + .get_mut(start..start + span.mask.len()) + .ok_or_else(|| vortex_err!("Pco contiguous output exceeds visible range"))?; + if !span.mask.all_false() { + let index = array.page_index()?; + let page = index + .pages + .get(span.page) + .ok_or_else(|| vortex_err!("Pco contiguous page is missing"))?; + let buffer = array + .pages + .get(span.page) + .ok_or_else(|| vortex_err!("Pco compressed page is missing"))?; + let chunk = array + .chunk_metas + .get(page.chunk) + .ok_or_else(|| vortex_err!("Pco chunk metadata is missing"))?; + let n = page.values.len(); + if span.mask.all_true() && span.values == (0..n) { + decoder.read_into(page.chunk, chunk.as_slice(), n, buffer.as_slice(), output)?; + } else { + if scratch.len() < n { + scratch.push_n(T::default(), n - scratch.len()); + } + decoder.read_into( + page.chunk, + chunk.as_slice(), + n, + buffer.as_slice(), + &mut scratch[..n], + )?; + let source = scratch + .get(span.values) + .ok_or_else(|| vortex_err!("Pco compact range exceeds page output"))?; + match &span.mask { + Mask::AllTrue(_) => output.copy_from_slice(source), + Mask::AllFalse(_) => {} + Mask::Values(mask) => { + let mut compact = 0; + for (start, end) in mask.bit_buffer().set_slices() { + let len = end - start; + output[start..end].copy_from_slice(&source[compact..compact + len]); + compact += len; + } + } + } + } + } + row = span.stop; + } + let validity = array + .unsliced_validity() + .slice(global_start..stop)? + .union_nullability(nullability); + Ok(PrimitiveArray::new(values.freeze(), validity)) +} + +pub(crate) fn slice( + array: ArrayView<'_, Pco>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType => { + stream_range::(array, 0..array.len(), array.dtype().nullability(), ctx) + }) +} diff --git a/vortex-file/examples/pco_access_profile.rs b/vortex-file/examples/pco_access_profile.rs new file mode 100644 index 00000000000..0572e4f451e --- /dev/null +++ b/vortex-file/examples/pco_access_profile.rs @@ -0,0 +1,398 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Public API resource harness. Run identical code on a candidate and a historical checkout. +//! Arguments: operation (write/full/slice/take/ranges/probe), pattern, rows, repetitions. +//! JSON lines retain cold and warmed iterations; Linux process accounting is read outside timing. + +use std::hint::black_box; +use std::time::Instant; + +use prost::Message; +use vortex_array::ArrayContext; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PiecewiseSequenceArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::buffer::BufferHandle; +use vortex_array::serde::SerializeOptions; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_pco::Pco; +use vortex_runend::RunEnd; +use vortex_session::VortexSession; + +#[derive(Clone, prost::Message)] +struct LegacyMetadata { + #[prost(bytes, tag = "1")] + header: Vec, + #[prost(message, repeated, tag = "2")] + chunks: Vec, +} + +fn legacy(array: &ArrayRef, session: &VortexSession) -> VortexResult { + let wire = session + .array_serialize(array)? + .vortex_expect("PCO serialization"); + let metadata = LegacyMetadata::decode(wire.metadata.as_slice())?.encode_to_vec(); + let buffers = wire + .buffers + .into_iter() + .map(BufferHandle::new_host) + .collect::>(); + ArrayPlugin::deserialize( + &Pco, + ArrayDeserialization::new( + wire.serialized_id, + array.dtype(), + array.len(), + &metadata, + &buffers, + &wire.children, + ), + session, + ) +} + +fn linux_accounting() -> (Option, Option) { + let ticks = std::fs::read_to_string("/proc/self/stat") + .ok() + .and_then(|stat| { + let fields = stat + .rsplit_once(')')? + .1 + .split_whitespace() + .collect::>(); + Some(fields.get(11)?.parse::().ok()? + fields.get(12)?.parse::().ok()?) + }); + let rss = std::fs::read_to_string("/proc/self/status") + .ok() + .and_then(|status| { + status + .lines() + .find(|line| line.starts_with("VmRSS:"))? + .split_whitespace() + .nth(1)? + .parse() + .ok() + }); + (ticks, rss) +} + +fn main() -> VortexResult<()> { + let args = std::env::args().collect::>(); + let mode = args.get(1).map(String::as_str).unwrap_or("slice"); + let pattern = args.get(2).map(String::as_str).unwrap_or("nullable"); + let rows = args + .get(3) + .map_or(Ok(8_388_608), |value| value.parse::()) + .map_err(|error| vortex_err!("invalid row count: {error}"))?; + let repetitions = args + .get(4) + .map_or(Ok(8), |value| value.parse::()) + .map_err(|error| vortex_err!("invalid repetition count: {error}"))?; + vortex_ensure!(rows >= 8192, "profile requires at least 8192 rows"); + let session = VortexSession::empty(); + vortex_file::register_default_encodings(&session); + let mut ctx = session.create_execution_ctx(); + let (validity_pattern, legacy_format) = pattern + .strip_suffix("-legacy") + .map_or((pattern, false), |pattern| (pattern, true)); + let validity = match validity_pattern { + "dense" => Validity::NonNullable, + "all-null" => Validity::AllInvalid, + "long-null" => Validity::from_iter((0..rows).map(|i| !(2048..=rows - 4096).contains(&i))), + "compressed" => { + let runs = rows.div_ceil(32768); + let ends = PrimitiveArray::from_iter((0..runs).map(|run| { + u64::try_from(((run + 1) * 32768).min(rows)).vortex_expect("run end fits u64") + })) + .into_array(); + let values = BoolArray::from_iter((0..runs).map(|run| run % 2 != 0)).into_array(); + Validity::Array(RunEnd::try_new(ends, values, &mut ctx)?.into_array()) + } + "nullable" => Validity::from_iter((0..rows).map(|i| i % 7 != 0)), + _ => return Err(vortex_err!("unknown pattern: {pattern}")), + }; + let input = PrimitiveArray::new( + (0..rows) + .map(|i| i64::try_from(i / 16).vortex_expect("fixture fits i64")) + .collect::>(), + validity, + ); + let encoded = if mode == "write" { + input.clone().into_array() + } else { + Pco::from_primitive( + input.as_view(), + 3, + if mode == "metadata" { 8192 } else { 1024 }, + &mut ctx, + )? + .into_array() + }; + let encoded = if legacy_format && mode != "write" { + legacy(&encoded, &session)? + } else { + encoded + }; + if mode == "metadata" { + let wire = session + .array_serialize(&encoded)? + .vortex_expect("PCO metadata"); + let stripped = LegacyMetadata::decode(wire.metadata.as_slice())?.encode_to_vec(); + let legacy = legacy(&encoded, &session)?; + let context = ArrayContext::empty(); + let options = SerializeOptions::default(); + let current = encoded.serialize(&context, &session, &options)?; + let historical = legacy.serialize(&context, &session, &options)?; + println!( + "{{\"rows\":{rows},\"page_values_limit\":8192,\"metadata_bytes\":{},\"legacy_metadata_bytes\":{},\"serialized_bytes\":{},\"legacy_serialized_bytes\":{},\"serialized_buffer_sizes\":[{}],\"legacy_serialized_buffer_sizes\":[{}]}}", + wire.metadata.len(), + stripped.len(), + current.iter().map(|buffer| buffer.len()).sum::(), + historical.iter().map(|buffer| buffer.len()).sum::(), + current + .iter() + .map(|buffer| buffer.len().to_string()) + .collect::>() + .join(","), + historical + .iter() + .map(|buffer| buffer.len().to_string()) + .collect::>() + .join(",") + ); + return Ok(()); + } + let take_indices = PrimitiveArray::from_iter( + (0..128usize).map(|i| u64::try_from((i * 104729) % rows).vortex_expect("index fits u64")), + ) + .into_array(); + let range_starts = (0..16usize) + .map(|i| u64::try_from((i * (rows / 17)) % (rows - 1024)).vortex_expect("range fits u64")) + .collect::>(); + let range_indices = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter(range_starts).into_array(), + PrimitiveArray::from_iter(vec![512u64; 16]).into_array(), + PrimitiveArray::from_iter(vec![1u64; 16]).into_array(), + 8192, + )? + .into_array(); + let large_indices = PrimitiveArray::from_iter( + (0..if mode == "take-large" { rows / 8 } else { 0 }) + .map(|i| u64::try_from((i * 104729) % rows).vortex_expect("index fits u64")), + ) + .into_array(); + let repeated_indices = + PrimitiveArray::from_iter((0..if mode == "take-repeated" { rows / 8 } else { 0 }).map( + |i| u64::try_from(((i % 16) * (rows / 17)) % rows).vortex_expect("index fits u64"), + )) + .into_array(); + let mixed_indices = PrimitiveArray::from_iter( + (0..if mode == "take-mixed" { rows / 8 } else { 0 }).map(|i| { + u64::try_from(if i < 2 { i } else { (i * 104729) % rows }) + .vortex_expect("index fits u64") + }), + ) + .into_array(); + let narrow_indices = PrimitiveArray::from_iter( + (0..if mode == "take-narrow" { + 100000usize + } else { + 0 + }) + .map(|i| if i % 2 == 0 { 0u64 } else { 127 }), + ) + .into_array(); + let repeated_spans = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter([ + 1u64, + u64::try_from(rows - 1).vortex_expect("index fits u64"), + ]) + .into_array(), + PrimitiveArray::from_iter([500000u64, 500000]).into_array(), + PrimitiveArray::from_iter([0u64, 0]).into_array(), + 1000000, + )? + .into_array(); + let contiguous_indices = PrimitiveArray::from_iter( + (0..if mode == "take-contiguous" { + rows - 2 + } else { + 0 + }) + .map(|i| u64::try_from(i + 1).vortex_expect("index fits u64")), + ) + .into_array(); + let contiguous_range = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter([1u64]).into_array(), + PrimitiveArray::from_iter([u64::try_from(rows - 2).vortex_expect("range fits u64")]) + .into_array(), + PrimitiveArray::from_iter([1u64]).into_array(), + rows - 2, + )? + .into_array(); + let wire = session + .array_serialize(&encoded)? + .vortex_expect("prepared array serialization"); + let mut probe = encoded.repeated_probe(); + for iteration in 0..repetitions { + let before = linux_accounting(); + let start = Instant::now(); + let output: Option = match mode { + "write" => Some(Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array()), + "full" => Some( + encoded + .clone() + .execute::(&mut ctx)? + .into_array(), + ), + "slice" => Some( + encoded + .slice(rows - 4096..rows - 2048)? + .execute::(&mut ctx)? + .into_array(), + ), + "take" => Some( + encoded + .take(take_indices.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "ranges" => Some( + encoded + .take(range_indices.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "slice-near-full" => Some( + encoded + .slice(1..rows - 1)? + .execute::(&mut ctx)? + .into_array(), + ), + "take-contiguous" => Some( + encoded + .take(contiguous_indices.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "take-range" => Some( + encoded + .take(contiguous_range.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "take-mixed" => Some( + encoded + .take(mixed_indices.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "take-narrow" => Some( + encoded + .slice(rows - 128..rows)? + .take(narrow_indices.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "take-spans" => Some( + encoded + .take(repeated_spans.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "take-tiny" => Some( + encoded + .take( + PrimitiveArray::from_iter([ + 0u64, + u64::try_from(rows - 1).vortex_expect("index fits u64"), + ]) + .into_array(), + )? + .execute::(&mut ctx)? + .into_array(), + ), + "take-large" => Some( + encoded + .take(large_indices.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "take-repeated" => Some( + encoded + .take(repeated_indices.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "open-only" | "open-first" => { + // Array-plugin open with compressed buffers and the validity child already + // available: includes protobuf decode/validation/index preparation, no file IO. + let buffers = wire + .buffers + .iter() + .cloned() + .map(BufferHandle::new_host) + .collect::>(); + let opened = ArrayPlugin::deserialize( + &Pco, + ArrayDeserialization::new( + wire.serialized_id, + encoded.dtype(), + encoded.len(), + &wire.metadata, + &buffers, + &wire.children, + ), + &session, + )?; + if mode == "open-first" { + black_box(opened.execute_scalar(rows - 2048, &mut ctx)?); + } + black_box(opened); + None + } + "probe-first" => { + black_box(encoded.probe().execute_scalar(rows - 2048, &mut ctx)?); + None + } + "probe" => { + for i in 0..1024 { + black_box(probe.execute_scalar(rows - 2048 + i % 256, &mut ctx)?); + } + None + } + _ => return Err(vortex_err!("unknown operation: {mode}")), + }; + black_box(&output); + let seconds = start.elapsed().as_secs_f64(); + let after = linux_accounting(); + // Keep results live and inspect a public scalar after the measurement. Independent value + // equivalence is tested by the public semantics/file suites, not by timing assertions. + if let Some(output) = output { + black_box(output.execute_scalar(0, &mut ctx)?); + } + let ticks = before + .0 + .zip(after.0) + .map(|(before, after)| after.saturating_sub(before)); + let ticks = ticks.map_or_else(|| "null".to_owned(), |value| value.to_string()); + let rss = after + .1 + .map_or_else(|| "null".to_owned(), |value| value.to_string()); + println!( + "{{\"operation\":\"{mode}\",\"pattern\":\"{pattern}\",\"rows\":{rows},\"iteration\":{iteration},\"wall_seconds\":{seconds},\"cpu_ticks\":{ticks},\"rss_kib\":{rss}}}" + ); + } + Ok(()) +} diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index 1c309112078..6d589461afa 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -129,7 +129,6 @@ pub use strategy::*; use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::session::ArraySessionExt; -use vortex_pco::Pco; use vortex_session::VortexSession; pub use writer::*; @@ -169,6 +168,7 @@ mod forever_constant { /// Registration covers reading: a session can decode every encoding registered here. The /// writer is gated separately by the editions enabled on its session. pub fn register_default_encodings(session: &VortexSession) { + vortex_pco::initialize(session); vortex_bytebool::initialize(session); vortex_fsst::initialize(session); vortex_onpair::initialize(session); @@ -178,7 +178,6 @@ pub fn register_default_encodings(session: &VortexSession) { { let arrays = session.arrays(); - arrays.register(Pco); if use_experimental_patches() { arrays.register(Patched); } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 7e15c104226..c4312825f92 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -174,7 +174,8 @@ async fn test_eof_values() { feature = "zstd", case::compact( BtrBlocksCompressorBuilder::from_session(&SESSION).with_compact(), - 55_344 + // PCO's optional logical-page metadata adds 64 bytes to this compact fixture. + 55_408 ) )] #[tokio::test] diff --git a/vortex-file/tests/pco_logical_pages_compat.rs b/vortex-file/tests/pco_logical_pages_compat.rs new file mode 100644 index 00000000000..6ff657ab51f --- /dev/null +++ b/vortex-file/tests/pco_logical_pages_compat.rs @@ -0,0 +1,194 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Public file harness also run, unchanged, against the pre-change reader by PCO compatibility CI. + +#![expect(clippy::tests_outside_test_module)] + +use std::sync::Arc; + +use futures::TryStreamExt; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::assert_arrays_eq; +use vortex_array::validity::Validity; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_file::OpenOptionsSessionExt; +use vortex_file::WriteOptionsSessionExt; +use vortex_io::session::RuntimeSession; +use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex_layout::session::LayoutSession; +use vortex_pco::Pco; +use vortex_runend::RunEnd; +use vortex_session::VortexSession; + +mod common; + +#[tokio::test] +async fn public_pco_compatibility() -> VortexResult<()> { + let session = VortexSession::empty() + .with::() + .with::(); + vortex_file::register_default_encodings(&session); + common::enable_all_registered_array_encodings(&session); + let mut ctx = session.create_execution_ctx(); + let original = PrimitiveArray::from_option_iter( + (0..131_072i64).map(|i| (!(1024..=120_000).contains(&i) && i % 7 != 0).then_some(i * 19)), + ); + let mode = std::env::var("PCO_COMPAT_MODE").unwrap_or_else(|_| "roundtrip".to_owned()); + let directory = std::env::var_os("PCO_COMPAT_DIR").map(std::path::PathBuf::from); + for (name, range) in [("unsliced", 0..131_072), ("normalized_slice", 777..125_000)] { + let encoded = Pco::from_primitive(original.as_view(), 3, 127, &mut ctx)?.into_array(); + // Explicit writer-owned normalization, never a hidden serde fallback: execute selected + // pages and recompress only the chosen logical rows using the caller's configuration. + let selected = encoded + .slice(range.clone())? + .execute::(&mut ctx)?; + let normalized = Pco::from_primitive(selected.as_view(), 3, 127, &mut ctx)?.into_array(); + let expected = + StructArray::from_fields(&[("values", original.slice(range)?)])?.into_array(); + let data = StructArray::from_fields(&[("values", normalized)])?.into_array(); + let path = directory + .as_ref() + .map(|directory| directory.join(format!("{name}.vortex"))); + let bytes = if mode == "read" { + std::fs::read( + path.as_ref() + .vortex_expect("read mode requires PCO_COMPAT_DIR"), + )? + } else { + let mut bytes = Vec::new(); + session + .write_options() + .with_strategy(Arc::new(FlatLayoutStrategy::default())) + .write(&mut bytes, data.to_array_stream()) + .await?; + if let Some(path) = &path { + std::fs::create_dir_all(path.parent().vortex_expect("fixture path has a parent"))?; + std::fs::write(path, &bytes)?; + } + bytes + }; + let file = session + .open_options() + .open_buffer(ByteBuffer::from(bytes))?; + let chunks = file.scan()?.into_stream()?.try_collect::>().await?; + let actual = ChunkedArray::try_new(chunks, expected.dtype().clone())?.into_array(); + assert_arrays_eq!(actual, expected, &mut ctx); + } + Ok(()) +} + +#[test] +fn compressed_validity_selection_oracle() -> VortexResult<()> { + let session = VortexSession::empty(); + vortex_file::register_default_encodings(&session); + let mut ctx = session.create_execution_ctx(); + let validity = Validity::Array( + RunEnd::try_new( + PrimitiveArray::from_iter([32768u64, 65536, 98304, 131072]).into_array(), + BoolArray::from_iter([false, true, false, true]).into_array(), + &mut ctx, + )? + .into_array(), + ); + let input = PrimitiveArray::new((0..131072i64).collect::>(), validity); + let expected = + PrimitiveArray::from_option_iter((0..131072i64).map(|i| (i / 32768 % 2 == 1).then_some(i))); + let encoded = Pco::from_primitive(input.as_view(), 3, 127, &mut ctx)?.into_array(); + let indices = + PrimitiveArray::from_iter([0u32, 32767, 32768, 65535, 65536, 98304, 131071, 32768]) + .into_array(); + assert_arrays_eq!( + encoded.take(indices.clone())?, + expected.take(indices)?, + &mut ctx + ); + assert_arrays_eq!( + encoded.slice(30001..130001)?, + expected.slice(30001..130001)?, + &mut ctx + ); + let mut probe = encoded.repeated_probe(); + for index in [0, 32767, 32768, 65535, 65536, 98304, 131071, 32768] { + assert_eq!( + probe.execute_scalar(index, &mut ctx)?, + expected.execute_scalar(index, &mut ctx)? + ); + } + Ok(()) +} + +/// Public byte-boundary diagnostic for the exact deterministic compact stock fixture. +/// This reports format overhead; it is not a timing or topology assertion. +#[cfg(feature = "zstd")] +#[tokio::test] +async fn public_stock_size_attribution() -> VortexResult<()> { + use rand::RngExt; + use rand::SeedableRng; + use rand::rngs::StdRng; + use vortex_array::stream::ArrayStreamExt; + use vortex_btrblocks::BtrBlocksCompressorBuilder; + use vortex_file::WriteStrategyBuilder; + + if std::env::var_os("PCO_SIZE_DIAGNOSTICS").is_none() { + return Ok(()); + } + let session = VortexSession::empty() + .with::() + .with::(); + vortex_file::register_default_encodings(&session); + common::enable_all_registered_array_encodings(&session); + let mut rng = StdRng::seed_from_u64(0); + let input = PrimitiveArray::from_iter((0..100000i64).map(|i| rng.random_range(i..=i + 10))) + .into_array(); + let mut bytes = Vec::new(); + session + .write_options() + .with_strategy( + WriteStrategyBuilder::from_session(&session) + .with_btrblocks_builder( + BtrBlocksCompressorBuilder::from_session(&session).with_compact(), + ) + .build(), + ) + .write(&mut bytes, input.clone().to_array_stream()) + .await?; + let len = bytes.len(); + let file = session + .open_options() + .open_buffer(ByteBuffer::from(bytes))?; + let segments = file.footer().segment_map(); + let payload = segments + .iter() + .map(|segment| u64::from(segment.length)) + .sum::(); + let end = segments + .iter() + .map(|segment| segment.offset + u64::from(segment.length)) + .max() + .unwrap_or(0); + println!( + "{{\"stock_file_bytes\":{len},\"segment_payload_bytes\":{payload},\"last_segment_end\":{end},\"trailing_format_bytes\":{},\"segments\":[{}]}}", + u64::try_from(len).vortex_expect("file fits u64") - end, + segments + .iter() + .map(|segment| format!( + "[{},{},{}]", + segment.offset, + segment.length, + segment.alignment.as_usize() + )) + .collect::>() + .join(",") + ); + let actual = file.scan()?.into_array_stream()?.read_all().await?; + assert_arrays_eq!(actual, input, &mut session.create_execution_ctx()); + Ok(()) +}