fix: sort extracted values before training a JSON scalar index#7819
Open
chakshu-dhannawat wants to merge 1 commit into
Open
fix: sort extracted values before training a JSON scalar index#7819chakshu-dhannawat wants to merge 1 commit into
chakshu-dhannawat wants to merge 1 commit into
Conversation
A JSON-path scalar index (target_index_type="btree") returned wrong results for float paths whose values are not exactly representable, e.g. range queries returned empty and equality missed the value. Reported in lance-format#7485. Root cause is ordering, not float precision. The scalar-index framework pre-sorts the input rows by the raw indexed column (the JSONB blob), but a JSON index re-extracts a different value: the value at `path`. The order of those extracted values does not match the raw-column order, so the stream handed to the btree trainer is not sorted by the value the btree actually stores. train_btree_index requires globally sorted input (analyze_batch takes the page min as values[0] and max as values[last]), so the page min/max end up wrong and range/equality scans skip the right pages. A plain (non-JSON) btree column is unaffected because its sort key is the indexed value itself. Re-sort the converted (value, row_id) stream by value in the JSON training path before handing it to the target trainer, restoring the sorted-input invariant the trainer relies on. The test values happened to make exact floats look fine (single-page in-page filtering masked it) while non-exact data exposed it, which is why the report reads as a float bug. Added a regression test that trains a JSON btree index over floats laid out so the extracted order differs from the row order, then asserts range and equality queries return the correct rows. It returns empty results without the fix and passes with it. Fixes lance-format#7485
github-actions
Bot
added
A-index
📝 WalkthroughWalkthroughJSON index training now sorts converted JSON-path values before building the target scalar index. A regression test covers unsorted float values and verifies btree range and equality query results. ChangesJSON index ordering
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant JsonIndexPlugin
participant SortExec
participant TargetScalarIndex
JsonIndexPlugin->>SortExec: Sort converted JSON-path values
SortExec-->>JsonIndexPlugin: Return sorted value and row_id stream
JsonIndexPlugin->>TargetScalarIndex: Train btree index
TargetScalarIndex-->>JsonIndexPlugin: Serve range and equality queries
Possibly related PRs
Suggested reviewers: 🚥 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: 1
🤖 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 1020-1138: Extend test_json_btree_index_unsorted_extracted_values
with a NULL JSON value and assert the indexed search behavior preserves
nulls-first ordering without affecting the existing non-null range and equality
results. Update the input batch and expected row mappings consistently, keeping
the regression focused on extracted-value sorting.
🪄 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: 75a484f9-8643-489b-a5f8-66b1009ce71b
📒 Files selected for processing (1)
rust/lance-index/src/scalar/json.rs
Comment on lines
+1020
to
+1138
|
|
||
| // Regression for https://github.com/lancedb/lance/issues/7485 | ||
| // | ||
| // A JSON-path btree index returned wrong results for float values whose row | ||
| // order (after the framework's sort on the raw JSON column) does not match | ||
| // the order of the extracted value. The btree trainer assumes its input is | ||
| // sorted by the stored value, so without re-sorting the extracted values the | ||
| // page min/max stats are wrong and range/equality queries miss rows. | ||
| #[tokio::test] | ||
| async fn test_json_btree_index_unsorted_extracted_values() { | ||
| use crate::metrics::NoOpMetricsCollector; | ||
| use crate::progress::noop_progress; | ||
| use crate::scalar::{SargableQuery, SearchResult, lance_format::LanceIndexStore}; | ||
| use arrow_array::{LargeBinaryArray, UInt64Array}; | ||
| use datafusion::physical_plan::stream::RecordBatchStreamAdapter; | ||
| use datafusion_common::ScalarValue; | ||
| use futures::stream; | ||
| use lance_core::cache::LanceCache; | ||
| use lance_core::utils::tempfile::TempObjDir; | ||
| use lance_io::object_store::ObjectStore; | ||
| use lance_select::RowAddrTreeMap; | ||
| use std::ops::Bound; | ||
|
|
||
| // Non-exact floats deliberately laid out so that the row order does not | ||
| // match ascending value order (10.5, 40.1, -3.2 sorts to -3.2, 10.5, 40.1). | ||
| let json_data = [ | ||
| r#"{"latitude": 10.5}"#, | ||
| r#"{"latitude": 40.1}"#, | ||
| r#"{"latitude": -3.2}"#, | ||
| ]; | ||
| let jsonb: Vec<Vec<u8>> = json_data | ||
| .iter() | ||
| .map(|s| s.parse::<jsonb::OwnedJsonb>().unwrap().to_vec()) | ||
| .collect(); | ||
|
|
||
| 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| Some(v.as_slice())).collect::<Vec<_>>(), | ||
| )) as ArrayRef, | ||
| Arc::new(UInt64Array::from(vec![0u64, 1, 2])) as ArrayRef, | ||
| ], | ||
| ) | ||
| .unwrap(); | ||
| let data = Box::pin(RecordBatchStreamAdapter::new( | ||
| schema.clone(), | ||
| stream::iter(vec![Ok(batch)]), | ||
| )) as SendableRecordBatchStream; | ||
|
|
||
| let tmpdir = TempObjDir::default(); | ||
| let store = Arc::new(LanceIndexStore::new( | ||
| Arc::new(ObjectStore::local()), | ||
| tmpdir.clone(), | ||
| Arc::new(LanceCache::no_cache()), | ||
| )); | ||
|
|
||
| let registry = IndexPluginRegistry::with_default_plugins(); | ||
| let plugin = registry.get_plugin_by_name("json").unwrap(); | ||
| let trainer = plugin.basic_trainer().unwrap(); | ||
| let params = r#"{"target_index_type":"btree","path":"latitude"}"#; | ||
| let request = trainer | ||
| .new_training_request( | ||
| params, | ||
| &Field::new(VALUE_COLUMN_NAME, DataType::LargeBinary, true), | ||
| ) | ||
| .unwrap(); | ||
| let created = trainer | ||
| .train_index(data, store.as_ref(), request, None, noop_progress()) | ||
| .await | ||
| .unwrap(); | ||
|
|
||
| let index = plugin | ||
| .load_index( | ||
| store.clone(), | ||
| &created.index_details, | ||
| None, | ||
| &LanceCache::no_cache(), | ||
| ) | ||
| .await | ||
| .unwrap(); | ||
|
|
||
| let search = |q: SargableQuery| { | ||
| let query = JsonQuery::new(Arc::new(q) as Arc<dyn AnyQuery>, "latitude".to_string()); | ||
| let index = index.clone(); | ||
| async move { index.search(&query, &NoOpMetricsCollector).await.unwrap() } | ||
| }; | ||
|
|
||
| // Range: latitude > 0 -> rows 0 (10.5) and 1 (40.1). | ||
| assert_eq!( | ||
| search(SargableQuery::Range( | ||
| Bound::Excluded(ScalarValue::Float64(Some(0.0))), | ||
| Bound::Unbounded, | ||
| )) | ||
| .await, | ||
| SearchResult::exact(RowAddrTreeMap::from_iter([0u64, 1])), | ||
| ); | ||
|
|
||
| // Equality on a non-exact float: latitude = 40.1 -> row 1. | ||
| assert_eq!( | ||
| search(SargableQuery::Equals(ScalarValue::Float64(Some(40.1)))).await, | ||
| SearchResult::exact(RowAddrTreeMap::from_iter([1u64])), | ||
| ); | ||
|
|
||
| // Lower bound that must exclude the smallest value: | ||
| // latitude >= 10.5 -> rows 0 (10.5) and 1 (40.1). | ||
| assert_eq!( | ||
| search(SargableQuery::Range( | ||
| Bound::Included(ScalarValue::Float64(Some(10.5))), | ||
| Bound::Unbounded, | ||
| )) | ||
| .await, | ||
| SearchResult::exact(RowAddrTreeMap::from_iter([0u64, 1])), | ||
| ); | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
Consider adding a NULL-value case to the regression test.
The test only covers non-null float values in a single fragment/batch. Per coding guidelines for index tests, NULL edge cases should be included, especially since the sort uses nulls_first: true and this is exactly the kind of ordering detail a regression test for this bug should also pin down.
As per coding guidelines, "Include multi-fragment dataset scenarios, NULL edge cases in index tests, vector-index recall assertions of at least 0.5...".
🤖 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 1020 - 1138, Extend
test_json_btree_index_unsorted_extracted_values with a NULL JSON value and
assert the indexed search behavior preserves nulls-first ordering without
affecting the existing non-null range and equality results. Update the input
batch and expected row mappings consistently, keeping the regression focused on
extracted-value sorting.
Source: Coding guidelines
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
A JSON-path scalar index (
IndexConfig(index_type="json", parameters={"target_index_type": "btree", "path": ...})) over a float path returns incorrect results when the values are not exactly representable infloat64: ranges come back empty and equality misses the value. Reported in #7485.Repro from the issue:
Root cause
It's an ordering bug, not float precision. The scalar-index framework pre-sorts the input rows by the raw indexed column (the JSONB blob), but a JSON index re-extracts a different value, the value at
path. The extracted values are not in the same order as the raw-column sort, so the stream handed to the btree trainer is not sorted by the value the btree stores.train_btree_indexrequires globally sorted input:analyze_batchrecords each page'sminasvalues[0]andmaxasvalues[values.len()-1]. With unsorted input those stats are wrong, so range/equality page selection skips the pages that actually hold the matching rows. A plain (non-JSON)float64btree column is unaffected because its sort key already is the indexed value.I confirmed the mechanism by probing the failing index: for
[10.5, 40.1, -3.2]the single page reportsmax = -3.2(the last row), so every> xwithx >= -3.2returns empty. The "only non-exact floats break" appearance is incidental: the exact-float cases in the report were already in ascending order (or single-page in-page filtering masked the bad stats), while the non-exact data happened to be unsorted.Fix
Re-sort the converted
(value, row_id)stream by value in the JSON training path before handing it to the target trainer, restoring the sorted-input invariant. This mirrors what the plain-column path already gets. Sorting is harmless for target index types that don't require it.Testing
Added
test_json_btree_index_unsorted_extracted_valuesinrust/lance-index/src/scalar/json.rs: it trains a JSON btree index over floats whose extracted order differs from the row order, then asserts> 0,= 40.1, and>= 10.5return the correct rows. The test returns empty results on the unfixed code and passes with the fix.cargo fmtandcargo clippy -p lance-index --tests -- -D warningsare clean; the existing scalar tests still pass.Summary by CodeRabbit