diff --git a/rust/lance-index/src/scalar/json.rs b/rust/lance-index/src/scalar/json.rs index 09678d6787e..ea99feb27a1 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,62 @@ 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 = 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 .train_index( converted_stream, @@ -837,8 +896,12 @@ 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 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 @@ -978,4 +1041,193 @@ 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() + } + + /// 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( + #[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() + }; + + let result = train_and_search_json_float_index(&docs, &query).await; + assert_eq!( + result, + SearchResult::exact(RowAddrTreeMap::from_iter(expected.iter().copied())), + "query {query:?}" + ); + } }