1use super::*;
18
19pub struct LimitedResult {
21 inner: Box<dyn QueryResult>,
23 limit: Option<usize>,
25 offset: usize,
27 returned_count: usize,
29 offset_applied: bool,
31 columns: Vec<String>,
33}
34
35impl LimitedResult {
36 pub fn new(inner: Box<dyn QueryResult>, limit: Option<usize>, offset: usize) -> Self {
38 let columns = inner.columns().to_vec();
39 Self {
40 inner,
41 limit,
42 offset,
43 returned_count: 0,
44 offset_applied: false,
45 columns,
46 }
47 }
48
49 pub fn with_limit(inner: Box<dyn QueryResult>, limit: usize) -> Self {
51 Self::new(inner, Some(limit), 0)
52 }
53
54 pub fn with_offset(inner: Box<dyn QueryResult>, offset: usize) -> Self {
56 Self::new(inner, None, offset)
57 }
58}
59
60impl QueryResult for LimitedResult {
61 fn columns(&self) -> &[String] {
62 &self.columns
63 }
64
65 fn columns_arc(&self) -> Option<CompactArc<Vec<String>>> {
66 self.inner.columns_arc()
67 }
68
69 fn next(&mut self) -> bool {
70 if !self.offset_applied {
72 for _ in 0..self.offset {
73 if !self.inner.next() {
74 self.offset_applied = true;
75 return false;
76 }
77 }
78 self.offset_applied = true;
79 }
80
81 if let Some(limit) = self.limit {
83 if self.returned_count >= limit {
84 return false;
85 }
86 }
87
88 if self.inner.next() {
90 self.returned_count += 1;
91 true
92 } else {
93 false
94 }
95 }
96
97 fn scan(&self, dest: &mut [Value]) -> Result<()> {
98 self.inner.scan(dest)
99 }
100
101 fn row(&self) -> &Row {
102 self.inner.row()
103 }
104
105 fn take_row(&mut self) -> Row {
106 self.inner.take_row()
107 }
108
109 fn take_deferred_row(&mut self) -> DeferredRow {
110 self.inner.take_deferred_row()
111 }
112
113 fn preserves_deferred_rows(&self) -> bool {
114 self.inner.preserves_deferred_rows()
115 }
116
117 fn ascending_nulls_last_ordering(&self) -> Option<Vec<usize>> {
118 self.inner.ascending_nulls_last_ordering()
119 }
120
121 fn close(&mut self) -> Result<()> {
122 self.inner.close()
123 }
124
125 fn rows_affected(&self) -> i64 {
126 self.inner.rows_affected()
127 }
128
129 fn last_insert_id(&self) -> i64 {
130 self.inner.last_insert_id()
131 }
132
133 fn last_error(&mut self) -> Option<radixdb_core::Error> {
134 self.inner.last_error()
135 }
136
137 fn estimated_count(&self) -> Option<usize> {
138 self.inner.estimated_count().map(|rows| {
139 let after_offset = if self.offset_applied {
140 rows
141 } else {
142 rows.saturating_sub(self.offset)
143 };
144 self.limit.map_or(after_offset, |limit| {
145 after_offset.min(limit.saturating_sub(self.returned_count))
146 })
147 })
148 }
149
150 fn with_aliases(self: Box<Self>, aliases: FxHashMap<String, String>) -> Box<dyn QueryResult> {
151 Box::new(AliasedResult::new(self, aliases))
152 }
153}
154
155pub struct OrderedResult {
157 inner: Box<dyn QueryResult>,
159}
160
161pub(super) const ORDERED_RUN_MAX_ROWS: usize = 131_072;
165pub(super) const ORDERED_RUN_MAX_BYTES: usize = 64 * 1024 * 1024;
166const ORDERED_SPILL_MAX_VALUE_BYTES: usize = 64 * 1024 * 1024;
167static ORDERED_SPILL_ID: AtomicU64 = AtomicU64::new(1);
168type OrderedRowComparator = dyn Fn(&Row, &Row) -> std::cmp::Ordering + Send;
169
170enum BoundedOrderedRows {
171 Memory {
172 rows: RowVec,
173 input_rows: usize,
174 peak_rows: usize,
175 peak_bytes: usize,
176 },
177 External {
178 paths: Vec<PathBuf>,
179 input_rows: usize,
180 peak_rows: usize,
181 peak_bytes: usize,
182 },
183}
184
185struct OrderedSpillRun {
186 path: PathBuf,
187 reader: BufReader<File>,
188 remaining: u64,
189 head: Option<Row>,
190}
191
192impl OrderedSpillRun {
193 fn open(path: PathBuf) -> Result<Self> {
194 let opened = (|| {
195 let file = File::open(&path).map_err(|error| {
196 Error::internal(format!(
197 "failed to open ORDER BY spill run {}: {error}",
198 path.display()
199 ))
200 })?;
201 let mut reader = BufReader::new(file);
202 let remaining = read_u64(&mut reader, "ORDER BY spill row count")?;
203 let mut run = Self {
204 path: path.clone(),
205 reader,
206 remaining,
207 head: None,
208 };
209 run.advance()?;
210 Ok(run)
211 })();
212 if opened.is_err() {
213 let _ = std::fs::remove_file(path);
214 }
215 opened
216 }
217
218 fn advance(&mut self) -> Result<()> {
219 self.head = if self.remaining == 0 {
220 None
221 } else {
222 self.remaining -= 1;
223 Some(read_spill_row(&mut self.reader)?)
224 };
225 Ok(())
226 }
227}
228
229impl Drop for OrderedSpillRun {
230 fn drop(&mut self) {
231 let _ = std::fs::remove_file(&self.path);
232 }
233}
234
235struct ExternalOrderedResult {
236 columns: CompactArc<Vec<String>>,
237 runs: Vec<OrderedSpillRun>,
238 compare: Box<OrderedRowComparator>,
239 current: Option<Row>,
240 remaining: usize,
241 closed: bool,
242 last_error: Option<Error>,
243}
244
245impl ExternalOrderedResult {
246 fn new<F>(
247 columns: Vec<String>,
248 paths: Vec<PathBuf>,
249 input_rows: usize,
250 compare: F,
251 ) -> Result<Self>
252 where
253 F: Fn(&Row, &Row) -> std::cmp::Ordering + Send + 'static,
254 {
255 let mut pending = paths.into_iter();
256 let mut runs = Vec::new();
257 while let Some(path) = pending.next() {
258 match OrderedSpillRun::open(path) {
259 Ok(run) => runs.push(run),
260 Err(error) => {
261 for path in pending {
262 let _ = std::fs::remove_file(path);
263 }
264 return Err(error);
265 }
266 }
267 }
268 Ok(Self {
269 columns: CompactArc::new(columns),
270 runs,
271 compare: Box::new(compare),
272 current: None,
273 remaining: input_rows,
274 closed: false,
275 last_error: None,
276 })
277 }
278}
279
280impl QueryResult for ExternalOrderedResult {
281 fn columns(&self) -> &[String] {
282 &self.columns
283 }
284
285 fn columns_arc(&self) -> Option<CompactArc<Vec<String>>> {
286 Some(CompactArc::clone(&self.columns))
287 }
288
289 fn next(&mut self) -> bool {
290 if self.closed || self.last_error.is_some() {
291 return false;
292 }
293 let mut best: Option<usize> = None;
294 for (index, run) in self.runs.iter().enumerate() {
295 let Some(candidate) = run.head.as_ref() else {
296 continue;
297 };
298 if best.is_none_or(|best_index| {
299 let best_row = self.runs[best_index]
300 .head
301 .as_ref()
302 .expect("selected ORDER BY run has a head row");
303 (self.compare)(candidate, best_row).is_lt()
304 }) {
305 best = Some(index);
306 }
307 }
308 let Some(best) = best else {
309 self.current = None;
310 return false;
311 };
312 self.current = self.runs[best].head.take();
313 if let Err(error) = self.runs[best].advance() {
314 self.last_error = Some(error);
315 self.current = None;
316 return false;
317 }
318 self.remaining = self.remaining.saturating_sub(1);
319 true
320 }
321
322 fn scan(&self, dest: &mut [Value]) -> Result<()> {
323 let row = self.row();
324 if dest.len() != row.len() {
325 return Err(Error::internal(format!(
326 "scan destination has {} values for ORDER BY row with {} columns",
327 dest.len(),
328 row.len()
329 )));
330 }
331 dest.clone_from_slice(row.as_slice());
332 Ok(())
333 }
334
335 fn row(&self) -> &Row {
336 self.current
337 .as_ref()
338 .expect("row() called without successful external ORDER BY next()")
339 }
340
341 fn take_row(&mut self) -> Row {
342 self.current
343 .take()
344 .expect("take_row() called without successful external ORDER BY next()")
345 }
346
347 fn close(&mut self) -> Result<()> {
348 self.closed = true;
349 self.current = None;
350 self.runs.clear();
351 Ok(())
352 }
353
354 fn rows_affected(&self) -> i64 {
355 0
356 }
357
358 fn last_insert_id(&self) -> i64 {
359 0
360 }
361
362 fn last_error(&mut self) -> Option<Error> {
363 self.last_error.take()
364 }
365
366 fn estimated_count(&self) -> Option<usize> {
367 Some(self.remaining)
368 }
369
370 fn with_aliases(self: Box<Self>, aliases: FxHashMap<String, String>) -> Box<dyn QueryResult> {
371 Box::new(AliasedResult::new(self, aliases))
372 }
373}
374
375#[derive(Clone, Copy)]
377pub struct RadixOrderSpec {
378 pub col_idx: usize,
380 pub ascending: bool,
382 pub nulls_first: Option<bool>,
387}
388
389fn ordered_spill_path() -> Result<PathBuf> {
390 let base = std::env::temp_dir();
391 for _ in 0..32 {
392 let ordinal = ORDERED_SPILL_ID.fetch_add(1, AtomicOrdering::Relaxed);
393 let path = base.join(format!(
394 "radixdb-order-{}-{ordinal}.run",
395 std::process::id()
396 ));
397 match OpenOptions::new().write(true).create_new(true).open(&path) {
398 Ok(file) => {
399 drop(file);
400 return Ok(path);
401 }
402 Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => continue,
403 Err(error) => {
404 return Err(Error::internal(format!(
405 "failed to create ORDER BY spill path {}: {error}",
406 path.display()
407 )))
408 }
409 }
410 }
411 Err(Error::internal(
412 "failed to allocate a unique ORDER BY spill path",
413 ))
414}
415
416fn write_u32(writer: &mut impl Write, value: usize, what: &str) -> Result<()> {
417 let value =
418 u32::try_from(value).map_err(|_| Error::invalid_argument(format!("{what} exceeds u32")))?;
419 writer
420 .write_all(&value.to_le_bytes())
421 .map_err(|error| Error::internal(format!("failed to write {what}: {error}")))
422}
423
424fn write_spill_row(writer: &mut impl Write, row: &Row) -> Result<()> {
425 write_u32(writer, row.len(), "ORDER BY spill column count")?;
426 let mut encoded = Vec::new();
427 for value in row.iter() {
428 encoded.clear();
429 radixdb_storage::mvcc::persistence::serialize_value_into(&mut encoded, value)?;
430 write_u32(writer, encoded.len(), "ORDER BY spill value length")?;
431 writer.write_all(&encoded).map_err(|error| {
432 Error::internal(format!("failed to write ORDER BY spill value: {error}"))
433 })?;
434 }
435 Ok(())
436}
437
438fn read_u32(reader: &mut impl Read, what: &str) -> Result<u32> {
439 let mut bytes = [0u8; 4];
440 reader
441 .read_exact(&mut bytes)
442 .map_err(|error| Error::internal(format!("failed to read {what}: {error}")))?;
443 Ok(u32::from_le_bytes(bytes))
444}
445
446fn read_u64(reader: &mut impl Read, what: &str) -> Result<u64> {
447 let mut bytes = [0u8; 8];
448 reader
449 .read_exact(&mut bytes)
450 .map_err(|error| Error::internal(format!("failed to read {what}: {error}")))?;
451 Ok(u64::from_le_bytes(bytes))
452}
453
454fn read_spill_row(reader: &mut impl Read) -> Result<Row> {
455 let columns = read_u32(reader, "ORDER BY spill column count")? as usize;
456 if columns > RetainedRowsBudget::DEFAULT_MAX_ROWS {
457 return Err(Error::internal(format!(
458 "ORDER BY spill row declares {columns} columns"
459 )));
460 }
461 let mut row = Row::with_capacity(columns);
462 for _ in 0..columns {
463 let value_len = read_u32(reader, "ORDER BY spill value length")? as usize;
464 if value_len > ORDERED_SPILL_MAX_VALUE_BYTES {
465 return Err(Error::internal(format!(
466 "ORDER BY spill value length {value_len} exceeds bounded decoder limit"
467 )));
468 }
469 let mut encoded = vec![0u8; value_len];
470 reader.read_exact(&mut encoded).map_err(|error| {
471 Error::internal(format!("failed to read ORDER BY spill value: {error}"))
472 })?;
473 row.push(radixdb_storage::mvcc::persistence::deserialize_value(
474 &encoded,
475 )?);
476 }
477 Ok(row)
478}
479
480fn write_ordered_run<F>(rows: &mut RowVec, compare: &F) -> Result<PathBuf>
481where
482 F: Fn(&Row, &Row) -> std::cmp::Ordering,
483{
484 rows.sort_unstable_by(|(_, left), (_, right)| compare(left, right));
485 let path = ordered_spill_path()?;
486 let result = (|| {
487 let file = OpenOptions::new()
488 .write(true)
489 .truncate(true)
490 .open(&path)
491 .map_err(|error| {
492 Error::internal(format!(
493 "failed to open ORDER BY spill run {}: {error}",
494 path.display()
495 ))
496 })?;
497 let mut writer = BufWriter::new(file);
498 writer
499 .write_all(&(rows.len() as u64).to_le_bytes())
500 .map_err(|error| Error::internal(format!("failed to write spill header: {error}")))?;
501 for (_, row) in rows.iter() {
502 write_spill_row(&mut writer, row)?;
503 }
504 writer
505 .flush()
506 .map_err(|error| Error::internal(format!("failed to flush ORDER BY spill: {error}")))
507 })();
508 if result.is_err() {
509 let _ = std::fs::remove_file(&path);
510 }
511 result.map(|()| path)
512}
513
514fn remove_ordered_runs(paths: impl IntoIterator<Item = PathBuf>) {
515 for path in paths {
516 let _ = std::fs::remove_file(path);
517 }
518}
519
520fn collect_bounded_ordered_rows<F>(
521 mut inner: Box<dyn QueryResult>,
522 compare: &F,
523) -> Result<(Vec<String>, BoundedOrderedRows)>
524where
525 F: Fn(&Row, &Row) -> std::cmp::Ordering,
526{
527 let columns = inner.columns().to_vec();
528 let mut rows = RowVec::with_capacity(ORDERED_RUN_MAX_ROWS.min(1024));
529 let mut budget = RetainedRowsBudget::with_limits(
530 "bounded ORDER BY run",
531 ORDERED_RUN_MAX_ROWS,
532 ORDERED_RUN_MAX_BYTES,
533 );
534 let mut paths = Vec::new();
535 let mut input_rows = 0usize;
536 let mut peak_rows = 0usize;
537 let mut peak_bytes = 0usize;
538
539 while inner.next() {
540 let row = inner.take_row();
541 if let Err(error) = budget.admit(&row) {
542 if rows.is_empty() {
543 return Err(error);
544 }
545 peak_rows = peak_rows.max(budget.peak_rows());
546 peak_bytes = peak_bytes.max(budget.peak_bytes());
547 let path = match write_ordered_run(&mut rows, compare) {
548 Ok(path) => path,
549 Err(error) => {
550 remove_ordered_runs(paths);
551 return Err(error);
552 }
553 };
554 paths.push(path);
555 rows.clear();
556 budget = RetainedRowsBudget::with_limits(
557 "bounded ORDER BY run",
558 ORDERED_RUN_MAX_ROWS,
559 ORDERED_RUN_MAX_BYTES,
560 );
561 if let Err(error) = budget.admit(&row) {
562 remove_ordered_runs(paths);
563 return Err(error);
564 }
565 }
566 rows.push((input_rows as i64, row));
567 input_rows = input_rows.saturating_add(1);
568 }
569 if let Some(error) = inner.last_error() {
570 remove_ordered_runs(paths);
571 return Err(error);
572 }
573 peak_rows = peak_rows.max(budget.peak_rows());
574 peak_bytes = peak_bytes.max(budget.peak_bytes());
575
576 if paths.is_empty() {
577 return Ok((
578 columns,
579 BoundedOrderedRows::Memory {
580 rows,
581 input_rows,
582 peak_rows,
583 peak_bytes,
584 },
585 ));
586 }
587 if !rows.is_empty() {
588 match write_ordered_run(&mut rows, compare) {
589 Ok(path) => paths.push(path),
590 Err(error) => {
591 remove_ordered_runs(paths);
592 return Err(error);
593 }
594 }
595 }
596 Ok((
597 columns,
598 BoundedOrderedRows::External {
599 paths,
600 input_rows,
601 peak_rows,
602 peak_bytes,
603 },
604 ))
605}
606
607impl OrderedResult {
608 pub fn new<F>(inner: Box<dyn QueryResult>, compare: F) -> Result<Self>
610 where
611 F: Fn(&Row, &Row) -> std::cmp::Ordering + Send + 'static,
612 {
613 let collect_started = radixdb_core::time_compat::Instant::now();
614 let (columns, bounded) = collect_bounded_ordered_rows(inner, &compare)?;
615 let collect_elapsed = collect_started.elapsed();
616 let inner: Box<dyn QueryResult> = match bounded {
617 BoundedOrderedRows::Memory {
618 mut rows,
619 input_rows,
620 peak_rows,
621 peak_bytes,
622 } => {
623 let finalize_started = radixdb_core::time_compat::Instant::now();
624 rows.sort_unstable_by(|(_, left), (_, right)| compare(left, right));
625 let finalize_elapsed = finalize_started.elapsed();
626 radixdb_storage::instrumentation::record_join_ordered_sort(
627 input_rows as u64,
628 0,
629 peak_rows as u64,
630 peak_bytes as u64,
631 collect_elapsed,
632 finalize_elapsed,
633 );
634 Box::new(ExecutorResult::new(columns, rows))
635 }
636 BoundedOrderedRows::External {
637 paths,
638 input_rows,
639 peak_rows,
640 peak_bytes,
641 } => {
642 let runs = paths.len();
643 let finalize_started = radixdb_core::time_compat::Instant::now();
644 let external = ExternalOrderedResult::new(columns, paths, input_rows, compare)?;
645 let finalize_elapsed = finalize_started.elapsed();
646 radixdb_storage::instrumentation::record_join_ordered_sort(
647 input_rows as u64,
648 runs as u64,
649 peak_rows as u64,
650 peak_bytes as u64,
651 collect_elapsed,
652 finalize_elapsed,
653 );
654 Box::new(external)
655 }
656 };
657 Ok(Self { inner })
658 }
659
660 pub fn new_radix<F>(
670 inner: Box<dyn QueryResult>,
671 order_specs: &[RadixOrderSpec],
672 fallback_compare: F,
673 ) -> Result<Self>
674 where
675 F: Fn(&Row, &Row) -> std::cmp::Ordering + Send + 'static,
676 {
677 let collect_started = radixdb_core::time_compat::Instant::now();
678 let (columns, bounded) = collect_bounded_ordered_rows(inner, &fallback_compare)?;
679 let collect_elapsed = collect_started.elapsed();
680 let BoundedOrderedRows::Memory {
681 mut rows,
682 input_rows,
683 peak_rows,
684 peak_bytes,
685 } = bounded
686 else {
687 let BoundedOrderedRows::External {
688 paths,
689 input_rows,
690 peak_rows,
691 peak_bytes,
692 } = bounded
693 else {
694 unreachable!()
695 };
696 let runs = paths.len();
697 let finalize_started = radixdb_core::time_compat::Instant::now();
698 let external =
699 ExternalOrderedResult::new(columns, paths, input_rows, fallback_compare)?;
700 let finalize_elapsed = finalize_started.elapsed();
701 radixdb_storage::instrumentation::record_join_ordered_sort(
702 input_rows as u64,
703 runs as u64,
704 peak_rows as u64,
705 peak_bytes as u64,
706 collect_elapsed,
707 finalize_elapsed,
708 );
709 return Ok(Self {
710 inner: Box::new(external),
711 });
712 };
713
714 let has_explicit_nulls_ordering = order_specs.iter().any(|s| s.nulls_first.is_some());
717
718 if !has_explicit_nulls_ordering {
719 if order_specs.len() == 1 {
721 let spec = &order_specs[0];
722 let finalize_started = radixdb_core::time_compat::Instant::now();
723 if Self::try_radix_sort_single_int(&mut rows, spec.col_idx, spec.ascending) {
724 let finalize_elapsed = finalize_started.elapsed();
725 radixdb_storage::instrumentation::record_join_ordered_sort(
726 input_rows as u64,
727 0,
728 peak_rows as u64,
729 peak_bytes as u64,
730 collect_elapsed,
731 finalize_elapsed,
732 );
733 return Ok(Self {
734 inner: Box::new(ExecutorResult::new(columns, rows)),
735 });
736 }
737 }
738
739 let finalize_started = radixdb_core::time_compat::Instant::now();
741 if order_specs.len() <= 4 && Self::try_radix_sort_multi_int(&mut rows, order_specs) {
742 let finalize_elapsed = finalize_started.elapsed();
743 radixdb_storage::instrumentation::record_join_ordered_sort(
744 input_rows as u64,
745 0,
746 peak_rows as u64,
747 peak_bytes as u64,
748 collect_elapsed,
749 finalize_elapsed,
750 );
751 return Ok(Self {
752 inner: Box::new(ExecutorResult::new(columns, rows)),
753 });
754 }
755 }
756
757 let finalize_started = radixdb_core::time_compat::Instant::now();
759 rows.sort_unstable_by(|(_, a), (_, b)| fallback_compare(a, b));
760 let finalize_elapsed = finalize_started.elapsed();
761 radixdb_storage::instrumentation::record_join_ordered_sort(
762 input_rows as u64,
763 0,
764 peak_rows as u64,
765 peak_bytes as u64,
766 collect_elapsed,
767 finalize_elapsed,
768 );
769
770 Ok(Self {
771 inner: Box::new(ExecutorResult::new(columns, rows)),
772 })
773 }
774
775 fn try_radix_sort_single_int(rows: &mut RowVec, col_idx: usize, ascending: bool) -> bool {
778 if Self::try_radix_sort_single_uuid(rows, col_idx, ascending) {
783 return true;
784 }
785
786 for (_, row) in rows.iter() {
788 match row.get(col_idx) {
789 Some(Value::Integer(value)) if *value != i64::MIN => continue,
790 _ => return false, }
792 }
793
794 if ascending {
797 radsort::sort_by_key(rows, |(_, row)| match row.get(col_idx) {
798 Some(Value::Integer(i)) => *i,
799 _ => unreachable!("radix admission requires non-null integers"),
800 });
801 } else {
802 radsort::sort_by_key(rows, |(_, row)| {
805 match row.get(col_idx) {
806 Some(Value::Integer(i)) => {
807 i.wrapping_neg().wrapping_sub(1)
809 }
810 _ => unreachable!("radix admission requires non-null integers"),
811 }
812 });
813 }
814
815 true
816 }
817
818 pub(super) fn try_radix_sort_single_uuid(
819 rows: &mut RowVec,
820 col_idx: usize,
821 ascending: bool,
822 ) -> bool {
823 if rows
824 .iter()
825 .any(|(_, row)| row.get(col_idx).and_then(Value::as_uuid_bytes).is_none())
826 {
827 return false;
828 }
829
830 let word = |row: &Row, range: std::ops::Range<usize>| {
831 let bytes = row
832 .get(col_idx)
833 .and_then(Value::as_uuid_bytes)
834 .expect("UUID radix admission validates every key");
835 let word = u64::from_be_bytes(
836 bytes[range]
837 .try_into()
838 .expect("UUID radix word is exactly eight bytes"),
839 );
840 if ascending {
841 word
842 } else {
843 !word
844 }
845 };
846
847 radsort::sort_by_key(rows, |(_, row)| word(row, 8..16));
850 radsort::sort_by_key(rows, |(_, row)| word(row, 0..8));
851 true
852 }
853
854 fn try_radix_sort_multi_int(rows: &mut RowVec, order_specs: &[RadixOrderSpec]) -> bool {
857 for (_, row) in rows.iter() {
859 for spec in order_specs {
860 match row.get(spec.col_idx) {
861 Some(Value::Integer(value)) if *value != i64::MIN => continue,
862 _ => return false,
863 }
864 }
865 }
866
867 for spec in order_specs.iter().rev() {
871 if spec.ascending {
872 radsort::sort_by_key(rows, |(_, row)| match row.get(spec.col_idx) {
873 Some(Value::Integer(i)) => *i,
874 _ => unreachable!("radix admission requires non-null integers"),
875 });
876 } else {
877 radsort::sort_by_key(rows, |(_, row)| match row.get(spec.col_idx) {
878 Some(Value::Integer(i)) => i.wrapping_neg().wrapping_sub(1),
879 _ => unreachable!("radix admission requires non-null integers"),
880 });
881 }
882 }
883
884 true
885 }
886}
887
888impl QueryResult for OrderedResult {
889 fn columns(&self) -> &[String] {
890 self.inner.columns()
891 }
892
893 fn columns_arc(&self) -> Option<CompactArc<Vec<String>>> {
894 self.inner.columns_arc()
895 }
896
897 fn next(&mut self) -> bool {
898 self.inner.next()
899 }
900
901 fn scan(&self, dest: &mut [Value]) -> Result<()> {
902 self.inner.scan(dest)
903 }
904
905 fn row(&self) -> &Row {
906 self.inner.row()
907 }
908
909 fn take_row(&mut self) -> Row {
910 self.inner.take_row()
911 }
912
913 fn close(&mut self) -> Result<()> {
914 self.inner.close()
915 }
916
917 fn rows_affected(&self) -> i64 {
918 0
919 }
920
921 fn last_insert_id(&self) -> i64 {
922 self.inner.last_insert_id()
923 }
924
925 fn last_error(&mut self) -> Option<Error> {
926 self.inner.last_error()
927 }
928
929 fn estimated_count(&self) -> Option<usize> {
930 self.inner.estimated_count()
931 }
932
933 fn with_aliases(self: Box<Self>, aliases: FxHashMap<String, String>) -> Box<dyn QueryResult> {
934 Box::new(AliasedResult::new(self, aliases))
935 }
936}
937
938pub struct TopNResult {
943 inner: ExecutorResult,
945 _budget: RetainedRowsBudget,
948}
949
950impl TopNResult {
951 pub fn new<F>(
962 inner: Box<dyn QueryResult>,
963 compare: F,
964 limit: usize,
965 offset: usize,
966 ) -> Result<Self>
967 where
968 F: Fn(&Row, &Row) -> std::cmp::Ordering + Clone,
969 {
970 Self::build(
971 inner,
972 compare,
973 limit,
974 offset,
975 RetainedRowsBudget::new("TOP-N"),
976 )
977 }
978
979 pub fn new_with_context<F>(
980 inner: Box<dyn QueryResult>,
981 compare: F,
982 limit: usize,
983 offset: usize,
984 ctx: &crate::context::ExecutionContext,
985 ) -> Result<Self>
986 where
987 F: Fn(&Row, &Row) -> std::cmp::Ordering + Clone,
988 {
989 Self::build(
990 inner,
991 compare,
992 limit,
993 offset,
994 RetainedRowsBudget::with_request_memory("TOP-N", ctx)?,
995 )
996 }
997
998 fn build<F>(
999 mut inner: Box<dyn QueryResult>,
1000 compare: F,
1001 limit: usize,
1002 offset: usize,
1003 mut budget: RetainedRowsBudget,
1004 ) -> Result<Self>
1005 where
1006 F: Fn(&Row, &Row) -> std::cmp::Ordering + Clone,
1007 {
1008 use std::collections::BinaryHeap;
1009
1010 let columns = inner.columns().to_vec();
1011 let heap_capacity = limit.saturating_add(offset);
1012 budget.ensure_capacity(heap_capacity)?;
1013
1014 if heap_capacity == 0 {
1016 return Ok(Self {
1017 inner: ExecutorResult::new(columns, RowVec::new()),
1018 _budget: budget,
1019 });
1020 }
1021
1022 let compare = std::sync::Arc::new(compare);
1024
1025 struct HeapRow<F: Fn(&Row, &Row) -> std::cmp::Ordering> {
1027 row: Row,
1028 compare: std::sync::Arc<F>,
1029 }
1030
1031 impl<F: Fn(&Row, &Row) -> std::cmp::Ordering> PartialEq for HeapRow<F> {
1032 fn eq(&self, other: &Self) -> bool {
1033 (self.compare)(&self.row, &other.row) == std::cmp::Ordering::Equal
1034 }
1035 }
1036
1037 impl<F: Fn(&Row, &Row) -> std::cmp::Ordering> Eq for HeapRow<F> {}
1038
1039 impl<F: Fn(&Row, &Row) -> std::cmp::Ordering> PartialOrd for HeapRow<F> {
1040 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
1041 Some(self.cmp(other))
1042 }
1043 }
1044
1045 impl<F: Fn(&Row, &Row) -> std::cmp::Ordering> Ord for HeapRow<F> {
1046 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
1047 (self.compare)(&self.row, &other.row)
1050 }
1051 }
1052
1053 let mut heap: BinaryHeap<HeapRow<F>> = BinaryHeap::with_capacity(heap_capacity + 1);
1054
1055 let mut input_rows = 0_u64;
1056 while inner.next() {
1057 input_rows = input_rows.saturating_add(1);
1058 let row = inner.take_row();
1059
1060 if heap.len() < heap_capacity {
1061 budget.admit(&row)?;
1062 heap.push(HeapRow {
1063 row,
1064 compare: std::sync::Arc::clone(&compare),
1065 });
1066 } else if let Some(worst) = heap.peek() {
1067 if compare(&row, &worst.row) == std::cmp::Ordering::Less {
1068 if let Some(removed) = heap.pop() {
1069 budget.release(&removed.row);
1070 }
1071 budget.admit(&row)?;
1072 heap.push(HeapRow {
1073 row,
1074 compare: std::sync::Arc::clone(&compare),
1075 });
1076 }
1077 }
1078 }
1079 if let Some(err) = inner.last_error() {
1080 return Err(err);
1081 }
1082
1083 let mut rows: Vec<Row> = heap.into_iter().map(|hr| hr.row).collect();
1085 rows.sort_unstable_by(|a, b| compare(a, b));
1086
1087 if offset > 0 && offset < rows.len() {
1089 for row in rows.drain(..offset) {
1090 budget.release(&row);
1091 }
1092 } else if offset >= rows.len() {
1093 for row in &rows {
1094 budget.release(row);
1095 }
1096 rows.clear();
1097 }
1098
1099 let result_rows: RowVec = rows
1101 .into_iter()
1102 .enumerate()
1103 .map(|(i, row)| (i as i64, row))
1104 .collect();
1105
1106 radixdb_storage::instrumentation::record_join_top_n(
1107 input_rows,
1108 budget.peak_rows() as u64,
1109 budget.peak_bytes() as u64,
1110 result_rows.len() as u64,
1111 );
1112
1113 Ok(Self {
1114 inner: ExecutorResult::new(columns, result_rows),
1115 _budget: budget,
1116 })
1117 }
1118
1119 pub fn from_rows_with_budget(
1120 columns: Vec<String>,
1121 rows: RowVec,
1122 budget: RetainedRowsBudget,
1123 ) -> Self {
1124 Self {
1125 inner: ExecutorResult::new(columns, rows),
1126 _budget: budget,
1127 }
1128 }
1129}
1130
1131impl QueryResult for TopNResult {
1132 fn columns(&self) -> &[String] {
1133 self.inner.columns()
1134 }
1135
1136 fn next(&mut self) -> bool {
1137 self.inner.next()
1138 }
1139
1140 fn scan(&self, dest: &mut [Value]) -> Result<()> {
1141 self.inner.scan(dest)
1142 }
1143
1144 fn row(&self) -> &Row {
1145 self.inner.row()
1146 }
1147
1148 fn take_row(&mut self) -> Row {
1149 self.inner.take_row()
1150 }
1151
1152 fn close(&mut self) -> Result<()> {
1153 self.inner.close()
1154 }
1155
1156 fn rows_affected(&self) -> i64 {
1157 0
1158 }
1159
1160 fn last_insert_id(&self) -> i64 {
1161 0
1162 }
1163
1164 fn with_aliases(self: Box<Self>, aliases: FxHashMap<String, String>) -> Box<dyn QueryResult> {
1165 Box::new(AliasedResult::new(self, aliases))
1166 }
1167}