diff --git a/rust/lance-index-core/src/scalar.rs b/rust/lance-index-core/src/scalar.rs index e209a32cdb8..49b8e949905 100644 --- a/rust/lance-index-core/src/scalar.rs +++ b/rust/lance-index-core/src/scalar.rs @@ -4,7 +4,7 @@ //! Abstract scalar index traits and types for Lance index plugins use arrow_array::{BooleanArray, RecordBatch, UInt64Array}; -use arrow_schema::Schema; +use arrow_schema::{DataType, Schema}; use async_trait::async_trait; use bytes::Bytes; use datafusion::physical_plan::SendableRecordBatchStream; @@ -555,6 +555,16 @@ pub trait ScalarIndex: Send + Sync + std::fmt::Debug + Index + DeepSizeOf { /// with the same configuration on another dataset. fn derive_index_params(&self) -> Result; + /// Returns the value type expected by [`Self::update`], when the index has + /// a durable type contract for its training data. + /// + /// Wrapper indices use this to transform new data to the same type as the + /// loaded index instead of inferring a potentially different type from an + /// update batch. Index types without such a contract may return `None`. + fn training_data_type(&self) -> Option { + None + } + /// Global `[min, max]` of the indexed column from index metadata, without a /// scan, or `None` if this index type cannot supply a sound bound. When /// `Some`, the range is a superset of live values (conservative under diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index ddbe6f3ed53..305df6ebae1 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -2368,6 +2368,10 @@ impl ScalarIndex for BTreeIndex { })?; Ok(ScalarIndexParams::for_builtin(BuiltinIndexType::BTree).with_params(¶ms)) } + + fn training_data_type(&self) -> Option { + Some(self.data_type.clone()) + } } struct BatchStats { diff --git a/rust/lance-index/src/scalar/json.rs b/rust/lance-index/src/scalar/json.rs index 00d1aaca333..315e4505412 100644 --- a/rust/lance-index/src/scalar/json.rs +++ b/rust/lance-index/src/scalar/json.rs @@ -15,7 +15,7 @@ use datafusion::{ execution::SendableRecordBatchStream, physical_plan::{ExecutionPlan, projection::ProjectionExec, sorts::sort::SortExec}, }; -use datafusion_common::{ScalarValue, config::ConfigOptions}; +use datafusion_common::{DataFusionError, ScalarValue, config::ConfigOptions}; use datafusion_expr::{Expr, Operator, ScalarUDF}; use datafusion_physical_expr::{ PhysicalExpr, PhysicalSortExpr, ScalarFunctionExpr, @@ -141,6 +141,22 @@ impl ScalarIndex for JsonIndex { dest_store: &dyn IndexStore, old_data_filter: Option, ) -> Result { + let target_criteria = self.target_index.update_criteria().data_criteria; + let target_type = self.target_index.training_data_type().ok_or_else(|| { + Error::not_supported(format!( + "JSON index updates for path '{}' require target index {} to report its training data type", + self.path, + self.target_index.index_type() + )) + })?; + let new_data = JsonIndexPlugin::extract_json(new_data, self.path.clone())?; + let new_data = + JsonIndexPlugin::convert_stream_by_type(new_data, target_type, self.path.clone())?; + let new_data = if target_criteria.ordering == TrainingOrdering::Values { + JsonIndexPlugin::sort_stream_by_value(new_data).await? + } else { + new_data + }; let target_created = self .target_index .update(new_data, dest_store, old_data_filter) @@ -158,11 +174,32 @@ impl ScalarIndex for JsonIndex { } fn update_criteria(&self) -> UpdateCriteria { - self.target_index.update_criteria() + let target_criteria = self.target_index.update_criteria(); + UpdateCriteria { + requires_old_data: target_criteria.requires_old_data, + data_criteria: json_scan_criteria(&target_criteria.data_criteria), + } } fn derive_index_params(&self) -> Result { - self.target_index.derive_index_params() + let target_params = self.target_index.derive_index_params()?; + let target_data_type = self + .target_index + .training_data_type() + .as_ref() + .map(JsonIndexTargetType::try_from) + .transpose()?; + let params = JsonIndexParameters { + target_index_type: target_params.index_type, + target_index_parameters: target_params.params, + target_data_type, + path: self.path.clone(), + }; + Ok(super::ScalarIndexParams::new("json".to_string()).with_params(¶ms)) + } + + fn training_data_type(&self) -> Option { + self.target_index.training_data_type() } } @@ -171,9 +208,49 @@ impl ScalarIndex for JsonIndex { pub struct JsonIndexParameters { target_index_type: String, target_index_parameters: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + target_data_type: Option, path: String, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +enum JsonIndexTargetType { + Boolean, + Int64, + Float64, + Utf8, + LargeBinary, +} + +impl TryFrom<&DataType> for JsonIndexTargetType { + type Error = Error; + + fn try_from(data_type: &DataType) -> Result { + match data_type { + DataType::Boolean => Ok(Self::Boolean), + DataType::Int64 => Ok(Self::Int64), + DataType::Float64 => Ok(Self::Float64), + DataType::Utf8 => Ok(Self::Utf8), + DataType::LargeBinary => Ok(Self::LargeBinary), + _ => Err(Error::not_supported(format!( + "JSON index target data type {data_type:?} cannot be preserved for rebuilds" + ))), + } + } +} + +impl From for DataType { + fn from(data_type: JsonIndexTargetType) -> Self { + match data_type { + JsonIndexTargetType::Boolean => Self::Boolean, + JsonIndexTargetType::Int64 => Self::Int64, + JsonIndexTargetType::Float64 => Self::Float64, + JsonIndexTargetType::Utf8 => Self::Utf8, + JsonIndexTargetType::LargeBinary => Self::LargeBinary, + } + } +} + // TODO: Do we really need to wrap the query or could we just return the target query directly? // // I think the only thing we really gain is a different format impl (e.g. it shows up as a json query @@ -360,25 +437,19 @@ pub struct JsonTrainingRequest { criteria: TrainingCriteria, } +fn json_scan_criteria(target_criteria: &TrainingCriteria) -> TrainingCriteria { + let mut criteria = target_criteria.clone(); + if criteria.ordering == TrainingOrdering::Values { + // The scanner can only order the raw JSON column. The JSON wrapper sorts + // the extracted path value before passing it to the target index. + criteria.ordering = TrainingOrdering::None; + } + criteria +} + impl JsonTrainingRequest { pub fn new(parameters: JsonIndexParameters, target_request: Box) -> Self { - let target_criteria = target_request.criteria(); - // The scanner can only sort its output by the raw JSON column, not by the value - // at `path` that this plugin extracts from it, so a `Values`-ordered scan here - // would sort by the wrong key and still need re-sorting after extraction. Ask - // for unordered input instead and let `train_index` sort the extracted value - // stream itself, once, right before handing it to the target trainer. - // - // Preserve `Addresses` so callers must still provide the ordering required by - // target trainers such as ZoneMap, even though the dataset scanner currently - // returns rows in row-address order by default. - let ordering = match target_criteria.ordering { - TrainingOrdering::Values => TrainingOrdering::None, - ordering => ordering, - }; - let mut criteria = TrainingCriteria::new(ordering); - criteria.needs_row_ids = target_criteria.needs_row_ids; - criteria.needs_row_addrs = target_criteria.needs_row_addrs; + let criteria = json_scan_criteria(target_request.criteria()); Self { parameters, target_request, @@ -414,11 +485,11 @@ impl JsonIndexPlugin { Ok(self.registry.lock().unwrap().as_ref().expect_ok()?.clone()) } - /// Extract JSON with type information using the new UDF - async fn extract_json_with_type_info( + /// Extract a JSON path and its type tag while preserving all row-location columns. + fn extract_json( data: SendableRecordBatchStream, path: String, - ) -> Result<(SendableRecordBatchStream, DataType)> { + ) -> Result { let input = Arc::new(OneShotExec::new(data)); let input_schema = input.schema(); let value_column_idx = input_schema @@ -452,75 +523,280 @@ impl JsonIndexPlugin { let project = ProjectionExec::try_new(exprs, input)?; let ctx = get_session_context(&LanceExecutionOptions::default()); - let mut stream = project.execute(0, ctx.task_ctx())?; + project.execute(0, ctx.task_ctx()).map_err(Into::into) + } - // Collect batches and determine type from first non-null value - let mut all_batches = Vec::new(); - let mut inferred_type: Option = None; + fn infer_type_from_batch(batch: &RecordBatch, path: &str) -> Result> { + let json_result_column = batch + .column_by_name("json_result") + .ok_or_else(|| Error::invalid_input_source("Missing json_result column".into()))?; + let struct_array = json_result_column + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::invalid_input_source("json_result is not a struct".into()))?; + let type_array = struct_array + .column_by_name("type_tag") + .ok_or_else(|| Error::invalid_input_source("Missing type_tag column in struct".into()))? + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::invalid_input_source("type_tag is not UInt8".into()))?; + + for type_tag in type_array.iter().flatten() { + let jsonb_type = JsonbType::from_u8(type_tag).ok_or_else(|| { + Error::invalid_input_source( + format!("JSON path '{path}' produced invalid type tag {type_tag}").into(), + ) + })?; + let data_type = match jsonb_type { + JsonbType::Null => continue, + JsonbType::Boolean => DataType::Boolean, + JsonbType::Int64 => DataType::Int64, + JsonbType::Float64 => DataType::Float64, + JsonbType::String => DataType::Utf8, + JsonbType::Array | JsonbType::Object => DataType::LargeBinary, + }; + return Ok(Some(data_type)); + } + Ok(None) + } + + /// Extract JSON and infer the target type from the first non-null value. + /// + /// Only the prefix needed for inference is buffered. Once a type is found, + /// the remainder of the input stays streaming. + async fn extract_json_with_type_info( + data: SendableRecordBatchStream, + path: String, + ) -> Result<(SendableRecordBatchStream, DataType)> { + let mut stream = Self::extract_json(data, path.clone())?; + let schema = stream.schema(); + let mut buffered_batches = Vec::new(); + let mut inferred_type = None; while let Some(batch_result) = stream.next().await { let batch = batch_result?; - - // Determine type from first non-null value if not yet set - if inferred_type.is_none() - && let Some(json_result_column) = batch.column_by_name("json_result") - && let Some(struct_array) = - json_result_column.as_any().downcast_ref::() - && let Some(type_array) = struct_array.column_by_name("type_tag") - && let Some(uint8_array) = type_array.as_any().downcast_ref::() - { - // Find first non-null value to determine type - for i in 0..uint8_array.len() { - if !uint8_array.is_null(i) { - let type_tag = uint8_array.value(i); - let jsonb_type = JsonbType::from_u8(type_tag).ok_or_else(|| { - Error::invalid_input_source( - format!("Invalid type tag: {}", type_tag).into(), - ) - })?; - - // Map JsonbType to Arrow DataType - inferred_type = Some(match jsonb_type { - JsonbType::Null => continue, // Skip null values - JsonbType::Boolean => DataType::Boolean, - JsonbType::Int64 => DataType::Int64, - JsonbType::Float64 => DataType::Float64, - JsonbType::String => DataType::Utf8, - JsonbType::Array => DataType::LargeBinary, - JsonbType::Object => DataType::LargeBinary, - }); - break; - } - } + inferred_type = Self::infer_type_from_batch(&batch, &path)?; + buffered_batches.push(batch); + if inferred_type.is_some() { + break; } - - all_batches.push(batch); } - // If no type was inferred (all nulls), default to String let inferred_type = inferred_type.unwrap_or(DataType::Utf8); - - // Recreate stream from collected batches - let schema = all_batches - .first() - .map(|b| b.schema()) - .ok_or_else(|| Error::invalid_input_source("No batches in stream".into()))?; - + let buffered = futures::stream::iter(buffered_batches.into_iter().map(Ok)); let recreated_stream = Box::pin(RecordBatchStreamAdapter::new( schema, - futures::stream::iter(all_batches.into_iter().map(Ok)), + buffered.chain(stream), )) as SendableRecordBatchStream; Ok((recreated_stream, inferred_type)) } - /// Convert the stream with JSONB values and type tags to properly typed values - async fn convert_stream_by_type( + fn validate_json_types( + binary_array: &LargeBinaryArray, + type_array: &UInt8Array, + target_type: &DataType, + path: &str, + ) -> Result<()> { + for index in 0..binary_array.len() { + if binary_array.is_null(index) { + continue; + } + if type_array.is_null(index) { + return Err(Error::invalid_input_source( + format!( + "JSON path '{path}' has a value at batch row {index} without a type tag" + ) + .into(), + )); + } + let type_tag = type_array.value(index); + let actual_type = JsonbType::from_u8(type_tag).ok_or_else(|| { + Error::invalid_input_source( + format!( + "JSON path '{path}' produced invalid type tag {type_tag} at batch row {index}" + ) + .into(), + ) + })?; + if actual_type == JsonbType::Null { + continue; + } + let is_compatible = match target_type { + DataType::Boolean => actual_type == JsonbType::Boolean, + DataType::Int64 => actual_type == JsonbType::Int64, + DataType::Float64 => { + matches!(actual_type, JsonbType::Int64 | JsonbType::Float64) + } + DataType::Utf8 => actual_type == JsonbType::String, + DataType::LargeBinary => { + matches!(actual_type, JsonbType::Array | JsonbType::Object) + } + _ => false, + }; + if !is_compatible { + return Err(Error::invalid_input_source( + format!( + "JSON path '{path}' expected {target_type:?}, but batch row {index} has JSON type {actual_type:?}" + ) + .into(), + )); + } + } + Ok(()) + } + + fn convert_batch_by_type( + batch: RecordBatch, + output_schema: Arc, + passthrough_indices: &[usize], + target_type: &DataType, + path: &str, + ) -> Result { + let json_result_column = batch + .column_by_name("json_result") + .ok_or_else(|| Error::invalid_input_source("Missing json_result column".into()))?; + + let struct_array = json_result_column + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::invalid_input_source("json_result is not a struct".into()))?; + + let value_array = struct_array + .column_by_name("value") + .ok_or_else(|| Error::invalid_input_source("Missing value column in struct".into()))?; + + let binary_array = value_array + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::invalid_input_source("value is not LargeBinary".into()))?; + let type_array = struct_array + .column_by_name("type_tag") + .ok_or_else(|| Error::invalid_input_source("Missing type_tag column in struct".into()))? + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::invalid_input_source("type_tag is not UInt8".into()))?; + + Self::validate_json_types(binary_array, type_array, target_type, path)?; + let is_null = |index| { + binary_array.is_null(index) + || (!type_array.is_null(index) + && type_array.value(index) == JsonbType::Null.as_u8()) + }; + + let converted_array: Arc = match target_type { + DataType::Boolean => { + let mut builder = + arrow_array::builder::BooleanBuilder::with_capacity(binary_array.len()); + for i in 0..binary_array.len() { + if is_null(i) { + builder.append_null(); + } else { + let raw_jsonb = jsonb::RawJsonb::new(binary_array.value(i)); + let value = jsonb::from_raw_jsonb::(&raw_jsonb).map_err(|error| { + Error::invalid_input_source( + format!( + "Failed to convert JSON path '{path}' at batch row {i} to Boolean: {error}" + ) + .into(), + ) + })?; + builder.append_value(value); + } + } + Arc::new(builder.finish()) + } + DataType::Int64 => { + let mut builder = + arrow_array::builder::Int64Builder::with_capacity(binary_array.len()); + for i in 0..binary_array.len() { + if is_null(i) { + builder.append_null(); + } else { + let raw_jsonb = jsonb::RawJsonb::new(binary_array.value(i)); + let value = jsonb::from_raw_jsonb::(&raw_jsonb).map_err(|error| { + Error::invalid_input_source( + format!( + "Failed to convert JSON path '{path}' at batch row {i} to Int64: {error}" + ) + .into(), + ) + })?; + builder.append_value(value); + } + } + Arc::new(builder.finish()) + } + DataType::Float64 => { + let mut builder = + arrow_array::builder::Float64Builder::with_capacity(binary_array.len()); + for i in 0..binary_array.len() { + if is_null(i) { + builder.append_null(); + } else { + let raw_jsonb = jsonb::RawJsonb::new(binary_array.value(i)); + let value = jsonb::from_raw_jsonb::(&raw_jsonb).map_err(|error| { + Error::invalid_input_source( + format!( + "Failed to convert JSON path '{path}' at batch row {i} to Float64: {error}" + ) + .into(), + ) + })?; + builder.append_value(value); + } + } + Arc::new(builder.finish()) + } + DataType::Utf8 => { + let mut builder = + arrow_array::builder::StringBuilder::with_capacity(binary_array.len(), 1024); + for i in 0..binary_array.len() { + if is_null(i) { + builder.append_null(); + } else { + let raw_jsonb = jsonb::RawJsonb::new(binary_array.value(i)); + let value = + jsonb::from_raw_jsonb::(&raw_jsonb).map_err(|error| { + Error::invalid_input_source( + format!( + "Failed to convert JSON path '{path}' at batch row {i} to Utf8: {error}" + ) + .into(), + ) + })?; + builder.append_value(value); + } + } + Arc::new(builder.finish()) + } + DataType::LargeBinary => Arc::new(LargeBinaryArray::from_iter( + (0..binary_array.len()).map(|i| (!is_null(i)).then(|| binary_array.value(i))), + )), + _ => { + return Err(Error::invalid_input_source( + format!("Unsupported JSON index target type: {target_type:?}").into(), + )); + } + }; + + let mut columns = Vec::with_capacity(output_schema.fields().len()); + columns.push(converted_array); + columns.extend( + passthrough_indices + .iter() + .map(|column_idx| batch.column(*column_idx).clone()), + ); + Ok(RecordBatch::try_new(output_schema, columns)?) + } + + /// Convert extracted JSON values one batch at a time to the target's trained type. + fn convert_stream_by_type( data: SendableRecordBatchStream, target_type: DataType, + path: String, ) -> Result { - let input = Arc::new(OneShotExec::new(data)); - let input_schema = input.schema(); + let input_schema = data.schema(); let passthrough_indices = input_schema .fields() .iter() @@ -536,165 +812,21 @@ impl JsonIndexPlugin { ) .collect::>(), )); - let ctx = get_session_context(&LanceExecutionOptions::default()); - let mut stream = input.execute(0, ctx.task_ctx())?; - - let mut converted_batches = Vec::new(); - - while let Some(batch_result) = stream.next().await { + let stream_schema = output_schema.clone(); + let converted = data.map(move |batch_result| { let batch = batch_result?; - - // Extract the struct column containing value and type_tag - let json_result_column = batch - .column_by_name("json_result") - .ok_or_else(|| Error::invalid_input_source("Missing json_result column".into()))?; - - let struct_array = json_result_column - .as_any() - .downcast_ref::() - .ok_or_else(|| Error::invalid_input_source("json_result is not a struct".into()))?; - - let value_array = struct_array.column_by_name("value").ok_or_else(|| { - Error::invalid_input_source("Missing value column in struct".into()) - })?; - - let binary_array = value_array - .as_any() - .downcast_ref::() - .ok_or_else(|| Error::invalid_input_source("value is not LargeBinary".into()))?; - - // Convert based on target type using serde deserialization - let converted_array: Arc = - match target_type { - DataType::Boolean => { - let mut builder = - arrow_array::builder::BooleanBuilder::with_capacity(binary_array.len()); - for i in 0..binary_array.len() { - if binary_array.is_null(i) { - builder.append_null(); - } else if let Some(bytes) = binary_array.value(i).into() { - let raw_jsonb = jsonb::RawJsonb::new(bytes); - // Try to deserialize directly to bool - match jsonb::from_raw_jsonb::(&raw_jsonb) { - Ok(bool_val) => builder.append_value(bool_val), - Err(e) => { - return Err(Error::invalid_input_source(format!( - "Failed to deserialize JSONB to bool at index {}: {}", - i, e - ) - .into())); - } - } - } else { - builder.append_null(); - } - } - Arc::new(builder.finish()) - } - DataType::Int64 => { - let mut builder = - arrow_array::builder::Int64Builder::with_capacity(binary_array.len()); - for i in 0..binary_array.len() { - if binary_array.is_null(i) { - builder.append_null(); - } else if let Some(bytes) = binary_array.value(i).into() { - let raw_jsonb = jsonb::RawJsonb::new(bytes); - // Try to deserialize directly to i64 - match jsonb::from_raw_jsonb::(&raw_jsonb) { - Ok(int_val) => builder.append_value(int_val), - Err(e) => { - return Err(Error::invalid_input_source(format!( - "Failed to deserialize JSONB to i64 at index {}: {}", - i, e - ) - .into())); - } - } - } else { - builder.append_null(); - } - } - Arc::new(builder.finish()) - } - DataType::Float64 => { - let mut builder = - arrow_array::builder::Float64Builder::with_capacity(binary_array.len()); - for i in 0..binary_array.len() { - if binary_array.is_null(i) { - builder.append_null(); - } else if let Some(bytes) = binary_array.value(i).into() { - let raw_jsonb = jsonb::RawJsonb::new(bytes); - // Try to deserialize directly to f64 (serde handles int->float conversion) - match jsonb::from_raw_jsonb::(&raw_jsonb) { - Ok(float_val) => builder.append_value(float_val), - Err(e) => { - return Err(Error::invalid_input_source(format!( - "Failed to deserialize JSONB to f64 at index {}: {}", - i, e - ) - .into())); - } - } - } else { - builder.append_null(); - } - } - Arc::new(builder.finish()) - } - DataType::Utf8 => { - let mut builder = arrow_array::builder::StringBuilder::with_capacity( - binary_array.len(), - 1024, - ); - for i in 0..binary_array.len() { - if binary_array.is_null(i) { - builder.append_null(); - } else if let Some(bytes) = binary_array.value(i).into() { - let raw_jsonb = jsonb::RawJsonb::new(bytes); - // Try to deserialize to String, or use to_string() for any type - match jsonb::from_raw_jsonb::(&raw_jsonb) { - Ok(str_val) => builder.append_value(&str_val), - Err(_) => { - // For non-string types, convert to string representation - builder.append_value(raw_jsonb.to_string()); - } - } - } else { - builder.append_null(); - } - } - Arc::new(builder.finish()) - } - DataType::LargeBinary => { - // Keep as binary for array/object types - value_array.clone() - } - _ => { - return Err(Error::invalid_input_source( - format!("Unsupported target type: {:?}", target_type).into(), - )); - } - }; - - let mut columns = Vec::with_capacity(output_schema.fields().len()); - columns.push(converted_array); - columns.extend( - passthrough_indices - .iter() - .map(|column_idx| batch.column(*column_idx).clone()), - ); - let new_batch = RecordBatch::try_new(output_schema.clone(), columns)?; - - converted_batches.push(new_batch); - } - - // Create stream from converted batches - if converted_batches.is_empty() { - return Err(Error::invalid_input_source("No batches to convert".into())); - } + Self::convert_batch_by_type( + batch, + output_schema.clone(), + &passthrough_indices, + &target_type, + &path, + ) + .map_err(DataFusionError::from) + }); Ok(Box::pin(RecordBatchStreamAdapter::new( - output_schema, - futures::stream::iter(converted_batches.into_iter().map(Ok)), + stream_schema, + converted, ))) } @@ -741,10 +873,13 @@ impl BasicTrainer for JsonIndexPlugin { )); } - // Initially use Utf8, will be refined during training with type inference - let target_type = DataType::Utf8; - let params = serde_json::from_str::(params)?; + // Initial builds infer the type from the data. Derived rebuild parameters + // carry the learned type so every new segment uses the same target schema. + let target_type = params + .target_data_type + .map(DataType::from) + .unwrap_or(DataType::Utf8); let registry = self.registry()?; let target_plugin = registry.get_plugin_by_name(¶ms.target_index_type)?; let target_trainer = target_plugin.basic_trainer().ok_or_else(|| { @@ -773,13 +908,20 @@ impl BasicTrainer for JsonIndexPlugin { .unwrap(); let path = request.parameters.path.clone(); - // Extract JSON with type information - let (data_stream, inferred_type) = - Self::extract_json_with_type_info(data, path.clone()).await?; + let (data_stream, target_type) = + if let Some(target_data_type) = request.parameters.target_data_type { + ( + Self::extract_json(data, path.clone())?, + DataType::from(target_data_type), + ) + } else { + Self::extract_json_with_type_info(data, path.clone()).await? + }; - // Convert the stream to properly typed values based on inferred type + // Initial builds use the inferred type; rebuilds use the learned target + // type carried in the derived parameters. let converted_stream = - Self::convert_stream_by_type(data_stream, inferred_type.clone()).await?; + Self::convert_stream_by_type(data_stream, target_type.clone(), path.clone())?; // `JsonTrainingRequest::criteria()` asked the scanner for unordered input (see // its constructor), since the scanner can only sort on the raw JSON column, not @@ -796,11 +938,9 @@ impl BasicTrainer for JsonIndexPlugin { converted_stream }; - // Update the target request with inferred type let registry = self.registry()?; let target_plugin = registry.get_plugin_by_name(&request.parameters.target_index_type)?; - // Create a new training request with the inferred type let target_trainer = target_plugin.basic_trainer().ok_or_else(|| { Error::invalid_input_source( format!("The '{}' index type does not support basic training, please refer to the index's documentation for more details on how to create this index.", request.parameters.target_index_type).into(), @@ -812,7 +952,7 @@ impl BasicTrainer for JsonIndexPlugin { .target_index_parameters .as_deref() .unwrap_or("{}"), - &Field::new("", inferred_type, true), + &Field::new("", target_type, true), )?; let target_index = target_trainer @@ -1147,6 +1287,203 @@ mod tests { (store, tmpdir) } + fn json_update_batch(json_docs: &[&str], row_ids: Vec) -> RecordBatch { + use arrow_array::{LargeBinaryArray, UInt64Array}; + + let jsonb = json_docs + .iter() + .map(|json| json.parse::().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), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(LargeBinaryArray::from_iter_values( + jsonb.iter().map(Vec::as_slice), + )), + Arc::new(UInt64Array::from(row_ids)), + ], + ) + .unwrap() + } + + fn json_update_stream(json_docs: &[&str], row_ids: Vec) -> SendableRecordBatchStream { + use futures::stream; + + let batch = json_update_batch(json_docs, row_ids); + let schema = batch.schema(); + Box::pin(RecordBatchStreamAdapter::new( + schema, + stream::iter([Ok(batch)]), + )) + } + + #[rstest] + #[case::all_null_delta( + &[r#"{"v": 1}"#], + &[r#"{"v": null}"#, r#"{"other": 2}"#], + SargableQuery::IsNull(), + vec![1, 2] + )] + #[case::integer_delta_for_float_index( + &[r#"{"v": 1.5}"#], + &[r#"{"v": 2}"#], + SargableQuery::Equals(ScalarValue::Float64(Some(2.0))), + vec![1] + )] + #[tokio::test] + async fn test_json_btree_update_uses_trained_target_type( + #[case] initial_docs: &[&str], + #[case] update_docs: &[&str], + #[case] query: SargableQuery, + #[case] expected_row_ids: Vec, + ) { + use crate::metrics::NoOpMetricsCollector; + use lance_select::RowAddrTreeMap; + + let (source_store, _source_dir) = local_json_index_store(); + let index = train_and_load_json_index( + source_store, + "btree", + TrainingOrdering::None, + "v", + initial_docs, + ) + .await; + let expected_type = if matches!(&query, SargableQuery::IsNull()) { + DataType::Int64 + } else { + DataType::Float64 + }; + assert_eq!(index.training_data_type(), Some(expected_type)); + + let (dest_store, _dest_dir) = local_json_index_store(); + let row_ids = + (initial_docs.len() as u64..(initial_docs.len() + update_docs.len()) as u64).collect(); + let created = index + .update( + json_update_stream(update_docs, row_ids), + dest_store.as_ref(), + None, + ) + .await + .unwrap(); + + let registry = IndexPluginRegistry::with_default_plugins(); + let plugin = registry.get_plugin_by_name("json").unwrap(); + let updated = plugin + .load_index( + dest_store, + &created.index_details, + None, + &LanceCache::no_cache(), + ) + .await + .unwrap(); + let result = updated + .search( + &JsonQuery::new(Arc::new(query), "v".to_string()), + &NoOpMetricsCollector, + ) + .await + .unwrap(); + assert_eq!( + result, + SearchResult::exact(RowAddrTreeMap::from_iter(expected_row_ids)) + ); + } + + #[tokio::test] + async fn test_json_conversion_is_streaming() { + use arrow_array::Int64Array; + use futures::stream; + + let first_batch = json_update_batch(&[r#"{"v": 1}"#], vec![0]); + let schema = first_batch.schema(); + let raw_stream = Box::pin(RecordBatchStreamAdapter::new( + schema, + stream::iter([ + Ok(first_batch), + Err(DataFusionError::Execution( + "second batch must not be polled for the first output".to_string(), + )), + ]), + )) as SendableRecordBatchStream; + + let extracted = JsonIndexPlugin::extract_json(raw_stream, "v".to_string()).unwrap(); + let mut converted = + JsonIndexPlugin::convert_stream_by_type(extracted, DataType::Int64, "v".to_string()) + .unwrap(); + + let first = converted.next().await.unwrap().unwrap(); + let values = first[VALUE_COLUMN_NAME] + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.value(0), 1); + let error = converted.next().await.unwrap().unwrap_err(); + assert!( + error + .to_string() + .contains("second batch must not be polled") + ); + } + + #[tokio::test] + async fn test_json_btree_update_reports_type_drift() { + let (source_store, _source_dir) = local_json_index_store(); + let index = train_and_load_json_index( + source_store, + "btree", + TrainingOrdering::None, + "v", + &[r#"{"v": 1}"#], + ) + .await; + let (dest_store, _dest_dir) = local_json_index_store(); + let error = index + .update( + json_update_stream(&[r#"{"v": true}"#], vec![1]), + dest_store.as_ref(), + None, + ) + .await + .err() + .expect("type drift should fail the update"); + let message = error.to_string(); + assert!(message.contains("JSON path 'v'"), "{message}"); + assert!(message.contains("expected Int64"), "{message}"); + assert!(message.contains("JSON type Boolean"), "{message}"); + } + + #[tokio::test] + async fn test_json_derived_params_preserve_wrapper() { + let (store, _tmpdir) = local_json_index_store(); + let index = train_and_load_json_index( + store, + "btree", + TrainingOrdering::None, + "v", + &[r#"{"v": 1}"#], + ) + .await; + + let derived = index.derive_index_params().unwrap(); + assert_eq!(derived.index_type, "json"); + let parameters: JsonIndexParameters = + serde_json::from_str(derived.params.as_deref().unwrap()).unwrap(); + assert_eq!(parameters.path, "v"); + assert_eq!(parameters.target_index_type, "btree"); + assert!(parameters.target_index_parameters.is_some()); + assert_eq!( + parameters.target_data_type, + Some(JsonIndexTargetType::Int64) + ); + } + /// Regression test for https://github.com/lance-format/lance/issues/7859. #[rstest] #[case::zonemap("zonemap", TrainingOrdering::Addresses)] @@ -1261,13 +1598,9 @@ mod tests { /// resulting btree still answers `IsNull` and non-null range/equality queries /// correctly with nulls mixed in and fed out of value order. /// - /// Row 1's `path` is missing entirely, which is what actually produces a null in the - /// extracted value column (`extract_json_path_with_type` returns `None`, which - /// `json_extract_with_type_impl` turns into an arrow-null). An explicit JSON `null` - /// literal at `path` (e.g. `{"v": null}`) is a different, pre-existing case that - /// `convert_stream_by_type` does not yet handle (it tries to deserialize the JSONB - /// `null` bytes as the inferred type and errors) -- unrelated to this fix, so it's - /// out of scope here. + /// Row 1's `path` contains an explicit JSON null. The extracted binary value still + /// contains JSONB bytes, so conversion must use the accompanying type tag to turn it + /// into an Arrow null before sorting and training the target index. #[tokio::test] async fn test_json_btree_index_null_at_path() { use crate::metrics::NoOpMetricsCollector; @@ -1281,10 +1614,10 @@ mod tests { TrainingOrdering::None, "v", &[ - r#"{"v": 40.1}"#, // row 0 - r#"{"other": 1}"#, // row 1: path missing -> null - r#"{"v": -3.2}"#, // row 2 - r#"{"v": 10.5}"#, // row 3 + r#"{"v": 40.1}"#, // row 0 + r#"{"v": null}"#, // row 1 + r#"{"v": -3.2}"#, // row 2 + r#"{"v": 10.5}"#, // row 3 ], ) .await; diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index aae1c302282..173fdec3d9b 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -4075,6 +4075,120 @@ async fn test_fts_phrase_query_preserves_stop_word_gaps() { assert!(!ids.contains(&3), "ids={ids:?}"); } +fn json_batch(values: Vec<&str>) -> RecordBatch { + let mut metadata = HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("json", DataType::Utf8, false).with_metadata(metadata), + ])); + RecordBatch::try_new(schema, vec![Arc::new(StringArray::from(values))]).unwrap() +} + +async fn json_btree_dataset(initial_values: Vec<&str>) -> Dataset { + let initial = json_batch(initial_values); + let initial_schema = initial.schema(); + let reader = RecordBatchIterator::new([Ok(initial)], initial_schema); + let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap(); + let params = ScalarIndexParams::new("json".to_string()).with_params(&serde_json::json!({ + "target_index_type": "btree", + "path": "val", + })); + dataset + .create_index( + &["json"], + IndexType::Scalar, + Some("json_idx".to_string()), + ¶ms, + false, + ) + .await + .unwrap(); + dataset +} + +#[rstest] +#[case::merge(false)] +#[case::append_rebuild(true)] +#[tokio::test] +async fn test_optimize_json_btree_index(#[case] append_rebuild: bool) { + let mut dataset = json_btree_dataset(vec![r#"{"val": 1000}"#]).await; + + for values in [ + vec![r#"{"val": null}"#, r#"{"val": 2000}"#], + vec![r#"{"other": 1}"#, r#"{"val": 3000}"#], + ] { + let batch = json_batch(values); + let schema = batch.schema(); + dataset + .append(RecordBatchIterator::new([Ok(batch)], schema), None) + .await + .unwrap(); + } + + let options = if append_rebuild { + OptimizeOptions::append() + } else { + OptimizeOptions::default() + }; + dataset.optimize_indices(&options).await.unwrap(); + + let indexed_fragments = dataset + .load_indices_by_name("json_idx") + .await + .unwrap() + .iter() + .flat_map(|index| index.fragment_bitmap.as_ref().unwrap().iter()) + .collect::>(); + assert_eq!(indexed_fragments, HashSet::from([0, 1, 2])); + + let result = dataset + .scan() + .filter("json_get_int(json, 'val') >= 2000") + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 2); +} + +#[tokio::test] +async fn test_optimize_append_json_btree_preserves_float_type() { + let mut dataset = json_btree_dataset(vec![r#"{"val": 1.5}"#]).await; + let appended = json_batch(vec![r#"{"val": 2}"#]); + let schema = appended.schema(); + dataset + .append(RecordBatchIterator::new([Ok(appended)], schema), None) + .await + .unwrap(); + dataset + .optimize_indices(&OptimizeOptions::append()) + .await + .unwrap(); + + let predicate = "json_get_float(json, 'val') = 2.0"; + let indexed = dataset + .scan() + .filter(predicate) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let mut baseline_scan = dataset.scan(); + baseline_scan.use_scalar_index(false); + let baseline = baseline_scan + .filter(predicate) + .unwrap() + .try_into_batch() + .await + .unwrap(); + + assert_eq!(baseline.num_rows(), 1); + assert_eq!(indexed.num_rows(), baseline.num_rows()); +} + async fn prepare_json_dataset() -> (Dataset, String) { let text_col = Arc::new(StringArray::from(vec![ r#"{