From e16b795affe7ccada511873c7e0e1a52dbc54d5c Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 29 Jun 2026 07:58:30 +0530 Subject: [PATCH 1/5] feat: stream cross-shard ORDER BY results with k-way merge --- pgdog/src/backend/pool/connection/binding.rs | 4 +- pgdog/src/backend/pool/connection/buffer.rs | 340 ++++++++++++++---- .../pool/connection/multi_shard/mod.rs | 89 ++++- .../pool/connection/multi_shard/test.rs | 77 +++- 4 files changed, 409 insertions(+), 101 deletions(-) diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index fc4981290..33be6a3b4 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -103,7 +103,7 @@ impl Binding { return Ok(message); } let mut read = false; - for server in shards.iter_mut() { + for (position, server) in shards.iter_mut().enumerate() { if !server.has_more_messages() { continue; } @@ -111,7 +111,7 @@ impl Binding { let message = server.read().await?; read = true; - if let Some(message) = state.forward(message)? { + if let Some(message) = state.forward_from(position, message)? { return Ok(message); } } diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index edea8fd85..22edd4fec 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -20,12 +20,212 @@ use pgdog_postgres_types::Datum; use super::Aggregates; +#[derive(Debug, Clone)] +struct OrderByMergeState { + columns: Vec, + per_shard: Vec>, + done: Vec, + ready: VecDeque, + offset: usize, + limit: Option, + skipped: usize, + emitted: usize, + exhausted: bool, +} + +impl OrderByMergeState { + fn new(shards: usize, columns: Vec, limit: &Limit) -> Self { + let mut per_shard = Vec::with_capacity(shards); + per_shard.resize_with(shards, VecDeque::new); + Self { + columns, + per_shard, + done: vec![false; shards], + ready: VecDeque::new(), + offset: limit.offset.unwrap_or(0), + limit: limit.limit, + skipped: 0, + emitted: 0, + exhausted: false, + } + } + + fn push_row(&mut self, shard: usize, row: DataRow, decoder: &Decoder) { + if self.exhausted { + return; + } + if let Some(queue) = self.per_shard.get_mut(shard) { + queue.push_back(row); + self.drain_ready(decoder); + } + } + + fn mark_done(&mut self, shard: usize, decoder: &Decoder) { + if let Some(done) = self.done.get_mut(shard) { + *done = true; + self.drain_ready(decoder); + } + } + + fn take_ready(&mut self) -> Option { + self.ready.pop_front().and_then(|row| row.message().ok()) + } + + fn output_rows(&self) -> usize { + self.emitted + } + + fn can_emit_next(&self) -> bool { + let mut has_candidate = false; + for (queue, done) in self.per_shard.iter().zip(self.done.iter()) { + if queue.is_empty() && !done { + return false; + } + has_candidate |= !queue.is_empty(); + } + has_candidate + } + + fn choose_shard(&self, decoder: &Decoder) -> Option { + let mut best: Option = None; + for (idx, queue) in self.per_shard.iter().enumerate() { + let Some(head) = queue.front() else { + continue; + }; + let replace = best + .and_then(|best_idx| { + self.per_shard + .get(best_idx) + .and_then(|best_queue| best_queue.front()) + .map(|best_head| { + compare_rows(head, best_head, &self.columns, decoder) == Ordering::Less + }) + }) + .unwrap_or(true); + if replace { + best = Some(idx); + } + } + best + } + + fn drain_ready(&mut self, decoder: &Decoder) { + if self.exhausted { + return; + } + loop { + if let Some(limit) = self.limit + && self.emitted >= limit + { + self.exhausted = true; + for queue in &mut self.per_shard { + queue.clear(); + } + return; + } + if !self.can_emit_next() { + return; + } + let Some(shard) = self.choose_shard(decoder) else { + return; + }; + let Some(row) = self.per_shard[shard].pop_front() else { + return; + }; + if self.skipped < self.offset { + self.skipped += 1; + continue; + } + self.emitted += 1; + self.ready.push_back(row); + } + } +} + +fn normalize_order_by(columns: &[OrderBy], decoder: &Decoder) -> Vec { + let mut normalized = vec![]; + for column in columns { + match column { + OrderBy::Asc(_) => normalized.push(column.clone()), + OrderBy::AscColumn(name) => { + if let Some(index) = decoder.rd().field_index(name) { + normalized.push(OrderBy::Asc(index + 1)); + } + // TODO: Error out instead of silently not sorting. + } + OrderBy::Desc(_) => normalized.push(column.clone()), + OrderBy::DescColumn(name) => { + if let Some(index) = decoder.rd().field_index(name) { + normalized.push(OrderBy::Desc(index + 1)); + } + // TODO: Error out instead of silently not sorting. + } + OrderBy::AscVectorL2(_, _) => normalized.push(column.clone()), + OrderBy::AscVectorL2Column(name, vector) => { + if let Some(index) = decoder.rd().field_index(name) { + normalized.push(OrderBy::AscVectorL2(index + 1, vector.clone())); + } + // TODO: Error out instead of silently not sorting. + } + }; + } + normalized +} + +fn compare_rows(a: &DataRow, b: &DataRow, cols: &[OrderBy], decoder: &Decoder) -> Ordering { + cols.iter() + .filter_map(|col| { + let index = col.index(); + let asc = col.asc(); + let index = index?; + let left = a.get_column(index, decoder); + let right = b.get_column(index, decoder); + + match (left, right) { + (Ok(Some(left)), Ok(Some(right))) => { + // Handle the special vector case. + if let OrderBy::AscVectorL2(_, vector) = col { + let left: Option = left.value.try_into().ok(); + let right: Option = right.value.try_into().ok(); + + if let (Some(left), Some(right)) = (left, right) { + let left = left.distance_l2(vector); + let right = right.distance_l2(vector); + + left.partial_cmp(&right) + } else { + Some(Ordering::Equal) + } + } else { + // FIXME(sage): We don't handle ASC NULLS FIRST or + // DESC NULLS LAST we should either error or add + // support rather than silently do the wrong sorting + match (&left.value, &right.value, asc) { + (Datum::Null, Datum::Null, _) => Some(Ordering::Equal), + (Datum::Null, _, true) => Some(Ordering::Greater), + (_, Datum::Null, true) => Some(Ordering::Less), + (Datum::Null, _, false) => Some(Ordering::Less), + (_, Datum::Null, false) => Some(Ordering::Greater), + (a, b, true) => a.partial_cmp(b), + (a, b, false) => b.partial_cmp(a), + } + } + } + + _ => Some(Ordering::Equal), + } + }) + .reduce(Ordering::then) + .unwrap_or(Ordering::Equal) +} + /// Sort and aggregate rows received from multiple shards. #[derive(Default, Debug, Clone)] pub(super) struct Buffer { buffer: VecDeque, full: bool, distinct: HashSet, + merge: Option, } impl Buffer { @@ -44,91 +244,54 @@ impl Buffer { self.full = true; } + pub(super) fn begin_order_by_merge( + &mut self, + shards: usize, + columns: &[OrderBy], + limit: &Limit, + decoder: &Decoder, + ) { + self.merge = Some(OrderByMergeState::new( + shards, + normalize_order_by(columns, decoder), + limit, + )); + } + + pub(super) fn merge_add( + &mut self, + shard: usize, + message: Message, + decoder: &Decoder, + ) -> Result<(), super::Error> { + let dr = DataRow::from_bytes(message.to_bytes())?; + if let Some(merge) = self.merge.as_mut() { + merge.push_row(shard, dr, decoder); + } else { + self.buffer.push_back(dr); + } + + Ok(()) + } + + pub(super) fn merge_done(&mut self, shard: usize, decoder: &Decoder) { + if let Some(merge) = self.merge.as_mut() { + merge.mark_done(shard, decoder); + } + } + pub(super) fn reset(&mut self) { self.buffer.clear(); self.full = false; + self.merge = None; } /// Sort the buffer. pub(super) fn sort(&mut self, columns: &[OrderBy], decoder: &Decoder) { - // Calculate column indices once, since - // fetching indices by name is O(number of columns). - let mut cols = vec![]; - for column in columns { - match column { - OrderBy::Asc(_) => cols.push(column.clone()), - OrderBy::AscColumn(name) => { - if let Some(index) = decoder.rd().field_index(name) { - cols.push(OrderBy::Asc(index + 1)); - } - // TODO: Error out instead of silently not sorting. - } - OrderBy::Desc(_) => cols.push(column.clone()), - OrderBy::DescColumn(name) => { - if let Some(index) = decoder.rd().field_index(name) { - cols.push(OrderBy::Desc(index + 1)); - } - // TODO: Error out instead of silently not sorting. - } - OrderBy::AscVectorL2(_, _) => cols.push(column.clone()), - OrderBy::AscVectorL2Column(name, vector) => { - if let Some(index) = decoder.rd().field_index(name) { - cols.push(OrderBy::AscVectorL2(index + 1, vector.clone())); - } - // TODO: Error out instead of silently not sorting. - } - }; - } - - // Sort rows. - let order_by = move |a: &DataRow, b: &DataRow| -> Ordering { - cols.iter() - .filter_map(|col| { - let index = col.index(); - let asc = col.asc(); - let index = index?; - let left = a.get_column(index, decoder); - let right = b.get_column(index, decoder); - - match (left, right) { - (Ok(Some(left)), Ok(Some(right))) => { - // Handle the special vector case. - if let OrderBy::AscVectorL2(_, vector) = col { - let left: Option = left.value.try_into().ok(); - let right: Option = right.value.try_into().ok(); - - if let (Some(left), Some(right)) = (left, right) { - let left = left.distance_l2(vector); - let right = right.distance_l2(vector); - - left.partial_cmp(&right) - } else { - Some(Ordering::Equal) - } - } else { - // FIXME(sage): We don't handle ASC NULLS FIRST or - // DESC NULLS LAST we should either error or add - // support rather than silently do the wrong sorting - match (&left.value, &right.value, asc) { - (Datum::Null, Datum::Null, _) => Some(Ordering::Equal), - (Datum::Null, _, true) => Some(Ordering::Greater), - (_, Datum::Null, true) => Some(Ordering::Less), - (Datum::Null, _, false) => Some(Ordering::Less), - (_, Datum::Null, false) => Some(Ordering::Greater), - (a, b, true) => a.partial_cmp(b), - (a, b, false) => b.partial_cmp(a), - } - } - } - - _ => Some(Ordering::Equal), - } - }) - .reduce(Ordering::then) - .unwrap_or(Ordering::Equal) - }; - - self.buffer.make_contiguous().sort_by(order_by); + let cols = normalize_order_by(columns, decoder); + self.buffer + .make_contiguous() + .sort_by(|a, b| compare_rows(a, b, &cols, decoder)); } /// Execute aggregate functions. @@ -209,6 +372,9 @@ impl Buffer { /// Take messages from buffer. pub(super) fn take(&mut self) -> Option { + if let Some(merge) = self.merge.as_mut() { + return merge.take_ready(); + } if self.full { self.buffer.pop_front().and_then(|s| s.message().ok()) } else { @@ -227,7 +393,23 @@ impl Buffer { } pub(super) fn len(&self) -> usize { - self.buffer.len() + if let Some(merge) = self.merge.as_ref() { + merge.ready.len() + } else { + self.buffer.len() + } + } + + pub(super) fn output_rows(&self) -> usize { + if let Some(merge) = self.merge.as_ref() { + merge.output_rows() + } else { + self.buffer.len() + } + } + + pub(super) fn merge_active(&self) -> bool { + self.merge.is_some() } #[allow(dead_code)] diff --git a/pgdog/src/backend/pool/connection/multi_shard/mod.rs b/pgdog/src/backend/pool/connection/multi_shard/mod.rs index 8c00ee2a2..d764f6f90 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/mod.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/mod.rs @@ -114,10 +114,29 @@ impl MultiShard { /// Check if the message should be sent to the client, skipped, /// or modified. + #[cfg(test)] pub(super) fn forward(&mut self, message: Message) -> Result, Error> { + self.forward_inner(None, message) + } + + /// Same as [`Self::forward`], but includes shard position. + pub(super) fn forward_from( + &mut self, + shard_position: usize, + message: Message, + ) -> Result, Error> { + self.forward_inner(Some(shard_position), message) + } + + fn forward_inner( + &mut self, + shard_position: Option, + message: Message, + ) -> Result, Error> { let mut forward = None; + let code = message.code(); - match message.code() { + match code { 'Z' => { self.counters.ready_for_query += 1; if message.transaction_error() { @@ -155,31 +174,42 @@ impl MultiShard { false }; self.counters.command_complete_count += 1; + if self.buffer.merge_active() + && let Some(shard) = shard_position + { + self.buffer.merge_done(shard, &self.decoder); + } if self .counters .command_complete_count .is_multiple_of(self.shards) { - self.buffer.full(); - - if !self.buffer.is_empty() { - self.buffer - .aggregate( - self.route.aggregate(), - &self.decoder, - self.route.aggregate_rewrite_plan(), - ) - .map_err(Error::from)?; - - self.buffer.sort(self.route.order_by(), &self.decoder); - self.buffer.distinct(self.route.distinct(), &self.decoder); - self.buffer.limit(self.route.limit()); + if self.buffer.merge_active() { + self.buffer.full(); + } else { + self.buffer.full(); + + if !self.buffer.is_empty() { + self.buffer + .aggregate( + self.route.aggregate(), + &self.decoder, + self.route.aggregate_rewrite_plan(), + ) + .map_err(Error::from)?; + + self.buffer.sort(self.route.order_by(), &self.decoder); + self.buffer.distinct(self.route.distinct(), &self.decoder); + self.buffer.limit(self.route.limit()); + } } if has_rows { - let rows = if self.should_buffer() { - self.buffer.len() + let rows = if self.buffer.merge_active() { + self.buffer.output_rows() + } else if self.should_buffer() { + self.buffer.output_rows() } else { self.counters.rows }; @@ -206,6 +236,14 @@ impl MultiShard { if self.counters.row_description == self.shards { // Only send it to the client once all shards sent it, // so we don't get early requests from clients. + if self.should_stream_order_by() { + self.buffer.begin_order_by_merge( + self.shards, + self.route.order_by(), + self.route.limit(), + &self.decoder, + ); + } let plan = self.route.aggregate_rewrite_plan(); if plan.is_noop() { forward = Some(message); @@ -241,7 +279,14 @@ impl MultiShard { self.counters.first_backend_data = message.source().backend_id(); } - if !self.should_buffer() + if self.buffer.merge_active() + && self.counters.row_description.is_multiple_of(self.shards) + && let Some(shard) = shard_position + { + self.buffer + .merge_add(shard, message, &self.decoder) + .map_err(Error::from)?; + } else if !self.should_buffer() && self.counters.row_description.is_multiple_of(self.shards) { if self.route.is_omnisharded() { @@ -338,6 +383,14 @@ impl MultiShard { self.shards > 1 && self.route.should_buffer() && !self.route.is_omnisharded() } + fn should_stream_order_by(&self) -> bool { + self.shards > 1 + && !self.route.order_by().is_empty() + && self.route.aggregate().is_empty() + && self.route.distinct().is_none() + && !self.route.is_omnisharded() + } + /// Multi-shard state is ready to send messages. pub(super) fn message(&mut self) -> Option { match self.buffer.take() { diff --git a/pgdog/src/backend/pool/connection/multi_shard/test.rs b/pgdog/src/backend/pool/connection/multi_shard/test.rs index 65eb3fe32..1ff5fbb3b 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/test.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/test.rs @@ -1,6 +1,6 @@ use crate::{ - frontend::router::parser::{Shard, ShardWithPriority}, - net::{DataRow, Field}, + frontend::router::parser::{Limit, OrderBy, Shard, ShardWithPriority}, + net::{DataRow, Field, Format}, }; use super::*; @@ -267,3 +267,76 @@ fn test_omni_data_rows_only_from_first_server() { .unwrap(); assert!(result.is_some()); // Should be forwarded } + +#[test] +fn test_order_by_k_way_merge_streams_before_command_complete() { + let route = Route::select( + ShardWithPriority::new_table(Shard::All), + vec![OrderBy::Asc(1)], + Default::default(), + Limit::default(), + None, + ); + let mut multi_shard = MultiShard::new(vec![0, 1], &route); + + let rd = RowDescription::new(&[Field::bigint("id")]); + let backend1 = BackendPid::for_test(1); + let backend2 = BackendPid::for_test(2); + + let rd_result = multi_shard + .forward_from(0, rd.message().unwrap().backend(backend1)) + .unwrap(); + assert!(rd_result.is_none()); + let rd_result = multi_shard + .forward_from(1, rd.message().unwrap().backend(backend2)) + .unwrap(); + assert!(rd_result.is_some()); + + let mut shard0_first = DataRow::new(); + shard0_first.add(1_i64); + multi_shard + .forward_from(0, shard0_first.message().unwrap().backend(backend1)) + .unwrap(); + assert!(multi_shard.message().is_none()); + + let mut shard1_first = DataRow::new(); + shard1_first.add(2_i64); + multi_shard + .forward_from(1, shard1_first.message().unwrap().backend(backend2)) + .unwrap(); + let next = DataRow::from_bytes(multi_shard.message().unwrap().to_bytes()).unwrap(); + assert_eq!(next.get::(0, Format::Text).unwrap(), 1_i64); + + let mut shard0_second = DataRow::new(); + shard0_second.add(3_i64); + multi_shard + .forward_from(0, shard0_second.message().unwrap().backend(backend1)) + .unwrap(); + let next = DataRow::from_bytes(multi_shard.message().unwrap().to_bytes()).unwrap(); + assert_eq!(next.get::(0, Format::Text).unwrap(), 2_i64); + + multi_shard + .forward_from( + 1, + CommandComplete::from_str("SELECT 1") + .message() + .unwrap() + .backend(backend2), + ) + .unwrap(); + let next = DataRow::from_bytes(multi_shard.message().unwrap().to_bytes()).unwrap(); + assert_eq!(next.get::(0, Format::Text).unwrap(), 3_i64); + + multi_shard + .forward_from( + 0, + CommandComplete::from_str("SELECT 2") + .message() + .unwrap() + .backend(backend1), + ) + .unwrap(); + + let cc = CommandComplete::from_bytes(multi_shard.message().unwrap().to_bytes()).unwrap(); + assert_eq!(cc.rows().unwrap(), Some(3)); +} From b81a76ca466a5e52690d8d8ab9502238fc19b949 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 29 Jun 2026 08:08:46 +0530 Subject: [PATCH 2/5] fmt fixes --- pgdog/src/backend/pool/connection/multi_shard/mod.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/pgdog/src/backend/pool/connection/multi_shard/mod.rs b/pgdog/src/backend/pool/connection/multi_shard/mod.rs index d764f6f90..2884bc498 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/mod.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/mod.rs @@ -206,9 +206,7 @@ impl MultiShard { } if has_rows { - let rows = if self.buffer.merge_active() { - self.buffer.output_rows() - } else if self.should_buffer() { + let rows = if self.buffer.merge_active() || self.should_buffer() { self.buffer.output_rows() } else { self.counters.rows From 9d53b8b8957f56b8b31d44cf26a6d07d118d4074 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 29 Jun 2026 08:18:16 +0530 Subject: [PATCH 3/5] add tests --- pgdog/src/backend/pool/connection/buffer.rs | 78 +++++++++++++++++++ .../pool/connection/multi_shard/test.rs | 68 +++++++++++++++- 2 files changed, 145 insertions(+), 1 deletion(-) diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index 22edd4fec..2bfd0e1dd 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -691,6 +691,84 @@ mod test { } } + #[test] + fn test_order_by_merge_waits_for_all_shards_heads() { + let mut buf = Buffer::default(); + let rd = RowDescription::new(&[Field::bigint("id")]); + let decoder = Decoder::from(&rd); + + buf.begin_order_by_merge( + 2, + &[OrderBy::Asc(1)], + &Limit { + limit: None, + offset: None, + }, + &decoder, + ); + + let mut row = DataRow::new(); + row.add(1_i64); + buf.merge_add(0, row.message().unwrap(), &decoder).unwrap(); + + // Can't emit yet: shard 1 has not produced a head row or finished. + assert!(buf.take().is_none()); + + let mut row = DataRow::new(); + row.add(2_i64); + buf.merge_add(1, row.message().unwrap(), &decoder).unwrap(); + let out = DataRow::from_bytes(buf.take().unwrap().to_bytes()).unwrap(); + assert_eq!(out.get::(0, Format::Text).unwrap(), 1_i64); + + // Still can't emit row 2: shard 0 is neither done nor has another head row. + buf.merge_done(1, &decoder); + assert!(buf.take().is_none()); + + // Once shard 0 is done, shard 1's remaining head can be emitted. + buf.merge_done(0, &decoder); + let out = DataRow::from_bytes(buf.take().unwrap().to_bytes()).unwrap(); + assert_eq!(out.get::(0, Format::Text).unwrap(), 2_i64); + } + + #[test] + fn test_order_by_merge_honors_offset_and_limit() { + let mut buf = Buffer::default(); + let rd = RowDescription::new(&[Field::bigint("id")]); + let decoder = Decoder::from(&rd); + + buf.begin_order_by_merge( + 2, + &[OrderBy::Asc(1)], + &Limit { + limit: Some(2), + offset: Some(1), + }, + &decoder, + ); + + for value in [1_i64, 3_i64] { + let mut row = DataRow::new(); + row.add(value); + buf.merge_add(0, row.message().unwrap(), &decoder).unwrap(); + } + for value in [2_i64, 4_i64] { + let mut row = DataRow::new(); + row.add(value); + buf.merge_add(1, row.message().unwrap(), &decoder).unwrap(); + } + + buf.merge_done(0, &decoder); + buf.merge_done(1, &decoder); + + // Global order is 1,2,3,4 -> offset 1 + limit 2 => 2,3 + let out = DataRow::from_bytes(buf.take().unwrap().to_bytes()).unwrap(); + assert_eq!(out.get::(0, Format::Text).unwrap(), 2_i64); + let out = DataRow::from_bytes(buf.take().unwrap().to_bytes()).unwrap(); + assert_eq!(out.get::(0, Format::Text).unwrap(), 3_i64); + assert!(buf.take().is_none()); + assert_eq!(buf.output_rows(), 2); + } + #[test] fn test_limit() { let mut buf = Buffer::default(); diff --git a/pgdog/src/backend/pool/connection/multi_shard/test.rs b/pgdog/src/backend/pool/connection/multi_shard/test.rs index 1ff5fbb3b..a3db62107 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/test.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/test.rs @@ -1,5 +1,5 @@ use crate::{ - frontend::router::parser::{Limit, OrderBy, Shard, ShardWithPriority}, + frontend::router::parser::{DistinctBy, Limit, OrderBy, Shard, ShardWithPriority}, net::{DataRow, Field, Format}, }; @@ -340,3 +340,69 @@ fn test_order_by_k_way_merge_streams_before_command_complete() { let cc = CommandComplete::from_bytes(multi_shard.message().unwrap().to_bytes()).unwrap(); assert_eq!(cc.rows().unwrap(), Some(3)); } + +#[test] +fn test_order_by_with_distinct_uses_buffered_path() { + let route = Route::select( + ShardWithPriority::new_table(Shard::All), + vec![OrderBy::Asc(1)], + Default::default(), + Limit::default(), + Some(DistinctBy::Row), + ); + let mut multi_shard = MultiShard::new(vec![0, 1], &route); + + let rd = RowDescription::new(&[Field::bigint("id")]); + let backend1 = BackendPid::for_test(1); + let backend2 = BackendPid::for_test(2); + + multi_shard + .forward_from(0, rd.message().unwrap().backend(backend1)) + .unwrap(); + let rd_result = multi_shard + .forward_from(1, rd.message().unwrap().backend(backend2)) + .unwrap(); + assert!(rd_result.is_some()); + + let mut dup0 = DataRow::new(); + dup0.add(1_i64); + multi_shard + .forward_from(0, dup0.message().unwrap().backend(backend1)) + .unwrap(); + + let mut dup1 = DataRow::new(); + dup1.add(1_i64); + multi_shard + .forward_from(1, dup1.message().unwrap().backend(backend2)) + .unwrap(); + + // DISTINCT should keep us on the buffered path: no rows stream early. + assert!(multi_shard.message().is_none()); + + multi_shard + .forward_from( + 0, + CommandComplete::from_str("SELECT 1") + .message() + .unwrap() + .backend(backend1), + ) + .unwrap(); + assert!(multi_shard.message().is_none()); + + multi_shard + .forward_from( + 1, + CommandComplete::from_str("SELECT 1") + .message() + .unwrap() + .backend(backend2), + ) + .unwrap(); + + let row = DataRow::from_bytes(multi_shard.message().unwrap().to_bytes()).unwrap(); + assert_eq!(row.get::(0, Format::Text).unwrap(), 1_i64); + + let cc = CommandComplete::from_bytes(multi_shard.message().unwrap().to_bytes()).unwrap(); + assert_eq!(cc.rows().unwrap(), Some(1)); +} From fa3672f140e72971039ae64f167c0a3b6ff8948a Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Thu, 2 Jul 2026 15:48:09 +0530 Subject: [PATCH 4/5] resolve comments --- pgdog/src/backend/pool/connection/binding.rs | 16 +- pgdog/src/backend/pool/connection/buffer.rs | 291 +---------------- .../pool/connection/multi_shard/mod.rs | 8 +- .../backend/pool/connection/order_by_merge.rs | 298 ++++++++++++++++++ 4 files changed, 316 insertions(+), 297 deletions(-) create mode 100644 pgdog/src/backend/pool/connection/order_by_merge.rs diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index 33be6a3b4..41d4c09d6 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -103,7 +103,7 @@ impl Binding { return Ok(message); } let mut read = false; - for (position, server) in shards.iter_mut().enumerate() { + for (shard_position, server) in shards.iter_mut().enumerate() { if !server.has_more_messages() { continue; } @@ -111,7 +111,7 @@ impl Binding { let message = server.read().await?; read = true; - if let Some(message) = state.forward_from(position, message)? { + if let Some(message) = state.forward_from(shard_position, message)? { return Ok(message); } } @@ -144,11 +144,11 @@ impl Binding { let mut shards_sent = servers.len(); let mut futures = Vec::new(); - for (position, server) in servers.iter_mut().enumerate() { + for (shard_position, server) in servers.iter_mut().enumerate() { // Map positional index to actual shard number. // When only a subset of shards is connected (Shard::Multi binding), // positional indices don't match actual shard numbers. - let shard = state.shard_index(position); + let shard = state.shard_index(shard_position); let send = match client_request.route().shard() { Shard::Direct(s) => { shards_sent = 1; @@ -205,8 +205,8 @@ impl Binding { Binding::MultiShard(servers, state) => { if !servers.is_empty() { let mut futures = Vec::new(); - for (position, server) in servers.iter_mut().enumerate() { - let shard = state.shard_index(position); + for (shard_position, server) in servers.iter_mut().enumerate() { + let shard = state.shard_index(shard_position); let send = match route.shard() { Shard::Direct(s) => *s == shard, Shard::Multi(shards) => shards.contains(&shard), @@ -235,8 +235,8 @@ impl Binding { match self { Binding::MultiShard(servers, state) => { for row in rows { - for (position, server) in servers.iter_mut().enumerate() { - let shard = state.shard_index(position); + for (shard_position, server) in servers.iter_mut().enumerate() { + let shard = state.shard_index(shard_position); match row.shard() { Shard::Direct(row_shard) => { if shard == *row_shard { diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index 2bfd0e1dd..33ad48b45 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -1,9 +1,6 @@ //! Buffer messages to sort and aggregate them later. -use std::{ - cmp::Ordering, - collections::{HashSet, VecDeque}, -}; +use std::collections::{HashSet, VecDeque}; use crate::{ frontend::router::parser::{ @@ -12,212 +9,14 @@ use crate::{ }, net::{ Decoder, - messages::{DataRow, FromBytes, Message, Protocol, ToBytes, Vector}, + messages::{DataRow, FromBytes, Message, Protocol, ToBytes}, }, }; -use pgdog_postgres_types::Datum; - use super::Aggregates; - -#[derive(Debug, Clone)] -struct OrderByMergeState { - columns: Vec, - per_shard: Vec>, - done: Vec, - ready: VecDeque, - offset: usize, - limit: Option, - skipped: usize, - emitted: usize, - exhausted: bool, -} - -impl OrderByMergeState { - fn new(shards: usize, columns: Vec, limit: &Limit) -> Self { - let mut per_shard = Vec::with_capacity(shards); - per_shard.resize_with(shards, VecDeque::new); - Self { - columns, - per_shard, - done: vec![false; shards], - ready: VecDeque::new(), - offset: limit.offset.unwrap_or(0), - limit: limit.limit, - skipped: 0, - emitted: 0, - exhausted: false, - } - } - - fn push_row(&mut self, shard: usize, row: DataRow, decoder: &Decoder) { - if self.exhausted { - return; - } - if let Some(queue) = self.per_shard.get_mut(shard) { - queue.push_back(row); - self.drain_ready(decoder); - } - } - - fn mark_done(&mut self, shard: usize, decoder: &Decoder) { - if let Some(done) = self.done.get_mut(shard) { - *done = true; - self.drain_ready(decoder); - } - } - - fn take_ready(&mut self) -> Option { - self.ready.pop_front().and_then(|row| row.message().ok()) - } - - fn output_rows(&self) -> usize { - self.emitted - } - - fn can_emit_next(&self) -> bool { - let mut has_candidate = false; - for (queue, done) in self.per_shard.iter().zip(self.done.iter()) { - if queue.is_empty() && !done { - return false; - } - has_candidate |= !queue.is_empty(); - } - has_candidate - } - - fn choose_shard(&self, decoder: &Decoder) -> Option { - let mut best: Option = None; - for (idx, queue) in self.per_shard.iter().enumerate() { - let Some(head) = queue.front() else { - continue; - }; - let replace = best - .and_then(|best_idx| { - self.per_shard - .get(best_idx) - .and_then(|best_queue| best_queue.front()) - .map(|best_head| { - compare_rows(head, best_head, &self.columns, decoder) == Ordering::Less - }) - }) - .unwrap_or(true); - if replace { - best = Some(idx); - } - } - best - } - - fn drain_ready(&mut self, decoder: &Decoder) { - if self.exhausted { - return; - } - loop { - if let Some(limit) = self.limit - && self.emitted >= limit - { - self.exhausted = true; - for queue in &mut self.per_shard { - queue.clear(); - } - return; - } - if !self.can_emit_next() { - return; - } - let Some(shard) = self.choose_shard(decoder) else { - return; - }; - let Some(row) = self.per_shard[shard].pop_front() else { - return; - }; - if self.skipped < self.offset { - self.skipped += 1; - continue; - } - self.emitted += 1; - self.ready.push_back(row); - } - } -} - -fn normalize_order_by(columns: &[OrderBy], decoder: &Decoder) -> Vec { - let mut normalized = vec![]; - for column in columns { - match column { - OrderBy::Asc(_) => normalized.push(column.clone()), - OrderBy::AscColumn(name) => { - if let Some(index) = decoder.rd().field_index(name) { - normalized.push(OrderBy::Asc(index + 1)); - } - // TODO: Error out instead of silently not sorting. - } - OrderBy::Desc(_) => normalized.push(column.clone()), - OrderBy::DescColumn(name) => { - if let Some(index) = decoder.rd().field_index(name) { - normalized.push(OrderBy::Desc(index + 1)); - } - // TODO: Error out instead of silently not sorting. - } - OrderBy::AscVectorL2(_, _) => normalized.push(column.clone()), - OrderBy::AscVectorL2Column(name, vector) => { - if let Some(index) = decoder.rd().field_index(name) { - normalized.push(OrderBy::AscVectorL2(index + 1, vector.clone())); - } - // TODO: Error out instead of silently not sorting. - } - }; - } - normalized -} - -fn compare_rows(a: &DataRow, b: &DataRow, cols: &[OrderBy], decoder: &Decoder) -> Ordering { - cols.iter() - .filter_map(|col| { - let index = col.index(); - let asc = col.asc(); - let index = index?; - let left = a.get_column(index, decoder); - let right = b.get_column(index, decoder); - - match (left, right) { - (Ok(Some(left)), Ok(Some(right))) => { - // Handle the special vector case. - if let OrderBy::AscVectorL2(_, vector) = col { - let left: Option = left.value.try_into().ok(); - let right: Option = right.value.try_into().ok(); - - if let (Some(left), Some(right)) = (left, right) { - let left = left.distance_l2(vector); - let right = right.distance_l2(vector); - - left.partial_cmp(&right) - } else { - Some(Ordering::Equal) - } - } else { - // FIXME(sage): We don't handle ASC NULLS FIRST or - // DESC NULLS LAST we should either error or add - // support rather than silently do the wrong sorting - match (&left.value, &right.value, asc) { - (Datum::Null, Datum::Null, _) => Some(Ordering::Equal), - (Datum::Null, _, true) => Some(Ordering::Greater), - (_, Datum::Null, true) => Some(Ordering::Less), - (Datum::Null, _, false) => Some(Ordering::Less), - (_, Datum::Null, false) => Some(Ordering::Greater), - (a, b, true) => a.partial_cmp(b), - (a, b, false) => b.partial_cmp(a), - } - } - } - - _ => Some(Ordering::Equal), - } - }) - .reduce(Ordering::then) - .unwrap_or(Ordering::Equal) -} +#[path = "order_by_merge.rs"] +mod order_by_merge; +use order_by_merge::{OrderByMergeState, compare_rows, normalize_order_by}; /// Sort and aggregate rows received from multiple shards. #[derive(Default, Debug, Clone)] @@ -394,7 +193,7 @@ impl Buffer { pub(super) fn len(&self) -> usize { if let Some(merge) = self.merge.as_ref() { - merge.ready.len() + merge.ready_len() } else { self.buffer.len() } @@ -691,84 +490,6 @@ mod test { } } - #[test] - fn test_order_by_merge_waits_for_all_shards_heads() { - let mut buf = Buffer::default(); - let rd = RowDescription::new(&[Field::bigint("id")]); - let decoder = Decoder::from(&rd); - - buf.begin_order_by_merge( - 2, - &[OrderBy::Asc(1)], - &Limit { - limit: None, - offset: None, - }, - &decoder, - ); - - let mut row = DataRow::new(); - row.add(1_i64); - buf.merge_add(0, row.message().unwrap(), &decoder).unwrap(); - - // Can't emit yet: shard 1 has not produced a head row or finished. - assert!(buf.take().is_none()); - - let mut row = DataRow::new(); - row.add(2_i64); - buf.merge_add(1, row.message().unwrap(), &decoder).unwrap(); - let out = DataRow::from_bytes(buf.take().unwrap().to_bytes()).unwrap(); - assert_eq!(out.get::(0, Format::Text).unwrap(), 1_i64); - - // Still can't emit row 2: shard 0 is neither done nor has another head row. - buf.merge_done(1, &decoder); - assert!(buf.take().is_none()); - - // Once shard 0 is done, shard 1's remaining head can be emitted. - buf.merge_done(0, &decoder); - let out = DataRow::from_bytes(buf.take().unwrap().to_bytes()).unwrap(); - assert_eq!(out.get::(0, Format::Text).unwrap(), 2_i64); - } - - #[test] - fn test_order_by_merge_honors_offset_and_limit() { - let mut buf = Buffer::default(); - let rd = RowDescription::new(&[Field::bigint("id")]); - let decoder = Decoder::from(&rd); - - buf.begin_order_by_merge( - 2, - &[OrderBy::Asc(1)], - &Limit { - limit: Some(2), - offset: Some(1), - }, - &decoder, - ); - - for value in [1_i64, 3_i64] { - let mut row = DataRow::new(); - row.add(value); - buf.merge_add(0, row.message().unwrap(), &decoder).unwrap(); - } - for value in [2_i64, 4_i64] { - let mut row = DataRow::new(); - row.add(value); - buf.merge_add(1, row.message().unwrap(), &decoder).unwrap(); - } - - buf.merge_done(0, &decoder); - buf.merge_done(1, &decoder); - - // Global order is 1,2,3,4 -> offset 1 + limit 2 => 2,3 - let out = DataRow::from_bytes(buf.take().unwrap().to_bytes()).unwrap(); - assert_eq!(out.get::(0, Format::Text).unwrap(), 2_i64); - let out = DataRow::from_bytes(buf.take().unwrap().to_bytes()).unwrap(); - assert_eq!(out.get::(0, Format::Text).unwrap(), 3_i64); - assert!(buf.take().is_none()); - assert_eq!(buf.output_rows(), 2); - } - #[test] fn test_limit() { let mut buf = Buffer::default(); diff --git a/pgdog/src/backend/pool/connection/multi_shard/mod.rs b/pgdog/src/backend/pool/connection/multi_shard/mod.rs index 2884bc498..ac4620f33 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/mod.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/mod.rs @@ -80,12 +80,12 @@ impl MultiShard { } } - /// Map a positional index to the actual shard number. - pub(super) fn shard_index(&self, position: usize) -> usize { + /// Map a position in the connected servers list to the actual shard number. + pub(super) fn shard_index(&self, shard_position: usize) -> usize { self.shard_indices - .get(position) + .get(shard_position) .copied() - .unwrap_or(position) + .unwrap_or(shard_position) } /// Update multi-shard state. diff --git a/pgdog/src/backend/pool/connection/order_by_merge.rs b/pgdog/src/backend/pool/connection/order_by_merge.rs new file mode 100644 index 000000000..4d6ae576d --- /dev/null +++ b/pgdog/src/backend/pool/connection/order_by_merge.rs @@ -0,0 +1,298 @@ +use std::{cmp::Ordering, collections::VecDeque}; + +use crate::{ + frontend::router::parser::{Limit, OrderBy}, + net::{ + Decoder, + messages::{DataRow, Message, Protocol, Vector}, + }, +}; + +use pgdog_postgres_types::Datum; + +#[derive(Debug, Clone)] +pub(super) struct OrderByMergeState { + columns: Vec, + per_shard: Vec>, + done: Vec, + ready: VecDeque, + offset: usize, + limit: Option, + skipped: usize, + emitted: usize, + exhausted: bool, +} + +impl OrderByMergeState { + pub(super) fn new(shards: usize, columns: Vec, limit: &Limit) -> Self { + let mut per_shard = Vec::with_capacity(shards); + per_shard.resize_with(shards, VecDeque::new); + Self { + columns, + per_shard, + done: vec![false; shards], + ready: VecDeque::new(), + offset: limit.offset.unwrap_or(0), + limit: limit.limit, + skipped: 0, + emitted: 0, + exhausted: false, + } + } + + pub(super) fn push_row(&mut self, shard: usize, row: DataRow, decoder: &Decoder) { + if self.exhausted { + return; + } + if let Some(queue) = self.per_shard.get_mut(shard) { + queue.push_back(row); + self.drain_ready(decoder); + } + } + + pub(super) fn mark_done(&mut self, shard: usize, decoder: &Decoder) { + if let Some(done) = self.done.get_mut(shard) { + *done = true; + self.drain_ready(decoder); + } + } + + pub(super) fn take_ready(&mut self) -> Option { + self.ready.pop_front().and_then(|row| row.message().ok()) + } + + pub(super) fn output_rows(&self) -> usize { + self.emitted + } + + pub(super) fn ready_len(&self) -> usize { + self.ready.len() + } + + fn can_emit_next(&self) -> bool { + let mut has_candidate = false; + for (queue, done) in self.per_shard.iter().zip(self.done.iter()) { + if queue.is_empty() && !done { + return false; + } + has_candidate |= !queue.is_empty(); + } + has_candidate + } + + fn choose_shard(&self, decoder: &Decoder) -> Option { + let mut best: Option = None; + for (idx, queue) in self.per_shard.iter().enumerate() { + let Some(head) = queue.front() else { + continue; + }; + let replace = best + .and_then(|best_idx| { + self.per_shard + .get(best_idx) + .and_then(|best_queue| best_queue.front()) + .map(|best_head| { + compare_rows(head, best_head, &self.columns, decoder) == Ordering::Less + }) + }) + .unwrap_or(true); + if replace { + best = Some(idx); + } + } + best + } + + fn drain_ready(&mut self, decoder: &Decoder) { + if self.exhausted { + return; + } + loop { + if let Some(limit) = self.limit + && self.emitted >= limit + { + self.exhausted = true; + for queue in &mut self.per_shard { + queue.clear(); + } + return; + } + if !self.can_emit_next() { + return; + } + let Some(shard) = self.choose_shard(decoder) else { + return; + }; + let Some(row) = self.per_shard[shard].pop_front() else { + return; + }; + if self.skipped < self.offset { + self.skipped += 1; + continue; + } + self.emitted += 1; + self.ready.push_back(row); + } + } +} + +pub(super) fn normalize_order_by(columns: &[OrderBy], decoder: &Decoder) -> Vec { + let mut normalized = vec![]; + for column in columns { + match column { + OrderBy::Asc(_) => normalized.push(column.clone()), + OrderBy::AscColumn(name) => { + if let Some(index) = decoder.rd().field_index(name) { + normalized.push(OrderBy::Asc(index + 1)); + } + // TODO: Error out instead of silently not sorting. + } + OrderBy::Desc(_) => normalized.push(column.clone()), + OrderBy::DescColumn(name) => { + if let Some(index) = decoder.rd().field_index(name) { + normalized.push(OrderBy::Desc(index + 1)); + } + // TODO: Error out instead of silently not sorting. + } + OrderBy::AscVectorL2(_, _) => normalized.push(column.clone()), + OrderBy::AscVectorL2Column(name, vector) => { + if let Some(index) = decoder.rd().field_index(name) { + normalized.push(OrderBy::AscVectorL2(index + 1, vector.clone())); + } + // TODO: Error out instead of silently not sorting. + } + }; + } + normalized +} + +pub(super) fn compare_rows( + a: &DataRow, + b: &DataRow, + cols: &[OrderBy], + decoder: &Decoder, +) -> Ordering { + cols.iter() + .filter_map(|col| { + let index = col.index(); + let asc = col.asc(); + let index = index?; + let left = a.get_column(index, decoder); + let right = b.get_column(index, decoder); + + match (left, right) { + (Ok(Some(left)), Ok(Some(right))) => { + // Handle the special vector case. + if let OrderBy::AscVectorL2(_, vector) = col { + let left: Option = left.value.try_into().ok(); + let right: Option = right.value.try_into().ok(); + + if let (Some(left), Some(right)) = (left, right) { + let left = left.distance_l2(vector); + let right = right.distance_l2(vector); + + left.partial_cmp(&right) + } else { + Some(Ordering::Equal) + } + } else { + // FIXME(sage): We don't handle ASC NULLS FIRST or + // DESC NULLS LAST we should either error or add + // support rather than silently do the wrong sorting + match (&left.value, &right.value, asc) { + (Datum::Null, Datum::Null, _) => Some(Ordering::Equal), + (Datum::Null, _, true) => Some(Ordering::Greater), + (_, Datum::Null, true) => Some(Ordering::Less), + (Datum::Null, _, false) => Some(Ordering::Less), + (_, Datum::Null, false) => Some(Ordering::Greater), + (a, b, true) => a.partial_cmp(b), + (a, b, false) => b.partial_cmp(a), + } + } + } + + _ => Some(Ordering::Equal), + } + }) + .reduce(Ordering::then) + .unwrap_or(Ordering::Equal) +} + +#[cfg(test)] +mod test { + use super::*; + use crate::net::messages::FromBytes; + use crate::net::{Field, Format, RowDescription}; + + #[test] + fn test_order_by_merge_honors_offset_and_limit() { + let rd = RowDescription::new(&[Field::bigint("id")]); + let decoder = Decoder::from(&rd); + let mut merge = OrderByMergeState::new( + 2, + normalize_order_by(&[OrderBy::Asc(1)], &decoder), + &Limit { + limit: Some(2), + offset: Some(1), + }, + ); + + for value in [1_i64, 3_i64] { + let mut row = DataRow::new(); + row.add(value); + merge.push_row(0, row, &decoder); + } + for value in [2_i64, 4_i64] { + let mut row = DataRow::new(); + row.add(value); + merge.push_row(1, row, &decoder); + } + + merge.mark_done(0, &decoder); + merge.mark_done(1, &decoder); + + // Global order is 1,2,3,4 -> offset 1 + limit 2 => 2,3 + let out = DataRow::from_bytes(merge.take_ready().unwrap().to_bytes()).unwrap(); + assert_eq!(out.get::(0, Format::Text).unwrap(), 2_i64); + let out = DataRow::from_bytes(merge.take_ready().unwrap().to_bytes()).unwrap(); + assert_eq!(out.get::(0, Format::Text).unwrap(), 3_i64); + assert!(merge.take_ready().is_none()); + assert_eq!(merge.output_rows(), 2); + } + + #[test] + fn test_order_by_merge_waits_for_all_shards_heads() { + let rd = RowDescription::new(&[Field::bigint("id")]); + let decoder = Decoder::from(&rd); + let mut merge = OrderByMergeState::new( + 2, + normalize_order_by(&[OrderBy::Asc(1)], &decoder), + &Limit { + limit: None, + offset: None, + }, + ); + + let mut row = DataRow::new(); + row.add(1_i64); + merge.push_row(0, row, &decoder); + + // Can't emit yet: shard 1 has not produced a head row or finished. + assert!(merge.take_ready().is_none()); + + let mut row = DataRow::new(); + row.add(2_i64); + merge.push_row(1, row, &decoder); + let out = DataRow::from_bytes(merge.take_ready().unwrap().to_bytes()).unwrap(); + assert_eq!(out.get::(0, Format::Text).unwrap(), 1_i64); + + // Still can't emit row 2: shard 0 is neither done nor has another head row. + merge.mark_done(1, &decoder); + assert!(merge.take_ready().is_none()); + + // Once shard 0 is done, shard 1's remaining head can be emitted. + merge.mark_done(0, &decoder); + let out = DataRow::from_bytes(merge.take_ready().unwrap().to_bytes()).unwrap(); + assert_eq!(out.get::(0, Format::Text).unwrap(), 2_i64); + } +} From 3e00ef44aace59ecae1fb178c0da47f27dcc9dae Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Thu, 2 Jul 2026 21:39:38 +0530 Subject: [PATCH 5/5] add integration test --- integration/rust/tests/integration/mod.rs | 1 + .../integration/order_by_cross_shard_merge.rs | 269 ++++++++++++++++++ .../backend/pool/connection/order_by_merge.rs | 2 +- 3 files changed, 271 insertions(+), 1 deletion(-) create mode 100644 integration/rust/tests/integration/order_by_cross_shard_merge.rs diff --git a/integration/rust/tests/integration/mod.rs b/integration/rust/tests/integration/mod.rs index ce9d2f971..f4940cae9 100644 --- a/integration/rust/tests/integration/mod.rs +++ b/integration/rust/tests/integration/mod.rs @@ -18,6 +18,7 @@ pub mod max; pub mod multi_set; pub mod notify; pub mod offset; +pub mod order_by_cross_shard_merge; pub mod partial_req; pub mod per_stmt_routing; pub mod prepared; diff --git a/integration/rust/tests/integration/order_by_cross_shard_merge.rs b/integration/rust/tests/integration/order_by_cross_shard_merge.rs new file mode 100644 index 000000000..db11102a9 --- /dev/null +++ b/integration/rust/tests/integration/order_by_cross_shard_merge.rs @@ -0,0 +1,269 @@ +use chrono::{DateTime, Duration, NaiveDate, NaiveDateTime, Utc}; +use rand::{Rng, SeedableRng, rngs::StdRng}; +use rust::setup::admin_sqlx; +use rust_decimal::Decimal; +use sqlx::{Executor, Row, postgres::PgPool, postgres::PgPoolOptions}; + +const TABLE: &str = "order_by_cross_shard_merge_test"; + +#[derive(Debug, Clone)] +struct TestRow { + id: i64, + v_text: String, + v_varchar: String, + v_bigint: i64, + v_int: i32, + v_numeric: Decimal, + v_timestamp: NaiveDateTime, + v_timestamptz: DateTime, + v_date: NaiveDate, +} + +#[tokio::test] +async fn cross_shard_order_by_randomized_types() -> Result<(), Box> { + let sharded = sharded_pool().await; + reset(&sharded).await?; + + let mut rng = StdRng::seed_from_u64(0xC0FFEE); + let mut rows = vec![]; + + // Insert deterministic random rows on each shard directly so each shard has data. + for shard in [0, 1] { + for i in 0..200_i64 { + let id = (shard as i64) * 10_000 + i + 1; + let ts = random_timestamp(&mut rng); + let row = TestRow { + id, + v_text: format!("txt_{:04}", rng.random_range(0..250)), + v_varchar: format!("varchar_{}_{}", rng.random_range(0..100), i % 11), + v_bigint: rng.random_range(-5_000_000..5_000_000), + v_int: rng.random_range(-100_000..100_000), + v_numeric: Decimal::new(rng.random_range(-999_999_999..999_999_999), 3), + v_timestamp: ts.naive_utc(), + v_timestamptz: ts, + v_date: ts.date_naive(), + }; + + insert_row(&sharded, shard, &row).await?; + rows.push(row); + } + } + + let total_rows = sharded + .fetch_one(format!("SELECT COUNT(*)::bigint AS n FROM {TABLE}").as_str()) + .await? + .get::("n"); + assert_eq!(total_rows, rows.len() as i64); + + assert_sorted(&sharded, &rows, "ORDER BY v_text ASC, id ASC", |a, b| { + a.v_text.cmp(&b.v_text).then(a.id.cmp(&b.id)) + }) + .await?; + assert_sorted(&sharded, &rows, "ORDER BY v_text DESC, id DESC", |a, b| { + b.v_text.cmp(&a.v_text).then(b.id.cmp(&a.id)) + }) + .await?; + + assert_sorted(&sharded, &rows, "ORDER BY v_varchar ASC, id ASC", |a, b| { + a.v_varchar.cmp(&b.v_varchar).then(a.id.cmp(&b.id)) + }) + .await?; + assert_sorted( + &sharded, + &rows, + "ORDER BY v_varchar DESC, id DESC", + |a, b| b.v_varchar.cmp(&a.v_varchar).then(b.id.cmp(&a.id)), + ) + .await?; + + assert_sorted(&sharded, &rows, "ORDER BY v_bigint ASC, id ASC", |a, b| { + a.v_bigint.cmp(&b.v_bigint).then(a.id.cmp(&b.id)) + }) + .await?; + assert_sorted( + &sharded, + &rows, + "ORDER BY v_bigint DESC, id DESC", + |a, b| b.v_bigint.cmp(&a.v_bigint).then(b.id.cmp(&a.id)), + ) + .await?; + + assert_sorted(&sharded, &rows, "ORDER BY v_int ASC, id ASC", |a, b| { + a.v_int.cmp(&b.v_int).then(a.id.cmp(&b.id)) + }) + .await?; + assert_sorted(&sharded, &rows, "ORDER BY v_int DESC, id DESC", |a, b| { + b.v_int.cmp(&a.v_int).then(b.id.cmp(&a.id)) + }) + .await?; + + assert_sorted(&sharded, &rows, "ORDER BY v_numeric ASC, id ASC", |a, b| { + a.v_numeric.cmp(&b.v_numeric).then(a.id.cmp(&b.id)) + }) + .await?; + assert_sorted( + &sharded, + &rows, + "ORDER BY v_numeric DESC, id DESC", + |a, b| b.v_numeric.cmp(&a.v_numeric).then(b.id.cmp(&a.id)), + ) + .await?; + + assert_sorted( + &sharded, + &rows, + "ORDER BY v_timestamp ASC, id ASC", + |a, b| a.v_timestamp.cmp(&b.v_timestamp).then(a.id.cmp(&b.id)), + ) + .await?; + assert_sorted( + &sharded, + &rows, + "ORDER BY v_timestamp DESC, id DESC", + |a, b| b.v_timestamp.cmp(&a.v_timestamp).then(b.id.cmp(&a.id)), + ) + .await?; + + assert_sorted( + &sharded, + &rows, + "ORDER BY v_timestamptz ASC, id ASC", + |a, b| a.v_timestamptz.cmp(&b.v_timestamptz).then(a.id.cmp(&b.id)), + ) + .await?; + assert_sorted( + &sharded, + &rows, + "ORDER BY v_timestamptz DESC, id DESC", + |a, b| b.v_timestamptz.cmp(&a.v_timestamptz).then(b.id.cmp(&a.id)), + ) + .await?; + + assert_sorted(&sharded, &rows, "ORDER BY v_date ASC, id ASC", |a, b| { + a.v_date.cmp(&b.v_date).then(a.id.cmp(&b.id)) + }) + .await?; + assert_sorted(&sharded, &rows, "ORDER BY v_date DESC, id DESC", |a, b| { + b.v_date.cmp(&a.v_date).then(b.id.cmp(&a.id)) + }) + .await?; + + cleanup(&sharded).await; + sharded.close().await; + Ok(()) +} + +async fn assert_sorted( + pool: &PgPool, + rows: &[TestRow], + order_by: &str, + cmp: impl Fn(&TestRow, &TestRow) -> std::cmp::Ordering, +) -> Result<(), sqlx::Error> { + // Include all sortable columns in the projection so ORDER BY keys are present + // in RowDescription for PgDog's current cross-shard sorting/merge logic. + let actual_ids: Vec = pool + .fetch_all( + format!( + "SELECT id, v_text, v_varchar, v_bigint, v_int, v_numeric, v_timestamp, v_timestamptz, v_date FROM {TABLE} {order_by}" + ) + .as_str(), + ) + .await? + .into_iter() + .map(|r| r.get::("id")) + .collect(); + + let mut expected = rows.to_vec(); + expected.sort_by(cmp); + let expected_ids: Vec = expected.into_iter().map(|r| r.id).collect(); + + assert_eq!( + actual_ids, expected_ids, + "ORDER BY mismatch for query: SELECT id FROM {TABLE} {order_by}" + ); + Ok(()) +} + +fn random_timestamp(rng: &mut StdRng) -> DateTime { + let base = NaiveDate::from_ymd_opt(2025, 1, 1) + .unwrap() + .and_hms_micro_opt(0, 0, 0, 0) + .unwrap() + .and_utc(); + base + Duration::seconds(rng.random_range(-5_000_000..5_000_000)) + + Duration::microseconds(rng.random_range(0..1_000_000)) +} + +async fn insert_row(pool: &PgPool, shard: usize, row: &TestRow) -> Result<(), sqlx::Error> { + sqlx::query( + format!( + "/* pgdog_shard: {shard} */ \ + INSERT INTO {TABLE} \ + (id, v_text, v_varchar, v_bigint, v_int, v_numeric, v_timestamp, v_timestamptz, v_date, customer_id) \ + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)" + ) + .as_str(), + ) + .bind(row.id) + .bind(&row.v_text) + .bind(&row.v_varchar) + .bind(row.v_bigint) + .bind(row.v_int) + .bind(row.v_numeric) + .bind(row.v_timestamp) + .bind(row.v_timestamptz) + .bind(row.v_date) + .bind(row.id) + .execute(pool) + .await?; + Ok(()) +} + +async fn reset(pool: &PgPool) -> Result<(), sqlx::Error> { + for shard in [0, 1] { + pool.execute(format!("/* pgdog_shard: {shard} */ DROP TABLE IF EXISTS {TABLE}").as_str()) + .await + .ok(); + } + for shard in [0, 1] { + pool.execute( + format!( + "/* pgdog_shard: {shard} */ \ + CREATE TABLE {TABLE} (\ + id BIGINT PRIMARY KEY, \ + v_text TEXT COLLATE \"C\" NOT NULL, \ + v_varchar VARCHAR(200) COLLATE \"C\" NOT NULL, \ + v_bigint BIGINT NOT NULL, \ + v_int INTEGER NOT NULL, \ + v_numeric NUMERIC NOT NULL, \ + v_timestamp TIMESTAMP NOT NULL, \ + v_timestamptz TIMESTAMPTZ NOT NULL, \ + v_date DATE NOT NULL, \ + customer_id BIGINT\ + )" + ) + .as_str(), + ) + .await?; + } + admin_sqlx().await.execute("RELOAD").await?; + Ok(()) +} + +async fn cleanup(pool: &PgPool) { + for shard in [0, 1] { + pool.execute(format!("/* pgdog_shard: {shard} */ DROP TABLE IF EXISTS {TABLE}").as_str()) + .await + .ok(); + } +} + +async fn sharded_pool() -> PgPool { + PgPoolOptions::new() + .max_connections(1) + .connect( + "postgres://pgdog:pgdog@127.0.0.1:6432/pgdog_sharded?application_name=sqlx_order_by", + ) + .await + .unwrap() +} diff --git a/pgdog/src/backend/pool/connection/order_by_merge.rs b/pgdog/src/backend/pool/connection/order_by_merge.rs index 4d6ae576d..1ec80fa36 100644 --- a/pgdog/src/backend/pool/connection/order_by_merge.rs +++ b/pgdog/src/backend/pool/connection/order_by_merge.rs @@ -221,7 +221,7 @@ pub(super) fn compare_rows( #[cfg(test)] mod test { use super::*; - use crate::net::messages::FromBytes; + use crate::net::messages::{FromBytes, ToBytes}; use crate::net::{Field, Format, RowDescription}; #[test]