fix: re-sort extracted JSON path values before training value-ordered indexes#7771
fix: re-sort extracted JSON path values before training value-ordered indexes#7771sohumt123 wants to merge 2 commits into
Conversation
… indexes A JSON-path scalar index whose target is a value-ordered index (e.g. btree) returned wrong results for non-exactly-representable float64 values: range predicates came back empty and equality missed rows. The training data reaching the JSON plugin is sorted by the raw JSONB bytes of the source column (that is all the scanner can order on), then the requested path is extracted and converted to a typed value column. The byte order of the JSON encoding does not match the numeric order of the extracted values for floats such as 40.1 or -3.2, so the typed column arrives out of order. The btree trainer assumes value ordering and derives each page's min/max from the first and last rows, so it recorded inverted statistics (min > max) and silently pruned matching pages. Re-sort the extracted value column ascending before handing it to the target trainer, gated on the target requiring TrainingOrdering::Values so ordering-agnostic targets (e.g. bitmap) are unaffected. Closes lance-format#7485
github-actions
Bot
added
A-index
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughJSON index training now sorts extracted values when the target trainer requires value ordering. New regression coverage validates JSON-path float btree indexes with non-exact floating-point values and adversarial feed orders. ChangesJSON index ordering
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant JsonIndexPlugin
participant SortExec
participant TargetTrainer
JsonIndexPlugin->>SortExec: Sort extracted values when ordering is required
SortExec->>TargetTrainer: Pass ordered training stream
JsonIndexPlugin->>TargetTrainer: Pass unchanged stream otherwise
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/lance-index/src/scalar/json.rs`:
- Around line 769-775: Track the bounded-memory limitation in the value-ordered
JSON index training path: convert_stream_by_type materializes all converted
batches before execute_plan runs SortExec, so use_spilling only constrains
sorting. Update the implementation or follow-up tracking so this path does not
imply end-to-end bounded memory, while preserving the existing sort behavior.
- Around line 1131-1184: Replace the manual loops over cases_a and cases_b with
rstest parameterized test cases, using named #[case::{name}(...)] entries for
each query and expected row set. Keep the existing
train_and_search_json_float_index invocation and exact SearchResult assertions,
including distinct case coverage and readable failure context.
- Around line 754-775: Update the sorting setup in the TrainingOrdering::Values
branch to pass a LexOrdering to SortExec::new, constructing it with
LexOrdering::new(vec![sort_expr]) instead of relying on [sort_expr].into().
Ensure the required LexOrdering import is available.
- Around line 1050-1063: Replace the manual schema and RecordBatch::try_new
construction in this test with arrow_array::record_batch!, building the
LargeBinary value column from jsonb and the UInt64 row-ID column from row_ids.
Preserve the existing column names, types, nullability, and row ordering.
- Around line 754-778: Update the converted_stream ordering logic to match
explicitly on target_request.criteria().ordering, including distinct branches
for Values, Addresses, and None. Preserve value sorting for
TrainingOrdering::Values, explicitly handle TrainingOrdering::Addresses
according to the address-ordering contract used by BloomFilter and ZoneMap, and
retain the current unsorted behavior only for None.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 6a77a916-2423-4e89-afe4-62fa1b672f8a
📒 Files selected for processing (1)
rust/lance-index/src/scalar/json.rs
| let converted_stream = if target_request.criteria().ordering == TrainingOrdering::Values { | ||
| let input = Arc::new(OneShotExec::new(converted_stream)); | ||
| let value_column_idx = input | ||
| .schema() | ||
| .column_with_name(VALUE_COLUMN_NAME) | ||
| .expect_ok()? | ||
| .0; | ||
| let sort_expr = PhysicalSortExpr { | ||
| expr: Arc::new(Column::new(VALUE_COLUMN_NAME, value_column_idx)), | ||
| options: SortOptions { | ||
| descending: false, | ||
| nulls_first: true, | ||
| }, | ||
| }; | ||
| let sorted = Arc::new(SortExec::new([sort_expr].into(), input)); | ||
| execute_plan( | ||
| sorted, | ||
| LanceExecutionOptions { | ||
| use_spilling: true, | ||
| ..Default::default() | ||
| }, | ||
| )? |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Does the DataFusion LexOrdering struct (post the "Convert LexOrdering type to struct" change) implement From<[PhysicalSortExpr; 1]> or From<Vec<PhysicalSortExpr>>, and what is the recommended way to build a single-column LexOrdering?
💡 Result:
The DataFusion LexOrdering struct does not implement From<[PhysicalSortExpr; 1]> or From<Vec> [1][2]. Following the change that converted LexOrdering from a type alias to a struct, the primary and recommended way to construct a LexOrdering is by using the LexOrdering::new(exprs) associated function [1]. LexOrdering::new takes an impl IntoIterator<Item = PhysicalSortExpr> and returns an Option [1]. It returns None if the provided iterator is empty, as LexOrdering is designed to be non-degenerate (i.e., it must contain at least one element) [1]. To build a single-column LexOrdering, you can pass a single-element collection (such as a vector or an array) to this function: // Recommended way to build a single-column LexOrdering let sort_expr = PhysicalSortExpr {... }; let ordering = LexOrdering::new(vec![sort_expr]).expect("LexOrdering requires at least one expression"); Alternatively, you can initialize a LexOrdering with multiple expressions, which is similarly handled via the same constructor [1]. Because LexOrdering is an opaque struct with private fields, you cannot instantiate it directly (e.g., via struct literal syntax) [1]. top_results: 1, 2, 3, 8
Citations:
- 1: https://docs.rs/datafusion/latest/datafusion/physical_expr/struct.LexOrdering.html
- 2: https://docs.rs/datafusion/latest/datafusion/physical_expr/index.html
🏁 Script executed:
set -euo pipefail
printf 'Pinned DataFusion versions in workspace:\n'
rg -n '^name = "datafusion(-physical-expr)?\"|^version = ' Cargo.lock Cargo.toml rust/**/Cargo.toml 2>/dev/null | sed -n '1,120p'
printf '\nCall site context:\n'
sed -n '744,782p' rust/lance-index/src/scalar/json.rsRepository: lance-format/lance
Length of output: 197
🏁 Script executed:
set -euo pipefail
echo 'Pinned DataFusion versions:'
git ls-files '**/Cargo.lock' '**/Cargo.toml' | sed -n '1,120p'
echo
rg -n '^(name = "datafusion(-physical-expr)?"|version = )' Cargo.lock rust/**/Cargo.toml Cargo.toml 2>/dev/null || true
echo
echo 'Relevant code:'
sed -n '744,782p' rust/lance-index/src/scalar/json.rsRepository: lance-format/lance
Length of output: 34765
🌐 Web query:
DataFusion 53.1.0 LexOrdering From<[PhysicalSortExpr; 1]> From<Vec<PhysicalSortExpr>> SortExec::new docs
💡 Result:
In DataFusion 53.1.0, the structures LexOrdering and SortExec serve distinct roles in managing and executing sorted data streams [1][2]. LexOrdering LexOrdering represents a lexicographical ordering, consisting of a vector of PhysicalSortExpr objects [1][3]. It enforces two primary invariants: 1. It is non-degenerate, containing at least one element [1][3]. 2. It is duplicate-free, meaning it does not contain multiple entries for the same column [1][3]. Because LexOrdering does not implement From<[PhysicalSortExpr; 1]> or From<Vec> directly, you should use its new constructor, which accepts an impl IntoIterator<Item = PhysicalSortExpr> and returns an Option (returning None if the input is empty) [1][3]. SortExec::new SortExec is the execution plan responsible for sorting data. Its constructor is defined as follows [2][4]: pub fn new(expr: LexOrdering, input: Arc) -> Self This function is infallible and creates a new sort execution plan that produces a single, sorted output partition [2][5]. When called, it automatically computes plan properties such as the common sort prefix [5][4]. Additional modifications, such as setting a fetch limit or preserving input partitioning, are typically handled via builder-style methods like.with_fetch(fetch) and.with_preserve_partitioning(true) [5][6][4].
Citations:
- 1: https://docs.rs/datafusion/latest/datafusion/physical_expr/struct.LexOrdering.html
- 2: https://docs.rs/datafusion/latest/datafusion/physical_plan/sorts/sort/struct.SortExec.html
- 3: https://github.com/apache/datafusion/blob/6b134dda/datafusion/physical-expr-common/src/sort_expr.rs
- 4: https://apache.googlesource.com/datafusion/+/refs/tags/48.0.0/datafusion/physical-plan/src/sorts/sort.rs
- 5: https://github.com/apache/datafusion/blob/6b134dda/datafusion/physical-plan/src/sorts/sort.rs
- 6: Clean up SortExec creation and add doc comments apache/datafusion#5889
Use LexOrdering::new(...) here. SortExec::new takes a LexOrdering, so [sort_expr].into() doesn’t match DataFusion 53.1.0; build the ordering with LexOrdering::new(vec![sort_expr]) instead.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance-index/src/scalar/json.rs` around lines 754 - 775, Update the
sorting setup in the TrainingOrdering::Values branch to pass a LexOrdering to
SortExec::new, constructing it with LexOrdering::new(vec![sort_expr]) instead of
relying on [sort_expr].into(). Ensure the required LexOrdering import is
available.
| let schema = Arc::new(Schema::new(vec![ | ||
| Field::new(VALUE_COLUMN_NAME, DataType::LargeBinary, true), | ||
| Field::new(ROW_ID, DataType::UInt64, false), | ||
| ])); | ||
| let batch = RecordBatch::try_new( | ||
| schema.clone(), | ||
| vec![ | ||
| Arc::new(LargeBinaryArray::from( | ||
| jsonb.iter().map(|v| v.as_deref()).collect::<Vec<_>>(), | ||
| )) as ArrayRef, | ||
| Arc::new(UInt64Array::from(row_ids)) as ArrayRef, | ||
| ], | ||
| ) | ||
| .unwrap(); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Does arrow-array's record_batch! macro support building a LargeBinary column from a Vec<&[u8]> variable in the arrow-rs version used by lance?
💡 Result:
Yes, the record_batch! macro in arrow-rs supports building a LargeBinary column from a variable (such as a Vec<&[u8]>) [1][2][3]. Following recent updates to the arrow-rs library (specifically in pull request #9522, merged in March 2026), the record_batch! macro was extended to accept both literal values and variables/expressions for array data, including support for LargeBinary types [1][4]. You can use it as follows: let large_binary_values = vec![b"some_data".as_slice]; let record_batch = record_batch!( ("my_col", LargeBinary, large_binary_values)).unwrap; In this syntax, the macro correctly handles the variable by utilizing LargeBinaryArray::from_vec [1][2][3]. As long as your version of arrow-rs includes these updates (which post-date March 2026), this functionality will be available [1][4].
Citations:
- 1: apache/arrow-rs@00ad7fc
- 2: https://arrow.apache.org/rust/src/arrow_array/record_batch.rs.html
- 3: https://github.com/apache/arrow-rs/blob/master/arrow-array/src/record_batch.rs
- 4: chore: extend record_batch macro to support variables and expressions apache/arrow-rs#9522
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant file and inspect the target range.
wc -l rust/lance-index/src/scalar/json.rs
sed -n '1015,1080p' rust/lance-index/src/scalar/json.rs
# Find any existing record_batch! usage in the repo for comparison.
rg -n "record_batch!\(" rust -g '*.rs' | head -n 40Repository: lance-format/lance
Length of output: 6921
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the surrounding function and test context.
sed -n '990,1085p' rust/lance-index/src/scalar/json.rs
# Check whether this file contains cfg(test) sections around the target lines.
rg -n "cfg\\(test\\)|record_batch!|RecordBatch::try_new|LargeBinaryArray::from" rust/lance-index/src/scalar/json.rsRepository: lance-format/lance
Length of output: 4492
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Show a few representative record_batch! usages, especially with non-primitive types.
sed -n '430,455p' rust/lance/src/dataset/tests/dataset_migrations.rs
sed -n '3490,3535p' rust/lance/src/dataset/write.rs
sed -n '3735,3765p' rust/lance/src/dataset/write.rsRepository: lance-format/lance
Length of output: 3694
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "record_batch!\\(.*LargeBinary|LargeBinary.*record_batch!|record_batch!\\(" rust -g '*.rs' | head -n 80Repository: lance-format/lance
Length of output: 6024
Use arrow_array::record_batch!() here
This test can build the LargeBinary and UInt64 columns directly from jsonb and row_ids, removing the manual Schema / RecordBatch::try_new boilerplate.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance-index/src/scalar/json.rs` around lines 1050 - 1063, Replace the
manual schema and RecordBatch::try_new construction in this test with
arrow_array::record_batch!, building the LargeBinary value column from jsonb and
the UInt64 row-ID column from row_ids. Preserve the existing column names,
types, nullability, and row ordering.
Source: Coding guidelines
Make the ordering decision in JsonIndexPlugin::train_index an explicit match on TrainingOrdering with distinct arms: - Values: re-sort the extracted value column by value (unchanged behavior). - Addresses: pass the stream through untouched. criteria() delegates the target's criteria verbatim, so the scanner already supplies address-ordered rows and the extraction/conversion pipeline preserves row order; re-sorting by value here would violate the Addresses contract. - None: leave unsorted. Document that use_spilling only bounds this SortExec, since the extraction and conversion stages already materialize all batches in memory upstream. Refactor the non-exact-float regression test to rstest parameterized #[case]s (repo test idiom), keeping the same train/search calls and exact SearchResult assertions and preserving the case A / case B feed-order asymmetry. Serialize the cases with a Mutex so their spilling sorts do not contend for the shared cached DataFusion memory pool, matching the original sequential behavior.
1 participant
Problem
A JSON-path scalar index (
index_type="json"withtarget_index_type="btree") over apa.json_()column returns incorrect filter results when the indexed path holds float64 values that are not exactly representable (e.g.40.1,-3.2). Range predicates come back empty and equality misses rows, while the same predicate over an un-indexed scan, an integer path, or exactly-representable floats is correct.Root cause
When training a JSON index, the scanner can only order the input on the raw JSONB bytes of the source column (
scan_training_datainrust/lance/src/index/scalar.rssorts withColumnOrdering::asc_nulls_firstover the binary column when the criteria ordering isValues). The JSON plugin then extracts the requested path and converts it to a typed value column, but hands that column to the target trainer without re-sorting it.For floats whose byte encoding order does not match their numeric order, the typed column arrives out of order. The btree trainer assumes value-ordering and takes each page's
min/maxfrom the first and last rows (btree.rs,min_val = value(0),max_val = value(num_rows - 1)), so it records inverted statistics (e.g.min = 40.1 > max = -3.2). Range pruning and equality lookups then discard pages that actually contain matches, producing silently wrong results.This is distinct from #7072 (path-based sub-parser routing); that fix corrected integer/equality routing but the float ordering problem remained.
Fix
In
JsonIndexPlugin::train_index, after the typed value stream is produced and before it is handed to the target trainer, re-sort it ascending by the value column when the target requires value-ordering (target_request.criteria().ordering == TrainingOrdering::Values). This reuses the existingSortExec/PhysicalSortExprpattern already used by the btree merge path (nulls_first: true,execute_planwithuse_spilling: true). Ordering-agnostic targets (e.g. bitmap) take the untouched stream, so they are unaffected.The change is confined to
rust/lance-index/src/scalar/json.rs.Testing
Added
test_json_float_btree_index_non_exact_floatsto the existing test module injson.rs. It trains a real JSON/btree index through the plugin trainer and searches the loaded index:[10.5, 40.1, -3.2]values, fed in raw-JSONB-byte order to mirrorscan_training_data. Asserts> 0 -> {0,1},>= 10.5 -> {0,1},= 40.1 -> {1},= 10.5 -> {0},< 100 -> {0,1,2}.[40.1 -> row0, -3.2 -> row1]fed unsorted, which fails pre-fix regardless of the exact JSONB byte encoding.Verified failing-before / passing-after with the test in place:
cargo fmt --all -- --checkandcargo clippy -p lance-index --tests -- -D warningsare clean.Notes
JsonTrainingRequest::criteria()toTrainingOrdering::None, which would let the upstream scanner skip the now-redundant raw-byte sort while this plugin does the value sort. That widens the blast radius, so I left it out.critical-fixlabel: it produces silent wrong query results rather than an error.Closes #7485
Summary by CodeRabbit