Skip to main content

vortex_pco/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::cmp;
5use std::fmt::Debug;
6use std::fmt::Display;
7use std::fmt::Formatter;
8use std::hash::Hash;
9use std::hash::Hasher;
10
11use pco::ChunkConfig;
12use pco::PagingSpec;
13use pco::data_types::Number;
14use pco::data_types::NumberType;
15use pco::errors::PcoError;
16use pco::match_number_enum;
17use pco::wrapped::ChunkDecompressor;
18use pco::wrapped::FileCompressor;
19use pco::wrapped::FileDecompressor;
20use prost::Message;
21use vortex_array::Array;
22use vortex_array::ArrayEq;
23use vortex_array::ArrayHash;
24use vortex_array::ArrayId;
25use vortex_array::ArrayParts;
26use vortex_array::ArrayRef;
27use vortex_array::ArrayView;
28use vortex_array::EqMode;
29use vortex_array::ExecutionCtx;
30use vortex_array::ExecutionResult;
31use vortex_array::IntoArray;
32use vortex_array::TypedArrayRef;
33use vortex_array::array_slots;
34use vortex_array::arrays::Primitive;
35use vortex_array::arrays::PrimitiveArray;
36use vortex_array::buffer::BufferHandle;
37use vortex_array::dtype::DType;
38use vortex_array::dtype::PType;
39use vortex_array::dtype::half;
40use vortex_array::scalar::Scalar;
41use vortex_array::serde::ArrayChildren;
42use vortex_array::validity::Validity;
43use vortex_array::vtable::OperationsVTable;
44use vortex_array::vtable::VTable;
45use vortex_array::vtable::ValidityVTable;
46use vortex_array::vtable::child_to_validity;
47use vortex_array::vtable::validity_to_child;
48use vortex_buffer::BufferMut;
49use vortex_buffer::ByteBuffer;
50use vortex_buffer::ByteBufferMut;
51use vortex_error::VortexError;
52use vortex_error::VortexResult;
53use vortex_error::vortex_bail;
54use vortex_error::vortex_ensure;
55use vortex_error::vortex_err;
56use vortex_session::VortexSession;
57use vortex_session::registry::CachedId;
58
59use crate::PcoChunkInfo;
60use crate::PcoMetadata;
61use crate::PcoPageInfo;
62
63// Overall approach here:
64// Chunk the array into Pco chunks (currently using the default recommended size
65// for good compression), and into finer-grained Pco pages. As we go, write each
66// ChunkMeta as a buffer, followed by each of that chunk's pages as a buffer. We
67// store metadata for each of these "components" (chunk or page). At
68// decompression time, we figure out which components we need to read and only
69// process those. We only compress and decompress valid values.
70
71// Visually, during decompression, we have an interval of pages we're
72// decompressing and a tighter interval of the slice we actually care about.
73// |=============values (all valid elements)==============|
74// |<-n_skipped_values->|----decompressed_values------|
75//                          |----slice_values----|
76//                          ^                    ^
77// |<---slice_value_start-->|<--slice_n_values-->|
78// We then insert these values to the correct position using a primitive array
79// constructor.
80
81const VALUES_PER_CHUNK: usize = pco::DEFAULT_MAX_PAGE_N;
82
83/// A [`Pco`]-encoded Vortex array.
84pub type PcoArray = Array<Pco>;
85
86impl ArrayHash for PcoData {
87    fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: EqMode) {
88        self.unsliced_n_rows.hash(state);
89        self.slice_start.hash(state);
90        self.slice_stop.hash(state);
91        // Hash chunk_metas and pages using pointer-based hashing
92        for chunk_meta in &self.chunk_metas {
93            chunk_meta.array_hash(state, accuracy);
94        }
95        for page in &self.pages {
96            page.array_hash(state, accuracy);
97        }
98    }
99}
100
101impl ArrayEq for PcoData {
102    fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool {
103        if self.unsliced_n_rows != other.unsliced_n_rows
104            || self.slice_start != other.slice_start
105            || self.slice_stop != other.slice_stop
106            || self.chunk_metas.len() != other.chunk_metas.len()
107            || self.pages.len() != other.pages.len()
108        {
109            return false;
110        }
111        for (a, b) in self.chunk_metas.iter().zip(&other.chunk_metas) {
112            if !a.array_eq(b, accuracy) {
113                return false;
114            }
115        }
116        for (a, b) in self.pages.iter().zip(&other.pages) {
117            if !a.array_eq(b, accuracy) {
118                return false;
119            }
120        }
121        true
122    }
123}
124
125impl VTable for Pco {
126    type TypedArrayData = PcoData;
127
128    type OperationsVTable = Self;
129    type ValidityVTable = Self;
130
131    fn id(&self) -> ArrayId {
132        static ID: CachedId = CachedId::new("vortex.pco");
133        *ID
134    }
135
136    fn validate(
137        &self,
138        data: &PcoData,
139        dtype: &DType,
140        len: usize,
141        slots: &[Option<ArrayRef>],
142    ) -> VortexResult<()> {
143        let validity = child_to_validity(
144            PcoSlotsView::from_slots(slots).validity,
145            dtype.nullability(),
146        );
147        data.validate(dtype, len, &validity)
148    }
149
150    fn nbuffers(array: ArrayView<'_, Self>) -> usize {
151        array.chunk_metas.len() + array.pages.len()
152    }
153
154    fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
155        if idx < array.chunk_metas.len() {
156            BufferHandle::new_host(array.chunk_metas[idx].clone())
157        } else {
158            let page_idx = idx - array.chunk_metas.len();
159            BufferHandle::new_host(array.pages[page_idx].clone())
160        }
161    }
162
163    fn buffer_name(array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
164        if idx < array.chunk_metas.len() {
165            Some(format!("chunk_meta_{idx}"))
166        } else {
167            Some(format!("page_{}", idx - array.chunk_metas.len()))
168        }
169    }
170
171    fn with_buffers(
172        &self,
173        array: ArrayView<'_, Self>,
174        buffers: &[BufferHandle],
175    ) -> VortexResult<ArrayParts<Self>> {
176        let mut data = array.data().clone();
177        let chunk_metas_len = data.metadata.chunks.len();
178        vortex_ensure!(buffers.len() >= chunk_metas_len);
179        data.chunk_metas = buffers[..chunk_metas_len]
180            .iter()
181            .map(|buffer| buffer.clone().try_to_host_sync())
182            .collect::<VortexResult<Vec<_>>>()?;
183        data.pages = buffers[chunk_metas_len..]
184            .iter()
185            .map(|buffer| buffer.clone().try_to_host_sync())
186            .collect::<VortexResult<Vec<_>>>()?;
187        Ok(
188            ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data)
189                .with_slots(array.slots().iter().cloned().collect()),
190        )
191    }
192
193    fn serialize(
194        array: ArrayView<'_, Self>,
195        _session: &VortexSession,
196    ) -> VortexResult<Option<Vec<u8>>> {
197        Ok(Some(array.metadata.clone().encode_to_vec()))
198    }
199
200    fn deserialize(
201        &self,
202        dtype: &DType,
203        len: usize,
204        metadata: &[u8],
205        buffers: &[BufferHandle],
206        children: &dyn ArrayChildren,
207        _session: &VortexSession,
208    ) -> VortexResult<ArrayParts<Self>> {
209        let metadata = PcoMetadata::decode(metadata)?;
210        let validity = if children.is_empty() {
211            Validity::from(dtype.nullability())
212        } else if children.len() == 1 {
213            let validity = children.get(0, &Validity::DTYPE, len)?;
214            Validity::Array(validity)
215        } else {
216            vortex_bail!("PcoArray expected 0 or 1 child, got {}", children.len());
217        };
218
219        vortex_ensure!(buffers.len() >= metadata.chunks.len());
220        let chunk_metas = buffers[..metadata.chunks.len()]
221            .iter()
222            .map(|b| b.clone().try_to_host_sync())
223            .collect::<VortexResult<Vec<_>>>()?;
224        let pages = buffers[metadata.chunks.len()..]
225            .iter()
226            .map(|b| b.clone().try_to_host_sync())
227            .collect::<VortexResult<Vec<_>>>()?;
228
229        let expected_n_pages = metadata
230            .chunks
231            .iter()
232            .map(|info| info.pages.len())
233            .sum::<usize>();
234        vortex_ensure!(pages.len() == expected_n_pages);
235
236        let slots = PcoSlots {
237            validity: validity_to_child(&validity, len),
238        }
239        .into_slots();
240        let data = PcoData::new(chunk_metas, pages, dtype.as_ptype(), metadata, len);
241        Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
242    }
243
244    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
245        PcoSlots::NAMES[idx].to_string()
246    }
247
248    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
249        let unsliced_validity = array.unsliced_validity();
250        Ok(ExecutionResult::done(
251            array
252                .data()
253                .decompress(&unsliced_validity, ctx)?
254                .into_array(),
255        ))
256    }
257
258    fn reduce_parent(
259        array: ArrayView<'_, Self>,
260        parent: &ArrayRef,
261        child_idx: usize,
262    ) -> VortexResult<Option<ArrayRef>> {
263        crate::rules::RULES.evaluate(array, parent, child_idx)
264    }
265}
266
267pub(crate) fn number_type_from_dtype(dtype: &DType) -> NumberType {
268    number_type_from_ptype(dtype.as_ptype())
269}
270
271pub(crate) fn number_type_from_ptype(ptype: PType) -> NumberType {
272    match ptype {
273        PType::F16 => NumberType::F16,
274        PType::F32 => NumberType::F32,
275        PType::F64 => NumberType::F64,
276        PType::I16 => NumberType::I16,
277        PType::I32 => NumberType::I32,
278        PType::I64 => NumberType::I64,
279        PType::U16 => NumberType::U16,
280        PType::U32 => NumberType::U32,
281        PType::U64 => NumberType::U64,
282        _ => unreachable!("PType not supported by Pco: {:?}", ptype),
283    }
284}
285
286fn collect_valid(
287    parray: ArrayView<'_, Primitive>,
288    ctx: &mut ExecutionCtx,
289) -> VortexResult<PrimitiveArray> {
290    let mask = parray
291        .array()
292        .validity()?
293        .execute_mask(parray.array().len(), ctx)?;
294    let result = parray
295        .array()
296        .filter(mask)?
297        .execute::<PrimitiveArray>(ctx)?;
298    Ok(result)
299}
300
301pub(crate) fn vortex_err_from_pco(err: PcoError) -> VortexError {
302    use pco::errors::ErrorKind::*;
303    match err.kind {
304        Io(io_kind) => VortexError::from(std::io::Error::new(io_kind, err.message)),
305        InvalidArgument => vortex_err!(InvalidArgument: "{}", err.message),
306        other => vortex_err!("Pco {:?} error: {}", other, err.message),
307    }
308}
309
310#[derive(Clone, Debug)]
311/// Pco array encoding marker.
312pub struct Pco;
313
314impl Pco {
315    pub(crate) fn try_new(
316        dtype: DType,
317        data: PcoData,
318        validity: Validity,
319    ) -> VortexResult<PcoArray> {
320        let len = data.len();
321        data.validate(&dtype, len, &validity)?;
322        let slots = PcoSlots {
323            validity: validity_to_child(&validity, data.unsliced_n_rows()),
324        }
325        .into_slots();
326        Ok(unsafe {
327            Array::from_parts_unchecked(ArrayParts::new(Pco, dtype, len, data).with_slots(slots))
328        })
329    }
330
331    /// Compress a primitive array using pcodec.
332    pub fn from_primitive(
333        parray: ArrayView<'_, Primitive>,
334        level: usize,
335        values_per_page: usize,
336        ctx: &mut ExecutionCtx,
337    ) -> VortexResult<PcoArray> {
338        let dtype = parray.dtype().clone();
339        let validity = parray.validity()?;
340        let data = PcoData::from_primitive(parray, level, values_per_page, ctx)?;
341        Self::try_new(dtype, data, validity)
342    }
343}
344
345#[array_slots(Pco)]
346pub struct PcoSlots {
347    /// The validity bitmap indicating which elements are non-null.
348    #[slot(0)]
349    pub validity: Option<ArrayRef>,
350}
351
352/// Additional typed accessors for Pco arrays.
353pub trait PcoArrayExt: PcoArraySlotsExt {
354    /// Reconstruct the unsliced [`Validity`] from the validity slot.
355    fn unsliced_validity(&self) -> Validity {
356        child_to_validity(
357            self.as_ref().slots()[PcoSlots::VALIDITY].as_ref(),
358            self.as_ref().dtype().nullability(),
359        )
360    }
361}
362impl<T: TypedArrayRef<Pco>> PcoArrayExt for T {}
363
364#[derive(Clone, Debug)]
365/// Encoding-specific data for a [`PcoArray`].
366pub struct PcoData {
367    pub(crate) chunk_metas: Vec<ByteBuffer>,
368    pub(crate) pages: Vec<ByteBuffer>,
369    pub(crate) metadata: PcoMetadata,
370    ptype: PType,
371    unsliced_n_rows: usize,
372    slice_start: usize,
373    slice_stop: usize,
374}
375
376impl Display for PcoData {
377    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
378        write!(
379            f,
380            "ptype: {}, nrows: {}, slice: {}..{}",
381            self.ptype, self.unsliced_n_rows, self.slice_start, self.slice_stop
382        )
383    }
384}
385
386impl PcoData {
387    /// Validate dtype, validity, slice, and Pco component invariants.
388    pub fn validate(&self, dtype: &DType, len: usize, validity: &Validity) -> VortexResult<()> {
389        let _ = number_type_from_ptype(self.ptype);
390        vortex_ensure!(
391            dtype.as_ptype() == self.ptype,
392            "expected ptype {}, got {}",
393            self.ptype,
394            dtype.as_ptype()
395        );
396        vortex_ensure!(
397            dtype.nullability() == validity.nullability(),
398            "expected nullability {}, got {}",
399            validity.nullability(),
400            dtype.nullability()
401        );
402        vortex_ensure!(
403            self.slice_start <= self.slice_stop && self.slice_stop <= self.unsliced_n_rows,
404            "invalid slice range {}..{} for {} rows",
405            self.slice_start,
406            self.slice_stop,
407            self.unsliced_n_rows
408        );
409        vortex_ensure!(
410            self.slice_stop - self.slice_start == len,
411            "expected len {len}, got {}",
412            self.slice_stop - self.slice_start
413        );
414        if let Some(validity_len) = validity.maybe_len() {
415            vortex_ensure!(
416                validity_len == self.unsliced_n_rows,
417                "expected validity len {}, got {}",
418                self.unsliced_n_rows,
419                validity_len
420            );
421        }
422        vortex_ensure!(
423            self.chunk_metas.len() == self.metadata.chunks.len(),
424            "expected {} chunk metas, got {}",
425            self.metadata.chunks.len(),
426            self.chunk_metas.len()
427        );
428        vortex_ensure!(
429            self.pages.len()
430                == self
431                    .metadata
432                    .chunks
433                    .iter()
434                    .map(|chunk| chunk.pages.len())
435                    .sum::<usize>(),
436            "page count does not match metadata"
437        );
438        Ok(())
439    }
440
441    /// Construct unsliced Pco data from chunk metadata, pages, and serialized metadata.
442    pub fn new(
443        chunk_metas: Vec<ByteBuffer>,
444        pages: Vec<ByteBuffer>,
445        ptype: PType,
446        metadata: PcoMetadata,
447        len: usize,
448    ) -> Self {
449        Self {
450            chunk_metas,
451            pages,
452            metadata,
453            ptype,
454            unsliced_n_rows: len,
455            slice_start: 0,
456            slice_stop: len,
457        }
458    }
459
460    /// Compress a primitive array into Pco data.
461    pub fn from_primitive(
462        parray: ArrayView<'_, Primitive>,
463        level: usize,
464        values_per_page: usize,
465        ctx: &mut ExecutionCtx,
466    ) -> VortexResult<Self> {
467        Self::from_primitive_with_values_per_chunk(
468            parray,
469            level,
470            VALUES_PER_CHUNK,
471            values_per_page,
472            ctx,
473        )
474    }
475
476    pub(crate) fn from_primitive_with_values_per_chunk(
477        parray: ArrayView<'_, Primitive>,
478        level: usize,
479        values_per_chunk: usize,
480        values_per_page: usize,
481        ctx: &mut ExecutionCtx,
482    ) -> VortexResult<Self> {
483        let number_type = number_type_from_dtype(parray.dtype());
484        let values_per_page = if values_per_page == 0 {
485            values_per_chunk
486        } else {
487            values_per_page
488        };
489
490        // perhaps one day we can make this more configurable
491        let chunk_config = ChunkConfig::default()
492            .with_compression_level(level)
493            .with_paging_spec(PagingSpec::EqualPagesUpTo(values_per_page));
494
495        let values = collect_valid(parray, ctx)?;
496        let n_values = values.len();
497
498        let fc = FileCompressor::default();
499        let mut header = vec![];
500        fc.write_header(&mut header).map_err(vortex_err_from_pco)?;
501
502        let mut chunk_meta_buffers = vec![]; // the Pco component
503        let mut chunk_infos = vec![]; // the Vortex metadata
504        let mut page_buffers = vec![];
505        for chunk_start in (0..n_values).step_by(values_per_chunk) {
506            let chunk_end = cmp::min(n_values, chunk_start + values_per_chunk);
507            let mut cc = match_number_enum!(
508                number_type,
509                NumberType<T> => {
510                    let values = values.to_buffer::<T>();
511                    let chunk = &values.as_slice()[chunk_start..chunk_end];
512                    fc
513                        .chunk_compressor(chunk, &chunk_config)
514                        .map_err(vortex_err_from_pco)?
515                }
516            );
517
518            let mut chunk_meta_buffer = ByteBufferMut::with_capacity(cc.meta_size_hint());
519            cc.write_meta(&mut chunk_meta_buffer)
520                .map_err(vortex_err_from_pco)?;
521            chunk_meta_buffers.push(chunk_meta_buffer.freeze());
522
523            let mut page_infos = vec![];
524            for (page_idx, page_n_values) in cc.n_per_page().into_iter().enumerate() {
525                let mut page = ByteBufferMut::with_capacity(cc.page_size_hint(page_idx));
526                cc.write_page(page_idx, &mut page)
527                    .map_err(vortex_err_from_pco)?;
528                page_buffers.push(page.freeze());
529                page_infos.push(PcoPageInfo {
530                    n_values: u32::try_from(page_n_values)?,
531                });
532            }
533            chunk_infos.push(PcoChunkInfo { pages: page_infos })
534        }
535
536        let metadata = PcoMetadata {
537            header,
538            chunks: chunk_infos,
539        };
540        Ok(PcoData::new(
541            chunk_meta_buffers,
542            page_buffers,
543            parray.dtype().as_ptype(),
544            metadata,
545            parray.len(),
546        ))
547    }
548
549    /// Downcast and compress an array into Pco data.
550    ///
551    /// # Errors
552    ///
553    /// Returns an error if the input is not a primitive array or compression fails.
554    pub fn from_array(
555        array: ArrayRef,
556        level: usize,
557        nums_per_page: usize,
558        ctx: &mut ExecutionCtx,
559    ) -> VortexResult<Self> {
560        let parray = array.try_downcast::<Primitive>().map_err(|a| {
561            vortex_err!(
562                "Pco can only encode primitive arrays, got {}",
563                a.encoding_id()
564            )
565        })?;
566        Self::from_primitive(parray.as_view(), level, nums_per_page, ctx)
567    }
568
569    /// Decompress this Pco data into a primitive array.
570    pub fn decompress(
571        &self,
572        unsliced_validity: &Validity,
573        ctx: &mut ExecutionCtx,
574    ) -> VortexResult<PrimitiveArray> {
575        // To start, we figure out which chunks and pages we need to decompress, and with
576        // what value offset into the first such page.
577        let number_type = number_type_from_ptype(self.ptype);
578        let values_byte_buffer = match_number_enum!(
579            number_type,
580            NumberType<T> => {
581              self.decompress_values_typed::<T>(unsliced_validity, ctx)?
582            }
583        );
584
585        Ok(PrimitiveArray::from_values_byte_buffer(
586            values_byte_buffer,
587            self.ptype,
588            unsliced_validity.slice(self.slice_start..self.slice_stop)?,
589            self.slice_stop - self.slice_start,
590            ctx,
591        ))
592    }
593
594    fn decompress_values_typed<T: Number>(
595        &self,
596        unsliced_validity: &Validity,
597        ctx: &mut ExecutionCtx,
598    ) -> VortexResult<ByteBuffer> {
599        // To start, we figure out what range of values we need to decompress.
600        let slice_value_indices = unsliced_validity
601            .execute_mask(self.unsliced_n_rows, ctx)?
602            .valid_counts_for_indices(&[self.slice_start, self.slice_stop]);
603        let slice_value_start = slice_value_indices[0];
604        let slice_value_stop = slice_value_indices[1];
605        let slice_n_values = slice_value_stop - slice_value_start;
606
607        // Then we decompress those pages into a buffer. Note that these values
608        // may exceed the bounds of the slice, so we need to slice later.
609        let (fd, _) =
610            FileDecompressor::new(self.metadata.header.as_slice()).map_err(vortex_err_from_pco)?;
611        let mut decompressed_values = BufferMut::<T>::with_capacity(slice_n_values);
612        let mut page_idx = 0;
613        let mut page_value_start = 0;
614        let mut n_skipped_values = 0;
615        for (chunk_info, chunk_meta) in self.metadata.chunks.iter().zip(&self.chunk_metas) {
616            // lazily initialize chunk decompressor
617            let mut chunk_decompressor: Option<ChunkDecompressor<T>> = None;
618            for page_info in &chunk_info.pages {
619                let page_n_values = page_info.n_values as usize;
620                let page_value_stop = page_value_start + page_n_values;
621
622                if page_value_start >= slice_value_stop {
623                    break;
624                }
625
626                if page_value_stop > slice_value_start {
627                    // we need this page
628                    let old_len = decompressed_values.len();
629                    let new_len = old_len + page_n_values;
630                    decompressed_values.reserve(page_n_values);
631                    unsafe {
632                        decompressed_values.set_len(new_len);
633                    }
634                    let page: &[u8] = self.pages[page_idx].as_ref();
635
636                    let mut cd = match chunk_decompressor.take() {
637                        Some(d) => d,
638                        None => {
639                            let (new_cd, _) = fd
640                                .chunk_decompressor(chunk_meta.as_ref())
641                                .map_err(vortex_err_from_pco)?;
642                            new_cd
643                        }
644                    };
645
646                    let mut pd = cd
647                        .page_decompressor(page, page_n_values)
648                        .map_err(vortex_err_from_pco)?;
649                    pd.read(&mut decompressed_values[old_len..new_len])
650                        .map_err(vortex_err_from_pco)?;
651
652                    chunk_decompressor = Some(cd);
653                } else {
654                    n_skipped_values += page_n_values;
655                }
656
657                page_value_start = page_value_stop;
658                page_idx += 1;
659            }
660        }
661
662        // Slice only the values requested.
663        let value_offset = slice_value_start - n_skipped_values;
664        Ok(decompressed_values
665            .freeze()
666            .slice(value_offset..value_offset + slice_n_values)
667            .into_byte_buffer())
668    }
669
670    pub(crate) fn _slice(&self, start: usize, stop: usize) -> Self {
671        PcoData {
672            slice_start: self.slice_start + start,
673            slice_stop: self.slice_start + stop,
674            ..self.clone()
675        }
676    }
677
678    /// Returns the number of elements in the array.
679    pub fn len(&self) -> usize {
680        self.slice_stop - self.slice_start
681    }
682
683    /// Returns `true` if the array contains no elements.
684    pub fn is_empty(&self) -> bool {
685        self.slice_stop == self.slice_start
686    }
687
688    pub(crate) fn slice_start(&self) -> usize {
689        self.slice_start
690    }
691
692    pub(crate) fn slice_stop(&self) -> usize {
693        self.slice_stop
694    }
695
696    pub(crate) fn unsliced_n_rows(&self) -> usize {
697        self.unsliced_n_rows
698    }
699}
700
701impl ValidityVTable<Pco> for Pco {
702    fn validity(array: ArrayView<'_, Pco>) -> VortexResult<Validity> {
703        array
704            .unsliced_validity()
705            .slice(array.slice_start()..array.slice_stop())
706    }
707}
708
709impl OperationsVTable<Pco> for Pco {
710    fn scalar_at(
711        array: ArrayView<'_, Pco>,
712        index: usize,
713        ctx: &mut ExecutionCtx,
714    ) -> VortexResult<Scalar> {
715        let unsliced_validity = array.unsliced_validity();
716        array
717            ._slice(index, index + 1)
718            .decompress(&unsliced_validity, ctx)?
719            .into_array()
720            .execute_scalar(0, ctx)
721    }
722}
723
724#[cfg(test)]
725mod tests {
726    use vortex_array::IntoArray;
727    use vortex_array::VortexSessionExecute;
728    use vortex_array::array_session;
729    use vortex_array::arrays::PrimitiveArray;
730    use vortex_array::assert_arrays_eq;
731    use vortex_array::validity::Validity;
732    use vortex_buffer::buffer;
733
734    use crate::Pco;
735
736    #[test]
737    fn test_slice_nullable() {
738        let mut ctx = array_session().create_execution_ctx();
739        // Create a nullable array with some nulls
740        let values = PrimitiveArray::new(
741            buffer![10u32, 20, 30, 40, 50, 60],
742            Validity::from_iter([false, true, true, true, true, false]),
743        );
744        let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx).unwrap();
745        assert_arrays_eq!(
746            pco,
747            PrimitiveArray::from_option_iter([
748                None,
749                Some(20u32),
750                Some(30),
751                Some(40),
752                Some(50),
753                None
754            ]),
755            &mut ctx
756        );
757
758        // Slice to get only the non-null values in the middle
759        let sliced = pco.slice(1..5).unwrap();
760        let expected =
761            PrimitiveArray::from_option_iter([Some(20u32), Some(30), Some(40), Some(50)])
762                .into_array();
763        assert_arrays_eq!(sliced, expected, &mut ctx);
764    }
765}