Skip to content

Commit 60d3bdc

Browse files
committed
[Sandbox] Reverse indexed-parquet segment iteration for opposite-direction sorts
Mirror of `ContextIndexSearcher.shouldUseTimeSeriesDescSortOptimization` for the indexed-parquet path. When `index.sort.field` is configured and the query's leading ORDER BY runs counter to the catalog's stored direction, reverse the per-shard segment iteration so a TopK above the scan pulls the highest-priority segment first and parquet page stats prune the rest. Reuses PR #22041's `ShardView.sort_fields` / `sort_orders` plumbing — no new Java/FFM surface. Only the indexed scan path (`IndexedTableProvider` / `QueryShardExec`) consumes them; the vanilla `ListingTable` path remains covered by #22041. ## Implementation `indexed_executor.rs`: - `analyze_top_sort(&LogicalPlan)` walks Projection/Limit/SubqueryAlias/ Distinct/Filter to find the leading top-of-plan ORDER BY. Returns None if the leading sort key isn't a plain `Expr::Column` (function on the sort key breaks catalog monotonicity). - `should_reverse_segments` triggers iff the catalog has a sort, the query's leading sort column matches `sort_fields[0]` (case-sensitive, matching PR #22041), and the directions differ. - `reverse_segment_iteration_order` flips `Vec<SegmentFileInfo>` in place. Per-segment `global_base` is deliberately NOT recomputed — see below. - Wired into `execute_indexed_with_context_inner` after `from_substrait_plan` decodes the LogicalPlan. `SessionContextHandle` carries `sort_fields` / `sort_orders` cloned from `ShardView` so both `create_session_context` and `create_session_context_indexed` populate the indexed-path consumer. ## QTF correctness — why `global_base` is preserved on reversal The `__row_id__` emitted by the indexed scan is `segment.global_base + position`. QTF's fetch phase (`api::fetch_by_row_ids`) goes through a different code path (`ShardTableProvider`, not `IndexedTableProvider`) and rebuilds segments via `build_segments(ShardView.object_metas, ShardView.writer_generations)` — always in catalog order. Fetch then partition_points the cached row IDs against catalog-order `global_base`s to find the owning segment. For the round trip to hold, query-phase IDs must match catalog-order bases. `reverse_segment_iteration_order` therefore flips iteration order ONLY — each segment retains its catalog-order `global_base`. Reversal becomes purely a scheduling decision; the row-id space is unchanged. ## Tests Rust unit tests in `indexed_executor::tests`: - `analyze_top_sort` (top-level desc, descends through Limit/Projection, None for no-sort, None for function sort key). - `should_reverse_segments` (matching field opposite direction, matching same direction, catalog desc query asc, no query sort, no catalog sort, non-leading catalog field, case-sensitive name). - `reverse_segments_preserves_global_base`, `reverse_segments_empty_is_noop`, `reverse_segments_single_keeps_its_base`. Rust e2e tests in `indexed_table/tests_e2e/sort_reverse_row_id.rs`: - `three_segments_row_ids_invariant_under_reversal`, `four_segments_row_ids_invariant_under_reversal` — pin the load-bearing invariant: same shard-global `__row_id__` set whether iteration is natural or reversed. Java IT in `sandbox/qa/analytics-engine-coordinator/.../SortReverseQtfIT.java` (parquet+lucene composite, 2 shards, 3 flushes → 2-4 segments per shard): - `testReversalTriggers_orderByTsDescLimit_correctResults` — captures the reversal log line via `MockLogAppender` on `RustLoggerBridge` AND asserts result correctness (top-k by ts DESC). - `testNoReversal_orderByTsAscLimit_correctResults` — `UnseenEventExpectation` proves reversal does NOT fire; results still correct. - `testReversal_keywordFilterDelegationPath_correctResults` — keyword filter shipped to Lucene perf peer + reversal; verifies QTF row-id round-trip composes with delegation. Signed-off-by: Arpit Bandejiya <abandeji@amazon.com>
1 parent ce2e409 commit 60d3bdc

9 files changed

Lines changed: 994 additions & 9 deletions

File tree

sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -333,9 +333,14 @@ pub struct ShardView {
333333
/// Index sort fields, in priority order. Sourced from the index's
334334
/// `index.sort.field` setting on the Java side. Empty when the index has
335335
/// no `index.sort.field` configured. Parallel to `sort_orders`.
336-
/// Used to build `ListingOptions.with_file_sort_order(...)` so DataFusion
337-
/// advertises `output_ordering` from the scan and enables the
338-
/// `sort_prefix` optimization on TopK / SortPreservingMerge.
336+
///
337+
/// Two consumers today:
338+
/// - Vanilla path: `ListingOptions.with_file_sort_order(...)` so DataFusion advertises
339+
/// `output_ordering` from the scan and the `sort_prefix` optimization fires on
340+
/// TopK / SortPreservingMerge.
341+
/// - Indexed path (`indexed_executor`): when the query's leading ORDER BY runs counter
342+
/// to catalog-snapshot order, the per-shard segment iteration is reversed so a TopK
343+
/// above us can prune via parquet page statistics.
339344
pub sort_fields: Vec<String>,
340345
/// Index sort directions per field — values: `"asc"` or `"desc"`.
341346
/// Parallel to `sort_fields`. Sourced from `index.sort.order`.

sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs

Lines changed: 303 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,8 @@ pub async fn execute_indexed_query(
176176
table_path: shard_view.table_path.clone(),
177177
object_metas: shard_view.object_metas.clone(),
178178
writer_generations: shard_view.writer_generations.clone(),
179+
sort_fields: shard_view.sort_fields.clone(),
180+
sort_orders: shard_view.sort_orders.clone(),
179181
query_context: crate::query_tracker::QueryTrackingContext::new(context_id, runtime.runtime_env.memory_pool.clone(), crate::query_tracker::QueryType::Shard),
180182
table_name: table_name.clone(),
181183
indexed_config: None, // derive classification from tree
@@ -199,6 +201,101 @@ pub async fn execute_indexed_query(
199201

200202
// ── Helpers ───────────────────────────────────────────────────────────
201203

204+
/// Result of walking a logical plan looking for the leading top-of-plan ORDER BY.
205+
///
206+
/// `column` is the bare column name (no qualifier — we compare against `index.sort.field`
207+
/// which is also unqualified). `descending` is `true` for `ORDER BY x DESC`.
208+
#[derive(Debug, Clone, PartialEq, Eq)]
209+
pub(crate) struct TopSort {
210+
pub column: String,
211+
pub descending: bool,
212+
}
213+
214+
/// Walk through the top of a logical plan to find the leading sort expression.
215+
///
216+
/// Descends through `Projection`, `Limit`, `SubqueryAlias`, `Distinct`, and `Filter` —
217+
/// nodes that don't reorder rows or rewrite sort key columns. On reaching
218+
/// `LogicalPlan::Sort`, returns the leading sort key's bare column name and direction.
219+
/// Returns `None` if the plan has no top-level Sort, or if the first Sort key is not
220+
/// a plain `Expr::Column` (e.g. `ORDER BY lower(x)` — we can't claim catalog monotonicity
221+
/// after a function).
222+
pub(crate) fn analyze_top_sort(plan: &datafusion::logical_expr::LogicalPlan) -> Option<TopSort> {
223+
use datafusion::logical_expr::{Expr, LogicalPlan};
224+
let mut current = plan;
225+
loop {
226+
match current {
227+
LogicalPlan::Sort(s) => {
228+
let leading = s.expr.first()?;
229+
let col = match &leading.expr {
230+
Expr::Column(c) => c.name.clone(),
231+
_ => return None,
232+
};
233+
return Some(TopSort {
234+
column: col,
235+
descending: !leading.asc,
236+
});
237+
}
238+
LogicalPlan::Projection(p) => current = p.input.as_ref(),
239+
LogicalPlan::Limit(l) => current = l.input.as_ref(),
240+
LogicalPlan::SubqueryAlias(a) => current = a.input.as_ref(),
241+
LogicalPlan::Distinct(d) => match d {
242+
datafusion::logical_expr::Distinct::All(input) => current = input.as_ref(),
243+
datafusion::logical_expr::Distinct::On(on) => current = on.input.as_ref(),
244+
},
245+
LogicalPlan::Filter(f) => current = f.input.as_ref(),
246+
_ => return None,
247+
}
248+
}
249+
}
250+
251+
/// Decide whether to flip segment iteration order for the indexed scan.
252+
///
253+
/// Returns `true` iff the catalog has a sort declaration, the query has a top-level
254+
/// ORDER BY whose leading key matches the catalog's leading sort field by name, and
255+
/// the query's direction is the **opposite** of the catalog's. In that case the
256+
/// segments — laid down newest-last by the writer — are in reverse order from what
257+
/// the query wants, so iterating them tail-first feeds the largest values to a
258+
/// `TopK` first and parquet page stats prune the rest.
259+
///
260+
/// All comparisons are case-sensitive on the field name (matching PR #22041's
261+
/// `Column::from_name`). Direction comparison is case-insensitive on `"asc"`/`"desc"`.
262+
pub(crate) fn should_reverse_segments(
263+
top_sort: Option<&TopSort>,
264+
sort_fields: &[String],
265+
sort_orders: &[String],
266+
) -> bool {
267+
let Some(top) = top_sort else { return false };
268+
let Some(catalog_field) = sort_fields.first() else { return false };
269+
let Some(catalog_order) = sort_orders.first() else { return false };
270+
if top.column != *catalog_field {
271+
return false;
272+
}
273+
let catalog_descending = catalog_order.eq_ignore_ascii_case("desc");
274+
top.descending != catalog_descending
275+
}
276+
277+
/// Reverse the iteration order of `segments` in place. Per-segment `global_base`
278+
/// values are deliberately **left untouched** so each segment keeps its
279+
/// catalog-order shard-global row ID space.
280+
///
281+
/// Why not recompute global_base after reversing?
282+
///
283+
/// `global_base` is the additive offset used to compute the QTF `__row_id__`:
284+
/// `id = segment.global_base + position`. The fetch phase (`api::fetch_by_row_ids`)
285+
/// rebuilds segments via `build_segments` against `ShardView.object_metas` (always
286+
/// in catalog order) and reverses the mapping back via `partition_point` on
287+
/// `global_base`. For the round trip to hold, both phases must agree on each
288+
/// segment's `global_base` — and the only way to guarantee that without changing
289+
/// the fetch path is to keep query-phase `global_base` at its catalog-order value.
290+
///
291+
/// Reversing only the iteration order — i.e. the order in which chunks/RGs are
292+
/// scheduled by `compute_assignments` — is the *whole* point of this optimization
293+
/// (newest-segment-first feeds TopK earlier and lets page stats prune older
294+
/// segments). The `global_base` values are unrelated to that pruning win.
295+
pub(crate) fn reverse_segment_iteration_order(segments: &mut [SegmentFileInfo]) {
296+
segments.reverse();
297+
}
298+
202299
/// Collect all `Predicate(expr)` leaves in DFS order. Used by the
203300
/// dispatcher to build a per-leaf `PruningPredicate` cache keyed by
204301
/// `Arc::as_ptr` identity.
@@ -422,6 +519,191 @@ mod tests {
422519
]);
423520
assert!(extract_single_collector_residual(&tree).is_none());
424521
}
522+
523+
// ── analyze_top_sort / should_reverse_segments ────────────────────
524+
525+
fn build_logical_plan(sql: &str) -> datafusion::logical_expr::LogicalPlan {
526+
use datafusion::execution::SessionStateBuilder;
527+
use datafusion::execution::context::SessionContext;
528+
let state = SessionStateBuilder::new().with_default_features().build();
529+
let ctx = SessionContext::new_with_state(state);
530+
let schema = Arc::new(Schema::new(vec![
531+
Field::new("id", DataType::Int64, false),
532+
Field::new("ts", DataType::Int64, false),
533+
Field::new("v", DataType::Int32, false),
534+
]));
535+
ctx.register_batch(
536+
"t",
537+
datafusion::arrow::record_batch::RecordBatch::new_empty(schema),
538+
)
539+
.expect("register_batch");
540+
let rt = tokio::runtime::Builder::new_current_thread()
541+
.enable_all()
542+
.build()
543+
.unwrap();
544+
let df = rt.block_on(ctx.sql(sql)).expect("sql");
545+
df.into_unoptimized_plan()
546+
}
547+
548+
#[test]
549+
fn analyze_top_sort_finds_leading_desc() {
550+
let plan = build_logical_plan("SELECT id FROM t ORDER BY id DESC");
551+
let ts = analyze_top_sort(&plan).expect("expected Sort");
552+
assert_eq!(ts.column, "id");
553+
assert!(ts.descending);
554+
}
555+
556+
#[test]
557+
fn analyze_top_sort_descends_through_limit() {
558+
let plan = build_logical_plan("SELECT id FROM t ORDER BY id ASC LIMIT 10");
559+
let ts = analyze_top_sort(&plan).expect("expected Sort");
560+
assert_eq!(ts.column, "id");
561+
assert!(!ts.descending);
562+
}
563+
564+
#[test]
565+
fn analyze_top_sort_descends_through_projection() {
566+
let plan = build_logical_plan("SELECT v FROM t ORDER BY ts DESC");
567+
let ts = analyze_top_sort(&plan).expect("expected Sort");
568+
assert_eq!(ts.column, "ts");
569+
assert!(ts.descending);
570+
}
571+
572+
#[test]
573+
fn analyze_top_sort_returns_none_when_no_sort() {
574+
let plan = build_logical_plan("SELECT id FROM t");
575+
assert!(analyze_top_sort(&plan).is_none());
576+
}
577+
578+
#[test]
579+
fn analyze_top_sort_returns_none_for_function_sort_key() {
580+
// `ORDER BY abs(v)` — leading sort key isn't a plain column. Catalog monotonicity
581+
// doesn't transfer through a function, so we conservatively decline to reverse.
582+
let plan = build_logical_plan("SELECT id FROM t ORDER BY abs(v)");
583+
assert!(analyze_top_sort(&plan).is_none());
584+
}
585+
586+
#[test]
587+
fn should_reverse_segments_matches_leading_field_opposite_direction() {
588+
let top = TopSort { column: "id".to_string(), descending: true };
589+
let fields = vec!["id".to_string()];
590+
let orders = vec!["asc".to_string()];
591+
assert!(should_reverse_segments(Some(&top), &fields, &orders));
592+
}
593+
594+
#[test]
595+
fn should_reverse_segments_matches_leading_field_same_direction() {
596+
let top = TopSort { column: "id".to_string(), descending: false };
597+
let fields = vec!["id".to_string()];
598+
let orders = vec!["asc".to_string()];
599+
assert!(!should_reverse_segments(Some(&top), &fields, &orders));
600+
}
601+
602+
#[test]
603+
fn should_reverse_segments_catalog_desc_query_asc() {
604+
let top = TopSort { column: "id".to_string(), descending: false };
605+
let fields = vec!["id".to_string()];
606+
let orders = vec!["desc".to_string()];
607+
assert!(should_reverse_segments(Some(&top), &fields, &orders));
608+
}
609+
610+
#[test]
611+
fn should_reverse_segments_no_query_sort() {
612+
let fields = vec!["id".to_string()];
613+
let orders = vec!["asc".to_string()];
614+
assert!(!should_reverse_segments(None, &fields, &orders));
615+
}
616+
617+
#[test]
618+
fn should_reverse_segments_no_catalog_sort() {
619+
let top = TopSort { column: "id".to_string(), descending: true };
620+
assert!(!should_reverse_segments(Some(&top), &[], &[]));
621+
}
622+
623+
#[test]
624+
fn should_reverse_segments_query_sort_on_non_leading_catalog_field() {
625+
// Catalog: [a ASC, b ASC]; query: ORDER BY b DESC. Segments are monotonic on `a`
626+
// (the leading key), not `b`. Reversing won't help — decline.
627+
let top = TopSort { column: "b".to_string(), descending: true };
628+
let fields = vec!["a".to_string(), "b".to_string()];
629+
let orders = vec!["asc".to_string(), "asc".to_string()];
630+
assert!(!should_reverse_segments(Some(&top), &fields, &orders));
631+
}
632+
633+
#[test]
634+
fn should_reverse_segments_field_name_case_sensitive() {
635+
// Match PR #22041 — `Column::from_name` is case-sensitive. If casing differs,
636+
// we don't claim the catalog ordering applies. Safe default: no reversal.
637+
let top = TopSort { column: "ID".to_string(), descending: true };
638+
let fields = vec!["id".to_string()];
639+
let orders = vec!["asc".to_string()];
640+
assert!(!should_reverse_segments(Some(&top), &fields, &orders));
641+
}
642+
643+
// ── reverse_segment_iteration_order ───────────────────────────────
644+
645+
fn dummy_segment(max_doc: i64, global_base: u64) -> SegmentFileInfo {
646+
use datafusion::parquet::file::metadata::{FileMetaData, ParquetMetaData};
647+
// Build a minimal ParquetMetaData. We never read it back in these tests.
648+
let schema = std::sync::Arc::new(
649+
datafusion::parquet::schema::types::SchemaDescriptor::new(
650+
std::sync::Arc::new(
651+
datafusion::parquet::schema::types::Type::group_type_builder("schema")
652+
.build()
653+
.unwrap(),
654+
),
655+
),
656+
);
657+
let file_meta = FileMetaData::new(0, 0, None, None, schema, None);
658+
let pq_meta = ParquetMetaData::new(file_meta, vec![]);
659+
SegmentFileInfo {
660+
writer_generation: global_base as i64 + 1, // arbitrary, just to vary
661+
max_doc,
662+
object_path: object_store::path::Path::from(format!("seg-{}.parquet", global_base)),
663+
parquet_size: 0,
664+
row_groups: vec![],
665+
metadata: std::sync::Arc::new(pq_meta),
666+
global_base,
667+
}
668+
}
669+
670+
#[test]
671+
fn reverse_segments_preserves_global_base() {
672+
// Original: A(max_doc=10, base=0), B(max_doc=20, base=10), C(max_doc=30, base=30).
673+
// Reversal must keep each segment's original `global_base` intact so QTF row IDs
674+
// emitted in query phase remain interpretable by the fetch phase (which always
675+
// computes catalog-order bases).
676+
let mut segs = vec![
677+
dummy_segment(10, 0),
678+
dummy_segment(20, 10),
679+
dummy_segment(30, 30),
680+
];
681+
reverse_segment_iteration_order(&mut segs);
682+
assert_eq!(segs.len(), 3);
683+
// New iteration order: C, B, A.
684+
assert_eq!(segs[0].max_doc, 30);
685+
assert_eq!(segs[0].global_base, 30); // C's catalog base, unchanged.
686+
assert_eq!(segs[1].max_doc, 20);
687+
assert_eq!(segs[1].global_base, 10); // B's catalog base, unchanged.
688+
assert_eq!(segs[2].max_doc, 10);
689+
assert_eq!(segs[2].global_base, 0); // A's catalog base, unchanged.
690+
}
691+
692+
#[test]
693+
fn reverse_segments_empty_is_noop() {
694+
let mut segs: Vec<SegmentFileInfo> = vec![];
695+
reverse_segment_iteration_order(&mut segs);
696+
assert!(segs.is_empty());
697+
}
698+
699+
#[test]
700+
fn reverse_segments_single_keeps_its_base() {
701+
let mut segs = vec![dummy_segment(42, 7)];
702+
reverse_segment_iteration_order(&mut segs);
703+
assert_eq!(segs.len(), 1);
704+
assert_eq!(segs[0].global_base, 7);
705+
assert_eq!(segs[0].max_doc, 42);
706+
}
425707
}
426708

