From 5e53f77282d7328974831900e657e0c6adcbce6a Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Wed, 21 Oct 2026 10:57:38 +0200 Subject: [PATCH 01/10] (feat) add sort keys to footer Sort keys indicate that a file is lexicographically sorted. The current implementation writes and reads sort keys present in a footer. One currently missing piece is that it doesn't assert that the rows are truly sorted, so a user could write an unsorted file and declare in the footer that the rows are sorted. To add this assertion, we could use an agg_fn that checks multiple columns at the same time and validates that the rows are sorted across all chunks. Signed-off-by: Joaquin Colacci --- docs/specs/file-format.md | 4 +- .../flatbuffers/vortex-file/footer.fbs | 22 +- .../src/flatbuffers/generated/footer.rs | 224 +++++++++++++++++- vortex-file/src/footer/file_layout.rs | 10 + vortex-file/src/footer/mod.rs | 23 +- vortex-file/src/footer/serializer.rs | 1 + vortex-file/src/sort_key/mod.rs | 137 +++++++++++ 7 files changed, 411 insertions(+), 10 deletions(-) create mode 100644 vortex-file/src/sort_key/mod.rs diff --git a/docs/specs/file-format.md b/docs/specs/file-format.md index 98c1b91c794..bb5dfc324ce 100644 --- a/docs/specs/file-format.md +++ b/docs/specs/file-format.md @@ -54,7 +54,7 @@ The postscript contains the locations of: 1. a `dtype` segment representing the top-level logical data type (i.e., schema) 2. a `layout` segment containing the root `Layout` 3. a `statistics` segment containing file-level per-field statistics (e.g., minima and maxima of each field/column, for whole-file pruning) -4. a `footer` segment containing a dictionary-encoded _segment map_, and other shared configuration such as compression and encryption schemes +4. a `footer` segment containing a dictionary-encoded _segment map_, compression and encryption settings, and an optional sort key 5. up to 16 user-defined `metadata` segments, each identified by a unique, non-empty UTF-8 key of at most 64 bytes The postscript carries a locator (offset, length, and alignment) for each metadata segment that is @@ -86,7 +86,7 @@ valid to store a `Float64` array, a `Boolean` array, or any other root data type The footer is a flat buffer serialized `Footer` object. This object contains all the information required to load the root `Layout` object into a usable `LayoutReader`. -For example, it contains the locations, compression schemes, encryption schemes, and required alignment of all segments in the file. +For example, it stores the location, compression scheme, encryption scheme, sort keys, and required alignment of each segment. :::{literalinclude} ../../vortex-file/flatbuffers/vortex-file/footer.fbs :start-after: [footer] diff --git a/vortex-file/flatbuffers/vortex-file/footer.fbs b/vortex-file/flatbuffers/vortex-file/footer.fbs index 636e8d4b6b7..8d7e1ec2671 100644 --- a/vortex-file/flatbuffers/vortex-file/footer.fbs +++ b/vortex-file/flatbuffers/vortex-file/footer.fbs @@ -21,7 +21,7 @@ include "vortex-layout/layout.fbs"; /// The segments pointed to by the postscript have inline compression and encryption specs to avoid /// the need to fetch encryption schemes up-front. /// -/// New fields must be appended to preserve FlatBuffers field-order compatibility with old readers. +/// Add new fields at the end so older readers can still read the footer. table Postscript { /// Segment containing the root `DType` flatbuffer. dtype: PostscriptSegment; @@ -113,8 +113,22 @@ table FileStatistics { partials: [ubyte]; } -/// The `Registry` object stores dictionary-encoded configuration for segments, -/// compression schemes, encryption schemes, etc. +/// A field path in the file's sort key. +table FieldPath { + /// Field names that make up the path. An empty path selects a non-struct root. + path: [string] (required); +} + +/// Sort key supplied by the caller. Fields are compared in order, ascending with nulls first. +/// Duplicate keys are allowed. +table SortKey { + fields: [FieldPath] (required); +} + +/// The `Footer` stores dictionary-encoded segment and encoding specifications, plus an optional +/// sort key for the file. +/// +/// Add new fields at the end so older readers can still read the footer. table Footer { // Dictionary-encoded array specs, up to u16::MAX. array_specs: [ArraySpec]; @@ -126,6 +140,8 @@ table Footer { compression_specs: [CompressionSpec]; // Dictionary-encoded encryption specs, up to u16::MAX. encryption_specs: [EncryptionSpec]; + /// Sort key for the rows in this file. If absent, the row order is unknown. + sort_key_spec: SortKey; } /// An `ArraySpec` describes the type of a particular array. diff --git a/vortex-file/src/flatbuffers/generated/footer.rs b/vortex-file/src/flatbuffers/generated/footer.rs index 9e0238c47e7..17276a73e04 100644 --- a/vortex-file/src/flatbuffers/generated/footer.rs +++ b/vortex-file/src/flatbuffers/generated/footer.rs @@ -348,7 +348,7 @@ pub enum PostscriptOffset {} /// The segments pointed to by the postscript have inline compression and encryption specs to avoid /// the need to fetch encryption schemes up-front. /// -/// New fields must be appended to preserve FlatBuffers field-order compatibility with old readers. +/// Add new fields at the end so older readers can still read the footer. pub struct Postscript<'a> { pub _tab: ::flatbuffers::Table<'a>, } @@ -1215,11 +1215,211 @@ impl ::core::fmt::Debug for FileStatistics<'_> { ds.finish() } } +pub enum FieldPathOffset {} +#[derive(Copy, Clone, PartialEq)] + +/// A field path in the file's sort key. +pub struct FieldPath<'a> { + pub _tab: ::flatbuffers::Table<'a>, +} + +impl<'a> ::flatbuffers::Follow<'a> for FieldPath<'a> { + type Inner = FieldPath<'a>; + #[inline] + unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { + Self { _tab: unsafe { ::flatbuffers::Table::new(buf, loc) } } + } +} + +impl<'a> FieldPath<'a> { + pub const VT_PATH: ::flatbuffers::VOffsetT = 4; + + #[inline] + pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { + FieldPath { _tab: table } + } + #[allow(unused_mut)] + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: ::flatbuffers::Allocator + 'bldr>( + _fbb: &'mut_bldr mut ::flatbuffers::FlatBufferBuilder<'bldr, A>, + args: &'args FieldPathArgs<'args> + ) -> ::flatbuffers::WIPOffset> { + let mut builder = FieldPathBuilder::new(_fbb); + if let Some(x) = args.path { builder.add_path(x); } + builder.finish() + } + + + /// Field names that make up the path. An empty path selects a non-struct root. + #[inline] + pub fn path(&self) -> ::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset<&'a str>> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset<&'a str>>>>(FieldPath::VT_PATH, None).unwrap()} + } +} + +impl ::flatbuffers::Verifiable for FieldPath<'_> { + #[inline] + fn run_verifier( + v: &mut ::flatbuffers::Verifier, pos: usize + ) -> Result<(), ::flatbuffers::InvalidFlatbuffer> { + v.visit_table(pos)? + .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset<&'_ str>>>>("path", Self::VT_PATH, true)? + .finish(); + Ok(()) + } +} +pub struct FieldPathArgs<'a> { + pub path: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset<&'a str>>>>, +} +impl<'a> Default for FieldPathArgs<'a> { + #[inline] + fn default() -> Self { + FieldPathArgs { + path: None, // required field + } + } +} + +pub struct FieldPathBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { + fbb_: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>, + start_: ::flatbuffers::WIPOffset<::flatbuffers::TableUnfinishedWIPOffset>, +} +impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> FieldPathBuilder<'a, 'b, A> { + #[inline] + pub fn add_path(&mut self, path: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset<&'b str>>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(FieldPath::VT_PATH, path); + } + #[inline] + pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> FieldPathBuilder<'a, 'b, A> { + let start = _fbb.start_table(); + FieldPathBuilder { + fbb_: _fbb, + start_: start, + } + } + #[inline] + pub fn finish(self) -> ::flatbuffers::WIPOffset> { + let o = self.fbb_.end_table(self.start_); + self.fbb_.required(o, FieldPath::VT_PATH,"path"); + ::flatbuffers::WIPOffset::new(o.value()) + } +} + +impl ::core::fmt::Debug for FieldPath<'_> { + fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { + let mut ds = f.debug_struct("FieldPath"); + ds.field("path", &self.path()); + ds.finish() + } +} +pub enum SortKeyOffset {} +#[derive(Copy, Clone, PartialEq)] + +/// Sort key supplied by the caller. Fields are compared in order, ascending with nulls first. +/// Duplicate keys are allowed. +pub struct SortKey<'a> { + pub _tab: ::flatbuffers::Table<'a>, +} + +impl<'a> ::flatbuffers::Follow<'a> for SortKey<'a> { + type Inner = SortKey<'a>; + #[inline] + unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { + Self { _tab: unsafe { ::flatbuffers::Table::new(buf, loc) } } + } +} + +impl<'a> SortKey<'a> { + pub const VT_FIELDS: ::flatbuffers::VOffsetT = 4; + + #[inline] + pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { + SortKey { _tab: table } + } + #[allow(unused_mut)] + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: ::flatbuffers::Allocator + 'bldr>( + _fbb: &'mut_bldr mut ::flatbuffers::FlatBufferBuilder<'bldr, A>, + args: &'args SortKeyArgs<'args> + ) -> ::flatbuffers::WIPOffset> { + let mut builder = SortKeyBuilder::new(_fbb); + if let Some(x) = args.fields { builder.add_fields(x); } + builder.finish() + } + + + #[inline] + pub fn fields(&self) -> ::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(SortKey::VT_FIELDS, None).unwrap()} + } +} + +impl ::flatbuffers::Verifiable for SortKey<'_> { + #[inline] + fn run_verifier( + v: &mut ::flatbuffers::Verifier, pos: usize + ) -> Result<(), ::flatbuffers::InvalidFlatbuffer> { + v.visit_table(pos)? + .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("fields", Self::VT_FIELDS, true)? + .finish(); + Ok(()) + } +} +pub struct SortKeyArgs<'a> { + pub fields: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, +} +impl<'a> Default for SortKeyArgs<'a> { + #[inline] + fn default() -> Self { + SortKeyArgs { + fields: None, // required field + } + } +} + +pub struct SortKeyBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { + fbb_: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>, + start_: ::flatbuffers::WIPOffset<::flatbuffers::TableUnfinishedWIPOffset>, +} +impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> SortKeyBuilder<'a, 'b, A> { + #[inline] + pub fn add_fields(&mut self, fields: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset>>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(SortKey::VT_FIELDS, fields); + } + #[inline] + pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> SortKeyBuilder<'a, 'b, A> { + let start = _fbb.start_table(); + SortKeyBuilder { + fbb_: _fbb, + start_: start, + } + } + #[inline] + pub fn finish(self) -> ::flatbuffers::WIPOffset> { + let o = self.fbb_.end_table(self.start_); + self.fbb_.required(o, SortKey::VT_FIELDS,"fields"); + ::flatbuffers::WIPOffset::new(o.value()) + } +} + +impl ::core::fmt::Debug for SortKey<'_> { + fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { + let mut ds = f.debug_struct("SortKey"); + ds.field("fields", &self.fields()); + ds.finish() + } +} pub enum FooterOffset {} #[derive(Copy, Clone, PartialEq)] -/// The `Registry` object stores dictionary-encoded configuration for segments, -/// compression schemes, encryption schemes, etc. +/// The `Footer` stores dictionary-encoded segment and encoding specifications, plus an optional +/// sort key for the file. +/// +/// Add new fields at the end so older readers can still read the footer. pub struct Footer<'a> { pub _tab: ::flatbuffers::Table<'a>, } @@ -1238,6 +1438,7 @@ impl<'a> Footer<'a> { pub const VT_SEGMENT_SPECS: ::flatbuffers::VOffsetT = 8; pub const VT_COMPRESSION_SPECS: ::flatbuffers::VOffsetT = 10; pub const VT_ENCRYPTION_SPECS: ::flatbuffers::VOffsetT = 12; + pub const VT_SORT_KEY_SPEC: ::flatbuffers::VOffsetT = 14; #[inline] pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { @@ -1249,6 +1450,7 @@ impl<'a> Footer<'a> { args: &'args FooterArgs<'args> ) -> ::flatbuffers::WIPOffset> { let mut builder = FooterBuilder::new(_fbb); + if let Some(x) = args.sort_key_spec { builder.add_sort_key_spec(x); } if let Some(x) = args.encryption_specs { builder.add_encryption_specs(x); } if let Some(x) = args.compression_specs { builder.add_compression_specs(x); } if let Some(x) = args.segment_specs { builder.add_segment_specs(x); } @@ -1293,6 +1495,14 @@ impl<'a> Footer<'a> { // which contains a valid value in this slot unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(Footer::VT_ENCRYPTION_SPECS, None)} } + /// Sort key for the rows in this file. If absent, the row order is unknown. + #[inline] + pub fn sort_key_spec(&self) -> Option> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset>(Footer::VT_SORT_KEY_SPEC, None)} + } } impl ::flatbuffers::Verifiable for Footer<'_> { @@ -1306,6 +1516,7 @@ impl ::flatbuffers::Verifiable for Footer<'_> { .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, SegmentSpec>>>("segment_specs", Self::VT_SEGMENT_SPECS, false)? .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("compression_specs", Self::VT_COMPRESSION_SPECS, false)? .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("encryption_specs", Self::VT_ENCRYPTION_SPECS, false)? + .visit_field::<::flatbuffers::ForwardsUOffset>("sort_key_spec", Self::VT_SORT_KEY_SPEC, false)? .finish(); Ok(()) } @@ -1316,6 +1527,7 @@ pub struct FooterArgs<'a> { pub segment_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, SegmentSpec>>>, pub compression_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, pub encryption_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, + pub sort_key_spec: Option<::flatbuffers::WIPOffset>>, } impl<'a> Default for FooterArgs<'a> { #[inline] @@ -1326,6 +1538,7 @@ impl<'a> Default for FooterArgs<'a> { segment_specs: None, compression_specs: None, encryption_specs: None, + sort_key_spec: None, } } } @@ -1356,6 +1569,10 @@ impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> FooterBuilder<'a, 'b, A> { self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(Footer::VT_ENCRYPTION_SPECS, encryption_specs); } #[inline] + pub fn add_sort_key_spec(&mut self, sort_key_spec: ::flatbuffers::WIPOffset>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset>(Footer::VT_SORT_KEY_SPEC, sort_key_spec); + } + #[inline] pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> FooterBuilder<'a, 'b, A> { let start = _fbb.start_table(); FooterBuilder { @@ -1378,6 +1595,7 @@ impl ::core::fmt::Debug for Footer<'_> { ds.field("segment_specs", &self.segment_specs()); ds.field("compression_specs", &self.compression_specs()); ds.field("encryption_specs", &self.encryption_specs()); + ds.field("sort_key_spec", &self.sort_key_spec()); ds.finish() } } diff --git a/vortex-file/src/footer/file_layout.rs b/vortex-file/src/footer/file_layout.rs index c848a3f207f..a8cad66710a 100644 --- a/vortex-file/src/footer/file_layout.rs +++ b/vortex-file/src/footer/file_layout.rs @@ -11,6 +11,7 @@ use std::sync::Arc; use flatbuffers::FlatBufferBuilder; use flatbuffers::WIPOffset; +use vortex_array::dtype::FieldPath; use vortex_array::flatbuffers::FlatBufferRoot; use vortex_array::flatbuffers::WriteFlatBuffer; use vortex_error::VortexResult; @@ -18,6 +19,7 @@ use vortex_session::registry::ReadContext; use crate::flatbuffers::footer as fb; use crate::footer::segment::SegmentSpec; +use crate::sort_key::write_sort_key; /// A writer for serializing a file layout to a FlatBuffer. /// @@ -30,6 +32,8 @@ pub(crate) struct FooterFlatBufferWriter { pub(crate) layout_ctx: ReadContext, /// Specifications for all segments in the file. pub(crate) segment_specs: Arc<[SegmentSpec]>, + /// Optional sort key, with fields listed in comparison order. + pub(crate) sort_key: Option>, } impl FlatBufferRoot for FooterFlatBufferWriter {} @@ -65,6 +69,11 @@ impl WriteFlatBuffer for FooterFlatBufferWriter { }) .collect::>(); let layout_specs = fbb.create_vector(layout_specs.as_slice()); + let sort_key_spec = self + .sort_key + .as_deref() + .map(|key| write_sort_key(fbb, key)) + .transpose()?; Ok(fb::Footer::create( fbb, @@ -74,6 +83,7 @@ impl WriteFlatBuffer for FooterFlatBufferWriter { layout_specs: Some(layout_specs), compression_specs: None, encryption_specs: None, + sort_key_spec, }, )) } diff --git a/vortex-file/src/footer/mod.rs b/vortex-file/src/footer/mod.rs index 57be23ec74e..1d793cb0f07 100644 --- a/vortex-file/src/footer/mod.rs +++ b/vortex-file/src/footer/mod.rs @@ -3,8 +3,8 @@ //! Vortex file footer metadata. //! -//! A footer contains the root layout, file-level statistics, the segment map, and the read contexts -//! needed to resolve array/layout encoding ids during deserialization. +//! A footer contains the root layout, file-level statistics and ordering, the segment map, and the +//! read contexts needed to resolve array/layout encoding ids during deserialization. //! //! The byte-level footer and postscript layout is part of the file-format spec; this module exposes //! the structured Rust representation and serializer/deserializer state machine. @@ -27,6 +27,7 @@ use itertools::Itertools; pub use segment::*; use vortex_array::ArrayId; use vortex_array::dtype::DType; +use vortex_array::dtype::FieldPath; use vortex_array::flatbuffers::FlatBuffer; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; @@ -39,6 +40,7 @@ use vortex_session::VortexSession; use vortex_session::registry::ReadContext; use crate::flatbuffers::footer as fb; +use crate::sort_key::read_sort_key; /// Maximum number of user-defined metadata segments. Keeps postscript bookkeeping small so the /// footer and required segments still fit the initial tail read. @@ -60,6 +62,7 @@ pub struct Footer { root_layout: LayoutRef, segments: Arc<[SegmentSpec]>, statistics: Option, + sort_key: Option>, metadata: Arc<[(String, SegmentSpec)]>, // `segments` followed by the `metadata` segments, which is the segment table a file is opened // with. Precomputed so opening a file with a cached footer does not copy every segment spec. @@ -89,6 +92,7 @@ impl Footer { segments_end: segments_end(&segments), segments, statistics, + sort_key: None, metadata: Arc::from([]), array_read_ctx, approx_byte_size: None, @@ -100,6 +104,11 @@ impl Footer { self } + pub(crate) fn with_sort_key(mut self, sort_key: Option>) -> Self { + self.sort_key = sort_key; + self + } + /// Read the [`Footer`] from a flatbuffer. pub(crate) fn from_flatbuffer( footer_bytes: &[u8], @@ -115,6 +124,10 @@ impl Footer { .sum(); let approx_byte_size = footer_bytes.len() + layout_bytes.len() + metadata_bytes; let fb_footer = root::(footer_bytes)?; + let sort_key = fb_footer + .sort_key_spec() + .map(|key| read_sort_key(key, &dtype)) + .transpose()?; // Create a LayoutContext from the registry. let layout_specs = fb_footer.layout_specs(); @@ -165,6 +178,7 @@ impl Footer { segments_with_metadata, segments, statistics, + sort_key, metadata, array_read_ctx, approx_byte_size: Some(approx_byte_size), @@ -186,6 +200,11 @@ impl Footer { self.statistics.as_ref() } + /// Returns the sort key supplied when the file was written. + pub fn sort_key(&self) -> Option<&[FieldPath]> { + self.sort_key.as_deref() + } + /// Returns the user-defined metadata segment locators stored in the postscript. pub fn metadata_segments(&self) -> impl Iterator { self.metadata diff --git a/vortex-file/src/footer/serializer.rs b/vortex-file/src/footer/serializer.rs index e72e58c0389..d9610cbc6c1 100644 --- a/vortex-file/src/footer/serializer.rs +++ b/vortex-file/src/footer/serializer.rs @@ -161,6 +161,7 @@ impl FooterSerializer { ctx: self.footer.array_read_ctx.clone(), layout_ctx: ReadContext::new(layout_ctx.to_ids()), segment_specs: Arc::clone(&self.footer.segments), + sort_key: self.footer.sort_key.clone(), }, )?; buffers.push(buffer); diff --git a/vortex-file/src/sort_key/mod.rs b/vortex-file/src/sort_key/mod.rs new file mode 100644 index 00000000000..91fb89e73f8 --- /dev/null +++ b/vortex-file/src/sort_key/mod.rs @@ -0,0 +1,137 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Sort keys for Vortex files. +//! +//! The current implementations in this file are for validating sort keys, +//! and for writing and reading sort keys to/from the footer. +//! +//! # About sort keys +//! +//! Sort keys are a compound key that describes how the rows in a file are sorted. +//! For example, these rows are lexicographically sorted by (user_id, city): +//! +//! ```text +//! user_id | city +//! --------+----------- +//! 1 | Bangkok +//! 1 | Mumbai +//! 2 | Tokyo +//! 3 | Copenhagen +//! ``` +//! +//! # Difference from sortedness stats +//! +//! The difference between `IsSorted` / `IsStrictSorted` and sort keys +//! is that the stats describe the sortedness of one column, while a compound key +//! requires sortedness across the file only for its first column. The remaining +//! columns only need to be sorted where all preceding columns are equal. +//! +//! The ordering must also hold across chunks. For example, each chunk below +//! is sorted by (user_id, city): +//! +//! ```text +//! user_id | city +//! --------+----------- +//! Chunk 1 1 | Bangkok +//! 2 | Tokyo +//! --------+----------- +//! Chunk 2 2 | Mumbai <- breaks the ordering +//! 3 | Copenhagen +//! ``` +//! +//! But together, they break the sort key: (2, Tokyo) comes before (2, Mumbai), +//! even though Mumbai should come first. + +use std::sync::Arc; + +use flatbuffers::FlatBufferBuilder; +use flatbuffers::WIPOffset; +use vortex_array::dtype::DType; +use vortex_array::dtype::Field; +use vortex_array::dtype::FieldPath; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_utils::aliases::hash_set::HashSet; + +use crate::flatbuffers::footer as fb; + +/// Check that the key is non-empty and its fields exist and appear only once. +pub(crate) fn validate_sort_key(dtype: &DType, fields: &[FieldPath]) -> VortexResult<()> { + vortex_ensure!( + !fields.is_empty(), + "A file sort key must contain at least one field" + ); + + let mut seen = HashSet::<&FieldPath>::default(); + for field in fields { + vortex_ensure!(seen.insert(field), "Duplicate sort field: {}", field); + + match (dtype, field.parts()) { + (DType::Struct(struct_fields, _), [Field::Name(name)]) => { + vortex_ensure!( + struct_fields.find(name).is_some(), + "Sort field {} does not exist in the file schema", + field + ); + } + // For a non-struct file, an empty path selects the entire root array as the sort field. + (dtype, []) if !dtype.is_struct() => {} + _ => vortex_bail!( + "Sort field {} must select a top-level field of a struct, or a non-struct root", + field + ), + } + } + + Ok(()) +} + +/// Reads and validates the footer's sort key. +pub(crate) fn read_sort_key(key: fb::SortKey<'_>, dtype: &DType) -> VortexResult> { + let fields: Arc<[_]> = key + .fields() + .iter() + .map(|field| field.path().iter().map(Field::from).collect()) + .collect(); + validate_sort_key(dtype, &fields)?; + + Ok(fields) +} + +/// Writes the sort key into the footer. +pub(crate) fn write_sort_key<'fb>( + fbb: &mut FlatBufferBuilder<'fb>, + fields: &[FieldPath], +) -> VortexResult>> { + let mut offsets = Vec::with_capacity(fields.len()); + + for field in fields { + let names = field + .parts() + .iter() + .map(|part| { + let name = part + .as_name() + .ok_or_else(|| vortex_err!("Invalid sort field path"))?; + Ok(fbb.create_string(name)) + }) + .collect::>>()?; + + let path = fbb.create_vector(&names); + offsets.push(fb::FieldPath::create( + fbb, + &fb::FieldPathArgs { path: Some(path) }, + )); + } + + let fields = fbb.create_vector(&offsets); + Ok(fb::SortKey::create( + fbb, + &fb::SortKeyArgs { + fields: Some(fields), + }, + )) +} From 15d6780587b1741da796c02571856b2c45b0b1da Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Wed, 21 Oct 2026 11:57:39 +0200 Subject: [PATCH 02/10] (fix) docs changes Unintended changes. Signed-off-by: Joaquin Colacci --- docs/specs/file-format.md | 4 ++-- vortex-file/flatbuffers/vortex-file/footer.fbs | 8 ++------ vortex-file/src/flatbuffers/generated/footer.rs | 8 ++------ vortex-file/src/footer/mod.rs | 4 ++-- 4 files changed, 8 insertions(+), 16 deletions(-) diff --git a/docs/specs/file-format.md b/docs/specs/file-format.md index bb5dfc324ce..4d8cf8c0229 100644 --- a/docs/specs/file-format.md +++ b/docs/specs/file-format.md @@ -54,7 +54,7 @@ The postscript contains the locations of: 1. a `dtype` segment representing the top-level logical data type (i.e., schema) 2. a `layout` segment containing the root `Layout` 3. a `statistics` segment containing file-level per-field statistics (e.g., minima and maxima of each field/column, for whole-file pruning) -4. a `footer` segment containing a dictionary-encoded _segment map_, compression and encryption settings, and an optional sort key +4. a `footer` segment containing a dictionary-encoded _segment map_, and other shared configuration such as compression, encryption schemes or sort keys 5. up to 16 user-defined `metadata` segments, each identified by a unique, non-empty UTF-8 key of at most 64 bytes The postscript carries a locator (offset, length, and alignment) for each metadata segment that is @@ -86,7 +86,7 @@ valid to store a `Float64` array, a `Boolean` array, or any other root data type The footer is a flat buffer serialized `Footer` object. This object contains all the information required to load the root `Layout` object into a usable `LayoutReader`. -For example, it stores the location, compression scheme, encryption scheme, sort keys, and required alignment of each segment. +For example, it contains the locations, compression schemes, encryption schemes, sort keys, and required alignment of all segments in the file. :::{literalinclude} ../../vortex-file/flatbuffers/vortex-file/footer.fbs :start-after: [footer] diff --git a/vortex-file/flatbuffers/vortex-file/footer.fbs b/vortex-file/flatbuffers/vortex-file/footer.fbs index 8d7e1ec2671..d6f3de99b43 100644 --- a/vortex-file/flatbuffers/vortex-file/footer.fbs +++ b/vortex-file/flatbuffers/vortex-file/footer.fbs @@ -20,8 +20,6 @@ include "vortex-layout/layout.fbs"; /// /// The segments pointed to by the postscript have inline compression and encryption specs to avoid /// the need to fetch encryption schemes up-front. -/// -/// Add new fields at the end so older readers can still read the footer. table Postscript { /// Segment containing the root `DType` flatbuffer. dtype: PostscriptSegment; @@ -125,10 +123,8 @@ table SortKey { fields: [FieldPath] (required); } -/// The `Footer` stores dictionary-encoded segment and encoding specifications, plus an optional -/// sort key for the file. -/// -/// Add new fields at the end so older readers can still read the footer. +/// The `Registry` object stores dictionary-encoded configuration for segments, +/// compression schemes, encryption schemes, etc. table Footer { // Dictionary-encoded array specs, up to u16::MAX. array_specs: [ArraySpec]; diff --git a/vortex-file/src/flatbuffers/generated/footer.rs b/vortex-file/src/flatbuffers/generated/footer.rs index 17276a73e04..68feb35c6e8 100644 --- a/vortex-file/src/flatbuffers/generated/footer.rs +++ b/vortex-file/src/flatbuffers/generated/footer.rs @@ -347,8 +347,6 @@ pub enum PostscriptOffset {} /// /// The segments pointed to by the postscript have inline compression and encryption specs to avoid /// the need to fetch encryption schemes up-front. -/// -/// Add new fields at the end so older readers can still read the footer. pub struct Postscript<'a> { pub _tab: ::flatbuffers::Table<'a>, } @@ -1416,10 +1414,8 @@ impl ::core::fmt::Debug for SortKey<'_> { pub enum FooterOffset {} #[derive(Copy, Clone, PartialEq)] -/// The `Footer` stores dictionary-encoded segment and encoding specifications, plus an optional -/// sort key for the file. -/// -/// Add new fields at the end so older readers can still read the footer. +/// The `Registry` object stores dictionary-encoded configuration for segments, +/// compression schemes, encryption schemes, etc. pub struct Footer<'a> { pub _tab: ::flatbuffers::Table<'a>, } diff --git a/vortex-file/src/footer/mod.rs b/vortex-file/src/footer/mod.rs index 1d793cb0f07..d509e5fd7ad 100644 --- a/vortex-file/src/footer/mod.rs +++ b/vortex-file/src/footer/mod.rs @@ -3,8 +3,8 @@ //! Vortex file footer metadata. //! -//! A footer contains the root layout, file-level statistics and ordering, the segment map, and the -//! read contexts needed to resolve array/layout encoding ids during deserialization. +//! A footer contains the root layout, file-level statistics, the segment map, and the read contexts +//! needed to resolve array/layout encoding ids during deserialization. //! //! The byte-level footer and postscript layout is part of the file-format spec; this module exposes //! the structured Rust representation and serializer/deserializer state machine. From 8db85a033efd425931cb30dbe9d54a2c00f19398 Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Wed, 21 Oct 2026 12:00:41 +0200 Subject: [PATCH 03/10] (fix) docs removed line Signed-off-by: Joaquin Colacci --- vortex-file/flatbuffers/vortex-file/footer.fbs | 2 ++ vortex-file/src/flatbuffers/generated/footer.rs | 2 ++ 2 files changed, 4 insertions(+) diff --git a/vortex-file/flatbuffers/vortex-file/footer.fbs b/vortex-file/flatbuffers/vortex-file/footer.fbs index d6f3de99b43..7b1b7b85707 100644 --- a/vortex-file/flatbuffers/vortex-file/footer.fbs +++ b/vortex-file/flatbuffers/vortex-file/footer.fbs @@ -20,6 +20,8 @@ include "vortex-layout/layout.fbs"; /// /// The segments pointed to by the postscript have inline compression and encryption specs to avoid /// the need to fetch encryption schemes up-front. +/// +/// New fields must be appended to preserve FlatBuffers field-order compatibility with old readers. table Postscript { /// Segment containing the root `DType` flatbuffer. dtype: PostscriptSegment; diff --git a/vortex-file/src/flatbuffers/generated/footer.rs b/vortex-file/src/flatbuffers/generated/footer.rs index 68feb35c6e8..94ee98d8cda 100644 --- a/vortex-file/src/flatbuffers/generated/footer.rs +++ b/vortex-file/src/flatbuffers/generated/footer.rs @@ -347,6 +347,8 @@ pub enum PostscriptOffset {} /// /// The segments pointed to by the postscript have inline compression and encryption specs to avoid /// the need to fetch encryption schemes up-front. +/// +/// New fields must be appended to preserve FlatBuffers field-order compatibility with old readers. pub struct Postscript<'a> { pub _tab: ::flatbuffers::Table<'a>, } From 7520833eb30dd66148d4f9eae162ba04545eda91 Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Wed, 21 Oct 2026 12:06:29 +0200 Subject: [PATCH 04/10] (fix) rename footer `SortKey` type to `SortKeySpec` for consistency Signed-off-by: Joaquin Colacci --- .../flatbuffers/vortex-file/footer.fbs | 4 +- .../src/flatbuffers/generated/footer.rs | 58 +++++++++---------- vortex-file/src/sort_key/mod.rs | 11 ++-- 3 files changed, 38 insertions(+), 35 deletions(-) diff --git a/vortex-file/flatbuffers/vortex-file/footer.fbs b/vortex-file/flatbuffers/vortex-file/footer.fbs index 7b1b7b85707..85a91bfa96b 100644 --- a/vortex-file/flatbuffers/vortex-file/footer.fbs +++ b/vortex-file/flatbuffers/vortex-file/footer.fbs @@ -121,7 +121,7 @@ table FieldPath { /// Sort key supplied by the caller. Fields are compared in order, ascending with nulls first. /// Duplicate keys are allowed. -table SortKey { +table SortKeySpec { fields: [FieldPath] (required); } @@ -139,7 +139,7 @@ table Footer { // Dictionary-encoded encryption specs, up to u16::MAX. encryption_specs: [EncryptionSpec]; /// Sort key for the rows in this file. If absent, the row order is unknown. - sort_key_spec: SortKey; + sort_key_spec: SortKeySpec; } /// An `ArraySpec` describes the type of a particular array. diff --git a/vortex-file/src/flatbuffers/generated/footer.rs b/vortex-file/src/flatbuffers/generated/footer.rs index 94ee98d8cda..d6230686be9 100644 --- a/vortex-file/src/flatbuffers/generated/footer.rs +++ b/vortex-file/src/flatbuffers/generated/footer.rs @@ -1314,36 +1314,36 @@ impl ::core::fmt::Debug for FieldPath<'_> { ds.finish() } } -pub enum SortKeyOffset {} +pub enum SortKeySpecOffset {} #[derive(Copy, Clone, PartialEq)] /// Sort key supplied by the caller. Fields are compared in order, ascending with nulls first. /// Duplicate keys are allowed. -pub struct SortKey<'a> { +pub struct SortKeySpec<'a> { pub _tab: ::flatbuffers::Table<'a>, } -impl<'a> ::flatbuffers::Follow<'a> for SortKey<'a> { - type Inner = SortKey<'a>; +impl<'a> ::flatbuffers::Follow<'a> for SortKeySpec<'a> { + type Inner = SortKeySpec<'a>; #[inline] unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { Self { _tab: unsafe { ::flatbuffers::Table::new(buf, loc) } } } } -impl<'a> SortKey<'a> { +impl<'a> SortKeySpec<'a> { pub const VT_FIELDS: ::flatbuffers::VOffsetT = 4; #[inline] pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { - SortKey { _tab: table } + SortKeySpec { _tab: table } } #[allow(unused_mut)] pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: ::flatbuffers::Allocator + 'bldr>( _fbb: &'mut_bldr mut ::flatbuffers::FlatBufferBuilder<'bldr, A>, - args: &'args SortKeyArgs<'args> - ) -> ::flatbuffers::WIPOffset> { - let mut builder = SortKeyBuilder::new(_fbb); + args: &'args SortKeySpecArgs<'args> + ) -> ::flatbuffers::WIPOffset> { + let mut builder = SortKeySpecBuilder::new(_fbb); if let Some(x) = args.fields { builder.add_fields(x); } builder.finish() } @@ -1354,11 +1354,11 @@ impl<'a> SortKey<'a> { // Safety: // Created from valid Table for this object // which contains a valid value in this slot - unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(SortKey::VT_FIELDS, None).unwrap()} + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(SortKeySpec::VT_FIELDS, None).unwrap()} } } -impl ::flatbuffers::Verifiable for SortKey<'_> { +impl ::flatbuffers::Verifiable for SortKeySpec<'_> { #[inline] fn run_verifier( v: &mut ::flatbuffers::Verifier, pos: usize @@ -1369,46 +1369,46 @@ impl ::flatbuffers::Verifiable for SortKey<'_> { Ok(()) } } -pub struct SortKeyArgs<'a> { +pub struct SortKeySpecArgs<'a> { pub fields: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, } -impl<'a> Default for SortKeyArgs<'a> { +impl<'a> Default for SortKeySpecArgs<'a> { #[inline] fn default() -> Self { - SortKeyArgs { + SortKeySpecArgs { fields: None, // required field } } } -pub struct SortKeyBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { +pub struct SortKeySpecBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { fbb_: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>, start_: ::flatbuffers::WIPOffset<::flatbuffers::TableUnfinishedWIPOffset>, } -impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> SortKeyBuilder<'a, 'b, A> { +impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> SortKeySpecBuilder<'a, 'b, A> { #[inline] pub fn add_fields(&mut self, fields: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset>>>) { - self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(SortKey::VT_FIELDS, fields); + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(SortKeySpec::VT_FIELDS, fields); } #[inline] - pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> SortKeyBuilder<'a, 'b, A> { + pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> SortKeySpecBuilder<'a, 'b, A> { let start = _fbb.start_table(); - SortKeyBuilder { + SortKeySpecBuilder { fbb_: _fbb, start_: start, } } #[inline] - pub fn finish(self) -> ::flatbuffers::WIPOffset> { + pub fn finish(self) -> ::flatbuffers::WIPOffset> { let o = self.fbb_.end_table(self.start_); - self.fbb_.required(o, SortKey::VT_FIELDS,"fields"); + self.fbb_.required(o, SortKeySpec::VT_FIELDS,"fields"); ::flatbuffers::WIPOffset::new(o.value()) } } -impl ::core::fmt::Debug for SortKey<'_> { +impl ::core::fmt::Debug for SortKeySpec<'_> { fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { - let mut ds = f.debug_struct("SortKey"); + let mut ds = f.debug_struct("SortKeySpec"); ds.field("fields", &self.fields()); ds.finish() } @@ -1495,11 +1495,11 @@ impl<'a> Footer<'a> { } /// Sort key for the rows in this file. If absent, the row order is unknown. #[inline] - pub fn sort_key_spec(&self) -> Option> { + pub fn sort_key_spec(&self) -> Option> { // Safety: // Created from valid Table for this object // which contains a valid value in this slot - unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset>(Footer::VT_SORT_KEY_SPEC, None)} + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset>(Footer::VT_SORT_KEY_SPEC, None)} } } @@ -1514,7 +1514,7 @@ impl ::flatbuffers::Verifiable for Footer<'_> { .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, SegmentSpec>>>("segment_specs", Self::VT_SEGMENT_SPECS, false)? .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("compression_specs", Self::VT_COMPRESSION_SPECS, false)? .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("encryption_specs", Self::VT_ENCRYPTION_SPECS, false)? - .visit_field::<::flatbuffers::ForwardsUOffset>("sort_key_spec", Self::VT_SORT_KEY_SPEC, false)? + .visit_field::<::flatbuffers::ForwardsUOffset>("sort_key_spec", Self::VT_SORT_KEY_SPEC, false)? .finish(); Ok(()) } @@ -1525,7 +1525,7 @@ pub struct FooterArgs<'a> { pub segment_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, SegmentSpec>>>, pub compression_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, pub encryption_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, - pub sort_key_spec: Option<::flatbuffers::WIPOffset>>, + pub sort_key_spec: Option<::flatbuffers::WIPOffset>>, } impl<'a> Default for FooterArgs<'a> { #[inline] @@ -1567,8 +1567,8 @@ impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> FooterBuilder<'a, 'b, A> { self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(Footer::VT_ENCRYPTION_SPECS, encryption_specs); } #[inline] - pub fn add_sort_key_spec(&mut self, sort_key_spec: ::flatbuffers::WIPOffset>) { - self.fbb_.push_slot_always::<::flatbuffers::WIPOffset>(Footer::VT_SORT_KEY_SPEC, sort_key_spec); + pub fn add_sort_key_spec(&mut self, sort_key_spec: ::flatbuffers::WIPOffset>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset>(Footer::VT_SORT_KEY_SPEC, sort_key_spec); } #[inline] pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> FooterBuilder<'a, 'b, A> { diff --git a/vortex-file/src/sort_key/mod.rs b/vortex-file/src/sort_key/mod.rs index 91fb89e73f8..1f86d9097a0 100644 --- a/vortex-file/src/sort_key/mod.rs +++ b/vortex-file/src/sort_key/mod.rs @@ -90,7 +90,10 @@ pub(crate) fn validate_sort_key(dtype: &DType, fields: &[FieldPath]) -> VortexRe } /// Reads and validates the footer's sort key. -pub(crate) fn read_sort_key(key: fb::SortKey<'_>, dtype: &DType) -> VortexResult> { +pub(crate) fn read_sort_key( + key: fb::SortKeySpec<'_>, + dtype: &DType, +) -> VortexResult> { let fields: Arc<[_]> = key .fields() .iter() @@ -105,7 +108,7 @@ pub(crate) fn read_sort_key(key: fb::SortKey<'_>, dtype: &DType) -> VortexResult pub(crate) fn write_sort_key<'fb>( fbb: &mut FlatBufferBuilder<'fb>, fields: &[FieldPath], -) -> VortexResult>> { +) -> VortexResult>> { let mut offsets = Vec::with_capacity(fields.len()); for field in fields { @@ -128,9 +131,9 @@ pub(crate) fn write_sort_key<'fb>( } let fields = fbb.create_vector(&offsets); - Ok(fb::SortKey::create( + Ok(fb::SortKeySpec::create( fbb, - &fb::SortKeyArgs { + &fb::SortKeySpecArgs { fields: Some(fields), }, )) From b5abb7e1a84d990d024460a401164a59993e3c5c Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Wed, 21 Oct 2026 13:53:54 +0200 Subject: [PATCH 05/10] (feat) accept multiple sort keys When reviewing the code, I noticed that footer properties are always vectors, so I stopped to think about whether a vec of sort keys would make sense, and it turns out that yes, even DataFusion can support multiple sort keys. So what I did was just update it from a single sort key to multiple valid sort keys. Also added SortKey and SortKeyRef as type aliases. Signed-off-by: Joaquin Colacci --- .../flatbuffers/vortex-file/footer.fbs | 5 ++- .../src/flatbuffers/generated/footer.rs | 24 ++++++------- vortex-file/src/footer/file_layout.rs | 23 ++++++++----- vortex-file/src/footer/mod.rs | 26 +++++++------- vortex-file/src/footer/serializer.rs | 2 +- vortex-file/src/lib.rs | 2 ++ vortex-file/src/sort_key/mod.rs | 34 +++++++++++-------- 7 files changed, 65 insertions(+), 51 deletions(-) diff --git a/vortex-file/flatbuffers/vortex-file/footer.fbs b/vortex-file/flatbuffers/vortex-file/footer.fbs index 85a91bfa96b..72dd8cf4b30 100644 --- a/vortex-file/flatbuffers/vortex-file/footer.fbs +++ b/vortex-file/flatbuffers/vortex-file/footer.fbs @@ -120,7 +120,6 @@ table FieldPath { } /// Sort key supplied by the caller. Fields are compared in order, ascending with nulls first. -/// Duplicate keys are allowed. table SortKeySpec { fields: [FieldPath] (required); } @@ -138,8 +137,8 @@ table Footer { compression_specs: [CompressionSpec]; // Dictionary-encoded encryption specs, up to u16::MAX. encryption_specs: [EncryptionSpec]; - /// Sort key for the rows in this file. If absent, the row order is unknown. - sort_key_spec: SortKeySpec; + /// Sort keys that each describe an ordering of the rows in this file. + sort_key_specs: [SortKeySpec]; } /// An `ArraySpec` describes the type of a particular array. diff --git a/vortex-file/src/flatbuffers/generated/footer.rs b/vortex-file/src/flatbuffers/generated/footer.rs index d6230686be9..c44dc1f1ae1 100644 --- a/vortex-file/src/flatbuffers/generated/footer.rs +++ b/vortex-file/src/flatbuffers/generated/footer.rs @@ -1318,7 +1318,6 @@ pub enum SortKeySpecOffset {} #[derive(Copy, Clone, PartialEq)] /// Sort key supplied by the caller. Fields are compared in order, ascending with nulls first. -/// Duplicate keys are allowed. pub struct SortKeySpec<'a> { pub _tab: ::flatbuffers::Table<'a>, } @@ -1436,7 +1435,7 @@ impl<'a> Footer<'a> { pub const VT_SEGMENT_SPECS: ::flatbuffers::VOffsetT = 8; pub const VT_COMPRESSION_SPECS: ::flatbuffers::VOffsetT = 10; pub const VT_ENCRYPTION_SPECS: ::flatbuffers::VOffsetT = 12; - pub const VT_SORT_KEY_SPEC: ::flatbuffers::VOffsetT = 14; + pub const VT_SORT_KEY_SPECS: ::flatbuffers::VOffsetT = 14; #[inline] pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { @@ -1448,7 +1447,7 @@ impl<'a> Footer<'a> { args: &'args FooterArgs<'args> ) -> ::flatbuffers::WIPOffset> { let mut builder = FooterBuilder::new(_fbb); - if let Some(x) = args.sort_key_spec { builder.add_sort_key_spec(x); } + if let Some(x) = args.sort_key_specs { builder.add_sort_key_specs(x); } if let Some(x) = args.encryption_specs { builder.add_encryption_specs(x); } if let Some(x) = args.compression_specs { builder.add_compression_specs(x); } if let Some(x) = args.segment_specs { builder.add_segment_specs(x); } @@ -1493,13 +1492,14 @@ impl<'a> Footer<'a> { // which contains a valid value in this slot unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(Footer::VT_ENCRYPTION_SPECS, None)} } - /// Sort key for the rows in this file. If absent, the row order is unknown. + /// Sort keys that each describe an ordering of the rows in this file. + /// If absent or empty, the row order is unknown. #[inline] - pub fn sort_key_spec(&self) -> Option> { + pub fn sort_key_specs(&self) -> Option<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>> { // Safety: // Created from valid Table for this object // which contains a valid value in this slot - unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset>(Footer::VT_SORT_KEY_SPEC, None)} + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(Footer::VT_SORT_KEY_SPECS, None)} } } @@ -1514,7 +1514,7 @@ impl ::flatbuffers::Verifiable for Footer<'_> { .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, SegmentSpec>>>("segment_specs", Self::VT_SEGMENT_SPECS, false)? .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("compression_specs", Self::VT_COMPRESSION_SPECS, false)? .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("encryption_specs", Self::VT_ENCRYPTION_SPECS, false)? - .visit_field::<::flatbuffers::ForwardsUOffset>("sort_key_spec", Self::VT_SORT_KEY_SPEC, false)? + .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("sort_key_specs", Self::VT_SORT_KEY_SPECS, false)? .finish(); Ok(()) } @@ -1525,7 +1525,7 @@ pub struct FooterArgs<'a> { pub segment_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, SegmentSpec>>>, pub compression_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, pub encryption_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, - pub sort_key_spec: Option<::flatbuffers::WIPOffset>>, + pub sort_key_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, } impl<'a> Default for FooterArgs<'a> { #[inline] @@ -1536,7 +1536,7 @@ impl<'a> Default for FooterArgs<'a> { segment_specs: None, compression_specs: None, encryption_specs: None, - sort_key_spec: None, + sort_key_specs: None, } } } @@ -1567,8 +1567,8 @@ impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> FooterBuilder<'a, 'b, A> { self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(Footer::VT_ENCRYPTION_SPECS, encryption_specs); } #[inline] - pub fn add_sort_key_spec(&mut self, sort_key_spec: ::flatbuffers::WIPOffset>) { - self.fbb_.push_slot_always::<::flatbuffers::WIPOffset>(Footer::VT_SORT_KEY_SPEC, sort_key_spec); + pub fn add_sort_key_specs(&mut self, sort_key_specs: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset>>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(Footer::VT_SORT_KEY_SPECS, sort_key_specs); } #[inline] pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> FooterBuilder<'a, 'b, A> { @@ -1593,7 +1593,7 @@ impl ::core::fmt::Debug for Footer<'_> { ds.field("segment_specs", &self.segment_specs()); ds.field("compression_specs", &self.compression_specs()); ds.field("encryption_specs", &self.encryption_specs()); - ds.field("sort_key_spec", &self.sort_key_spec()); + ds.field("sort_key_specs", &self.sort_key_specs()); ds.finish() } } diff --git a/vortex-file/src/footer/file_layout.rs b/vortex-file/src/footer/file_layout.rs index a8cad66710a..a982eb6e486 100644 --- a/vortex-file/src/footer/file_layout.rs +++ b/vortex-file/src/footer/file_layout.rs @@ -11,7 +11,6 @@ use std::sync::Arc; use flatbuffers::FlatBufferBuilder; use flatbuffers::WIPOffset; -use vortex_array::dtype::FieldPath; use vortex_array::flatbuffers::FlatBufferRoot; use vortex_array::flatbuffers::WriteFlatBuffer; use vortex_error::VortexResult; @@ -19,6 +18,7 @@ use vortex_session::registry::ReadContext; use crate::flatbuffers::footer as fb; use crate::footer::segment::SegmentSpec; +use crate::sort_key::SortKeyRef; use crate::sort_key::write_sort_key; /// A writer for serializing a file layout to a FlatBuffer. @@ -32,8 +32,8 @@ pub(crate) struct FooterFlatBufferWriter { pub(crate) layout_ctx: ReadContext, /// Specifications for all segments in the file. pub(crate) segment_specs: Arc<[SegmentSpec]>, - /// Optional sort key, with fields listed in comparison order. - pub(crate) sort_key: Option>, + /// Declared sort keys, with each key's fields listed in comparison order. + pub(crate) sort_keys: Arc<[SortKeyRef]>, } impl FlatBufferRoot for FooterFlatBufferWriter {} @@ -69,11 +69,16 @@ impl WriteFlatBuffer for FooterFlatBufferWriter { }) .collect::>(); let layout_specs = fbb.create_vector(layout_specs.as_slice()); - let sort_key_spec = self - .sort_key - .as_deref() - .map(|key| write_sort_key(fbb, key)) - .transpose()?; + let sort_key_specs = if self.sort_keys.is_empty() { + None + } else { + let specs = self + .sort_keys + .iter() + .map(|key| write_sort_key(fbb, key)) + .collect::>>()?; + Some(fbb.create_vector(&specs)) + }; Ok(fb::Footer::create( fbb, @@ -83,7 +88,7 @@ impl WriteFlatBuffer for FooterFlatBufferWriter { layout_specs: Some(layout_specs), compression_specs: None, encryption_specs: None, - sort_key_spec, + sort_key_specs, }, )) } diff --git a/vortex-file/src/footer/mod.rs b/vortex-file/src/footer/mod.rs index d509e5fd7ad..d7b20651a17 100644 --- a/vortex-file/src/footer/mod.rs +++ b/vortex-file/src/footer/mod.rs @@ -27,7 +27,6 @@ use itertools::Itertools; pub use segment::*; use vortex_array::ArrayId; use vortex_array::dtype::DType; -use vortex_array::dtype::FieldPath; use vortex_array::flatbuffers::FlatBuffer; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; @@ -40,6 +39,7 @@ use vortex_session::VortexSession; use vortex_session::registry::ReadContext; use crate::flatbuffers::footer as fb; +use crate::sort_key::SortKeyRef; use crate::sort_key::read_sort_key; /// Maximum number of user-defined metadata segments. Keeps postscript bookkeeping small so the @@ -62,7 +62,7 @@ pub struct Footer { root_layout: LayoutRef, segments: Arc<[SegmentSpec]>, statistics: Option, - sort_key: Option>, + sort_keys: Arc<[SortKeyRef]>, metadata: Arc<[(String, SegmentSpec)]>, // `segments` followed by the `metadata` segments, which is the segment table a file is opened // with. Precomputed so opening a file with a cached footer does not copy every segment spec. @@ -92,7 +92,7 @@ impl Footer { segments_end: segments_end(&segments), segments, statistics, - sort_key: None, + sort_keys: Arc::from([]), metadata: Arc::from([]), array_read_ctx, approx_byte_size: None, @@ -104,8 +104,8 @@ impl Footer { self } - pub(crate) fn with_sort_key(mut self, sort_key: Option>) -> Self { - self.sort_key = sort_key; + pub(crate) fn with_sort_keys(mut self, sort_keys: Arc<[SortKeyRef]>) -> Self { + self.sort_keys = sort_keys; self } @@ -124,10 +124,12 @@ impl Footer { .sum(); let approx_byte_size = footer_bytes.len() + layout_bytes.len() + metadata_bytes; let fb_footer = root::(footer_bytes)?; - let sort_key = fb_footer - .sort_key_spec() + let sort_keys = fb_footer + .sort_key_specs() + .iter() + .flat_map(|keys| keys.iter()) .map(|key| read_sort_key(key, &dtype)) - .transpose()?; + .collect::>>()?; // Create a LayoutContext from the registry. let layout_specs = fb_footer.layout_specs(); @@ -178,7 +180,7 @@ impl Footer { segments_with_metadata, segments, statistics, - sort_key, + sort_keys, metadata, array_read_ctx, approx_byte_size: Some(approx_byte_size), @@ -200,9 +202,9 @@ impl Footer { self.statistics.as_ref() } - /// Returns the sort key supplied when the file was written. - pub fn sort_key(&self) -> Option<&[FieldPath]> { - self.sort_key.as_deref() + /// Returns the declared compound sort keys, with each key's fields in comparison order. + pub fn sort_keys(&self) -> &[SortKeyRef] { + &self.sort_keys } /// Returns the user-defined metadata segment locators stored in the postscript. diff --git a/vortex-file/src/footer/serializer.rs b/vortex-file/src/footer/serializer.rs index d9610cbc6c1..647e03a4854 100644 --- a/vortex-file/src/footer/serializer.rs +++ b/vortex-file/src/footer/serializer.rs @@ -161,7 +161,7 @@ impl FooterSerializer { ctx: self.footer.array_read_ctx.clone(), layout_ctx: ReadContext::new(layout_ctx.to_ids()), segment_specs: Arc::clone(&self.footer.segments), - sort_key: self.footer.sort_key.clone(), + sort_keys: Arc::clone(&self.footer.sort_keys), }, )?; buffers.push(buffer); diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index 393e69d6c14..b1b0553afbf 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -126,6 +126,8 @@ pub use file::*; pub use footer::*; pub use forever_constant::*; pub use open::*; +pub use sort_key::SortKey; +pub use sort_key::SortKeyRef; pub use strategy::*; use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; diff --git a/vortex-file/src/sort_key/mod.rs b/vortex-file/src/sort_key/mod.rs index 1f86d9097a0..7d9555f0931 100644 --- a/vortex-file/src/sort_key/mod.rs +++ b/vortex-file/src/sort_key/mod.rs @@ -9,17 +9,22 @@ //! # About sort keys //! //! Sort keys are a compound key that describes how the rows in a file are sorted. -//! For example, these rows are lexicographically sorted by (user_id, city): +//! A vortex file can store one or more sorting keys, each of which must hold +//! for the same rows. For example, these rows are lexicographically sorted by +//! both (user_id, city) and (timestamp, user_id): //! //! ```text -//! user_id | city -//! --------+----------- -//! 1 | Bangkok -//! 1 | Mumbai -//! 2 | Tokyo -//! 3 | Copenhagen +//! user_id | city | timestamp +//! --------+------------+---------- +//! 1 | Bangkok | 10 +//! 1 | Mumbai | 20 +//! 2 | Tokyo | 20 +//! 3 | Copenhagen | 30 //! ``` //! +//! The first sort key orders cities within each user. The second key orders user IDs +//! within each timestamp. +//! //! # Difference from sortedness stats //! //! The difference between `IsSorted` / `IsStrictSorted` and sort keys @@ -58,8 +63,12 @@ use vortex_utils::aliases::hash_set::HashSet; use crate::flatbuffers::footer as fb; +/// Represents a compound sort key, with fields listed in comparison order. +pub type SortKey = [FieldPath]; +pub type SortKeyRef = Arc; + /// Check that the key is non-empty and its fields exist and appear only once. -pub(crate) fn validate_sort_key(dtype: &DType, fields: &[FieldPath]) -> VortexResult<()> { +pub(crate) fn validate_sort_key(dtype: &DType, fields: &SortKey) -> VortexResult<()> { vortex_ensure!( !fields.is_empty(), "A file sort key must contain at least one field" @@ -90,11 +99,8 @@ pub(crate) fn validate_sort_key(dtype: &DType, fields: &[FieldPath]) -> VortexRe } /// Reads and validates the footer's sort key. -pub(crate) fn read_sort_key( - key: fb::SortKeySpec<'_>, - dtype: &DType, -) -> VortexResult> { - let fields: Arc<[_]> = key +pub(crate) fn read_sort_key(key: fb::SortKeySpec<'_>, dtype: &DType) -> VortexResult { + let fields: SortKeyRef = key .fields() .iter() .map(|field| field.path().iter().map(Field::from).collect()) @@ -107,7 +113,7 @@ pub(crate) fn read_sort_key( /// Writes the sort key into the footer. pub(crate) fn write_sort_key<'fb>( fbb: &mut FlatBufferBuilder<'fb>, - fields: &[FieldPath], + fields: &SortKey, ) -> VortexResult>> { let mut offsets = Vec::with_capacity(fields.len()); From ff64e956b8fda35b60065aa7b8b33818eab0546b Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Thu, 29 Oct 2026 21:30:14 +0100 Subject: [PATCH 06/10] (fix) cleanup vortex-file footer The previous work was incorrectly writing the sort keys in the footer as a spec rather than writing them as an aggregate function. This commit cleans up the previous incorrect work. Signed-off-by: Joaquin Colacci --- .../flatbuffers/vortex-file/footer.fbs | 13 -- .../src/flatbuffers/generated/footer.rs | 216 ------------------ vortex-file/src/footer/file_layout.rs | 15 -- vortex-file/src/footer/mod.rs | 21 -- vortex-file/src/footer/serializer.rs | 1 - vortex-file/src/lib.rs | 2 - vortex-file/src/sort_key/mod.rs | 146 ------------ 7 files changed, 414 deletions(-) delete mode 100644 vortex-file/src/sort_key/mod.rs diff --git a/vortex-file/flatbuffers/vortex-file/footer.fbs b/vortex-file/flatbuffers/vortex-file/footer.fbs index 72dd8cf4b30..636e8d4b6b7 100644 --- a/vortex-file/flatbuffers/vortex-file/footer.fbs +++ b/vortex-file/flatbuffers/vortex-file/footer.fbs @@ -113,17 +113,6 @@ table FileStatistics { partials: [ubyte]; } -/// A field path in the file's sort key. -table FieldPath { - /// Field names that make up the path. An empty path selects a non-struct root. - path: [string] (required); -} - -/// Sort key supplied by the caller. Fields are compared in order, ascending with nulls first. -table SortKeySpec { - fields: [FieldPath] (required); -} - /// The `Registry` object stores dictionary-encoded configuration for segments, /// compression schemes, encryption schemes, etc. table Footer { @@ -137,8 +126,6 @@ table Footer { compression_specs: [CompressionSpec]; // Dictionary-encoded encryption specs, up to u16::MAX. encryption_specs: [EncryptionSpec]; - /// Sort keys that each describe an ordering of the rows in this file. - sort_key_specs: [SortKeySpec]; } /// An `ArraySpec` describes the type of a particular array. diff --git a/vortex-file/src/flatbuffers/generated/footer.rs b/vortex-file/src/flatbuffers/generated/footer.rs index c44dc1f1ae1..9e0238c47e7 100644 --- a/vortex-file/src/flatbuffers/generated/footer.rs +++ b/vortex-file/src/flatbuffers/generated/footer.rs @@ -1215,203 +1215,6 @@ impl ::core::fmt::Debug for FileStatistics<'_> { ds.finish() } } -pub enum FieldPathOffset {} -#[derive(Copy, Clone, PartialEq)] - -/// A field path in the file's sort key. -pub struct FieldPath<'a> { - pub _tab: ::flatbuffers::Table<'a>, -} - -impl<'a> ::flatbuffers::Follow<'a> for FieldPath<'a> { - type Inner = FieldPath<'a>; - #[inline] - unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { - Self { _tab: unsafe { ::flatbuffers::Table::new(buf, loc) } } - } -} - -impl<'a> FieldPath<'a> { - pub const VT_PATH: ::flatbuffers::VOffsetT = 4; - - #[inline] - pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { - FieldPath { _tab: table } - } - #[allow(unused_mut)] - pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: ::flatbuffers::Allocator + 'bldr>( - _fbb: &'mut_bldr mut ::flatbuffers::FlatBufferBuilder<'bldr, A>, - args: &'args FieldPathArgs<'args> - ) -> ::flatbuffers::WIPOffset> { - let mut builder = FieldPathBuilder::new(_fbb); - if let Some(x) = args.path { builder.add_path(x); } - builder.finish() - } - - - /// Field names that make up the path. An empty path selects a non-struct root. - #[inline] - pub fn path(&self) -> ::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset<&'a str>> { - // Safety: - // Created from valid Table for this object - // which contains a valid value in this slot - unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset<&'a str>>>>(FieldPath::VT_PATH, None).unwrap()} - } -} - -impl ::flatbuffers::Verifiable for FieldPath<'_> { - #[inline] - fn run_verifier( - v: &mut ::flatbuffers::Verifier, pos: usize - ) -> Result<(), ::flatbuffers::InvalidFlatbuffer> { - v.visit_table(pos)? - .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset<&'_ str>>>>("path", Self::VT_PATH, true)? - .finish(); - Ok(()) - } -} -pub struct FieldPathArgs<'a> { - pub path: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset<&'a str>>>>, -} -impl<'a> Default for FieldPathArgs<'a> { - #[inline] - fn default() -> Self { - FieldPathArgs { - path: None, // required field - } - } -} - -pub struct FieldPathBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { - fbb_: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>, - start_: ::flatbuffers::WIPOffset<::flatbuffers::TableUnfinishedWIPOffset>, -} -impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> FieldPathBuilder<'a, 'b, A> { - #[inline] - pub fn add_path(&mut self, path: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset<&'b str>>>) { - self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(FieldPath::VT_PATH, path); - } - #[inline] - pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> FieldPathBuilder<'a, 'b, A> { - let start = _fbb.start_table(); - FieldPathBuilder { - fbb_: _fbb, - start_: start, - } - } - #[inline] - pub fn finish(self) -> ::flatbuffers::WIPOffset> { - let o = self.fbb_.end_table(self.start_); - self.fbb_.required(o, FieldPath::VT_PATH,"path"); - ::flatbuffers::WIPOffset::new(o.value()) - } -} - -impl ::core::fmt::Debug for FieldPath<'_> { - fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { - let mut ds = f.debug_struct("FieldPath"); - ds.field("path", &self.path()); - ds.finish() - } -} -pub enum SortKeySpecOffset {} -#[derive(Copy, Clone, PartialEq)] - -/// Sort key supplied by the caller. Fields are compared in order, ascending with nulls first. -pub struct SortKeySpec<'a> { - pub _tab: ::flatbuffers::Table<'a>, -} - -impl<'a> ::flatbuffers::Follow<'a> for SortKeySpec<'a> { - type Inner = SortKeySpec<'a>; - #[inline] - unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { - Self { _tab: unsafe { ::flatbuffers::Table::new(buf, loc) } } - } -} - -impl<'a> SortKeySpec<'a> { - pub const VT_FIELDS: ::flatbuffers::VOffsetT = 4; - - #[inline] - pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { - SortKeySpec { _tab: table } - } - #[allow(unused_mut)] - pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: ::flatbuffers::Allocator + 'bldr>( - _fbb: &'mut_bldr mut ::flatbuffers::FlatBufferBuilder<'bldr, A>, - args: &'args SortKeySpecArgs<'args> - ) -> ::flatbuffers::WIPOffset> { - let mut builder = SortKeySpecBuilder::new(_fbb); - if let Some(x) = args.fields { builder.add_fields(x); } - builder.finish() - } - - - #[inline] - pub fn fields(&self) -> ::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>> { - // Safety: - // Created from valid Table for this object - // which contains a valid value in this slot - unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(SortKeySpec::VT_FIELDS, None).unwrap()} - } -} - -impl ::flatbuffers::Verifiable for SortKeySpec<'_> { - #[inline] - fn run_verifier( - v: &mut ::flatbuffers::Verifier, pos: usize - ) -> Result<(), ::flatbuffers::InvalidFlatbuffer> { - v.visit_table(pos)? - .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("fields", Self::VT_FIELDS, true)? - .finish(); - Ok(()) - } -} -pub struct SortKeySpecArgs<'a> { - pub fields: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, -} -impl<'a> Default for SortKeySpecArgs<'a> { - #[inline] - fn default() -> Self { - SortKeySpecArgs { - fields: None, // required field - } - } -} - -pub struct SortKeySpecBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { - fbb_: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>, - start_: ::flatbuffers::WIPOffset<::flatbuffers::TableUnfinishedWIPOffset>, -} -impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> SortKeySpecBuilder<'a, 'b, A> { - #[inline] - pub fn add_fields(&mut self, fields: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset>>>) { - self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(SortKeySpec::VT_FIELDS, fields); - } - #[inline] - pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> SortKeySpecBuilder<'a, 'b, A> { - let start = _fbb.start_table(); - SortKeySpecBuilder { - fbb_: _fbb, - start_: start, - } - } - #[inline] - pub fn finish(self) -> ::flatbuffers::WIPOffset> { - let o = self.fbb_.end_table(self.start_); - self.fbb_.required(o, SortKeySpec::VT_FIELDS,"fields"); - ::flatbuffers::WIPOffset::new(o.value()) - } -} - -impl ::core::fmt::Debug for SortKeySpec<'_> { - fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { - let mut ds = f.debug_struct("SortKeySpec"); - ds.field("fields", &self.fields()); - ds.finish() - } -} pub enum FooterOffset {} #[derive(Copy, Clone, PartialEq)] @@ -1435,7 +1238,6 @@ impl<'a> Footer<'a> { pub const VT_SEGMENT_SPECS: ::flatbuffers::VOffsetT = 8; pub const VT_COMPRESSION_SPECS: ::flatbuffers::VOffsetT = 10; pub const VT_ENCRYPTION_SPECS: ::flatbuffers::VOffsetT = 12; - pub const VT_SORT_KEY_SPECS: ::flatbuffers::VOffsetT = 14; #[inline] pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { @@ -1447,7 +1249,6 @@ impl<'a> Footer<'a> { args: &'args FooterArgs<'args> ) -> ::flatbuffers::WIPOffset> { let mut builder = FooterBuilder::new(_fbb); - if let Some(x) = args.sort_key_specs { builder.add_sort_key_specs(x); } if let Some(x) = args.encryption_specs { builder.add_encryption_specs(x); } if let Some(x) = args.compression_specs { builder.add_compression_specs(x); } if let Some(x) = args.segment_specs { builder.add_segment_specs(x); } @@ -1492,15 +1293,6 @@ impl<'a> Footer<'a> { // which contains a valid value in this slot unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(Footer::VT_ENCRYPTION_SPECS, None)} } - /// Sort keys that each describe an ordering of the rows in this file. - /// If absent or empty, the row order is unknown. - #[inline] - pub fn sort_key_specs(&self) -> Option<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>> { - // Safety: - // Created from valid Table for this object - // which contains a valid value in this slot - unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(Footer::VT_SORT_KEY_SPECS, None)} - } } impl ::flatbuffers::Verifiable for Footer<'_> { @@ -1514,7 +1306,6 @@ impl ::flatbuffers::Verifiable for Footer<'_> { .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, SegmentSpec>>>("segment_specs", Self::VT_SEGMENT_SPECS, false)? .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("compression_specs", Self::VT_COMPRESSION_SPECS, false)? .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("encryption_specs", Self::VT_ENCRYPTION_SPECS, false)? - .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("sort_key_specs", Self::VT_SORT_KEY_SPECS, false)? .finish(); Ok(()) } @@ -1525,7 +1316,6 @@ pub struct FooterArgs<'a> { pub segment_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, SegmentSpec>>>, pub compression_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, pub encryption_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, - pub sort_key_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, } impl<'a> Default for FooterArgs<'a> { #[inline] @@ -1536,7 +1326,6 @@ impl<'a> Default for FooterArgs<'a> { segment_specs: None, compression_specs: None, encryption_specs: None, - sort_key_specs: None, } } } @@ -1567,10 +1356,6 @@ impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> FooterBuilder<'a, 'b, A> { self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(Footer::VT_ENCRYPTION_SPECS, encryption_specs); } #[inline] - pub fn add_sort_key_specs(&mut self, sort_key_specs: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset>>>) { - self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(Footer::VT_SORT_KEY_SPECS, sort_key_specs); - } - #[inline] pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> FooterBuilder<'a, 'b, A> { let start = _fbb.start_table(); FooterBuilder { @@ -1593,7 +1378,6 @@ impl ::core::fmt::Debug for Footer<'_> { ds.field("segment_specs", &self.segment_specs()); ds.field("compression_specs", &self.compression_specs()); ds.field("encryption_specs", &self.encryption_specs()); - ds.field("sort_key_specs", &self.sort_key_specs()); ds.finish() } } diff --git a/vortex-file/src/footer/file_layout.rs b/vortex-file/src/footer/file_layout.rs index a982eb6e486..c848a3f207f 100644 --- a/vortex-file/src/footer/file_layout.rs +++ b/vortex-file/src/footer/file_layout.rs @@ -18,8 +18,6 @@ use vortex_session::registry::ReadContext; use crate::flatbuffers::footer as fb; use crate::footer::segment::SegmentSpec; -use crate::sort_key::SortKeyRef; -use crate::sort_key::write_sort_key; /// A writer for serializing a file layout to a FlatBuffer. /// @@ -32,8 +30,6 @@ pub(crate) struct FooterFlatBufferWriter { pub(crate) layout_ctx: ReadContext, /// Specifications for all segments in the file. pub(crate) segment_specs: Arc<[SegmentSpec]>, - /// Declared sort keys, with each key's fields listed in comparison order. - pub(crate) sort_keys: Arc<[SortKeyRef]>, } impl FlatBufferRoot for FooterFlatBufferWriter {} @@ -69,16 +65,6 @@ impl WriteFlatBuffer for FooterFlatBufferWriter { }) .collect::>(); let layout_specs = fbb.create_vector(layout_specs.as_slice()); - let sort_key_specs = if self.sort_keys.is_empty() { - None - } else { - let specs = self - .sort_keys - .iter() - .map(|key| write_sort_key(fbb, key)) - .collect::>>()?; - Some(fbb.create_vector(&specs)) - }; Ok(fb::Footer::create( fbb, @@ -88,7 +74,6 @@ impl WriteFlatBuffer for FooterFlatBufferWriter { layout_specs: Some(layout_specs), compression_specs: None, encryption_specs: None, - sort_key_specs, }, )) } diff --git a/vortex-file/src/footer/mod.rs b/vortex-file/src/footer/mod.rs index d7b20651a17..57be23ec74e 100644 --- a/vortex-file/src/footer/mod.rs +++ b/vortex-file/src/footer/mod.rs @@ -39,8 +39,6 @@ use vortex_session::VortexSession; use vortex_session::registry::ReadContext; use crate::flatbuffers::footer as fb; -use crate::sort_key::SortKeyRef; -use crate::sort_key::read_sort_key; /// Maximum number of user-defined metadata segments. Keeps postscript bookkeeping small so the /// footer and required segments still fit the initial tail read. @@ -62,7 +60,6 @@ pub struct Footer { root_layout: LayoutRef, segments: Arc<[SegmentSpec]>, statistics: Option, - sort_keys: Arc<[SortKeyRef]>, metadata: Arc<[(String, SegmentSpec)]>, // `segments` followed by the `metadata` segments, which is the segment table a file is opened // with. Precomputed so opening a file with a cached footer does not copy every segment spec. @@ -92,7 +89,6 @@ impl Footer { segments_end: segments_end(&segments), segments, statistics, - sort_keys: Arc::from([]), metadata: Arc::from([]), array_read_ctx, approx_byte_size: None, @@ -104,11 +100,6 @@ impl Footer { self } - pub(crate) fn with_sort_keys(mut self, sort_keys: Arc<[SortKeyRef]>) -> Self { - self.sort_keys = sort_keys; - self - } - /// Read the [`Footer`] from a flatbuffer. pub(crate) fn from_flatbuffer( footer_bytes: &[u8], @@ -124,12 +115,6 @@ impl Footer { .sum(); let approx_byte_size = footer_bytes.len() + layout_bytes.len() + metadata_bytes; let fb_footer = root::(footer_bytes)?; - let sort_keys = fb_footer - .sort_key_specs() - .iter() - .flat_map(|keys| keys.iter()) - .map(|key| read_sort_key(key, &dtype)) - .collect::>>()?; // Create a LayoutContext from the registry. let layout_specs = fb_footer.layout_specs(); @@ -180,7 +165,6 @@ impl Footer { segments_with_metadata, segments, statistics, - sort_keys, metadata, array_read_ctx, approx_byte_size: Some(approx_byte_size), @@ -202,11 +186,6 @@ impl Footer { self.statistics.as_ref() } - /// Returns the declared compound sort keys, with each key's fields in comparison order. - pub fn sort_keys(&self) -> &[SortKeyRef] { - &self.sort_keys - } - /// Returns the user-defined metadata segment locators stored in the postscript. pub fn metadata_segments(&self) -> impl Iterator { self.metadata diff --git a/vortex-file/src/footer/serializer.rs b/vortex-file/src/footer/serializer.rs index 647e03a4854..e72e58c0389 100644 --- a/vortex-file/src/footer/serializer.rs +++ b/vortex-file/src/footer/serializer.rs @@ -161,7 +161,6 @@ impl FooterSerializer { ctx: self.footer.array_read_ctx.clone(), layout_ctx: ReadContext::new(layout_ctx.to_ids()), segment_specs: Arc::clone(&self.footer.segments), - sort_keys: Arc::clone(&self.footer.sort_keys), }, )?; buffers.push(buffer); diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index b1b0553afbf..393e69d6c14 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -126,8 +126,6 @@ pub use file::*; pub use footer::*; pub use forever_constant::*; pub use open::*; -pub use sort_key::SortKey; -pub use sort_key::SortKeyRef; pub use strategy::*; use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; diff --git a/vortex-file/src/sort_key/mod.rs b/vortex-file/src/sort_key/mod.rs deleted file mode 100644 index 7d9555f0931..00000000000 --- a/vortex-file/src/sort_key/mod.rs +++ /dev/null @@ -1,146 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Sort keys for Vortex files. -//! -//! The current implementations in this file are for validating sort keys, -//! and for writing and reading sort keys to/from the footer. -//! -//! # About sort keys -//! -//! Sort keys are a compound key that describes how the rows in a file are sorted. -//! A vortex file can store one or more sorting keys, each of which must hold -//! for the same rows. For example, these rows are lexicographically sorted by -//! both (user_id, city) and (timestamp, user_id): -//! -//! ```text -//! user_id | city | timestamp -//! --------+------------+---------- -//! 1 | Bangkok | 10 -//! 1 | Mumbai | 20 -//! 2 | Tokyo | 20 -//! 3 | Copenhagen | 30 -//! ``` -//! -//! The first sort key orders cities within each user. The second key orders user IDs -//! within each timestamp. -//! -//! # Difference from sortedness stats -//! -//! The difference between `IsSorted` / `IsStrictSorted` and sort keys -//! is that the stats describe the sortedness of one column, while a compound key -//! requires sortedness across the file only for its first column. The remaining -//! columns only need to be sorted where all preceding columns are equal. -//! -//! The ordering must also hold across chunks. For example, each chunk below -//! is sorted by (user_id, city): -//! -//! ```text -//! user_id | city -//! --------+----------- -//! Chunk 1 1 | Bangkok -//! 2 | Tokyo -//! --------+----------- -//! Chunk 2 2 | Mumbai <- breaks the ordering -//! 3 | Copenhagen -//! ``` -//! -//! But together, they break the sort key: (2, Tokyo) comes before (2, Mumbai), -//! even though Mumbai should come first. - -use std::sync::Arc; - -use flatbuffers::FlatBufferBuilder; -use flatbuffers::WIPOffset; -use vortex_array::dtype::DType; -use vortex_array::dtype::Field; -use vortex_array::dtype::FieldPath; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; -use vortex_error::vortex_err; -use vortex_utils::aliases::hash_set::HashSet; - -use crate::flatbuffers::footer as fb; - -/// Represents a compound sort key, with fields listed in comparison order. -pub type SortKey = [FieldPath]; -pub type SortKeyRef = Arc; - -/// Check that the key is non-empty and its fields exist and appear only once. -pub(crate) fn validate_sort_key(dtype: &DType, fields: &SortKey) -> VortexResult<()> { - vortex_ensure!( - !fields.is_empty(), - "A file sort key must contain at least one field" - ); - - let mut seen = HashSet::<&FieldPath>::default(); - for field in fields { - vortex_ensure!(seen.insert(field), "Duplicate sort field: {}", field); - - match (dtype, field.parts()) { - (DType::Struct(struct_fields, _), [Field::Name(name)]) => { - vortex_ensure!( - struct_fields.find(name).is_some(), - "Sort field {} does not exist in the file schema", - field - ); - } - // For a non-struct file, an empty path selects the entire root array as the sort field. - (dtype, []) if !dtype.is_struct() => {} - _ => vortex_bail!( - "Sort field {} must select a top-level field of a struct, or a non-struct root", - field - ), - } - } - - Ok(()) -} - -/// Reads and validates the footer's sort key. -pub(crate) fn read_sort_key(key: fb::SortKeySpec<'_>, dtype: &DType) -> VortexResult { - let fields: SortKeyRef = key - .fields() - .iter() - .map(|field| field.path().iter().map(Field::from).collect()) - .collect(); - validate_sort_key(dtype, &fields)?; - - Ok(fields) -} - -/// Writes the sort key into the footer. -pub(crate) fn write_sort_key<'fb>( - fbb: &mut FlatBufferBuilder<'fb>, - fields: &SortKey, -) -> VortexResult>> { - let mut offsets = Vec::with_capacity(fields.len()); - - for field in fields { - let names = field - .parts() - .iter() - .map(|part| { - let name = part - .as_name() - .ok_or_else(|| vortex_err!("Invalid sort field path"))?; - Ok(fbb.create_string(name)) - }) - .collect::>>()?; - - let path = fbb.create_vector(&names); - offsets.push(fb::FieldPath::create( - fbb, - &fb::FieldPathArgs { path: Some(path) }, - )); - } - - let fields = fbb.create_vector(&offsets); - Ok(fb::SortKeySpec::create( - fbb, - &fb::SortKeySpecArgs { - fields: Some(fields), - }, - )) -} From 3fcdefc10845a137fb5950f6ff18635e0fc1c3cc Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Fri, 30 Oct 2026 16:28:52 +0100 Subject: [PATCH 07/10] (feat) add sort order aggregate fn This commit adds the SortOrder aggregate fn into the `vortex-array` crate. The new `SortOrder` basically reuses the sort already implemented for fastpaths, so it has the same limitations (only ASC and nulls first). The `SortOptions` are ser/de using protobuf, and there is no support for strictness, but I could simple be added. I tried to add enough documentation to understand how the pair check works. Signed-off-by: Joaquin Colacci --- vortex-array/proto/expr.proto | 15 + vortex-array/src/aggregate_fn/fns/mod.rs | 1 + .../src/aggregate_fn/fns/sort_order/mod.rs | 1264 +++++++++++++++++ .../src/proto/generated/vortex.expr.rs | 52 + 4 files changed, 1332 insertions(+) create mode 100644 vortex-array/src/aggregate_fn/fns/sort_order/mod.rs diff --git a/vortex-array/proto/expr.proto b/vortex-array/proto/expr.proto index b24cd5e2c8f..489b6323d4a 100644 --- a/vortex-array/proto/expr.proto +++ b/vortex-array/proto/expr.proto @@ -32,6 +32,21 @@ message NumericalAggregateOpts { bool skip_nans = 1; } +// A declared lexicographic ordering, checked by `vortex.sort_order`. +message SortOrderOpts { + repeated SortField fields = 1; + + message SortField { + string name = 1; + SortDirection direction = 2; + } + + enum SortDirection { + ASC = 0; + DESC = 1; + } +} + // Options for `vortex.uncompressed_size_in_bytes`: whether the size is an approximation from the // buffers the array holds (the default) or the actual bytes its rows use. message UncompressedSizeInBytesOpts { diff --git a/vortex-array/src/aggregate_fn/fns/mod.rs b/vortex-array/src/aggregate_fn/fns/mod.rs index 2a0119820b1..e63e9bebe7a 100644 --- a/vortex-array/src/aggregate_fn/fns/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mod.rs @@ -20,6 +20,7 @@ pub mod min; pub mod min_max; pub mod nan_count; pub mod null_count; +pub mod sort_order; pub mod sum; pub mod sum_v2; pub mod uncompressed_size_in_bytes; diff --git a/vortex-array/src/aggregate_fn/fns/sort_order/mod.rs b/vortex-array/src/aggregate_fn/fns/sort_order/mod.rs new file mode 100644 index 00000000000..0b2bf416257 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/sort_order/mod.rs @@ -0,0 +1,1264 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Checks that rows are lexicographically sorted by a declared sequence of fields. +//! +//! # About sort keys +//! +//! A sort key is a compound key that describes how the rows in a file are sorted. +//! For example, these rows are lexicographically sorted by +//! both (user_id, city) and (timestamp, user_id): +//! +//! ```text +//! user_id | city | timestamp +//! --------+------------+---------- +//! 1 | Bangkok | 10 +//! 1 | Mumbai | 20 +//! 2 | Tokyo | 20 +//! 3 | Copenhagen | 30 +//! ``` +//! +//! The first sort key orders cities within each user. The second key orders user IDs +//! within each timestamp. +//! +//! # Difference from sortedness stats +//! +//! The difference between `IsSorted` / `IsStrictSorted` and sort keys +//! is that they describe the sortedness of one field, while a compound key +//! requires sortedness across the file only for its first field. The remaining +//! fields only need to be sorted where all preceding fields are equal. +//! +//! The ordering must also hold across chunks. For example, each chunk below +//! is sorted by (user_id, city): +//! +//! ```text +//! user_id | city +//! --------+----------- +//! Chunk 1 1 | Bangkok +//! 2 | Tokyo +//! --------+----------- +//! Chunk 2 2 | Mumbai <- breaks the ordering +//! 3 | Copenhagen +//! ``` +//! +//! But together, they break the sort key: (2, Tokyo) comes before (2, Mumbai), +//! even though Mumbai should come first. +//! +//! ## Notes +//! +//! Only `ASC` is supported for now. Whole-field checks reuse [`IsSorted`], +//! which only supports `ASC`. `DESC` is reserved in the options but rejected for now. +//! `DESC` support needs direction-aware checks to avoid reversing and materializing +//! the input. + +use std::cmp::Ordering; +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; +use std::sync::LazyLock; + +use prost::Message; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; +use vortex_utils::aliases::hash_set::HashSet; + +use crate::ArrayRef; +use crate::Columnar; +use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateArgs; +use crate::aggregate_fn::AggregateFnId; +use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::fns::is_sorted::IsSorted; +use crate::aggregate_fn::fns::is_sorted::IsSortedOptions; +use crate::aggregate_fn::fns::is_sorted::is_sorted; +use crate::arrays::BoolArray; +use crate::arrays::bool::BoolArrayExt; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::FieldName; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::StructFields; +use crate::proto::expr as pb; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::operators::Operator; + +/// The comparison direction of a field in a declared sort key. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)] +pub enum SortDirection { + /// Ascending order, with nulls first. + #[default] + Asc, + /// Descending order, reserved for future support and currently unsupported. + Desc, +} + +/// Options for [`SortOrder`]. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct SortOrderOptions { + /// Top-level struct fields and their directions, in priority order. + /// Currently only ascending directions are supported. + pub fields: Vec<(FieldName, SortDirection)>, +} + +impl SortOrderOptions { + /// Declare a lexicographic ordering over the given fields. + pub fn new(fields: impl IntoIterator, SortDirection)>) -> Self { + Self { + fields: fields + .into_iter() + .map(|(name, direction)| (name.into(), direction)) + .collect(), + } + } + + /// Validates that fields are non-empty, have no duplicates, + /// and use only `ASC`. + fn validate_fields(&self) -> VortexResult<()> { + vortex_ensure!( + !self.fields.is_empty(), + "A sort key must contain at least one field" + ); + + let mut seen = HashSet::<&FieldName>::default(); + for (name, direction) in &self.fields { + vortex_ensure!( + *direction == SortDirection::Asc, + "SortOrder currently supports ASC only" + ); + vortex_ensure!(seen.insert(name), "Duplicate sort field: {}", name); + } + + Ok(()) + } + + /// Validates the declared ordering against a struct input dtype. + /// + /// Dtypes must be supported by [`IsSorted`] (not strict). + pub fn validate(&self, dtype: &DType) -> VortexResult<()> { + self.validate_fields()?; + + let fields = dtype + .as_struct_fields_opt() + .ok_or_else(|| vortex_err!("SortOrder requires a struct input"))?; + + for (name, _) in &self.fields { + let dtype = fields.field(name).ok_or_else(|| { + vortex_err!("Sort field {} does not exist in the input dtype", name) + })?; + vortex_ensure!( + IsSorted + .return_dtype(&IsSortedOptions { strict: false }, &dtype) + .is_some(), + "Sort field {} has unsupported dtype {}", + name, + dtype + ); + } + + Ok(()) + } + + fn key_dtype(&self, dtype: &DType) -> VortexResult { + self.validate(dtype)?; + + let fields = dtype + .as_struct_fields_opt() + .vortex_expect("validated struct"); + + let dtypes = self + .fields + .iter() + .map(|(name, _)| { + fields + .field(name) + .vortex_expect("validated field") + .as_nullable() + }) + .collect(); + Ok(DType::Struct( + StructFields::new( + self.fields.iter().map(|(name, _)| name.clone()).collect(), + dtypes, + ), + Nullability::Nullable, + )) + } + + /// Compares lexicographically two scalars (rows) in `ASC` order with nulls first. + /// + /// Used for cross-chunks and cross-partial checks, where only the boundary rows + /// are compared. Intra-chunks checks use [`check_chunk_sortedness`] instead. + fn compare_sort_keys(&self, previous: &Scalar, next: &Scalar) -> VortexResult { + for index in 0..self.fields.len() { + let previous = previous + .as_struct() + .field_by_idx(index) + .vortex_expect("key field"); + + let next = next + .as_struct() + .field_by_idx(index) + .vortex_expect("key field"); + + let ordering = match (previous.is_null(), next.is_null()) { + (true, true) => Ordering::Equal, + (true, false) => Ordering::Less, + (false, true) => Ordering::Greater, + (false, false) => previous.partial_cmp(&next).ok_or_else(|| { + vortex_err!("Sort field {} cannot be compared", self.fields[index].0) + })?, + }; + + // Tie-break + if ordering != Ordering::Equal { + return Ok(ordering); + } + } + + Ok(Ordering::Equal) + } +} + +impl Display for SortOrderOptions { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + for (index, (name, direction)) in self.fields.iter().enumerate() { + if index > 0 { + write!(f, ", ")?; + } + + let direction = match direction { + SortDirection::Asc => "ASC", + SortDirection::Desc => "DESC", + }; + + write!(f, "{} {}", name, direction)?; + } + Ok(()) + } +} + +/// Checks that the rows inside one batch are sorted by the key fields (ASC, nulls +/// first). Equal neighbours are fine. +/// +/// `sort_fields` are the key columns in priority order, all the same length (at least +/// 2 rows), and there must be at least one. Only rows *within* the batch are compared. +/// The seam with the previous batch is checked separately by +/// [`SortOrderOptions::compare_sort_keys`]. +/// +/// # How it works +/// +/// A batch is sorted if every pair of neighbouring rows is in order. +/// For a compound key, "in order" depends on the field: +/// +/// - On the first field, `row[i + 1]` must not be smaller than `row[i]`. +/// - On later fields, that only matters if the two rows tied on every field before +/// it. If they differ on an earlier field, they belong to different groups and the +/// later field doesn't constrain them. +/// +/// So we keep a mask `tied` with one entry per pair. `tied[i]` is true while `row[i]` and `row[i + 1]` +/// agree on every field looked at so far. +/// +/// Example, sorted by `(country, city)`: +/// +/// ```text +/// row | country | city +/// ----+---------+------- +/// 0 | DE | Berlin +/// 1 | DE | Munich +/// 2 | FR | Lyon +/// 3 | FR | Paris +/// 4 | FR | Nice <- Paris > Nice, so this batch is not sorted +/// ``` +/// +/// There are 4 pairs, `(0,1) (1,2) (2,3) (3,4)`, and `tied` starts as `[T, T, T, T]`. +/// +/// **Field 0 (`country`).** Nothing is tied-off yet, so the whole column must simply +/// be sorted. We hand it to [`is_sorted`] to get its per-encoding fast paths. Then +/// we narrow `tied` to the pairs with equal countries: +/// +/// ```text +/// pair: (0,1) (1,2) (2,3) (3,4) +/// equal: T F T T (DE=DE, DE!=FR, FR=FR, FR=FR) +/// tied: T F T T +/// ``` +/// +/// **Field 1 (`city`).** Now we look for pairs where `city[i] > city[i + 1]`, but +/// only among tied pairs: +/// +/// ```text +/// pair: (0,1) (1,2) (2,3) (3,4) +/// cities: B < M M > L L < P P > N +/// out of order: F T F T +/// tied: T F T T +/// both: F F F T <- violation, return false +/// ``` +/// +/// Pair `(1,2)` is out of order too, but it crosses from DE to FR, so it starts a +/// new group and is ignored. +/// +/// After each field, `tied` is narrowed to the pairs that are also equal on that +/// field. If no pairs are tied anymore, later fields can't matter and we stop early. +/// +/// # Nulls +/// +/// Nulls sort first, and two nulls count as equal. So "out of order" also includes a +/// valid value followed by a null, and "equal" also includes null followed by null. +/// +/// These rules must match `compare_sort_keys`. Together the two define the result for the +/// whole input. +fn check_chunk_sortedness(sort_fields: &[ArrayRef], ctx: &mut ExecutionCtx) -> VortexResult { + let len = sort_fields[0].len(); + debug_assert!(len >= 2); + + let pair_count = len - 1; + let mut tied = Mask::new_true(pair_count); + + for (index, field) in sort_fields.iter().enumerate() { + let all_pairs_tied = tied.all_true(); + let is_last_field = index + 1 == sort_fields.len(); + + // If every pair is tied, reuse the encoding-specific whole-column check. + if all_pairs_tied { + if !is_sorted(field, ctx)? { + return Ok(false); + } + + if is_last_field { + return Ok(true); + } + } + + // Align each row with the next row. For city in the example above: + // pair: (0,1) (1,2) (2,3) (3,4) + // previous: Berlin Munich Lyon Paris + // next: Munich Lyon Paris Nice + let previous = field.slice(0..pair_count)?; + let next = field.slice(1..len)?; + + // For each pair, mark whether the previous and next values are non-null. + let previous_valid = previous.validity()?.execute_mask(pair_count, ctx)?; + let next_valid = next.validity()?.execute_mask(pair_count, ctx)?; + + // Otherwise, this field only constrains pairs tied on all earlier fields. + if !all_pairs_tied { + // Given that `SortOrder` uses only `ASC` and `NULLS` first, + // a valid value followed by null is also a violation. + let null_violation = &previous_valid & &!&next_valid; + let out_of_order = previous + .binary(next.clone(), Operator::Gt)? + .execute::(ctx)? + .to_mask_fill_null_false(ctx) + | &null_violation; + + if (&tied & &out_of_order).true_count() != 0 { + return Ok(false); + } + + if is_last_field { + return Ok(true); + } + } + + // Later fields matter only for pairs that are also equal on this field. + // Both values are null when neither side is valid. + let both_null = !(&previous_valid | &next_valid); + let equal_pairs = previous + .binary(next, Operator::Eq)? + .execute::(ctx)? + .to_mask_fill_null_false(ctx) + | &both_null; + + tied = tied & &equal_pairs; + if tied.all_false() { + break; + } + } + + Ok(true) +} + +/// Checks a declared lexicographic ordering across all accumulated batches. +/// +/// It is not `strict`, so equal row keys are allowed. Only `ASC` is supported for now, +/// DESC remains reserved in the options. +#[derive(Clone, Debug)] +pub struct SortOrder; + +/// Running verdict and boundary keys for [`SortOrder`]. +pub struct SortOrderPartial { + sorted: bool, + // None means no rows. A row whose key fields are null still has a non-null key tuple. + first: Option, + last: Option, +} + +impl SortOrderPartial { + /// The state of a group with no accumulated values. + fn empty() -> Self { + Self { + sorted: true, + first: None, + last: None, + } + } +} + +static NAMES: LazyLock = + LazyLock::new(|| FieldNames::from(["is_sorted", "first", "last"])); + +impl AggregateFnVTable for SortOrder { + type Options = SortOrderOptions; + type Partial = SortOrderPartial; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.sort_order"); + *ID + } + + /// Serializes [`SortOrderOptions`] as protobuf. + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + options.validate_fields()?; + Ok(Some( + pb::SortOrderOpts { + fields: options + .fields + .iter() + .map(|(name, direction)| pb::sort_order_opts::SortField { + name: name.to_string(), + direction: match direction { + SortDirection::Asc => pb::sort_order_opts::SortDirection::Asc, + SortDirection::Desc => pb::sort_order_opts::SortDirection::Desc, + } as i32, + }) + .collect(), + } + .encode_to_vec(), + )) + } + + /// Deserializes [`SortOrderOptions`] from protobuf. + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + let proto = pb::SortOrderOpts::decode(metadata)?; + let fields = proto + .fields + .into_iter() + .map(|field| { + let direction = match pb::sort_order_opts::SortDirection::try_from(field.direction)? + { + pb::sort_order_opts::SortDirection::Asc => SortDirection::Asc, + pb::sort_order_opts::SortDirection::Desc => SortDirection::Desc, + }; + Ok((field.name, direction)) + }) + .collect::>>()?; + let options = SortOrderOptions::new(fields); + options.validate_fields()?; + Ok(options) + } + + fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + options.validate(input_dtype).ok()?; + Some(DType::Bool(Nullability::NonNullable)) + } + + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + let key = options.key_dtype(input_dtype).ok()?; + + Some(DType::Struct( + StructFields::new( + NAMES.clone(), + vec![DType::Bool(Nullability::NonNullable), key.clone(), key], + ), + Nullability::Nullable, + )) + } + + fn empty_partial( + &self, + _args: AggregateArgs<'_, Self::Options>, + ) -> VortexResult { + Ok(SortOrderPartial::empty()) + } + + fn partial_from_scalar( + &self, + _args: AggregateArgs<'_, Self::Options>, + scalar: &Scalar, + ) -> VortexResult { + // A null struct means the producing accumulator was empty. + if scalar.is_null() { + return Ok(SortOrderPartial::empty()); + } + + let fields = scalar.as_struct(); + let first = fields.field_by_idx(1).vortex_expect("first key"); + let last = fields.field_by_idx(2).vortex_expect("last key"); + + vortex_ensure!( + !first.is_null() && !last.is_null(), + "Non-empty sortedness partials require boundary keys" + ); + + Ok(SortOrderPartial { + sorted: fields + .field_by_idx(0) + .vortex_expect("sorted verdict") + .as_bool() + .value() + .ok_or_else(|| vortex_err!("Sortedness partials require a non-null verdict"))?, + first: Some(first), + last: Some(last), + }) + } + + fn merge_partials( + &self, + args: AggregateArgs<'_, Self::Options>, + mut acc: Self::Partial, + partial: Self::Partial, + ) -> VortexResult { + if !acc.sorted { + return Ok(acc); + } + + acc.sorted &= partial.sorted; + + // Given that we only handle `ASC` for now, + // the compare only checks that the incoming partial first key + // is greater than the last acc. key. For `DESC` it would be the inverse. + if acc.sorted + && let (Some(previous), Some(next)) = (&acc.last, &partial.first) + { + acc.sorted = args.options.compare_sort_keys(previous, next)? != Ordering::Greater; + } + + if acc.first.is_none() { + acc.first = partial.first; + } + + if partial.last.is_some() { + acc.last = partial.last; + } + + Ok(acc) + } + + fn to_scalar( + &self, + args: AggregateArgs<'_, Self::Options>, + partial: &Self::Partial, + ) -> VortexResult { + let Some(first) = &partial.first else { + return Ok(Scalar::null(args.partial_dtype.clone())); + }; + + Ok(Scalar::struct_( + args.partial_dtype.clone(), + vec![ + Scalar::bool(partial.sorted, Nullability::NonNullable), + first.clone(), + partial + .last + .clone() + .vortex_expect("last key accompanies first key"), + ], + )) + } + + /// As with `IsSorted`, an unsorted partial makes the global result unsorted (saturated). + #[inline] + fn is_saturated( + &self, + _args: AggregateArgs<'_, Self::Options>, + partial: &Self::Partial, + ) -> bool { + !partial.sorted + } + + /// Handled here instead of in `accumulate` so we can read just the key fields from + /// the batch's original encoding, rather than decoding the whole struct to + /// [`Columnar`] first. `accumulate` would give the same result, only slower. + fn try_accumulate( + &self, + args: AggregateArgs<'_, Self::Options>, + partial: &mut Self::Partial, + batch: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + if batch.is_empty() || !partial.sorted { + return Ok(true); + } + + let sort_fields = args + .options + .fields + .iter() + .map(|(name, _)| batch.get_item(name.clone())) + .collect::>>()?; + + let key_dtype = args + .partial_dtype + .as_struct_fields_opt() + .and_then(|fields| fields.field_by_index(1)) + .vortex_expect("SortOrder partial contains the first key dtype"); + + let first = Scalar::struct_( + key_dtype.clone(), + sort_fields + .iter() + .map(|field| field.execute_scalar(0, ctx).map(Scalar::into_nullable)) + .collect::>>()?, + ); + + let last = Scalar::struct_( + key_dtype, + sort_fields + .iter() + .map(|field| { + field + .execute_scalar(batch.len() - 1, ctx) + .map(Scalar::into_nullable) + }) + .collect::>>()?, + ); + + if let Some(previous) = &partial.last { + partial.sorted = args.options.compare_sort_keys(previous, &first)? != Ordering::Greater; + } + + // Check fields sortedness + if partial.sorted && batch.len() > 1 { + partial.sorted = check_chunk_sortedness(&sort_fields, ctx)?; + } + + if partial.first.is_none() { + partial.first = Some(first); + } + + partial.last = Some(last); + + Ok(true) + } + + fn accumulate( + &self, + _args: AggregateArgs<'_, Self::Options>, + _partial: &mut Self::Partial, + _batch: &Columnar, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + unreachable!("SortOrder::try_accumulate handles all arrays") + } + + fn finalize( + &self, + _args: AggregateArgs<'_, Self::Options>, + partials: ArrayRef, + ) -> VortexResult { + partials + .get_item(NAMES.get(0).vortex_expect("out of bounds").clone())? + .fill_null(Scalar::bool(true, Nullability::NonNullable)) + } + + fn finalize_scalar( + &self, + _args: AggregateArgs<'_, Self::Options>, + partial: &Self::Partial, + ) -> VortexResult { + Ok(Scalar::bool(partial.sorted, Nullability::NonNullable)) + } +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use prost::Message; + use rstest::rstest; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_buffer::buffer_mut; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + use vortex_error::vortex_bail; + use vortex_session::VortexSession; + + use super::SortDirection; + use super::SortOrder; + use super::SortOrderOptions; + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateFnVTable; + use crate::aggregate_fn::DynAccumulator; + use crate::aggregate_fn::fns::is_sorted::is_sorted; + use crate::arrays::BoolArray; + use crate::arrays::ConstantArray; + use crate::arrays::DecimalArray; + use crate::arrays::DictArray; + use crate::arrays::PiecewiseSequenceArray; + use crate::arrays::PrimitiveArray; + use crate::arrays::RunEnd; + use crate::arrays::StructArray; + use crate::arrays::VarBinViewArray; + use crate::builtins::ArrayBuiltins; + use crate::dtype::DType; + use crate::dtype::DecimalDType; + use crate::dtype::Nullability; + use crate::proto::expr as pb; + use crate::scalar::Scalar; + use crate::validity::Validity; + + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + + fn order_lines(values: &[(i32, i32)]) -> VortexResult { + // Unsorted metadata before the key fields catches accidental field selection. + let is_expedited = + BoolArray::from_iter(values.iter().enumerate().map(|(index, _)| index % 2 == 0)) + .into_array(); + let order_ids = PrimitiveArray::from_iter(values.iter().map(|row| row.0)).into_array(); + let line_item_ids = PrimitiveArray::from_iter(values.iter().map(|row| row.1)).into_array(); + + Ok(StructArray::from_fields(&[ + ("is_expedited", is_expedited), + ("order_id", order_ids), + ("line_item_id", line_item_ids), + ])? + .into_array()) + } + + fn order_line_options() -> SortOrderOptions { + SortOrderOptions::new([ + ("order_id", SortDirection::Asc), + ("line_item_id", SortDirection::Asc), + ]) + } + + // Each tuple is (order_id, line_item_id); each inner vector is one input batch. + #[rstest] + #[case::secondary_restarts( + vec![ + vec![(1, 7), (1, 9)], + vec![(2, 0)], + ], + true, + )] + #[case::wrong_secondary( + vec![vec![(1, 9), (1, 7)]], + false, + )] + #[case::bad_boundary( + vec![ + vec![(1, 4), (1, 7)], + vec![(1, 6), (2, 9)], + ], + false, + )] + #[case::wrong_primary( + vec![vec![(2, 7), (1, 4)]], + false, + )] + #[case::duplicates( + vec![vec![(1, 7)], vec![(1, 7)]], + true, + )] + #[case::empty(vec![vec![]], true)] + #[case::singleton(vec![vec![(1, 7)]], true)] + #[case::empty_between( + vec![ + vec![(1, 7)], + vec![], + vec![(1, 6)], + ], + false, + )] + fn checks_declared_key( + #[case] chunks: Vec>, + #[case] expected: bool, + #[values(false, true)] combine_partials: bool, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let dtype = order_lines(&[])?.dtype().clone(); + let options = order_line_options(); + let mut accumulator = Accumulator::try_new(SortOrder, options.clone(), dtype.clone())?; + + for chunk in chunks { + let array = order_lines(&chunk)?; + + if combine_partials { + let mut batch_accumulator = + Accumulator::try_new(SortOrder, options.clone(), dtype.clone())?; + batch_accumulator.accumulate(&array, &mut ctx)?; + let partial = batch_accumulator.partial_scalar()?; + + accumulator.combine_partials(&partial)?; + } else { + accumulator.accumulate(&array, &mut ctx)?; + } + } + + assert_eq!( + accumulator.finish()?, + Scalar::bool(expected, Nullability::NonNullable) + ); + Ok(()) + } + + #[rstest] + #[case::order_ids_restart_per_line_item( + vec![(7, 1), (9, 1), (0, 2)], + true, + )] + #[case::line_item_ids_decrease( + vec![(1, 2), (2, 1)], + false, + )] + #[case::order_ids_decrease_within_line_item( + vec![(9, 1), (7, 1)], + false, + )] + fn field_priority_follows_declared_order( + #[case] values: Vec<(i32, i32)>, + #[case] expected: bool, + #[values(false, true)] one_row_batches: bool, + ) -> VortexResult<()> { + check_batches( + order_lines(&values)?, + SortOrderOptions::new([ + ("line_item_id", SortDirection::Asc), + ("order_id", SortDirection::Asc), + ]), + one_row_batches, + expected, + ) + } + + #[rstest] + #[case::nulls_first(vec![None, Some(2), Some(3)], true)] + #[case::nulls_last(vec![Some(2), Some(3), None], false)] + #[case::null_in_middle(vec![Some(2), None, Some(3)], false)] + #[case::all_null(vec![None, None], true)] + fn nulls_must_be_first( + #[case] values: Vec>, + #[case] expected: bool, + #[values(false, true)] one_row_batches: bool, + ) -> VortexResult<()> { + let array = StructArray::from_fields(&[( + "order_id", + PrimitiveArray::from_option_iter(values).into_array(), + )])? + .into_array(); + + check_batches( + array, + SortOrderOptions::new([("order_id", SortDirection::Asc)]), + one_row_batches, + expected, + ) + } + + fn check_batches( + array: ArrayRef, + options: SortOrderOptions, + one_row_batches: bool, + expected: bool, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let mut accumulator = Accumulator::try_new(SortOrder, options, array.dtype().clone())?; + + if one_row_batches { + for index in 0..array.len() { + let batch = array.slice(index..index + 1)?; + accumulator.accumulate(&batch, &mut ctx)?; + } + } else { + accumulator.accumulate(&array, &mut ctx)?; + } + + assert_eq!( + accumulator.finish()?, + Scalar::bool(expected, Nullability::NonNullable) + ); + Ok(()) + } + + #[rstest] + #[case::null_first( + vec![(100, -100), (2, 7), (3, 9)], + vec![false, true, true], + true, + )] + #[case::null_last( + vec![(2, 7), (3, 9), (-100, 100)], + vec![true, true, false], + false, + )] + fn parent_nulls_mask_sort_fields( + #[case] values: Vec<(i32, i32)>, + #[case] valid: Vec, + #[case] expected: bool, + #[values(false, true)] one_row_batches: bool, + ) -> VortexResult<()> { + let array = order_lines(&values)?.mask(BoolArray::from_iter(valid).into_array())?; + + check_batches(array, order_line_options(), one_row_batches, expected) + } + + #[rstest] + #[case::nan_at_end(buffer![1.0, 2.0, f64::NAN], true)] + #[case::nan_before_finite(buffer![1.0, f64::NAN, 2.0], false)] + #[case::signed_zero(buffer![-0.0, 0.0], true)] + fn float_order_matches_across_batches( + #[case] prices: Buffer, + #[case] expected: bool, + #[values(false, true)] one_row_batches: bool, + ) -> VortexResult<()> { + let array = StructArray::from_fields(&[("price", prices.into_array())])?.into_array(); + + check_batches( + array, + SortOrderOptions::new([("price", SortDirection::Asc)]), + one_row_batches, + expected, + ) + } + + #[rstest] + #[case::all_prefix_fields( + buffer![1i32, 1, 2, 2], // customer IDs + buffer![1i32, 2, 2, 2], // order IDs + buffer![9i32, 1, 9, 10], // line-item IDs + true, + )] + #[case::wrong_third_field( + buffer![1i32, 1, 2, 2], + buffer![1i32, 2, 2, 2], + buffer![9i32, 1, 10, 9], + false, + )] + fn groups_use_the_whole_prefix( + #[case] customer_ids: Buffer, + #[case] order_ids: Buffer, + #[case] line_item_ids: Buffer, + #[case] expected: bool, + #[values(false, true)] one_row_batches: bool, + ) -> VortexResult<()> { + let array = StructArray::from_fields(&[ + ("customer_id", customer_ids.into_array()), + ("order_id", order_ids.into_array()), + ("line_item_id", line_item_ids.into_array()), + ])? + .into_array(); + + check_batches( + array, + SortOrderOptions::new([ + ("customer_id", SortDirection::Asc), + ("order_id", SortDirection::Asc), + ("line_item_id", SortDirection::Asc), + ]), + one_row_batches, + expected, + ) + } + + #[rstest] + #[case::null_prefix( + PrimitiveArray::from_option_iter([None::, None]).into_array(), + )] + #[case::nan_prefix(buffer![f64::NAN; 2].into_array())] + fn equal_null_and_nan_prefixes_require_sorted_secondary( + #[case] prefix: ArrayRef, + ) -> VortexResult<()> { + let array = StructArray::from_fields(&[ + ("price", prefix), + ("line_item_id", buffer![2i32, 1].into_array()), + ])? + .into_array(); + + check_batches( + array, + SortOrderOptions::new([ + ("price", SortDirection::Asc), + ("line_item_id", SortDirection::Asc), + ]), + false, + false, + ) + } + + #[rstest] + #[case::sorted(buffer![1i32, 2, 2, 3].into_array())] + #[case::unsorted(buffer![1i32, 3, 2].into_array())] + fn single_field_matches_is_sorted(#[case] values: ArrayRef) -> VortexResult<()> { + let expected = is_sorted(&values, &mut SESSION.create_execution_ctx())?; + let array = StructArray::from_fields(&[("order_id", values)])?.into_array(); + + check_batches( + array, + SortOrderOptions::new([("order_id", SortDirection::Asc)]), + false, + expected, + ) + } + + #[test] + fn options_roundtrip() -> VortexResult<()> { + let options = order_line_options(); + let bytes = SortOrder + .serialize(&options)? + .vortex_expect("serializable options"); + + assert_eq!(SortOrder.deserialize(&bytes, &SESSION)?, options); + assert_eq!(options.to_string(), "order_id ASC, line_item_id ASC"); + + assert!( + SortOrder + .zone_stat_default(&DType::from(crate::dtype::PType::I32)) + .is_none() + ); + + Ok(()) + } + + #[rstest] + #[case::empty(vec![])] + #[case::duplicate(vec![("order_id", 0), ("order_id", 0)])] + #[case::unknown_direction(vec![("order_id", 0), ("line_item_id", 99)])] + #[case::descending_primary(vec![("order_id", 1), ("line_item_id", 0)])] + #[case::descending_secondary(vec![("order_id", 0), ("line_item_id", 1)])] + fn invalid_serialized_options_are_rejected(#[case] fields: Vec<(&str, i32)>) { + let bytes = pb::SortOrderOpts { + fields: fields + .into_iter() + .map(|(name, direction)| pb::sort_order_opts::SortField { + name: name.into(), + direction, + }) + .collect(), + } + .encode_to_vec(); + + assert!(SortOrder.deserialize(&bytes, &SESSION).is_err()); + } + + #[rstest] + #[case::empty(vec![])] + #[case::missing(vec!["missing_order_id"])] + #[case::duplicate(vec!["order_id", "order_id"])] + fn invalid_keys_are_rejected(#[case] fields: Vec<&str>) -> VortexResult<()> { + let dtype = order_lines(&[])?.dtype().clone(); + let options = + SortOrderOptions::new(fields.into_iter().map(|name| (name, SortDirection::Asc))); + + assert!(options.validate(&dtype).is_err()); + assert!(Accumulator::try_new(SortOrder, options, dtype).is_err()); + + Ok(()) + } + + #[test] + fn unsupported_input_and_field_are_rejected() -> VortexResult<()> { + let options = SortOrderOptions::new([("order_id", SortDirection::Asc)]); + let primitive_input = buffer![1i32].into_array(); + let nested_input = + StructArray::from_fields(&[("order_id", order_lines(&[(1, 2)])?)])?.into_array(); + + assert!(options.validate(primitive_input.dtype()).is_err()); + assert!(options.validate(nested_input.dtype()).is_err()); + + Ok(()) + } + + #[test] + fn descending_is_reserved_but_rejected() -> VortexResult<()> { + let dtype = order_lines(&[])?.dtype().clone(); + let options = SortOrderOptions::new([ + ("order_id", SortDirection::Asc), + ("line_item_id", SortDirection::Desc), + ]); + let error = options + .validate(&dtype) + .expect_err("DESC is not implemented"); + + assert_eq!(options.to_string(), "order_id ASC, line_item_id DESC"); + assert!(error.to_string().contains("ASC only")); + assert!(Accumulator::try_new(SortOrder, options.clone(), dtype).is_err()); + assert!(SortOrder.serialize(&options).is_err()); + + Ok(()) + } + + #[rstest] + #[case::distinct_zero(buffer![-0.0, 0.0], true)] + #[case::equal_zero(buffer![-0.0; 2], false)] + fn float_prefix_equality_matches_total_order( + #[case] prices: Buffer, + #[case] expected: bool, + #[values(false, true)] one_row_batches: bool, + ) -> VortexResult<()> { + let array = StructArray::from_fields(&[ + ("price", prices.into_array()), + ("line_item_id", buffer![2i32, 1].into_array()), + ])? + .into_array(); + + check_batches( + array, + SortOrderOptions::new([ + ("price", SortDirection::Asc), + ("line_item_id", SortDirection::Asc), + ]), + one_row_batches, + expected, + ) + } + + #[rstest] + #[case::nulls_first(vec![None, Some(2), None, Some(1)], true)] + #[case::nulls_last(vec![Some(2), None, None, Some(1)], false)] + #[case::descending_values(vec![None, Some(2), Some(2), Some(1)], false)] + fn secondary_nulls_in_tied_groups( + #[case] values: Vec>, + #[case] expected: bool, + #[values(false, true)] one_row_batches: bool, + ) -> VortexResult<()> { + let array = StructArray::from_fields(&[ + ("order_id", buffer![1i32, 1, 2, 2].into_array()), + ( + "line_item_id", + PrimitiveArray::from_option_iter(values).into_array(), + ), + ])? + .into_array(); + + check_batches(array, order_line_options(), one_row_batches, expected) + } + + #[rstest] + #[case::utf8( + VarBinViewArray::from_iter_str(["Berlin", "Paris", "Athens", "Berlin"]).into_array(), + true, + )] + #[case::bad_utf8( + VarBinViewArray::from_iter_str(["Paris", "Berlin", "Athens", "Berlin"]).into_array(), + false, + )] + #[case::binary( + VarBinViewArray::from_iter_bin([b"B2", b"C3", b"A1", b"B2"]).into_array(), + true, + )] + #[case::boolean( + BoolArray::from_iter([false, true, false, true]).into_array(), + true, + )] + #[case::bad_boolean( + BoolArray::from_iter([true, false, false, true]).into_array(), + false, + )] + #[case::decimal( + DecimalArray::new( + buffer![1i128, 2, 0, 1], + DecimalDType::new(10, 0), + Validity::NonNullable, + ) + .into_array(), + true, + )] + #[case::bad_decimal( + DecimalArray::new( + buffer![2i128, 1, 0, 1], + DecimalDType::new(10, 0), + Validity::NonNullable, + ) + .into_array(), + false, + )] + fn typed_secondary_fields( + #[case] values: ArrayRef, + #[case] expected: bool, + #[values(false, true)] one_row_batches: bool, + ) -> VortexResult<()> { + let array = StructArray::from_fields(&[ + ("order_id", buffer![1i32, 1, 2, 2].into_array()), + ("item_value", values), + ])? + .into_array(); + + check_batches( + array, + SortOrderOptions::new([ + ("order_id", SortDirection::Asc), + ("item_value", SortDirection::Asc), + ]), + one_row_batches, + expected, + ) + } + + #[rstest] + fn encoded_and_sliced_fields( + #[values("constant", "dict", "runend", "sequence")] encoding: &str, + #[values(false, true)] sliced: bool, + #[values(false, true)] one_row_batches: bool, + #[values(false, true)] reverse_secondary: bool, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let order_ids = match encoding { + "constant" => ConstantArray::new(1u64, 6).into_array(), + "dict" => DictArray::try_new( + buffer![0u32, 1, 1, 2, 2, 3].into_array(), + buffer![0u64, 1, 2, 3].into_array(), + )? + .into_array(), + "runend" => RunEnd::try_new( + buffer![1u64, 3, 5, 6].into_array(), + buffer![0u64, 1, 2, 3].into_array(), + &mut ctx, + )? + .into_array(), + "sequence" => PiecewiseSequenceArray::try_new( + ConstantArray::new(0u64, 1).into_array(), + ConstantArray::new(6u64, 1).into_array(), + ConstantArray::new(1u64, 1).into_array(), + 6, + )? + .into_array(), + _ => vortex_bail!("Unknown encoding fixture: {}", encoding), + }; + + let mut line_item_ids = if encoding == "constant" { + buffer_mut![0i32, 1, 2, 3, 4, 5] + } else { + buffer_mut![99i32, 1, 2, 0, 1, 99] + }; + if reverse_secondary { + line_item_ids.swap(1, 2); + } + + let mut array = StructArray::from_fields(&[ + ("order_id", order_ids), + ("line_item_id", line_item_ids.into_array()), + ])? + .into_array(); + + if sliced { + array = array.slice(1..5)?; + } + + // The sequence has distinct order IDs, so line-item IDs are unconstrained. + let expected = !reverse_secondary || encoding == "sequence"; + check_batches(array, order_line_options(), one_row_batches, expected) + } +} diff --git a/vortex-array/src/proto/generated/vortex.expr.rs b/vortex-array/src/proto/generated/vortex.expr.rs index 2faf3be3b1a..1588d1e8982 100644 --- a/vortex-array/src/proto/generated/vortex.expr.rs +++ b/vortex-array/src/proto/generated/vortex.expr.rs @@ -31,6 +31,58 @@ pub struct NumericalAggregateOpts { #[prost(bool, tag = "1")] pub skip_nans: bool, } +/// A declared lexicographic ordering, checked by `vortex.sort_order`. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct SortOrderOpts { + #[prost(message, repeated, tag = "1")] + pub fields: ::prost::alloc::vec::Vec, +} +/// Nested message and enum types in `SortOrderOpts`. +pub mod sort_order_opts { + #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] + pub struct SortField { + #[prost(string, tag = "1")] + pub name: ::prost::alloc::string::String, + #[prost(enumeration = "SortDirection", tag = "2")] + pub direction: i32, + } + #[derive( + Clone, + Copy, + Debug, + PartialEq, + Eq, + Hash, + PartialOrd, + Ord, + ::prost::Enumeration + )] + #[repr(i32)] + pub enum SortDirection { + Asc = 0, + Desc = 1, + } + impl SortDirection { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Asc => "ASC", + Self::Desc => "DESC", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "ASC" => Some(Self::Asc), + "DESC" => Some(Self::Desc), + _ => None, + } + } + } +} /// Options for `vortex.uncompressed_size_in_bytes`: whether the size is an approximation from the /// buffers the array holds (the default) or the actual bytes its rows use. #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] From 4b5e13ec3ce5b38049bd3fae59077d2b9260b873 Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Fri, 30 Oct 2026 16:37:07 +0100 Subject: [PATCH 08/10] (bench) add bench to measure sortedness/pairness checks I think the most important to measure is for utf-8 and the single column bench to assert that not much on top of the fastest case (single column/field sort), similar to `IsSorted`. Signed-off-by: Joaquin Colacci --- vortex-array/Cargo.toml | 4 + vortex-array/benches/sort_order.rs | 312 +++++++++++++++++++++++++++++ 2 files changed, 316 insertions(+) create mode 100644 vortex-array/benches/sort_order.rs diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 58cc4e2d9d7..dc86a4187be 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -360,3 +360,7 @@ harness = false [[bench]] name = "run_end_decode_primitive" harness = false + +[[bench]] +name = "sort_order" +harness = false diff --git a/vortex-array/benches/sort_order.rs b/vortex-array/benches/sort_order.rs new file mode 100644 index 00000000000..5575ba8b162 --- /dev/null +++ b/vortex-array/benches/sort_order.rs @@ -0,0 +1,312 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Sort order benchmark. +//! +//! Covers order and line-item keys with dense and sparse ties, batch boundaries, +//! nullable and UTF-8 fields, and constant and dictionary encodings. Also compares +//! single-field `SortOrder` with `IsSorted` on the same ascending order IDs. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use mimalloc::MiMalloc; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; +use vortex_array::aggregate_fn::fns::is_sorted::IsSortedOptions; +use vortex_array::aggregate_fn::fns::sort_order::SortDirection; +use vortex_array::aggregate_fn::fns::sort_order::SortOrder; +use vortex_array::aggregate_fn::fns::sort_order::SortOrderOptions; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_buffer::Buffer; +use vortex_session::VortexSession; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); +const ROWS: u32 = 65_536; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +fn order_line_options() -> SortOrderOptions { + SortOrderOptions::new([ + ("order_id", SortDirection::Asc), + ("line_item_id", SortDirection::Asc), + ]) +} + +fn measure( + bencher: Bencher, + vtable: V, + options: V::Options, + make_array: impl Fn() -> ArrayRef + Sync, +) { + bencher + .counter(ItemsCount::new(ROWS)) + .with_inputs(|| { + let array = make_array(); + let accumulator = + Accumulator::try_new(vtable.clone(), options.clone(), array.dtype().clone()) + .unwrap(); + (array, accumulator, SESSION.create_execution_ctx()) + }) + .bench_values(|(array, mut accumulator, mut ctx)| { + accumulator.accumulate(&array, &mut ctx).unwrap(); + let result = accumulator.finish().unwrap(); + assert_eq!(result.as_bool().value(), Some(true)); + divan::black_box(result) + }); +} + +// Slice during setup so the measurement covers accumulation and boundary comparisons. +fn measure_batches(bencher: Bencher, batch_rows: usize, make_array: impl Fn() -> ArrayRef + Sync) { + bencher + .counter(ItemsCount::new(ROWS)) + .with_inputs(|| { + let array = make_array(); + let accumulator = + Accumulator::try_new(SortOrder, order_line_options(), array.dtype().clone()) + .unwrap(); + let batches = (0..array.len()) + .step_by(batch_rows) + .map(|start| { + let end = (start + batch_rows).min(array.len()); + array.slice(start..end).unwrap() + }) + .collect::>(); + (batches, accumulator, SESSION.create_execution_ctx()) + }) + .bench_values(|(batches, mut accumulator, mut ctx)| { + for batch in &batches { + accumulator.accumulate(batch, &mut ctx).unwrap(); + } + let result = accumulator.finish().unwrap(); + assert_eq!(result.as_bool().value(), Some(true)); + divan::black_box(result) + }); +} + +#[divan::bench(args = [1, 2, 16, 65_536])] +fn order_lines_by_order_size(bencher: Bencher, items_per_order: u32) { + let order_ids = (0..ROWS) + .map(|index| u64::from(index / items_per_order)) + .collect::>(); + let line_item_ids = (0..ROWS) + .map(|index| u64::from(index % items_per_order)) + .collect::>(); + measure(bencher, SortOrder, order_line_options(), || { + StructArray::from_fields(&[ + ("order_id", order_ids.clone().into_array()), + ("line_item_id", line_item_ids.clone().into_array()), + ]) + .unwrap() + .into_array() + }); +} + +#[divan::bench] +fn repeated_order_line_key(bencher: Bencher) { + measure(bencher, SortOrder, order_line_options(), || { + StructArray::from_fields(&[ + ( + "order_id", + ConstantArray::new(1u64, ROWS as usize).into_array(), + ), + ( + "line_item_id", + ConstantArray::new(2u64, ROWS as usize).into_array(), + ), + ]) + .unwrap() + .into_array() + }); +} + +// Customer IDs repeat every 1024 rows. Order IDs preserve ties every 2048 rows. +// This leaves sparse tied pairs for the line-item IDs to compare. +#[divan::bench] +fn customer_order_lines_sparse_ties(bencher: Bencher) { + let customer_ids = (0..ROWS) + .map(|index| u64::from(index - index / 1024)) + .collect::>(); + let order_ids = (0..ROWS) + .map(|index| u64::from(index - index / 2048)) + .collect::>(); + let line_item_ids = (0..ROWS).map(u64::from).collect::>(); + measure( + bencher, + SortOrder, + SortOrderOptions::new([ + ("customer_id", SortDirection::Asc), + ("order_id", SortDirection::Asc), + ("line_item_id", SortDirection::Asc), + ]), + || { + StructArray::from_fields(&[ + ("customer_id", customer_ids.clone().into_array()), + ("order_id", order_ids.clone().into_array()), + ("line_item_id", line_item_ids.clone().into_array()), + ]) + .unwrap() + .into_array() + }, + ); +} + +// A batch size of 63 splits orders, so some boundary comparisons reach line-item IDs. +#[divan::bench(args = [63, 64, 1024, 16_384])] +fn order_lines_across_batches(bencher: Bencher, batch_rows: usize) { + let order_ids = (0..ROWS) + .map(|index| u64::from(index / 16)) + .collect::>(); + let line_item_ids = (0..ROWS) + .map(|index| u64::from(index % 16)) + .collect::>(); + measure_batches(bencher, batch_rows, || { + StructArray::from_fields(&[ + ("order_id", order_ids.clone().into_array()), + ("line_item_id", line_item_ids.clone().into_array()), + ]) + .unwrap() + .into_array() + }); +} + +// Rows without an order ID are sorted by line-item ID. Each order starts with a null item ID. +#[divan::bench] +fn nullable_order_line_keys(bencher: Bencher) { + let order_ids = (0..ROWS) + .map(|index| (index >= 1024).then(|| u64::from(index / 16))) + .collect::>(); + let line_item_ids = (0..ROWS) + .map(|index| { + if index < 1024 { + Some(u64::from(index)) + } else { + (index % 16 != 0).then(|| u64::from(index % 16)) + } + }) + .collect::>(); + measure(bencher, SortOrder, order_line_options(), || { + StructArray::from_fields(&[ + ( + "order_id", + PrimitiveArray::from_option_iter(order_ids.clone()).into_array(), + ), + ( + "line_item_id", + PrimitiveArray::from_option_iter(line_item_ids.clone()).into_array(), + ), + ]) + .unwrap() + .into_array() + }); +} + +// Zero-padding keeps lexicographic and numeric ordering consistent. +#[divan::bench] +fn customer_order_string_keys(bencher: Bencher) { + let customer_ids = (0..ROWS) + .map(|index| format!("customer-{:05}", index / 16)) + .collect::>(); + let order_ids = (0..ROWS) + .map(|index| format!("order-{:05}", index % 16)) + .collect::>(); + measure( + bencher, + SortOrder, + SortOrderOptions::new([ + ("customer_id", SortDirection::Asc), + ("order_id", SortDirection::Asc), + ]), + || { + StructArray::from_fields(&[ + ( + "customer_id", + VarBinViewArray::from_iter_str(customer_ids.iter()).into_array(), + ), + ( + "order_id", + VarBinViewArray::from_iter_str(order_ids.iter()).into_array(), + ), + ]) + .unwrap() + .into_array() + }, + ); +} + +#[divan::bench] +fn order_lines_with_dictionary_item_ids(bencher: Bencher) { + let order_ids = (0..ROWS) + .map(|index| u64::from(index / 16)) + .collect::>(); + let line_item_codes = (0..ROWS).map(|index| index % 16).collect::>(); + let line_item_dictionary = (0..16u64).collect::>(); + measure(bencher, SortOrder, order_line_options(), || { + StructArray::from_fields(&[ + ("order_id", order_ids.clone().into_array()), + ( + "line_item_id", + DictArray::try_new( + line_item_codes.clone().into_array(), + line_item_dictionary.clone().into_array(), + ) + .unwrap() + .into_array(), + ), + ]) + .unwrap() + .into_array() + }); +} + +/// The following bench is to compare against the `IsSorted` +/// aggregate fn. It shouldn't deviate much from it, given +/// that it is a single column/field. +mod single_order_id { + use super::*; + + fn order_ids() -> Buffer { + (0..ROWS).map(u64::from).collect() + } + + #[divan::bench] + fn sort_order(bencher: Bencher) { + let order_ids = order_ids(); + measure( + bencher, + SortOrder, + SortOrderOptions::new([("order_id", SortDirection::Asc)]), + || { + StructArray::from_fields(&[("order_id", order_ids.clone().into_array())]) + .unwrap() + .into_array() + }, + ); + } + + #[divan::bench] + fn is_sorted(bencher: Bencher) { + let order_ids = order_ids(); + measure(bencher, IsSorted, IsSortedOptions { strict: false }, || { + order_ids.clone().into_array() + }); + } +} From 83c21d3d83be1d5c5c0431e4f6b5c94bd564c7a3 Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Fri, 30 Oct 2026 16:40:13 +0100 Subject: [PATCH 09/10] (fix) test file write, and add the aggfn to the unstable edition Similar to other new agg fns, they need to be added to the unstable edition. Also, this commit adds a test to check that the aggregate fn is correctly written into the file stats, but it is not present in the layout zones. Signed-off-by: Joaquin Colacci --- docs/specs/file-format.md | 12 +- .../src/declarations/unstable/mod.rs | 1 + vortex-file/src/tests.rs | 113 ++++++++++++++++++ vortex-file/src/writer.rs | 17 +++ .../editions/unstable/unstable2026.10.0.toml | 2 + 5 files changed, 139 insertions(+), 6 deletions(-) diff --git a/docs/specs/file-format.md b/docs/specs/file-format.md index 4d8cf8c0229..1e47e9c1e05 100644 --- a/docs/specs/file-format.md +++ b/docs/specs/file-format.md @@ -19,11 +19,11 @@ definition that allows efficiently querying the layout. Other considerations for the Vortex file format include: -* File compatibility. Editions constrain which serialized components a writer can use. A writer +- File compatibility. Editions constrain which serialized components a writer can use. A writer built with newer Vortex crates can target an older edition that its intended readers support. -* Fine-grained encryption. -* Efficient access for both local disk and cloud storage. -* Minimal overhead reading few columns or rows from wide or long arrays. +- Fine-grained encryption. +- Efficient access for both local disk and cloud storage. +- Minimal overhead reading few columns or rows from wide or long arrays. ## File Specification @@ -54,7 +54,7 @@ The postscript contains the locations of: 1. a `dtype` segment representing the top-level logical data type (i.e., schema) 2. a `layout` segment containing the root `Layout` 3. a `statistics` segment containing file-level per-field statistics (e.g., minima and maxima of each field/column, for whole-file pruning) -4. a `footer` segment containing a dictionary-encoded _segment map_, and other shared configuration such as compression, encryption schemes or sort keys +4. a `footer` segment containing a dictionary-encoded _segment map_, and other shared configuration such as compression and encryption schemes 5. up to 16 user-defined `metadata` segments, each identified by a unique, non-empty UTF-8 key of at most 64 bytes The postscript carries a locator (offset, length, and alignment) for each metadata segment that is @@ -86,7 +86,7 @@ valid to store a `Float64` array, a `Boolean` array, or any other root data type The footer is a flat buffer serialized `Footer` object. This object contains all the information required to load the root `Layout` object into a usable `LayoutReader`. -For example, it contains the locations, compression schemes, encryption schemes, sort keys, and required alignment of all segments in the file. +For example, it contains the locations, compression schemes, encryption schemes, and required alignment of all segments in the file. :::{literalinclude} ../../vortex-file/flatbuffers/vortex-file/footer.fbs :start-after: [footer] diff --git a/vortex-edition/src/declarations/unstable/mod.rs b/vortex-edition/src/declarations/unstable/mod.rs index 23733d78027..2f015ce2cc1 100644 --- a/vortex-edition/src/declarations/unstable/mod.rs +++ b/vortex-edition/src/declarations/unstable/mod.rs @@ -40,6 +40,7 @@ pub static DECLARATION: EditionDeclaration = EditionDeclaration { EditionMember::array(&"vortex.union"), EditionMember::layout(&"vortex.list"), EditionMember::aggregate(&"vortex.bloom_filter.sbbf"), + EditionMember::aggregate(&"vortex.sort_order"), EditionMember::aggregate(&"vortex.sum_v2"), ], }; diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index ce4d7e54ba5..01ff961f17d 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -4,6 +4,7 @@ #![expect(clippy::cast_possible_truncation)] use std::fs; use std::iter; +use std::num::NonZeroUsize; use std::sync::Arc; use std::sync::LazyLock; @@ -23,7 +24,11 @@ use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::sort_order::SortDirection; +use vortex_array::aggregate_fn::fns::sort_order::SortOrder; +use vortex_array::aggregate_fn::fns::sort_order::SortOrderOptions; use vortex_array::aggregate_fn::fns::sum_v2::SumV2; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ConstantArray; @@ -41,6 +46,7 @@ use vortex_array::assert_arrays_eq; use vortex_array::builders::MapBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::FieldNames; use vortex_array::dtype::FieldPath; use vortex_array::dtype::MapDType; use vortex_array::dtype::Nullability; @@ -108,6 +114,8 @@ use vortex_layout::layouts::struct_::StructStrategy; use vortex_layout::layouts::table::TableStrategy; use vortex_layout::layouts::zoned::LegacyStats; use vortex_layout::layouts::zoned::Zoned; +use vortex_layout::layouts::zoned::writer::ZonedLayoutOptions; +use vortex_layout::layouts::zoned::writer::ZonedStrategy; use vortex_layout::scan::scan_builder::ScanBuilder; use vortex_layout::scan::split_by::DEFAULT_MAX_SPLIT_ROWS; use vortex_layout::scan::split_by::SplitBy; @@ -2353,6 +2361,111 @@ async fn test_writer_with_statistics() -> VortexResult<()> { Ok(()) } +/// Assert that `SortOrder` is written in the +/// file stats, and also check that nothing was written +/// as zone stats. +#[tokio::test] +async fn file_footer_contains_sort_order() -> VortexResult<()> { + let session = VortexSession::empty() + .with::() + .with::() + .with::(); + + crate::register_default_encodings(&session); + session.aggregate_fns().register(SortOrder); + for declaration in EDITION_DECLARATIONS { + session.register_edition(declaration)?; + } + session.enable_edition(CORE_2026_08_3)?; + session.enable_edition(UNSTABLE)?; + + let array = StructArray::new( + FieldNames::from(["numbers"]), + [buffer![1u32, 2, 3, 4, 5].into_array()], + 5, + Validity::AllValid, + ) + .into_array(); + + let aggregate = SortOrder.bind(SortOrderOptions::new([("numbers", SortDirection::Asc)])); + + // The default strategy skips zone maps beneath nullable structs. Write zones at the + // root so the assertions below inspect stored zone statistics alongside file statistics. + let zone_len = NonZeroUsize::new(2).vortex_expect("non-zero zone length"); + let strategy = ZonedStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + FlatLayoutStrategy::default(), + ZonedLayoutOptions { + block_size: zone_len, + ..Default::default() + }, + ); + let mut buf = ByteBufferMut::empty(); + let mut writer = session + .write_options() + .with_strategy(Arc::new(strategy)) + .with_file_statistics(vec![aggregate.clone()]) + .writer(&mut buf, array.dtype().clone()); + + for start in (0..array.len()).step_by(zone_len.get()) { + let end = (start + zone_len.get()).min(array.len()); + writer.push(array.slice(start..end)?).await?; + } + + let summary = writer.finish().await?; + let statistics = summary + .footer() + .statistics() + .vortex_expect("statistics were written"); + let (root_stats, _) = statistics + .get_by_path(&FieldPath::root()) + .vortex_expect("root statistics were written"); + + assert_eq!( + root_stats.get(&aggregate), + Precision::exact(Scalar::bool(true, Nullability::NonNullable)), + ); + + // Reopen layout and check serialized zones metadata. + let file = session.open_options().open_buffer(buf)?; + let zoned = file + .footer() + .layout() + .as_opt::() + .vortex_expect("the test must write zone maps"); + assert_eq!(zoned.nzones(), 3); + + let sort_order_id = aggregate.id(); + assert!( + zoned + .present_aggregates() + .iter() + .all(|name| !name.starts_with(sort_order_id.as_ref())), + "zone metadata must not include any SortOrder aggregate", + ); + + let zones = zoned.slot(1)?.vortex_expect("stored zone statistics"); + assert_eq!(zones.row_count(), 3); + + let fields = zones + .dtype() + .as_struct_fields_opt() + .vortex_expect("zone statistics table is a struct"); + assert!( + !fields.names().is_empty(), + "zone statistics must actually be stored" + ); + assert!( + fields + .names() + .iter() + .all(|name| !name.as_ref().starts_with(sort_order_id.as_ref())), + "zone statistics must not contain any SortOrder partial", + ); + + Ok(()) +} + #[tokio::test] async fn file_sum_is_absent_when_a_chunk_overflows() -> VortexResult<()> { let dtype = DType::Struct( diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 14b43531695..4057bb5166a 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -20,6 +20,7 @@ use vortex_array::ArrayContext; use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::fns::sort_order::SortOrder; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; use vortex_array::iter::ArrayIterator; @@ -282,6 +283,22 @@ impl VortexWriteOptions { validate_dtype_editions(&self.session, &dtype)?; } + // TODO (joacoc): could this validation be moved out of here? + // + // `SortOrder` supports non-nullable structs, but storing the result requires + // a root statistics entry, and the current file-stats layout omits it. + for aggregate in self.file_statistics.iter().flatten() { + if let Some(options) = aggregate.as_opt::() { + options.validate(&dtype)?; + if !dtype.is_nullable() { + vortex_bail!( + "File sort keys for non-nullable struct roots are not supported yet: \ + the statistics layout has no root entry" + ); + } + } + } + let (mut ptr, eof) = SequenceId::root().split(); let stream = SequentialStreamAdapter::new( diff --git a/vortex/editions/unstable/unstable2026.10.0.toml b/vortex/editions/unstable/unstable2026.10.0.toml index 4810ff50be1..cb286f65eba 100644 --- a/vortex/editions/unstable/unstable2026.10.0.toml +++ b/vortex/editions/unstable/unstable2026.10.0.toml @@ -26,6 +26,7 @@ layouts = [ dtypes = [] aggregates = [ "vortex.bloom_filter.sbbf", + "vortex.sort_order", "vortex.sum_v2", ] @@ -46,5 +47,6 @@ layouts = [ dtypes = [] aggregates = [ "vortex.bloom_filter.sbbf", + "vortex.sort_order", "vortex.sum_v2", ] From ad63b14b6f1be55042a0ba6170574bf2cc6f8226 Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Fri, 30 Oct 2026 17:18:53 +0100 Subject: [PATCH 10/10] (rename) `SortOrder` to `IsCompoundSorted` Given that already exists `IsSorted` and `IsStrictSorted` I found appropiate to call it `IsCompoundSorted`. Signed-off-by: Joaquin Colacci --- vortex-array/Cargo.toml | 2 +- .../{sort_order.rs => is_compound_sorted.rs} | 36 ++--- vortex-array/proto/expr.proto | 4 +- .../{sort_order => is_compound_sorted}/mod.rs | 126 +++++++++--------- vortex-array/src/aggregate_fn/fns/mod.rs | 2 +- .../src/proto/generated/vortex.expr.rs | 10 +- .../src/declarations/unstable/mod.rs | 2 +- vortex-file/src/tests.rs | 25 ++-- vortex-file/src/writer.rs | 6 +- .../editions/unstable/unstable2026.10.0.toml | 4 +- 10 files changed, 110 insertions(+), 107 deletions(-) rename vortex-array/benches/{sort_order.rs => is_compound_sorted.rs} (90%) rename vortex-array/src/aggregate_fn/fns/{sort_order => is_compound_sorted}/mod.rs (90%) diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index dc86a4187be..87b55c09f79 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -362,5 +362,5 @@ name = "run_end_decode_primitive" harness = false [[bench]] -name = "sort_order" +name = "is_compound_sorted" harness = false diff --git a/vortex-array/benches/sort_order.rs b/vortex-array/benches/is_compound_sorted.rs similarity index 90% rename from vortex-array/benches/sort_order.rs rename to vortex-array/benches/is_compound_sorted.rs index 5575ba8b162..edebd7ee987 100644 --- a/vortex-array/benches/sort_order.rs +++ b/vortex-array/benches/is_compound_sorted.rs @@ -5,7 +5,7 @@ //! //! Covers order and line-item keys with dense and sparse ties, batch boundaries, //! nullable and UTF-8 fields, and constant and dictionary encodings. Also compares -//! single-field `SortOrder` with `IsSorted` on the same ascending order IDs. +//! single-field `IsCompoundSorted` with `IsSorted` on the same ascending order IDs. #![expect(clippy::unwrap_used)] @@ -20,11 +20,11 @@ use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::Accumulator; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::fns::is_compound_sorted::IsCompoundSorted; +use vortex_array::aggregate_fn::fns::is_compound_sorted::IsCompoundSortedOptions; +use vortex_array::aggregate_fn::fns::is_compound_sorted::SortDirection; use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; use vortex_array::aggregate_fn::fns::is_sorted::IsSortedOptions; -use vortex_array::aggregate_fn::fns::sort_order::SortDirection; -use vortex_array::aggregate_fn::fns::sort_order::SortOrder; -use vortex_array::aggregate_fn::fns::sort_order::SortOrderOptions; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; @@ -44,8 +44,8 @@ fn main() { divan::main(); } -fn order_line_options() -> SortOrderOptions { - SortOrderOptions::new([ +fn order_line_options() -> IsCompoundSortedOptions { + IsCompoundSortedOptions::new([ ("order_id", SortDirection::Asc), ("line_item_id", SortDirection::Asc), ]) @@ -81,7 +81,7 @@ fn measure_batches(bencher: Bencher, batch_rows: usize, make_array: impl Fn() -> .with_inputs(|| { let array = make_array(); let accumulator = - Accumulator::try_new(SortOrder, order_line_options(), array.dtype().clone()) + Accumulator::try_new(IsCompoundSorted, order_line_options(), array.dtype().clone()) .unwrap(); let batches = (0..array.len()) .step_by(batch_rows) @@ -110,7 +110,7 @@ fn order_lines_by_order_size(bencher: Bencher, items_per_order: u32) { let line_item_ids = (0..ROWS) .map(|index| u64::from(index % items_per_order)) .collect::>(); - measure(bencher, SortOrder, order_line_options(), || { + measure(bencher, IsCompoundSorted, order_line_options(), || { StructArray::from_fields(&[ ("order_id", order_ids.clone().into_array()), ("line_item_id", line_item_ids.clone().into_array()), @@ -122,7 +122,7 @@ fn order_lines_by_order_size(bencher: Bencher, items_per_order: u32) { #[divan::bench] fn repeated_order_line_key(bencher: Bencher) { - measure(bencher, SortOrder, order_line_options(), || { + measure(bencher, IsCompoundSorted, order_line_options(), || { StructArray::from_fields(&[ ( "order_id", @@ -151,8 +151,8 @@ fn customer_order_lines_sparse_ties(bencher: Bencher) { let line_item_ids = (0..ROWS).map(u64::from).collect::>(); measure( bencher, - SortOrder, - SortOrderOptions::new([ + IsCompoundSorted, + IsCompoundSortedOptions::new([ ("customer_id", SortDirection::Asc), ("order_id", SortDirection::Asc), ("line_item_id", SortDirection::Asc), @@ -203,7 +203,7 @@ fn nullable_order_line_keys(bencher: Bencher) { } }) .collect::>(); - measure(bencher, SortOrder, order_line_options(), || { + measure(bencher, IsCompoundSorted, order_line_options(), || { StructArray::from_fields(&[ ( "order_id", @@ -230,8 +230,8 @@ fn customer_order_string_keys(bencher: Bencher) { .collect::>(); measure( bencher, - SortOrder, - SortOrderOptions::new([ + IsCompoundSorted, + IsCompoundSortedOptions::new([ ("customer_id", SortDirection::Asc), ("order_id", SortDirection::Asc), ]), @@ -259,7 +259,7 @@ fn order_lines_with_dictionary_item_ids(bencher: Bencher) { .collect::>(); let line_item_codes = (0..ROWS).map(|index| index % 16).collect::>(); let line_item_dictionary = (0..16u64).collect::>(); - measure(bencher, SortOrder, order_line_options(), || { + measure(bencher, IsCompoundSorted, order_line_options(), || { StructArray::from_fields(&[ ("order_id", order_ids.clone().into_array()), ( @@ -288,12 +288,12 @@ mod single_order_id { } #[divan::bench] - fn sort_order(bencher: Bencher) { + fn is_compound_sorted(bencher: Bencher) { let order_ids = order_ids(); measure( bencher, - SortOrder, - SortOrderOptions::new([("order_id", SortDirection::Asc)]), + IsCompoundSorted, + IsCompoundSortedOptions::new([("order_id", SortDirection::Asc)]), || { StructArray::from_fields(&[("order_id", order_ids.clone().into_array())]) .unwrap() diff --git a/vortex-array/proto/expr.proto b/vortex-array/proto/expr.proto index 489b6323d4a..82ef80257be 100644 --- a/vortex-array/proto/expr.proto +++ b/vortex-array/proto/expr.proto @@ -32,8 +32,8 @@ message NumericalAggregateOpts { bool skip_nans = 1; } -// A declared lexicographic ordering, checked by `vortex.sort_order`. -message SortOrderOpts { +// A declared lexicographic ordering, checked by `vortex.is_compound_sorted`. +message IsCompoundSortedOpts { repeated SortField fields = 1; message SortField { diff --git a/vortex-array/src/aggregate_fn/fns/sort_order/mod.rs b/vortex-array/src/aggregate_fn/fns/is_compound_sorted/mod.rs similarity index 90% rename from vortex-array/src/aggregate_fn/fns/sort_order/mod.rs rename to vortex-array/src/aggregate_fn/fns/is_compound_sorted/mod.rs index 0b2bf416257..36a01e12301 100644 --- a/vortex-array/src/aggregate_fn/fns/sort_order/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_compound_sorted/mod.rs @@ -98,15 +98,15 @@ pub enum SortDirection { Desc, } -/// Options for [`SortOrder`]. +/// Options for [`IsCompoundSorted`]. #[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub struct SortOrderOptions { +pub struct IsCompoundSortedOptions { /// Top-level struct fields and their directions, in priority order. /// Currently only ascending directions are supported. pub fields: Vec<(FieldName, SortDirection)>, } -impl SortOrderOptions { +impl IsCompoundSortedOptions { /// Declare a lexicographic ordering over the given fields. pub fn new(fields: impl IntoIterator, SortDirection)>) -> Self { Self { @@ -129,7 +129,7 @@ impl SortOrderOptions { for (name, direction) in &self.fields { vortex_ensure!( *direction == SortDirection::Asc, - "SortOrder currently supports ASC only" + "IsCompoundSorted currently supports ASC only" ); vortex_ensure!(seen.insert(name), "Duplicate sort field: {}", name); } @@ -145,7 +145,7 @@ impl SortOrderOptions { let fields = dtype .as_struct_fields_opt() - .ok_or_else(|| vortex_err!("SortOrder requires a struct input"))?; + .ok_or_else(|| vortex_err!("IsCompoundSorted requires a struct input"))?; for (name, _) in &self.fields { let dtype = fields.field(name).ok_or_else(|| { @@ -225,7 +225,7 @@ impl SortOrderOptions { } } -impl Display for SortOrderOptions { +impl Display for IsCompoundSortedOptions { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { for (index, (name, direction)) in self.fields.iter().enumerate() { if index > 0 { @@ -249,7 +249,7 @@ impl Display for SortOrderOptions { /// `sort_fields` are the key columns in priority order, all the same length (at least /// 2 rows), and there must be at least one. Only rows *within* the batch are compared. /// The seam with the previous batch is checked separately by -/// [`SortOrderOptions::compare_sort_keys`]. +/// [`IsCompoundSortedOptions::compare_sort_keys`]. /// /// # How it works /// @@ -347,7 +347,7 @@ fn check_chunk_sortedness(sort_fields: &[ArrayRef], ctx: &mut ExecutionCtx) -> V // Otherwise, this field only constrains pairs tied on all earlier fields. if !all_pairs_tied { - // Given that `SortOrder` uses only `ASC` and `NULLS` first, + // Given that `IsCompoundSorted` uses only `ASC` and `NULLS` first, // a valid value followed by null is also a violation. let null_violation = &previous_valid & &!&next_valid; let out_of_order = previous @@ -388,17 +388,17 @@ fn check_chunk_sortedness(sort_fields: &[ArrayRef], ctx: &mut ExecutionCtx) -> V /// It is not `strict`, so equal row keys are allowed. Only `ASC` is supported for now, /// DESC remains reserved in the options. #[derive(Clone, Debug)] -pub struct SortOrder; +pub struct IsCompoundSorted; -/// Running verdict and boundary keys for [`SortOrder`]. -pub struct SortOrderPartial { +/// Running verdict and boundary keys for [`IsCompoundSorted`]. +pub struct IsCompoundSortedPartial { sorted: bool, // None means no rows. A row whose key fields are null still has a non-null key tuple. first: Option, last: Option, } -impl SortOrderPartial { +impl IsCompoundSortedPartial { /// The state of a group with no accumulated values. fn empty() -> Self { Self { @@ -412,28 +412,28 @@ impl SortOrderPartial { static NAMES: LazyLock = LazyLock::new(|| FieldNames::from(["is_sorted", "first", "last"])); -impl AggregateFnVTable for SortOrder { - type Options = SortOrderOptions; - type Partial = SortOrderPartial; +impl AggregateFnVTable for IsCompoundSorted { + type Options = IsCompoundSortedOptions; + type Partial = IsCompoundSortedPartial; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.sort_order"); + static ID: CachedId = CachedId::new("vortex.is_compound_sorted"); *ID } - /// Serializes [`SortOrderOptions`] as protobuf. + /// Serializes [`IsCompoundSortedOptions`] as protobuf. fn serialize(&self, options: &Self::Options) -> VortexResult>> { options.validate_fields()?; Ok(Some( - pb::SortOrderOpts { + pb::IsCompoundSortedOpts { fields: options .fields .iter() - .map(|(name, direction)| pb::sort_order_opts::SortField { + .map(|(name, direction)| pb::is_compound_sorted_opts::SortField { name: name.to_string(), direction: match direction { - SortDirection::Asc => pb::sort_order_opts::SortDirection::Asc, - SortDirection::Desc => pb::sort_order_opts::SortDirection::Desc, + SortDirection::Asc => pb::is_compound_sorted_opts::SortDirection::Asc, + SortDirection::Desc => pb::is_compound_sorted_opts::SortDirection::Desc, } as i32, }) .collect(), @@ -442,26 +442,26 @@ impl AggregateFnVTable for SortOrder { )) } - /// Deserializes [`SortOrderOptions`] from protobuf. + /// Deserializes [`IsCompoundSortedOptions`] from protobuf. fn deserialize( &self, metadata: &[u8], _session: &VortexSession, ) -> VortexResult { - let proto = pb::SortOrderOpts::decode(metadata)?; + let proto = pb::IsCompoundSortedOpts::decode(metadata)?; let fields = proto .fields .into_iter() .map(|field| { - let direction = match pb::sort_order_opts::SortDirection::try_from(field.direction)? - { - pb::sort_order_opts::SortDirection::Asc => SortDirection::Asc, - pb::sort_order_opts::SortDirection::Desc => SortDirection::Desc, - }; + let direction = + match pb::is_compound_sorted_opts::SortDirection::try_from(field.direction)? { + pb::is_compound_sorted_opts::SortDirection::Asc => SortDirection::Asc, + pb::is_compound_sorted_opts::SortDirection::Desc => SortDirection::Desc, + }; Ok((field.name, direction)) }) .collect::>>()?; - let options = SortOrderOptions::new(fields); + let options = IsCompoundSortedOptions::new(fields); options.validate_fields()?; Ok(options) } @@ -487,7 +487,7 @@ impl AggregateFnVTable for SortOrder { &self, _args: AggregateArgs<'_, Self::Options>, ) -> VortexResult { - Ok(SortOrderPartial::empty()) + Ok(IsCompoundSortedPartial::empty()) } fn partial_from_scalar( @@ -497,7 +497,7 @@ impl AggregateFnVTable for SortOrder { ) -> VortexResult { // A null struct means the producing accumulator was empty. if scalar.is_null() { - return Ok(SortOrderPartial::empty()); + return Ok(IsCompoundSortedPartial::empty()); } let fields = scalar.as_struct(); @@ -509,7 +509,7 @@ impl AggregateFnVTable for SortOrder { "Non-empty sortedness partials require boundary keys" ); - Ok(SortOrderPartial { + Ok(IsCompoundSortedPartial { sorted: fields .field_by_idx(0) .vortex_expect("sorted verdict") @@ -610,7 +610,7 @@ impl AggregateFnVTable for SortOrder { .partial_dtype .as_struct_fields_opt() .and_then(|fields| fields.field_by_index(1)) - .vortex_expect("SortOrder partial contains the first key dtype"); + .vortex_expect("IsCompoundSorted partial contains the first key dtype"); let first = Scalar::struct_( key_dtype.clone(), @@ -657,7 +657,7 @@ impl AggregateFnVTable for SortOrder { _batch: &Columnar, _ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - unreachable!("SortOrder::try_accumulate handles all arrays") + unreachable!("IsCompoundSorted::try_accumulate handles all arrays") } fn finalize( @@ -693,9 +693,9 @@ mod tests { use vortex_error::vortex_bail; use vortex_session::VortexSession; + use super::IsCompoundSorted; + use super::IsCompoundSortedOptions; use super::SortDirection; - use super::SortOrder; - use super::SortOrderOptions; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; @@ -738,8 +738,8 @@ mod tests { .into_array()) } - fn order_line_options() -> SortOrderOptions { - SortOrderOptions::new([ + fn order_line_options() -> IsCompoundSortedOptions { + IsCompoundSortedOptions::new([ ("order_id", SortDirection::Asc), ("line_item_id", SortDirection::Asc), ]) @@ -791,14 +791,15 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let dtype = order_lines(&[])?.dtype().clone(); let options = order_line_options(); - let mut accumulator = Accumulator::try_new(SortOrder, options.clone(), dtype.clone())?; + let mut accumulator = + Accumulator::try_new(IsCompoundSorted, options.clone(), dtype.clone())?; for chunk in chunks { let array = order_lines(&chunk)?; if combine_partials { let mut batch_accumulator = - Accumulator::try_new(SortOrder, options.clone(), dtype.clone())?; + Accumulator::try_new(IsCompoundSorted, options.clone(), dtype.clone())?; batch_accumulator.accumulate(&array, &mut ctx)?; let partial = batch_accumulator.partial_scalar()?; @@ -835,7 +836,7 @@ mod tests { ) -> VortexResult<()> { check_batches( order_lines(&values)?, - SortOrderOptions::new([ + IsCompoundSortedOptions::new([ ("line_item_id", SortDirection::Asc), ("order_id", SortDirection::Asc), ]), @@ -862,7 +863,7 @@ mod tests { check_batches( array, - SortOrderOptions::new([("order_id", SortDirection::Asc)]), + IsCompoundSortedOptions::new([("order_id", SortDirection::Asc)]), one_row_batches, expected, ) @@ -870,12 +871,13 @@ mod tests { fn check_batches( array: ArrayRef, - options: SortOrderOptions, + options: IsCompoundSortedOptions, one_row_batches: bool, expected: bool, ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let mut accumulator = Accumulator::try_new(SortOrder, options, array.dtype().clone())?; + let mut accumulator = + Accumulator::try_new(IsCompoundSorted, options, array.dtype().clone())?; if one_row_batches { for index in 0..array.len() { @@ -928,7 +930,7 @@ mod tests { check_batches( array, - SortOrderOptions::new([("price", SortDirection::Asc)]), + IsCompoundSortedOptions::new([("price", SortDirection::Asc)]), one_row_batches, expected, ) @@ -963,7 +965,7 @@ mod tests { check_batches( array, - SortOrderOptions::new([ + IsCompoundSortedOptions::new([ ("customer_id", SortDirection::Asc), ("order_id", SortDirection::Asc), ("line_item_id", SortDirection::Asc), @@ -989,7 +991,7 @@ mod tests { check_batches( array, - SortOrderOptions::new([ + IsCompoundSortedOptions::new([ ("price", SortDirection::Asc), ("line_item_id", SortDirection::Asc), ]), @@ -1007,7 +1009,7 @@ mod tests { check_batches( array, - SortOrderOptions::new([("order_id", SortDirection::Asc)]), + IsCompoundSortedOptions::new([("order_id", SortDirection::Asc)]), false, expected, ) @@ -1016,15 +1018,15 @@ mod tests { #[test] fn options_roundtrip() -> VortexResult<()> { let options = order_line_options(); - let bytes = SortOrder + let bytes = IsCompoundSorted .serialize(&options)? .vortex_expect("serializable options"); - assert_eq!(SortOrder.deserialize(&bytes, &SESSION)?, options); + assert_eq!(IsCompoundSorted.deserialize(&bytes, &SESSION)?, options); assert_eq!(options.to_string(), "order_id ASC, line_item_id ASC"); assert!( - SortOrder + IsCompoundSorted .zone_stat_default(&DType::from(crate::dtype::PType::I32)) .is_none() ); @@ -1039,10 +1041,10 @@ mod tests { #[case::descending_primary(vec![("order_id", 1), ("line_item_id", 0)])] #[case::descending_secondary(vec![("order_id", 0), ("line_item_id", 1)])] fn invalid_serialized_options_are_rejected(#[case] fields: Vec<(&str, i32)>) { - let bytes = pb::SortOrderOpts { + let bytes = pb::IsCompoundSortedOpts { fields: fields .into_iter() - .map(|(name, direction)| pb::sort_order_opts::SortField { + .map(|(name, direction)| pb::is_compound_sorted_opts::SortField { name: name.into(), direction, }) @@ -1050,7 +1052,7 @@ mod tests { } .encode_to_vec(); - assert!(SortOrder.deserialize(&bytes, &SESSION).is_err()); + assert!(IsCompoundSorted.deserialize(&bytes, &SESSION).is_err()); } #[rstest] @@ -1060,17 +1062,17 @@ mod tests { fn invalid_keys_are_rejected(#[case] fields: Vec<&str>) -> VortexResult<()> { let dtype = order_lines(&[])?.dtype().clone(); let options = - SortOrderOptions::new(fields.into_iter().map(|name| (name, SortDirection::Asc))); + IsCompoundSortedOptions::new(fields.into_iter().map(|name| (name, SortDirection::Asc))); assert!(options.validate(&dtype).is_err()); - assert!(Accumulator::try_new(SortOrder, options, dtype).is_err()); + assert!(Accumulator::try_new(IsCompoundSorted, options, dtype).is_err()); Ok(()) } #[test] fn unsupported_input_and_field_are_rejected() -> VortexResult<()> { - let options = SortOrderOptions::new([("order_id", SortDirection::Asc)]); + let options = IsCompoundSortedOptions::new([("order_id", SortDirection::Asc)]); let primitive_input = buffer![1i32].into_array(); let nested_input = StructArray::from_fields(&[("order_id", order_lines(&[(1, 2)])?)])?.into_array(); @@ -1084,7 +1086,7 @@ mod tests { #[test] fn descending_is_reserved_but_rejected() -> VortexResult<()> { let dtype = order_lines(&[])?.dtype().clone(); - let options = SortOrderOptions::new([ + let options = IsCompoundSortedOptions::new([ ("order_id", SortDirection::Asc), ("line_item_id", SortDirection::Desc), ]); @@ -1094,8 +1096,8 @@ mod tests { assert_eq!(options.to_string(), "order_id ASC, line_item_id DESC"); assert!(error.to_string().contains("ASC only")); - assert!(Accumulator::try_new(SortOrder, options.clone(), dtype).is_err()); - assert!(SortOrder.serialize(&options).is_err()); + assert!(Accumulator::try_new(IsCompoundSorted, options.clone(), dtype).is_err()); + assert!(IsCompoundSorted.serialize(&options).is_err()); Ok(()) } @@ -1116,7 +1118,7 @@ mod tests { check_batches( array, - SortOrderOptions::new([ + IsCompoundSortedOptions::new([ ("price", SortDirection::Asc), ("line_item_id", SortDirection::Asc), ]), @@ -1198,7 +1200,7 @@ mod tests { check_batches( array, - SortOrderOptions::new([ + IsCompoundSortedOptions::new([ ("order_id", SortDirection::Asc), ("item_value", SortDirection::Asc), ]), diff --git a/vortex-array/src/aggregate_fn/fns/mod.rs b/vortex-array/src/aggregate_fn/fns/mod.rs index e63e9bebe7a..26150843920 100644 --- a/vortex-array/src/aggregate_fn/fns/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mod.rs @@ -11,6 +11,7 @@ pub mod bounded_min; pub mod count; pub mod first; mod fixed_size_bytes; +pub mod is_compound_sorted; pub mod is_constant; pub mod is_sorted; pub mod last; @@ -20,7 +21,6 @@ pub mod min; pub mod min_max; pub mod nan_count; pub mod null_count; -pub mod sort_order; pub mod sum; pub mod sum_v2; pub mod uncompressed_size_in_bytes; diff --git a/vortex-array/src/proto/generated/vortex.expr.rs b/vortex-array/src/proto/generated/vortex.expr.rs index 1588d1e8982..28238d89cc6 100644 --- a/vortex-array/src/proto/generated/vortex.expr.rs +++ b/vortex-array/src/proto/generated/vortex.expr.rs @@ -31,14 +31,14 @@ pub struct NumericalAggregateOpts { #[prost(bool, tag = "1")] pub skip_nans: bool, } -/// A declared lexicographic ordering, checked by `vortex.sort_order`. +/// A declared lexicographic ordering, checked by `vortex.is_compound_sorted`. #[derive(Clone, PartialEq, ::prost::Message)] -pub struct SortOrderOpts { +pub struct IsCompoundSortedOpts { #[prost(message, repeated, tag = "1")] - pub fields: ::prost::alloc::vec::Vec, + pub fields: ::prost::alloc::vec::Vec, } -/// Nested message and enum types in `SortOrderOpts`. -pub mod sort_order_opts { +/// Nested message and enum types in `IsCompoundSortedOpts`. +pub mod is_compound_sorted_opts { #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct SortField { #[prost(string, tag = "1")] diff --git a/vortex-edition/src/declarations/unstable/mod.rs b/vortex-edition/src/declarations/unstable/mod.rs index 2f015ce2cc1..93b3c4e3749 100644 --- a/vortex-edition/src/declarations/unstable/mod.rs +++ b/vortex-edition/src/declarations/unstable/mod.rs @@ -40,7 +40,7 @@ pub static DECLARATION: EditionDeclaration = EditionDeclaration { EditionMember::array(&"vortex.union"), EditionMember::layout(&"vortex.list"), EditionMember::aggregate(&"vortex.bloom_filter.sbbf"), - EditionMember::aggregate(&"vortex.sort_order"), + EditionMember::aggregate(&"vortex.is_compound_sorted"), EditionMember::aggregate(&"vortex.sum_v2"), ], }; diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 01ff961f17d..a9c76505443 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -24,9 +24,9 @@ use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::sort_order::SortDirection; -use vortex_array::aggregate_fn::fns::sort_order::SortOrder; -use vortex_array::aggregate_fn::fns::sort_order::SortOrderOptions; +use vortex_array::aggregate_fn::fns::is_compound_sorted::IsCompoundSorted; +use vortex_array::aggregate_fn::fns::is_compound_sorted::IsCompoundSortedOptions; +use vortex_array::aggregate_fn::fns::is_compound_sorted::SortDirection; use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::arrays::BoolArray; @@ -2361,18 +2361,18 @@ async fn test_writer_with_statistics() -> VortexResult<()> { Ok(()) } -/// Assert that `SortOrder` is written in the +/// Assert that `IsCompoundSorted` is written in the /// file stats, and also check that nothing was written /// as zone stats. #[tokio::test] -async fn file_footer_contains_sort_order() -> VortexResult<()> { +async fn file_footer_contains_is_compound_sorted() -> VortexResult<()> { let session = VortexSession::empty() .with::() .with::() .with::(); crate::register_default_encodings(&session); - session.aggregate_fns().register(SortOrder); + session.aggregate_fns().register(IsCompoundSorted); for declaration in EDITION_DECLARATIONS { session.register_edition(declaration)?; } @@ -2387,7 +2387,8 @@ async fn file_footer_contains_sort_order() -> VortexResult<()> { ) .into_array(); - let aggregate = SortOrder.bind(SortOrderOptions::new([("numbers", SortDirection::Asc)])); + let aggregate = + IsCompoundSorted.bind(IsCompoundSortedOptions::new([("numbers", SortDirection::Asc)])); // The default strategy skips zone maps beneath nullable structs. Write zones at the // root so the assertions below inspect stored zone statistics alongside file statistics. @@ -2435,13 +2436,13 @@ async fn file_footer_contains_sort_order() -> VortexResult<()> { .vortex_expect("the test must write zone maps"); assert_eq!(zoned.nzones(), 3); - let sort_order_id = aggregate.id(); + let is_compound_sorted_id = aggregate.id(); assert!( zoned .present_aggregates() .iter() - .all(|name| !name.starts_with(sort_order_id.as_ref())), - "zone metadata must not include any SortOrder aggregate", + .all(|name| !name.starts_with(is_compound_sorted_id.as_ref())), + "zone metadata must not include any IsCompoundSorted aggregate", ); let zones = zoned.slot(1)?.vortex_expect("stored zone statistics"); @@ -2459,8 +2460,8 @@ async fn file_footer_contains_sort_order() -> VortexResult<()> { fields .names() .iter() - .all(|name| !name.as_ref().starts_with(sort_order_id.as_ref())), - "zone statistics must not contain any SortOrder partial", + .all(|name| !name.as_ref().starts_with(is_compound_sorted_id.as_ref())), + "zone statistics must not contain any IsCompoundSorted partial", ); Ok(()) diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 4057bb5166a..05c0d4ab742 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -20,7 +20,7 @@ use vortex_array::ArrayContext; use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::sort_order::SortOrder; +use vortex_array::aggregate_fn::fns::is_compound_sorted::IsCompoundSorted; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; use vortex_array::iter::ArrayIterator; @@ -285,10 +285,10 @@ impl VortexWriteOptions { // TODO (joacoc): could this validation be moved out of here? // - // `SortOrder` supports non-nullable structs, but storing the result requires + // `IsCompoundSorted` supports non-nullable structs, but storing the result requires // a root statistics entry, and the current file-stats layout omits it. for aggregate in self.file_statistics.iter().flatten() { - if let Some(options) = aggregate.as_opt::() { + if let Some(options) = aggregate.as_opt::() { options.validate(&dtype)?; if !dtype.is_nullable() { vortex_bail!( diff --git a/vortex/editions/unstable/unstable2026.10.0.toml b/vortex/editions/unstable/unstable2026.10.0.toml index cb286f65eba..55a5f462429 100644 --- a/vortex/editions/unstable/unstable2026.10.0.toml +++ b/vortex/editions/unstable/unstable2026.10.0.toml @@ -26,7 +26,7 @@ layouts = [ dtypes = [] aggregates = [ "vortex.bloom_filter.sbbf", - "vortex.sort_order", + "vortex.is_compound_sorted", "vortex.sum_v2", ] @@ -47,6 +47,6 @@ layouts = [ dtypes = [] aggregates = [ "vortex.bloom_filter.sbbf", - "vortex.sort_order", + "vortex.is_compound_sorted", "vortex.sum_v2", ]