From e162373df43614e9b6e922eb0fe965262972097d Mon Sep 17 00:00:00 2001 From: Sohum Trivedi Date: Mon, 13 Jul 2026 22:48:14 -0700 Subject: [PATCH 1/2] fix: re-sort extracted JSON path values before training value-ordered 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 #7485 --- rust/lance-index/src/scalar/json.rs | 215 +++++++++++++++++++++++++++- 1 file changed, 210 insertions(+), 5 deletions(-) diff --git a/rust/lance-index/src/scalar/json.rs b/rust/lance-index/src/scalar/json.rs index 09678d6787e..dab66f4591f 100644 --- a/rust/lance-index/src/scalar/json.rs +++ b/rust/lance-index/src/scalar/json.rs @@ -8,22 +8,24 @@ use std::{ }; use arrow_array::{Array, LargeBinaryArray, RecordBatch, StructArray, UInt8Array}; -use arrow_schema::{DataType, Field, Field as ArrowField, Schema}; +use arrow_schema::{DataType, Field, Field as ArrowField, Schema, SortOptions}; use async_trait::async_trait; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::{ execution::SendableRecordBatchStream, - physical_plan::{ExecutionPlan, projection::ProjectionExec}, + physical_plan::{ExecutionPlan, projection::ProjectionExec, sorts::sort::SortExec}, }; use datafusion_common::{ScalarValue, config::ConfigOptions}; use datafusion_expr::{Expr, Operator, ScalarUDF}; use datafusion_physical_expr::{ - PhysicalExpr, ScalarFunctionExpr, + PhysicalExpr, PhysicalSortExpr, ScalarFunctionExpr, expressions::{Column, Literal}, }; use futures::StreamExt; use lance_core::deepsize::DeepSizeOf; -use lance_datafusion::exec::{LanceExecutionOptions, OneShotExec, get_session_context}; +use lance_datafusion::exec::{ + LanceExecutionOptions, OneShotExec, execute_plan, get_session_context, +}; use lance_datafusion::udf::json::JsonbType; use prost::Message; use roaring::RoaringBitmap; @@ -40,7 +42,8 @@ use crate::{ UpdateCriteria, expression::{IndexedExpression, ScalarIndexExpr, ScalarIndexSearch, ScalarQueryParser}, registry::{ - BasicTrainer, ScalarIndexPlugin, TrainingCriteria, TrainingRequest, VALUE_COLUMN_NAME, + BasicTrainer, ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest, + VALUE_COLUMN_NAME, }, }, }; @@ -739,6 +742,41 @@ impl BasicTrainer for JsonIndexPlugin { &Field::new("", inferred_type, true), )?; + // The training data reaching this plugin is sorted by the raw JSONB bytes of + // the source JSON column (that is what the scanner can order on), not by the + // extracted value. For values whose byte order does not match their numeric + // order (e.g. non-exactly-representable floats such as 40.1 or -3.2), the + // extracted column handed to the target trainer would be out of order. Targets + // that assume value ordering (e.g. btree, whose per-page min/max come from the + // first/last row) would then record inverted statistics and silently return + // wrong query results. Re-sort by the extracted value here so the value-ordering + // contract is honored. Ordering-agnostic targets (e.g. bitmap) are left as-is. + 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() + }, + )? + } else { + converted_stream + }; + let target_index = target_trainer .train_index( converted_stream, @@ -837,6 +875,7 @@ impl ScalarIndexPlugin for JsonIndexPlugin { #[cfg(test)] mod tests { use super::*; + use crate::scalar::SargableQuery; use arrow_array::{ArrayRef, RecordBatch}; use arrow_schema::{DataType, Field, Schema}; use std::sync::Arc; @@ -978,4 +1017,170 @@ mod tests { assert_eq!(inferred_type, DataType::Utf8); } + + /// Build a JSON-path btree index over the given `(json_doc, row_id)` pairs, + /// fed to the trainer in exactly the given order, then run `query` against it. + /// + /// This mirrors what the dataset layer does: it hands the JSON binary column to + /// the JSON plugin trainer. The caller controls the feed order so we can exercise + /// the case where the extracted values do not arrive in numeric order. + async fn train_and_search_json_float_index( + docs: &[(&str, u64)], + query: &SargableQuery, + ) -> SearchResult { + use crate::metrics::NoOpMetricsCollector; + use crate::progress::noop_progress; + use crate::registry::IndexPluginRegistry; + use crate::scalar::lance_format::LanceIndexStore; + use arrow_array::{LargeBinaryArray, UInt64Array}; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use futures::stream; + use lance_core::cache::LanceCache; + use lance_core::utils::tempfile::TempObjDir; + use lance_io::object_store::ObjectStore; + + let path = "$.latitude"; + + let jsonb: Vec>> = docs + .iter() + .map(|(doc, _)| Some(doc.parse::().unwrap().to_vec())) + .collect(); + let row_ids: Vec = docs.iter().map(|(_, id)| *id).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| v.as_deref()).collect::>(), + )) as ArrayRef, + Arc::new(UInt64Array::from(row_ids)) as ArrayRef, + ], + ) + .unwrap(); + let data_stream = Box::pin(RecordBatchStreamAdapter::new( + schema, + stream::iter(vec![Ok(batch)]), + )) as SendableRecordBatchStream; + + // Keep the temp dir alive for the duration of the test. + 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 = format!(r#"{{"target_index_type": "btree", "path": "{path}"}}"#); + let request = trainer + .new_training_request(¶ms, &Field::new("", DataType::LargeBinary, true)) + .unwrap(); + + let created = trainer + .train_index(data_stream, 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 json_query = JsonQuery::new(Arc::new(query.clone()), path.to_string()); + index + .search(&json_query, &NoOpMetricsCollector) + .await + .unwrap() + } + + /// Regression test for the JSON-path float index returning wrong results for + /// non-exactly-representable float64 values (issue #7485). The extracted value + /// column must be re-sorted numerically before it reaches the btree trainer; + /// otherwise the per-page min/max invert and range/equality searches miss rows. + #[tokio::test] + async fn test_json_float_btree_index_non_exact_floats() { + use lance_select::RowAddrTreeMap; + use std::ops::Bound; + + let f64 = |v: f64| ScalarValue::Float64(Some(v)); + + // Case A: the three values from the issue's repro. We feed them sorted by the + // raw JSONB bytes to faithfully simulate `scan_training_data`'s asc_nulls_first + // sort over the raw binary column. row0=10.5, row1=40.1, row2=-3.2. + let docs_a = &[ + (r#"{"latitude": 10.5}"#, 0u64), + (r#"{"latitude": 40.1}"#, 1u64), + (r#"{"latitude": -3.2}"#, 2u64), + ][..]; + let byte_key = |doc: &str| doc.parse::().unwrap().to_vec(); + let mut sorted_a = docs_a.to_vec(); + sorted_a.sort_by_key(|(doc, _)| byte_key(doc)); + + let cases_a: &[(SargableQuery, Vec)] = &[ + // > 0 -> the two positive values + ( + SargableQuery::Range(Bound::Excluded(f64(0.0)), Bound::Unbounded), + vec![0, 1], + ), + // >= 10.5 -> the two values at or above 10.5 + ( + SargableQuery::Range(Bound::Included(f64(10.5)), Bound::Unbounded), + vec![0, 1], + ), + // = 40.1 (non-exact) -> row1 only + (SargableQuery::Equals(f64(40.1)), vec![1]), + // = 10.5 (exact) -> row0; guards against silent all-null extraction + (SargableQuery::Equals(f64(10.5)), vec![0]), + // < 100 -> all three rows + ( + SargableQuery::Range(Bound::Unbounded, Bound::Excluded(f64(100.0))), + vec![0, 1, 2], + ), + ]; + for (query, expected) in cases_a { + let result = train_and_search_json_float_index(&sorted_a, query).await; + assert_eq!( + result, + SearchResult::exact(RowAddrTreeMap::from_iter(expected.iter().copied())), + "case A query {query:?}" + ); + } + + // Case B: a deterministic inversion that fails regardless of the exact JSONB + // byte encoding. Feeding [40.1 -> row0, -3.2 -> row1] in that order records an + // inverted page min/max (min=40.1 > max=-3.2) unless the value column is + // re-sorted before training. + let docs_b = &[ + (r#"{"latitude": 40.1}"#, 0u64), + (r#"{"latitude": -3.2}"#, 1u64), + ][..]; + let cases_b: &[(SargableQuery, Vec)] = &[ + ( + SargableQuery::Range(Bound::Excluded(f64(0.0)), Bound::Unbounded), + vec![0], + ), + (SargableQuery::Equals(f64(40.1)), vec![0]), + (SargableQuery::Equals(f64(-3.2)), vec![1]), + ]; + for (query, expected) in cases_b { + let result = train_and_search_json_float_index(docs_b, query).await; + assert_eq!( + result, + SearchResult::exact(RowAddrTreeMap::from_iter(expected.iter().copied())), + "case B query {query:?}" + ); + } + } } From 181b5d65920b10748244b5144caaed704dfdeb53 Mon Sep 17 00:00:00 2001 From: Sohum Trivedi Date: Tue, 14 Jul 2026 09:21:30 -0700 Subject: [PATCH 2/2] Address review feedback on JSON value-ordering training path 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. --- rust/lance-index/src/scalar/json.rs | 237 +++++++++++++++++----------- 1 file changed, 142 insertions(+), 95 deletions(-) diff --git a/rust/lance-index/src/scalar/json.rs b/rust/lance-index/src/scalar/json.rs index dab66f4591f..ea99feb27a1 100644 --- a/rust/lance-index/src/scalar/json.rs +++ b/rust/lance-index/src/scalar/json.rs @@ -751,30 +751,51 @@ impl BasicTrainer for JsonIndexPlugin { // first/last row) would then record inverted statistics and silently return // wrong query results. Re-sort by the extracted value here so the value-ordering // contract is honored. Ordering-agnostic targets (e.g. bitmap) are left as-is. - 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() - }, - )? - } else { - converted_stream + let converted_stream = match 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)); + // `use_spilling` only bounds this SortExec. It does not make training + // memory-bounded overall: `extract_json_with_type_info` and + // `convert_stream_by_type` already collect every batch into a Vec before + // re-streaming, so the extracted/converted data is fully materialized in + // memory upstream of this sort regardless of this flag. + execute_plan( + sorted, + LanceExecutionOptions { + use_spilling: true, + ..Default::default() + }, + )? + } + TrainingOrdering::Addresses => { + // `JsonTrainingRequest::criteria()` delegates the target's criteria + // verbatim, so when the target wants address ordering the scanner already + // supplies the raw JSON rows in address order, and the extraction + // (a ProjectionExec) and per-batch conversion preserve that row order. + // Re-sorting by value here would violate the Addresses contract, so pass + // the stream through untouched. (Note: address-ordered targets such as + // zonemap/bloomfilter additionally require a row-addr column that this + // pipeline does not yet project - only [value, row_id] - an orthogonal, + // pre-existing limitation independent of ordering.) + converted_stream + } + TrainingOrdering::None => { + // No ordering requested; leave the stream unsorted. + converted_stream + } }; let target_index = target_trainer @@ -878,6 +899,9 @@ mod tests { use crate::scalar::SargableQuery; use arrow_array::{ArrayRef, RecordBatch}; use arrow_schema::{DataType, Field, Schema}; + use lance_select::RowAddrTreeMap; + use rstest::rstest; + use std::ops::Bound; use std::sync::Arc; // Note: The old test_detect_json_value_type test has been removed as we now use @@ -1105,82 +1129,105 @@ mod tests { .unwrap() } + /// Serializes the parameterized cases below. Each case runs a spilling `SortExec` + /// that reserves a non-spillable merge buffer from the process-wide cached DataFusion + /// memory pool (see `get_session_context`); running all cases in parallel would + /// contend for that shared pool and intermittently exhaust it. The original + /// single-`#[tokio::test]` version ran the cases sequentially, so hold this guard for + /// the duration of each case to preserve that one-at-a-time behavior. + static FLOAT_INDEX_CASE_GUARD: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + + /// Case A docs: the three values from the issue's repro. Fed sorted by the raw + /// JSONB bytes (see the test body) to faithfully simulate `scan_training_data`'s + /// asc_nulls_first sort over the raw binary column. row0=10.5, row1=40.1, row2=-3.2. + const DOCS_A: &[(&str, u64)] = &[ + (r#"{"latitude": 10.5}"#, 0u64), + (r#"{"latitude": 40.1}"#, 1u64), + (r#"{"latitude": -3.2}"#, 2u64), + ]; + + /// Case B docs: a deterministic inversion that fails regardless of the exact JSONB + /// byte encoding. Feeding [40.1 -> row0, -3.2 -> row1] in that order records an + /// inverted page min/max (min=40.1 > max=-3.2) unless the value column is re-sorted + /// before training. Fed unsorted on purpose - this asymmetry with case A is the + /// regression being covered. + const DOCS_B: &[(&str, u64)] = &[ + (r#"{"latitude": 40.1}"#, 0u64), + (r#"{"latitude": -3.2}"#, 1u64), + ]; + /// Regression test for the JSON-path float index returning wrong results for /// non-exactly-representable float64 values (issue #7485). The extracted value /// column must be re-sorted numerically before it reaches the btree trainer; /// otherwise the per-page min/max invert and range/equality searches miss rows. + /// + /// Case A (`sort_by_raw_jsonb = true`) is fed byte-sorted to simulate the scanner; + /// case B is fed in raw order to force a page min/max inversion. + #[rstest] + // Case A: > 0 -> the two positive values + #[case::a_gt_zero( + DOCS_A, + true, + SargableQuery::Range(Bound::Excluded(ScalarValue::Float64(Some(0.0))), Bound::Unbounded), + vec![0, 1] + )] + // Case A: >= 10.5 -> the two values at or above 10.5 + #[case::a_gte_10_5( + DOCS_A, + true, + SargableQuery::Range(Bound::Included(ScalarValue::Float64(Some(10.5))), Bound::Unbounded), + vec![0, 1] + )] + // Case A: = 40.1 (non-exact) -> row1 only + #[case::a_eq_40_1(DOCS_A, true, SargableQuery::Equals(ScalarValue::Float64(Some(40.1))), vec![1])] + // Case A: = 10.5 (exact) -> row0; guards against silent all-null extraction + #[case::a_eq_10_5(DOCS_A, true, SargableQuery::Equals(ScalarValue::Float64(Some(10.5))), vec![0])] + // Case A: < 100 -> all three rows + #[case::a_lt_100( + DOCS_A, + true, + SargableQuery::Range(Bound::Unbounded, Bound::Excluded(ScalarValue::Float64(Some(100.0)))), + vec![0, 1, 2] + )] + // Case B: > 0 -> row0 only + #[case::b_gt_zero( + DOCS_B, + false, + SargableQuery::Range(Bound::Excluded(ScalarValue::Float64(Some(0.0))), Bound::Unbounded), + vec![0] + )] + // Case B: = 40.1 (non-exact) -> row0 + #[case::b_eq_40_1(DOCS_B, false, SargableQuery::Equals(ScalarValue::Float64(Some(40.1))), vec![0])] + // Case B: = -3.2 (non-exact) -> row1 + #[case::b_eq_neg_3_2(DOCS_B, false, SargableQuery::Equals(ScalarValue::Float64(Some(-3.2))), vec![1])] #[tokio::test] - async fn test_json_float_btree_index_non_exact_floats() { - use lance_select::RowAddrTreeMap; - use std::ops::Bound; - - let f64 = |v: f64| ScalarValue::Float64(Some(v)); - - // Case A: the three values from the issue's repro. We feed them sorted by the - // raw JSONB bytes to faithfully simulate `scan_training_data`'s asc_nulls_first - // sort over the raw binary column. row0=10.5, row1=40.1, row2=-3.2. - let docs_a = &[ - (r#"{"latitude": 10.5}"#, 0u64), - (r#"{"latitude": 40.1}"#, 1u64), - (r#"{"latitude": -3.2}"#, 2u64), - ][..]; - let byte_key = |doc: &str| doc.parse::().unwrap().to_vec(); - let mut sorted_a = docs_a.to_vec(); - sorted_a.sort_by_key(|(doc, _)| byte_key(doc)); - - let cases_a: &[(SargableQuery, Vec)] = &[ - // > 0 -> the two positive values - ( - SargableQuery::Range(Bound::Excluded(f64(0.0)), Bound::Unbounded), - vec![0, 1], - ), - // >= 10.5 -> the two values at or above 10.5 - ( - SargableQuery::Range(Bound::Included(f64(10.5)), Bound::Unbounded), - vec![0, 1], - ), - // = 40.1 (non-exact) -> row1 only - (SargableQuery::Equals(f64(40.1)), vec![1]), - // = 10.5 (exact) -> row0; guards against silent all-null extraction - (SargableQuery::Equals(f64(10.5)), vec![0]), - // < 100 -> all three rows - ( - SargableQuery::Range(Bound::Unbounded, Bound::Excluded(f64(100.0))), - vec![0, 1, 2], - ), - ]; - for (query, expected) in cases_a { - let result = train_and_search_json_float_index(&sorted_a, query).await; - assert_eq!( - result, - SearchResult::exact(RowAddrTreeMap::from_iter(expected.iter().copied())), - "case A query {query:?}" - ); - } + async fn test_json_float_btree_index_non_exact_floats( + #[case] docs: &'static [(&'static str, u64)], + #[case] sort_by_raw_jsonb: bool, + #[case] query: SargableQuery, + #[case] expected: Vec, + ) { + // Serialize the spilling sorts so they do not contend for the shared cached + // memory pool; held until the end of the test body. + let _guard = FLOAT_INDEX_CASE_GUARD.lock().await; + + // Case A simulates `scan_training_data`'s asc_nulls_first sort over the raw JSONB + // bytes; case B is fed in raw order to force a page min/max inversion. Preserve + // that asymmetry - it is what exercises the #7485 regression. + let docs = if sort_by_raw_jsonb { + let byte_key = |doc: &str| doc.parse::().unwrap().to_vec(); + let mut sorted = docs.to_vec(); + sorted.sort_by_key(|(doc, _)| byte_key(doc)); + sorted + } else { + docs.to_vec() + }; - // Case B: a deterministic inversion that fails regardless of the exact JSONB - // byte encoding. Feeding [40.1 -> row0, -3.2 -> row1] in that order records an - // inverted page min/max (min=40.1 > max=-3.2) unless the value column is - // re-sorted before training. - let docs_b = &[ - (r#"{"latitude": 40.1}"#, 0u64), - (r#"{"latitude": -3.2}"#, 1u64), - ][..]; - let cases_b: &[(SargableQuery, Vec)] = &[ - ( - SargableQuery::Range(Bound::Excluded(f64(0.0)), Bound::Unbounded), - vec![0], - ), - (SargableQuery::Equals(f64(40.1)), vec![0]), - (SargableQuery::Equals(f64(-3.2)), vec![1]), - ]; - for (query, expected) in cases_b { - let result = train_and_search_json_float_index(docs_b, query).await; - assert_eq!( - result, - SearchResult::exact(RowAddrTreeMap::from_iter(expected.iter().copied())), - "case B query {query:?}" - ); - } + let result = train_and_search_json_float_index(&docs, &query).await; + assert_eq!( + result, + SearchResult::exact(RowAddrTreeMap::from_iter(expected.iter().copied())), + "query {query:?}" + ); } }