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