1use 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
63const VALUES_PER_CHUNK: usize = pco::DEFAULT_MAX_PAGE_N;
82
83pub 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 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)]
311pub 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 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 #[slot(0)]
349 pub validity: Option<ArrayRef>,
350}
351
352pub trait PcoArrayExt: PcoArraySlotsExt {
354 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)]
365pub 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 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 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 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 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![]; let mut chunk_infos = vec![]; 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 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 pub fn decompress(
571 &self,
572 unsliced_validity: &Validity,
573 ctx: &mut ExecutionCtx,
574 ) -> VortexResult<PrimitiveArray> {
575 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 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 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 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 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 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 pub fn len(&self) -> usize {
680 self.slice_stop - self.slice_start
681 }
682
683 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 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 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}