427709
/// Instruction-based indexed execution path. Consumes a pre-configured SessionContextHandle
@@ -508,6 +790,8 @@ async unsafe fn execute_indexed_with_context_inner(
508790
let table_path = handle.table_path;
509791
let object_metas = handle.object_metas;
510792
let writer_generations = handle.writer_generations;
793+
let sort_fields = handle.sort_fields;
794+
let sort_orders = handle.sort_orders;
511795
let query_context = handle.query_context;
512796
let io_handle = handle.io_handle;
513797
// Extract context_id early so it can be captured by the per-segment closures
@@ -550,6 +834,25 @@ async unsafe fn execute_indexed_with_context_inner(
550834
.map_err(|e| DataFusionError::Execution(format!("decode substrait: {}", e)))?;
551835
let logical_plan = from_substrait_plan(&ctx.state(), &plan).await?;
552836

837+
// Sort-aware segment iteration. Mirror of `ContextIndexSearcher.shouldUseTimeSeriesDescSortOptimization`
838+
// for the indexed-parquet path. When the index has `index.sort.field` and the query's leading
839+
// ORDER BY runs counter to the catalog's stored direction, reverse the segment vector so a
840+
// TopK above us pulls the highest-priority segment first and parquet page stats prune the rest.
841+
//
842+
// QTF safe: `reverse_segment_iteration_order` deliberately does NOT recompute `global_base`
843+
// — each segment retains its catalog-order base, so the row IDs query phase emits are still
844+
// interpretable by `api::fetch_by_row_ids` (which builds its own segments from
845+
// `ShardView.object_metas` in catalog order).
846+
let mut segments = segments;
847+
if should_reverse_segments(analyze_top_sort(&logical_plan).as_ref(), &sort_fields, &sort_orders) {
848+
log_debug!(
849+
"indexed_executor: reversing segment iteration (catalog leading sort={:?} {:?}, query opposite)",
850+
sort_fields.first(),
851+
sort_orders.first()
852+
);
853+
reverse_segment_iteration_order(&mut segments);
854+
}
855+
553856
let emit_row_ids = requests_row_ids;
554857
let filter_expr = extract_filter_expr(&logical_plan);
555858
let extraction = match filter_expr {

sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ mod qtf_fetch_phase;
4949
mod row_id_emission;
5050
mod row_id_strategies;
5151
mod schema_drift;
52+
mod sort_reverse_row_id;
5253
mod streaming_at_scale;
5354

5455
// ── Test fixture: parquet table with 16 rows ────────────────────────

0 commit comments

Comments
 (0)