1pub mod proto;
2pub mod prune;
3pub mod server;
4
5mod aggregate;
6mod builder;
7mod codec;
8mod diagnostics;
9mod filter;
10mod predicate;
11mod scan;
12mod schema;
13mod types;
14mod writer;
15
16pub use schema::KvSchema;
17pub use server::{sql_connect_stack, SqlConnect, SqlServer};
18pub use types::default_orders_index_specs;
19pub use types::{
20 CellValue, IndexBackfillEvent, IndexBackfillOptions, IndexBackfillReport, IndexLayout,
21 IndexSpec, TableColumnConfig,
22};
23pub use writer::{BatchReceipt, BatchWriter, PreparedBatch, TableWriter};
24
25#[cfg(test)]
26mod tests {
27 #![allow(refining_impl_trait)]
28
29 use super::aggregate::*;
30 use super::builder::*;
31 use super::codec::*;
32 use super::diagnostics::*;
33 use super::filter::*;
34 use super::predicate::*;
35 use super::scan::*;
36 use super::types::*;
37 use super::writer::*;
38 use super::*;
39 use commonware_codec::Encode;
40 use datafusion::arrow::array::{
41 BinaryViewArray, Float64Array, Int64Array, LargeBinaryArray, LargeStringArray,
42 StringViewArray,
43 };
44 use datafusion::arrow::datatypes::{i256, DataType, TimeUnit};
45 use datafusion::arrow::record_batch::RecordBatch;
46 use datafusion::common::{config::ConfigOptions, ScalarValue};
47 use datafusion::logical_expr::{Expr, Operator};
48 use datafusion::physical_optimizer::limit_pushdown::LimitPushdown;
49 use datafusion::physical_optimizer::PhysicalOptimizerRule;
50 use datafusion::physical_plan::limit::GlobalLimitExec;
51 use datafusion::physical_plan::ExecutionPlan;
52 use datafusion::prelude::SessionContext;
53 use exoware_sdk::keys::{Key, Prefix};
54 use exoware_sdk::kv_codec::{
55 canonicalize_reduced_group_values, decode_stored_row, encode_reduced_group_key,
56 eval_predicate, KvReducedValue, StoredRow,
57 };
58 use exoware_sdk::{
59 PrefixedStoreClient, RangeReduceOp, RangeReduceRequest, StoreBatchUpload, StoreClient,
60 };
61 use std::collections::{BTreeMap, HashSet};
62 use std::ops::Bound::{Included, Unbounded};
63 use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering as AtomicOrdering};
64 use std::sync::{Arc, Mutex};
65 use std::time::Duration;
66
67 use axum::Router;
68 use bytes::Bytes;
69 use connectrpc::{Chain, ConnectError, ConnectRpcService, RequestContext as Context};
70 use exoware_sdk::common::kv::v1::Entry as ProtoEntry;
71 use exoware_sdk::connect_compression_registry;
72 use exoware_sdk::kv_codec::{eval_expr, expr_needs_value};
73 use exoware_sdk::log::ingest::v1::{
74 PutResponse as ProtoPutResponse, Service as IngestService,
75 ServiceServer as IngestServiceServer,
76 };
77 use exoware_sdk::store::query::v1::{
78 GetManyEntry as ProtoGetManyEntry, GetManyFrame as ProtoGetManyFrame,
79 GetResponse as ProtoGetResponse, RangeFrame as ProtoRangeFrame,
80 ReduceResponse as ProtoReduceResponse, Service as QueryService,
81 ServiceServer as QueryServiceServer,
82 };
83 use exoware_sdk::RangeMode;
84 use exoware_sdk::{
85 parse_range_traversal_direction, to_domain_reduce_request, to_proto_optional_reduced_value,
86 to_proto_reduced_value, RangeTraversalDirection, RangeTraversalModeError,
87 };
88 use exoware_sdk::{RangeReduceGroup, RangeReduceResponse, RangeReduceResult};
89 use futures::{stream, TryStreamExt};
90 use tokio::sync::{mpsc, oneshot, Notify};
91
92 fn assert_explain_includes_query_stats_surface(
94 explain: &str,
95 surface: QueryStatsExplainSurface,
96 ) {
97 let expected = format!("query_stats={}", format_query_stats_explain(surface));
98 assert!(
99 explain.contains(&expected),
100 "expected EXPLAIN output to include `{expected}`\n{explain}"
101 );
102 }
103
104 fn simple_int64_model(prefix: u8) -> TableModel {
105 let config = KvTableConfig::new(
106 prefix,
107 vec![TableColumnConfig::new("id", DataType::Int64, false)],
108 vec!["id".to_string()],
109 vec![],
110 )
111 .unwrap();
112 TableModel::from_config(&config).unwrap()
113 }
114
115 fn key_payload(prefix: &Prefix, key: &Key, offset: usize, len: usize) -> Vec<u8> {
116 let payload = prefix.strip(key).expect("key payload");
117 payload[offset..offset + len].to_vec()
118 }
119
120 fn primary_payload(model: &TableModel, key: &Key, offset: usize, len: usize) -> Vec<u8> {
121 key_payload(&model.primary_key_prefix, key, offset, len)
122 }
123
124 fn index_payload(spec: &ResolvedIndexSpec, key: &Key, offset: usize, len: usize) -> Vec<u8> {
125 key_payload(&spec.prefix, key, offset, len)
126 }
127
128 fn matches_primary_key(table_prefix: u8, key: &Key) -> bool {
129 primary_key_prefix(table_prefix)
130 .expect("primary prefix")
131 .matches(key)
132 }
133
134 fn matches_secondary_index_key(table_prefix: u8, index_id: u8, key: &Key) -> bool {
135 secondary_index_prefix(table_prefix, index_id)
136 .expect("secondary prefix")
137 .matches(key)
138 }
139
140 fn test_model() -> (TableModel, Vec<ResolvedIndexSpec>) {
141 let config = KvTableConfig::new(
142 0,
143 vec![
144 TableColumnConfig::new("region", DataType::Utf8, false),
145 TableColumnConfig::new("customer_id", DataType::Int64, false),
146 TableColumnConfig::new("order_id", DataType::Int64, false),
147 TableColumnConfig::new("amount_cents", DataType::Int64, false),
148 TableColumnConfig::new("status", DataType::Utf8, false),
149 ],
150 vec!["order_id".to_string()],
151 vec![
152 IndexSpec::new(
153 "region_customer",
154 vec!["region".to_string(), "customer_id".to_string()],
155 )
156 .expect("valid"),
157 IndexSpec::new(
158 "status_customer",
159 vec!["status".to_string(), "customer_id".to_string()],
160 )
161 .expect("valid"),
162 ],
163 )
164 .expect("valid config");
165 let model = TableModel::from_config(&config).expect("model");
166 let specs = model
167 .resolve_index_specs(&config.index_specs)
168 .expect("specs");
169 (model, specs)
170 }
171
172 fn zorder_test_model() -> (TableModel, Vec<ResolvedIndexSpec>) {
173 let config = KvTableConfig::new(
174 0,
175 vec![
176 TableColumnConfig::new("x", DataType::Int64, false),
177 TableColumnConfig::new("y", DataType::Int64, false),
178 TableColumnConfig::new("id", DataType::Int64, false),
179 TableColumnConfig::new("value", DataType::Int64, false),
180 ],
181 vec!["id".to_string()],
182 vec![
183 IndexSpec::new("xy_lex", vec!["x".to_string(), "y".to_string()])
184 .expect("valid")
185 .with_cover_columns(vec!["value".to_string()]),
186 IndexSpec::z_order("xy_z", vec!["x".to_string(), "y".to_string()])
187 .expect("valid")
188 .with_cover_columns(vec!["value".to_string()]),
189 ],
190 )
191 .expect("valid config");
192 let model = TableModel::from_config(&config).expect("model");
193 let specs = model
194 .resolve_index_specs(&config.index_specs)
195 .expect("specs");
196 (model, specs)
197 }
198
199 #[derive(Clone)]
200 struct MockState {
201 kv: Arc<Mutex<BTreeMap<Key, Bytes>>>,
202 range_calls: Arc<AtomicUsize>,
203 range_reduce_calls: Arc<AtomicUsize>,
204 sequence_number: Arc<AtomicU64>,
205 }
206
207 #[derive(Debug)]
208 struct MockGroupedReduceState {
209 group_values: Vec<Option<KvReducedValue>>,
210 states: Vec<PartialAggregateState>,
211 }
212
213 type MockReduceRow = (Vec<Option<KvReducedValue>>, Vec<Option<KvReducedValue>>);
214
215 fn extract_mock_reduce_row(
216 key: &Key,
217 value: &Bytes,
218 request: &RangeReduceRequest,
219 ) -> Option<MockReduceRow> {
220 let needs_value = request
221 .group_by
222 .iter()
223 .chain(
224 request
225 .reducers
226 .iter()
227 .filter_map(|reducer| reducer.expr.as_ref()),
228 )
229 .any(expr_needs_value)
230 || request
231 .filter
232 .as_ref()
233 .is_some_and(exoware_sdk::kv_codec::predicate_needs_value);
234 let archived = if needs_value {
235 decode_stored_row(value.as_ref()).ok()
236 } else {
237 None
238 };
239
240 if let Some(filter) = &request.filter {
241 if !eval_predicate(key, archived.as_ref(), filter).ok()? {
242 return None;
243 }
244 }
245
246 let mut group_values = Vec::with_capacity(request.group_by.len());
247 for expr in &request.group_by {
248 let extracted_value = eval_expr(key, archived.as_ref(), expr).ok()?;
249 group_values.push(extracted_value);
250 }
251 canonicalize_reduced_group_values(&mut group_values);
252
253 let mut reducer_values = Vec::with_capacity(request.reducers.len());
254 for reducer in &request.reducers {
255 let extracted_value = match (&reducer.expr, archived.as_ref()) {
256 (None, _) => None,
257 (Some(expr), _) => eval_expr(key, archived.as_ref(), expr).ok()?,
258 };
259 reducer_values.push(extracted_value);
260 }
261
262 Some((group_values, reducer_values))
263 }
264
265 #[allow(clippy::result_large_err)]
266 fn ensure_min_sequence_number(
267 token: &Arc<AtomicU64>,
268 required: Option<u64>,
269 ) -> Result<(), ConnectError> {
270 let current = token.load(AtomicOrdering::Relaxed);
271 if let Some(required) = required {
272 if current < required {
273 return Err(ConnectError::aborted(format!(
274 "consistency_not_ready: required={required}, current={current}"
275 )));
276 }
277 }
278 Ok(())
279 }
280
281 fn proto_range_entries_frame(results: Vec<(Key, Vec<u8>)>) -> ProtoRangeFrame {
282 ProtoRangeFrame {
283 results: results
284 .into_iter()
285 .map(|(key, value)| ProtoEntry {
286 key: key.to_vec(),
287 value: value.into(),
288 ..Default::default()
289 })
290 .collect(),
291 detail: Some(query_detail(7)).into(),
292 ..Default::default()
293 }
294 }
295
296 fn query_detail(sequence_number: u64) -> exoware_sdk::store::query::v1::Detail {
297 exoware_sdk::store::query::v1::Detail {
298 sequence_number,
299 extra: Default::default(),
300 ..Default::default()
301 }
302 }
303
304 fn final_range_detail_frame(sequence_number: u64) -> ProtoRangeFrame {
305 ProtoRangeFrame {
306 detail: Some(query_detail(sequence_number)).into(),
307 ..Default::default()
308 }
309 }
310
311 fn final_get_many_detail_frame(sequence_number: u64) -> ProtoGetManyFrame {
312 ProtoGetManyFrame {
313 detail: Some(query_detail(sequence_number)).into(),
314 ..Default::default()
315 }
316 }
317
318 #[derive(Clone)]
319 struct MockIngestConnect {
320 state: MockState,
321 }
322
323 impl IngestService for MockIngestConnect {
324 async fn put(
325 &self,
326 _ctx: Context,
327 request: buffa::view::OwnedView<exoware_sdk::log::ingest::v1::PutRequestView<'static>>,
328 ) -> connectrpc::ServiceResult<ProtoPutResponse> {
329 let mut parsed = Vec::<(Key, Bytes)>::new();
330 let wire = request.bytes();
331 for kv in request.kvs.iter() {
332 parsed.push((wire.slice_ref(kv.key), wire.slice_ref(kv.value)));
333 }
334 let mut guard = self.state.kv.lock().expect("kv mutex poisoned");
335 for (key, value) in parsed.iter() {
336 guard.insert(key.clone(), value.clone());
337 }
338 let seq = self
339 .state
340 .sequence_number
341 .fetch_add(1, AtomicOrdering::SeqCst)
342 + 1;
343 connectrpc::Response::ok(ProtoPutResponse {
344 sequence_number: seq,
345 ..Default::default()
346 })
347 }
348 }
349
350 #[derive(Clone)]
351 struct MockQueryConnect {
352 state: MockState,
353 }
354
355 impl QueryService for MockQueryConnect {
356 async fn get(
357 &self,
358 _ctx: Context,
359 request: buffa::view::OwnedView<exoware_sdk::store::query::v1::GetRequestView<'static>>,
360 ) -> connectrpc::ServiceResult<ProtoGetResponse> {
361 ensure_min_sequence_number(&self.state.sequence_number, request.min_sequence_number)?;
362 let key: Key = request.bytes().slice_ref(request.key);
363 let guard = self.state.kv.lock().expect("kv mutex poisoned");
364 let value = guard.get(&key).cloned();
365 let token = self.state.sequence_number.load(AtomicOrdering::Relaxed);
366 connectrpc::Response::ok(ProtoGetResponse {
367 value,
368 detail: Some(query_detail(token)).into(),
369 ..Default::default()
370 })
371 }
372
373 async fn range(
374 &self,
375 _ctx: Context,
376 request: buffa::view::OwnedView<
377 exoware_sdk::store::query::v1::RangeRequestView<'static>,
378 >,
379 ) -> connectrpc::ServiceResult<connectrpc::ServiceStream<ProtoRangeFrame>> {
380 ensure_min_sequence_number(&self.state.sequence_number, request.min_sequence_number)?;
381 self.state.range_calls.fetch_add(1, AtomicOrdering::SeqCst);
382
383 let wire = request.bytes();
384 let start_key: Key = wire.slice_ref(request.start);
385 let end_key: Key = wire.slice_ref(request.end);
386 let limit = request.limit.map(|v| v as usize).unwrap_or(usize::MAX);
387 let batch_size = usize::try_from(request.batch_size).unwrap_or(usize::MAX);
388 if batch_size == 0 {
389 return Err(ConnectError::invalid_argument(
390 "invalid batch_size: expected positive integer",
391 ));
392 }
393
394 let mode = match parse_range_traversal_direction(request.mode) {
395 Ok(RangeTraversalDirection::Forward) => RangeMode::Forward,
396 Ok(RangeTraversalDirection::Reverse) => RangeMode::Reverse,
397 Err(RangeTraversalModeError::UnknownWireValue(v)) => {
398 return Err(ConnectError::invalid_argument(format!(
399 "unknown TraversalMode enum value {v}"
400 )));
401 }
402 };
403
404 let state = self.state.clone();
405 let guard = state.kv.lock().expect("kv mutex poisoned");
406 let range: (std::ops::Bound<&Key>, std::ops::Bound<&Key>) = (
408 Included(&start_key),
409 if end_key.is_empty() {
410 Unbounded
411 } else {
412 Included(&end_key)
413 },
414 );
415 let range_iter = guard.range::<Key, _>(range);
416 let iter: Box<dyn Iterator<Item = (&Key, &Bytes)> + Send> = match mode {
417 RangeMode::Forward => Box::new(range_iter),
418 RangeMode::Reverse => Box::new(range_iter.rev()),
419 };
420 let mut results: Vec<ProtoEntry> = Vec::new();
421 for (key, value) in iter.take(limit) {
422 results.push(ProtoEntry {
423 key: key.to_vec(),
424 value: value.clone(),
425 ..Default::default()
426 });
427 }
428 drop(guard);
429 let token = state.sequence_number.load(AtomicOrdering::Relaxed);
430 let batch = batch_size.max(1);
431 let mut frames: Vec<Result<ProtoRangeFrame, ConnectError>> = Vec::new();
432 let mut emitted_frame = false;
433 for chunk in results.chunks(batch) {
434 frames.push(Ok(ProtoRangeFrame {
435 results: chunk.to_vec(),
436 detail: Some(query_detail(token)).into(),
437 ..Default::default()
438 }));
439 emitted_frame = true;
440 }
441 if !emitted_frame {
442 frames.push(Ok(final_range_detail_frame(token)));
443 }
444 Ok(connectrpc::Response::stream(stream::iter(frames)))
445 }
446
447 async fn get_many(
448 &self,
449 _ctx: Context,
450 request: buffa::view::OwnedView<
451 exoware_sdk::store::query::v1::GetManyRequestView<'static>,
452 >,
453 ) -> connectrpc::ServiceResult<connectrpc::ServiceStream<ProtoGetManyFrame>> {
454 ensure_min_sequence_number(&self.state.sequence_number, request.min_sequence_number)?;
455 let batch_size = usize::try_from(request.batch_size)
456 .unwrap_or(usize::MAX)
457 .max(1);
458 let guard = self.state.kv.lock().expect("kv mutex poisoned");
459 let mut entries: Vec<ProtoGetManyEntry> = Vec::new();
460 let wire = request.bytes();
461 for key_bytes in request.keys.iter() {
462 let key: Key = wire.slice_ref(key_bytes);
463 let value = guard.get(&key).cloned();
464 entries.push(ProtoGetManyEntry {
465 key: key.to_vec(),
466 value,
467 ..Default::default()
468 });
469 }
470 drop(guard);
471 let token = self.state.sequence_number.load(AtomicOrdering::Relaxed);
472 let mut frames: Vec<Result<ProtoGetManyFrame, ConnectError>> = Vec::new();
473 let mut emitted_frame = false;
474 for chunk in entries.chunks(batch_size) {
475 frames.push(Ok(ProtoGetManyFrame {
476 results: chunk.to_vec(),
477 detail: Some(query_detail(token)).into(),
478 ..Default::default()
479 }));
480 emitted_frame = true;
481 }
482 if !emitted_frame {
483 frames.push(Ok(final_get_many_detail_frame(token)));
484 }
485 Ok(connectrpc::Response::stream(stream::iter(frames)))
486 }
487
488 async fn reduce(
489 &self,
490 _ctx: Context,
491 request: buffa::view::OwnedView<
492 exoware_sdk::store::query::v1::ReduceRequestView<'static>,
493 >,
494 ) -> connectrpc::ServiceResult<ProtoReduceResponse> {
495 ensure_min_sequence_number(&self.state.sequence_number, request.min_sequence_number)?;
496 self.state
497 .range_reduce_calls
498 .fetch_add(1, AtomicOrdering::SeqCst);
499 let owned = request.to_owned_message();
500 let start_key: Key = owned.start.clone().into();
501 let end_key: Key = owned.end.clone().into();
502 let reduce_req = owned
503 .params
504 .as_option()
505 .ok_or_else(|| ConnectError::invalid_argument("missing range reduce params"))?;
506 let domain_request =
507 to_domain_reduce_request(reduce_req).map_err(ConnectError::invalid_argument)?;
508
509 let state = self.state.clone();
510 let guard = state.kv.lock().expect("kv mutex poisoned");
511 let mut states = domain_request.group_by.is_empty().then(|| {
512 domain_request
513 .reducers
514 .iter()
515 .map(|reducer| PartialAggregateState::from_op(reducer.op))
516 .collect::<Vec<_>>()
517 });
518 let mut grouped = BTreeMap::<Vec<u8>, MockGroupedReduceState>::new();
519
520 let range: (std::ops::Bound<&Key>, std::ops::Bound<&Key>) = (
521 Included(&start_key),
522 if end_key.is_empty() {
523 Unbounded
524 } else {
525 Included(&end_key)
526 },
527 );
528 for (key, value) in guard.range::<Key, _>(range) {
529 let Some((group_values, reducer_values)) =
530 extract_mock_reduce_row(key, value, &domain_request)
531 else {
532 continue;
533 };
534 if domain_request.group_by.is_empty() {
535 let states = states.as_mut().expect("scalar states");
536 for ((state, reducer), value) in states
537 .iter_mut()
538 .zip(domain_request.reducers.iter())
539 .zip(reducer_values)
540 {
541 match reducer.op {
542 RangeReduceOp::CountAll => state
543 .merge_partial(reducer.op, Some(&KvReducedValue::UInt64(1)))
544 .map_err(|e| ConnectError::internal(e.to_string()))?,
545 RangeReduceOp::CountField => {
546 let partial =
547 KvReducedValue::UInt64(if value.is_some() { 1 } else { 0 });
548 state
549 .merge_partial(reducer.op, Some(&partial))
550 .map_err(|e| ConnectError::internal(e.to_string()))?
551 }
552 _ => state
553 .merge_partial(reducer.op, value.as_ref())
554 .map_err(|e| ConnectError::internal(e.to_string()))?,
555 }
556 }
557 } else {
558 let group_key = encode_reduced_group_key(&group_values);
559 let group =
560 grouped
561 .entry(group_key)
562 .or_insert_with(|| MockGroupedReduceState {
563 group_values: group_values.clone(),
564 states: domain_request
565 .reducers
566 .iter()
567 .map(|reducer| PartialAggregateState::from_op(reducer.op))
568 .collect(),
569 });
570 for ((state, reducer), value) in group
571 .states
572 .iter_mut()
573 .zip(domain_request.reducers.iter())
574 .zip(reducer_values)
575 {
576 match reducer.op {
577 RangeReduceOp::CountAll => state
578 .merge_partial(reducer.op, Some(&KvReducedValue::UInt64(1)))
579 .map_err(|e| ConnectError::internal(e.to_string()))?,
580 RangeReduceOp::CountField => {
581 let partial =
582 KvReducedValue::UInt64(if value.is_some() { 1 } else { 0 });
583 state
584 .merge_partial(reducer.op, Some(&partial))
585 .map_err(|e| ConnectError::internal(e.to_string()))?
586 }
587 _ => state
588 .merge_partial(reducer.op, value.as_ref())
589 .map_err(|e| ConnectError::internal(e.to_string()))?,
590 }
591 }
592 }
593 }
594
595 let response = if let Some(states) = states {
596 RangeReduceResponse {
597 results: states
598 .iter()
599 .map(|state| RangeReduceResult {
600 value: match state {
601 PartialAggregateState::Count(count) => {
602 Some(KvReducedValue::UInt64(*count))
603 }
604 PartialAggregateState::Sum(value)
605 | PartialAggregateState::Min(value)
606 | PartialAggregateState::Max(value) => value.clone(),
607 },
608 })
609 .collect(),
610 groups: Vec::new(),
611 }
612 } else {
613 RangeReduceResponse {
614 results: Vec::new(),
615 groups: grouped
616 .into_values()
617 .map(|group| RangeReduceGroup {
618 group_values: group.group_values,
619 results: group
620 .states
621 .into_iter()
622 .map(|state| RangeReduceResult {
623 value: match state {
624 PartialAggregateState::Count(count) => {
625 Some(KvReducedValue::UInt64(count))
626 }
627 PartialAggregateState::Sum(value)
628 | PartialAggregateState::Min(value)
629 | PartialAggregateState::Max(value) => value,
630 },
631 })
632 .collect(),
633 })
634 .collect(),
635 }
636 };
637 drop(guard);
638 let token = state.sequence_number.load(AtomicOrdering::Relaxed);
639 connectrpc::Response::ok(ProtoReduceResponse {
640 results: response
641 .results
642 .into_iter()
643 .map(|result| exoware_sdk::store::query::v1::RangeReduceResult {
644 value: result.value.map(to_proto_reduced_value).into(),
645 ..Default::default()
646 })
647 .collect(),
648 groups: response
649 .groups
650 .into_iter()
651 .map(|group| {
652 let group_values_present: Vec<bool> =
653 group.group_values.iter().map(|v| v.is_some()).collect();
654 exoware_sdk::store::query::v1::RangeReduceGroup {
655 group_values: group
656 .group_values
657 .into_iter()
658 .map(to_proto_optional_reduced_value)
659 .collect(),
660 group_values_present,
661 results: group
662 .results
663 .into_iter()
664 .map(|result| exoware_sdk::store::query::v1::RangeReduceResult {
665 value: result.value.map(to_proto_reduced_value).into(),
666 ..Default::default()
667 })
668 .collect(),
669 ..Default::default()
670 }
671 })
672 .collect(),
673 detail: Some(query_detail(token)).into(),
674 ..Default::default()
675 })
676 }
677 }
678
679 async fn spawn_mock_server(state: MockState) -> (String, oneshot::Sender<()>) {
680 let connect = ConnectRpcService::new(Chain(
681 IngestServiceServer::new(MockIngestConnect {
682 state: state.clone(),
683 }),
684 QueryServiceServer::new(MockQueryConnect { state }),
685 ))
686 .with_compression(connect_compression_registry());
687 let app = Router::new().fallback_service(connect);
688
689 let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
690 .await
691 .expect("bind mock server");
692 let addr = listener.local_addr().expect("local addr");
693 let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
694 tokio::spawn(async move {
695 axum::serve(listener, app)
696 .with_graceful_shutdown(async move {
697 let _ = shutdown_rx.await;
698 })
699 .await
700 .expect("mock server should run");
701 });
702 (format!("http://{addr}"), shutdown_tx)
703 }
704
705 fn assert_count_scalar(batch: &RecordBatch, col_idx: usize, row_idx: usize, expected: u64) {
706 let scalar = ScalarValue::try_from_array(batch.column(col_idx), row_idx)
707 .expect("count scalar should decode");
708 match scalar {
709 ScalarValue::UInt64(Some(value)) => assert_eq!(value, expected),
710 ScalarValue::Int64(Some(value)) => assert_eq!(value, expected as i64),
711 other => panic!("unexpected count scalar: {other:?}"),
712 }
713 }
714
715 fn collect_i64_column(batches: &[RecordBatch], col_idx: usize) -> Vec<i64> {
716 let mut values = Vec::new();
717 for batch in batches {
718 for row_idx in 0..batch.num_rows() {
719 match ScalarValue::try_from_array(batch.column(col_idx), row_idx)
720 .expect("int64 scalar should decode")
721 {
722 ScalarValue::Int64(Some(value)) => values.push(value),
723 other => panic!("unexpected int64 scalar: {other:?}"),
724 }
725 }
726 }
727 values
728 }
729
730 async fn explain_plan_rows(ctx: &SessionContext, sql: &str) -> Vec<(String, String)> {
731 let batches = ctx
732 .sql(&format!("EXPLAIN {sql}"))
733 .await
734 .expect("explain query")
735 .collect()
736 .await
737 .expect("explain collect");
738 let mut rows = Vec::new();
739 for batch in batches {
740 for row_idx in 0..batch.num_rows() {
741 let plan_type = scalar_to_string(
742 &ScalarValue::try_from_array(batch.column(0), row_idx).expect("plan type"),
743 )
744 .expect("plan type string");
745 let plan = scalar_to_string(
746 &ScalarValue::try_from_array(batch.column(1), row_idx).expect("plan"),
747 )
748 .expect("plan string");
749 rows.push((plan_type, plan));
750 }
751 }
752 rows
753 }
754
755 fn physical_plan_text(rows: &[(String, String)]) -> String {
756 rows.iter()
757 .filter(|(plan_type, _)| plan_type.contains("physical_plan"))
758 .map(|(_, plan)| plan.as_str())
759 .collect::<Vec<_>>()
760 .join("\n")
761 }
762
763 #[tokio::test]
764 async fn explain_reports_full_scan_like_primary_key_scan() {
765 let state = MockState {
766 kv: Arc::new(Mutex::new(BTreeMap::new())),
767 range_calls: Arc::new(AtomicUsize::new(0)),
768 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
769 sequence_number: Arc::new(AtomicU64::new(0)),
770 };
771 let (base_url, shutdown_tx) = spawn_mock_server(state).await;
772 let client = StoreClient::new(&base_url);
773
774 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
775 .table(
776 "orders",
777 vec![
778 TableColumnConfig::new("id", DataType::Int64, false),
779 TableColumnConfig::new("status", DataType::Utf8, false),
780 TableColumnConfig::new("amount_cents", DataType::Int64, false),
781 ],
782 vec!["id".to_string()],
783 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
784 .expect("valid")
785 .with_cover_columns(vec!["amount_cents".to_string()])],
786 )
787 .expect("schema");
788 let ctx = SessionContext::new();
789 schema.register_all(&ctx).expect("register");
790
791 let explain =
792 physical_plan_text(&explain_plan_rows(&ctx, "SELECT id, status FROM orders").await);
793 assert!(explain.contains("KvScanExec:"));
794 assert!(explain.contains("mode=primary_key"));
795 assert!(explain.contains("predicate=<none>"));
796 assert!(explain.contains("row_recheck=false"));
797 assert!(explain.contains("full_scan_like=true"));
798 assert_explain_includes_query_stats_surface(
799 &explain,
800 QueryStatsExplainSurface::StreamedRangeDetail,
801 );
802
803 let _ = shutdown_tx.send(());
804 }
805
806 #[tokio::test]
807 async fn explain_reports_secondary_index_scan_and_row_recheck() {
808 let state = MockState {
809 kv: Arc::new(Mutex::new(BTreeMap::new())),
810 range_calls: Arc::new(AtomicUsize::new(0)),
811 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
812 sequence_number: Arc::new(AtomicU64::new(0)),
813 };
814 let (base_url, shutdown_tx) = spawn_mock_server(state).await;
815 let client = StoreClient::new(&base_url);
816
817 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
818 .table(
819 "orders",
820 vec![
821 TableColumnConfig::new("id", DataType::Int64, false),
822 TableColumnConfig::new("status", DataType::Utf8, false),
823 TableColumnConfig::new("amount_cents", DataType::Int64, false),
824 ],
825 vec!["id".to_string()],
826 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
827 .expect("valid")
828 .with_cover_columns(vec!["amount_cents".to_string()])],
829 )
830 .expect("schema");
831 let ctx = SessionContext::new();
832 schema.register_all(&ctx).expect("register");
833
834 let explain = physical_plan_text(
835 &explain_plan_rows(
836 &ctx,
837 "SELECT id, status, amount_cents FROM orders \
838 WHERE status = 'open' AND amount_cents >= 5",
839 )
840 .await,
841 );
842 assert!(explain.contains("KvScanExec:"));
843 assert!(explain.contains("mode=secondary_index(status_idx, lexicographic)"));
844 assert!(explain.contains("predicate=status = 'open' AND amount_cents >= 5"));
845 assert!(explain.contains("exact=false"));
846 assert!(explain.contains("row_recheck=true"));
847 assert!(explain.contains("full_scan_like=false"));
848 assert_explain_includes_query_stats_surface(
849 &explain,
850 QueryStatsExplainSurface::StreamedRangeDetail,
851 );
852
853 let _ = shutdown_tx.send(());
854 }
855
856 #[tokio::test]
857 async fn explain_reports_zorder_secondary_index_scan() {
858 let state = MockState {
859 kv: Arc::new(Mutex::new(BTreeMap::new())),
860 range_calls: Arc::new(AtomicUsize::new(0)),
861 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
862 sequence_number: Arc::new(AtomicU64::new(0)),
863 };
864 let (base_url, shutdown_tx) = spawn_mock_server(state).await;
865 let client = StoreClient::new(&base_url);
866
867 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
868 .table(
869 "points",
870 vec![
871 TableColumnConfig::new("x", DataType::Int64, false),
872 TableColumnConfig::new("y", DataType::Int64, false),
873 TableColumnConfig::new("id", DataType::Int64, false),
874 TableColumnConfig::new("value", DataType::Int64, false),
875 ],
876 vec!["id".to_string()],
877 vec![
878 IndexSpec::z_order("xy_z", vec!["x".to_string(), "y".to_string()])
879 .expect("valid")
880 .with_cover_columns(vec!["value".to_string()]),
881 ],
882 )
883 .expect("schema");
884 let ctx = SessionContext::new();
885 schema.register_all(&ctx).expect("register");
886
887 let explain = physical_plan_text(
888 &explain_plan_rows(
889 &ctx,
890 "SELECT id, value FROM points \
891 WHERE x >= 1 AND x <= 2 AND y >= 1 AND y <= 2",
892 )
893 .await,
894 );
895 assert!(explain.contains("KvScanExec:"));
896 assert!(explain.contains("mode=secondary_index(xy_z, z_order)"));
897 assert!(explain.contains("exact=false"));
898 assert!(explain.contains("row_recheck=true"));
899
900 let _ = shutdown_tx.send(());
901 }
902
903 #[tokio::test]
904 async fn explain_reports_aggregate_pushdown_access_path_details() {
905 let state = MockState {
906 kv: Arc::new(Mutex::new(BTreeMap::new())),
907 range_calls: Arc::new(AtomicUsize::new(0)),
908 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
909 sequence_number: Arc::new(AtomicU64::new(0)),
910 };
911 let (base_url, shutdown_tx) = spawn_mock_server(state).await;
912 let client = StoreClient::new(&base_url);
913
914 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
915 .table(
916 "orders",
917 vec![
918 TableColumnConfig::new("id", DataType::Int64, false),
919 TableColumnConfig::new("status", DataType::Utf8, false),
920 TableColumnConfig::new("amount_cents", DataType::Int64, false),
921 ],
922 vec!["id".to_string()],
923 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
924 .expect("valid")
925 .with_cover_columns(vec!["amount_cents".to_string()])],
926 )
927 .expect("schema");
928 let ctx = SessionContext::new();
929 schema.register_all(&ctx).expect("register");
930
931 let explain = physical_plan_text(
932 &explain_plan_rows(
933 &ctx,
934 "SELECT status, SUM(amount_cents) AS total_cents \
935 FROM orders WHERE status = 'open' GROUP BY status",
936 )
937 .await,
938 );
939 assert!(explain.contains("KvAggregateExec:"));
940 assert!(explain.contains("grouped=true"));
941 assert!(explain.contains("job0{mode=secondary_index(status_idx, lexicographic)"));
942 assert!(explain.contains("predicate=status = 'open'"));
943 assert!(explain.contains("exact=true"));
944 assert!(explain.contains("row_recheck=false"));
945 assert_explain_includes_query_stats_surface(
946 &explain,
947 QueryStatsExplainSurface::RangeReduceDetail,
948 );
949
950 let _ = shutdown_tx.send(());
951 }
952
953 #[test]
954 fn index_spec_constructor_sets_name_and_keys() {
955 let spec = IndexSpec::new(
956 "status_customer",
957 vec!["status".to_string(), "customer_id".to_string()],
958 )
959 .expect("valid index spec");
960 assert_eq!(spec.name(), "status_customer");
961 assert_eq!(spec.key_columns(), &["status", "customer_id"]);
962 assert!(spec.cover_columns().is_empty());
963 }
964
965 #[test]
966 fn index_spec_cover_columns_are_configurable_in_code() {
967 let spec = IndexSpec::new("status_customer", vec!["status".to_string()])
968 .expect("valid")
969 .with_cover_columns(vec!["amount_cents".to_string()]);
970 assert_eq!(spec.key_columns(), &["status"]);
971 assert_eq!(spec.cover_columns(), &["amount_cents"]);
972 }
973
974 #[test]
975 fn describe_in_list_places_truncation_ellipsis_inside_parentheses() {
976 let rendered = describe_in_list((1..=6).map(|v| v.to_string()));
977 assert_eq!(rendered, "IN (1, 2, 3, 4, 5, ...)");
978 }
979
980 #[test]
981 fn normalize_sum_case_then_one_uses_countall_optimization() {
982 let (model, _) = test_model();
983 let argument = normalize_case_then_expr(
984 AggregatePushdownFunction::Sum,
985 &Expr::Literal(ScalarValue::Int64(Some(1)), None),
986 &model,
987 )
988 .expect("normalize");
989 assert_eq!(argument, AggregatePushdownArgument::CountAll);
990 }
991
992 #[test]
993 fn normalize_count_case_then_literal_uses_countall_optimization() {
994 use datafusion::logical_expr::col;
995
996 let (model, _) = test_model();
997 let case_expr = Expr::Case(datafusion::logical_expr::expr::Case {
998 expr: None,
999 when_then_expr: vec![(
1000 Box::new(col("status").eq(Expr::Literal(
1001 ScalarValue::Utf8(Some("open".to_string())),
1002 None,
1003 ))),
1004 Box::new(Expr::Literal(
1005 ScalarValue::Utf8(Some("yes".to_string())),
1006 None,
1007 )),
1008 )],
1009 else_expr: Some(Box::new(Expr::Literal(ScalarValue::Utf8(None), None))),
1010 });
1011
1012 let (func, argument, filter) =
1013 normalize_count_aggregate_argument(&case_expr, &model).expect("normalize");
1014 assert_eq!(func, AggregatePushdownFunction::Count);
1015 assert_eq!(argument, AggregatePushdownArgument::CountAll);
1016 assert!(filter.is_some());
1017 }
1018
1019 #[test]
1020 fn reduced_value_to_scalar_preserves_timestamp_timezone_label() {
1021 let tz: Arc<str> = Arc::from("America/New_York");
1022 let scalar = reduced_value_to_scalar(
1023 Some(KvReducedValue::Timestamp(1_700_000_000_000_000)),
1024 &DataType::Timestamp(TimeUnit::Microsecond, Some(tz.clone())),
1025 )
1026 .expect("timestamp scalar");
1027 assert_eq!(
1028 scalar,
1029 ScalarValue::TimestampMicrosecond(Some(1_700_000_000_000_000), Some(tz))
1030 );
1031 }
1032
1033 #[test]
1034 fn index_spec_cover_pk_column_is_rejected() {
1035 let config = KvTableConfig::new(
1036 0,
1037 vec![
1038 TableColumnConfig::new("id", DataType::Int64, false),
1039 TableColumnConfig::new("status", DataType::Utf8, false),
1040 ],
1041 vec!["id".to_string()],
1042 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
1043 .expect("valid")
1044 .with_cover_columns(vec!["id".to_string()])],
1045 )
1046 .expect("valid config");
1047 let model = TableModel::from_config(&config).expect("model");
1048 let err = model
1049 .resolve_index_specs(&config.index_specs)
1050 .expect_err("covering a PK column must be rejected");
1051 assert!(err.contains("primary key column"));
1052 }
1053
1054 #[test]
1055 fn access_plan_requires_cover_columns_for_index_scan() {
1056 let (model, _) = test_model();
1057 let predicate = QueryPredicate::default();
1058 let projection = Some(vec![
1059 *model.columns_by_name.get("order_id").unwrap(),
1060 *model.columns_by_name.get("amount_cents").unwrap(),
1061 ]);
1062 let plan = ScanAccessPlan::new(&model, &projection, &predicate);
1063
1064 let no_cover = IndexSpec::new("status_idx", vec!["status".to_string()]).unwrap();
1065 let with_cover = IndexSpec::new("status_idx", vec!["status".to_string()])
1066 .unwrap()
1067 .with_cover_columns(vec!["amount_cents".to_string()]);
1068 let no_cover_resolved = model.resolve_index_specs(&[no_cover]).unwrap();
1069 let with_cover_resolved = model.resolve_index_specs(&[with_cover]).unwrap();
1070
1071 assert!(!plan.index_covers_required_non_pk(&no_cover_resolved[0]));
1072 assert!(plan.index_covers_required_non_pk(&with_cover_resolved[0]));
1073 }
1074
1075 #[test]
1076 fn choose_index_plan_prefers_longer_prefix() {
1077 let (model, specs) = test_model();
1078 let region_idx = *model.columns_by_name.get("region").unwrap();
1079 let customer_idx = *model.columns_by_name.get("customer_id").unwrap();
1080 let mut predicate = QueryPredicate::default();
1081 predicate.constraints.insert(
1082 region_idx,
1083 PredicateConstraint::StringEq("us-east".to_string()),
1084 );
1085 predicate.constraints.insert(
1086 customer_idx,
1087 PredicateConstraint::IntRange {
1088 min: Some(10),
1089 max: Some(20),
1090 },
1091 );
1092 let plan = predicate
1093 .choose_index_plan(&model, &specs)
1094 .expect("plan")
1095 .expect("exists");
1096 assert_eq!(plan.spec_idx, 0);
1097 assert_eq!(plan.constrained_prefix_len, 2);
1098 }
1099
1100 #[test]
1101 fn choose_index_plan_prefers_covering_index_when_prefix_strength_ties() {
1102 let config = KvTableConfig::new(
1103 0,
1104 vec![
1105 TableColumnConfig::new("id", DataType::Int64, false),
1106 TableColumnConfig::new("status", DataType::Utf8, false),
1107 TableColumnConfig::new("amount_cents", DataType::Int64, false),
1108 ],
1109 vec!["id".to_string()],
1110 vec![
1111 IndexSpec::new("status_plain", vec!["status".to_string()]).expect("valid"),
1112 IndexSpec::new("status_covering", vec!["status".to_string()])
1113 .expect("valid")
1114 .with_cover_columns(vec!["amount_cents".to_string()]),
1115 ],
1116 )
1117 .expect("config");
1118 let model = TableModel::from_config(&config).expect("model");
1119 let specs = model
1120 .resolve_index_specs(&config.index_specs)
1121 .expect("specs");
1122 let status_idx = *model.columns_by_name.get("status").unwrap();
1123 let amount_idx = *model.columns_by_name.get("amount_cents").unwrap();
1124 let mut predicate = QueryPredicate::default();
1125 predicate.constraints.insert(
1126 status_idx,
1127 PredicateConstraint::StringEq("open".to_string()),
1128 );
1129 predicate.constraints.insert(
1130 amount_idx,
1131 PredicateConstraint::IntRange {
1132 min: Some(10),
1133 max: None,
1134 },
1135 );
1136
1137 let plan = predicate
1138 .choose_index_plan(&model, &specs)
1139 .expect("plan")
1140 .expect("exists");
1141 assert_eq!(specs[plan.spec_idx].name, "status_covering");
1142 }
1143
1144 #[test]
1145 fn choose_index_plan_prefers_zorder_for_multi_column_box_constraints() {
1146 let (model, specs) = zorder_test_model();
1147 let x_idx = *model.columns_by_name.get("x").unwrap();
1148 let y_idx = *model.columns_by_name.get("y").unwrap();
1149 let mut predicate = QueryPredicate::default();
1150 predicate.constraints.insert(
1151 x_idx,
1152 PredicateConstraint::IntRange {
1153 min: Some(1),
1154 max: Some(2),
1155 },
1156 );
1157 predicate.constraints.insert(
1158 y_idx,
1159 PredicateConstraint::IntRange {
1160 min: Some(1),
1161 max: Some(2),
1162 },
1163 );
1164
1165 let plan = predicate
1166 .choose_index_plan(&model, &specs)
1167 .expect("plan")
1168 .expect("exists");
1169 assert_eq!(specs[plan.spec_idx].name, "xy_z");
1170 assert_eq!(specs[plan.spec_idx].layout, IndexLayout::ZOrder);
1171 assert_eq!(plan.constrained_column_count, 2);
1172 }
1173
1174 #[test]
1175 fn secondary_index_key_round_trip() {
1176 let (model, specs) = test_model();
1177 let row = KvRow {
1178 values: vec![
1179 CellValue::Utf8("us-east".to_string()),
1180 CellValue::Int64(42),
1181 CellValue::Int64(9001),
1182 CellValue::Int64(1500),
1183 CellValue::Utf8("open".to_string()),
1184 ],
1185 };
1186 let key = encode_secondary_index_key(model.table_prefix, &specs[0], &model, &row)
1187 .expect("encode");
1188 let decoded = decode_secondary_index_key(model.table_prefix, &specs[0], &model, &key)
1189 .expect("decode");
1190 let region_idx = *model.columns_by_name.get("region").unwrap();
1191 let customer_idx = *model.columns_by_name.get("customer_id").unwrap();
1192 assert!(matches!(
1193 decoded.values.get(®ion_idx),
1194 Some(CellValue::Utf8(v)) if v == "us-east"
1195 ));
1196 assert!(matches!(
1197 decoded.values.get(&customer_idx),
1198 Some(CellValue::Int64(v)) if *v == 42
1199 ));
1200 assert!(matches!(
1201 &decoded.primary_key_values[0],
1202 CellValue::Int64(9001)
1203 ));
1204 let expected_pk = encode_primary_key_from_row(model.table_prefix, &row, &model)
1205 .expect("primary key should encode");
1206 assert_eq!(decoded.primary_key, expected_pk);
1207 }
1208
1209 #[test]
1210 fn zorder_secondary_index_key_round_trip() {
1211 let (model, specs) = zorder_test_model();
1212 let row = KvRow {
1213 values: vec![
1214 CellValue::Int64(2),
1215 CellValue::Int64(1),
1216 CellValue::Int64(42),
1217 CellValue::Int64(900),
1218 ],
1219 };
1220 let key = encode_secondary_index_key(model.table_prefix, &specs[1], &model, &row)
1221 .expect("encode");
1222 let decoded = decode_secondary_index_key(model.table_prefix, &specs[1], &model, &key)
1223 .expect("decode");
1224 let x_idx = *model.columns_by_name.get("x").unwrap();
1225 let y_idx = *model.columns_by_name.get("y").unwrap();
1226 assert!(matches!(
1227 decoded.values.get(&x_idx),
1228 Some(CellValue::Int64(v)) if *v == 2
1229 ));
1230 assert!(matches!(
1231 decoded.values.get(&y_idx),
1232 Some(CellValue::Int64(v)) if *v == 1
1233 ));
1234 assert!(matches!(
1235 &decoded.primary_key_values[0],
1236 CellValue::Int64(42)
1237 ));
1238 }
1239
1240 #[test]
1241 fn table_config_supports_non_orders_schema() {
1242 let config = KvTableConfig::new(
1243 0,
1244 vec![
1245 TableColumnConfig::new("tenant", DataType::Utf8, false),
1246 TableColumnConfig::new("id", DataType::Int64, false),
1247 TableColumnConfig::new("score", DataType::Int64, false),
1248 ],
1249 vec!["id".to_string()],
1250 vec![IndexSpec::new(
1251 "tenant_score",
1252 vec!["tenant".to_string(), "score".to_string()],
1253 )
1254 .expect("valid")],
1255 )
1256 .expect("schema agnostic config should be valid");
1257 assert_eq!(config.primary_key_columns, vec!["id".to_string()]);
1258 assert_eq!(config.columns.len(), 3);
1259 }
1260
1261 #[test]
1262 fn table_config_accepts_float64_column() {
1263 let config = KvTableConfig::new(
1264 0,
1265 vec![
1266 TableColumnConfig::new("id", DataType::Int64, false),
1267 TableColumnConfig::new("price", DataType::Float64, false),
1268 ],
1269 vec!["id".to_string()],
1270 vec![],
1271 )
1272 .expect("Float64 column should be accepted");
1273 assert_eq!(config.columns.len(), 2);
1274 }
1275
1276 #[test]
1277 fn table_config_accepts_boolean_column() {
1278 let config = KvTableConfig::new(
1279 0,
1280 vec![
1281 TableColumnConfig::new("id", DataType::Int64, false),
1282 TableColumnConfig::new("active", DataType::Boolean, false),
1283 ],
1284 vec!["id".to_string()],
1285 vec![],
1286 )
1287 .expect("Boolean column should be accepted");
1288 assert_eq!(config.columns.len(), 2);
1289 }
1290
1291 #[test]
1292 fn build_projected_batch_uses_large_utf8_type() {
1293 let config = KvTableConfig::new(
1294 0,
1295 vec![
1296 TableColumnConfig::new("id", DataType::Int64, false),
1297 TableColumnConfig::new("name", DataType::LargeUtf8, false),
1298 ],
1299 vec!["id".to_string()],
1300 vec![],
1301 )
1302 .unwrap();
1303 let model = TableModel::from_config(&config).unwrap();
1304 let rows = vec![KvRow {
1305 values: vec![CellValue::Int64(1), CellValue::Utf8("hello".to_string())],
1306 }];
1307 let batch = build_projected_batch(&rows, &model, &model.schema, &None).unwrap();
1308 assert_eq!(batch.column(1).data_type(), &DataType::LargeUtf8);
1309 let values = batch
1310 .column(1)
1311 .as_any()
1312 .downcast_ref::<LargeStringArray>()
1313 .expect("must build LargeStringArray");
1314 assert_eq!(values.value(0), "hello");
1315 }
1316
1317 #[test]
1318 fn build_projected_batch_uses_utf8_view_type() {
1319 let config = KvTableConfig::new(
1320 0,
1321 vec![
1322 TableColumnConfig::new("id", DataType::Int64, false),
1323 TableColumnConfig::new("name", DataType::Utf8View, false),
1324 ],
1325 vec!["id".to_string()],
1326 vec![],
1327 )
1328 .unwrap();
1329 let model = TableModel::from_config(&config).unwrap();
1330 let rows = vec![KvRow {
1331 values: vec![CellValue::Int64(1), CellValue::Utf8("hello".to_string())],
1332 }];
1333 let batch = build_projected_batch(&rows, &model, &model.schema, &None).unwrap();
1334 assert_eq!(batch.column(1).data_type(), &DataType::Utf8View);
1335 let values = batch
1336 .column(1)
1337 .as_any()
1338 .downcast_ref::<StringViewArray>()
1339 .expect("must build StringViewArray");
1340 assert_eq!(values.value(0), "hello");
1341 }
1342
1343 #[test]
1344 fn build_projected_batch_uses_large_binary_type() {
1345 let config = KvTableConfig::new(
1346 0,
1347 vec![
1348 TableColumnConfig::new("id", DataType::Int64, false),
1349 TableColumnConfig::new("payload", DataType::LargeBinary, false),
1350 ],
1351 vec!["id".to_string()],
1352 vec![],
1353 )
1354 .unwrap();
1355 let model = TableModel::from_config(&config).unwrap();
1356 let rows = vec![KvRow {
1357 values: vec![CellValue::Int64(1), CellValue::Binary(vec![1, 2, 3])],
1358 }];
1359 let batch = build_projected_batch(&rows, &model, &model.schema, &None).unwrap();
1360 assert_eq!(batch.column(1).data_type(), &DataType::LargeBinary);
1361 let values = batch
1362 .column(1)
1363 .as_any()
1364 .downcast_ref::<LargeBinaryArray>()
1365 .expect("must build LargeBinaryArray");
1366 assert_eq!(values.value(0), &[1, 2, 3]);
1367 }
1368
1369 #[test]
1370 fn build_projected_batch_uses_binary_view_type() {
1371 let config = KvTableConfig::new(
1372 0,
1373 vec![
1374 TableColumnConfig::new("id", DataType::Int64, false),
1375 TableColumnConfig::new("payload", DataType::BinaryView, false),
1376 ],
1377 vec!["id".to_string()],
1378 vec![],
1379 )
1380 .unwrap();
1381 let model = TableModel::from_config(&config).unwrap();
1382 let rows = vec![KvRow {
1383 values: vec![CellValue::Int64(1), CellValue::Binary(vec![1, 2, 3])],
1384 }];
1385 let batch = build_projected_batch(&rows, &model, &model.schema, &None).unwrap();
1386 assert_eq!(batch.column(1).data_type(), &DataType::BinaryView);
1387 let values = batch
1388 .column(1)
1389 .as_any()
1390 .downcast_ref::<BinaryViewArray>()
1391 .expect("must build BinaryViewArray");
1392 assert_eq!(values.value(0), &[1, 2, 3]);
1393 }
1394
1395 #[test]
1396 fn binary_column_rejected_in_primary_key() {
1397 let error = KvTableConfig::new(
1398 0,
1399 vec![
1400 TableColumnConfig::new("id", DataType::Binary, false),
1401 TableColumnConfig::new("data", DataType::Utf8, true),
1402 ],
1403 vec!["id".to_string()],
1404 vec![],
1405 )
1406 .expect_err("Binary primary key must be rejected");
1407 assert!(
1408 error.contains("must be Int64, UInt64, Utf8, or FixedSizeBinary"),
1409 "unexpected error: {error}"
1410 );
1411 }
1412
1413 #[test]
1414 fn binary_column_rejected_in_index() {
1415 let result = KvTableConfig::new(
1416 0,
1417 vec![
1418 TableColumnConfig::new("id", DataType::Int64, false),
1419 TableColumnConfig::new("body", DataType::Binary, false),
1420 ],
1421 vec!["id".to_string()],
1422 vec![IndexSpec::new("body_idx", vec!["body".to_string()]).unwrap()],
1423 );
1424 assert!(
1425 result.is_err() || {
1426 let config = result.unwrap();
1427 let model = TableModel::from_config(&config).unwrap();
1428 model.resolve_index_specs(&config.index_specs).is_err()
1429 }
1430 );
1431 }
1432
1433 #[test]
1436 fn binary_filters_are_not_pushed_down() {
1437 let config = KvTableConfig::new(
1438 0,
1439 vec![
1440 TableColumnConfig::new("id", DataType::Int64, false),
1441 TableColumnConfig::new("body", DataType::Binary, false),
1442 ],
1443 vec!["id".to_string()],
1444 vec![],
1445 )
1446 .unwrap();
1447 let model = TableModel::from_config(&config).unwrap();
1448
1449 use datafusion::logical_expr::col;
1450 let filter = col("body").eq(Expr::Literal(
1451 ScalarValue::Binary(Some(vec![0xAB, 0xCD])),
1452 None,
1453 ));
1454
1455 assert!(!QueryPredicate::supports_filter(&filter, &model));
1456 let pred = QueryPredicate::from_filters(&[filter], &model);
1457 assert!(pred.constraints.is_empty());
1458 }
1459
1460 #[test]
1463 fn binary_has_no_kv_aggregate_field_kind() {
1464 assert!(kv_field_kind(ColumnKind::Binary).is_none());
1465 }
1466
1467 #[test]
1470 fn binary_cells_decode_from_archived_non_pk_values() {
1471 let config = KvTableConfig::new(
1472 0,
1473 vec![
1474 TableColumnConfig::new("id", DataType::UInt64, false),
1475 TableColumnConfig::new("body", DataType::Binary, false),
1476 ],
1477 vec!["id".to_string()],
1478 vec![],
1479 )
1480 .unwrap();
1481 let model = TableModel::from_config(&config).unwrap();
1482 let body_idx = *model.columns_by_name.get("body").expect("body column");
1483
1484 let bodies: Vec<Vec<u8>> = vec![vec![], vec![0xAB], vec![0xCD; 300]];
1485 for body in bodies {
1486 let row = KvRow {
1487 values: vec![CellValue::UInt64(1), CellValue::Binary(body.clone())],
1488 };
1489 let encoded = encode_base_row_value(&row, &model).expect("encode row");
1490 let stored = decode_stored_row(&encoded).expect("decode stored row");
1491 let cell = cell_value_from_archived_non_pk(
1492 &model.columns[body_idx],
1493 stored.values[body_idx].as_ref(),
1494 )
1495 .expect("decode archived binary")
1496 .expect("value is not null");
1497 assert!(
1498 matches!(&cell, CellValue::Binary(v) if *v == body),
1499 "wrong cell for {}-byte body: {cell:?}",
1500 body.len()
1501 );
1502 }
1503 }
1504
1505 #[test]
1506 fn f64_ordered_encoding_preserves_order() {
1507 let values = [
1508 f64::NEG_INFINITY,
1509 f64::MIN,
1510 -1000.0,
1511 -1.0,
1512 -0.001,
1513 0.0,
1514 0.001,
1515 1.0,
1516 1000.0,
1517 f64::MAX,
1518 f64::INFINITY,
1519 ];
1520 let encoded: Vec<[u8; 8]> = values.iter().map(|v| encode_f64_ordered(*v)).collect();
1521 for i in 0..encoded.len() - 1 {
1522 assert!(
1523 encoded[i] < encoded[i + 1],
1524 "encode_f64_ordered({}) >= encode_f64_ordered({})",
1525 values[i],
1526 values[i + 1]
1527 );
1528 }
1529 }
1530
1531 #[test]
1532 fn f64_ordered_encoding_round_trip() {
1533 let values = [
1534 f64::MIN,
1535 -42.5,
1536 -0.0,
1537 0.0,
1538 3.125,
1539 f64::MAX,
1540 f64::INFINITY,
1541 f64::NEG_INFINITY,
1542 ];
1543 for v in values {
1544 let encoded = encode_f64_ordered(v);
1545 let decoded = decode_f64_ordered(encoded);
1546 assert!(
1547 v.to_bits() == decoded.to_bits(),
1548 "round-trip failed for {v}: got {decoded}"
1549 );
1550 }
1551 }
1552
1553 fn mixed_model() -> (TableModel, Vec<ResolvedIndexSpec>) {
1554 let config = KvTableConfig::new(
1555 0,
1556 vec![
1557 TableColumnConfig::new("id", DataType::Int64, false),
1558 TableColumnConfig::new("label", DataType::Utf8, false),
1559 TableColumnConfig::new("score", DataType::Float64, false),
1560 TableColumnConfig::new("active", DataType::Boolean, false),
1561 ],
1562 vec!["id".to_string()],
1563 vec![
1564 IndexSpec::new(
1565 "active_score",
1566 vec!["active".to_string(), "score".to_string()],
1567 )
1568 .expect("valid"),
1569 IndexSpec::new("label_idx", vec!["label".to_string()]).expect("valid"),
1570 ],
1571 )
1572 .expect("valid config");
1573 let model = TableModel::from_config(&config).expect("model");
1574 let specs = model
1575 .resolve_index_specs(&config.index_specs)
1576 .expect("specs");
1577 (model, specs)
1578 }
1579
1580 #[test]
1581 fn secondary_index_key_round_trip_with_float64_and_boolean() {
1582 let (model, specs) = mixed_model();
1583 let row = KvRow {
1584 values: vec![
1585 CellValue::Int64(100),
1586 CellValue::Utf8("hello".to_string()),
1587 CellValue::Float64(3.125),
1588 CellValue::Boolean(true),
1589 ],
1590 };
1591 let key = encode_secondary_index_key(model.table_prefix, &specs[0], &model, &row)
1592 .expect("encode");
1593 let decoded = decode_secondary_index_key(model.table_prefix, &specs[0], &model, &key)
1594 .expect("decode");
1595 let active_idx = *model.columns_by_name.get("active").unwrap();
1596 let score_idx = *model.columns_by_name.get("score").unwrap();
1597 assert!(matches!(
1598 decoded.values.get(&active_idx),
1599 Some(CellValue::Boolean(true))
1600 ));
1601 assert!(
1602 matches!(decoded.values.get(&score_idx), Some(CellValue::Float64(v)) if (*v - 3.125).abs() < f64::EPSILON)
1603 );
1604 assert!(matches!(
1605 &decoded.primary_key_values[0],
1606 CellValue::Int64(100)
1607 ));
1608 }
1609
1610 #[test]
1611 fn base_row_round_trip_with_float64_and_boolean() {
1612 let (model, _specs) = mixed_model();
1613 let row = KvRow {
1614 values: vec![
1615 CellValue::Int64(42),
1616 CellValue::Utf8("world".to_string()),
1617 CellValue::Float64(-99.5),
1618 CellValue::Boolean(false),
1619 ],
1620 };
1621 let encoded = encode_base_row_value(&row, &model).expect("encode");
1622 let decoded =
1623 decode_base_row(vec![CellValue::Int64(42)], &encoded, &model).expect("decode");
1624 assert!(matches!(&decoded.values[0], CellValue::Int64(42)));
1625 assert!(matches!(&decoded.values[1], CellValue::Utf8(v) if v == "world"));
1626 assert!(
1627 matches!(&decoded.values[2], CellValue::Float64(v) if (*v - (-99.5)).abs() < f64::EPSILON)
1628 );
1629 assert!(matches!(&decoded.values[3], CellValue::Boolean(false)));
1630 }
1631
1632 #[test]
1633 fn predicate_bool_eq_matches() {
1634 let (model, _specs) = mixed_model();
1635 let active_idx = *model.columns_by_name.get("active").unwrap();
1636 let mut pred = QueryPredicate::default();
1637 pred.constraints
1638 .insert(active_idx, PredicateConstraint::BoolEq(true));
1639 let row_true = KvRow {
1640 values: vec![
1641 CellValue::Int64(1),
1642 CellValue::Utf8("a".to_string()),
1643 CellValue::Float64(1.0),
1644 CellValue::Boolean(true),
1645 ],
1646 };
1647 let row_false = KvRow {
1648 values: vec![
1649 CellValue::Int64(2),
1650 CellValue::Utf8("b".to_string()),
1651 CellValue::Float64(2.0),
1652 CellValue::Boolean(false),
1653 ],
1654 };
1655 assert!(pred.matches_row(&row_true));
1656 assert!(!pred.matches_row(&row_false));
1657 }
1658
1659 #[test]
1660 fn predicate_float_range_matches() {
1661 let (model, _specs) = mixed_model();
1662 let score_idx = *model.columns_by_name.get("score").unwrap();
1663 let mut pred = QueryPredicate::default();
1664 pred.constraints.insert(
1665 score_idx,
1666 PredicateConstraint::FloatRange {
1667 min: Some((2.0, true)),
1668 max: Some((5.0, false)),
1669 },
1670 );
1671 let make_row = |score: f64| KvRow {
1672 values: vec![
1673 CellValue::Int64(1),
1674 CellValue::Utf8("a".to_string()),
1675 CellValue::Float64(score),
1676 CellValue::Boolean(true),
1677 ],
1678 };
1679 assert!(!pred.matches_row(&make_row(1.99)));
1680 assert!(pred.matches_row(&make_row(2.0)));
1681 assert!(pred.matches_row(&make_row(3.5)));
1682 assert!(pred.matches_row(&make_row(4.99)));
1683 assert!(!pred.matches_row(&make_row(5.0)));
1684 assert!(!pred.matches_row(&make_row(5.01)));
1685 }
1686
1687 #[test]
1688 fn float_range_rejects_nan_row_value() {
1689 let constraint = PredicateConstraint::FloatRange {
1690 min: Some((0.0, true)),
1691 max: Some((10.0, true)),
1692 };
1693 assert!(!matches_constraint(
1694 &CellValue::Float64(f64::NAN),
1695 &constraint
1696 ));
1697 }
1698
1699 #[test]
1700 fn index_plan_with_boolean_prefix() {
1701 let (model, specs) = mixed_model();
1702 let active_idx = *model.columns_by_name.get("active").unwrap();
1703 let mut pred = QueryPredicate::default();
1704 pred.constraints
1705 .insert(active_idx, PredicateConstraint::BoolEq(true));
1706 let plan = pred
1707 .choose_index_plan(&model, &specs)
1708 .expect("plan")
1709 .expect("should find index");
1710 assert_eq!(plan.spec_idx, 0);
1711 assert_eq!(plan.constrained_prefix_len, 1);
1712 }
1713
1714 #[test]
1715 fn float_constraint_contradiction() {
1716 let mut lo: Option<(f64, bool)> = None;
1717 let mut hi: Option<(f64, bool)> = None;
1718 let mut contradiction = false;
1719 apply_float_constraint(&mut lo, &mut hi, Operator::Gt, 10.0, &mut contradiction);
1720 assert!(!contradiction);
1721 apply_float_constraint(&mut lo, &mut hi, Operator::Lt, 5.0, &mut contradiction);
1722 assert!(contradiction);
1723 }
1724
1725 #[test]
1726 fn float_constraint_eq_then_range_contradicts() {
1727 let mut lo: Option<(f64, bool)> = None;
1728 let mut hi: Option<(f64, bool)> = None;
1729 let mut contradiction = false;
1730 apply_float_constraint(&mut lo, &mut hi, Operator::Eq, 5.0, &mut contradiction);
1731 assert!(!contradiction);
1732 apply_float_constraint(&mut lo, &mut hi, Operator::Gt, 5.0, &mut contradiction);
1733 assert!(contradiction);
1734 }
1735
1736 #[test]
1737 fn float_nan_literal_comparison_marks_contradiction() {
1738 let config = KvTableConfig::new(
1739 0,
1740 vec![
1741 TableColumnConfig::new("id", DataType::Int64, false),
1742 TableColumnConfig::new("score", DataType::Float64, false),
1743 ],
1744 vec!["id".to_string()],
1745 vec![],
1746 )
1747 .unwrap();
1748 let model = TableModel::from_config(&config).unwrap();
1749
1750 use datafusion::logical_expr::col;
1751 let filter = col("score").gt(Expr::Literal(ScalarValue::Float64(Some(f64::NAN)), None));
1752 assert!(QueryPredicate::supports_filter(&filter, &model));
1753
1754 let pred = QueryPredicate::from_filters(&[filter], &model);
1755 assert!(
1756 pred.contradiction,
1757 "comparison with NaN literal must produce contradiction"
1758 );
1759 }
1760
1761 #[test]
1762 fn null_predicate_merges_are_order_independent() {
1763 use datafusion::logical_expr::col;
1764 let config = KvTableConfig::new(
1765 0,
1766 vec![
1767 TableColumnConfig::new("id", DataType::Int64, false),
1768 TableColumnConfig::new("label", DataType::Utf8, true),
1769 ],
1770 vec!["id".to_string()],
1771 vec![],
1772 )
1773 .unwrap();
1774 let model = TableModel::from_config(&config).unwrap();
1775 let label_idx = *model.columns_by_name.get("label").unwrap();
1776 let eq_foo = col("label").eq(Expr::Literal(
1777 ScalarValue::Utf8(Some("foo".to_string())),
1778 None,
1779 ));
1780 let is_null = col("label").is_null();
1781 let is_not_null = col("label").is_not_null();
1782 let row_foo = KvRow {
1783 values: vec![CellValue::Int64(1), CellValue::Utf8("foo".to_string())],
1784 };
1785
1786 for filters in [[&eq_foo, &is_not_null], [&is_not_null, &eq_foo]] {
1789 let owned: Vec<Expr> = filters.iter().map(|e| (*e).clone()).collect();
1790 let pred = QueryPredicate::from_filters(&owned, &model);
1791 assert!(!pred.contradiction);
1792 assert!(matches!(
1793 pred.constraints.get(&label_idx),
1794 Some(PredicateConstraint::StringEq(s)) if s == "foo"
1795 ));
1796 assert!(pred.matches_row(&row_foo));
1797 }
1798
1799 for filters in [[&eq_foo, &is_null], [&is_null, &eq_foo]] {
1801 let owned: Vec<Expr> = filters.iter().map(|e| (*e).clone()).collect();
1802 assert!(QueryPredicate::from_filters(&owned, &model).contradiction);
1803 }
1804 }
1805
1806 #[test]
1807 fn non_nullable_null_predicates_simplify() {
1808 use datafusion::logical_expr::col;
1809 let config = KvTableConfig::new(
1810 0,
1811 vec![
1812 TableColumnConfig::new("id", DataType::Int64, false),
1813 TableColumnConfig::new("label", DataType::Utf8, true),
1814 ],
1815 vec!["id".to_string()],
1816 vec![],
1817 )
1818 .unwrap();
1819 let model = TableModel::from_config(&config).unwrap();
1820
1821 let is_null = col("id").is_null();
1822 assert!(QueryPredicate::supports_filter(&is_null, &model));
1823 assert!(QueryPredicate::from_filters(&[is_null], &model).contradiction);
1824
1825 let is_not_null = col("id").is_not_null();
1826 assert!(QueryPredicate::supports_filter(&is_not_null, &model));
1827 let pred = QueryPredicate::from_filters(&[is_not_null], &model);
1828 assert!(!pred.contradiction);
1829 assert!(pred.constraints.is_empty());
1830 }
1831
1832 #[test]
1833 fn null_literal_comparisons_are_contradictions() {
1834 use datafusion::logical_expr::col;
1835 let config = KvTableConfig::new(
1836 0,
1837 vec![
1838 TableColumnConfig::new("id", DataType::Int64, false),
1839 TableColumnConfig::new("label", DataType::Utf8, true),
1840 ],
1841 vec!["id".to_string()],
1842 vec![],
1843 )
1844 .unwrap();
1845 let model = TableModel::from_config(&config).unwrap();
1846
1847 for filter in [
1848 col("label").eq(Expr::Literal(ScalarValue::Utf8(None), None)),
1849 col("id").gt(Expr::Literal(ScalarValue::Int64(None), None)),
1850 ] {
1851 assert!(QueryPredicate::supports_filter(&filter, &model));
1852 assert!(QueryPredicate::from_filters(&[filter], &model).contradiction);
1853 }
1854 }
1855
1856 #[test]
1857 fn in_predicate_merges_with_comparisons_order_independent() {
1858 use datafusion::logical_expr::{col, in_list};
1859 let config = KvTableConfig::new(
1860 0,
1861 vec![
1862 TableColumnConfig::new("id", DataType::Int64, false),
1863 TableColumnConfig::new("label", DataType::Utf8, true),
1864 TableColumnConfig::new("score", DataType::Int64, true),
1865 TableColumnConfig::new("version", DataType::UInt64, true),
1866 TableColumnConfig::new("hash", DataType::FixedSizeBinary(2), true),
1867 ],
1868 vec!["id".to_string()],
1869 vec![],
1870 )
1871 .unwrap();
1872 let model = TableModel::from_config(&config).unwrap();
1873 let label_idx = *model.columns_by_name.get("label").unwrap();
1874 let score_idx = *model.columns_by_name.get("score").unwrap();
1875 let version_idx = *model.columns_by_name.get("version").unwrap();
1876 let hash_idx = *model.columns_by_name.get("hash").unwrap();
1877
1878 let label_in = in_list(
1879 col("label"),
1880 vec![
1881 Expr::Literal(ScalarValue::Utf8(Some("foo".to_string())), None),
1882 Expr::Literal(ScalarValue::Utf8(Some("bar".to_string())), None),
1883 ],
1884 false,
1885 );
1886 let label_eq = col("label").eq(Expr::Literal(
1887 ScalarValue::Utf8(Some("foo".to_string())),
1888 None,
1889 ));
1890 for filters in [[&label_in, &label_eq], [&label_eq, &label_in]] {
1891 let owned: Vec<Expr> = filters.iter().map(|e| (*e).clone()).collect();
1892 let pred = QueryPredicate::from_filters(&owned, &model);
1893 assert!(!pred.contradiction);
1894 assert!(matches!(
1895 pred.constraints.get(&label_idx),
1896 Some(PredicateConstraint::StringEq(v)) if v == "foo"
1897 ));
1898 }
1899
1900 let label_miss = col("label").eq(Expr::Literal(
1901 ScalarValue::Utf8(Some("baz".to_string())),
1902 None,
1903 ));
1904 for filters in [[&label_in, &label_miss], [&label_miss, &label_in]] {
1905 let owned: Vec<Expr> = filters.iter().map(|e| (*e).clone()).collect();
1906 assert!(QueryPredicate::from_filters(&owned, &model).contradiction);
1907 }
1908
1909 let score_in = in_list(
1910 col("score"),
1911 vec![
1912 Expr::Literal(ScalarValue::Int64(Some(1)), None),
1913 Expr::Literal(ScalarValue::Int64(Some(2)), None),
1914 Expr::Literal(ScalarValue::Int64(Some(3)), None),
1915 ],
1916 false,
1917 );
1918 let score_gt = col("score").gt(Expr::Literal(ScalarValue::Int64(Some(1)), None));
1919 for filters in [[&score_in, &score_gt], [&score_gt, &score_in]] {
1920 let owned: Vec<Expr> = filters.iter().map(|e| (*e).clone()).collect();
1921 let pred = QueryPredicate::from_filters(&owned, &model);
1922 assert!(!pred.contradiction);
1923 assert!(matches!(
1924 pred.constraints.get(&score_idx),
1925 Some(PredicateConstraint::IntIn(v)) if v == &vec![2, 3]
1926 ));
1927 }
1928
1929 let version_in = in_list(
1930 col("version"),
1931 vec![
1932 Expr::Literal(ScalarValue::UInt64(Some(1)), None),
1933 Expr::Literal(ScalarValue::UInt64(Some(2)), None),
1934 Expr::Literal(ScalarValue::UInt64(Some(3)), None),
1935 ],
1936 false,
1937 );
1938 let version_lt = col("version").lt(Expr::Literal(ScalarValue::UInt64(Some(3)), None));
1939 for filters in [[&version_in, &version_lt], [&version_lt, &version_in]] {
1940 let owned: Vec<Expr> = filters.iter().map(|e| (*e).clone()).collect();
1941 let pred = QueryPredicate::from_filters(&owned, &model);
1942 assert!(!pred.contradiction);
1943 assert!(matches!(
1944 pred.constraints.get(&version_idx),
1945 Some(PredicateConstraint::UInt64In(v)) if v == &vec![1, 2]
1946 ));
1947 }
1948
1949 let hash_in = in_list(
1950 col("hash"),
1951 vec![
1952 Expr::Literal(
1953 ScalarValue::FixedSizeBinary(2, Some(vec![0xAA, 0xAA])),
1954 None,
1955 ),
1956 Expr::Literal(
1957 ScalarValue::FixedSizeBinary(2, Some(vec![0xBB, 0xBB])),
1958 None,
1959 ),
1960 ],
1961 false,
1962 );
1963 let hash_eq = col("hash").eq(Expr::Literal(
1964 ScalarValue::FixedSizeBinary(2, Some(vec![0xAA, 0xAA])),
1965 None,
1966 ));
1967 for filters in [[&hash_in, &hash_eq], [&hash_eq, &hash_in]] {
1968 let owned: Vec<Expr> = filters.iter().map(|e| (*e).clone()).collect();
1969 let pred = QueryPredicate::from_filters(&owned, &model);
1970 assert!(!pred.contradiction);
1971 assert!(matches!(
1972 pred.constraints.get(&hash_idx),
1973 Some(PredicateConstraint::FixedBinaryEq(v)) if v == &vec![0xAA, 0xAA]
1974 ));
1975 }
1976 }
1977
1978 #[test]
1979 fn empty_and_null_in_lists_are_exact() {
1980 use datafusion::logical_expr::{col, in_list};
1981 let config = KvTableConfig::new(
1982 0,
1983 vec![
1984 TableColumnConfig::new("id", DataType::Int64, false),
1985 TableColumnConfig::new("label", DataType::Utf8, true),
1986 TableColumnConfig::new("score", DataType::Float64, true),
1987 TableColumnConfig::new("version", DataType::UInt64, true),
1988 TableColumnConfig::new("hash", DataType::FixedSizeBinary(2), true),
1989 ],
1990 vec!["id".to_string()],
1991 vec![],
1992 )
1993 .unwrap();
1994 let model = TableModel::from_config(&config).unwrap();
1995 let id_idx = *model.columns_by_name.get("id").unwrap();
1996
1997 for filter in [
1998 in_list(col("id"), vec![], false),
1999 in_list(
2000 col("id"),
2001 vec![Expr::Literal(ScalarValue::Int64(None), None)],
2002 false,
2003 ),
2004 in_list(
2005 col("label"),
2006 vec![Expr::Literal(ScalarValue::Utf8(None), None)],
2007 false,
2008 ),
2009 in_list(
2010 col("score"),
2011 vec![Expr::Literal(ScalarValue::Float64(None), None)],
2012 false,
2013 ),
2014 in_list(
2015 col("version"),
2016 vec![Expr::Literal(ScalarValue::UInt64(None), None)],
2017 false,
2018 ),
2019 in_list(
2020 col("hash"),
2021 vec![Expr::Literal(ScalarValue::FixedSizeBinary(2, None), None)],
2022 false,
2023 ),
2024 ] {
2025 assert!(QueryPredicate::supports_filter(&filter, &model));
2026 assert!(QueryPredicate::from_filters(&[filter], &model).contradiction);
2027 }
2028
2029 let filter = in_list(
2030 col("id"),
2031 vec![
2032 Expr::Literal(ScalarValue::Int64(Some(7)), None),
2033 Expr::Literal(ScalarValue::Int64(None), None),
2034 ],
2035 false,
2036 );
2037 assert!(QueryPredicate::supports_filter(&filter, &model));
2038 let pred = QueryPredicate::from_filters(&[filter], &model);
2039 assert!(!pred.contradiction);
2040 assert!(matches!(
2041 pred.constraints.get(&id_idx),
2042 Some(PredicateConstraint::IntRange {
2043 min: Some(7),
2044 max: Some(7)
2045 })
2046 ));
2047 }
2048
2049 #[test]
2050 fn table_config_accepts_date32_column() {
2051 let config = KvTableConfig::new(
2052 0,
2053 vec![
2054 TableColumnConfig::new("id", DataType::Int64, false),
2055 TableColumnConfig::new("created", DataType::Date32, false),
2056 ],
2057 vec!["id".to_string()],
2058 vec![],
2059 )
2060 .expect("Date32 column should be accepted");
2061 assert_eq!(config.columns.len(), 2);
2062 }
2063
2064 #[test]
2065 fn table_config_accepts_timestamp_column() {
2066 let config = KvTableConfig::new(
2067 0,
2068 vec![
2069 TableColumnConfig::new("id", DataType::Int64, false),
2070 TableColumnConfig::new(
2071 "ts",
2072 DataType::Timestamp(TimeUnit::Microsecond, None),
2073 false,
2074 ),
2075 ],
2076 vec!["id".to_string()],
2077 vec![],
2078 )
2079 .expect("Timestamp column should be accepted");
2080 let schema = config.to_schema();
2081 assert!(matches!(
2082 schema.field(1).data_type(),
2083 DataType::Timestamp(TimeUnit::Microsecond, _)
2084 ));
2085 }
2086
2087 #[test]
2088 fn table_config_normalizes_timestamp_to_microsecond() {
2089 let config = KvTableConfig::new(
2090 0,
2091 vec![
2092 TableColumnConfig::new("id", DataType::Int64, false),
2093 TableColumnConfig::new(
2094 "ts",
2095 DataType::Timestamp(TimeUnit::Nanosecond, None),
2096 false,
2097 ),
2098 ],
2099 vec!["id".to_string()],
2100 vec![],
2101 )
2102 .expect("Nanosecond timestamp should be accepted");
2103 let schema = config.to_schema();
2104 assert!(matches!(
2105 schema.field(1).data_type(),
2106 DataType::Timestamp(TimeUnit::Microsecond, _)
2107 ));
2108 }
2109
2110 #[test]
2111 fn table_config_accepts_decimal128_column() {
2112 let config = KvTableConfig::new(
2113 0,
2114 vec![
2115 TableColumnConfig::new("id", DataType::Int64, false),
2116 TableColumnConfig::new("price", DataType::Decimal128(10, 2), false),
2117 ],
2118 vec!["id".to_string()],
2119 vec![],
2120 )
2121 .expect("Decimal128 column should be accepted");
2122 assert_eq!(config.columns.len(), 2);
2123 }
2124
2125 #[test]
2126 fn table_config_accepts_list_column() {
2127 use datafusion::arrow::datatypes::Field;
2128
2129 let config = KvTableConfig::new(
2130 0,
2131 vec![
2132 TableColumnConfig::new("id", DataType::Int64, false),
2133 TableColumnConfig::new(
2134 "tags",
2135 DataType::List(Arc::new(Field::new("item", DataType::Utf8, false))),
2136 false,
2137 ),
2138 ],
2139 vec!["id".to_string()],
2140 vec![],
2141 )
2142 .expect("List<Utf8> column should be accepted");
2143 assert_eq!(config.columns.len(), 2);
2144 }
2145
2146 #[test]
2147 fn list_column_rejected_in_index() {
2148 use datafusion::arrow::datatypes::Field;
2149
2150 let result = KvTableConfig::new(
2151 0,
2152 vec![
2153 TableColumnConfig::new("id", DataType::Int64, false),
2154 TableColumnConfig::new(
2155 "tags",
2156 DataType::List(Arc::new(Field::new("item", DataType::Utf8, false))),
2157 false,
2158 ),
2159 ],
2160 vec!["id".to_string()],
2161 vec![IndexSpec::new("tags_idx", vec!["tags".to_string()]).unwrap()],
2162 );
2163 assert!(
2164 result.is_err() || {
2165 let config = result.unwrap();
2166 let model = TableModel::from_config(&config).unwrap();
2167 model.resolve_index_specs(&config.index_specs).is_err()
2168 }
2169 );
2170 }
2171
2172 #[test]
2173 fn i32_ordered_encoding_round_trip() {
2174 let values = [i32::MIN, -1000, -1, 0, 1, 1000, i32::MAX];
2175 for v in values {
2176 assert_eq!(decode_i32_ordered(encode_i32_ordered(v)), v);
2177 }
2178 let encoded: Vec<[u8; 4]> = values.iter().map(|v| encode_i32_ordered(*v)).collect();
2179 for i in 0..encoded.len() - 1 {
2180 assert!(encoded[i] < encoded[i + 1]);
2181 }
2182 }
2183
2184 #[test]
2185 fn i128_ordered_encoding_round_trip() {
2186 let values = [i128::MIN, -1, 0, 1, 1234567890123456789, i128::MAX];
2187 for v in values {
2188 assert_eq!(decode_i128_ordered(encode_i128_ordered(v)), v);
2189 }
2190 let encoded: Vec<[u8; 16]> = values.iter().map(|v| encode_i128_ordered(*v)).collect();
2191 for i in 0..encoded.len() - 1 {
2192 assert!(encoded[i] < encoded[i + 1]);
2193 }
2194 }
2195
2196 fn extended_model() -> (TableModel, Vec<ResolvedIndexSpec>) {
2197 let config = KvTableConfig::new(
2198 0,
2199 vec![
2200 TableColumnConfig::new("id", DataType::Int64, false),
2201 TableColumnConfig::new("created", DataType::Date32, false),
2202 TableColumnConfig::new(
2203 "ts",
2204 DataType::Timestamp(TimeUnit::Microsecond, None),
2205 false,
2206 ),
2207 TableColumnConfig::new("price", DataType::Decimal128(10, 2), false),
2208 TableColumnConfig::new("label", DataType::Utf8, false),
2209 ],
2210 vec!["id".to_string()],
2211 vec![
2212 IndexSpec::new(
2213 "date_label",
2214 vec!["created".to_string(), "label".to_string()],
2215 )
2216 .expect("valid"),
2217 IndexSpec::new("price_idx", vec!["price".to_string()]).expect("valid"),
2218 ],
2219 )
2220 .expect("valid config");
2221 let model = TableModel::from_config(&config).expect("model");
2222 let specs = model
2223 .resolve_index_specs(&config.index_specs)
2224 .expect("specs");
2225 (model, specs)
2226 }
2227
2228 #[test]
2229 fn secondary_index_key_round_trip_date32_and_decimal128() {
2230 let (model, specs) = extended_model();
2231 let row = KvRow {
2232 values: vec![
2233 CellValue::Int64(42),
2234 CellValue::Date32(19000),
2235 CellValue::Timestamp(1_700_000_000_000_000),
2236 CellValue::Decimal128(123456),
2237 CellValue::Utf8("hello".to_string()),
2238 ],
2239 };
2240 let key = encode_secondary_index_key(model.table_prefix, &specs[0], &model, &row)
2241 .expect("encode");
2242 let decoded = decode_secondary_index_key(model.table_prefix, &specs[0], &model, &key)
2243 .expect("decode");
2244 let created_idx = *model.columns_by_name.get("created").unwrap();
2245 let label_idx = *model.columns_by_name.get("label").unwrap();
2246 assert!(matches!(
2247 decoded.values.get(&created_idx),
2248 Some(CellValue::Date32(19000))
2249 ));
2250 assert!(matches!(
2251 decoded.values.get(&label_idx),
2252 Some(CellValue::Utf8(v)) if v == "hello"
2253 ));
2254 assert!(matches!(
2255 &decoded.primary_key_values[0],
2256 CellValue::Int64(42)
2257 ));
2258
2259 let key2 = encode_secondary_index_key(model.table_prefix, &specs[1], &model, &row)
2260 .expect("encode");
2261 let decoded2 = decode_secondary_index_key(model.table_prefix, &specs[1], &model, &key2)
2262 .expect("decode");
2263 let price_idx = *model.columns_by_name.get("price").unwrap();
2264 assert!(matches!(
2265 decoded2.values.get(&price_idx),
2266 Some(CellValue::Decimal128(123456))
2267 ));
2268 assert!(matches!(
2269 &decoded2.primary_key_values[0],
2270 CellValue::Int64(42)
2271 ));
2272 }
2273
2274 #[test]
2275 fn base_row_round_trip_with_date32_timestamp_decimal128() {
2276 let (model, _specs) = extended_model();
2277 let row = KvRow {
2278 values: vec![
2279 CellValue::Int64(7),
2280 CellValue::Date32(19500),
2281 CellValue::Timestamp(1_700_000_000_000_000),
2282 CellValue::Decimal128(-9876543),
2283 CellValue::Utf8("world".to_string()),
2284 ],
2285 };
2286 let encoded = encode_base_row_value(&row, &model).expect("encode");
2287 let decoded = decode_base_row(vec![CellValue::Int64(7)], &encoded, &model).expect("decode");
2288 assert!(matches!(&decoded.values[0], CellValue::Int64(7)));
2289 assert!(matches!(&decoded.values[1], CellValue::Date32(19500)));
2290 assert!(matches!(
2291 &decoded.values[2],
2292 CellValue::Timestamp(1_700_000_000_000_000)
2293 ));
2294 assert!(matches!(
2295 &decoded.values[3],
2296 CellValue::Decimal128(-9876543)
2297 ));
2298 assert!(matches!(&decoded.values[4], CellValue::Utf8(v) if v == "world"));
2299 }
2300
2301 #[test]
2302 fn base_row_round_trip_with_list() {
2303 use datafusion::arrow::datatypes::Field;
2304
2305 let config = KvTableConfig::new(
2306 0,
2307 vec![
2308 TableColumnConfig::new("id", DataType::Int64, false),
2309 TableColumnConfig::new(
2310 "tags",
2311 DataType::List(Arc::new(Field::new("item", DataType::Utf8, false))),
2312 false,
2313 ),
2314 TableColumnConfig::new(
2315 "scores",
2316 DataType::List(Arc::new(Field::new("item", DataType::Int64, false))),
2317 false,
2318 ),
2319 ],
2320 vec!["id".to_string()],
2321 vec![],
2322 )
2323 .expect("valid");
2324 let model = TableModel::from_config(&config).expect("model");
2325 let row = KvRow {
2326 values: vec![
2327 CellValue::Int64(1),
2328 CellValue::List(vec![
2329 CellValue::Utf8("a".to_string()),
2330 CellValue::Utf8("b".to_string()),
2331 ]),
2332 CellValue::List(vec![CellValue::Int64(10), CellValue::Int64(20)]),
2333 ],
2334 };
2335 let encoded = encode_base_row_value(&row, &model).expect("encode");
2336 let decoded = decode_base_row(vec![CellValue::Int64(1)], &encoded, &model).expect("decode");
2337 assert!(matches!(&decoded.values[0], CellValue::Int64(1)));
2338 match &decoded.values[1] {
2339 CellValue::List(items) => {
2340 assert_eq!(items.len(), 2);
2341 assert!(matches!(&items[0], CellValue::Utf8(v) if v == "a"));
2342 assert!(matches!(&items[1], CellValue::Utf8(v) if v == "b"));
2343 }
2344 _ => panic!("expected List"),
2345 }
2346 match &decoded.values[2] {
2347 CellValue::List(items) => {
2348 assert_eq!(items.len(), 2);
2349 assert!(matches!(&items[0], CellValue::Int64(10)));
2350 assert!(matches!(&items[1], CellValue::Int64(20)));
2351 }
2352 _ => panic!("expected List"),
2353 }
2354 }
2355
2356 #[test]
2357 fn decimal128_constraint_range() {
2358 let mut min: Option<i128> = None;
2359 let mut max: Option<i128> = None;
2360 let mut contradiction = false;
2361 apply_decimal128_constraint(&mut min, &mut max, Operator::GtEq, 100, &mut contradiction);
2362 assert!(!contradiction);
2363 apply_decimal128_constraint(&mut min, &mut max, Operator::LtEq, 200, &mut contradiction);
2364 assert!(!contradiction);
2365 assert_eq!(min, Some(100));
2366 assert_eq!(max, Some(200));
2367 assert!(in_i128_bounds(150, min, max));
2368 assert!(!in_i128_bounds(99, min, max));
2369 assert!(!in_i128_bounds(201, min, max));
2370 }
2371
2372 #[test]
2373 fn decimal256_gt_max_is_contradiction() {
2374 let mut min: Option<i256> = None;
2375 let mut max: Option<i256> = None;
2376 let mut contradiction = false;
2377 apply_i256_constraint(
2378 &mut min,
2379 &mut max,
2380 Operator::Gt,
2381 i256::MAX,
2382 &mut contradiction,
2383 );
2384 assert!(contradiction);
2385 assert_eq!(min, None);
2386 assert_eq!(max, None);
2387 }
2388
2389 #[test]
2390 fn decimal256_lt_min_is_contradiction() {
2391 let mut min: Option<i256> = None;
2392 let mut max: Option<i256> = None;
2393 let mut contradiction = false;
2394 apply_i256_constraint(
2395 &mut min,
2396 &mut max,
2397 Operator::Lt,
2398 i256::MIN,
2399 &mut contradiction,
2400 );
2401 assert!(contradiction);
2402 assert_eq!(min, None);
2403 assert_eq!(max, None);
2404 }
2405
2406 #[test]
2407 fn date32_index_bound_clamps_on_i64_overflow() {
2408 let config = KvTableConfig::new(
2409 0,
2410 vec![
2411 TableColumnConfig::new("id", DataType::Int64, false),
2412 TableColumnConfig::new("created", DataType::Date32, false),
2413 ],
2414 vec!["id".to_string()],
2415 vec![IndexSpec::new("created_idx", vec!["created".to_string()]).unwrap()],
2416 )
2417 .unwrap();
2418 let model = TableModel::from_config(&config).unwrap();
2419 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
2420
2421 let created_idx = *model.columns_by_name.get("created").unwrap();
2422 let mut pred = QueryPredicate::default();
2423 pred.constraints.insert(
2424 created_idx,
2425 PredicateConstraint::IntRange {
2426 min: Some(i32::MAX as i64 + 1),
2427 max: None,
2428 },
2429 );
2430
2431 let start = pred
2432 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, false)
2433 .unwrap();
2434 let end = pred
2435 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, true)
2436 .unwrap();
2437
2438 assert!(
2439 start <= end,
2440 "lower bound must not exceed upper bound (was wrapping via as i32)"
2441 );
2442
2443 let lower_payload = specs[0].prefix.strip(&start).unwrap();
2444 let encoded_lower: [u8; 4] = lower_payload[..4].try_into().unwrap();
2445 let decoded_lower = decode_i32_ordered(encoded_lower);
2446 assert_eq!(
2447 decoded_lower,
2448 i32::MAX,
2449 "out-of-range i64 must clamp to i32::MAX, not wrap"
2450 );
2451 }
2452
2453 #[test]
2454 fn overlong_utf8_index_prefix_errors_instead_of_panicking() {
2455 let config = KvTableConfig::new(
2456 0,
2457 vec![
2458 TableColumnConfig::new("s", DataType::Utf8, false),
2459 TableColumnConfig::new("b", DataType::FixedSizeBinary(16), false),
2460 TableColumnConfig::new("pk", DataType::Int64, false),
2461 ],
2462 vec!["pk".to_string()],
2463 vec![IndexSpec::new("s_b", vec!["s".to_string(), "b".to_string()]).unwrap()],
2464 )
2465 .unwrap();
2466 let model = TableModel::from_config(&config).unwrap();
2467 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
2468
2469 let s_idx = *model.columns_by_name.get("s").unwrap();
2474 let mut pred = QueryPredicate::default();
2475 pred.constraints
2476 .insert(s_idx, PredicateConstraint::StringEq("a".repeat(245)));
2477 assert!(pred
2478 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, true)
2479 .is_err());
2480 assert!(pred
2481 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, false)
2482 .is_err());
2483 }
2484
2485 #[test]
2486 fn overlong_utf8_index_prefix_with_fixed_pk_errors() {
2487 let config = KvTableConfig::new(
2491 0,
2492 vec![
2493 TableColumnConfig::new("s", DataType::Utf8, false),
2494 TableColumnConfig::new("pk", DataType::FixedSizeBinary(16), false),
2495 ],
2496 vec!["pk".to_string()],
2497 vec![IndexSpec::new("s_idx", vec!["s".to_string()]).unwrap()],
2498 )
2499 .unwrap();
2500 let model = TableModel::from_config(&config).unwrap();
2501 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
2502
2503 let s_idx = *model.columns_by_name.get("s").unwrap();
2504 let mut pred = QueryPredicate::default();
2505 pred.constraints
2506 .insert(s_idx, PredicateConstraint::StringEq("a".repeat(245)));
2507 assert!(pred
2508 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, true)
2509 .is_err());
2510 }
2511
2512 #[test]
2513 fn short_utf8_index_prefix_still_produces_valid_bounds() {
2514 let config = KvTableConfig::new(
2515 0,
2516 vec![
2517 TableColumnConfig::new("s", DataType::Utf8, false),
2518 TableColumnConfig::new("b", DataType::FixedSizeBinary(16), false),
2519 TableColumnConfig::new("pk", DataType::Int64, false),
2520 ],
2521 vec!["pk".to_string()],
2522 vec![IndexSpec::new("s_b", vec!["s".to_string(), "b".to_string()]).unwrap()],
2523 )
2524 .unwrap();
2525 let model = TableModel::from_config(&config).unwrap();
2526 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
2527
2528 let s_idx = *model.columns_by_name.get("s").unwrap();
2531 let mut pred = QueryPredicate::default();
2532 pred.constraints
2533 .insert(s_idx, PredicateConstraint::StringEq("hello".to_string()));
2534 let start = pred
2535 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, false)
2536 .expect("lower bound should encode");
2537 let end = pred
2538 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, true)
2539 .expect("upper bound should encode");
2540 assert!(start <= end, "lower bound must not exceed upper bound");
2541 }
2542
2543 #[test]
2550 fn index_lower_bound_includes_boundary_row() {
2551 let config = KvTableConfig::new(
2552 0,
2553 vec![
2554 TableColumnConfig::new("s", DataType::Utf8, false),
2555 TableColumnConfig::new("id", DataType::Int64, false),
2556 ],
2557 vec!["id".to_string()],
2558 vec![IndexSpec::new("s_idx", vec!["s".to_string()]).unwrap()],
2559 )
2560 .unwrap();
2561 let model = TableModel::from_config(&config).unwrap();
2562 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
2563
2564 let s_idx = *model.columns_by_name.get("s").unwrap();
2566 let id_idx = *model.columns_by_name.get("id").unwrap();
2567 let mut pred = QueryPredicate::default();
2568 pred.constraints
2569 .insert(s_idx, PredicateConstraint::StringEq("a".to_string()));
2570 pred.constraints.insert(
2571 id_idx,
2572 PredicateConstraint::IntRange {
2573 min: Some(5),
2574 max: None,
2575 },
2576 );
2577 let start = pred
2578 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, false)
2579 .expect("lower bound");
2580 let end = pred
2581 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, true)
2582 .expect("upper bound");
2583
2584 let mut payload = Vec::new();
2587 payload.extend_from_slice(&encode_string_variable("a").unwrap());
2588 payload.extend_from_slice(&encode_i64_ordered(5));
2589 let row_key = specs[0].prefix.encode(&payload).unwrap();
2590
2591 assert!(start <= end);
2592 assert!(
2593 row_key >= start && row_key <= end,
2594 "boundary row must fall inside the index range"
2595 );
2596 }
2597
2598 #[test]
2599 fn long_utf8_equality_produces_usable_index_bounds() {
2600 let config = KvTableConfig::new(
2601 0,
2602 vec![
2603 TableColumnConfig::new("s", DataType::Utf8, false),
2604 TableColumnConfig::new("id", DataType::Int64, false),
2605 ],
2606 vec!["id".to_string()],
2607 vec![IndexSpec::new("s_idx", vec!["s".to_string()]).unwrap()],
2608 )
2609 .unwrap();
2610 let model = TableModel::from_config(&config).unwrap();
2611 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
2612
2613 let needle = "a".repeat(30);
2617 let s_idx = *model.columns_by_name.get("s").unwrap();
2618 let mut pred = QueryPredicate::default();
2619 pred.constraints
2620 .insert(s_idx, PredicateConstraint::StringEq(needle.clone()));
2621 let start = pred
2622 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, false)
2623 .expect("lower bound");
2624 let end = pred
2625 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, true)
2626 .expect("upper bound");
2627
2628 let mut payload = Vec::new();
2629 payload.extend_from_slice(&encode_string_variable(&needle).unwrap());
2630 payload.extend_from_slice(&encode_i64_ordered(7));
2631 let row_key = specs[0].prefix.encode(&payload).unwrap();
2632
2633 assert!(start <= end);
2634 assert!(row_key >= start && row_key <= end);
2635 }
2636
2637 #[test]
2638 fn index_lower_bound_includes_empty_utf8_pk_row() {
2639 let config = KvTableConfig::new(
2640 0,
2641 vec![
2642 TableColumnConfig::new("a", DataType::Int64, false),
2643 TableColumnConfig::new("k", DataType::Utf8, false),
2644 ],
2645 vec!["k".to_string()],
2646 vec![IndexSpec::new("a_idx", vec!["a".to_string()]).unwrap()],
2647 )
2648 .unwrap();
2649 let model = TableModel::from_config(&config).unwrap();
2650 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
2651
2652 let a_idx = *model.columns_by_name.get("a").unwrap();
2655 let mut pred = QueryPredicate::default();
2656 pred.constraints.insert(
2657 a_idx,
2658 PredicateConstraint::IntRange {
2659 min: Some(5),
2660 max: Some(5),
2661 },
2662 );
2663 let start = pred
2664 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, false)
2665 .expect("lower bound");
2666 let end = pred
2667 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, true)
2668 .expect("upper bound");
2669
2670 let mut payload = Vec::new();
2671 payload.extend_from_slice(&encode_i64_ordered(5));
2672 payload.extend_from_slice(&encode_string_variable("").unwrap());
2673 let row_key = specs[0].prefix.encode(&payload).unwrap();
2674
2675 assert!(start <= end);
2676 assert!(row_key >= start && row_key <= end);
2677 }
2678
2679 #[test]
2680 fn zorder_index_range_includes_short_utf8_pk_rows() {
2681 let config = KvTableConfig::new(
2686 0,
2687 vec![
2688 TableColumnConfig::new("name", DataType::Utf8, false),
2689 TableColumnConfig::new("x", DataType::Int64, false),
2690 TableColumnConfig::new("y", DataType::Int64, false),
2691 TableColumnConfig::new("value", DataType::Int64, false),
2692 ],
2693 vec!["name".to_string()],
2694 vec![
2695 IndexSpec::z_order("xy_z", vec!["x".to_string(), "y".to_string()])
2696 .expect("valid")
2697 .with_cover_columns(vec!["value".to_string()]),
2698 ],
2699 )
2700 .expect("valid config");
2701 let model = TableModel::from_config(&config).expect("model");
2702 let specs = model
2703 .resolve_index_specs(&config.index_specs)
2704 .expect("specs");
2705 let x_idx = *model.columns_by_name.get("x").unwrap();
2706 let y_idx = *model.columns_by_name.get("y").unwrap();
2707
2708 let mut pred = QueryPredicate::default();
2710 pred.constraints.insert(
2711 x_idx,
2712 PredicateConstraint::IntRange {
2713 min: Some(5),
2714 max: Some(5),
2715 },
2716 );
2717 pred.constraints.insert(
2718 y_idx,
2719 PredicateConstraint::IntRange {
2720 min: Some(7),
2721 max: Some(7),
2722 },
2723 );
2724 let ranges = pred
2725 .expand_zorder_index_ranges(model.table_prefix, &model, &specs[0])
2726 .expect("ranges");
2727 assert_eq!(ranges.len(), 1);
2728
2729 for name in ["", "bob"] {
2732 let row = KvRow {
2733 values: vec![
2734 CellValue::Utf8(name.to_string()),
2735 CellValue::Int64(5),
2736 CellValue::Int64(7),
2737 CellValue::Int64(100),
2738 ],
2739 };
2740 let stored = encode_secondary_index_key(model.table_prefix, &specs[0], &model, &row)
2741 .expect("stored");
2742 assert!(
2743 ranges[0].start <= stored && stored <= ranges[0].end,
2744 "zorder range for x=5,y=7 must include stored entry with pk '{name}'"
2745 );
2746 }
2747 }
2748
2749 #[test]
2750 fn utf8_pk_behind_lex_index_range_includes_short_pk_rows() {
2751 let config = KvTableConfig::new(
2756 0,
2757 vec![
2758 TableColumnConfig::new("name", DataType::Utf8, false),
2759 TableColumnConfig::new("age", DataType::Int64, false),
2760 ],
2761 vec!["name".to_string()],
2762 vec![IndexSpec::new("age_idx", vec!["age".to_string()]).unwrap()],
2763 )
2764 .unwrap();
2765 let model = TableModel::from_config(&config).unwrap();
2766 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
2767 let age_idx = *model.columns_by_name.get("age").unwrap();
2768
2769 let mut pred = QueryPredicate::default();
2770 pred.constraints.insert(
2771 age_idx,
2772 PredicateConstraint::IntRange {
2773 min: Some(5),
2774 max: Some(5),
2775 },
2776 );
2777 let start = pred
2778 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, false)
2779 .unwrap();
2780 let end = pred
2781 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, true)
2782 .unwrap();
2783
2784 for name in ["", "bob"] {
2785 let row = KvRow {
2786 values: vec![CellValue::Utf8(name.to_string()), CellValue::Int64(5)],
2787 };
2788 let stored = encode_secondary_index_key(model.table_prefix, &specs[0], &model, &row)
2789 .expect("stored index key encodes");
2790 assert!(
2791 start <= stored && stored <= end,
2792 "index range for age=5 must include stored entry with pk '{name}'"
2793 );
2794 }
2795 }
2796
2797 #[test]
2798 fn multi_column_index_utf8_prefix_range_includes_min_suffix_rows() {
2799 let config = KvTableConfig::new(
2803 0,
2804 vec![
2805 TableColumnConfig::new("id", DataType::Int64, false),
2806 TableColumnConfig::new("a", DataType::Utf8, false),
2807 TableColumnConfig::new("b", DataType::Int64, false),
2808 ],
2809 vec!["id".to_string()],
2810 vec![IndexSpec::new("ab_idx", vec!["a".to_string(), "b".to_string()]).unwrap()],
2811 )
2812 .unwrap();
2813 let model = TableModel::from_config(&config).unwrap();
2814 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
2815 let a_idx = *model.columns_by_name.get("a").unwrap();
2816
2817 let mut pred = QueryPredicate::default();
2818 pred.constraints
2819 .insert(a_idx, PredicateConstraint::StringEq("x".to_string()));
2820 let start = pred
2821 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, false)
2822 .unwrap();
2823 let end = pred
2824 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, true)
2825 .unwrap();
2826
2827 for (b, id) in [(i64::MIN, i64::MIN), (0, 0), (i64::MAX, i64::MAX)] {
2828 let row = KvRow {
2829 values: vec![
2830 CellValue::Int64(id),
2831 CellValue::Utf8("x".to_string()),
2832 CellValue::Int64(b),
2833 ],
2834 };
2835 let stored = encode_secondary_index_key(model.table_prefix, &specs[0], &model, &row)
2836 .expect("stored index key encodes");
2837 assert!(
2838 start <= stored && stored <= end,
2839 "index range for a='x' must include stored entry with b={b}, id={id}"
2840 );
2841 }
2842
2843 let row = KvRow {
2844 values: vec![
2845 CellValue::Int64(0),
2846 CellValue::Utf8("y".to_string()),
2847 CellValue::Int64(0),
2848 ],
2849 };
2850 let excluded = encode_secondary_index_key(model.table_prefix, &specs[0], &model, &row)
2851 .expect("stored index key encodes");
2852 assert!(
2853 excluded > end,
2854 "index range for a='x' must exclude entries for a='y'"
2855 );
2856 }
2857
2858 #[test]
2859 fn variable_width_column_rejected_in_zorder_index() {
2860 let config = KvTableConfig::new(
2865 0,
2866 vec![
2867 TableColumnConfig::new("id", DataType::Int64, false),
2868 TableColumnConfig::new("x", DataType::Int64, false),
2869 TableColumnConfig::new("region", DataType::Utf8, false),
2870 ],
2871 vec!["id".to_string()],
2872 vec![IndexSpec::z_order("xr_z", vec!["x".to_string(), "region".to_string()]).unwrap()],
2873 )
2874 .unwrap();
2875 let model = TableModel::from_config(&config).unwrap();
2876 let result = model.resolve_index_specs(&config.index_specs);
2877 assert!(result.is_err());
2878 assert!(result.unwrap_err().contains("z-order key column"));
2879
2880 let lex_config = KvTableConfig::new(
2882 0,
2883 vec![
2884 TableColumnConfig::new("id", DataType::Int64, false),
2885 TableColumnConfig::new("x", DataType::Int64, false),
2886 TableColumnConfig::new("region", DataType::Utf8, false),
2887 ],
2888 vec!["id".to_string()],
2889 vec![IndexSpec::new("r_idx", vec!["region".to_string()]).unwrap()],
2890 )
2891 .unwrap();
2892 let lex_model = TableModel::from_config(&lex_config).unwrap();
2893 assert!(lex_model
2894 .resolve_index_specs(&lex_config.index_specs)
2895 .is_ok());
2896 }
2897
2898 #[test]
2899 fn utf8_pk_point_query_range_includes_short_string_row() {
2900 let config = KvTableConfig::new(
2905 0,
2906 vec![
2907 TableColumnConfig::new("name", DataType::Utf8, false),
2908 TableColumnConfig::new("age", DataType::Int64, false),
2909 ],
2910 vec!["name".to_string()],
2911 vec![],
2912 )
2913 .unwrap();
2914 let model = TableModel::from_config(&config).unwrap();
2915 let name_idx = *model.columns_by_name.get("name").unwrap();
2916
2917 for name in [
2918 "alice",
2919 "",
2920 "a",
2921 "exactly15chars!",
2922 "longer-than-sixteen-bytes",
2923 ] {
2924 let value = CellValue::Utf8(name.to_string());
2925 let stored = encode_primary_key(0, &[&value], &model).expect("stored key encodes");
2926
2927 let mut pred = QueryPredicate::default();
2928 pred.constraints
2929 .insert(name_idx, PredicateConstraint::StringEq(name.to_string()));
2930 let ranges = pred.primary_key_ranges(&model).unwrap();
2931 assert_eq!(
2932 ranges.len(),
2933 1,
2934 "point query on '{name}' must yield one range"
2935 );
2936 assert!(
2937 ranges[0].start <= stored,
2938 "lower bound must not exceed stored key for '{name}'"
2939 );
2940 assert!(
2941 stored <= ranges[0].end,
2942 "stored key must not exceed upper bound for '{name}'"
2943 );
2944 assert_eq!(
2945 ranges[0].start, stored,
2946 "point lower bound must equal the stored key for '{name}'"
2947 );
2948 }
2949
2950 let mut pred = QueryPredicate::default();
2952 pred.constraints
2953 .insert(name_idx, PredicateConstraint::StringEq("al".to_string()));
2954 let ranges = pred.primary_key_ranges(&model).unwrap();
2955 let alice = CellValue::Utf8("alice".to_string());
2956 let alice_key = encode_primary_key(0, &[&alice], &model).expect("stored key encodes");
2957 assert!(
2958 alice_key > ranges[0].end,
2959 "range for 'al' must exclude the stored key for 'alice'"
2960 );
2961 }
2962
2963 #[test]
2964 fn utf8_pk_prefix_query_range_includes_composite_rows() {
2965 let config = KvTableConfig::new(
2969 0,
2970 vec![
2971 TableColumnConfig::new("name", DataType::Utf8, false),
2972 TableColumnConfig::new("id", DataType::Int64, false),
2973 ],
2974 vec!["name".to_string(), "id".to_string()],
2975 vec![],
2976 )
2977 .unwrap();
2978 let model = TableModel::from_config(&config).unwrap();
2979 let name_idx = *model.columns_by_name.get("name").unwrap();
2980
2981 let mut pred = QueryPredicate::default();
2982 pred.constraints
2983 .insert(name_idx, PredicateConstraint::StringEq("al".to_string()));
2984 let ranges = pred.primary_key_ranges(&model).unwrap();
2985 assert_eq!(ranges.len(), 1);
2986
2987 for id in [i64::MIN, 0, 42, i64::MAX] {
2988 let name = CellValue::Utf8("al".to_string());
2989 let id_value = CellValue::Int64(id);
2990 let stored =
2991 encode_primary_key(0, &[&name, &id_value], &model).expect("stored key encodes");
2992 assert!(
2993 ranges[0].start <= stored && stored <= ranges[0].end,
2994 "range for name='al' must include stored key with id={id}"
2995 );
2996 }
2997
2998 let other = CellValue::Utf8("alx".to_string());
2999 let id_value = CellValue::Int64(0);
3000 let excluded =
3001 encode_primary_key(0, &[&other, &id_value], &model).expect("stored key encodes");
3002 assert!(
3003 excluded > ranges[0].end,
3004 "range for name='al' must exclude rows with name='alx'"
3005 );
3006 }
3007
3008 #[test]
3009 fn utf8_index_point_query_range_includes_min_pk_row() {
3010 let config = KvTableConfig::new(
3016 0,
3017 vec![
3018 TableColumnConfig::new("id", DataType::Int64, false),
3019 TableColumnConfig::new("region", DataType::Utf8, false),
3020 ],
3021 vec!["id".to_string()],
3022 vec![IndexSpec::new("region_idx", vec!["region".to_string()]).unwrap()],
3023 )
3024 .unwrap();
3025 let model = TableModel::from_config(&config).unwrap();
3026 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
3027 let region_idx = *model.columns_by_name.get("region").unwrap();
3028
3029 let mut pred = QueryPredicate::default();
3030 pred.constraints
3031 .insert(region_idx, PredicateConstraint::StringEq("us".to_string()));
3032 let start = pred
3033 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, false)
3034 .unwrap();
3035 let end = pred
3036 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 1, true)
3037 .unwrap();
3038
3039 for id in [i64::MIN, 0, i64::MAX] {
3040 let row = KvRow {
3041 values: vec![CellValue::Int64(id), CellValue::Utf8("us".to_string())],
3042 };
3043 let stored = encode_secondary_index_key(model.table_prefix, &specs[0], &model, &row)
3044 .expect("stored index key encodes");
3045 assert!(
3046 start <= stored && stored <= end,
3047 "index range for region='us' must include stored entry with id={id}"
3048 );
3049 }
3050
3051 let row = KvRow {
3052 values: vec![CellValue::Int64(0), CellValue::Utf8("ut".to_string())],
3053 };
3054 let excluded = encode_secondary_index_key(model.table_prefix, &specs[0], &model, &row)
3055 .expect("stored index key encodes");
3056 assert!(
3057 excluded > end,
3058 "index range for region='us' must exclude entries for region='ut'"
3059 );
3060 }
3061
3062 #[tokio::test]
3063 async fn utf8_pk_point_query_returns_row() {
3064 let state = MockState {
3068 kv: Arc::new(Mutex::new(BTreeMap::new())),
3069 range_calls: Arc::new(AtomicUsize::new(0)),
3070 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
3071 sequence_number: Arc::new(AtomicU64::new(0)),
3072 };
3073 let (base_url, shutdown_tx) = spawn_mock_server(state).await;
3074 let client = StoreClient::new(&base_url);
3075
3076 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
3077 .table(
3078 "users",
3079 vec![
3080 TableColumnConfig::new("name", DataType::Utf8, false),
3081 TableColumnConfig::new("age", DataType::Int64, false),
3082 ],
3083 vec!["name".to_string()],
3084 vec![],
3085 )
3086 .expect("schema");
3087 let mut writer = schema.batch_writer();
3088 for (name, age) in [("alice", 30i64), ("bob", 25), ("", 99)] {
3089 writer
3090 .insert(
3091 "users",
3092 vec![CellValue::Utf8(name.to_string()), CellValue::Int64(age)],
3093 )
3094 .expect("row");
3095 }
3096 writer.flush().await.expect("flush");
3097
3098 let ctx = SessionContext::new();
3099 schema.register_all(&ctx).expect("register");
3100
3101 for (name, expected_age) in [("alice", 30i64), ("bob", 25), ("", 99)] {
3102 let df = ctx
3103 .sql(&format!("SELECT age FROM users WHERE name = '{name}'"))
3104 .await
3105 .expect("plan");
3106 let batches = df.collect().await.expect("collect");
3107 let ages: Vec<i64> = batches
3108 .iter()
3109 .flat_map(|b| {
3110 b.column(0)
3111 .as_any()
3112 .downcast_ref::<datafusion::arrow::array::Int64Array>()
3113 .unwrap()
3114 .iter()
3115 .map(|v| v.unwrap())
3116 })
3117 .collect();
3118 assert_eq!(
3119 ages,
3120 vec![expected_age],
3121 "point query on name='{name}' must return exactly its row"
3122 );
3123 }
3124
3125 let df = ctx
3126 .sql("SELECT age FROM users WHERE name = 'al'")
3127 .await
3128 .expect("plan");
3129 let batches = df.collect().await.expect("collect");
3130 let total: usize = batches.iter().map(|b| b.num_rows()).sum();
3131 assert_eq!(total, 0, "prefix of a stored name must match nothing");
3132
3133 let _ = shutdown_tx.send(());
3134 }
3135
3136 #[tokio::test]
3137 async fn utf8_index_point_query_returns_min_pk_row() {
3138 let state = MockState {
3143 kv: Arc::new(Mutex::new(BTreeMap::new())),
3144 range_calls: Arc::new(AtomicUsize::new(0)),
3145 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
3146 sequence_number: Arc::new(AtomicU64::new(0)),
3147 };
3148 let (base_url, shutdown_tx) = spawn_mock_server(state).await;
3149 let client = StoreClient::new(&base_url);
3150
3151 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
3152 .table(
3153 "orders_nc",
3154 vec![
3155 TableColumnConfig::new("id", DataType::Int64, false),
3156 TableColumnConfig::new("region", DataType::Utf8, false),
3157 TableColumnConfig::new("amount", DataType::Int64, false),
3158 ],
3159 vec!["id".to_string()],
3160 vec![IndexSpec::new("region_idx", vec!["region".to_string()]).expect("valid")],
3161 )
3162 .expect("schema")
3163 .table(
3164 "orders_cov",
3165 vec![
3166 TableColumnConfig::new("id", DataType::Int64, false),
3167 TableColumnConfig::new("region", DataType::Utf8, false),
3168 TableColumnConfig::new("amount", DataType::Int64, false),
3169 ],
3170 vec!["id".to_string()],
3171 vec![IndexSpec::new("region_idx", vec!["region".to_string()])
3172 .expect("valid")
3173 .with_cover_columns(vec!["amount".to_string()])],
3174 )
3175 .expect("schema");
3176 let mut writer = schema.batch_writer();
3177 for table in ["orders_nc", "orders_cov"] {
3178 writer
3179 .insert(
3180 table,
3181 vec![
3182 CellValue::Int64(i64::MIN),
3183 CellValue::Utf8("us".to_string()),
3184 CellValue::Int64(7),
3185 ],
3186 )
3187 .expect("row");
3188 writer
3189 .insert(
3190 table,
3191 vec![
3192 CellValue::Int64(1),
3193 CellValue::Utf8("eu".to_string()),
3194 CellValue::Int64(8),
3195 ],
3196 )
3197 .expect("row");
3198 }
3199 writer.flush().await.expect("flush");
3200
3201 let ctx = SessionContext::new();
3202 schema.register_all(&ctx).expect("register");
3203
3204 for table in ["orders_nc", "orders_cov"] {
3205 let df = ctx
3206 .sql(&format!("SELECT amount FROM {table} WHERE region = 'us'"))
3207 .await
3208 .expect("plan");
3209 let batches = df.collect().await.expect("collect");
3210 let amounts: Vec<i64> = batches
3211 .iter()
3212 .flat_map(|b| {
3213 b.column(0)
3214 .as_any()
3215 .downcast_ref::<datafusion::arrow::array::Int64Array>()
3216 .unwrap()
3217 .iter()
3218 .map(|v| v.unwrap())
3219 })
3220 .collect();
3221 assert_eq!(
3222 amounts,
3223 vec![7],
3224 "index lookup on {table} must return the id=i64::MIN row"
3225 );
3226 }
3227
3228 let _ = shutdown_tx.send(());
3229 }
3230
3231 #[test]
3232 fn timestamp_nanos_gt_uses_floor_division() {
3233 let micros = timestamp_scalar_to_micros_for_op(
3234 &ScalarValue::TimestampNanosecond(Some(-1500), None),
3235 Operator::Gt,
3236 )
3237 .unwrap();
3238 assert_eq!(micros, -2, "Gt on -1500ns should floor to -2us");
3239
3240 let mut min: Option<i64> = None;
3241 let mut max: Option<i64> = None;
3242 let mut contradiction = false;
3243 apply_int_constraint(&mut min, &mut max, Operator::Gt, micros, &mut contradiction);
3244 assert_eq!(min, Some(-1), "Gt(-2us) + 1 = min -1us");
3245
3246 let row_at_minus_1 = CellValue::Timestamp(-1);
3247 assert!(matches_constraint(
3248 &row_at_minus_1,
3249 &PredicateConstraint::IntRange { min, max }
3250 ));
3251 }
3252
3253 #[test]
3254 fn timestamp_nanos_lteq_uses_floor_division() {
3255 let micros = timestamp_scalar_to_micros_for_op(
3256 &ScalarValue::TimestampNanosecond(Some(-1500), None),
3257 Operator::LtEq,
3258 )
3259 .unwrap();
3260 assert_eq!(micros, -2, "LtEq on -1500ns should floor to -2us");
3261
3262 let row_at_minus_1 = CellValue::Timestamp(-1);
3263 assert!(
3264 !matches_constraint(
3265 &row_at_minus_1,
3266 &PredicateConstraint::IntRange {
3267 min: None,
3268 max: Some(micros)
3269 }
3270 ),
3271 "-1us (-1000ns) > -1500ns, must not satisfy <= -1500ns"
3272 );
3273 }
3274
3275 #[test]
3276 fn timestamp_nanos_gteq_uses_ceil_division() {
3277 let micros = timestamp_scalar_to_micros_for_op(
3278 &ScalarValue::TimestampNanosecond(Some(-1500), None),
3279 Operator::GtEq,
3280 )
3281 .unwrap();
3282 assert_eq!(micros, -1, "GtEq on -1500ns should ceil to -1us");
3283 }
3284
3285 #[test]
3286 fn timestamp_nanos_lt_uses_ceil_division() {
3287 let micros = timestamp_scalar_to_micros_for_op(
3288 &ScalarValue::TimestampNanosecond(Some(-1500), None),
3289 Operator::Lt,
3290 )
3291 .unwrap();
3292 assert_eq!(micros, -1, "Lt on -1500ns should ceil to -1us");
3293
3294 let mut min: Option<i64> = None;
3295 let mut max: Option<i64> = None;
3296 let mut contradiction = false;
3297 apply_int_constraint(&mut min, &mut max, Operator::Lt, micros, &mut contradiction);
3298 assert_eq!(max, Some(-2), "Lt(-1us) - 1 = max -2us");
3299 }
3300
3301 #[test]
3302 fn timestamp_nanos_eq_non_aligned_is_contradiction() {
3303 let result = timestamp_scalar_to_micros_for_op(
3304 &ScalarValue::TimestampNanosecond(Some(-1500), None),
3305 Operator::Eq,
3306 );
3307 assert!(
3308 result.is_none(),
3309 "non-aligned ns Eq must produce contradiction"
3310 );
3311 }
3312
3313 #[test]
3314 fn timestamp_nanos_exact_multiple_is_unchanged() {
3315 for op in [
3316 Operator::Eq,
3317 Operator::Gt,
3318 Operator::GtEq,
3319 Operator::Lt,
3320 Operator::LtEq,
3321 ] {
3322 let micros = timestamp_scalar_to_micros_for_op(
3323 &ScalarValue::TimestampNanosecond(Some(-2000), None),
3324 op,
3325 )
3326 .unwrap();
3327 assert_eq!(micros, -2, "exact multiple -2000ns = -2us for {op:?}");
3328 }
3329 }
3330
3331 #[test]
3332 fn float64_index_bounds_include_infinity() {
3333 let config = KvTableConfig::new(
3334 0,
3335 vec![
3336 TableColumnConfig::new("id", DataType::Int64, false),
3337 TableColumnConfig::new("val", DataType::Float64, false),
3338 ],
3339 vec!["id".to_string()],
3340 vec![IndexSpec::new("val_idx", vec!["val".to_string()]).unwrap()],
3341 )
3342 .unwrap();
3343 let model = TableModel::from_config(&config).unwrap();
3344 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
3345
3346 let pred = QueryPredicate::default();
3347 let start = pred
3348 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 0, false)
3349 .unwrap();
3350 let end = pred
3351 .encode_index_bound_key(model.table_prefix, &model, &specs[0], 0, true)
3352 .unwrap();
3353
3354 let neg_inf_row = KvRow {
3355 values: vec![CellValue::Int64(1), CellValue::Float64(f64::NEG_INFINITY)],
3356 };
3357 let pos_inf_row = KvRow {
3358 values: vec![CellValue::Int64(2), CellValue::Float64(f64::INFINITY)],
3359 };
3360
3361 let neg_inf_key =
3362 encode_secondary_index_key(model.table_prefix, &specs[0], &model, &neg_inf_row)
3363 .unwrap();
3364 let pos_inf_key =
3365 encode_secondary_index_key(model.table_prefix, &specs[0], &model, &pos_inf_row)
3366 .unwrap();
3367
3368 assert!(
3369 neg_inf_key >= start,
3370 "NEG_INFINITY row key must be within scan start bound"
3371 );
3372 assert!(
3373 pos_inf_key <= end,
3374 "INFINITY row key must be within scan end bound"
3375 );
3376 }
3377
3378 #[test]
3379 fn distinct_table_prefixes_produce_non_overlapping_pk_ranges() {
3380 let range_a = primary_key_prefix_range(1);
3381 let range_b = primary_key_prefix_range(2);
3382 assert!(
3383 range_a.end < range_b.start,
3384 "table prefix 1 pk range must be entirely below table prefix 2"
3385 );
3386 }
3387
3388 #[test]
3389 fn distinct_table_prefixes_isolate_primary_keys() {
3390 let model_1 = simple_int64_model(1);
3391 let model_2 = simple_int64_model(2);
3392 let pk = CellValue::Int64(42);
3393 let key_a = encode_primary_key(1, &[&pk], &model_1).expect("pk key encodes");
3394 let key_b = encode_primary_key(2, &[&pk], &model_2).expect("pk key encodes");
3395 assert_ne!(key_a, key_b, "same PK under different prefixes must differ");
3396 assert!(
3397 decode_primary_key(1, &key_a, &model_1).is_some(),
3398 "key_a must decode under prefix 1"
3399 );
3400 assert!(
3401 decode_primary_key(2, &key_a, &model_2).is_none(),
3402 "key_a must NOT decode under prefix 2"
3403 );
3404 assert!(
3405 decode_primary_key(2, &key_b, &model_2).is_some(),
3406 "key_b must decode under prefix 2"
3407 );
3408 }
3409
3410 #[test]
3411 fn distinct_table_prefixes_isolate_secondary_keys() {
3412 let config_a = KvTableConfig::new(
3413 10,
3414 vec![
3415 TableColumnConfig::new("id", DataType::Int64, false),
3416 TableColumnConfig::new("name", DataType::Utf8, false),
3417 ],
3418 vec!["id".to_string()],
3419 vec![IndexSpec::new("name_idx", vec!["name".to_string()]).unwrap()],
3420 )
3421 .unwrap();
3422 let config_b = KvTableConfig::new(
3423 11,
3424 vec![
3425 TableColumnConfig::new("id", DataType::Int64, false),
3426 TableColumnConfig::new("name", DataType::Utf8, false),
3427 ],
3428 vec!["id".to_string()],
3429 vec![IndexSpec::new("name_idx", vec!["name".to_string()]).unwrap()],
3430 )
3431 .unwrap();
3432
3433 let model_a = TableModel::from_config(&config_a).unwrap();
3434 let specs_a = model_a.resolve_index_specs(&config_a.index_specs).unwrap();
3435 let model_b = TableModel::from_config(&config_b).unwrap();
3436 let specs_b = model_b.resolve_index_specs(&config_b.index_specs).unwrap();
3437
3438 let row = KvRow {
3439 values: vec![CellValue::Int64(1), CellValue::Utf8("alice".to_string())],
3440 };
3441 let key_a =
3442 encode_secondary_index_key(model_a.table_prefix, &specs_a[0], &model_a, &row).unwrap();
3443 let key_b =
3444 encode_secondary_index_key(model_b.table_prefix, &specs_b[0], &model_b, &row).unwrap();
3445
3446 assert_ne!(
3447 key_a, key_b,
3448 "same row under different prefixes must differ"
3449 );
3450 assert!(
3451 decode_secondary_index_key(model_a.table_prefix, &specs_a[0], &model_a, &key_a)
3452 .is_some()
3453 );
3454 assert!(
3455 decode_secondary_index_key(model_a.table_prefix, &specs_a[0], &model_a, &key_b)
3456 .is_none(),
3457 "key from table B must not decode under table A's prefix"
3458 );
3459 }
3460
3461 #[test]
3462 fn table_prefix_stored_in_model() {
3463 let config = KvTableConfig::new(
3464 12,
3465 vec![TableColumnConfig::new("id", DataType::Int64, false)],
3466 vec!["id".to_string()],
3467 vec![],
3468 )
3469 .unwrap();
3470 let model = TableModel::from_config(&config).unwrap();
3471 assert_eq!(model.table_prefix, 12);
3472 }
3473
3474 #[test]
3475 fn codec_layout_places_payload_bytes_after_family_prefix() {
3476 let config = KvTableConfig::new(
3477 0,
3478 vec![
3479 TableColumnConfig::new("id", DataType::FixedSizeBinary(16), false),
3480 TableColumnConfig::new("bucket", DataType::FixedSizeBinary(16), false),
3481 ],
3482 vec!["id".to_string()],
3483 vec![IndexSpec::new("bucket_idx", vec!["bucket".to_string()]).unwrap()],
3484 )
3485 .unwrap();
3486 let model = TableModel::from_config(&config).unwrap();
3487 let spec = model
3488 .resolve_index_specs(&config.index_specs)
3489 .unwrap()
3490 .remove(0);
3491
3492 let primary_prefix = primary_key_prefix(model.table_prefix).unwrap();
3493 let index_prefix = secondary_index_prefix(model.table_prefix, spec.id).unwrap();
3494
3495 let mut primary_payload_bytes = HashSet::new();
3496 let mut secondary_payload_bytes = HashSet::new();
3497
3498 for first_byte in 0u8..=255 {
3499 let mut id = vec![0u8; 16];
3500 id[0] = first_byte;
3501 let mut bucket = vec![0u8; 16];
3502 bucket[0] = first_byte;
3503
3504 let pk = CellValue::FixedBinary(id.clone());
3505 let primary = encode_primary_key(model.table_prefix, &[&pk], &model).unwrap();
3506 assert_eq!(
3509 &primary[..FAMILY_PREFIX_LEN],
3510 primary_prefix.as_bytes().as_ref()
3511 );
3512 primary_payload_bytes.insert(primary[FAMILY_PREFIX_LEN]);
3513
3514 let row = KvRow {
3515 values: vec![
3516 CellValue::FixedBinary(id),
3517 CellValue::FixedBinary(bucket.clone()),
3518 ],
3519 };
3520 let index =
3521 encode_secondary_index_key(model.table_prefix, &spec, &model, &row).unwrap();
3522 assert_eq!(
3523 &index[..FAMILY_PREFIX_LEN],
3524 index_prefix.as_bytes().as_ref()
3525 );
3526 secondary_payload_bytes.insert(index[FAMILY_PREFIX_LEN]);
3527 }
3528
3529 assert_eq!(primary_payload_bytes.len(), 256);
3533 assert_eq!(secondary_payload_bytes.len(), 256);
3534 }
3535
3536 #[test]
3537 fn kv_schema_auto_assigns_sequential_prefixes() {
3538 let client = StoreClient::new("http://localhost:10000");
3539 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
3540 .table(
3541 "alpha",
3542 vec![TableColumnConfig::new("id", DataType::Int64, false)],
3543 vec!["id".to_string()],
3544 vec![],
3545 )
3546 .unwrap()
3547 .table(
3548 "beta",
3549 vec![
3550 TableColumnConfig::new("id", DataType::Int64, false),
3551 TableColumnConfig::new("name", DataType::Utf8, false),
3552 ],
3553 vec!["id".to_string()],
3554 vec![],
3555 )
3556 .unwrap()
3557 .table(
3558 "gamma",
3559 vec![TableColumnConfig::new("id", DataType::Int64, false)],
3560 vec!["id".to_string()],
3561 vec![],
3562 )
3563 .unwrap();
3564
3565 assert_eq!(schema.table_count(), 3);
3566 }
3567
3568 #[test]
3569 fn kv_schema_allows_max_codec_table_count_and_rejects_overflow() {
3570 let client = StoreClient::new("http://localhost:10000");
3571 let mut schema = KvSchema::new(PrefixedStoreClient::empty(client));
3572 for idx in 0..MAX_TABLES {
3573 schema = schema
3574 .table(
3575 format!("t{idx}"),
3576 vec![TableColumnConfig::new("id", DataType::Int64, false)],
3577 vec!["id".to_string()],
3578 vec![],
3579 )
3580 .expect("tables up to codec capacity should be accepted");
3581 }
3582 assert_eq!(schema.table_count(), MAX_TABLES);
3583
3584 let overflow = schema.table(
3585 "overflow",
3586 vec![TableColumnConfig::new("id", DataType::Int64, false)],
3587 vec!["id".to_string()],
3588 vec![],
3589 );
3590 match overflow {
3591 Ok(_) => panic!("overflow table should be rejected"),
3592 Err(err) => assert!(
3593 err.contains(&format!(
3594 "too many tables for key layout (max {MAX_TABLES})"
3595 )),
3596 "overflow table should be rejected with key-layout error"
3597 ),
3598 }
3599 }
3600
3601 #[test]
3602 fn sequential_prefixes_produce_non_overlapping_pk_ranges() {
3603 let range_a = primary_key_prefix_range(0);
3604 let range_b = primary_key_prefix_range(1);
3605 let range_c = primary_key_prefix_range(2);
3606 assert!(range_a.end < range_b.start);
3607 assert!(range_b.end < range_c.start);
3608 }
3609
3610 #[test]
3611 fn sequential_prefixes_isolate_primary_keys() {
3612 let model_0 = simple_int64_model(0);
3613 let model_1 = simple_int64_model(1);
3614 let pk = CellValue::Int64(42);
3615 let key_a = encode_primary_key(0, &[&pk], &model_0).expect("pk key encodes");
3616 let key_b = encode_primary_key(1, &[&pk], &model_1).expect("pk key encodes");
3617 assert_ne!(key_a, key_b);
3618 assert!(decode_primary_key(0, &key_a, &model_0).is_some());
3619 assert!(decode_primary_key(1, &key_a, &model_1).is_none());
3620 assert!(decode_primary_key(0, &key_b, &model_0).is_none());
3621 assert!(decode_primary_key(1, &key_b, &model_1).is_some());
3622 }
3623
3624 #[test]
3625 fn sequential_prefixes_isolate_secondary_keys() {
3626 let config_a = KvTableConfig::new(
3627 0,
3628 vec![
3629 TableColumnConfig::new("id", DataType::Int64, false),
3630 TableColumnConfig::new("name", DataType::Utf8, false),
3631 ],
3632 vec!["id".to_string()],
3633 vec![IndexSpec::new("name_idx", vec!["name".to_string()]).unwrap()],
3634 )
3635 .unwrap();
3636 let config_b = KvTableConfig::new(
3637 1,
3638 vec![
3639 TableColumnConfig::new("id", DataType::Int64, false),
3640 TableColumnConfig::new("name", DataType::Utf8, false),
3641 ],
3642 vec!["id".to_string()],
3643 vec![IndexSpec::new("name_idx", vec!["name".to_string()]).unwrap()],
3644 )
3645 .unwrap();
3646
3647 let model_a = TableModel::from_config(&config_a).unwrap();
3648 let specs_a = model_a.resolve_index_specs(&config_a.index_specs).unwrap();
3649 let model_b = TableModel::from_config(&config_b).unwrap();
3650 let specs_b = model_b.resolve_index_specs(&config_b.index_specs).unwrap();
3651
3652 let row = KvRow {
3653 values: vec![CellValue::Int64(1), CellValue::Utf8("alice".to_string())],
3654 };
3655 let key_a =
3656 encode_secondary_index_key(model_a.table_prefix, &specs_a[0], &model_a, &row).unwrap();
3657 let key_b =
3658 encode_secondary_index_key(model_b.table_prefix, &specs_b[0], &model_b, &row).unwrap();
3659 assert_ne!(key_a, key_b);
3660 assert!(
3661 decode_secondary_index_key(model_a.table_prefix, &specs_a[0], &model_a, &key_b)
3662 .is_none(),
3663 "key from prefix 1 must not decode under prefix 0"
3664 );
3665 }
3666
3667 #[tokio::test]
3668 async fn kv_schema_register_all_enables_join() {
3669 let ctx = SessionContext::new();
3670 let client = StoreClient::new("http://localhost:10000");
3671
3672 let result = KvSchema::new(PrefixedStoreClient::empty(client))
3673 .table(
3674 "customers",
3675 vec![
3676 TableColumnConfig::new("customer_id", DataType::Int64, false),
3677 TableColumnConfig::new("name", DataType::Utf8, false),
3678 ],
3679 vec!["customer_id".to_string()],
3680 vec![],
3681 )
3682 .unwrap()
3683 .table(
3684 "orders",
3685 vec![
3686 TableColumnConfig::new("order_id", DataType::Int64, false),
3687 TableColumnConfig::new("customer_id", DataType::Int64, false),
3688 TableColumnConfig::new("amount", DataType::Int64, false),
3689 ],
3690 vec!["order_id".to_string()],
3691 vec![IndexSpec::new("cust_idx", vec!["customer_id".to_string()]).unwrap()],
3692 )
3693 .unwrap()
3694 .register_all(&ctx);
3695
3696 assert!(
3697 result.is_ok(),
3698 "register_all must succeed: {:?}",
3699 result.err()
3700 );
3701
3702 let plan = ctx
3703 .sql(
3704 "SELECT c.name, o.order_id, o.amount \
3705 FROM orders o \
3706 JOIN customers c ON o.customer_id = c.customer_id",
3707 )
3708 .await;
3709 assert!(
3710 plan.is_ok(),
3711 "JOIN query must plan successfully: {:?}",
3712 plan.err()
3713 );
3714 }
3715
3716 #[tokio::test]
3717 async fn kv_schema_three_way_join() {
3718 let ctx = SessionContext::new();
3719 let client = StoreClient::new("http://localhost:10000");
3720
3721 KvSchema::new(PrefixedStoreClient::empty(client))
3722 .table(
3723 "products",
3724 vec![
3725 TableColumnConfig::new("product_id", DataType::Int64, false),
3726 TableColumnConfig::new("name", DataType::Utf8, false),
3727 TableColumnConfig::new("price", DataType::Int64, false),
3728 ],
3729 vec!["product_id".to_string()],
3730 vec![],
3731 )
3732 .unwrap()
3733 .table(
3734 "line_items",
3735 vec![
3736 TableColumnConfig::new("item_id", DataType::Int64, false),
3737 TableColumnConfig::new("order_id", DataType::Int64, false),
3738 TableColumnConfig::new("product_id", DataType::Int64, false),
3739 TableColumnConfig::new("qty", DataType::Int64, false),
3740 ],
3741 vec!["item_id".to_string()],
3742 vec![
3743 IndexSpec::new("prod_idx", vec!["product_id".to_string()]).unwrap(),
3744 IndexSpec::new("order_idx", vec!["order_id".to_string()]).unwrap(),
3745 ],
3746 )
3747 .unwrap()
3748 .table(
3749 "orders",
3750 vec![
3751 TableColumnConfig::new("order_id", DataType::Int64, false),
3752 TableColumnConfig::new("customer", DataType::Utf8, false),
3753 ],
3754 vec!["order_id".to_string()],
3755 vec![],
3756 )
3757 .unwrap()
3758 .register_all(&ctx)
3759 .unwrap();
3760
3761 let plan = ctx
3762 .sql(
3763 "SELECT o.customer, p.name, li.qty \
3764 FROM line_items li \
3765 JOIN products p ON li.product_id = p.product_id \
3766 JOIN orders o ON li.order_id = o.order_id",
3767 )
3768 .await;
3769 assert!(plan.is_ok(), "three-way JOIN must plan: {:?}", plan.err());
3770 }
3771
3772 #[test]
3773 fn kv_schema_orders_table_convenience() {
3774 let client = StoreClient::new("http://localhost:10000");
3775 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
3776 .orders_table(
3777 "my_orders",
3778 vec![IndexSpec::new(
3779 "region_customer",
3780 vec!["region".to_string(), "customer_id".to_string()],
3781 )
3782 .unwrap()],
3783 )
3784 .unwrap();
3785 assert_eq!(schema.table_count(), 1);
3786 }
3787
3788 #[test]
3789 fn nullable_column_accepted_in_config() {
3790 let config = KvTableConfig::new(
3791 0,
3792 vec![
3793 TableColumnConfig::new("id", DataType::Int64, false),
3794 TableColumnConfig::new("name", DataType::Utf8, true),
3795 ],
3796 vec!["id".to_string()],
3797 vec![],
3798 );
3799 assert!(config.is_ok());
3800 }
3801
3802 #[test]
3803 fn nullable_column_rejected_in_index() {
3804 let config = KvTableConfig::new(
3805 0,
3806 vec![
3807 TableColumnConfig::new("id", DataType::Int64, false),
3808 TableColumnConfig::new("name", DataType::Utf8, true),
3809 ],
3810 vec!["id".to_string()],
3811 vec![IndexSpec::new("name_idx", vec!["name".to_string()]).unwrap()],
3812 )
3813 .unwrap();
3814 let model = TableModel::from_config(&config).unwrap();
3815 let result = model.resolve_index_specs(&config.index_specs);
3816 assert!(result.is_err());
3817 assert!(result.unwrap_err().contains("nullable"));
3818 }
3819
3820 #[test]
3821 fn base_row_round_trip_with_null() {
3822 let config = KvTableConfig::new(
3823 0,
3824 vec![
3825 TableColumnConfig::new("id", DataType::Int64, false),
3826 TableColumnConfig::new("label", DataType::Utf8, true),
3827 TableColumnConfig::new("score", DataType::Int64, true),
3828 ],
3829 vec!["id".to_string()],
3830 vec![],
3831 )
3832 .unwrap();
3833 let model = TableModel::from_config(&config).unwrap();
3834 let row = KvRow {
3835 values: vec![CellValue::Int64(1), CellValue::Null, CellValue::Int64(42)],
3836 };
3837 let encoded = encode_base_row_value(&row, &model).unwrap();
3838 let decoded = decode_base_row(vec![CellValue::Int64(1)], &encoded, &model).unwrap();
3839 assert!(matches!(&decoded.values[0], CellValue::Int64(1)));
3840 assert!(matches!(&decoded.values[1], CellValue::Null));
3841 assert!(matches!(&decoded.values[2], CellValue::Int64(42)));
3842 }
3843
3844 #[test]
3845 fn null_does_not_match_equality_constraint() {
3846 assert!(!matches_constraint(
3847 &CellValue::Null,
3848 &PredicateConstraint::StringEq("x".to_string())
3849 ));
3850 assert!(!matches_constraint(
3851 &CellValue::Null,
3852 &PredicateConstraint::IntRange {
3853 min: Some(0),
3854 max: Some(10)
3855 }
3856 ));
3857 }
3858
3859 #[test]
3860 fn is_null_constraint_matches() {
3861 assert!(matches_constraint(
3862 &CellValue::Null,
3863 &PredicateConstraint::IsNull
3864 ));
3865 assert!(!matches_constraint(
3866 &CellValue::Utf8("x".to_string()),
3867 &PredicateConstraint::IsNull
3868 ));
3869 assert!(!matches_constraint(
3870 &CellValue::Null,
3871 &PredicateConstraint::IsNotNull
3872 ));
3873 assert!(matches_constraint(
3874 &CellValue::Int64(5),
3875 &PredicateConstraint::IsNotNull
3876 ));
3877 }
3878
3879 #[test]
3880 fn string_in_constraint_matches() {
3881 let constraint =
3882 PredicateConstraint::StringIn(vec!["us-east".to_string(), "us-west".to_string()]);
3883 assert!(matches_constraint(
3884 &CellValue::Utf8("us-east".to_string()),
3885 &constraint,
3886 ));
3887 assert!(matches_constraint(
3888 &CellValue::Utf8("us-west".to_string()),
3889 &constraint,
3890 ));
3891 assert!(!matches_constraint(
3892 &CellValue::Utf8("eu-central".to_string()),
3893 &constraint,
3894 ));
3895 }
3896
3897 #[test]
3898 fn int_in_constraint_matches() {
3899 let constraint = PredicateConstraint::IntIn(vec![1, 2, 3]);
3900 assert!(matches_constraint(&CellValue::Int64(1), &constraint));
3901 assert!(matches_constraint(&CellValue::Int64(3), &constraint));
3902 assert!(!matches_constraint(&CellValue::Int64(4), &constraint));
3903 }
3904
3905 #[test]
3906 fn oversized_utf8_primary_key_constraint_is_empty_terminal_range() {
3907 let too_large = "a".repeat(exoware_sdk::keys::MAX_KEY_LEN + 1);
3908 match primary_key_range_constraint(
3909 ColumnKind::Utf8,
3910 &PredicateConstraint::StringEq(too_large),
3911 ) {
3912 PrimaryKeyRangeConstraint::Terminal(spans) => assert!(
3913 spans.is_empty(),
3914 "unencodable UTF-8 PK equality should be impossible"
3915 ),
3916 other => panic!("unexpected oversized UTF-8 PK constraint: {other:?}"),
3917 }
3918 }
3919
3920 #[test]
3921 fn in_predicate_generates_multiple_index_ranges() {
3922 let (model, specs) = test_model();
3923 let region_idx = *model.columns_by_name.get("region").unwrap();
3924 let mut pred = QueryPredicate::default();
3925 pred.constraints.insert(
3926 region_idx,
3927 PredicateConstraint::StringIn(vec!["us-east".to_string(), "us-west".to_string()]),
3928 );
3929 let plan = pred
3930 .choose_index_plan(&model, &specs)
3931 .expect("plan")
3932 .expect("should find index");
3933 assert_eq!(plan.ranges.len(), 2);
3934 }
3935
3936 #[test]
3937 fn int_in_generates_multiple_pk_ranges() {
3938 let (model, _specs) = test_model();
3939 let mut pred = QueryPredicate::default();
3940 pred.constraints.insert(
3941 model.primary_key_indices[0],
3942 PredicateConstraint::IntIn(vec![200, 100, 300]),
3943 );
3944 let ranges = pred.primary_key_ranges(&model).unwrap();
3945 assert_eq!(ranges.len(), 3);
3946 let expected_starts = [100, 200, 300]
3947 .into_iter()
3948 .map(|id| {
3949 let value = CellValue::Int64(id);
3950 encode_primary_key_bound(model.table_prefix, &[&value], &model, false)
3951 .expect("primary key lower bound")
3952 })
3953 .collect::<Vec<_>>();
3954 assert_eq!(
3955 ranges
3956 .iter()
3957 .map(|range| range.start.clone())
3958 .collect::<Vec<_>>(),
3959 expected_starts,
3960 "primary_key_ranges should sort terminal IN ranges by encoded start key"
3961 );
3962 }
3963
3964 #[test]
3965 fn duplicate_int_in_values_deduplicated() {
3966 let (model, _specs) = test_model();
3967 let filter = Expr::InList(datafusion::logical_expr::expr::InList {
3969 expr: Box::new(Expr::Column(datafusion::common::Column::new_unqualified(
3970 "order_id",
3971 ))),
3972 list: vec![
3973 Expr::Literal(ScalarValue::Int64(Some(5)), None),
3974 Expr::Literal(ScalarValue::Int64(Some(5)), None),
3975 Expr::Literal(ScalarValue::Int64(Some(10)), None),
3976 ],
3977 negated: false,
3978 });
3979 let pred = QueryPredicate::from_filters(&[filter], &model);
3980 let ranges = pred.primary_key_ranges(&model).unwrap();
3981 assert_eq!(
3982 ranges.len(),
3983 2,
3984 "duplicate IN values must be deduped, producing 2 ranges not 3"
3985 );
3986 }
3987
3988 #[test]
3989 fn duplicate_uint64_in_values_deduplicated() {
3990 let config = KvTableConfig::new(
3991 0,
3992 vec![
3993 TableColumnConfig::new("id", DataType::UInt64, false),
3994 TableColumnConfig::new("name", DataType::Utf8, false),
3995 ],
3996 vec!["id".to_string()],
3997 vec![],
3998 )
3999 .unwrap();
4000 let model = TableModel::from_config(&config).unwrap();
4001 let filter = Expr::InList(datafusion::logical_expr::expr::InList {
4002 expr: Box::new(Expr::Column(datafusion::common::Column::new_unqualified(
4003 "id",
4004 ))),
4005 list: vec![
4006 Expr::Literal(ScalarValue::UInt64(Some(100)), None),
4007 Expr::Literal(ScalarValue::UInt64(Some(100)), None),
4008 Expr::Literal(ScalarValue::UInt64(Some(200)), None),
4009 ],
4010 negated: false,
4011 });
4012 let pred = QueryPredicate::from_filters(&[filter], &model);
4013 let ranges = pred.primary_key_ranges(&model).unwrap();
4014 assert_eq!(
4015 ranges.len(),
4016 2,
4017 "duplicate UInt64 IN values must be deduped"
4018 );
4019 }
4020
4021 #[test]
4022 fn duplicate_fixed_binary_in_values_deduplicated() {
4023 let config = KvTableConfig::new(
4024 0,
4025 vec![
4026 TableColumnConfig::new("hash", DataType::FixedSizeBinary(16), false),
4027 TableColumnConfig::new("val", DataType::Int64, false),
4028 ],
4029 vec!["hash".to_string()],
4030 vec![],
4031 )
4032 .unwrap();
4033 let model = TableModel::from_config(&config).unwrap();
4034 let dup_val = vec![0xAA; 16];
4035 let other_val = vec![0xBB; 16];
4036 let filter = Expr::InList(datafusion::logical_expr::expr::InList {
4037 expr: Box::new(Expr::Column(datafusion::common::Column::new_unqualified(
4038 "hash",
4039 ))),
4040 list: vec![
4041 Expr::Literal(
4042 ScalarValue::FixedSizeBinary(16, Some(dup_val.clone())),
4043 None,
4044 ),
4045 Expr::Literal(ScalarValue::FixedSizeBinary(16, Some(dup_val)), None),
4046 Expr::Literal(ScalarValue::FixedSizeBinary(16, Some(other_val)), None),
4047 ],
4048 negated: false,
4049 });
4050 let pred = QueryPredicate::from_filters(&[filter], &model);
4051 let ranges = pred.primary_key_ranges(&model).unwrap();
4052 assert_eq!(
4053 ranges.len(),
4054 2,
4055 "duplicate FixedBinary IN values must be deduped"
4056 );
4057 }
4058
4059 #[test]
4060 fn composite_utf8_primary_key_prefix_overflow_produces_no_ranges() {
4061 let config = KvTableConfig::new(
4062 0,
4063 vec![
4064 TableColumnConfig::new("a", DataType::Utf8, false),
4065 TableColumnConfig::new("b", DataType::Utf8, false),
4066 ],
4067 vec!["a".to_string(), "b".to_string()],
4068 vec![],
4069 )
4070 .unwrap();
4071 let model = TableModel::from_config(&config).unwrap();
4072 let capacity = model.primary_key_prefix.max_payload_len();
4073 let a = "a".repeat(capacity - 10);
4074 let b = "b".repeat(10);
4075 let a_width = encode_string_variable(&a).expect("a should encode").len();
4076 let b_width = encode_string_variable(&b).expect("b should encode").len();
4077 assert!(a_width <= capacity, "first component should fit alone");
4078 assert!(
4079 a_width + b_width > capacity,
4080 "composite UTF-8 prefix should overflow"
4081 );
4082
4083 let mut pred = QueryPredicate::default();
4084 pred.constraints.insert(0, PredicateConstraint::StringEq(a));
4085 pred.constraints.insert(1, PredicateConstraint::StringEq(b));
4086
4087 let ranges = pred.primary_key_ranges(&model).unwrap();
4088 assert!(
4089 ranges.is_empty(),
4090 "overflowing composite UTF-8 PK equality can never match"
4091 );
4092 }
4093
4094 #[test]
4095 fn or_equalities_extracted_as_in_list() {
4096 let (model, _) = test_model();
4097 let expr = Expr::BinaryExpr(datafusion::logical_expr::BinaryExpr {
4098 left: Box::new(Expr::BinaryExpr(datafusion::logical_expr::BinaryExpr {
4099 left: Box::new(Expr::Column(datafusion::common::Column::new_unqualified(
4100 "region",
4101 ))),
4102 op: Operator::Eq,
4103 right: Box::new(Expr::Literal(
4104 ScalarValue::Utf8(Some("us-east".to_string())),
4105 None,
4106 )),
4107 })),
4108 op: Operator::Or,
4109 right: Box::new(Expr::BinaryExpr(datafusion::logical_expr::BinaryExpr {
4110 left: Box::new(Expr::Column(datafusion::common::Column::new_unqualified(
4111 "region",
4112 ))),
4113 op: Operator::Eq,
4114 right: Box::new(Expr::Literal(
4115 ScalarValue::Utf8(Some("us-west".to_string())),
4116 None,
4117 )),
4118 })),
4119 });
4120 let result = extract_or_in_column(&expr, &model);
4121 assert!(result.is_some());
4122 let (col, vals) = result.unwrap();
4123 assert_eq!(col, "region");
4124 assert_eq!(vals.len(), 2);
4125 }
4126
4127 #[test]
4128 fn or_equalities_on_float64_are_not_pushdown_supported() {
4129 let config = KvTableConfig::new(
4130 0,
4131 vec![
4132 TableColumnConfig::new("id", DataType::Int64, false),
4133 TableColumnConfig::new("score", DataType::Float64, false),
4134 ],
4135 vec!["id".to_string()],
4136 vec![],
4137 )
4138 .unwrap();
4139 let model = TableModel::from_config(&config).unwrap();
4140
4141 use datafusion::logical_expr::col;
4142 let filter = col("score")
4143 .eq(Expr::Literal(ScalarValue::Float64(Some(1.0)), None))
4144 .or(col("score").eq(Expr::Literal(ScalarValue::Float64(Some(2.0)), None)));
4145
4146 assert!(
4147 !QueryPredicate::supports_filter(&filter, &model),
4148 "OR-equality pushdown should be disabled for Float64 because apply_in_list cannot enforce it"
4149 );
4150
4151 let pred = QueryPredicate::from_filters(&[filter], &model);
4152 assert!(!pred.contradiction);
4153 assert!(
4154 pred.constraints.is_empty(),
4155 "unsupported OR predicate must not contribute pushdown constraints"
4156 );
4157 }
4158
4159 #[test]
4160 fn batch_writer_encodes_rows_across_tables() {
4161 let client = StoreClient::new("http://localhost:10000");
4162 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
4163 .table(
4164 "customers",
4165 vec![
4166 TableColumnConfig::new("customer_id", DataType::Int64, false),
4167 TableColumnConfig::new("name", DataType::Utf8, false),
4168 ],
4169 vec!["customer_id".to_string()],
4170 vec![],
4171 )
4172 .unwrap()
4173 .table(
4174 "orders",
4175 vec![
4176 TableColumnConfig::new("order_id", DataType::Int64, false),
4177 TableColumnConfig::new("customer_id", DataType::Int64, false),
4178 TableColumnConfig::new("amount", DataType::Int64, false),
4179 ],
4180 vec!["order_id".to_string()],
4181 vec![IndexSpec::new("cust_idx", vec!["customer_id".to_string()]).unwrap()],
4182 )
4183 .unwrap();
4184
4185 let mut batch = schema.batch_writer();
4186 batch
4187 .insert(
4188 "customers",
4189 vec![CellValue::Int64(1), CellValue::Utf8("Alice".to_string())],
4190 )
4191 .unwrap();
4192 batch
4193 .insert(
4194 "orders",
4195 vec![
4196 CellValue::Int64(100),
4197 CellValue::Int64(1),
4198 CellValue::Int64(4999),
4199 ],
4200 )
4201 .unwrap();
4202 batch
4203 .insert(
4204 "orders",
4205 vec![
4206 CellValue::Int64(101),
4207 CellValue::Int64(1),
4208 CellValue::Int64(2999),
4209 ],
4210 )
4211 .unwrap();
4212
4213 assert_eq!(batch.pending_count(), 5);
4215 }
4216
4217 #[test]
4218 fn batch_writer_rejects_unknown_table() {
4219 let client = StoreClient::new("http://localhost:10000");
4220 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
4221 .table(
4222 "t1",
4223 vec![TableColumnConfig::new("id", DataType::Int64, false)],
4224 vec!["id".to_string()],
4225 vec![],
4226 )
4227 .unwrap();
4228
4229 let mut batch = schema.batch_writer();
4230 let result = batch.insert("nonexistent", vec![CellValue::Int64(1)]);
4231 assert!(result.is_err());
4232 assert!(result.unwrap_err().contains("unknown table"));
4233 }
4234
4235 #[test]
4236 fn batch_writer_rejects_wrong_column_count() {
4237 let client = StoreClient::new("http://localhost:10000");
4238 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
4239 .table(
4240 "t1",
4241 vec![
4242 TableColumnConfig::new("id", DataType::Int64, false),
4243 TableColumnConfig::new("name", DataType::Utf8, false),
4244 ],
4245 vec!["id".to_string()],
4246 vec![],
4247 )
4248 .unwrap();
4249
4250 let mut batch = schema.batch_writer();
4251 let result = batch.insert("t1", vec![CellValue::Int64(1)]);
4252 assert!(result.is_err());
4253 assert!(result.unwrap_err().contains("expected 2"));
4254 }
4255
4256 #[test]
4257 fn batch_writer_rejects_non_pk_type_mismatch() {
4258 let client = StoreClient::new("http://localhost:10000");
4259 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
4260 .table(
4261 "t1",
4262 vec![
4263 TableColumnConfig::new("id", DataType::Int64, false),
4264 TableColumnConfig::new("amount", DataType::Int64, false),
4265 ],
4266 vec!["id".to_string()],
4267 vec![],
4268 )
4269 .unwrap();
4270
4271 let mut batch = schema.batch_writer();
4272 let result = batch.insert(
4273 "t1",
4274 vec![CellValue::Int64(1), CellValue::Utf8("bad".to_string())],
4275 );
4276 assert!(result.is_err());
4277 assert!(
4278 result.unwrap_err().contains("type mismatch"),
4279 "non-PK schema-invalid values must be rejected at insert-time"
4280 );
4281 }
4282
4283 #[test]
4284 fn batch_writer_entries_use_distinct_table_prefixes() {
4285 let client = StoreClient::new("http://localhost:10000");
4286 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
4287 .table(
4288 "a",
4289 vec![TableColumnConfig::new("id", DataType::Int64, false)],
4290 vec!["id".to_string()],
4291 vec![],
4292 )
4293 .unwrap()
4294 .table(
4295 "b",
4296 vec![TableColumnConfig::new("id", DataType::Int64, false)],
4297 vec!["id".to_string()],
4298 vec![],
4299 )
4300 .unwrap();
4301
4302 let mut batch = schema.batch_writer();
4303 batch.insert("a", vec![CellValue::Int64(42)]).unwrap();
4304 batch.insert("b", vec![CellValue::Int64(42)]).unwrap();
4305
4306 assert_eq!(batch.pending_count(), 2);
4307 assert_ne!(
4308 batch.pending_keys[0], batch.pending_keys[1],
4309 "same PK in different tables must produce different keys"
4310 );
4311 assert_ne!(
4312 batch.pending_keys[0][0], batch.pending_keys[1][0],
4313 "table prefix byte must differ"
4314 );
4315 }
4316
4317 #[tokio::test]
4318 async fn batch_writer_trait_failure_requeues_prepared_before_new_pending() {
4319 let client = StoreClient::new("http://localhost:10000");
4320 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
4321 .table(
4322 "t",
4323 vec![TableColumnConfig::new("id", DataType::Int64, false)],
4324 vec!["id".to_string()],
4325 vec![],
4326 )
4327 .unwrap();
4328
4329 let mut batch = schema.batch_writer();
4330 batch.insert("t", vec![CellValue::Int64(1)]).unwrap();
4331 let prepared = batch.prepare_flush().unwrap().expect("prepared row");
4332 assert_eq!(prepared.request_id(), 0);
4333 assert_eq!(prepared.entry_count(), 1);
4334
4335 batch.insert("t", vec![CellValue::Int64(2)]).unwrap();
4336 StoreBatchUpload::mark_upload_failed(&batch, prepared, "commit failed".to_string()).await;
4337 assert_eq!(batch.pending_count(), 2);
4338
4339 let retry = batch.prepare_flush().unwrap().expect("retry row");
4340 assert_eq!(retry.request_id(), 0);
4341 assert_eq!(retry.entry_count(), 1);
4342 let next = batch.prepare_flush().unwrap().expect("new pending row");
4343 assert_eq!(next.request_id(), 1);
4344 assert_eq!(next.entry_count(), 1);
4345 }
4346
4347 #[test]
4348 fn batch_writer_supports_nullable_columns() {
4349 let client = StoreClient::new("http://localhost:10000");
4350 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
4351 .table(
4352 "t",
4353 vec![
4354 TableColumnConfig::new("id", DataType::Int64, false),
4355 TableColumnConfig::new("note", DataType::Utf8, true),
4356 ],
4357 vec!["id".to_string()],
4358 vec![],
4359 )
4360 .unwrap();
4361
4362 let mut batch = schema.batch_writer();
4363 batch
4364 .insert("t", vec![CellValue::Int64(1), CellValue::Null])
4365 .unwrap();
4366 assert_eq!(batch.pending_count(), 1);
4367 }
4368
4369 #[test]
4370 fn non_nullable_column_rejects_null_in_batch_writer() {
4371 let client = StoreClient::new("http://localhost:10000");
4372 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
4373 .table(
4374 "t",
4375 vec![
4376 TableColumnConfig::new("id", DataType::Int64, false),
4377 TableColumnConfig::new("name", DataType::Utf8, false),
4378 TableColumnConfig::new("note", DataType::Utf8, true),
4379 ],
4380 vec!["id".to_string()],
4381 vec![],
4382 )
4383 .unwrap();
4384
4385 let mut batch = schema.batch_writer();
4387 let result = batch.insert(
4388 "t",
4389 vec![
4390 CellValue::Int64(1),
4391 CellValue::Null,
4392 CellValue::Utf8("ok".to_string()),
4393 ],
4394 );
4395 assert!(result.is_err());
4396 assert!(
4397 result.unwrap_err().contains("not nullable"),
4398 "error should mention non-nullable constraint"
4399 );
4400
4401 let mut batch = schema.batch_writer();
4403 batch
4404 .insert(
4405 "t",
4406 vec![
4407 CellValue::Int64(1),
4408 CellValue::Utf8("Alice".to_string()),
4409 CellValue::Null,
4410 ],
4411 )
4412 .unwrap();
4413 assert_eq!(batch.pending_count(), 1);
4414
4415 let mut batch = schema.batch_writer();
4417 batch
4418 .insert(
4419 "t",
4420 vec![
4421 CellValue::Int64(1),
4422 CellValue::Utf8("Alice".to_string()),
4423 CellValue::Utf8("hello".to_string()),
4424 ],
4425 )
4426 .unwrap();
4427 assert_eq!(batch.pending_count(), 1);
4428 }
4429
4430 #[test]
4431 fn uint64_column_accepted() {
4432 let config = KvTableConfig::new(
4433 0,
4434 vec![
4435 TableColumnConfig::new("id", DataType::UInt64, false),
4436 TableColumnConfig::new("name", DataType::Utf8, false),
4437 ],
4438 vec!["id".to_string()],
4439 vec![],
4440 );
4441 assert!(config.is_ok());
4442 }
4443
4444 #[test]
4445 fn uint64_primary_key_round_trip() {
4446 let config = KvTableConfig::new(
4447 0,
4448 vec![
4449 TableColumnConfig::new("id", DataType::UInt64, false),
4450 TableColumnConfig::new("label", DataType::Utf8, false),
4451 ],
4452 vec!["id".to_string()],
4453 vec![],
4454 )
4455 .unwrap();
4456 let model = TableModel::from_config(&config).unwrap();
4457 let row = KvRow {
4458 values: vec![
4459 CellValue::UInt64(u64::MAX),
4460 CellValue::Utf8("max".to_string()),
4461 ],
4462 };
4463 let encoded = encode_base_row_value(&row, &model).unwrap();
4464 let pk = row
4465 .primary_key_values(&model)
4466 .into_iter()
4467 .cloned()
4468 .collect::<Vec<_>>();
4469 let decoded = decode_base_row(pk, &encoded, &model).unwrap();
4470 assert!(matches!(&decoded.values[0], CellValue::UInt64(v) if *v == u64::MAX));
4471 assert!(matches!(&decoded.values[1], CellValue::Utf8(v) if v == "max"));
4472 }
4473
4474 #[test]
4475 fn string_primary_key_accepted() {
4476 let config = KvTableConfig::new(
4477 0,
4478 vec![
4479 TableColumnConfig::new("code", DataType::Utf8, false),
4480 TableColumnConfig::new("value", DataType::Int64, false),
4481 ],
4482 vec!["code".to_string()],
4483 vec![],
4484 );
4485 assert!(config.is_ok());
4486 }
4487
4488 #[test]
4489 fn fixed_binary_primary_key_round_trip() {
4490 let config = KvTableConfig::new(
4491 0,
4492 vec![
4493 TableColumnConfig::new("hash", DataType::FixedSizeBinary(32), false),
4494 TableColumnConfig::new("amount", DataType::Int64, false),
4495 ],
4496 vec!["hash".to_string()],
4497 vec![],
4498 )
4499 .unwrap();
4500 let model = TableModel::from_config(&config).unwrap();
4501 let hash_val = vec![0xABu8; 32];
4502 let row = KvRow {
4503 values: vec![
4504 CellValue::FixedBinary(hash_val.clone()),
4505 CellValue::Int64(100),
4506 ],
4507 };
4508 let encoded = encode_base_row_value(&row, &model).unwrap();
4509 let pk = row
4510 .primary_key_values(&model)
4511 .into_iter()
4512 .cloned()
4513 .collect::<Vec<_>>();
4514 let decoded = decode_base_row(pk, &encoded, &model).unwrap();
4515 assert!(matches!(&decoded.values[0], CellValue::FixedBinary(v) if *v == hash_val));
4516 }
4517
4518 #[test]
4519 fn fixed_binary_key_rejects_wrong_length() {
4520 let config = KvTableConfig::new(
4521 0,
4522 vec![
4523 TableColumnConfig::new("hash", DataType::FixedSizeBinary(16), false),
4524 TableColumnConfig::new("amount", DataType::Int64, false),
4525 ],
4526 vec!["hash".to_string()],
4527 vec![],
4528 )
4529 .unwrap();
4530 let model = TableModel::from_config(&config).unwrap();
4531
4532 let short_row = KvRow {
4534 values: vec![CellValue::FixedBinary(vec![0xAB; 10]), CellValue::Int64(1)],
4535 };
4536 let result = encode_primary_key_from_row(model.table_prefix, &short_row, &model);
4537 assert!(result.is_err());
4538 assert!(
4539 result.unwrap_err().contains("requires exactly 16 bytes"),
4540 "should mention exact width requirement"
4541 );
4542
4543 let long_row = KvRow {
4545 values: vec![CellValue::FixedBinary(vec![0xCD; 20]), CellValue::Int64(2)],
4546 };
4547 let result = encode_primary_key_from_row(model.table_prefix, &long_row, &model);
4548 assert!(result.is_err());
4549 assert!(result.unwrap_err().contains("requires exactly 16 bytes"));
4550
4551 let ok_row = KvRow {
4553 values: vec![CellValue::FixedBinary(vec![0xEF; 16]), CellValue::Int64(3)],
4554 };
4555 assert!(encode_primary_key_from_row(model.table_prefix, &ok_row, &model).is_ok());
4556 }
4557
4558 #[test]
4559 fn fixed_binary_index_key_rejects_wrong_length() {
4560 let config = KvTableConfig::new(
4561 0,
4562 vec![
4563 TableColumnConfig::new("id", DataType::Int64, false),
4564 TableColumnConfig::new("tag", DataType::FixedSizeBinary(8), false),
4565 ],
4566 vec!["id".to_string()],
4567 vec![IndexSpec::new("tag_idx", vec!["tag".to_string()]).unwrap()],
4568 )
4569 .unwrap();
4570 let model = TableModel::from_config(&config).unwrap();
4571 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
4572
4573 let bad_row = KvRow {
4575 values: vec![CellValue::Int64(1), CellValue::FixedBinary(vec![0x01; 4])],
4576 };
4577 let result = encode_secondary_index_key(model.table_prefix, &specs[0], &model, &bad_row);
4578 assert!(result.is_err());
4579 assert!(result.unwrap_err().contains("requires exactly 8 bytes"));
4580
4581 let ok_row = KvRow {
4583 values: vec![CellValue::Int64(1), CellValue::FixedBinary(vec![0x02; 8])],
4584 };
4585 assert!(encode_secondary_index_key(model.table_prefix, &specs[0], &model, &ok_row).is_ok());
4586 }
4587
4588 #[test]
4589 fn decimal256_column_round_trip() {
4590 let config = KvTableConfig::new(
4591 0,
4592 vec![
4593 TableColumnConfig::new("id", DataType::Int64, false),
4594 TableColumnConfig::new("balance", DataType::Decimal256(76, 0), false),
4595 ],
4596 vec!["id".to_string()],
4597 vec![],
4598 )
4599 .unwrap();
4600 let model = TableModel::from_config(&config).unwrap();
4601 let big_val = i256::from(123456789012345i64);
4602 let row = KvRow {
4603 values: vec![CellValue::Int64(1), CellValue::Decimal256(big_val)],
4604 };
4605 let encoded = encode_base_row_value(&row, &model).unwrap();
4606 let pk = row
4607 .primary_key_values(&model)
4608 .into_iter()
4609 .cloned()
4610 .collect::<Vec<_>>();
4611 let decoded = decode_base_row(pk, &encoded, &model).unwrap();
4612 assert!(matches!(&decoded.values[1], CellValue::Decimal256(v) if *v == big_val));
4613 }
4614
4615 #[test]
4616 fn float64_primary_key_rejected() {
4617 let config = KvTableConfig::new(
4618 0,
4619 vec![TableColumnConfig::new("id", DataType::Float64, false)],
4620 vec!["id".to_string()],
4621 vec![],
4622 );
4623 assert!(config.is_err());
4624 }
4625
4626 #[test]
4627 fn i256_ordered_encoding_round_trip() {
4628 let values = [
4629 i256::from_i128(i128::MIN),
4630 i256::from(-1i64),
4631 i256::from(0i64),
4632 i256::from(1i64),
4633 i256::from_i128(i128::MAX),
4634 ];
4635 for v in values {
4636 assert_eq!(decode_i256_ordered(encode_i256_ordered(v)), v);
4637 }
4638 let encoded: Vec<[u8; 32]> = values.iter().map(|v| encode_i256_ordered(*v)).collect();
4639 for i in 0..encoded.len() - 1 {
4640 assert!(encoded[i] < encoded[i + 1]);
4641 }
4642 }
4643
4644 #[test]
4645 fn uint64_primary_key_encode_decode() {
4646 let config = KvTableConfig::new(
4647 5,
4648 vec![
4649 TableColumnConfig::new("id", DataType::UInt64, false),
4650 TableColumnConfig::new("name", DataType::Utf8, false),
4651 ],
4652 vec!["id".to_string()],
4653 vec![],
4654 )
4655 .unwrap();
4656 let model = TableModel::from_config(&config).unwrap();
4657 let pk = CellValue::UInt64(12345);
4658 let key = encode_primary_key(5, &[&pk], &model).expect("pk key encodes");
4659 let decoded = decode_primary_key(5, &key, &model).unwrap();
4660 assert!(matches!(&decoded[0], CellValue::UInt64(12345)));
4661 }
4662
4663 #[test]
4664 fn utf8_primary_key_encode_decode() {
4665 let config = KvTableConfig::new(
4666 3,
4667 vec![
4668 TableColumnConfig::new("code", DataType::Utf8, false),
4669 TableColumnConfig::new("val", DataType::Int64, false),
4670 ],
4671 vec!["code".to_string()],
4672 vec![],
4673 )
4674 .unwrap();
4675 let model = TableModel::from_config(&config).unwrap();
4676 let pk = CellValue::Utf8("HELLO".to_string());
4677 let key = encode_primary_key(3, &[&pk], &model).expect("pk key encodes");
4678 let decoded = decode_primary_key(3, &key, &model).unwrap();
4679 assert!(matches!(&decoded[0], CellValue::Utf8(v) if v == "HELLO"));
4680 }
4681
4682 #[test]
4683 fn fixed_binary_primary_key_encode_decode() {
4684 let config = KvTableConfig::new(
4685 7,
4686 vec![
4687 TableColumnConfig::new("hash", DataType::FixedSizeBinary(16), false),
4688 TableColumnConfig::new("val", DataType::Int64, false),
4689 ],
4690 vec!["hash".to_string()],
4691 vec![],
4692 )
4693 .unwrap();
4694 let model = TableModel::from_config(&config).unwrap();
4695 let data = vec![0xDE, 0xAD, 0xBE, 0xEF, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
4696 let pk = CellValue::FixedBinary(data.clone());
4697 let key = encode_primary_key(7, &[&pk], &model).expect("pk key encodes");
4698 let decoded = decode_primary_key(7, &key, &model).unwrap();
4699 assert!(matches!(&decoded[0], CellValue::FixedBinary(v) if *v == data));
4700 }
4701
4702 #[test]
4703 fn secondary_index_with_uint64_column() {
4704 let config = KvTableConfig::new(
4705 0,
4706 vec![
4707 TableColumnConfig::new("id", DataType::Int64, false),
4708 TableColumnConfig::new("counter", DataType::UInt64, false),
4709 ],
4710 vec!["id".to_string()],
4711 vec![IndexSpec::new("counter_idx", vec!["counter".to_string()]).unwrap()],
4712 )
4713 .unwrap();
4714 let model = TableModel::from_config(&config).unwrap();
4715 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
4716 let row = KvRow {
4717 values: vec![CellValue::Int64(1), CellValue::UInt64(999)],
4718 };
4719 let key = encode_secondary_index_key(model.table_prefix, &specs[0], &model, &row).unwrap();
4720 let decoded =
4721 decode_secondary_index_key(model.table_prefix, &specs[0], &model, &key).unwrap();
4722 let counter_idx = *model.columns_by_name.get("counter").unwrap();
4723 assert!(matches!(
4724 decoded.values.get(&counter_idx),
4725 Some(CellValue::UInt64(999))
4726 ));
4727 assert!(matches!(
4728 &decoded.primary_key_values[0],
4729 CellValue::Int64(1)
4730 ));
4731 }
4732
4733 #[test]
4734 fn secondary_index_with_decimal256_column() {
4735 let config = KvTableConfig::new(
4736 0,
4737 vec![
4738 TableColumnConfig::new("id", DataType::Int64, false),
4739 TableColumnConfig::new("big_val", DataType::Decimal256(76, 0), false),
4740 ],
4741 vec!["id".to_string()],
4742 vec![IndexSpec::new("big_idx", vec!["big_val".to_string()]).unwrap()],
4743 )
4744 .unwrap();
4745 let model = TableModel::from_config(&config).unwrap();
4746 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
4747 let val = i256::from(42i64);
4748 let row = KvRow {
4749 values: vec![CellValue::Int64(1), CellValue::Decimal256(val)],
4750 };
4751 let key = encode_secondary_index_key(model.table_prefix, &specs[0], &model, &row).unwrap();
4752 let decoded =
4753 decode_secondary_index_key(model.table_prefix, &specs[0], &model, &key).unwrap();
4754 let big_idx = *model.columns_by_name.get("big_val").unwrap();
4755 assert!(matches!(
4756 decoded.values.get(&big_idx),
4757 Some(CellValue::Decimal256(v)) if *v == val
4758 ));
4759 }
4760
4761 #[test]
4766 fn composite_pk_config_accepted() {
4767 let config = KvTableConfig::new(
4768 0,
4769 vec![
4770 TableColumnConfig::new("entity", DataType::FixedSizeBinary(32), false),
4771 TableColumnConfig::new("version", DataType::UInt64, false),
4772 TableColumnConfig::new("data", DataType::Utf8, true),
4773 ],
4774 vec!["entity".to_string(), "version".to_string()],
4775 vec![],
4776 );
4777 assert!(config.is_ok());
4778 let c = config.unwrap();
4779 assert_eq!(c.primary_key_columns, vec!["entity", "version"]);
4780 }
4781
4782 #[test]
4783 fn composite_pk_rejects_unsupported_type() {
4784 let result = KvTableConfig::new(
4785 0,
4786 vec![
4787 TableColumnConfig::new("entity", DataType::FixedSizeBinary(32), false),
4788 TableColumnConfig::new("score", DataType::Float64, false),
4789 ],
4790 vec!["entity".to_string(), "score".to_string()],
4791 vec![],
4792 );
4793 assert!(result.is_err());
4794 let err = result.unwrap_err();
4795 assert!(
4796 err.contains("must be Int64") || err.contains("must be"),
4797 "expected PK type error, got: {err}"
4798 );
4799 }
4800
4801 #[test]
4802 fn composite_pk_rejects_too_wide() {
4803 let config = KvTableConfig::new(
4804 0,
4805 vec![
4806 TableColumnConfig::new("big", DataType::FixedSizeBinary(60), false),
4807 TableColumnConfig::new("ver", DataType::UInt64, false),
4808 ],
4809 vec!["big".to_string(), "ver".to_string()],
4810 vec![],
4811 )
4812 .expect("variable-length keys should allow wider composite PKs");
4813 let model = TableModel::from_config(&config).expect("model");
4814 assert_eq!(model.primary_key_width, 68);
4815 }
4816
4817 #[test]
4818 fn composite_pk_encode_decode_round_trip() {
4819 let config = KvTableConfig::new(
4820 1,
4821 vec![
4822 TableColumnConfig::new("entity", DataType::FixedSizeBinary(32), false),
4823 TableColumnConfig::new("version", DataType::UInt64, false),
4824 TableColumnConfig::new("title", DataType::Utf8, true),
4825 ],
4826 vec!["entity".to_string(), "version".to_string()],
4827 vec![],
4828 )
4829 .unwrap();
4830 let model = TableModel::from_config(&config).unwrap();
4831
4832 let entity = vec![0xAA; 32];
4833 let pk_entity = CellValue::FixedBinary(entity.clone());
4834 let pk_version = CellValue::UInt64(42);
4835 let key =
4836 encode_primary_key(1, &[&pk_entity, &pk_version], &model).expect("pk key encodes");
4837
4838 let decoded = decode_primary_key(1, &key, &model).unwrap();
4839 assert_eq!(decoded.len(), 2);
4840 assert!(matches!(&decoded[0], CellValue::FixedBinary(v) if *v == entity));
4841 assert!(matches!(&decoded[1], CellValue::UInt64(42)));
4842 }
4843
4844 #[test]
4845 fn composite_pk_version_sort_order() {
4846 let config = KvTableConfig::new(
4847 0,
4848 vec![
4849 TableColumnConfig::new("entity", DataType::FixedSizeBinary(32), false),
4850 TableColumnConfig::new("version", DataType::UInt64, false),
4851 ],
4852 vec!["entity".to_string(), "version".to_string()],
4853 vec![],
4854 )
4855 .unwrap();
4856 let model = TableModel::from_config(&config).unwrap();
4857
4858 let entity = vec![0xBB; 32];
4859 let pk_entity = CellValue::FixedBinary(entity.clone());
4860
4861 let key_v1 = encode_primary_key(0, &[&pk_entity, &CellValue::UInt64(1)], &model)
4862 .expect("pk key encodes");
4863 let key_v10 = encode_primary_key(0, &[&pk_entity, &CellValue::UInt64(10)], &model)
4864 .expect("pk key encodes");
4865 let key_v100 = encode_primary_key(0, &[&pk_entity, &CellValue::UInt64(100)], &model)
4866 .expect("pk key encodes");
4867
4868 assert!(key_v1 < key_v10);
4870 assert!(key_v10 < key_v100);
4871 }
4872
4873 #[test]
4874 fn composite_pk_value_excludes_all_pk_columns() {
4875 let config = KvTableConfig::new(
4876 0,
4877 vec![
4878 TableColumnConfig::new("entity", DataType::FixedSizeBinary(16), false),
4879 TableColumnConfig::new("version", DataType::UInt64, false),
4880 TableColumnConfig::new("data", DataType::Utf8, true),
4881 ],
4882 vec!["entity".to_string(), "version".to_string()],
4883 vec![],
4884 )
4885 .unwrap();
4886 let model = TableModel::from_config(&config).unwrap();
4887
4888 let row = KvRow {
4889 values: vec![
4890 CellValue::FixedBinary(vec![0xCC; 16]),
4891 CellValue::UInt64(7),
4892 CellValue::Utf8("hello".to_string()),
4893 ],
4894 };
4895 let encoded = encode_base_row_value(&row, &model).unwrap();
4896 let decoded = decode_base_row(
4898 vec![CellValue::FixedBinary(vec![0xCC; 16]), CellValue::UInt64(7)],
4899 &encoded,
4900 &model,
4901 )
4902 .unwrap();
4903 assert!(matches!(&decoded.values[0], CellValue::FixedBinary(v) if v.len() == 16));
4904 assert!(matches!(&decoded.values[1], CellValue::UInt64(7)));
4905 assert!(matches!(&decoded.values[2], CellValue::Utf8(v) if v == "hello"));
4906 }
4907
4908 #[test]
4909 fn composite_pk_secondary_index_appends_all_pk_columns() {
4910 let config = KvTableConfig::new(
4911 0,
4912 vec![
4913 TableColumnConfig::new("entity", DataType::FixedSizeBinary(16), false),
4914 TableColumnConfig::new("version", DataType::UInt64, false),
4915 TableColumnConfig::new("tag", DataType::Int64, false),
4916 ],
4917 vec!["entity".to_string(), "version".to_string()],
4918 vec![IndexSpec::new("tag_idx", vec!["tag".to_string()]).unwrap()],
4919 )
4920 .unwrap();
4921 let model = TableModel::from_config(&config).unwrap();
4922 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
4923
4924 let entity_data = vec![0xDD; 16];
4925 let row = KvRow {
4926 values: vec![
4927 CellValue::FixedBinary(entity_data.clone()),
4928 CellValue::UInt64(99),
4929 CellValue::Int64(42),
4930 ],
4931 };
4932 let key = encode_secondary_index_key(model.table_prefix, &specs[0], &model, &row).unwrap();
4933 let decoded =
4934 decode_secondary_index_key(model.table_prefix, &specs[0], &model, &key).unwrap();
4935
4936 assert_eq!(decoded.primary_key_values.len(), 2);
4937 assert!(matches!(
4938 &decoded.primary_key_values[0],
4939 CellValue::FixedBinary(v) if *v == entity_data
4940 ));
4941 assert!(matches!(
4942 &decoded.primary_key_values[1],
4943 CellValue::UInt64(99)
4944 ));
4945 let tag_idx = *model.columns_by_name.get("tag").unwrap();
4946 assert!(matches!(
4947 decoded.values.get(&tag_idx),
4948 Some(CellValue::Int64(42))
4949 ));
4950 }
4951
4952 #[test]
4953 fn table_versioned_convenience() {
4954 let client = StoreClient::new("http://localhost:10000");
4955 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
4956 .table_versioned(
4957 "documents",
4958 vec![
4959 TableColumnConfig::new("doc_id", DataType::FixedSizeBinary(32), false),
4960 TableColumnConfig::new("version", DataType::UInt64, false),
4961 TableColumnConfig::new("title", DataType::Utf8, false),
4962 ],
4963 "doc_id",
4964 "version",
4965 vec![],
4966 )
4967 .unwrap();
4968 assert_eq!(schema.table_count(), 1);
4969 }
4970
4971 #[test]
4972 fn single_column_pk_backward_compat() {
4973 let config = KvTableConfig::new(
4975 0,
4976 vec![
4977 TableColumnConfig::new("id", DataType::Int64, false),
4978 TableColumnConfig::new("name", DataType::Utf8, true),
4979 ],
4980 vec!["id".to_string()],
4981 vec![],
4982 )
4983 .unwrap();
4984 let model = TableModel::from_config(&config).unwrap();
4985 assert_eq!(model.primary_key_indices.len(), 1);
4986 assert_eq!(model.primary_key_indices[0], 0);
4987 assert_eq!(model.primary_key_width, 8);
4988
4989 let pk = CellValue::Int64(42);
4990 let key = encode_primary_key(0, &[&pk], &model).expect("pk key encodes");
4991 let decoded = decode_primary_key(0, &key, &model).unwrap();
4992 assert_eq!(decoded.len(), 1);
4993 assert!(matches!(&decoded[0], CellValue::Int64(42)));
4994 }
4995
4996 #[test]
4997 fn partial_prefix_upper_bound_fills_trailing_pk_bytes() {
4998 let config = KvTableConfig::new(
5003 0,
5004 vec![
5005 TableColumnConfig::new("entity", DataType::FixedSizeBinary(16), false),
5006 TableColumnConfig::new("version", DataType::UInt64, false),
5007 ],
5008 vec!["entity".to_string(), "version".to_string()],
5009 vec![],
5010 )
5011 .unwrap();
5012 let model = TableModel::from_config(&config).unwrap();
5013 assert_eq!(model.primary_key_width, 24); let entity = CellValue::FixedBinary(vec![0xAA; 16]);
5016 let upper =
5018 encode_primary_key_bound(0, &[&entity], &model, true).expect("pk bound encodes");
5019
5020 assert_eq!(primary_payload(&model, &upper, 0, 16), vec![0xAA; 16]);
5022 assert_eq!(
5024 primary_payload(&model, &upper, 16, 8),
5025 vec![0xFF; 8],
5026 "trailing PK column (version) must be 0xFF for upper bound"
5027 );
5028 assert!(primary_payload(
5030 &model,
5031 &upper,
5032 24,
5033 model.primary_key_prefix.max_payload_len() - 24
5034 )
5035 .iter()
5036 .all(|&b| b == 0xFF));
5037
5038 let lower =
5040 encode_primary_key_bound(0, &[&entity], &model, false).expect("pk bound encodes");
5041 assert_eq!(primary_payload(&model, &lower, 0, 16), vec![0xAA; 16]);
5042 assert_eq!(
5043 lower,
5044 model
5045 .primary_key_prefix
5046 .encode(&[0xAA; 16])
5047 .expect("encode entity prefix"),
5048 "lower bound must be exactly the encoded prefix"
5049 );
5050 }
5051
5052 #[test]
5057 fn composite_pk_range_pushdown_entity_eq_version_lte() {
5058 let config = KvTableConfig::new(
5062 0,
5063 vec![
5064 TableColumnConfig::new("entity", DataType::FixedSizeBinary(16), false),
5065 TableColumnConfig::new("version", DataType::UInt64, false),
5066 TableColumnConfig::new("data", DataType::Utf8, true),
5067 ],
5068 vec!["entity".to_string(), "version".to_string()],
5069 vec![],
5070 )
5071 .unwrap();
5072 let model = TableModel::from_config(&config).unwrap();
5073
5074 let mut pred = QueryPredicate::default();
5076 pred.constraints
5077 .insert(0, PredicateConstraint::FixedBinaryEq(vec![0xCC; 16]));
5078 pred.constraints.insert(
5079 1,
5080 PredicateConstraint::UInt64Range {
5081 min: None,
5082 max: Some(42),
5083 },
5084 );
5085
5086 let ranges = pred.primary_key_ranges(&model).unwrap();
5087 assert_eq!(ranges.len(), 1, "should produce exactly one range");
5088
5089 let range = &ranges[0];
5090
5091 let expected_start = encode_primary_key(
5093 0,
5094 &[
5095 &CellValue::FixedBinary(vec![0xCC; 16]),
5096 &CellValue::UInt64(0),
5097 ],
5098 &model,
5099 )
5100 .expect("pk key encodes");
5101 assert_eq!(
5102 range.start, expected_start,
5103 "start should be entity=CC..CC, version=0"
5104 );
5105
5106 let expected_end_prefix = encode_primary_key(
5108 0,
5109 &[
5110 &CellValue::FixedBinary(vec![0xCC; 16]),
5111 &CellValue::UInt64(42),
5112 ],
5113 &model,
5114 )
5115 .expect("pk key encodes");
5116 assert_eq!(
5118 primary_payload(&model, &range.end, 0, model.primary_key_width),
5119 primary_payload(&model, &expected_end_prefix, 0, model.primary_key_width),
5120 "end prefix should be entity=CC..CC, version=42"
5121 );
5122 assert!(
5124 primary_payload(
5125 &model,
5126 &range.end,
5127 model.primary_key_width,
5128 model.primary_key_prefix.max_payload_len() - model.primary_key_width
5129 )
5130 .iter()
5131 .all(|&b| b == 0xFF),
5132 "end trailing bytes should be 0xFF"
5133 );
5134
5135 let full_range = primary_key_prefix_range(0);
5137 assert_ne!(
5138 range.start, full_range.start,
5139 "range must not be a full table scan"
5140 );
5141 }
5142
5143 #[test]
5144 fn composite_pk_range_pushdown_entity_eq_only() {
5145 let config = KvTableConfig::new(
5149 0,
5150 vec![
5151 TableColumnConfig::new("entity", DataType::FixedSizeBinary(16), false),
5152 TableColumnConfig::new("version", DataType::UInt64, false),
5153 ],
5154 vec!["entity".to_string(), "version".to_string()],
5155 vec![],
5156 )
5157 .unwrap();
5158 let model = TableModel::from_config(&config).unwrap();
5159
5160 let mut pred = QueryPredicate::default();
5161 pred.constraints
5162 .insert(0, PredicateConstraint::FixedBinaryEq(vec![0xDD; 16]));
5163
5164 let ranges = pred.primary_key_ranges(&model).unwrap();
5165 assert_eq!(ranges.len(), 1);
5166
5167 let range = &ranges[0];
5168 assert_eq!(primary_payload(&model, &range.start, 0, 16), vec![0xDD; 16]);
5170 assert_eq!(primary_payload(&model, &range.end, 0, 16), vec![0xDD; 16]);
5172 assert!(
5173 primary_payload(
5174 &model,
5175 &range.end,
5176 16,
5177 model.primary_key_prefix.max_payload_len() - 16
5178 )
5179 .iter()
5180 .all(|&b| b == 0xFF),
5181 "after entity bytes, everything should be 0xFF"
5182 );
5183 }
5184
5185 #[test]
5186 fn fixed_binary_eq_constraint_extracted() {
5187 let config = KvTableConfig::new(
5188 0,
5189 vec![
5190 TableColumnConfig::new("entity", DataType::FixedSizeBinary(16), false),
5191 TableColumnConfig::new("version", DataType::UInt64, false),
5192 ],
5193 vec!["entity".to_string(), "version".to_string()],
5194 vec![],
5195 )
5196 .unwrap();
5197 let model = TableModel::from_config(&config).unwrap();
5198
5199 use datafusion::logical_expr::col;
5201 let entity_literal =
5202 Expr::Literal(ScalarValue::FixedSizeBinary(16, Some(vec![0xAA; 16])), None);
5203 let filter = col("entity").eq(entity_literal);
5204
5205 assert!(
5206 QueryPredicate::supports_filter(&filter, &model),
5207 "FixedSizeBinary equality should be supported"
5208 );
5209
5210 let pred = QueryPredicate::from_filters(&[filter], &model);
5211 assert!(
5212 matches!(
5213 pred.constraints.get(&0),
5214 Some(PredicateConstraint::FixedBinaryEq(v)) if *v == vec![0xAA; 16]
5215 ),
5216 "should extract FixedBinaryEq constraint"
5217 );
5218 }
5219
5220 #[test]
5221 fn uint64_range_constraint_extracted() {
5222 let config = KvTableConfig::new(
5223 0,
5224 vec![
5225 TableColumnConfig::new("version", DataType::UInt64, false),
5226 TableColumnConfig::new("data", DataType::Utf8, true),
5227 ],
5228 vec!["version".to_string()],
5229 vec![],
5230 )
5231 .unwrap();
5232 let model = TableModel::from_config(&config).unwrap();
5233
5234 use datafusion::logical_expr::col;
5235 let filter = col("version").lt_eq(Expr::Literal(ScalarValue::UInt64(Some(42)), None));
5236
5237 assert!(
5238 QueryPredicate::supports_filter(&filter, &model),
5239 "UInt64 range should be supported"
5240 );
5241
5242 let pred = QueryPredicate::from_filters(&[filter], &model);
5243 assert!(
5244 matches!(
5245 pred.constraints.get(&0),
5246 Some(PredicateConstraint::UInt64Range {
5247 min: None,
5248 max: Some(42)
5249 })
5250 ),
5251 "should extract UInt64Range with max=42"
5252 );
5253 }
5254
5255 #[test]
5256 fn uint64_range_constraint_supports_values_above_i64_max() {
5257 let config = KvTableConfig::new(
5258 0,
5259 vec![
5260 TableColumnConfig::new("version", DataType::UInt64, false),
5261 TableColumnConfig::new("data", DataType::Utf8, true),
5262 ],
5263 vec!["version".to_string()],
5264 vec![],
5265 )
5266 .unwrap();
5267 let model = TableModel::from_config(&config).unwrap();
5268
5269 let threshold = (1u64 << 63) + 5;
5270 use datafusion::logical_expr::col;
5271 let filter =
5272 col("version").gt_eq(Expr::Literal(ScalarValue::UInt64(Some(threshold)), None));
5273
5274 assert!(QueryPredicate::supports_filter(&filter, &model));
5275
5276 let pred = QueryPredicate::from_filters(&[filter], &model);
5277 assert!(matches!(
5278 pred.constraints.get(&0),
5279 Some(PredicateConstraint::UInt64Range {
5280 min: Some(v),
5281 max: None
5282 }) if *v == threshold
5283 ));
5284 }
5285
5286 #[test]
5287 fn unsupported_uint64_comparison_does_not_force_contradiction() {
5288 let config = KvTableConfig::new(
5289 0,
5290 vec![
5291 TableColumnConfig::new("version", DataType::UInt64, false),
5292 TableColumnConfig::new("data", DataType::Utf8, true),
5293 ],
5294 vec!["version".to_string()],
5295 vec![],
5296 )
5297 .unwrap();
5298 let model = TableModel::from_config(&config).unwrap();
5299
5300 use datafusion::logical_expr::col;
5301 let unsupported = col("version").gt(Expr::Literal(ScalarValue::Int64(Some(-1)), None));
5302
5303 assert!(
5304 !QueryPredicate::supports_filter(&unsupported, &model),
5305 "negative Int64 literal on UInt64 column should not be pushdown-supported"
5306 );
5307
5308 let pred = QueryPredicate::from_filters(&[unsupported], &model);
5309 assert!(
5310 !pred.contradiction,
5311 "unsupported filter must not collapse scan to empty result"
5312 );
5313 assert!(
5314 pred.constraints.is_empty(),
5315 "unsupported filter must not contribute pushed constraints"
5316 );
5317 }
5318
5319 #[test]
5320 fn unsupported_uint64_comparison_in_and_keeps_supported_sibling() {
5321 let config = KvTableConfig::new(
5322 0,
5323 vec![
5324 TableColumnConfig::new("version", DataType::UInt64, false),
5325 TableColumnConfig::new("data", DataType::Utf8, true),
5326 ],
5327 vec!["version".to_string()],
5328 vec![],
5329 )
5330 .unwrap();
5331 let model = TableModel::from_config(&config).unwrap();
5332
5333 use datafusion::logical_expr::col;
5334 let supported = col("version").gt_eq(Expr::Literal(ScalarValue::UInt64(Some(10)), None));
5335 let unsupported = col("version").gt(Expr::Literal(ScalarValue::Int64(Some(-1)), None));
5336 let filter = supported.and(unsupported);
5337
5338 assert!(
5339 !QueryPredicate::supports_filter(&filter, &model),
5340 "mixed AND should not be marked fully pushdown-supported"
5341 );
5342
5343 let pred = QueryPredicate::from_filters(&[filter], &model);
5344 assert!(!pred.contradiction);
5345 assert!(matches!(
5346 pred.constraints.get(&0),
5347 Some(PredicateConstraint::UInt64Range {
5348 min: Some(10),
5349 max: None
5350 })
5351 ));
5352 }
5353
5354 #[test]
5355 fn uint64_in_list_pushdown() {
5356 let config = KvTableConfig::new(
5357 0,
5358 vec![
5359 TableColumnConfig::new("version", DataType::UInt64, false),
5360 TableColumnConfig::new("data", DataType::Utf8, true),
5361 ],
5362 vec!["version".to_string()],
5363 vec![],
5364 )
5365 .unwrap();
5366 let model = TableModel::from_config(&config).unwrap();
5367
5368 use datafusion::logical_expr::{col, in_list};
5369 let filter = in_list(
5370 col("version"),
5371 vec![
5372 Expr::Literal(ScalarValue::UInt64(Some(1)), None),
5373 Expr::Literal(ScalarValue::UInt64(Some(5)), None),
5374 Expr::Literal(ScalarValue::UInt64(Some(10)), None),
5375 ],
5376 false,
5377 );
5378
5379 assert!(QueryPredicate::supports_filter(&filter, &model));
5380 let pred = QueryPredicate::from_filters(&[filter], &model);
5381 assert!(
5382 matches!(pred.constraints.get(&0), Some(PredicateConstraint::UInt64In(v)) if v.len() == 3),
5383 "should extract UInt64In with 3 values"
5384 );
5385 }
5386
5387 #[test]
5388 fn uint64_in_list_pushdown_supports_values_above_i64_max() {
5389 let config = KvTableConfig::new(
5390 0,
5391 vec![
5392 TableColumnConfig::new("version", DataType::UInt64, false),
5393 TableColumnConfig::new("data", DataType::Utf8, true),
5394 ],
5395 vec!["version".to_string()],
5396 vec![],
5397 )
5398 .unwrap();
5399 let model = TableModel::from_config(&config).unwrap();
5400
5401 let huge = 1u64 << 63;
5402 use datafusion::logical_expr::{col, in_list};
5403 let filter = in_list(
5404 col("version"),
5405 vec![
5406 Expr::Literal(ScalarValue::UInt64(Some(1)), None),
5407 Expr::Literal(ScalarValue::UInt64(Some(huge)), None),
5408 ],
5409 false,
5410 );
5411
5412 assert!(QueryPredicate::supports_filter(&filter, &model));
5413 let pred = QueryPredicate::from_filters(&[filter], &model);
5414 assert!(matches!(
5415 pred.constraints.get(&0),
5416 Some(PredicateConstraint::UInt64In(v)) if v.contains(&huge) && v.len() == 2
5417 ));
5418 }
5419
5420 #[test]
5421 fn fixed_binary_in_list_pushdown() {
5422 let config = KvTableConfig::new(
5423 0,
5424 vec![
5425 TableColumnConfig::new("entity", DataType::FixedSizeBinary(16), false),
5426 TableColumnConfig::new("data", DataType::Utf8, true),
5427 ],
5428 vec!["entity".to_string()],
5429 vec![],
5430 )
5431 .unwrap();
5432 let model = TableModel::from_config(&config).unwrap();
5433
5434 use datafusion::logical_expr::{col, in_list};
5435 let filter = in_list(
5436 col("entity"),
5437 vec![
5438 Expr::Literal(ScalarValue::FixedSizeBinary(16, Some(vec![0xAA; 16])), None),
5439 Expr::Literal(ScalarValue::FixedSizeBinary(16, Some(vec![0xBB; 16])), None),
5440 ],
5441 false,
5442 );
5443
5444 assert!(QueryPredicate::supports_filter(&filter, &model));
5445 let pred = QueryPredicate::from_filters(&[filter], &model);
5446 assert!(
5447 matches!(
5448 pred.constraints.get(&0),
5449 Some(PredicateConstraint::FixedBinaryIn(v)) if v.len() == 2
5450 ),
5451 "should extract FixedBinaryIn with 2 values"
5452 );
5453
5454 let ranges = pred.primary_key_ranges(&model).unwrap();
5456 assert_eq!(ranges.len(), 2, "should produce one range per entity");
5457 }
5458
5459 #[test]
5460 fn decimal256_range_pushdown() {
5461 let config = KvTableConfig::new(
5462 0,
5463 vec![
5464 TableColumnConfig::new("id", DataType::Int64, false),
5465 TableColumnConfig::new("big_val", DataType::Decimal256(76, 0), false),
5466 ],
5467 vec!["id".to_string()],
5468 vec![IndexSpec::new("big_idx", vec!["big_val".to_string()]).unwrap()],
5469 )
5470 .unwrap();
5471 let model = TableModel::from_config(&config).unwrap();
5472
5473 use datafusion::logical_expr::col;
5474 let filter = col("big_val").gt_eq(Expr::Literal(
5475 ScalarValue::Decimal256(Some(i256::from(100i64)), 76, 0),
5476 None,
5477 ));
5478
5479 assert!(
5480 QueryPredicate::supports_filter(&filter, &model),
5481 "Decimal256 range should be supported"
5482 );
5483
5484 let pred = QueryPredicate::from_filters(&[filter], &model);
5485 let big_idx = *model.columns_by_name.get("big_val").unwrap();
5486 assert!(
5487 matches!(
5488 pred.constraints.get(&big_idx),
5489 Some(PredicateConstraint::Decimal256Range {
5490 min: Some(_),
5491 max: None
5492 })
5493 ),
5494 "should extract Decimal256Range with min=100, no max"
5495 );
5496
5497 let val_in = CellValue::Decimal256(i256::from(200i64));
5499 let val_out = CellValue::Decimal256(i256::from(50i64));
5500 let constraint = pred.constraints.get(&big_idx).unwrap();
5501 assert!(matches_constraint(&val_in, constraint));
5502 assert!(!matches_constraint(&val_out, constraint));
5503 }
5504
5505 #[test]
5506 fn uint64_constraint_matching_does_not_wrap_large_values() {
5507 let gt_zero = PredicateConstraint::UInt64Range {
5508 min: Some(1),
5509 max: None,
5510 };
5511 assert!(matches_constraint(&CellValue::UInt64(1u64 << 63), >_zero));
5512 assert!(!matches_constraint(&CellValue::UInt64(0), >_zero));
5513
5514 let in_list = PredicateConstraint::UInt64In(vec![1, 2, 3]);
5515 assert!(matches_constraint(&CellValue::UInt64(2), &in_list));
5516 assert!(!matches_constraint(
5517 &CellValue::UInt64(1u64 << 63),
5518 &in_list
5519 ));
5520 }
5521
5522 #[test]
5523 fn uint64_empty_range_produces_no_pk_ranges() {
5524 let config = KvTableConfig::new(
5525 0,
5526 vec![TableColumnConfig::new("version", DataType::UInt64, false)],
5527 vec!["version".to_string()],
5528 vec![],
5529 )
5530 .unwrap();
5531 let model = TableModel::from_config(&config).unwrap();
5532 let mut pred = QueryPredicate::default();
5533 pred.constraints.insert(
5534 0,
5535 PredicateConstraint::UInt64Range {
5536 min: Some(10),
5537 max: Some(9),
5538 },
5539 );
5540
5541 let ranges = pred.primary_key_ranges(&model).unwrap();
5542 assert!(ranges.is_empty());
5543 }
5544
5545 #[test]
5546 fn utf8_primary_key_encoding_supports_unicode_and_long_values() {
5547 let config = KvTableConfig::new(
5548 0,
5549 vec![TableColumnConfig::new("id", DataType::Utf8, false)],
5550 vec!["id".to_string()],
5551 vec![],
5552 )
5553 .unwrap();
5554 let model = TableModel::from_config(&config).unwrap();
5555
5556 let row_non_ascii = KvRow {
5557 values: vec![CellValue::Utf8("naive-cafe-e9".replace("e9", "\u{00E9}"))],
5558 };
5559 let key_non_ascii = encode_primary_key_from_row(model.table_prefix, &row_non_ascii, &model)
5560 .expect("non-ascii PK should encode");
5561 let decoded_non_ascii = decode_primary_key(model.table_prefix, &key_non_ascii, &model)
5562 .expect("non-ascii PK should decode");
5563 assert!(matches!(
5564 decoded_non_ascii.as_slice(),
5565 [CellValue::Utf8(value)] if value == "naive-cafe-\u{00E9}"
5566 ));
5567
5568 let row_too_long = KvRow {
5569 values: vec![CellValue::Utf8("abcdefghijklmnopq".to_string())],
5570 };
5571 let key_too_long = encode_primary_key_from_row(model.table_prefix, &row_too_long, &model)
5572 .expect("long UTF-8 PK should encode");
5573 let decoded_too_long = decode_primary_key(model.table_prefix, &key_too_long, &model)
5574 .expect("long UTF-8 PK should decode");
5575 assert!(matches!(
5576 decoded_too_long.as_slice(),
5577 [CellValue::Utf8(value)] if value == "abcdefghijklmnopq"
5578 ));
5579 }
5580
5581 #[test]
5582 fn utf8_primary_key_encodes_at_max_key_payload_and_rejects_overflow() {
5583 let config = KvTableConfig::new(
5584 0,
5585 vec![TableColumnConfig::new("id", DataType::Utf8, false)],
5586 vec!["id".to_string()],
5587 vec![],
5588 )
5589 .unwrap();
5590 let model = TableModel::from_config(&config).unwrap();
5591 let max_payload = model.primary_key_prefix.max_payload_len();
5592 let max_value = "a".repeat(max_payload - 1);
5593 let overflow_value = "a".repeat(max_payload);
5594
5595 let key = encode_primary_key_from_row(
5596 model.table_prefix,
5597 &KvRow {
5598 values: vec![CellValue::Utf8(max_value.clone())],
5599 },
5600 &model,
5601 )
5602 .expect("max-length UTF-8 PK should encode");
5603 assert_eq!(key.len(), exoware_sdk::keys::MAX_KEY_LEN);
5604 let decoded = decode_primary_key(model.table_prefix, &key, &model)
5605 .expect("max-length PK should decode");
5606 assert!(matches!(
5607 decoded.as_slice(),
5608 [CellValue::Utf8(value)] if value == &max_value
5609 ));
5610
5611 let err = encode_primary_key_from_row(
5612 model.table_prefix,
5613 &KvRow {
5614 values: vec![CellValue::Utf8(overflow_value)],
5615 },
5616 &model,
5617 )
5618 .expect_err("UTF-8 PK exceeding key payload should be rejected");
5619 assert!(err.contains("failed to encode primary key: key length 255 exceeds max 254"));
5620 }
5621
5622 #[test]
5623 fn utf8_primary_key_round_trips_embedded_nul() {
5624 let config = KvTableConfig::new(
5625 0,
5626 vec![TableColumnConfig::new("id", DataType::Utf8, false)],
5627 vec!["id".to_string()],
5628 vec![],
5629 )
5630 .unwrap();
5631 let model = TableModel::from_config(&config).unwrap();
5632 let row = KvRow {
5633 values: vec![CellValue::Utf8("AB\0CD".to_string())],
5634 };
5635
5636 let key = encode_primary_key_from_row(model.table_prefix, &row, &model)
5637 .expect("embedded NUL in key text must encode");
5638 let decoded =
5639 decode_primary_key(model.table_prefix, &key, &model).expect("embedded NUL must decode");
5640 assert!(matches!(
5641 decoded.as_slice(),
5642 [CellValue::Utf8(value)] if value == "AB\0CD"
5643 ));
5644 }
5645
5646 #[test]
5647 fn utf8_index_key_round_trips_embedded_nul() {
5648 let config = KvTableConfig::new(
5649 0,
5650 vec![
5651 TableColumnConfig::new("id", DataType::Int64, false),
5652 TableColumnConfig::new("tag", DataType::Utf8, false),
5653 ],
5654 vec!["id".to_string()],
5655 vec![IndexSpec::new("tag_idx", vec!["tag".to_string()]).unwrap()],
5656 )
5657 .unwrap();
5658 let model = TableModel::from_config(&config).unwrap();
5659 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
5660 let row = KvRow {
5661 values: vec![CellValue::Int64(1), CellValue::Utf8("AB\0CD".to_string())],
5662 };
5663
5664 let key = encode_secondary_index_key(model.table_prefix, &specs[0], &model, &row)
5665 .expect("embedded NUL in index key text must encode");
5666 let decoded = decode_secondary_index_key(model.table_prefix, &specs[0], &model, &key)
5667 .expect("embedded NUL index key must decode");
5668 assert!(matches!(
5669 decoded.values.get(&1),
5670 Some(CellValue::Utf8(value)) if value == "AB\0CD"
5671 ));
5672 }
5673
5674 #[test]
5675 fn secondary_index_with_long_utf8_primary_key_encodes_at_max_payload_and_rejects_overflow() {
5676 let config = KvTableConfig::new(
5677 0,
5678 vec![
5679 TableColumnConfig::new("id", DataType::Utf8, false),
5680 TableColumnConfig::new("tag", DataType::Utf8, false),
5681 ],
5682 vec!["id".to_string()],
5683 vec![IndexSpec::new("tag_idx", vec!["tag".to_string()]).unwrap()],
5684 )
5685 .unwrap();
5686 let model = TableModel::from_config(&config).unwrap();
5687 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
5688 let spec = &specs[0];
5689 let max_payload = spec.prefix.max_payload_len();
5690 let max_tag = "t".to_string();
5691 let max_id = "i".repeat(max_payload - encode_string_variable(&max_tag).unwrap().len() - 1);
5692 let overflow_id = format!("{max_id}x");
5693
5694 let key = encode_secondary_index_key(
5695 model.table_prefix,
5696 spec,
5697 &model,
5698 &KvRow {
5699 values: vec![
5700 CellValue::Utf8(max_id.clone()),
5701 CellValue::Utf8(max_tag.clone()),
5702 ],
5703 },
5704 )
5705 .expect("secondary key at max payload should encode");
5706 assert_eq!(key.len(), exoware_sdk::keys::MAX_KEY_LEN);
5707 let decoded =
5708 decode_secondary_index_key(model.table_prefix, spec, &model, &key).expect("decode");
5709 assert!(matches!(
5710 decoded.values.get(&1),
5711 Some(CellValue::Utf8(value)) if value == &max_tag
5712 ));
5713 assert!(matches!(
5714 decoded.primary_key_values.as_slice(),
5715 [CellValue::Utf8(value)] if value == &max_id
5716 ));
5717
5718 let err = encode_secondary_index_key(
5719 model.table_prefix,
5720 spec,
5721 &model,
5722 &KvRow {
5723 values: vec![CellValue::Utf8(overflow_id), CellValue::Utf8(max_tag)],
5724 },
5725 )
5726 .expect_err("secondary key exceeding max payload should be rejected");
5727 assert!(
5728 err.contains("failed to encode index 'tag_idx' key: key length 255 exceeds max 254")
5729 );
5730 }
5731
5732 #[test]
5733 fn secondary_index_from_parts_with_long_utf8_primary_key_rejects_overflow() {
5734 let config = KvTableConfig::new(
5735 0,
5736 vec![
5737 TableColumnConfig::new("id", DataType::Utf8, false),
5738 TableColumnConfig::new("tag", DataType::Utf8, false),
5739 ],
5740 vec!["id".to_string()],
5741 vec![IndexSpec::new("tag_idx", vec!["tag".to_string()]).unwrap()],
5742 )
5743 .unwrap();
5744 let model = TableModel::from_config(&config).unwrap();
5745 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
5746 let spec = &specs[0];
5747 let max_payload = spec.prefix.max_payload_len();
5748 let max_tag = "t".to_string();
5749 let max_id = "i".repeat(max_payload - encode_string_variable(&max_tag).unwrap().len() - 1);
5750 let overflow_id = format!("{max_id}x");
5751 let max_row = KvRow {
5752 values: vec![
5753 CellValue::Utf8(max_id.clone()),
5754 CellValue::Utf8(max_tag.clone()),
5755 ],
5756 };
5757 let encoded_row = encode_base_row_value(&max_row, &model).expect("encode row");
5758 let archived = decode_stored_row(&encoded_row).expect("archive row");
5759
5760 let key = encode_secondary_index_key_from_parts(
5761 model.table_prefix,
5762 spec,
5763 &model,
5764 &[CellValue::Utf8(max_id.clone())],
5765 &archived,
5766 )
5767 .expect("backfill path should encode max payload");
5768 assert_eq!(key.len(), exoware_sdk::keys::MAX_KEY_LEN);
5769
5770 let err = encode_secondary_index_key_from_parts(
5771 model.table_prefix,
5772 spec,
5773 &model,
5774 &[CellValue::Utf8(overflow_id)],
5775 &archived,
5776 )
5777 .expect_err("backfill path overflow should be rejected");
5778 assert!(err
5779 .to_string()
5780 .contains("failed to encode index 'tag_idx' key: key length 255 exceeds max 254"));
5781 }
5782
5783 #[test]
5784 fn primary_key_type_mismatch_returns_error_instead_of_panicking() {
5785 let config = KvTableConfig::new(
5786 0,
5787 vec![TableColumnConfig::new("id", DataType::UInt64, false)],
5788 vec!["id".to_string()],
5789 vec![],
5790 )
5791 .unwrap();
5792 let model = TableModel::from_config(&config).unwrap();
5793 let row = KvRow {
5794 values: vec![CellValue::Int64(7)],
5795 };
5796
5797 let err = encode_primary_key_from_row(model.table_prefix, &row, &model)
5798 .expect_err("mismatched PK type should return an error");
5799 assert!(err.contains("type mismatch while encoding key value"));
5800 }
5801
5802 #[test]
5803 fn choose_index_plan_uses_fixed_binary_leading_constraint() {
5804 let config = KvTableConfig::new(
5805 0,
5806 vec![
5807 TableColumnConfig::new("id", DataType::Int64, false),
5808 TableColumnConfig::new("entity", DataType::FixedSizeBinary(16), false),
5809 ],
5810 vec!["id".to_string()],
5811 vec![IndexSpec::new("entity_idx", vec!["entity".to_string()]).unwrap()],
5812 )
5813 .unwrap();
5814 let model = TableModel::from_config(&config).unwrap();
5815 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
5816
5817 use datafusion::logical_expr::col;
5818 let filter = col("entity").eq(Expr::Literal(
5819 ScalarValue::FixedSizeBinary(16, Some(vec![0xAB; 16])),
5820 None,
5821 ));
5822 let pred = QueryPredicate::from_filters(&[filter], &model);
5823 let plan = pred
5824 .choose_index_plan(&model, &specs)
5825 .unwrap()
5826 .expect("fixed-binary equality should choose an index");
5827
5828 assert_eq!(plan.constrained_prefix_len, 1);
5829 assert_eq!(plan.ranges.len(), 1);
5830 let range = &plan.ranges[0];
5831 assert_eq!(
5832 index_payload(&specs[0], &range.start, 0, 16),
5833 vec![0xAB; 16]
5834 );
5835 assert_eq!(index_payload(&specs[0], &range.end, 0, 16), vec![0xAB; 16]);
5836 }
5837
5838 #[test]
5839 fn choose_index_plan_uses_decimal256_leading_constraint() {
5840 let config = KvTableConfig::new(
5841 0,
5842 vec![
5843 TableColumnConfig::new("id", DataType::Int64, false),
5844 TableColumnConfig::new("big_val", DataType::Decimal256(76, 0), false),
5845 ],
5846 vec!["id".to_string()],
5847 vec![IndexSpec::new("big_idx", vec!["big_val".to_string()]).unwrap()],
5848 )
5849 .unwrap();
5850 let model = TableModel::from_config(&config).unwrap();
5851 let specs = model.resolve_index_specs(&config.index_specs).unwrap();
5852
5853 use datafusion::logical_expr::col;
5854 let filter = col("big_val").gt_eq(Expr::Literal(
5855 ScalarValue::Decimal256(Some(i256::from(100i64)), 76, 0),
5856 None,
5857 ));
5858 let pred = QueryPredicate::from_filters(&[filter], &model);
5859 let plan = pred
5860 .choose_index_plan(&model, &specs)
5861 .unwrap()
5862 .expect("decimal256 range should choose an index");
5863
5864 assert_eq!(plan.constrained_prefix_len, 1);
5865 assert_eq!(plan.ranges.len(), 1);
5866 let range = &plan.ranges[0];
5867 assert_eq!(
5868 index_payload(&specs[0], &range.start, 0, 32),
5869 encode_i256_ordered(i256::from(100i64)).to_vec()
5870 );
5871 }
5872
5873 #[tokio::test]
5874 async fn backfill_added_indexes_writes_entries_for_existing_rows() {
5875 let state = MockState {
5876 kv: Arc::new(Mutex::new(BTreeMap::new())),
5877 range_calls: Arc::new(AtomicUsize::new(0)),
5878 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
5879 sequence_number: Arc::new(AtomicU64::new(0)),
5880 };
5881 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
5882 let client = StoreClient::new(&base_url);
5883
5884 let seed_schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
5885 .table(
5886 "orders",
5887 vec![
5888 TableColumnConfig::new("id", DataType::Int64, false),
5889 TableColumnConfig::new("status", DataType::Utf8, false),
5890 TableColumnConfig::new("amount_cents", DataType::Int64, false),
5891 ],
5892 vec!["id".to_string()],
5893 vec![],
5894 )
5895 .expect("seed schema");
5896 let mut writer = seed_schema.batch_writer();
5897 for i in 0..6i64 {
5898 writer
5899 .insert(
5900 "orders",
5901 vec![
5902 CellValue::Int64(i),
5903 CellValue::Utf8(if i % 2 == 0 { "open" } else { "closed" }.to_string()),
5904 CellValue::Int64(i * 10),
5905 ],
5906 )
5907 .expect("seed row");
5908 }
5909 writer.flush().await.expect("seed flush");
5910
5911 {
5912 let guard = state.kv.lock().expect("kv mutex poisoned");
5913 let base_rows = guard
5914 .keys()
5915 .filter(|key| matches_primary_key(0, key))
5916 .count();
5917 let index_rows = guard
5918 .keys()
5919 .filter(|key| matches_secondary_index_key(0, 1, key))
5920 .count();
5921 assert_eq!(base_rows, 6);
5922 assert_eq!(index_rows, 0);
5923 }
5924
5925 let backfill_schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
5926 .table(
5927 "orders",
5928 vec![
5929 TableColumnConfig::new("id", DataType::Int64, false),
5930 TableColumnConfig::new("status", DataType::Utf8, false),
5931 TableColumnConfig::new("amount_cents", DataType::Int64, false),
5932 ],
5933 vec!["id".to_string()],
5934 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
5935 .expect("valid index")
5936 .with_cover_columns(vec!["amount_cents".to_string()])],
5937 )
5938 .expect("backfill schema");
5939 let report = backfill_schema
5940 .backfill_added_indexes_with_options(
5941 "orders",
5942 &[],
5943 IndexBackfillOptions {
5944 row_batch_size: 2,
5945 start_from_primary_key: None,
5946 },
5947 )
5948 .await
5949 .expect("backfill should succeed");
5950 assert_eq!(report.scanned_rows, 6);
5951 assert_eq!(report.indexes_backfilled, 1);
5952 assert_eq!(report.index_entries_written, 6);
5953
5954 {
5955 let guard = state.kv.lock().expect("kv mutex poisoned");
5956 let index_rows = guard
5957 .keys()
5958 .filter(|key| matches_secondary_index_key(0, 1, key))
5959 .count();
5960 assert_eq!(index_rows, 6);
5961 let (_, sample_value) = guard
5962 .iter()
5963 .find(|(key, _)| matches_secondary_index_key(0, 1, key))
5964 .expect("backfill should create index entry");
5965 let archived = decode_stored_row(sample_value.as_ref())
5966 .expect("covering value must be valid codec");
5967 assert_eq!(archived.values.len(), 3);
5968 }
5969
5970 let _ = shutdown_tx.send(());
5971 }
5972
5973 #[tokio::test]
5974 async fn backfill_added_indexes_writes_zorder_entries_for_existing_rows() {
5975 let state = MockState {
5976 kv: Arc::new(Mutex::new(BTreeMap::new())),
5977 range_calls: Arc::new(AtomicUsize::new(0)),
5978 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
5979 sequence_number: Arc::new(AtomicU64::new(0)),
5980 };
5981 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
5982 let client = StoreClient::new(&base_url);
5983
5984 let seed_schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
5985 .table(
5986 "points",
5987 vec![
5988 TableColumnConfig::new("x", DataType::Int64, false),
5989 TableColumnConfig::new("y", DataType::Int64, false),
5990 TableColumnConfig::new("id", DataType::Int64, false),
5991 TableColumnConfig::new("value", DataType::Int64, false),
5992 ],
5993 vec!["id".to_string()],
5994 vec![],
5995 )
5996 .expect("seed schema");
5997 let mut writer = seed_schema.batch_writer();
5998 for (x, y, id, value) in [(1, 1, 11, 110), (1, 2, 12, 120), (2, 1, 21, 210)] {
5999 writer
6000 .insert(
6001 "points",
6002 vec![
6003 CellValue::Int64(x),
6004 CellValue::Int64(y),
6005 CellValue::Int64(id),
6006 CellValue::Int64(value),
6007 ],
6008 )
6009 .expect("seed row");
6010 }
6011 writer.flush().await.expect("seed flush");
6012
6013 let backfill_schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
6014 .table(
6015 "points",
6016 vec![
6017 TableColumnConfig::new("x", DataType::Int64, false),
6018 TableColumnConfig::new("y", DataType::Int64, false),
6019 TableColumnConfig::new("id", DataType::Int64, false),
6020 TableColumnConfig::new("value", DataType::Int64, false),
6021 ],
6022 vec!["id".to_string()],
6023 vec![
6024 IndexSpec::z_order("xy_z", vec!["x".to_string(), "y".to_string()])
6025 .expect("valid index")
6026 .with_cover_columns(vec!["value".to_string()]),
6027 ],
6028 )
6029 .expect("backfill schema");
6030 let report = backfill_schema
6031 .backfill_added_indexes_with_options(
6032 "points",
6033 &[],
6034 IndexBackfillOptions {
6035 row_batch_size: 2,
6036 start_from_primary_key: None,
6037 },
6038 )
6039 .await
6040 .expect("backfill should succeed");
6041 assert_eq!(report.scanned_rows, 3);
6042 assert_eq!(report.index_entries_written, 3);
6043
6044 let guard = state.kv.lock().expect("kv mutex poisoned");
6045 let index_entry = guard
6046 .keys()
6047 .find(|key| matches_secondary_index_key(0, 1, key))
6048 .cloned()
6049 .expect("z-order backfill should create index entry");
6050 let config = KvTableConfig::new(
6051 0,
6052 vec![
6053 TableColumnConfig::new("x", DataType::Int64, false),
6054 TableColumnConfig::new("y", DataType::Int64, false),
6055 TableColumnConfig::new("id", DataType::Int64, false),
6056 TableColumnConfig::new("value", DataType::Int64, false),
6057 ],
6058 vec!["id".to_string()],
6059 vec![
6060 IndexSpec::z_order("xy_z", vec!["x".to_string(), "y".to_string()]).expect("valid"),
6061 ],
6062 )
6063 .expect("config");
6064 let model = TableModel::from_config(&config).expect("model");
6065 let spec = model
6066 .resolve_index_specs(&config.index_specs)
6067 .expect("specs")
6068 .remove(0);
6069 let decoded = decode_secondary_index_key(model.table_prefix, &spec, &model, &index_entry)
6070 .expect("decode z-order key");
6071 let x_idx = *model.columns_by_name.get("x").unwrap();
6072 let y_idx = *model.columns_by_name.get("y").unwrap();
6073 assert!(matches!(
6074 decoded.values.get(&x_idx),
6075 Some(CellValue::Int64(_))
6076 ));
6077 assert!(matches!(
6078 decoded.values.get(&y_idx),
6079 Some(CellValue::Int64(_))
6080 ));
6081
6082 let _ = shutdown_tx.send(());
6083 }
6084
6085 #[tokio::test]
6086 async fn backfill_added_indexes_requires_append_only_index_evolution() {
6087 let client = StoreClient::new("http://127.0.0.1:1");
6088 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
6089 .table(
6090 "orders",
6091 vec![
6092 TableColumnConfig::new("id", DataType::Int64, false),
6093 TableColumnConfig::new("status", DataType::Utf8, false),
6094 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6095 ],
6096 vec!["id".to_string()],
6097 vec![
6098 IndexSpec::new("status_idx", vec!["status".to_string()]).expect("valid"),
6099 IndexSpec::new("amount_idx", vec!["amount_cents".to_string()]).expect("valid"),
6100 ],
6101 )
6102 .expect("schema");
6103
6104 let previous_specs =
6105 vec![IndexSpec::new("amount_idx", vec!["amount_cents".to_string()]).expect("valid")];
6106 let err = schema
6107 .backfill_added_indexes("orders", &previous_specs)
6108 .await
6109 .expect_err("non-append-only evolution should be rejected");
6110 assert!(err
6111 .to_string()
6112 .contains("index evolution must be append-only"));
6113 }
6114
6115 #[tokio::test]
6116 async fn backfill_added_indexes_is_noop_when_no_new_indexes() {
6117 let client = StoreClient::new("http://127.0.0.1:1");
6118 let existing = IndexSpec::new("status_idx", vec!["status".to_string()])
6119 .expect("valid")
6120 .with_cover_columns(vec!["amount_cents".to_string()]);
6121 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
6122 .table(
6123 "orders",
6124 vec![
6125 TableColumnConfig::new("id", DataType::Int64, false),
6126 TableColumnConfig::new("status", DataType::Utf8, false),
6127 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6128 ],
6129 vec!["id".to_string()],
6130 vec![existing.clone()],
6131 )
6132 .expect("schema");
6133
6134 let report = schema
6135 .backfill_added_indexes("orders", &[existing])
6136 .await
6137 .expect("no-op backfill should succeed");
6138 assert_eq!(report, IndexBackfillReport::default());
6139 }
6140
6141 #[tokio::test]
6142 async fn backfill_added_indexes_rejects_zero_row_batch_size() {
6143 let client = StoreClient::new("http://127.0.0.1:1");
6144 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
6145 .table(
6146 "orders",
6147 vec![
6148 TableColumnConfig::new("id", DataType::Int64, false),
6149 TableColumnConfig::new("status", DataType::Utf8, false),
6150 ],
6151 vec!["id".to_string()],
6152 vec![IndexSpec::new("status_idx", vec!["status".to_string()]).expect("valid")],
6153 )
6154 .expect("schema");
6155 let err = schema
6156 .backfill_added_indexes_with_options(
6157 "orders",
6158 &[],
6159 IndexBackfillOptions {
6160 row_batch_size: 0,
6161 start_from_primary_key: None,
6162 },
6163 )
6164 .await
6165 .expect_err("row_batch_size=0 should fail");
6166 assert!(err.to_string().contains("row_batch_size must be > 0"));
6167 }
6168
6169 #[tokio::test]
6170 async fn backfill_added_indexes_emits_progress_events() {
6171 let state = MockState {
6172 kv: Arc::new(Mutex::new(BTreeMap::new())),
6173 range_calls: Arc::new(AtomicUsize::new(0)),
6174 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
6175 sequence_number: Arc::new(AtomicU64::new(0)),
6176 };
6177 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
6178 let client = StoreClient::new(&base_url);
6179
6180 let seed_schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
6181 .table(
6182 "orders",
6183 vec![
6184 TableColumnConfig::new("id", DataType::Int64, false),
6185 TableColumnConfig::new("status", DataType::Utf8, false),
6186 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6187 ],
6188 vec!["id".to_string()],
6189 vec![],
6190 )
6191 .expect("seed schema");
6192 let mut writer = seed_schema.batch_writer();
6193 for i in 0..5i64 {
6194 writer
6195 .insert(
6196 "orders",
6197 vec![
6198 CellValue::Int64(i),
6199 CellValue::Utf8("open".to_string()),
6200 CellValue::Int64(i * 10),
6201 ],
6202 )
6203 .expect("seed row");
6204 }
6205 writer.flush().await.expect("seed flush");
6206
6207 let backfill_schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
6208 .table(
6209 "orders",
6210 vec![
6211 TableColumnConfig::new("id", DataType::Int64, false),
6212 TableColumnConfig::new("status", DataType::Utf8, false),
6213 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6214 ],
6215 vec!["id".to_string()],
6216 vec![IndexSpec::new("status_idx", vec!["status".to_string()]).expect("valid")],
6217 )
6218 .expect("backfill schema");
6219
6220 let (progress_tx, mut progress_rx) = mpsc::unbounded_channel();
6221 let report = backfill_schema
6222 .backfill_added_indexes_with_options_and_progress(
6223 "orders",
6224 &[],
6225 IndexBackfillOptions {
6226 row_batch_size: 2,
6227 start_from_primary_key: None,
6228 },
6229 Some(&progress_tx),
6230 )
6231 .await
6232 .expect("backfill should succeed");
6233 drop(progress_tx);
6234
6235 let mut saw_started = false;
6236 let mut saw_completed = false;
6237 let mut progress_events = 0usize;
6238 while let Some(event) = progress_rx.recv().await {
6239 match event {
6240 IndexBackfillEvent::Started {
6241 table_name,
6242 indexes_backfilled,
6243 row_batch_size,
6244 ..
6245 } => {
6246 saw_started = true;
6247 assert_eq!(table_name, "orders");
6248 assert_eq!(indexes_backfilled, 1);
6249 assert_eq!(row_batch_size, 2);
6250 }
6251 IndexBackfillEvent::Progress {
6252 scanned_rows,
6253 index_entries_written,
6254 ..
6255 } => {
6256 progress_events += 1;
6257 assert!(scanned_rows >= 1);
6258 assert_eq!(scanned_rows, index_entries_written);
6259 }
6260 IndexBackfillEvent::Completed {
6261 report: completed_report,
6262 } => {
6263 saw_completed = true;
6264 assert_eq!(completed_report, report);
6265 }
6266 }
6267 }
6268 assert!(saw_started);
6269 assert!(saw_completed);
6270 assert!(progress_events >= 1);
6271
6272 let _ = shutdown_tx.send(());
6273 }
6274
6275 #[tokio::test]
6276 async fn backfill_added_indexes_can_resume_from_primary_key() {
6277 let state = MockState {
6278 kv: Arc::new(Mutex::new(BTreeMap::new())),
6279 range_calls: Arc::new(AtomicUsize::new(0)),
6280 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
6281 sequence_number: Arc::new(AtomicU64::new(0)),
6282 };
6283 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
6284 let client = StoreClient::new(&base_url);
6285
6286 let seed_schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
6287 .table(
6288 "orders",
6289 vec![
6290 TableColumnConfig::new("id", DataType::Int64, false),
6291 TableColumnConfig::new("status", DataType::Utf8, false),
6292 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6293 ],
6294 vec!["id".to_string()],
6295 vec![],
6296 )
6297 .expect("seed schema");
6298 let mut writer = seed_schema.batch_writer();
6299 for i in 0..6i64 {
6300 writer
6301 .insert(
6302 "orders",
6303 vec![
6304 CellValue::Int64(i),
6305 CellValue::Utf8("open".to_string()),
6306 CellValue::Int64(i * 10),
6307 ],
6308 )
6309 .expect("seed row");
6310 }
6311 writer.flush().await.expect("seed flush");
6312
6313 let backfill_schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
6314 .table(
6315 "orders",
6316 vec![
6317 TableColumnConfig::new("id", DataType::Int64, false),
6318 TableColumnConfig::new("status", DataType::Utf8, false),
6319 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6320 ],
6321 vec!["id".to_string()],
6322 vec![IndexSpec::new("status_idx", vec!["status".to_string()]).expect("valid")],
6323 )
6324 .expect("backfill schema");
6325
6326 let config = KvTableConfig::new(
6327 0,
6328 vec![
6329 TableColumnConfig::new("id", DataType::Int64, false),
6330 TableColumnConfig::new("status", DataType::Utf8, false),
6331 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6332 ],
6333 vec!["id".to_string()],
6334 vec![],
6335 )
6336 .expect("valid config");
6337 let model = TableModel::from_config(&config).expect("model");
6338 let resume_value = CellValue::Int64(3);
6339 let resume_key =
6340 encode_primary_key(model.table_prefix, &[&resume_value], &model).expect("resume key");
6341
6342 let report = backfill_schema
6343 .backfill_added_indexes_with_options(
6344 "orders",
6345 &[],
6346 IndexBackfillOptions {
6347 row_batch_size: 2,
6348 start_from_primary_key: Some(resume_key.clone()),
6349 },
6350 )
6351 .await
6352 .expect("resume backfill should succeed");
6353 assert_eq!(report.scanned_rows, 3);
6354 assert_eq!(report.index_entries_written, 3);
6355
6356 {
6357 let guard = state.kv.lock().expect("kv mutex poisoned");
6358 let index_rows = guard
6359 .keys()
6360 .filter(|key| matches_secondary_index_key(0, 1, key))
6361 .count();
6362 assert_eq!(index_rows, 3);
6363 }
6364
6365 let resume_payload = model
6366 .primary_key_prefix
6367 .strip(&resume_key)
6368 .expect("resume payload");
6369 let wrong_prefix = secondary_index_prefix(model.table_prefix, 1)
6370 .expect("secondary prefix")
6371 .encode(&resume_payload)
6372 .expect("wrong prefix key");
6373 let err = backfill_schema
6374 .backfill_added_indexes_with_options(
6375 "orders",
6376 &[],
6377 IndexBackfillOptions {
6378 row_batch_size: 2,
6379 start_from_primary_key: Some(wrong_prefix),
6380 },
6381 )
6382 .await
6383 .expect_err("wrong key prefix must be rejected");
6384 assert!(err.to_string().contains("primary-key prefix"));
6385
6386 let _ = shutdown_tx.send(());
6387 }
6388
6389 #[tokio::test]
6390 async fn covering_index_scan_fails_closed_when_covering_payload_missing() {
6391 let state = MockState {
6392 kv: Arc::new(Mutex::new(BTreeMap::new())),
6393 range_calls: Arc::new(AtomicUsize::new(0)),
6394 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
6395 sequence_number: Arc::new(AtomicU64::new(0)),
6396 };
6397 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
6398 let client = StoreClient::new(&base_url);
6399
6400 let schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
6401 .table(
6402 "orders",
6403 vec![
6404 TableColumnConfig::new("id", DataType::Int64, false),
6405 TableColumnConfig::new("status", DataType::Utf8, false),
6406 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6407 ],
6408 vec!["id".to_string()],
6409 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
6410 .expect("valid")
6411 .with_cover_columns(vec!["amount_cents".to_string()])],
6412 )
6413 .expect("schema");
6414 let mut writer = schema.batch_writer();
6415 for id in 0..4i64 {
6416 writer
6417 .insert(
6418 "orders",
6419 vec![
6420 CellValue::Int64(id),
6421 CellValue::Utf8("open".to_string()),
6422 CellValue::Int64(id * 10),
6423 ],
6424 )
6425 .expect("row");
6426 }
6427 writer.flush().await.expect("flush");
6428
6429 {
6430 let mut guard = state.kv.lock().expect("kv mutex poisoned");
6431 let key = guard
6432 .keys()
6433 .find(|key| matches_secondary_index_key(0, 1, key))
6434 .expect("index row should exist")
6435 .clone();
6436 guard.insert(key, Bytes::new());
6437 }
6438
6439 let ctx = SessionContext::new();
6440 schema.register_all(&ctx).expect("register");
6441 let df = ctx
6442 .sql("SELECT amount_cents FROM orders WHERE status = 'open'")
6443 .await
6444 .expect("query should plan");
6445 let err = df
6446 .collect()
6447 .await
6448 .expect_err("missing covering payload must fail closed");
6449 assert!(err
6450 .to_string()
6451 .contains("secondary index entry missing covering payload"));
6452
6453 let _ = shutdown_tx.send(());
6454 }
6455
6456 #[tokio::test]
6457 async fn covering_index_scan_fails_closed_when_covering_payload_is_corrupt() {
6458 let state = MockState {
6459 kv: Arc::new(Mutex::new(BTreeMap::new())),
6460 range_calls: Arc::new(AtomicUsize::new(0)),
6461 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
6462 sequence_number: Arc::new(AtomicU64::new(0)),
6463 };
6464 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
6465 let client = StoreClient::new(&base_url);
6466
6467 let schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
6468 .table(
6469 "orders",
6470 vec![
6471 TableColumnConfig::new("id", DataType::Int64, false),
6472 TableColumnConfig::new("status", DataType::Utf8, false),
6473 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6474 ],
6475 vec!["id".to_string()],
6476 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
6477 .expect("valid")
6478 .with_cover_columns(vec!["amount_cents".to_string()])],
6479 )
6480 .expect("schema");
6481 let mut writer = schema.batch_writer();
6482 for id in 0..4i64 {
6483 writer
6484 .insert(
6485 "orders",
6486 vec![
6487 CellValue::Int64(id),
6488 CellValue::Utf8("open".to_string()),
6489 CellValue::Int64(id * 10),
6490 ],
6491 )
6492 .expect("row");
6493 }
6494 writer.flush().await.expect("flush");
6495
6496 {
6497 let mut guard = state.kv.lock().expect("kv mutex poisoned");
6498 let key = guard
6499 .keys()
6500 .find(|key| matches_secondary_index_key(0, 1, key))
6501 .expect("index row should exist")
6502 .clone();
6503 guard.insert(key, Bytes::from_static(b"not-codec"));
6504 }
6505
6506 let ctx = SessionContext::new();
6507 schema.register_all(&ctx).expect("register");
6508 let df = ctx
6509 .sql("SELECT amount_cents FROM orders WHERE status = 'open'")
6510 .await
6511 .expect("query should plan");
6512 let err = df
6513 .collect()
6514 .await
6515 .expect_err("corrupt covering payload must fail closed");
6516 assert!(err.to_string().contains("invalid covering index payload"));
6517
6518 let _ = shutdown_tx.send(());
6519 }
6520
6521 #[tokio::test]
6522 async fn non_covering_index_uses_point_lookup_instead_of_full_scan() {
6523 let state = MockState {
6524 kv: Arc::new(Mutex::new(BTreeMap::new())),
6525 range_calls: Arc::new(AtomicUsize::new(0)),
6526 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
6527 sequence_number: Arc::new(AtomicU64::new(0)),
6528 };
6529 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
6530 let client = StoreClient::new(&base_url);
6531
6532 let schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
6533 .table(
6534 "orders",
6535 vec![
6536 TableColumnConfig::new("id", DataType::Int64, false),
6537 TableColumnConfig::new("status", DataType::Utf8, false),
6538 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6539 TableColumnConfig::new("notes", DataType::Utf8, true),
6540 ],
6541 vec!["id".to_string()],
6542 vec![IndexSpec::new("status_idx", vec!["status".to_string()]).expect("valid")],
6543 )
6544 .expect("schema");
6545 let mut writer = schema.batch_writer();
6546 writer
6547 .insert(
6548 "orders",
6549 vec![
6550 CellValue::Int64(1),
6551 CellValue::Utf8("open".to_string()),
6552 CellValue::Int64(100),
6553 CellValue::Utf8("first".to_string()),
6554 ],
6555 )
6556 .expect("row");
6557 writer
6558 .insert(
6559 "orders",
6560 vec![
6561 CellValue::Int64(2),
6562 CellValue::Utf8("closed".to_string()),
6563 CellValue::Int64(200),
6564 CellValue::Utf8("second".to_string()),
6565 ],
6566 )
6567 .expect("row");
6568 writer
6569 .insert(
6570 "orders",
6571 vec![
6572 CellValue::Int64(3),
6573 CellValue::Utf8("open".to_string()),
6574 CellValue::Int64(300),
6575 CellValue::Utf8("third".to_string()),
6576 ],
6577 )
6578 .expect("row");
6579 writer.flush().await.expect("flush");
6580
6581 let ctx = SessionContext::new();
6582 schema.register_all(&ctx).expect("register");
6583
6584 let df = ctx
6585 .sql("SELECT id, notes FROM orders WHERE status = 'open' ORDER BY id")
6586 .await
6587 .expect("plan");
6588 let batches = df.collect().await.expect("non-covering index lookup");
6589 let ids: Vec<i64> = batches
6590 .iter()
6591 .flat_map(|b| {
6592 b.column(0)
6593 .as_any()
6594 .downcast_ref::<datafusion::arrow::array::Int64Array>()
6595 .unwrap()
6596 .iter()
6597 .map(|v| v.unwrap())
6598 })
6599 .collect();
6600 let notes: Vec<String> = batches
6601 .iter()
6602 .flat_map(|b| {
6603 b.column(1)
6604 .as_any()
6605 .downcast_ref::<datafusion::arrow::array::StringArray>()
6606 .unwrap()
6607 .iter()
6608 .map(|v| v.unwrap().to_string())
6609 })
6610 .collect();
6611 assert_eq!(ids, vec![1, 3]);
6612 assert_eq!(notes, vec!["first", "third"]);
6613
6614 let _ = shutdown_tx.send(());
6615 }
6616
6617 #[tokio::test]
6618 async fn backfill_resume_cursor_can_continue_without_skips_or_duplicates() {
6619 let state = MockState {
6620 kv: Arc::new(Mutex::new(BTreeMap::new())),
6621 range_calls: Arc::new(AtomicUsize::new(0)),
6622 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
6623 sequence_number: Arc::new(AtomicU64::new(0)),
6624 };
6625 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
6626 let client = StoreClient::new(&base_url);
6627
6628 let seed_schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
6629 .table(
6630 "orders",
6631 vec![
6632 TableColumnConfig::new("id", DataType::Int64, false),
6633 TableColumnConfig::new("status", DataType::Utf8, false),
6634 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6635 ],
6636 vec!["id".to_string()],
6637 vec![],
6638 )
6639 .expect("seed schema");
6640 let mut writer = seed_schema.batch_writer();
6641 for i in 0..8i64 {
6642 writer
6643 .insert(
6644 "orders",
6645 vec![
6646 CellValue::Int64(i),
6647 CellValue::Utf8("open".to_string()),
6648 CellValue::Int64(i * 10),
6649 ],
6650 )
6651 .expect("seed row");
6652 }
6653 writer.flush().await.expect("seed flush");
6654
6655 let backfill_schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
6656 .table(
6657 "orders",
6658 vec![
6659 TableColumnConfig::new("id", DataType::Int64, false),
6660 TableColumnConfig::new("status", DataType::Utf8, false),
6661 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6662 ],
6663 vec!["id".to_string()],
6664 vec![IndexSpec::new("status_idx", vec!["status".to_string()]).expect("valid")],
6665 )
6666 .expect("backfill schema");
6667
6668 let task_schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
6669 .table(
6670 "orders",
6671 vec![
6672 TableColumnConfig::new("id", DataType::Int64, false),
6673 TableColumnConfig::new("status", DataType::Utf8, false),
6674 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6675 ],
6676 vec!["id".to_string()],
6677 vec![IndexSpec::new("status_idx", vec!["status".to_string()]).expect("valid")],
6678 )
6679 .expect("task schema");
6680 let (progress_tx, mut progress_rx) = mpsc::unbounded_channel();
6681 let handle = tokio::spawn(async move {
6682 task_schema
6683 .backfill_added_indexes_with_options_and_progress(
6684 "orders",
6685 &[],
6686 IndexBackfillOptions {
6687 row_batch_size: 2,
6688 start_from_primary_key: None,
6689 },
6690 Some(&progress_tx),
6691 )
6692 .await
6693 });
6694
6695 let mut resume_cursor = None;
6696 while let Some(event) = progress_rx.recv().await {
6697 if let IndexBackfillEvent::Progress { next_cursor, .. } = event {
6698 resume_cursor = next_cursor;
6699 break;
6700 }
6701 }
6702 handle.abort();
6703 let resume_cursor =
6704 resume_cursor.expect("first progress event should provide resume cursor");
6705
6706 let report = backfill_schema
6707 .backfill_added_indexes_with_options(
6708 "orders",
6709 &[],
6710 IndexBackfillOptions {
6711 row_batch_size: 2,
6712 start_from_primary_key: Some(resume_cursor),
6713 },
6714 )
6715 .await
6716 .expect("resume backfill should succeed");
6717 assert_eq!(report.scanned_rows, 6);
6718
6719 let guard = state.kv.lock().expect("kv mutex poisoned");
6720 let base_rows = guard
6721 .keys()
6722 .filter(|key| matches_primary_key(0, key))
6723 .count();
6724 let index_rows = guard
6725 .keys()
6726 .filter(|key| matches_secondary_index_key(0, 1, key))
6727 .count();
6728 assert_eq!(base_rows, 8);
6729 assert_eq!(
6730 index_rows, 8,
6731 "resume should backfill each row exactly once"
6732 );
6733
6734 let _ = shutdown_tx.send(());
6735 }
6736
6737 #[tokio::test]
6738 async fn concurrent_writes_during_backfill_preserve_index_correctness() {
6739 let state = MockState {
6740 kv: Arc::new(Mutex::new(BTreeMap::new())),
6741 range_calls: Arc::new(AtomicUsize::new(0)),
6742 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
6743 sequence_number: Arc::new(AtomicU64::new(0)),
6744 };
6745 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
6746 let client = StoreClient::new(&base_url);
6747
6748 let seed_schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
6749 .table(
6750 "orders",
6751 vec![
6752 TableColumnConfig::new("id", DataType::Int64, false),
6753 TableColumnConfig::new("status", DataType::Utf8, false),
6754 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6755 ],
6756 vec!["id".to_string()],
6757 vec![],
6758 )
6759 .expect("seed schema");
6760 let mut seed_writer = seed_schema.batch_writer();
6761 for i in 0..40i64 {
6762 seed_writer
6763 .insert(
6764 "orders",
6765 vec![
6766 CellValue::Int64(i),
6767 CellValue::Utf8(if i % 2 == 0 { "open" } else { "closed" }.to_string()),
6768 CellValue::Int64(i * 10),
6769 ],
6770 )
6771 .expect("seed row");
6772 }
6773 seed_writer.flush().await.expect("seed flush");
6774
6775 let backfill_schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
6776 .table(
6777 "orders",
6778 vec![
6779 TableColumnConfig::new("id", DataType::Int64, false),
6780 TableColumnConfig::new("status", DataType::Utf8, false),
6781 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6782 ],
6783 vec!["id".to_string()],
6784 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
6785 .expect("valid")
6786 .with_cover_columns(vec!["amount_cents".to_string()])],
6787 )
6788 .expect("backfill schema");
6789
6790 let task_schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
6791 .table(
6792 "orders",
6793 vec![
6794 TableColumnConfig::new("id", DataType::Int64, false),
6795 TableColumnConfig::new("status", DataType::Utf8, false),
6796 TableColumnConfig::new("amount_cents", DataType::Int64, false),
6797 ],
6798 vec!["id".to_string()],
6799 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
6800 .expect("valid")
6801 .with_cover_columns(vec!["amount_cents".to_string()])],
6802 )
6803 .expect("task schema");
6804 let (progress_tx, mut progress_rx) = mpsc::unbounded_channel();
6805 let handle = tokio::spawn(async move {
6806 task_schema
6807 .backfill_added_indexes_with_options_and_progress(
6808 "orders",
6809 &[],
6810 IndexBackfillOptions {
6811 row_batch_size: 5,
6812 start_from_primary_key: None,
6813 },
6814 Some(&progress_tx),
6815 )
6816 .await
6817 });
6818
6819 while let Some(event) = progress_rx.recv().await {
6820 if matches!(event, IndexBackfillEvent::Progress { .. }) {
6821 break;
6822 }
6823 }
6824
6825 let mut concurrent_writer = backfill_schema.batch_writer();
6826 for id in [100i64, 101i64] {
6827 concurrent_writer
6828 .insert(
6829 "orders",
6830 vec![
6831 CellValue::Int64(id),
6832 CellValue::Utf8("open".to_string()),
6833 CellValue::Int64(id * 10),
6834 ],
6835 )
6836 .expect("concurrent row");
6837 }
6838 concurrent_writer.flush().await.expect("concurrent flush");
6839
6840 let report = handle
6841 .await
6842 .expect("backfill task join")
6843 .expect("backfill result");
6844 assert!(
6845 report.scanned_rows >= 40,
6846 "backfill should at least scan the original historical rows"
6847 );
6848
6849 let guard = state.kv.lock().expect("kv mutex poisoned");
6850 let base_rows = guard
6851 .keys()
6852 .filter(|key| matches_primary_key(0, key))
6853 .count();
6854 let index_rows = guard
6855 .keys()
6856 .filter(|key| matches_secondary_index_key(0, 1, key))
6857 .count();
6858 assert_eq!(base_rows, 42);
6859 assert_eq!(
6860 index_rows, 42,
6861 "historical backfill plus concurrent indexed writes should leave one index row per base row"
6862 );
6863
6864 let _ = shutdown_tx.send(());
6865 }
6866
6867 #[derive(Clone)]
6868 struct DeferredChunkRangeHarness {
6869 first_chunk_sent: Arc<Notify>,
6870 release_second_chunk: Arc<Notify>,
6871 first_frame: ProtoRangeFrame,
6872 second_frame: ProtoRangeFrame,
6873 }
6874
6875 impl QueryService for DeferredChunkRangeHarness {
6876 async fn get(
6877 &self,
6878 _ctx: Context,
6879 _request: buffa::view::OwnedView<
6880 exoware_sdk::store::query::v1::GetRequestView<'static>,
6881 >,
6882 ) -> connectrpc::ServiceResult<ProtoGetResponse> {
6883 Err(ConnectError::unimplemented("test harness"))
6884 }
6885
6886 async fn get_many(
6887 &self,
6888 _ctx: Context,
6889 _request: buffa::view::OwnedView<
6890 exoware_sdk::store::query::v1::GetManyRequestView<'static>,
6891 >,
6892 ) -> connectrpc::ServiceResult<connectrpc::ServiceStream<ProtoGetManyFrame>> {
6893 Err(ConnectError::unimplemented("test harness"))
6894 }
6895
6896 async fn range(
6897 &self,
6898 _ctx: Context,
6899 _request: buffa::view::OwnedView<
6900 exoware_sdk::store::query::v1::RangeRequestView<'static>,
6901 >,
6902 ) -> connectrpc::ServiceResult<connectrpc::ServiceStream<ProtoRangeFrame>> {
6903 let first_chunk_sent = self.first_chunk_sent.clone();
6904 let release_second_chunk = self.release_second_chunk.clone();
6905 let first_frame = self.first_frame.clone();
6906 let second_frame = self.second_frame.clone();
6907 let stream = stream::try_unfold(0u8, move |state| {
6908 let first_chunk_sent = first_chunk_sent.clone();
6909 let release_second_chunk = release_second_chunk.clone();
6910 let first_frame = first_frame.clone();
6911 let second_frame = second_frame.clone();
6912 async move {
6913 match state {
6914 0 => {
6915 first_chunk_sent.notify_one();
6916 Ok(Some((first_frame, 1)))
6917 }
6918 1 => {
6919 release_second_chunk.notified().await;
6920 Ok(Some((second_frame, 2)))
6921 }
6922 _ => Ok(None),
6923 }
6924 }
6925 });
6926 Ok(connectrpc::Response::stream(stream))
6927 }
6928
6929 async fn reduce(
6930 &self,
6931 _ctx: Context,
6932 _request: buffa::view::OwnedView<
6933 exoware_sdk::store::query::v1::ReduceRequestView<'static>,
6934 >,
6935 ) -> connectrpc::ServiceResult<ProtoReduceResponse> {
6936 Err(ConnectError::unimplemented("test harness"))
6937 }
6938 }
6939
6940 #[derive(Clone)]
6941 struct ObservedLimitRangeHarness {
6942 release_second_chunk: Arc<Notify>,
6943 observed_limit: Arc<AtomicUsize>,
6944 observed_mode: Arc<AtomicUsize>,
6945 first_frame: ProtoRangeFrame,
6946 second_frame: ProtoRangeFrame,
6947 }
6948
6949 impl QueryService for ObservedLimitRangeHarness {
6950 async fn get(
6951 &self,
6952 _ctx: Context,
6953 _request: buffa::view::OwnedView<
6954 exoware_sdk::store::query::v1::GetRequestView<'static>,
6955 >,
6956 ) -> connectrpc::ServiceResult<ProtoGetResponse> {
6957 Err(ConnectError::unimplemented("test harness"))
6958 }
6959
6960 async fn get_many(
6961 &self,
6962 _ctx: Context,
6963 _request: buffa::view::OwnedView<
6964 exoware_sdk::store::query::v1::GetManyRequestView<'static>,
6965 >,
6966 ) -> connectrpc::ServiceResult<connectrpc::ServiceStream<ProtoGetManyFrame>> {
6967 Err(ConnectError::unimplemented("test harness"))
6968 }
6969
6970 async fn range(
6971 &self,
6972 _ctx: Context,
6973 request: buffa::view::OwnedView<
6974 exoware_sdk::store::query::v1::RangeRequestView<'static>,
6975 >,
6976 ) -> connectrpc::ServiceResult<connectrpc::ServiceStream<ProtoRangeFrame>> {
6977 let limit = request.limit.map(|v| v as usize).unwrap_or(usize::MAX);
6978 self.observed_limit.store(limit, AtomicOrdering::SeqCst);
6979 let mode = match parse_range_traversal_direction(request.mode) {
6980 Ok(RangeTraversalDirection::Forward) => 0,
6981 Ok(RangeTraversalDirection::Reverse) => 1,
6982 Err(_) => usize::MAX,
6983 };
6984 self.observed_mode.store(mode, AtomicOrdering::SeqCst);
6985 let release_second_chunk = self.release_second_chunk.clone();
6986 let first_frame = self.first_frame.clone();
6987 let second_frame = self.second_frame.clone();
6988 let stream = stream::try_unfold(0u8, move |state| {
6989 let release_second_chunk = release_second_chunk.clone();
6990 let first_frame = first_frame.clone();
6991 let second_frame = second_frame.clone();
6992 async move {
6993 match state {
6994 0 => Ok(Some((first_frame, 1))),
6995 1 => {
6996 if limit > 1 {
6997 release_second_chunk.notified().await;
6998 Ok(Some((second_frame, 2)))
6999 } else {
7000 Ok(None)
7001 }
7002 }
7003 _ => Ok(None),
7004 }
7005 }
7006 });
7007 Ok(connectrpc::Response::stream(stream))
7008 }
7009
7010 async fn reduce(
7011 &self,
7012 _ctx: Context,
7013 _request: buffa::view::OwnedView<
7014 exoware_sdk::store::query::v1::ReduceRequestView<'static>,
7015 >,
7016 ) -> connectrpc::ServiceResult<ProtoReduceResponse> {
7017 Err(ConnectError::unimplemented("test harness"))
7018 }
7019 }
7020
7021 #[derive(Clone)]
7022 struct ObservedLimitIndexRangeHarness {
7023 observed_limit: Arc<AtomicUsize>,
7024 entries_frame: ProtoRangeFrame,
7025 }
7026
7027 impl QueryService for ObservedLimitIndexRangeHarness {
7028 async fn get(
7029 &self,
7030 _ctx: Context,
7031 _request: buffa::view::OwnedView<
7032 exoware_sdk::store::query::v1::GetRequestView<'static>,
7033 >,
7034 ) -> connectrpc::ServiceResult<ProtoGetResponse> {
7035 Err(ConnectError::unimplemented("test harness"))
7036 }
7037
7038 async fn get_many(
7039 &self,
7040 _ctx: Context,
7041 _request: buffa::view::OwnedView<
7042 exoware_sdk::store::query::v1::GetManyRequestView<'static>,
7043 >,
7044 ) -> connectrpc::ServiceResult<connectrpc::ServiceStream<ProtoGetManyFrame>> {
7045 Err(ConnectError::unimplemented("test harness"))
7046 }
7047
7048 async fn range(
7049 &self,
7050 _ctx: Context,
7051 request: buffa::view::OwnedView<
7052 exoware_sdk::store::query::v1::RangeRequestView<'static>,
7053 >,
7054 ) -> connectrpc::ServiceResult<connectrpc::ServiceStream<ProtoRangeFrame>> {
7055 let limit = request
7056 .limit
7057 .map(|v| {
7058 if v == u32::MAX {
7059 usize::MAX
7060 } else {
7061 v as usize
7062 }
7063 })
7064 .unwrap_or(usize::MAX);
7065 self.observed_limit.store(limit, AtomicOrdering::SeqCst);
7066 let entries_frame = self.entries_frame.clone();
7067 Ok(connectrpc::Response::stream(stream::iter(vec![Ok(
7068 entries_frame,
7069 )])))
7070 }
7071
7072 async fn reduce(
7073 &self,
7074 _ctx: Context,
7075 _request: buffa::view::OwnedView<
7076 exoware_sdk::store::query::v1::ReduceRequestView<'static>,
7077 >,
7078 ) -> connectrpc::ServiceResult<ProtoReduceResponse> {
7079 Err(ConnectError::unimplemented("test harness"))
7080 }
7081 }
7082
7083 #[tokio::test]
7084 async fn kv_scan_streaming_range_reads_emit_first_batch_before_full_range_completes() {
7085 let model = Arc::new(simple_int64_model(0));
7086 let first_chunk_sent = Arc::new(Notify::new());
7087 let release_second_chunk = Arc::new(Notify::new());
7088
7089 let encoded_row = (StoredRow { values: vec![None] }).encode().to_vec();
7090
7091 let first_results = {
7092 let mut results = Vec::with_capacity(BATCH_FLUSH_ROWS);
7093 for id in 0..BATCH_FLUSH_ROWS {
7094 let key =
7095 encode_primary_key(model.table_prefix, &[&CellValue::Int64(id as i64)], &model)
7096 .expect("primary key");
7097 results.push((key, encoded_row.clone()));
7098 }
7099 results
7100 };
7101 let first_frame = proto_range_entries_frame(first_results);
7102
7103 let second_results = {
7104 let key = encode_primary_key(
7105 model.table_prefix,
7106 &[&CellValue::Int64(BATCH_FLUSH_ROWS as i64)],
7107 &model,
7108 )
7109 .expect("primary key");
7110 vec![(key, encoded_row)]
7111 };
7112 let second_frame = proto_range_entries_frame(second_results);
7113
7114 let harness = DeferredChunkRangeHarness {
7115 first_chunk_sent: first_chunk_sent.clone(),
7116 release_second_chunk: release_second_chunk.clone(),
7117 first_frame,
7118 second_frame,
7119 };
7120 let connect = ConnectRpcService::new(QueryServiceServer::new(harness))
7121 .with_compression(connect_compression_registry());
7122 let app = Router::new().fallback_service(connect);
7123
7124 let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
7125 .await
7126 .expect("bind test listener");
7127 let url = format!("http://{}", listener.local_addr().expect("listener addr"));
7128 tokio::spawn(async move {
7129 axum::serve(listener, app).await.expect("serve test app");
7130 });
7131
7132 let client = StoreClient::new(&url);
7133 let scan = KvScanExec::new(
7134 PrefixedStoreClient::empty(client),
7135 model.clone(),
7136 Arc::new(Vec::new()),
7137 QueryPredicate::default(),
7138 None,
7139 model.schema.clone(),
7140 None,
7141 );
7142
7143 let session_ctx = SessionContext::new();
7144 let mut stream = scan
7145 .execute(0, session_ctx.task_ctx())
7146 .expect("scan execute should start");
7147
7148 tokio::time::timeout(Duration::from_secs(1), first_chunk_sent.notified())
7149 .await
7150 .expect("server should send first range frame");
7151 let first_batch = tokio::time::timeout(Duration::from_millis(200), stream.try_next())
7152 .await
7153 .expect("first record batch should arrive before the second stream chunk is released")
7154 .expect("stream poll should succeed")
7155 .expect("expected first record batch");
7156 assert_eq!(first_batch.num_rows(), BATCH_FLUSH_ROWS);
7157
7158 release_second_chunk.notify_one();
7159
7160 let second_batch = stream
7161 .try_next()
7162 .await
7163 .expect("second poll should succeed")
7164 .expect("expected second record batch");
7165 assert_eq!(second_batch.num_rows(), 1);
7166 assert!(
7167 stream
7168 .try_next()
7169 .await
7170 .expect("stream completion poll")
7171 .is_none(),
7172 "stream should finish after the second batch"
7173 );
7174 }
7175
7176 #[tokio::test]
7177 async fn kv_scan_sql_limit_is_pushed_upstream_on_exact_streaming_scan() {
7178 let release_second_chunk = Arc::new(Notify::new());
7179 let observed_limit = Arc::new(AtomicUsize::new(0));
7180 let observed_mode = Arc::new(AtomicUsize::new(usize::MAX));
7181 let model = simple_int64_model(0);
7182
7183 let encoded_row = (StoredRow { values: vec![None] }).encode().to_vec();
7184
7185 let first_key = encode_primary_key(model.table_prefix, &[&CellValue::Int64(1)], &model)
7186 .expect("first primary key");
7187 let second_key = encode_primary_key(model.table_prefix, &[&CellValue::Int64(2)], &model)
7188 .expect("second primary key");
7189
7190 let first_frame = proto_range_entries_frame(vec![(first_key, encoded_row.clone())]);
7191 let second_frame = proto_range_entries_frame(vec![(second_key, encoded_row)]);
7192
7193 let harness = ObservedLimitRangeHarness {
7194 release_second_chunk: release_second_chunk.clone(),
7195 observed_limit: observed_limit.clone(),
7196 observed_mode,
7197 first_frame,
7198 second_frame,
7199 };
7200 let connect = ConnectRpcService::new(QueryServiceServer::new(harness))
7201 .with_compression(connect_compression_registry());
7202 let app = Router::new().fallback_service(connect);
7203
7204 let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
7205 .await
7206 .expect("bind test listener");
7207 let url = format!("http://{}", listener.local_addr().expect("listener addr"));
7208 tokio::spawn(async move {
7209 axum::serve(listener, app).await.expect("serve test app");
7210 });
7211
7212 let client = StoreClient::new(&url);
7213 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
7214 .table(
7215 "items",
7216 vec![TableColumnConfig::new("id", DataType::Int64, false)],
7217 vec!["id".to_string()],
7218 vec![],
7219 )
7220 .expect("schema");
7221 let ctx = SessionContext::new();
7222 schema.register_all(&ctx).expect("register");
7223 let batches = tokio::time::timeout(Duration::from_millis(200), async {
7224 ctx.sql("SELECT id FROM items LIMIT 1")
7225 .await
7226 .expect("query")
7227 .collect()
7228 .await
7229 .expect("collect")
7230 })
7231 .await
7232 .expect("query with LIMIT 1 should finish without waiting for a delayed second chunk");
7233
7234 assert_eq!(
7235 batches.iter().map(|batch| batch.num_rows()).sum::<usize>(),
7236 1
7237 );
7238 assert_eq!(
7239 observed_limit.load(AtomicOrdering::SeqCst),
7240 1,
7241 "exact streaming scan should push SQL LIMIT upstream"
7242 );
7243 release_second_chunk.notify_one();
7244 }
7245
7246 #[test]
7247 fn kv_scan_physical_limit_pushdown_sets_leaf_fetch() {
7248 let model = Arc::new(simple_int64_model(0));
7249 let scan = Arc::new(KvScanExec::new(
7250 PrefixedStoreClient::empty(StoreClient::new("http://127.0.0.1:0")),
7251 model.clone(),
7252 Arc::new(Vec::new()),
7253 QueryPredicate::default(),
7254 None,
7255 model.schema.clone(),
7256 None,
7257 ));
7258 let limited: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(scan, 0, Some(1)));
7259
7260 let optimized = LimitPushdown::new()
7261 .optimize(limited, &ConfigOptions::new())
7262 .expect("limit pushdown");
7263
7264 let scan = optimized
7265 .as_any()
7266 .downcast_ref::<KvScanExec>()
7267 .expect("physical limit should be converted to a scan fetch");
7268 assert_eq!(scan.fetch(), Some(1));
7269 }
7270
7271 #[tokio::test]
7272 async fn kv_scan_activity_desc_order_pushes_reverse_range_limit() {
7273 let release_second_chunk = Arc::new(Notify::new());
7274 let observed_limit = Arc::new(AtomicUsize::new(0));
7275 let observed_mode = Arc::new(AtomicUsize::new(usize::MAX));
7276 let config = KvTableConfig::new(
7277 0,
7278 vec![
7279 TableColumnConfig::new("account", DataType::UInt64, false),
7280 TableColumnConfig::new("height", DataType::UInt64, false),
7281 TableColumnConfig::new("index", DataType::UInt64, false),
7282 TableColumnConfig::new("role", DataType::UInt64, false),
7283 ],
7284 vec![
7285 "account".to_string(),
7286 "height".to_string(),
7287 "index".to_string(),
7288 "role".to_string(),
7289 ],
7290 vec![],
7291 )
7292 .expect("config");
7293 let model = TableModel::from_config(&config).expect("model");
7294 let encoded_row = (StoredRow {
7295 values: vec![None, None, None, None],
7296 })
7297 .encode()
7298 .to_vec();
7299
7300 let receiver_key = encode_primary_key(
7301 model.table_prefix,
7302 &[
7303 &CellValue::UInt64(7),
7304 &CellValue::UInt64(9),
7305 &CellValue::UInt64(4),
7306 &CellValue::UInt64(1),
7307 ],
7308 &model,
7309 )
7310 .expect("receiver primary key");
7311 let sender_key = encode_primary_key(
7312 model.table_prefix,
7313 &[
7314 &CellValue::UInt64(7),
7315 &CellValue::UInt64(9),
7316 &CellValue::UInt64(4),
7317 &CellValue::UInt64(0),
7318 ],
7319 &model,
7320 )
7321 .expect("sender primary key");
7322
7323 let first_frame = proto_range_entries_frame(vec![(receiver_key, encoded_row.clone())]);
7324 let second_frame = proto_range_entries_frame(vec![(sender_key, encoded_row)]);
7325
7326 let harness = ObservedLimitRangeHarness {
7327 release_second_chunk: release_second_chunk.clone(),
7328 observed_limit: observed_limit.clone(),
7329 observed_mode: observed_mode.clone(),
7330 first_frame,
7331 second_frame,
7332 };
7333 let connect = ConnectRpcService::new(QueryServiceServer::new(harness))
7334 .with_compression(connect_compression_registry());
7335 let app = Router::new().fallback_service(connect);
7336
7337 let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
7338 .await
7339 .expect("bind test listener");
7340 let url = format!("http://{}", listener.local_addr().expect("listener addr"));
7341 tokio::spawn(async move {
7342 axum::serve(listener, app).await.expect("serve test app");
7343 });
7344
7345 let client = StoreClient::new(&url);
7346 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
7347 .table(
7348 "tx_activity",
7349 vec![
7350 TableColumnConfig::new("account", DataType::UInt64, false),
7351 TableColumnConfig::new("height", DataType::UInt64, false),
7352 TableColumnConfig::new("index", DataType::UInt64, false),
7353 TableColumnConfig::new("role", DataType::UInt64, false),
7354 ],
7355 vec![
7356 "account".to_string(),
7357 "height".to_string(),
7358 "index".to_string(),
7359 "role".to_string(),
7360 ],
7361 vec![],
7362 )
7363 .expect("schema");
7364 let ctx = SessionContext::new();
7365 schema.register_all(&ctx).expect("register");
7366
7367 let batches = tokio::time::timeout(Duration::from_millis(200), async {
7368 ctx.sql(
7369 "SELECT height, index, role \
7370 FROM tx_activity \
7371 WHERE account = 7 \
7372 ORDER BY height DESC, index DESC, role DESC \
7373 LIMIT 1",
7374 )
7375 .await
7376 .expect("query")
7377 .collect()
7378 .await
7379 .expect("collect")
7380 })
7381 .await
7382 .expect("activity DESC LIMIT query should not wait for a delayed second chunk");
7383
7384 assert_eq!(
7385 batches.iter().map(|batch| batch.num_rows()).sum::<usize>(),
7386 1
7387 );
7388 assert_eq!(
7389 observed_mode.load(AtomicOrdering::SeqCst),
7390 1,
7391 "activity DESC primary-key order should use reverse range traversal"
7392 );
7393 assert_eq!(
7394 observed_limit.load(AtomicOrdering::SeqCst),
7395 1,
7396 "activity DESC LIMIT should push the top-K limit to the range request"
7397 );
7398 release_second_chunk.notify_one();
7399 }
7400
7401 #[tokio::test]
7402 async fn kv_scan_activity_mixed_order_does_not_push_range_limit() {
7403 let state = MockState {
7404 kv: Arc::new(Mutex::new(BTreeMap::new())),
7405 range_calls: Arc::new(AtomicUsize::new(0)),
7406 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
7407 sequence_number: Arc::new(AtomicU64::new(0)),
7408 };
7409 let (base_url, shutdown_tx) = spawn_mock_server(state).await;
7410 let client = StoreClient::new(&base_url);
7411
7412 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
7413 .table(
7414 "tx_activity",
7415 vec![
7416 TableColumnConfig::new("account", DataType::UInt64, false),
7417 TableColumnConfig::new("height", DataType::UInt64, false),
7418 TableColumnConfig::new("index", DataType::UInt64, false),
7419 TableColumnConfig::new("role", DataType::UInt64, false),
7420 ],
7421 vec![
7422 "account".to_string(),
7423 "height".to_string(),
7424 "index".to_string(),
7425 "role".to_string(),
7426 ],
7427 vec![],
7428 )
7429 .expect("schema");
7430 let ctx = SessionContext::new();
7431 schema.register_all(&ctx).expect("register");
7432
7433 let explain = physical_plan_text(
7434 &explain_plan_rows(
7435 &ctx,
7436 "SELECT height, index, role \
7437 FROM tx_activity \
7438 WHERE account = 7 \
7439 ORDER BY height DESC, index DESC, role ASC \
7440 LIMIT 1",
7441 )
7442 .await,
7443 );
7444 assert!(explain.contains("KvScanExec:"));
7445 assert!(
7446 explain.contains("KvScanExec: limit=None"),
7447 "mixed direction order cannot be a bounded key-order scan:\n{explain}"
7448 );
7449 assert!(
7450 !explain.contains("KvScanExec: limit=Some(1)"),
7451 "mixed direction order incorrectly pushed top-K into the scan:\n{explain}"
7452 );
7453
7454 let _ = shutdown_tx.send(());
7455 }
7456
7457 #[tokio::test]
7458 async fn kv_topk_sort_pushdown_stops_at_filter_exec() {
7459 let state = MockState {
7460 kv: Arc::new(Mutex::new(BTreeMap::new())),
7461 range_calls: Arc::new(AtomicUsize::new(0)),
7462 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
7463 sequence_number: Arc::new(AtomicU64::new(0)),
7464 };
7465 let (base_url, shutdown_tx) = spawn_mock_server(state).await;
7466 let client = StoreClient::new(&base_url);
7467
7468 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
7469 .table(
7470 "tx_activity",
7471 vec![
7472 TableColumnConfig::new("account", DataType::UInt64, false),
7473 TableColumnConfig::new("height", DataType::UInt64, false),
7474 TableColumnConfig::new("index", DataType::UInt64, false),
7475 TableColumnConfig::new("role", DataType::UInt64, false),
7476 ],
7477 vec![
7478 "account".to_string(),
7479 "height".to_string(),
7480 "index".to_string(),
7481 "role".to_string(),
7482 ],
7483 vec![],
7484 )
7485 .expect("schema");
7486 let ctx = SessionContext::new();
7487 schema.register_all(&ctx).expect("register");
7488
7489 let explain = physical_plan_text(
7490 &explain_plan_rows(
7491 &ctx,
7492 "SELECT height, index, role \
7493 FROM tx_activity \
7494 WHERE account = 7 AND height + 0 > 1 \
7495 ORDER BY height DESC, index DESC, role DESC \
7496 LIMIT 1",
7497 )
7498 .await,
7499 );
7500 assert!(
7501 explain.contains("FilterExec"),
7502 "unsupported filter should remain above scan:\n{explain}"
7503 );
7504 assert!(
7505 explain.contains("KvScanExec: limit=None"),
7506 "top-K fetch must not cross a row-dropping filter:\n{explain}"
7507 );
7508 assert!(
7509 !explain.contains("KvScanExec: limit=Some(1)"),
7510 "top-K fetch crossed a filter into the scan:\n{explain}"
7511 );
7512
7513 let _ = shutdown_tx.send(());
7514 }
7515
7516 #[tokio::test]
7517 async fn kv_topk_reverse_scan_keeps_reading_after_later_pk_filter_rejects_first_key() {
7518 let state = MockState {
7519 kv: Arc::new(Mutex::new(BTreeMap::new())),
7520 range_calls: Arc::new(AtomicUsize::new(0)),
7521 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
7522 sequence_number: Arc::new(AtomicU64::new(0)),
7523 };
7524 let (base_url, shutdown_tx) = spawn_mock_server(state).await;
7525 let client = StoreClient::new(&base_url);
7526
7527 let schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
7528 .table(
7529 "tx_activity",
7530 vec![
7531 TableColumnConfig::new("account", DataType::UInt64, false),
7532 TableColumnConfig::new("height", DataType::UInt64, false),
7533 TableColumnConfig::new("index", DataType::UInt64, false),
7534 TableColumnConfig::new("role", DataType::UInt64, false),
7535 ],
7536 vec![
7537 "account".to_string(),
7538 "height".to_string(),
7539 "index".to_string(),
7540 "role".to_string(),
7541 ],
7542 vec![],
7543 )
7544 .expect("schema");
7545 let mut writer = schema.batch_writer();
7546 writer
7547 .insert(
7548 "tx_activity",
7549 vec![
7550 CellValue::UInt64(8),
7551 CellValue::UInt64(0),
7552 CellValue::UInt64(0),
7553 CellValue::UInt64(0),
7554 ],
7555 )
7556 .expect("non-matching higher key");
7557 writer
7558 .insert(
7559 "tx_activity",
7560 vec![
7561 CellValue::UInt64(7),
7562 CellValue::UInt64(9),
7563 CellValue::UInt64(4),
7564 CellValue::UInt64(1),
7565 ],
7566 )
7567 .expect("matching lower key");
7568 writer.flush().await.expect("seed rows");
7569
7570 let ctx = SessionContext::new();
7571 schema.register_all(&ctx).expect("register");
7572 let sql = "SELECT account, height \
7573 FROM tx_activity \
7574 WHERE account >= 7 AND height = 9 \
7575 ORDER BY account DESC \
7576 LIMIT 1";
7577 let explain = physical_plan_text(&explain_plan_rows(&ctx, sql).await);
7578 assert!(
7579 explain.contains("KvScanExec: limit=Some(1), direction=Reverse"),
7580 "test must exercise reverse top-K scan pushdown:\n{explain}"
7581 );
7582
7583 let batches = ctx
7584 .sql(sql)
7585 .await
7586 .expect("query")
7587 .collect()
7588 .await
7589 .expect("collect");
7590 assert_eq!(
7591 batches.iter().map(|batch| batch.num_rows()).sum::<usize>(),
7592 1
7593 );
7594 let batch = batches
7595 .iter()
7596 .find(|batch| batch.num_rows() > 0)
7597 .expect("non-empty result batch");
7598 assert_eq!(
7599 ScalarValue::try_from_array(batch.column(0), 0).expect("account scalar"),
7600 ScalarValue::UInt64(Some(7))
7601 );
7602 assert_eq!(
7603 ScalarValue::try_from_array(batch.column(1), 0).expect("height scalar"),
7604 ScalarValue::UInt64(Some(9))
7605 );
7606
7607 let _ = shutdown_tx.send(());
7608 }
7609
7610 #[tokio::test]
7611 async fn kv_topk_reverse_scan_orders_multiple_primary_key_ranges() {
7612 let state = MockState {
7613 kv: Arc::new(Mutex::new(BTreeMap::new())),
7614 range_calls: Arc::new(AtomicUsize::new(0)),
7615 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
7616 sequence_number: Arc::new(AtomicU64::new(0)),
7617 };
7618 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
7619 let client = StoreClient::new(&base_url);
7620
7621 let schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
7622 .table(
7623 "events",
7624 vec![TableColumnConfig::new("id", DataType::Int64, false)],
7625 vec!["id".to_string()],
7626 vec![],
7627 )
7628 .expect("schema");
7629 let mut writer = schema.batch_writer();
7630 for id in [1, 3, 10, 20] {
7631 writer
7632 .insert("events", vec![CellValue::Int64(id)])
7633 .expect("row");
7634 }
7635 writer.flush().await.expect("seed rows");
7636
7637 let ctx = SessionContext::new();
7638 schema.register_all(&ctx).expect("register");
7639 let sql = "SELECT id \
7640 FROM events \
7641 WHERE id IN (3, 20, 1, 10) \
7642 ORDER BY id DESC \
7643 LIMIT 3";
7644 let explain = physical_plan_text(&explain_plan_rows(&ctx, sql).await);
7645 assert!(
7646 explain.contains("KvScanExec: limit=Some(3), direction=Reverse"),
7647 "test must exercise reverse top-K scan pushdown over IN ranges:\n{explain}"
7648 );
7649
7650 state.range_calls.store(0, AtomicOrdering::SeqCst);
7651 let batches = ctx
7652 .sql(sql)
7653 .await
7654 .expect("query")
7655 .collect()
7656 .await
7657 .expect("collect");
7658 assert_eq!(collect_i64_column(&batches, 0), vec![20, 10, 3]);
7659 assert_eq!(
7660 state.range_calls.load(AtomicOrdering::SeqCst),
7661 3,
7662 "reverse scan should stop after the three highest IN ranges"
7663 );
7664
7665 let _ = shutdown_tx.send(());
7666 }
7667
7668 #[tokio::test]
7669 async fn kv_topk_reverse_scan_applies_offset_over_multiple_primary_key_ranges() {
7670 let state = MockState {
7671 kv: Arc::new(Mutex::new(BTreeMap::new())),
7672 range_calls: Arc::new(AtomicUsize::new(0)),
7673 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
7674 sequence_number: Arc::new(AtomicU64::new(0)),
7675 };
7676 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
7677 let client = StoreClient::new(&base_url);
7678
7679 let schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
7680 .table(
7681 "events",
7682 vec![TableColumnConfig::new("id", DataType::Int64, false)],
7683 vec!["id".to_string()],
7684 vec![],
7685 )
7686 .expect("schema");
7687 let mut writer = schema.batch_writer();
7688 for id in [1, 3, 10, 20] {
7689 writer
7690 .insert("events", vec![CellValue::Int64(id)])
7691 .expect("row");
7692 }
7693 writer.flush().await.expect("seed rows");
7694
7695 let ctx = SessionContext::new();
7696 schema.register_all(&ctx).expect("register");
7697 let sql = "SELECT id \
7698 FROM events \
7699 WHERE id IN (3, 20, 1, 10) \
7700 ORDER BY id DESC \
7701 LIMIT 2 OFFSET 1";
7702 let explain = physical_plan_text(&explain_plan_rows(&ctx, sql).await);
7703 assert!(
7704 explain.contains("KvScanExec: limit=Some(3), direction=Reverse"),
7705 "OFFSET requires fetching offset + limit rows from the reverse scan:\n{explain}"
7706 );
7707
7708 state.range_calls.store(0, AtomicOrdering::SeqCst);
7709 let batches = ctx
7710 .sql(sql)
7711 .await
7712 .expect("query")
7713 .collect()
7714 .await
7715 .expect("collect");
7716 assert_eq!(collect_i64_column(&batches, 0), vec![10, 3]);
7717 assert_eq!(
7718 state.range_calls.load(AtomicOrdering::SeqCst),
7719 3,
7720 "reverse OFFSET scan should fetch only offset + limit IN ranges"
7721 );
7722
7723 let _ = shutdown_tx.send(());
7724 }
7725
7726 #[tokio::test]
7727 async fn kv_scan_impossible_utf8_primary_key_equality_skips_scan() {
7728 let state = MockState {
7729 kv: Arc::new(Mutex::new(BTreeMap::new())),
7730 range_calls: Arc::new(AtomicUsize::new(0)),
7731 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
7732 sequence_number: Arc::new(AtomicU64::new(0)),
7733 };
7734 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
7735 let client = StoreClient::new(&base_url);
7736
7737 let schema = KvSchema::new(PrefixedStoreClient::empty(client.clone()))
7738 .table(
7739 "items",
7740 vec![TableColumnConfig::new("id", DataType::Utf8, false)],
7741 vec!["id".to_string()],
7742 vec![],
7743 )
7744 .expect("schema");
7745 let mut writer = schema.batch_writer();
7746 writer
7747 .insert("items", vec![CellValue::Utf8("present".to_string())])
7748 .expect("row");
7749 writer.flush().await.expect("seed rows");
7750
7751 let ctx = SessionContext::new();
7752 schema.register_all(&ctx).expect("register");
7753
7754 let oversized = "a".repeat(
7755 primary_key_prefix(0)
7756 .expect("primary-key prefix")
7757 .max_payload_len(),
7758 );
7759 let sql = format!("SELECT id FROM items WHERE id = '{oversized}'");
7760
7761 state.range_calls.store(0, AtomicOrdering::SeqCst);
7762 let batches = ctx
7763 .sql(&sql)
7764 .await
7765 .expect("query")
7766 .collect()
7767 .await
7768 .expect("collect");
7769 assert_eq!(
7770 batches.iter().map(|batch| batch.num_rows()).sum::<usize>(),
7771 0,
7772 "an unencodable PK equality can never match any row"
7773 );
7774 assert_eq!(
7775 state.range_calls.load(AtomicOrdering::SeqCst),
7776 0,
7777 "an impossible PK equality must not issue a range scan"
7778 );
7779
7780 let _ = shutdown_tx.send(());
7781 }
7782
7783 #[tokio::test]
7784 async fn kv_scan_index_limit_does_not_push_upstream_when_seen_dedup_can_drop_entries() {
7785 let observed_limit = Arc::new(AtomicUsize::new(0));
7786 let config = KvTableConfig::new(
7787 0,
7788 vec![
7789 TableColumnConfig::new("id", DataType::Int64, false),
7790 TableColumnConfig::new("status", DataType::Utf8, false),
7791 TableColumnConfig::new("amount_cents", DataType::Int64, false),
7792 ],
7793 vec!["id".to_string()],
7794 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
7795 .expect("valid")
7796 .with_cover_columns(vec!["status".to_string(), "amount_cents".to_string()])],
7797 )
7798 .expect("config");
7799 let model = TableModel::from_config(&config).expect("model");
7800 let spec = model
7801 .resolve_index_specs(&config.index_specs)
7802 .expect("specs")
7803 .into_iter()
7804 .next()
7805 .expect("status index spec");
7806 let stale_row = KvRow {
7807 values: vec![
7808 CellValue::Int64(7),
7809 CellValue::Utf8("closed".to_string()),
7810 CellValue::Int64(10),
7811 ],
7812 };
7813 let current_row = KvRow {
7814 values: vec![
7815 CellValue::Int64(7),
7816 CellValue::Utf8("open".to_string()),
7817 CellValue::Int64(10),
7818 ],
7819 };
7820 let unique_row = KvRow {
7821 values: vec![
7822 CellValue::Int64(8),
7823 CellValue::Utf8("open".to_string()),
7824 CellValue::Int64(20),
7825 ],
7826 };
7827 let stale_key = encode_secondary_index_key(model.table_prefix, &spec, &model, &stale_row)
7828 .expect("stale index key");
7829 let current_key =
7830 encode_secondary_index_key(model.table_prefix, &spec, &model, ¤t_row)
7831 .expect("current index key");
7832 let unique_key = encode_secondary_index_key(model.table_prefix, &spec, &model, &unique_row)
7833 .expect("unique index key");
7834 let stale_payload =
7835 encode_secondary_index_value(&stale_row, &model, &spec).expect("stale payload");
7836 let current_payload =
7837 encode_secondary_index_value(¤t_row, &model, &spec).expect("current payload");
7838 let unique_payload =
7839 encode_secondary_index_value(&unique_row, &model, &spec).expect("unique payload");
7840
7841 let entries_frame = proto_range_entries_frame(vec![
7842 (stale_key, stale_payload),
7843 (current_key, current_payload),
7844 (unique_key, unique_payload),
7845 ]);
7846 let harness = ObservedLimitIndexRangeHarness {
7847 observed_limit: observed_limit.clone(),
7848 entries_frame,
7849 };
7850 let connect = ConnectRpcService::new(QueryServiceServer::new(harness))
7851 .with_compression(connect_compression_registry());
7852 let app = Router::new().fallback_service(connect);
7853
7854 let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
7855 .await
7856 .expect("bind test listener");
7857 let url = format!("http://{}", listener.local_addr().expect("listener addr"));
7858 tokio::spawn(async move {
7859 axum::serve(listener, app).await.expect("serve test app");
7860 });
7861
7862 let client = StoreClient::new(&url);
7863 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
7864 .table(
7865 "orders",
7866 vec![
7867 TableColumnConfig::new("id", DataType::Int64, false),
7868 TableColumnConfig::new("status", DataType::Utf8, false),
7869 TableColumnConfig::new("amount_cents", DataType::Int64, false),
7870 ],
7871 vec!["id".to_string()],
7872 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
7873 .expect("valid")
7874 .with_cover_columns(vec!["status".to_string(), "amount_cents".to_string()])],
7875 )
7876 .expect("schema");
7877 let ctx = SessionContext::new();
7878 schema.register_all(&ctx).expect("register");
7879
7880 let batches = ctx
7881 .sql(
7882 "SELECT id, amount_cents \
7883 FROM orders \
7884 WHERE status IN ('open', 'closed') \
7885 LIMIT 2",
7886 )
7887 .await
7888 .expect("query")
7889 .collect()
7890 .await
7891 .expect("collect");
7892
7893 assert_eq!(
7894 batches.iter().map(|batch| batch.num_rows()).sum::<usize>(),
7895 2
7896 );
7897 assert_eq!(
7898 observed_limit.load(AtomicOrdering::SeqCst),
7899 usize::MAX,
7900 "index streaming scans should not push SQL LIMIT upstream while seen-dedup can drop duplicate primary keys"
7901 );
7902 }
7903
7904 #[tokio::test]
7905 async fn zorder_covering_index_scan_filters_false_positive_morton_span_rows() {
7906 let state = MockState {
7907 kv: Arc::new(Mutex::new(BTreeMap::new())),
7908 range_calls: Arc::new(AtomicUsize::new(0)),
7909 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
7910 sequence_number: Arc::new(AtomicU64::new(0)),
7911 };
7912 let (base_url, shutdown_tx) = spawn_mock_server(state).await;
7913 let client = StoreClient::new(&base_url);
7914
7915 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
7916 .table(
7917 "points",
7918 vec![
7919 TableColumnConfig::new("x", DataType::Int64, false),
7920 TableColumnConfig::new("y", DataType::Int64, false),
7921 TableColumnConfig::new("id", DataType::Int64, false),
7922 TableColumnConfig::new("value", DataType::Int64, false),
7923 ],
7924 vec!["id".to_string()],
7925 vec![
7926 IndexSpec::z_order("xy_z", vec!["x".to_string(), "y".to_string()])
7927 .expect("valid")
7928 .with_cover_columns(vec!["value".to_string()]),
7929 ],
7930 )
7931 .expect("schema");
7932
7933 let mut writer = schema.batch_writer();
7934 for (x, y, id, value) in [
7935 (0, 2, 2, 20),
7936 (1, 1, 11, 110),
7937 (1, 2, 12, 120),
7938 (2, 1, 21, 210),
7939 (2, 2, 22, 220),
7940 (3, 0, 30, 300),
7941 ] {
7942 writer
7943 .insert(
7944 "points",
7945 vec![
7946 CellValue::Int64(x),
7947 CellValue::Int64(y),
7948 CellValue::Int64(id),
7949 CellValue::Int64(value),
7950 ],
7951 )
7952 .expect("row");
7953 }
7954 writer.flush().await.expect("flush");
7955
7956 let ctx = SessionContext::new();
7957 schema.register_all(&ctx).expect("register");
7958
7959 let batches = ctx
7960 .sql(
7961 "SELECT id, value FROM points \
7962 WHERE x >= 1 AND x <= 2 AND y >= 1 AND y <= 2 \
7963 ORDER BY id",
7964 )
7965 .await
7966 .expect("query")
7967 .collect()
7968 .await
7969 .expect("collect");
7970
7971 let mut rows = Vec::new();
7972 for batch in &batches {
7973 let ids = batch
7974 .column(0)
7975 .as_any()
7976 .downcast_ref::<Int64Array>()
7977 .expect("id int64");
7978 let values = batch
7979 .column(1)
7980 .as_any()
7981 .downcast_ref::<Int64Array>()
7982 .expect("value int64");
7983 for row_idx in 0..batch.num_rows() {
7984 rows.push((ids.value(row_idx), values.value(row_idx)));
7985 }
7986 }
7987 assert_eq!(rows, vec![(11, 110), (12, 120), (21, 210), (22, 220)]);
7988
7989 let _ = shutdown_tx.send(());
7990 }
7991
7992 #[tokio::test]
7993 async fn aggregate_pushdown_uses_range_reduce_for_supported_global_aggregates() {
7994 let state = MockState {
7995 kv: Arc::new(Mutex::new(BTreeMap::new())),
7996 range_calls: Arc::new(AtomicUsize::new(0)),
7997 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
7998 sequence_number: Arc::new(AtomicU64::new(0)),
7999 };
8000 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
8001 let client = StoreClient::new(&base_url);
8002
8003 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
8004 .table(
8005 "orders",
8006 vec![
8007 TableColumnConfig::new("id", DataType::Int64, false),
8008 TableColumnConfig::new("status", DataType::Utf8, false),
8009 TableColumnConfig::new("amount_cents", DataType::Int64, false),
8010 ],
8011 vec!["id".to_string()],
8012 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
8013 .expect("valid")
8014 .with_cover_columns(vec!["amount_cents".to_string()])],
8015 )
8016 .expect("schema");
8017
8018 let mut writer = schema.batch_writer();
8019 for (id, status, amount) in [
8020 (1, "open", 10),
8021 (2, "closed", 15),
8022 (3, "open", 30),
8023 (4, "closed", 40),
8024 ] {
8025 writer
8026 .insert(
8027 "orders",
8028 vec![
8029 CellValue::Int64(id),
8030 CellValue::Utf8(status.to_string()),
8031 CellValue::Int64(amount),
8032 ],
8033 )
8034 .expect("row");
8035 }
8036 writer.flush().await.expect("flush");
8037
8038 let ctx = SessionContext::new();
8039 schema.register_all(&ctx).expect("register");
8040
8041 state.range_calls.store(0, AtomicOrdering::SeqCst);
8042 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
8043
8044 let df = ctx
8045 .sql(
8046 "SELECT COUNT(*) AS row_count, SUM(amount_cents) AS total_cents, \
8047 AVG(amount_cents) AS avg_cents \
8048 FROM orders WHERE status = 'open'",
8049 )
8050 .await
8051 .expect("query");
8052 let batches = df.collect().await.expect("collect");
8053
8054 assert_eq!(batches.len(), 1);
8055 let batch = &batches[0];
8056 let row_count = ScalarValue::try_from_array(batch.column(0), 0).expect("row_count scalar");
8057 let total = batch
8058 .column(1)
8059 .as_any()
8060 .downcast_ref::<Int64Array>()
8061 .expect("sum int64")
8062 .value(0);
8063 let avg = batch
8064 .column(2)
8065 .as_any()
8066 .downcast_ref::<Float64Array>()
8067 .expect("avg float64")
8068 .value(0);
8069 assert!(matches!(
8070 row_count,
8071 ScalarValue::UInt64(Some(2)) | ScalarValue::Int64(Some(2))
8072 ));
8073 assert_eq!(total, 40);
8074 assert_eq!(avg, 20.0);
8075 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
8076 assert!(
8077 state.range_reduce_calls.load(AtomicOrdering::SeqCst) >= 1,
8078 "supported aggregate should use range reduction path"
8079 );
8080
8081 let _ = shutdown_tx.send(());
8082 }
8083
8084 #[tokio::test]
8086 async fn primary_key_inclusive_upper_bound_streaming_scan_uses_range() {
8087 let state = MockState {
8088 kv: Arc::new(Mutex::new(BTreeMap::new())),
8089 range_calls: Arc::new(AtomicUsize::new(0)),
8090 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
8091 sequence_number: Arc::new(AtomicU64::new(0)),
8092 };
8093 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
8094 let client = StoreClient::new(&base_url);
8095
8096 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
8097 .table(
8098 "inc_pk",
8099 vec![
8100 TableColumnConfig::new("id", DataType::Int64, false),
8101 TableColumnConfig::new("amount", DataType::Int64, false),
8102 ],
8103 vec!["id".to_string()],
8104 vec![],
8105 )
8106 .expect("schema");
8107
8108 let mut writer = schema.batch_writer();
8109 for id in 1i64..=5i64 {
8110 writer
8111 .insert(
8112 "inc_pk",
8113 vec![CellValue::Int64(id), CellValue::Int64(id * 100)],
8114 )
8115 .expect("row");
8116 }
8117 writer.flush().await.expect("flush");
8118
8119 let ctx = SessionContext::new();
8120 schema.register_all(&ctx).expect("register");
8121
8122 state.range_calls.store(0, AtomicOrdering::SeqCst);
8123 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
8124
8125 let batches = ctx
8126 .sql("SELECT id FROM inc_pk WHERE id <= 3 ORDER BY id")
8127 .await
8128 .expect("lte query")
8129 .collect()
8130 .await
8131 .expect("collect");
8132 let mut ids = Vec::new();
8133 for batch in &batches {
8134 let col = batch
8135 .column(0)
8136 .as_any()
8137 .downcast_ref::<Int64Array>()
8138 .expect("id");
8139 for i in 0..batch.num_rows() {
8140 ids.push(col.value(i));
8141 }
8142 }
8143 assert_eq!(ids, vec![1, 2, 3], "id <= 3 must include id 3");
8144 assert!(
8145 state.range_calls.load(AtomicOrdering::SeqCst) >= 1,
8146 "PK bounded scan should call range"
8147 );
8148 assert_eq!(
8149 state.range_reduce_calls.load(AtomicOrdering::SeqCst),
8150 0,
8151 "streaming scan must not use range_reduce"
8152 );
8153
8154 state.range_calls.store(0, AtomicOrdering::SeqCst);
8155 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
8156
8157 let batches = ctx
8158 .sql("SELECT id FROM inc_pk WHERE id BETWEEN 2 AND 4 ORDER BY id")
8159 .await
8160 .expect("between query")
8161 .collect()
8162 .await
8163 .expect("collect");
8164 ids.clear();
8165 for batch in &batches {
8166 let col = batch
8167 .column(0)
8168 .as_any()
8169 .downcast_ref::<Int64Array>()
8170 .expect("id");
8171 for i in 0..batch.num_rows() {
8172 ids.push(col.value(i));
8173 }
8174 }
8175 assert_eq!(ids, vec![2, 3, 4], "BETWEEN must include both endpoints");
8176 assert!(state.range_calls.load(AtomicOrdering::SeqCst) >= 1);
8177 assert_eq!(state.range_reduce_calls.load(AtomicOrdering::SeqCst), 0);
8178
8179 let _ = shutdown_tx.send(());
8180 }
8181
8182 #[tokio::test]
8183 async fn primary_key_inclusive_upper_bound_scalar_aggregates_use_range_reduce() {
8184 let state = MockState {
8185 kv: Arc::new(Mutex::new(BTreeMap::new())),
8186 range_calls: Arc::new(AtomicUsize::new(0)),
8187 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
8188 sequence_number: Arc::new(AtomicU64::new(0)),
8189 };
8190 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
8191 let client = StoreClient::new(&base_url);
8192
8193 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
8194 .table(
8195 "inc_pk",
8196 vec![
8197 TableColumnConfig::new("id", DataType::Int64, false),
8198 TableColumnConfig::new("amount", DataType::Int64, false),
8199 ],
8200 vec!["id".to_string()],
8201 vec![],
8202 )
8203 .expect("schema");
8204
8205 let mut writer = schema.batch_writer();
8206 for id in 1i64..=5i64 {
8207 writer
8208 .insert(
8209 "inc_pk",
8210 vec![CellValue::Int64(id), CellValue::Int64(id * 100)],
8211 )
8212 .expect("row");
8213 }
8214 writer.flush().await.expect("flush");
8215
8216 let ctx = SessionContext::new();
8217 schema.register_all(&ctx).expect("register");
8218
8219 state.range_calls.store(0, AtomicOrdering::SeqCst);
8220 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
8221
8222 let batches = ctx
8223 .sql("SELECT COUNT(*) AS c, SUM(amount) AS s FROM inc_pk WHERE id <= 3")
8224 .await
8225 .expect("lte agg")
8226 .collect()
8227 .await
8228 .expect("collect");
8229 assert_eq!(batches.len(), 1);
8230 let batch = &batches[0];
8231 let c = ScalarValue::try_from_array(batch.column(0), 0).expect("count");
8232 assert!(
8233 matches!(
8234 c,
8235 ScalarValue::UInt64(Some(3)) | ScalarValue::Int64(Some(3))
8236 ),
8237 "count should include id=3"
8238 );
8239 assert_eq!(
8240 batch
8241 .column(1)
8242 .as_any()
8243 .downcast_ref::<Int64Array>()
8244 .expect("sum")
8245 .value(0),
8246 100 + 200 + 300
8247 );
8248 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
8249 assert!(
8250 state.range_reduce_calls.load(AtomicOrdering::SeqCst) >= 1,
8251 "scalar aggregate on PK range should use range_reduce"
8252 );
8253
8254 state.range_calls.store(0, AtomicOrdering::SeqCst);
8255 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
8256
8257 let batches = ctx
8258 .sql("SELECT SUM(amount) AS s FROM inc_pk WHERE id BETWEEN 2 AND 4")
8259 .await
8260 .expect("between agg")
8261 .collect()
8262 .await
8263 .expect("collect");
8264 assert_eq!(batches.len(), 1);
8265 let batch = &batches[0];
8266 assert_eq!(
8267 batch
8268 .column(0)
8269 .as_any()
8270 .downcast_ref::<Int64Array>()
8271 .expect("sum")
8272 .value(0),
8273 200 + 300 + 400
8274 );
8275 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
8276 assert!(
8277 state.range_reduce_calls.load(AtomicOrdering::SeqCst) >= 1,
8278 "BETWEEN aggregate should use range_reduce"
8279 );
8280
8281 let _ = shutdown_tx.send(());
8282 }
8283
8284 #[tokio::test]
8285 async fn aggregate_pushdown_uses_zorder_index_with_worker_filter() {
8286 let state = MockState {
8287 kv: Arc::new(Mutex::new(BTreeMap::new())),
8288 range_calls: Arc::new(AtomicUsize::new(0)),
8289 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
8290 sequence_number: Arc::new(AtomicU64::new(0)),
8291 };
8292 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
8293 let client = StoreClient::new(&base_url);
8294
8295 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
8296 .table(
8297 "points",
8298 vec![
8299 TableColumnConfig::new("x", DataType::Int64, false),
8300 TableColumnConfig::new("y", DataType::Int64, false),
8301 TableColumnConfig::new("id", DataType::Int64, false),
8302 TableColumnConfig::new("value", DataType::Int64, false),
8303 ],
8304 vec!["id".to_string()],
8305 vec![
8306 IndexSpec::z_order("xy_z", vec!["x".to_string(), "y".to_string()])
8307 .expect("valid")
8308 .with_cover_columns(vec!["value".to_string()]),
8309 ],
8310 )
8311 .expect("schema");
8312
8313 let mut writer = schema.batch_writer();
8314 for (x, y, id, value) in [
8315 (0, 2, 2, 20),
8316 (1, 1, 11, 110),
8317 (1, 2, 12, 120),
8318 (2, 1, 21, 210),
8319 (2, 2, 22, 220),
8320 (3, 0, 30, 300),
8321 ] {
8322 writer
8323 .insert(
8324 "points",
8325 vec![
8326 CellValue::Int64(x),
8327 CellValue::Int64(y),
8328 CellValue::Int64(id),
8329 CellValue::Int64(value),
8330 ],
8331 )
8332 .expect("row");
8333 }
8334 writer.flush().await.expect("flush");
8335
8336 let ctx = SessionContext::new();
8337 schema.register_all(&ctx).expect("register");
8338
8339 state.range_calls.store(0, AtomicOrdering::SeqCst);
8340 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
8341
8342 let batches = ctx
8343 .sql(
8344 "SELECT COUNT(*) AS row_count, SUM(value) AS total_value \
8345 FROM points \
8346 WHERE x >= 1 AND x <= 2 AND y >= 1 AND y <= 2",
8347 )
8348 .await
8349 .expect("query")
8350 .collect()
8351 .await
8352 .expect("collect");
8353
8354 assert_eq!(batches.len(), 1);
8355 let batch = &batches[0];
8356 let row_count = ScalarValue::try_from_array(batch.column(0), 0).expect("row_count scalar");
8357 let total = batch
8358 .column(1)
8359 .as_any()
8360 .downcast_ref::<Int64Array>()
8361 .expect("sum int64")
8362 .value(0);
8363 assert!(matches!(
8364 row_count,
8365 ScalarValue::UInt64(Some(4)) | ScalarValue::Int64(Some(4))
8366 ));
8367 assert_eq!(total, 660);
8368 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
8369 assert!(
8370 state.range_reduce_calls.load(AtomicOrdering::SeqCst) >= 1,
8371 "z-order aggregate should use range reduction path"
8372 );
8373
8374 let _ = shutdown_tx.send(());
8375 }
8376
8377 #[tokio::test]
8378 async fn aggregate_pushdown_avg_merges_sum_and_count_across_multiple_ranges() {
8379 let state = MockState {
8380 kv: Arc::new(Mutex::new(BTreeMap::new())),
8381 range_calls: Arc::new(AtomicUsize::new(0)),
8382 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
8383 sequence_number: Arc::new(AtomicU64::new(0)),
8384 };
8385 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
8386 let client = StoreClient::new(&base_url);
8387
8388 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
8389 .table(
8390 "orders",
8391 vec![
8392 TableColumnConfig::new("id", DataType::Int64, false),
8393 TableColumnConfig::new("status", DataType::Utf8, false),
8394 TableColumnConfig::new("amount_cents", DataType::Int64, false),
8395 ],
8396 vec!["id".to_string()],
8397 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
8398 .expect("valid")
8399 .with_cover_columns(vec!["amount_cents".to_string()])],
8400 )
8401 .expect("schema");
8402
8403 let mut writer = schema.batch_writer();
8404 for (id, status, amount) in [
8405 (1, "open", 10),
8406 (2, "open", 20),
8407 (3, "closed", 100),
8408 (4, "pending", 1_000),
8409 ] {
8410 writer
8411 .insert(
8412 "orders",
8413 vec![
8414 CellValue::Int64(id),
8415 CellValue::Utf8(status.to_string()),
8416 CellValue::Int64(amount),
8417 ],
8418 )
8419 .expect("row");
8420 }
8421 writer.flush().await.expect("flush");
8422
8423 let ctx = SessionContext::new();
8424 schema.register_all(&ctx).expect("register");
8425
8426 state.range_calls.store(0, AtomicOrdering::SeqCst);
8427 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
8428
8429 let batches = ctx
8430 .sql(
8431 "SELECT AVG(amount_cents) AS avg_cents \
8432 FROM orders \
8433 WHERE status IN ('open', 'closed')",
8434 )
8435 .await
8436 .expect("query")
8437 .collect()
8438 .await
8439 .expect("collect");
8440
8441 assert_eq!(batches.len(), 1);
8442 let batch = &batches[0];
8443 let avg = batch
8444 .column(0)
8445 .as_any()
8446 .downcast_ref::<Float64Array>()
8447 .expect("avg float64")
8448 .value(0);
8449 let expected = 130.0 / 3.0;
8450 assert!(
8451 (avg - expected).abs() < 1e-12,
8452 "AVG should merge SUM+COUNT across unequal-count ranges: got {avg}, expected {expected}"
8453 );
8454 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
8455 assert_eq!(
8456 state.range_reduce_calls.load(AtomicOrdering::SeqCst),
8457 2,
8458 "status IN (...) should expand to two pushed reduction ranges"
8459 );
8460
8461 let _ = shutdown_tx.send(());
8462 }
8463
8464 #[tokio::test]
8465 async fn aggregate_pushdown_supports_filtered_global_aggregates() {
8466 let state = MockState {
8467 kv: Arc::new(Mutex::new(BTreeMap::new())),
8468 range_calls: Arc::new(AtomicUsize::new(0)),
8469 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
8470 sequence_number: Arc::new(AtomicU64::new(0)),
8471 };
8472 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
8473 let client = StoreClient::new(&base_url);
8474
8475 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
8476 .table(
8477 "orders",
8478 vec![
8479 TableColumnConfig::new("id", DataType::Int64, false),
8480 TableColumnConfig::new("status", DataType::Utf8, false),
8481 TableColumnConfig::new("amount_cents", DataType::Int64, false),
8482 ],
8483 vec!["id".to_string()],
8484 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
8485 .expect("valid")
8486 .with_cover_columns(vec!["amount_cents".to_string()])],
8487 )
8488 .expect("schema");
8489
8490 let mut writer = schema.batch_writer();
8491 for (id, status, amount) in [
8492 (1, "open", 10),
8493 (2, "closed", 15),
8494 (3, "open", 30),
8495 (4, "closed", 40),
8496 ] {
8497 writer
8498 .insert(
8499 "orders",
8500 vec![
8501 CellValue::Int64(id),
8502 CellValue::Utf8(status.to_string()),
8503 CellValue::Int64(amount),
8504 ],
8505 )
8506 .expect("row");
8507 }
8508 writer.flush().await.expect("flush");
8509
8510 let ctx = SessionContext::new();
8511 schema.register_all(&ctx).expect("register");
8512
8513 state.range_calls.store(0, AtomicOrdering::SeqCst);
8514 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
8515
8516 let query = "SELECT COUNT(*) FILTER (WHERE status = 'open') AS open_count, \
8517 COUNT(*) FILTER (WHERE status = 'closed') AS closed_count, \
8518 AVG(amount_cents) FILTER (WHERE status = 'closed') AS closed_avg \
8519 FROM orders";
8520 let batches = ctx
8521 .sql(query)
8522 .await
8523 .expect("query")
8524 .collect()
8525 .await
8526 .expect("collect");
8527
8528 assert_eq!(batches.len(), 1);
8529 let batch = &batches[0];
8530 assert_count_scalar(batch, 0, 0, 2);
8531 assert_count_scalar(batch, 1, 0, 2);
8532 let closed_avg = batch
8533 .column(2)
8534 .as_any()
8535 .downcast_ref::<Float64Array>()
8536 .expect("avg float64")
8537 .value(0);
8538 assert_eq!(closed_avg, 27.5);
8539 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
8540 assert!(
8541 state.range_reduce_calls.load(AtomicOrdering::SeqCst) >= 3,
8542 "filtered aggregate pushdown should use dedicated reduction jobs"
8543 );
8544
8545 let _ = shutdown_tx.send(());
8546 }
8547
8548 #[tokio::test]
8549 async fn aggregate_pushdown_supports_case_filtered_global_aggregates() {
8550 let state = MockState {
8551 kv: Arc::new(Mutex::new(BTreeMap::new())),
8552 range_calls: Arc::new(AtomicUsize::new(0)),
8553 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
8554 sequence_number: Arc::new(AtomicU64::new(0)),
8555 };
8556 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
8557 let client = StoreClient::new(&base_url);
8558
8559 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
8560 .table(
8561 "orders",
8562 vec![
8563 TableColumnConfig::new("id", DataType::Int64, false),
8564 TableColumnConfig::new("status", DataType::Utf8, false),
8565 TableColumnConfig::new("amount_cents", DataType::Int64, false),
8566 ],
8567 vec!["id".to_string()],
8568 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
8569 .expect("valid")
8570 .with_cover_columns(vec!["amount_cents".to_string()])],
8571 )
8572 .expect("schema");
8573
8574 let mut writer = schema.batch_writer();
8575 for (id, status, amount) in [
8576 (1, "open", 10),
8577 (2, "closed", 15),
8578 (3, "open", 30),
8579 (4, "closed", 40),
8580 ] {
8581 writer
8582 .insert(
8583 "orders",
8584 vec![
8585 CellValue::Int64(id),
8586 CellValue::Utf8(status.to_string()),
8587 CellValue::Int64(amount),
8588 ],
8589 )
8590 .expect("row");
8591 }
8592 writer.flush().await.expect("flush");
8593
8594 let ctx = SessionContext::new();
8595 schema.register_all(&ctx).expect("register");
8596
8597 state.range_calls.store(0, AtomicOrdering::SeqCst);
8598 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
8599
8600 let query = "SELECT SUM(CASE status WHEN 'open' THEN amount_cents END) AS open_total, \
8601 COUNT(CASE status WHEN 'closed' THEN 1 END) AS closed_count, \
8602 AVG(CASE WHEN status = 'closed' THEN amount_cents END) AS closed_avg \
8603 FROM orders";
8604 let batches = ctx
8605 .sql(query)
8606 .await
8607 .expect("query")
8608 .collect()
8609 .await
8610 .expect("collect");
8611
8612 assert_eq!(batches.len(), 1);
8613 let batch = &batches[0];
8614 assert_eq!(
8615 ScalarValue::try_from_array(batch.column(0), 0).expect("sum scalar"),
8616 ScalarValue::Int64(Some(40))
8617 );
8618 assert_count_scalar(batch, 1, 0, 2);
8619 let closed_avg = batch
8620 .column(2)
8621 .as_any()
8622 .downcast_ref::<Float64Array>()
8623 .expect("avg float64")
8624 .value(0);
8625 assert_eq!(closed_avg, 27.5);
8626 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
8627 assert!(
8628 state.range_reduce_calls.load(AtomicOrdering::SeqCst) >= 3,
8629 "case-based conditional aggregates should use reduction jobs"
8630 );
8631
8632 let _ = shutdown_tx.send(());
8633 }
8634
8635 #[tokio::test]
8636 async fn aggregate_pushdown_supports_casted_group_and_aggregate_expressions() {
8637 let state = MockState {
8638 kv: Arc::new(Mutex::new(BTreeMap::new())),
8639 range_calls: Arc::new(AtomicUsize::new(0)),
8640 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
8641 sequence_number: Arc::new(AtomicU64::new(0)),
8642 };
8643 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
8644 let client = StoreClient::new(&base_url);
8645
8646 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
8647 .table(
8648 "orders",
8649 vec![
8650 TableColumnConfig::new("id", DataType::Int64, false),
8651 TableColumnConfig::new("status", DataType::Utf8, false),
8652 TableColumnConfig::new("amount_cents", DataType::Int64, false),
8653 ],
8654 vec!["id".to_string()],
8655 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
8656 .expect("valid")
8657 .with_cover_columns(vec!["amount_cents".to_string()])],
8658 )
8659 .expect("schema");
8660
8661 let mut writer = schema.batch_writer();
8662 for (id, status, amount) in [
8663 (1, "open", 10),
8664 (2, "open", 30),
8665 (3, "closed", 15),
8666 (4, "closed", 40),
8667 ] {
8668 writer
8669 .insert(
8670 "orders",
8671 vec![
8672 CellValue::Int64(id),
8673 CellValue::Utf8(status.to_string()),
8674 CellValue::Int64(amount),
8675 ],
8676 )
8677 .expect("row");
8678 }
8679 writer.flush().await.expect("flush");
8680
8681 let ctx = SessionContext::new();
8682 schema.register_all(&ctx).expect("register");
8683
8684 state.range_calls.store(0, AtomicOrdering::SeqCst);
8685 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
8686
8687 let batches = ctx
8688 .sql(
8689 "SELECT CAST(status AS VARCHAR) AS status_text, \
8690 SUM(CAST(amount_cents AS DOUBLE)) AS total_cents \
8691 FROM orders \
8692 GROUP BY CAST(status AS VARCHAR) \
8693 ORDER BY status_text",
8694 )
8695 .await
8696 .expect("query")
8697 .collect()
8698 .await
8699 .expect("collect");
8700
8701 assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
8702 let batch = &batches[0];
8703 let status = ScalarValue::try_from_array(batch.column(0), 0).expect("status scalar");
8704 assert_eq!(scalar_to_string(&status).as_deref(), Some("closed"));
8705 let closed_total = batch
8706 .column(1)
8707 .as_any()
8708 .downcast_ref::<Float64Array>()
8709 .expect("sum float64")
8710 .value(0);
8711 let open_total = batch
8712 .column(1)
8713 .as_any()
8714 .downcast_ref::<Float64Array>()
8715 .expect("sum float64")
8716 .value(1);
8717 assert_eq!(closed_total, 55.0);
8718 assert_eq!(open_total, 40.0);
8719 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
8720 assert!(
8721 state.range_reduce_calls.load(AtomicOrdering::SeqCst) >= 1,
8722 "casted grouped aggregates should stay on the reduction path"
8723 );
8724
8725 let _ = shutdown_tx.send(());
8726 }
8727
8728 #[tokio::test]
8729 async fn aggregate_pushdown_supports_computed_aggregate_inputs() {
8730 let state = MockState {
8731 kv: Arc::new(Mutex::new(BTreeMap::new())),
8732 range_calls: Arc::new(AtomicUsize::new(0)),
8733 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
8734 sequence_number: Arc::new(AtomicU64::new(0)),
8735 };
8736 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
8737 let client = StoreClient::new(&base_url);
8738
8739 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
8740 .table(
8741 "orders",
8742 vec![
8743 TableColumnConfig::new("id", DataType::Int64, false),
8744 TableColumnConfig::new("price_cents", DataType::Int64, false),
8745 TableColumnConfig::new("qty", DataType::Int64, false),
8746 TableColumnConfig::new("duration_ms", DataType::Int64, false),
8747 ],
8748 vec!["id".to_string()],
8749 vec![],
8750 )
8751 .expect("schema");
8752
8753 let mut writer = schema.batch_writer();
8754 for (id, price, qty, duration_ms) in [(1, 10, 2, 500), (2, 15, 3, 2500), (3, 7, 4, 1000)] {
8755 writer
8756 .insert(
8757 "orders",
8758 vec![
8759 CellValue::Int64(id),
8760 CellValue::Int64(price),
8761 CellValue::Int64(qty),
8762 CellValue::Int64(duration_ms),
8763 ],
8764 )
8765 .expect("row");
8766 }
8767 writer.flush().await.expect("flush");
8768
8769 let ctx = SessionContext::new();
8770 schema.register_all(&ctx).expect("register");
8771
8772 state.range_calls.store(0, AtomicOrdering::SeqCst);
8773 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
8774
8775 let batches = ctx
8776 .sql(
8777 "SELECT SUM(price_cents * qty) AS total_revenue, \
8778 AVG(duration_ms / 1e3) AS avg_seconds \
8779 FROM orders",
8780 )
8781 .await
8782 .expect("query")
8783 .collect()
8784 .await
8785 .expect("collect");
8786
8787 assert_eq!(batches.len(), 1);
8788 let batch = &batches[0];
8789 assert_eq!(
8790 ScalarValue::try_from_array(batch.column(0), 0).expect("sum scalar"),
8791 ScalarValue::Int64(Some(93))
8792 );
8793 let avg_seconds = batch
8794 .column(1)
8795 .as_any()
8796 .downcast_ref::<Float64Array>()
8797 .expect("avg float64")
8798 .value(0);
8799 assert!((avg_seconds - (4.0 / 3.0)).abs() < 1e-12);
8800 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
8801 assert!(
8802 state.range_reduce_calls.load(AtomicOrdering::SeqCst) >= 2,
8803 "computed aggregate inputs should use reduction jobs"
8804 );
8805
8806 let _ = shutdown_tx.send(());
8807 }
8808
8809 #[tokio::test]
8810 async fn aggregate_pushdown_supports_add_and_subtract_inputs() {
8811 let state = MockState {
8812 kv: Arc::new(Mutex::new(BTreeMap::new())),
8813 range_calls: Arc::new(AtomicUsize::new(0)),
8814 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
8815 sequence_number: Arc::new(AtomicU64::new(0)),
8816 };
8817 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
8818 let client = StoreClient::new(&base_url);
8819
8820 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
8821 .table(
8822 "orders",
8823 vec![
8824 TableColumnConfig::new("id", DataType::Int64, false),
8825 TableColumnConfig::new("price_cents", DataType::Int64, false),
8826 TableColumnConfig::new("fee_cents", DataType::Int64, false),
8827 TableColumnConfig::new("discount_cents", DataType::Int64, false),
8828 ],
8829 vec!["id".to_string()],
8830 vec![],
8831 )
8832 .expect("schema");
8833
8834 let mut writer = schema.batch_writer();
8835 for (id, price, fee, discount) in [(1, 10, 2, 1), (2, 15, 3, 4), (3, 7, 1, 2)] {
8836 writer
8837 .insert(
8838 "orders",
8839 vec![
8840 CellValue::Int64(id),
8841 CellValue::Int64(price),
8842 CellValue::Int64(fee),
8843 CellValue::Int64(discount),
8844 ],
8845 )
8846 .expect("row");
8847 }
8848 writer.flush().await.expect("flush");
8849
8850 let ctx = SessionContext::new();
8851 schema.register_all(&ctx).expect("register");
8852
8853 state.range_calls.store(0, AtomicOrdering::SeqCst);
8854 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
8855
8856 let batches = ctx
8857 .sql(
8858 "SELECT SUM(price_cents + fee_cents) AS gross_plus_fee, \
8859 SUM(price_cents - discount_cents) AS net_total \
8860 FROM orders",
8861 )
8862 .await
8863 .expect("query")
8864 .collect()
8865 .await
8866 .expect("collect");
8867
8868 assert_eq!(batches.len(), 1);
8869 let batch = &batches[0];
8870 assert_eq!(
8871 ScalarValue::try_from_array(batch.column(0), 0).expect("sum scalar"),
8872 ScalarValue::Int64(Some(38))
8873 );
8874 assert_eq!(
8875 ScalarValue::try_from_array(batch.column(1), 0).expect("sum scalar"),
8876 ScalarValue::Int64(Some(25))
8877 );
8878 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
8879 assert!(
8880 state.range_reduce_calls.load(AtomicOrdering::SeqCst) >= 2,
8881 "add/sub aggregate inputs should use reduction jobs"
8882 );
8883
8884 let _ = shutdown_tx.send(());
8885 }
8886
8887 #[tokio::test]
8888 async fn aggregate_pushdown_supports_case_filtered_computed_aggregates() {
8889 let state = MockState {
8890 kv: Arc::new(Mutex::new(BTreeMap::new())),
8891 range_calls: Arc::new(AtomicUsize::new(0)),
8892 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
8893 sequence_number: Arc::new(AtomicU64::new(0)),
8894 };
8895 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
8896 let client = StoreClient::new(&base_url);
8897
8898 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
8899 .table(
8900 "orders",
8901 vec![
8902 TableColumnConfig::new("id", DataType::Int64, false),
8903 TableColumnConfig::new("status", DataType::Utf8, false),
8904 TableColumnConfig::new("price_cents", DataType::Int64, false),
8905 TableColumnConfig::new("qty", DataType::Int64, false),
8906 ],
8907 vec!["id".to_string()],
8908 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
8909 .expect("valid")
8910 .with_cover_columns(vec!["price_cents".to_string(), "qty".to_string()])],
8911 )
8912 .expect("schema");
8913
8914 let mut writer = schema.batch_writer();
8915 for (id, status, price, qty) in [
8916 (1, "open", 10, 2),
8917 (2, "closed", 99, 1),
8918 (3, "open", 15, 3),
8919 (4, "closed", 7, 4),
8920 ] {
8921 writer
8922 .insert(
8923 "orders",
8924 vec![
8925 CellValue::Int64(id),
8926 CellValue::Utf8(status.to_string()),
8927 CellValue::Int64(price),
8928 CellValue::Int64(qty),
8929 ],
8930 )
8931 .expect("row");
8932 }
8933 writer.flush().await.expect("flush");
8934
8935 let ctx = SessionContext::new();
8936 schema.register_all(&ctx).expect("register");
8937
8938 state.range_calls.store(0, AtomicOrdering::SeqCst);
8939 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
8940
8941 let batches = ctx
8942 .sql(
8943 "SELECT SUM(CASE WHEN status = 'open' THEN price_cents * qty END) \
8944 AS open_revenue \
8945 FROM orders",
8946 )
8947 .await
8948 .expect("query")
8949 .collect()
8950 .await
8951 .expect("collect");
8952
8953 assert_eq!(batches.len(), 1);
8954 let batch = &batches[0];
8955 assert_eq!(
8956 ScalarValue::try_from_array(batch.column(0), 0).expect("sum scalar"),
8957 ScalarValue::Int64(Some(65))
8958 );
8959 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
8960 assert!(
8961 state.range_reduce_calls.load(AtomicOrdering::SeqCst) >= 1,
8962 "case-filtered computed aggregate should use reduction jobs"
8963 );
8964
8965 let _ = shutdown_tx.send(());
8966 }
8967
8968 #[tokio::test]
8969 async fn aggregate_pushdown_does_not_rewrite_sum_case_else_zero_semantics() {
8970 let state = MockState {
8971 kv: Arc::new(Mutex::new(BTreeMap::new())),
8972 range_calls: Arc::new(AtomicUsize::new(0)),
8973 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
8974 sequence_number: Arc::new(AtomicU64::new(0)),
8975 };
8976 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
8977 let client = StoreClient::new(&base_url);
8978
8979 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
8980 .table(
8981 "orders",
8982 vec![
8983 TableColumnConfig::new("id", DataType::Int64, false),
8984 TableColumnConfig::new("region", DataType::Utf8, false),
8985 TableColumnConfig::new("status", DataType::Utf8, false),
8986 TableColumnConfig::new("amount_cents", DataType::Int64, false),
8987 ],
8988 vec!["id".to_string()],
8989 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
8990 .expect("valid")
8991 .with_cover_columns(vec!["region".to_string(), "amount_cents".to_string()])],
8992 )
8993 .expect("schema");
8994
8995 let mut writer = schema.batch_writer();
8996 for (id, region, status, amount) in [
8997 (1, "east", "open", 10),
8998 (2, "east", "closed", 20),
8999 (3, "west", "closed", 30),
9000 ] {
9001 writer
9002 .insert(
9003 "orders",
9004 vec![
9005 CellValue::Int64(id),
9006 CellValue::Utf8(region.to_string()),
9007 CellValue::Utf8(status.to_string()),
9008 CellValue::Int64(amount),
9009 ],
9010 )
9011 .expect("row");
9012 }
9013 writer.flush().await.expect("flush");
9014
9015 let ctx = SessionContext::new();
9016 schema.register_all(&ctx).expect("register");
9017
9018 state.range_calls.store(0, AtomicOrdering::SeqCst);
9019 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
9020
9021 let batches = ctx
9022 .sql(
9023 "SELECT region, \
9024 SUM(CASE WHEN status = 'open' THEN amount_cents ELSE 0 END) AS open_total \
9025 FROM orders \
9026 GROUP BY region \
9027 ORDER BY region",
9028 )
9029 .await
9030 .expect("query")
9031 .collect()
9032 .await
9033 .expect("collect");
9034
9035 assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
9036 let batch = &batches[0];
9037 assert_eq!(
9038 ScalarValue::try_from_array(batch.column(0), 0).expect("region scalar"),
9039 ScalarValue::Utf8(Some("east".to_string()))
9040 );
9041 assert_eq!(
9042 ScalarValue::try_from_array(batch.column(1), 0).expect("sum scalar"),
9043 ScalarValue::Int64(Some(10))
9044 );
9045 assert_eq!(
9046 ScalarValue::try_from_array(batch.column(0), 1).expect("region scalar"),
9047 ScalarValue::Utf8(Some("west".to_string()))
9048 );
9049 assert_eq!(
9050 ScalarValue::try_from_array(batch.column(1), 1).expect("sum scalar"),
9051 ScalarValue::Int64(Some(0))
9052 );
9053 assert_eq!(
9054 state.range_reduce_calls.load(AtomicOrdering::SeqCst),
9055 0,
9056 "SUM(CASE ... ELSE 0 END) must not push down because FILTER rewrite changes semantics"
9057 );
9058
9059 let _ = shutdown_tx.send(());
9060 }
9061
9062 #[tokio::test]
9063 async fn aggregate_pushdown_supports_computed_group_keys() {
9064 let state = MockState {
9065 kv: Arc::new(Mutex::new(BTreeMap::new())),
9066 range_calls: Arc::new(AtomicUsize::new(0)),
9067 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
9068 sequence_number: Arc::new(AtomicU64::new(0)),
9069 };
9070 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
9071 let client = StoreClient::new(&base_url);
9072
9073 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
9074 .table(
9075 "events",
9076 vec![
9077 TableColumnConfig::new("id", DataType::Int64, false),
9078 TableColumnConfig::new("country", DataType::Utf8, false),
9079 TableColumnConfig::new(
9080 "occurred_at",
9081 DataType::Timestamp(TimeUnit::Microsecond, None),
9082 false,
9083 ),
9084 TableColumnConfig::new("amount_cents", DataType::Int64, false),
9085 ],
9086 vec!["id".to_string()],
9087 vec![],
9088 )
9089 .expect("schema");
9090
9091 let day_micros = 86_400_000_000i64;
9092 let day0 = 1_700_000_000_000_000i64;
9093 let day1 = day0 + day_micros;
9094 let day0_bucket = day0.div_euclid(day_micros) * day_micros;
9095 let day1_bucket = day1.div_euclid(day_micros) * day_micros;
9096 let mut writer = schema.batch_writer();
9097 for (id, country, occurred_at, amount) in [
9098 (1, "East", day0 + 111, 10),
9099 (2, "east", day0 + 222, 30),
9100 (3, "West", day1 + 333, 7),
9101 ] {
9102 writer
9103 .insert(
9104 "events",
9105 vec![
9106 CellValue::Int64(id),
9107 CellValue::Utf8(country.to_string()),
9108 CellValue::Timestamp(occurred_at),
9109 CellValue::Int64(amount),
9110 ],
9111 )
9112 .expect("row");
9113 }
9114 writer.flush().await.expect("flush");
9115
9116 let ctx = SessionContext::new();
9117 schema.register_all(&ctx).expect("register");
9118
9119 state.range_calls.store(0, AtomicOrdering::SeqCst);
9120 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
9121
9122 let batches = ctx
9123 .sql(
9124 "SELECT lower(country) AS country_norm, \
9125 date_trunc('day', occurred_at) AS day_bucket, \
9126 SUM(amount_cents) AS total_cents \
9127 FROM events \
9128 GROUP BY lower(country), date_trunc('day', occurred_at) \
9129 ORDER BY country_norm, day_bucket",
9130 )
9131 .await
9132 .expect("query")
9133 .collect()
9134 .await
9135 .expect("collect");
9136
9137 assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
9138 let batch = &batches[0];
9139 assert_eq!(
9140 scalar_to_string(
9141 &ScalarValue::try_from_array(batch.column(0), 0).expect("country scalar")
9142 )
9143 .as_deref(),
9144 Some("east")
9145 );
9146 assert_eq!(
9147 ScalarValue::try_from_array(batch.column(1), 0).expect("day scalar"),
9148 ScalarValue::TimestampMicrosecond(Some(day0_bucket), None)
9149 );
9150 assert_eq!(
9151 ScalarValue::try_from_array(batch.column(2), 0).expect("sum scalar"),
9152 ScalarValue::Int64(Some(40))
9153 );
9154 assert_eq!(
9155 scalar_to_string(
9156 &ScalarValue::try_from_array(batch.column(0), 1).expect("country scalar")
9157 )
9158 .as_deref(),
9159 Some("west")
9160 );
9161 assert_eq!(
9162 ScalarValue::try_from_array(batch.column(1), 1).expect("day scalar"),
9163 ScalarValue::TimestampMicrosecond(Some(day1_bucket), None)
9164 );
9165 assert_eq!(
9166 ScalarValue::try_from_array(batch.column(2), 1).expect("sum scalar"),
9167 ScalarValue::Int64(Some(7))
9168 );
9169 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
9170 assert!(
9171 state.range_reduce_calls.load(AtomicOrdering::SeqCst) >= 1,
9172 "computed group keys should use grouped reduction path"
9173 );
9174
9175 let _ = shutdown_tx.send(());
9176 }
9177
9178 #[tokio::test]
9179 async fn aggregate_pushdown_supports_group_by_queries() {
9180 let state = MockState {
9181 kv: Arc::new(Mutex::new(BTreeMap::new())),
9182 range_calls: Arc::new(AtomicUsize::new(0)),
9183 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
9184 sequence_number: Arc::new(AtomicU64::new(0)),
9185 };
9186 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
9187 let client = StoreClient::new(&base_url);
9188
9189 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
9190 .table(
9191 "orders",
9192 vec![
9193 TableColumnConfig::new("id", DataType::Int64, false),
9194 TableColumnConfig::new("status", DataType::Utf8, false),
9195 TableColumnConfig::new("amount_cents", DataType::Int64, false),
9196 ],
9197 vec!["id".to_string()],
9198 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
9199 .expect("valid")
9200 .with_cover_columns(vec!["amount_cents".to_string()])],
9201 )
9202 .expect("schema");
9203
9204 let mut writer = schema.batch_writer();
9205 for (id, status, amount) in [
9206 (1, "open", 10),
9207 (2, "open", 30),
9208 (3, "closed", 15),
9209 (4, "closed", 40),
9210 ] {
9211 writer
9212 .insert(
9213 "orders",
9214 vec![
9215 CellValue::Int64(id),
9216 CellValue::Utf8(status.to_string()),
9217 CellValue::Int64(amount),
9218 ],
9219 )
9220 .expect("row");
9221 }
9222 writer.flush().await.expect("flush");
9223
9224 let ctx = SessionContext::new();
9225 schema.register_all(&ctx).expect("register");
9226
9227 state.range_calls.store(0, AtomicOrdering::SeqCst);
9228 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
9229
9230 let batches = ctx
9231 .sql(
9232 "SELECT status, COUNT(*) AS row_count, SUM(amount_cents) AS total_cents \
9233 FROM orders GROUP BY status ORDER BY status",
9234 )
9235 .await
9236 .expect("query")
9237 .collect()
9238 .await
9239 .expect("collect");
9240
9241 assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
9242 let batch = &batches[0];
9243 assert_eq!(
9244 ScalarValue::try_from_array(batch.column(0), 0).expect("status scalar"),
9245 ScalarValue::Utf8(Some("closed".to_string()))
9246 );
9247 assert_count_scalar(batch, 1, 0, 2);
9248 assert_eq!(
9249 ScalarValue::try_from_array(batch.column(2), 0).expect("sum scalar"),
9250 ScalarValue::Int64(Some(55))
9251 );
9252 assert_eq!(
9253 ScalarValue::try_from_array(batch.column(0), 1).expect("status scalar"),
9254 ScalarValue::Utf8(Some("open".to_string()))
9255 );
9256 assert_count_scalar(batch, 1, 1, 2);
9257 assert_eq!(
9258 ScalarValue::try_from_array(batch.column(2), 1).expect("sum scalar"),
9259 ScalarValue::Int64(Some(40))
9260 );
9261 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
9262 assert!(
9263 state.range_reduce_calls.load(AtomicOrdering::SeqCst) >= 2,
9264 "group-by aggregate should use grouped range reduction path"
9265 );
9266
9267 let _ = shutdown_tx.send(());
9268 }
9269
9270 #[tokio::test]
9271 async fn aggregate_pushdown_group_by_float_canonicalizes_signed_zero() {
9272 let state = MockState {
9273 kv: Arc::new(Mutex::new(BTreeMap::new())),
9274 range_calls: Arc::new(AtomicUsize::new(0)),
9275 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
9276 sequence_number: Arc::new(AtomicU64::new(0)),
9277 };
9278 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
9279 let client = StoreClient::new(&base_url);
9280
9281 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
9282 .table(
9283 "metrics",
9284 vec![
9285 TableColumnConfig::new("id", DataType::Int64, false),
9286 TableColumnConfig::new("score", DataType::Float64, false),
9287 ],
9288 vec!["id".to_string()],
9289 vec![],
9290 )
9291 .expect("schema");
9292
9293 let mut writer = schema.batch_writer();
9294 for (id, score) in [(1, -0.0), (2, 0.0), (3, 1.5)] {
9295 writer
9296 .insert(
9297 "metrics",
9298 vec![CellValue::Int64(id), CellValue::Float64(score)],
9299 )
9300 .expect("row");
9301 }
9302 writer.flush().await.expect("flush");
9303
9304 let ctx = SessionContext::new();
9305 schema.register_all(&ctx).expect("register");
9306
9307 state.range_calls.store(0, AtomicOrdering::SeqCst);
9308 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
9309
9310 let batches = ctx
9311 .sql(
9312 "SELECT score, COUNT(*) AS row_count \
9313 FROM metrics GROUP BY score ORDER BY row_count DESC, score",
9314 )
9315 .await
9316 .expect("query")
9317 .collect()
9318 .await
9319 .expect("collect");
9320
9321 assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
9322 let batch = &batches[0];
9323 let top_score = batch
9324 .column(0)
9325 .as_any()
9326 .downcast_ref::<Float64Array>()
9327 .expect("score float64")
9328 .value(0);
9329 assert_eq!(top_score.to_bits(), 0.0f64.to_bits());
9330 assert_count_scalar(batch, 1, 0, 2);
9331 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
9332 assert!(
9333 state.range_reduce_calls.load(AtomicOrdering::SeqCst) >= 1,
9334 "float group-by aggregate should stay on grouped reduction path"
9335 );
9336
9337 let _ = shutdown_tx.send(());
9338 }
9339
9340 #[tokio::test]
9341 async fn aggregate_pushdown_supports_filtered_group_by_queries() {
9342 let state = MockState {
9343 kv: Arc::new(Mutex::new(BTreeMap::new())),
9344 range_calls: Arc::new(AtomicUsize::new(0)),
9345 range_reduce_calls: Arc::new(AtomicUsize::new(0)),
9346 sequence_number: Arc::new(AtomicU64::new(0)),
9347 };
9348 let (base_url, shutdown_tx) = spawn_mock_server(state.clone()).await;
9349 let client = StoreClient::new(&base_url);
9350
9351 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
9352 .table(
9353 "orders",
9354 vec![
9355 TableColumnConfig::new("id", DataType::Int64, false),
9356 TableColumnConfig::new("region", DataType::Utf8, false),
9357 TableColumnConfig::new("status", DataType::Utf8, false),
9358 TableColumnConfig::new("amount_cents", DataType::Int64, false),
9359 ],
9360 vec!["id".to_string()],
9361 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
9362 .expect("valid")
9363 .with_cover_columns(vec!["region".to_string(), "amount_cents".to_string()])],
9364 )
9365 .expect("schema");
9366
9367 let mut writer = schema.batch_writer();
9368 for (id, region, status, amount) in [
9369 (1, "east", "open", 10),
9370 (2, "east", "closed", 20),
9371 (3, "west", "open", 30),
9372 (4, "north", "closed", 40),
9373 ] {
9374 writer
9375 .insert(
9376 "orders",
9377 vec![
9378 CellValue::Int64(id),
9379 CellValue::Utf8(region.to_string()),
9380 CellValue::Utf8(status.to_string()),
9381 CellValue::Int64(amount),
9382 ],
9383 )
9384 .expect("row");
9385 }
9386 writer.flush().await.expect("flush");
9387
9388 let ctx = SessionContext::new();
9389 schema.register_all(&ctx).expect("register");
9390
9391 state.range_calls.store(0, AtomicOrdering::SeqCst);
9392 state.range_reduce_calls.store(0, AtomicOrdering::SeqCst);
9393
9394 let batches = ctx
9395 .sql(
9396 "SELECT region, \
9397 COUNT(*) FILTER (WHERE status = 'open') AS open_count, \
9398 SUM(amount_cents) FILTER (WHERE status = 'closed') AS closed_total \
9399 FROM orders \
9400 GROUP BY region \
9401 ORDER BY region",
9402 )
9403 .await
9404 .expect("query")
9405 .collect()
9406 .await
9407 .expect("collect");
9408
9409 assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 3);
9410 let batch = &batches[0];
9411
9412 assert_eq!(
9413 ScalarValue::try_from_array(batch.column(0), 0).expect("region scalar"),
9414 ScalarValue::Utf8(Some("east".to_string()))
9415 );
9416 assert_count_scalar(batch, 1, 0, 1);
9417 assert_eq!(
9418 ScalarValue::try_from_array(batch.column(2), 0).expect("sum scalar"),
9419 ScalarValue::Int64(Some(20))
9420 );
9421
9422 assert_eq!(
9423 ScalarValue::try_from_array(batch.column(0), 1).expect("region scalar"),
9424 ScalarValue::Utf8(Some("north".to_string()))
9425 );
9426 assert_count_scalar(batch, 1, 1, 0);
9427 assert_eq!(
9428 ScalarValue::try_from_array(batch.column(2), 1).expect("sum scalar"),
9429 ScalarValue::Int64(Some(40))
9430 );
9431
9432 assert_eq!(
9433 ScalarValue::try_from_array(batch.column(0), 2).expect("region scalar"),
9434 ScalarValue::Utf8(Some("west".to_string()))
9435 );
9436 assert_count_scalar(batch, 1, 2, 1);
9437 assert_eq!(
9438 ScalarValue::try_from_array(batch.column(2), 2).expect("sum scalar"),
9439 ScalarValue::Int64(None)
9440 );
9441
9442 assert_eq!(state.range_calls.load(AtomicOrdering::SeqCst), 0);
9443 assert!(
9444 state.range_reduce_calls.load(AtomicOrdering::SeqCst) >= 3,
9445 "filtered group-by aggregate should use grouped reduction plus seed job"
9446 );
9447
9448 let _ = shutdown_tx.send(());
9449 }
9450
9451 mod e2e {
9452 use super::*;
9453 use datafusion::prelude::SessionContext;
9454 use exoware_sdk::StoreClient;
9455
9456 struct TestServers {
9457 ingest_url: String,
9458 query_url: String,
9459 }
9460
9461 impl TestServers {
9462 fn client(&self) -> StoreClient {
9463 StoreClient::builder()
9464 .health_url(&self.query_url)
9465 .ingest_url(&self.ingest_url)
9466 .query_url(&self.query_url)
9467 .compact_url(&self.ingest_url)
9468 .stream_url(&self.query_url)
9469 .build()
9470 .expect("test server URLs are set")
9471 }
9472 }
9473
9474 async fn spawn_e2e_servers() -> TestServers {
9475 let (_task, url) = exoware_simulator::open_temp()
9476 .await
9477 .expect("spawn simulator");
9478 TestServers {
9479 ingest_url: url.clone(),
9480 query_url: url,
9481 }
9482 }
9483
9484 #[tokio::test]
9485 async fn sql_insert_and_select_through_real_ingest_query_workers() {
9486 let servers = spawn_e2e_servers().await;
9487 let client = servers.client();
9488
9489 let schema = KvSchema::new(PrefixedStoreClient::empty(client))
9490 .table(
9491 "orders",
9492 vec![
9493 TableColumnConfig::new("id", DataType::Int64, false),
9494 TableColumnConfig::new("status", DataType::Utf8, false),
9495 TableColumnConfig::new("amount_cents", DataType::Int64, false),
9496 ],
9497 vec!["id".to_string()],
9498 vec![IndexSpec::new("status_idx", vec!["status".to_string()])
9499 .expect("valid")
9500 .with_cover_columns(vec!["amount_cents".to_string()])],
9501 )
9502 .expect("schema");
9503
9504 let mut writer = schema.batch_writer();
9505 for (id, status, amount) in [
9506 (1i64, "open", 100i64),
9507 (2, "closed", 200),
9508 (3, "open", 300),
9509 (4, "closed", 400),
9510 (5, "open", 500),
9511 ] {
9512 writer
9513 .insert(
9514 "orders",
9515 vec![
9516 CellValue::Int64(id),
9517 CellValue::Utf8(status.to_string()),
9518 CellValue::Int64(amount),
9519 ],
9520 )
9521 .expect("insert row");
9522 }
9523 writer.flush().await.expect("flush batch");
9524
9525 let ctx = SessionContext::new();
9526 schema.register_all(&ctx).expect("register tables");
9527
9528 let batches = ctx
9529 .sql("SELECT id, amount_cents FROM orders ORDER BY id")
9530 .await
9531 .expect("full scan query")
9532 .collect()
9533 .await
9534 .expect("collect full scan");
9535 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
9536 assert_eq!(total_rows, 5, "all 5 rows returned from full scan");
9537
9538 let mut ids = Vec::new();
9539 let mut amounts = Vec::new();
9540 for batch in &batches {
9541 let id_col = batch
9542 .column(0)
9543 .as_any()
9544 .downcast_ref::<Int64Array>()
9545 .expect("id column");
9546 let amt_col = batch
9547 .column(1)
9548 .as_any()
9549 .downcast_ref::<Int64Array>()
9550 .expect("amount column");
9551 for i in 0..batch.num_rows() {
9552 ids.push(id_col.value(i));
9553 amounts.push(amt_col.value(i));
9554 }
9555 }
9556 assert_eq!(ids, vec![1, 2, 3, 4, 5]);
9557 assert_eq!(amounts, vec![100, 200, 300, 400, 500]);
9558
9559 let filtered = ctx
9560 .sql(
9561 "SELECT id, amount_cents FROM orders \
9562 WHERE status = 'open' ORDER BY id",
9563 )
9564 .await
9565 .expect("filtered query")
9566 .collect()
9567 .await
9568 .expect("collect filtered");
9569 let mut filtered_ids = Vec::new();
9570 let mut filtered_amounts = Vec::new();
9571 for batch in &filtered {
9572 let id_col = batch
9573 .column(0)
9574 .as_any()
9575 .downcast_ref::<Int64Array>()
9576 .expect("id column");
9577 let amt_col = batch
9578 .column(1)
9579 .as_any()
9580 .downcast_ref::<Int64Array>()
9581 .expect("amount column");
9582 for i in 0..batch.num_rows() {
9583 filtered_ids.push(id_col.value(i));
9584 filtered_amounts.push(amt_col.value(i));
9585 }
9586 }
9587 assert_eq!(filtered_ids, vec![1, 3, 5]);
9588 assert_eq!(filtered_amounts, vec![100, 300, 500]);
9589
9590 let agg = ctx
9591 .sql(
9592 "SELECT COUNT(*) AS cnt, SUM(amount_cents) AS total \
9593 FROM orders WHERE status = 'open'",
9594 )
9595 .await
9596 .expect("aggregate query")
9597 .collect()
9598 .await
9599 .expect("collect aggregate");
9600 assert_eq!(agg.len(), 1);
9601 let batch = &agg[0];
9602 assert_eq!(batch.num_rows(), 1);
9603 assert_count_scalar(batch, 0, 0, 3);
9604 let total = ScalarValue::try_from_array(batch.column(1), 0).expect("sum scalar");
9605 match total {
9606 ScalarValue::Int64(Some(v)) => assert_eq!(v, 900),
9607 other => panic!("unexpected sum type: {other:?}"),
9608 }
9609 }
9610 }
9611}