lance_index/scalar/
json.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{
5    collections::HashMap,
6    ops::Bound,
7    sync::{Arc, Mutex},
8};
9
10use arrow_array::{Array, LargeBinaryArray, RecordBatch, StructArray, UInt8Array};
11use arrow_schema::{DataType, Field, Field as ArrowField, Schema};
12use async_trait::async_trait;
13use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
14use datafusion::{
15    execution::SendableRecordBatchStream,
16    physical_plan::{ExecutionPlan, projection::ProjectionExec},
17};
18use datafusion_common::{ScalarValue, config::ConfigOptions};
19use datafusion_expr::{Expr, Operator, ScalarUDF};
20use datafusion_physical_expr::{
21    PhysicalExpr, ScalarFunctionExpr,
22    expressions::{Column, Literal},
23};
24use deepsize::DeepSizeOf;
25use futures::StreamExt;
26use lance_datafusion::exec::{LanceExecutionOptions, OneShotExec, get_session_context};
27use lance_datafusion::udf::json::JsonbType;
28use prost::Message;
29use roaring::RoaringBitmap;
30use serde::{Deserialize, Serialize};
31
32use lance_core::{Error, ROW_ID, Result, cache::LanceCache, error::LanceOptionExt};
33
34use crate::{
35    Index, IndexType,
36    frag_reuse::FragReuseIndex,
37    metrics::MetricsCollector,
38    registry::IndexPluginRegistry,
39    scalar::{
40        AnyQuery, CreatedIndex, IndexStore, ScalarIndex, SearchResult, UpdateCriteria,
41        expression::{IndexedExpression, ScalarIndexExpr, ScalarIndexSearch, ScalarQueryParser},
42        registry::{ScalarIndexPlugin, TrainingCriteria, TrainingRequest, VALUE_COLUMN_NAME},
43    },
44};
45
46const JSON_INDEX_VERSION: u32 = 0;
47
48/// A JSON index that indexes a field in a JSON column
49///
50/// The underlying index can be any other type of scalar index
51#[derive(Debug)]
52pub struct JsonIndex {
53    target_index: Arc<dyn ScalarIndex>,
54    path: String,
55}
56
57impl JsonIndex {
58    pub fn new(target_index: Arc<dyn ScalarIndex>, path: String) -> Self {
59        Self { target_index, path }
60    }
61}
62
63impl DeepSizeOf for JsonIndex {
64    fn deep_size_of_children(&self, context: &mut deepsize::Context) -> usize {
65        self.target_index.deep_size_of_children(context) + self.path.deep_size_of_children(context)
66    }
67}
68
69#[async_trait]
70impl Index for JsonIndex {
71    fn as_any(&self) -> &dyn std::any::Any {
72        self
73    }
74
75    fn as_index(self: Arc<Self>) -> Arc<dyn Index> {
76        self
77    }
78
79    fn as_vector_index(self: Arc<Self>) -> Result<Arc<dyn crate::vector::VectorIndex>> {
80        unimplemented!()
81    }
82
83    fn index_type(&self) -> IndexType {
84        // TODO: This causes the index to appear as btree in list_indices call.  Need better logic
85        // in list_indices to use details instead of index_type.
86        IndexType::Scalar
87    }
88
89    async fn prewarm(&self) -> Result<()> {
90        self.target_index.prewarm().await
91    }
92
93    fn statistics(&self) -> Result<serde_json::Value> {
94        todo!()
95    }
96
97    async fn calculate_included_frags(&self) -> Result<RoaringBitmap> {
98        self.target_index.calculate_included_frags().await
99    }
100}
101
102#[async_trait]
103impl ScalarIndex for JsonIndex {
104    async fn search(
105        &self,
106        query: &dyn AnyQuery,
107        metrics: &dyn MetricsCollector,
108    ) -> Result<SearchResult> {
109        let query = query.as_any().downcast_ref::<JsonQuery>().unwrap();
110        self.target_index
111            .search(query.target_query.as_ref(), metrics)
112            .await
113    }
114
115    fn can_remap(&self) -> bool {
116        self.target_index.can_remap()
117    }
118
119    async fn remap(
120        &self,
121        mapping: &HashMap<u64, Option<u64>>,
122        dest_store: &dyn IndexStore,
123    ) -> Result<CreatedIndex> {
124        let target_created = self.target_index.remap(mapping, dest_store).await?;
125        let json_details = crate::pb::JsonIndexDetails {
126            path: self.path.clone(),
127            target_details: Some(target_created.index_details),
128        };
129        Ok(CreatedIndex {
130            index_details: prost_types::Any::from_msg(&json_details)?,
131            // TODO: We should store the target index version in the details
132            index_version: JSON_INDEX_VERSION,
133        })
134    }
135
136    async fn update(
137        &self,
138        new_data: SendableRecordBatchStream,
139        dest_store: &dyn IndexStore,
140        valid_old_fragments: Option<&RoaringBitmap>,
141    ) -> Result<CreatedIndex> {
142        let target_created = self
143            .target_index
144            .update(new_data, dest_store, valid_old_fragments)
145            .await?;
146        let json_details = crate::pb::JsonIndexDetails {
147            path: self.path.clone(),
148            target_details: Some(target_created.index_details),
149        };
150        Ok(CreatedIndex {
151            index_details: prost_types::Any::from_msg(&json_details)?,
152            // TODO: We should store the target index version in the details
153            index_version: JSON_INDEX_VERSION,
154        })
155    }
156
157    fn update_criteria(&self) -> UpdateCriteria {
158        self.target_index.update_criteria()
159    }
160
161    fn derive_index_params(&self) -> Result<super::ScalarIndexParams> {
162        self.target_index.derive_index_params()
163    }
164}
165
166/// Parameters for a [`JsonIndex`]
167#[derive(Debug, Serialize, Deserialize)]
168pub struct JsonIndexParameters {
169    target_index_type: String,
170    target_index_parameters: Option<String>,
171    path: String,
172}
173
174// TODO: Do we really need to wrap the query or could we just return the target query directly?
175//
176// I think the only thing we really gain is a different format impl (e.g. it shows up as a json query
177// in the explain plan) but I don't know if that helps the user much.
178#[derive(Debug, Clone)]
179pub struct JsonQuery {
180    target_query: Arc<dyn AnyQuery>,
181    path: String,
182}
183
184impl JsonQuery {
185    pub fn new(target_query: Arc<dyn AnyQuery>, path: String) -> Self {
186        Self { target_query, path }
187    }
188}
189
190impl PartialEq for JsonQuery {
191    fn eq(&self, other: &Self) -> bool {
192        self.target_query.dyn_eq(other.target_query.as_ref()) && self.path == other.path
193    }
194}
195
196impl AnyQuery for JsonQuery {
197    fn as_any(&self) -> &dyn std::any::Any {
198        self
199    }
200
201    fn format(&self, col: &str) -> String {
202        format!("Json({}->{})", self.target_query.format(col), self.path)
203    }
204
205    fn to_expr(&self, _col: String) -> Expr {
206        todo!()
207    }
208
209    fn dyn_eq(&self, other: &dyn AnyQuery) -> bool {
210        match other.as_any().downcast_ref::<Self>() {
211            Some(o) => self == o,
212            None => false,
213        }
214    }
215}
216
217#[derive(Debug)]
218pub struct JsonQueryParser {
219    path: String,
220    target_parser: Box<dyn ScalarQueryParser>,
221}
222
223impl JsonQueryParser {
224    pub fn new(path: String, target_parser: Box<dyn ScalarQueryParser>) -> Self {
225        Self {
226            path,
227            target_parser,
228        }
229    }
230
231    fn wrap_search(&self, target_expr: IndexedExpression) -> IndexedExpression {
232        if let Some(scalar_query) = target_expr.scalar_query {
233            let scalar_query = match scalar_query {
234                ScalarIndexExpr::Query(ScalarIndexSearch {
235                    column,
236                    index_name,
237                    query,
238                    needs_recheck,
239                }) => ScalarIndexExpr::Query(ScalarIndexSearch {
240                    column,
241                    index_name,
242                    query: Arc::new(JsonQuery::new(query, self.path.clone())),
243                    needs_recheck,
244                }),
245                // This code path should only be hit on leaf expr
246                _ => unreachable!(),
247            };
248            IndexedExpression {
249                scalar_query: Some(scalar_query),
250                refine_expr: target_expr.refine_expr,
251            }
252        } else {
253            target_expr
254        }
255    }
256}
257
258impl ScalarQueryParser for JsonQueryParser {
259    fn visit_between(
260        &self,
261        column: &str,
262        low: &Bound<ScalarValue>,
263        high: &Bound<ScalarValue>,
264    ) -> Option<IndexedExpression> {
265        self.target_parser
266            .visit_between(column, low, high)
267            .map(|target_expr| self.wrap_search(target_expr))
268    }
269    fn visit_in_list(&self, column: &str, in_list: &[ScalarValue]) -> Option<IndexedExpression> {
270        self.target_parser
271            .visit_in_list(column, in_list)
272            .map(|target_expr| self.wrap_search(target_expr))
273    }
274    fn visit_is_bool(&self, column: &str, value: bool) -> Option<IndexedExpression> {
275        self.target_parser
276            .visit_is_bool(column, value)
277            .map(|target_expr| self.wrap_search(target_expr))
278    }
279    fn visit_is_null(&self, column: &str) -> Option<IndexedExpression> {
280        self.target_parser
281            .visit_is_null(column)
282            .map(|target_expr| self.wrap_search(target_expr))
283    }
284    fn visit_comparison(
285        &self,
286        column: &str,
287        value: &ScalarValue,
288        op: &Operator,
289    ) -> Option<IndexedExpression> {
290        self.target_parser
291            .visit_comparison(column, value, op)
292            .map(|target_expr| self.wrap_search(target_expr))
293    }
294    fn visit_scalar_function(
295        &self,
296        column: &str,
297        data_type: &DataType,
298        func: &ScalarUDF,
299        args: &[Expr],
300    ) -> Option<IndexedExpression> {
301        self.target_parser
302            .visit_scalar_function(column, data_type, func, args)
303            .map(|target_expr| self.wrap_search(target_expr))
304    }
305
306    // TODO: maybe we should address it by https://github.com/lance-format/lance/issues/4624
307    fn is_valid_reference(&self, func: &Expr, _data_type: &DataType) -> Option<DataType> {
308        match func {
309            Expr::ScalarFunction(udf) => {
310                // Support multiple JSON extraction functions
311                let json_functions = [
312                    "json_extract",
313                    "json_get",
314                    "json_get_int",
315                    "json_get_float",
316                    "json_get_bool",
317                    "json_get_string",
318                ];
319                if !json_functions.contains(&udf.name()) {
320                    return None;
321                }
322                if udf.args.len() != 2 {
323                    return None;
324                }
325                // We already know index 0 is a column reference to the column so we just need to
326                // ensure that index 1 matches our path
327                match &udf.args[1] {
328                    Expr::Literal(ScalarValue::Utf8(Some(path)), _) => {
329                        if path == &self.path {
330                            // Return the appropriate type based on the function
331                            match udf.name() {
332                                "json_get_int" => Some(DataType::Int64),
333                                "json_get_float" => Some(DataType::Float64),
334                                "json_get_bool" => Some(DataType::Boolean),
335                                "json_get_string" | "json_extract" => Some(DataType::Utf8),
336                                _ => None,
337                            }
338                        } else {
339                            None
340                        }
341                    }
342                    _ => None,
343                }
344            }
345            _ => None,
346        }
347    }
348}
349
350pub struct JsonTrainingRequest {
351    parameters: JsonIndexParameters,
352    target_request: Box<dyn TrainingRequest>,
353}
354
355impl JsonTrainingRequest {
356    pub fn new(parameters: JsonIndexParameters, target_request: Box<dyn TrainingRequest>) -> Self {
357        Self {
358            parameters,
359            target_request,
360        }
361    }
362}
363
364impl TrainingRequest for JsonTrainingRequest {
365    fn as_any(&self) -> &dyn std::any::Any {
366        self
367    }
368
369    fn criteria(&self) -> &TrainingCriteria {
370        self.target_request.criteria()
371    }
372}
373
374/// Plugin implementation for a [`JsonIndex`]
375#[derive(Default)]
376pub struct JsonIndexPlugin {
377    registry: Mutex<Option<Arc<IndexPluginRegistry>>>,
378}
379
380impl std::fmt::Debug for JsonIndexPlugin {
381    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
382        write!(f, "JsonIndexPlugin")
383    }
384}
385
386impl JsonIndexPlugin {
387    fn registry(&self) -> Result<Arc<IndexPluginRegistry>> {
388        Ok(self.registry.lock().unwrap().as_ref().expect_ok()?.clone())
389    }
390
391    /// Extract JSON with type information using the new UDF
392    async fn extract_json_with_type_info(
393        data: SendableRecordBatchStream,
394        path: String,
395    ) -> Result<(SendableRecordBatchStream, DataType)> {
396        let input = Arc::new(OneShotExec::new(data));
397        let input_schema = input.schema();
398        let value_column_idx = input_schema
399            .column_with_name(VALUE_COLUMN_NAME)
400            .expect_ok()?
401            .0;
402        let row_id_column_idx = input_schema.column_with_name(ROW_ID).expect_ok()?.0;
403
404        // Call json_extract_with_type UDF
405        let exprs = vec![
406            (
407                Arc::new(ScalarFunctionExpr::try_new(
408                    Arc::new(lance_datafusion::udf::json::json_extract_with_type_udf()),
409                    vec![
410                        Arc::new(Column::new(VALUE_COLUMN_NAME, value_column_idx)),
411                        Arc::new(Literal::new(ScalarValue::Utf8(Some(path)))),
412                    ],
413                    &input_schema,
414                    Arc::new(ConfigOptions::default()),
415                )?) as Arc<dyn PhysicalExpr>,
416                "json_result".to_string(),
417            ),
418            (
419                Arc::new(Column::new(ROW_ID, row_id_column_idx)) as Arc<dyn PhysicalExpr>,
420                ROW_ID.to_string(),
421            ),
422        ];
423
424        let project = ProjectionExec::try_new(exprs, input)?;
425        let ctx = get_session_context(&LanceExecutionOptions::default());
426        let mut stream = project.execute(0, ctx.task_ctx())?;
427
428        // Collect batches and determine type from first non-null value
429        let mut all_batches = Vec::new();
430        let mut inferred_type: Option<DataType> = None;
431
432        while let Some(batch_result) = stream.next().await {
433            let batch = batch_result?;
434
435            // Determine type from first non-null value if not yet set
436            if inferred_type.is_none()
437                && let Some(json_result_column) = batch.column_by_name("json_result")
438                && let Some(struct_array) =
439                    json_result_column.as_any().downcast_ref::<StructArray>()
440                && let Some(type_array) = struct_array.column_by_name("type_tag")
441                && let Some(uint8_array) = type_array.as_any().downcast_ref::<UInt8Array>()
442            {
443                // Find first non-null value to determine type
444                for i in 0..uint8_array.len() {
445                    if !uint8_array.is_null(i) {
446                        let type_tag = uint8_array.value(i);
447                        let jsonb_type = JsonbType::from_u8(type_tag).ok_or_else(|| {
448                            Error::invalid_input_source(
449                                format!("Invalid type tag: {}", type_tag).into(),
450                            )
451                        })?;
452
453                        // Map JsonbType to Arrow DataType
454                        inferred_type = Some(match jsonb_type {
455                            JsonbType::Null => continue, // Skip null values
456                            JsonbType::Boolean => DataType::Boolean,
457                            JsonbType::Int64 => DataType::Int64,
458                            JsonbType::Float64 => DataType::Float64,
459                            JsonbType::String => DataType::Utf8,
460                            JsonbType::Array => DataType::LargeBinary,
461                            JsonbType::Object => DataType::LargeBinary,
462                        });
463                        break;
464                    }
465                }
466            }
467
468            all_batches.push(batch);
469        }
470
471        // If no type was inferred (all nulls), default to String
472        let inferred_type = inferred_type.unwrap_or(DataType::Utf8);
473
474        // Recreate stream from collected batches
475        let schema = all_batches
476            .first()
477            .map(|b| b.schema())
478            .ok_or_else(|| Error::invalid_input_source("No batches in stream".into()))?;
479
480        let recreated_stream = Box::pin(RecordBatchStreamAdapter::new(
481            schema,
482            futures::stream::iter(all_batches.into_iter().map(Ok)),
483        )) as SendableRecordBatchStream;
484
485        Ok((recreated_stream, inferred_type))
486    }
487
488    /// Convert the stream with JSONB values and type tags to properly typed values
489    async fn convert_stream_by_type(
490        data: SendableRecordBatchStream,
491        target_type: DataType,
492    ) -> Result<SendableRecordBatchStream> {
493        let input = Arc::new(OneShotExec::new(data));
494        let _input_schema = input.schema();
495        let ctx = get_session_context(&LanceExecutionOptions::default());
496        let mut stream = input.execute(0, ctx.task_ctx())?;
497
498        let mut converted_batches = Vec::new();
499
500        while let Some(batch_result) = stream.next().await {
501            let batch = batch_result?;
502
503            // Extract the struct column containing value and type_tag
504            let json_result_column = batch
505                .column_by_name("json_result")
506                .ok_or_else(|| Error::invalid_input_source("Missing json_result column".into()))?;
507
508            let struct_array = json_result_column
509                .as_any()
510                .downcast_ref::<StructArray>()
511                .ok_or_else(|| Error::invalid_input_source("json_result is not a struct".into()))?;
512
513            let value_array = struct_array.column_by_name("value").ok_or_else(|| {
514                Error::invalid_input_source("Missing value column in struct".into())
515            })?;
516
517            let binary_array = value_array
518                .as_any()
519                .downcast_ref::<LargeBinaryArray>()
520                .ok_or_else(|| Error::invalid_input_source("value is not LargeBinary".into()))?;
521
522            // Convert based on target type using serde deserialization
523            let converted_array: Arc<dyn Array> =
524                match target_type {
525                    DataType::Boolean => {
526                        let mut builder =
527                            arrow_array::builder::BooleanBuilder::with_capacity(binary_array.len());
528                        for i in 0..binary_array.len() {
529                            if binary_array.is_null(i) {
530                                builder.append_null();
531                            } else if let Some(bytes) = binary_array.value(i).into() {
532                                let raw_jsonb = jsonb::RawJsonb::new(bytes);
533                                // Try to deserialize directly to bool
534                                match jsonb::from_raw_jsonb::<bool>(&raw_jsonb) {
535                                    Ok(bool_val) => builder.append_value(bool_val),
536                                    Err(e) => {
537                                        return Err(Error::invalid_input_source(format!(
538                                        "Failed to deserialize JSONB to bool at index {}: {}",
539                                        i, e
540                                    )
541                                    .into()));
542                                    }
543                                }
544                            } else {
545                                builder.append_null();
546                            }
547                        }
548                        Arc::new(builder.finish())
549                    }
550                    DataType::Int64 => {
551                        let mut builder =
552                            arrow_array::builder::Int64Builder::with_capacity(binary_array.len());
553                        for i in 0..binary_array.len() {
554                            if binary_array.is_null(i) {
555                                builder.append_null();
556                            } else if let Some(bytes) = binary_array.value(i).into() {
557                                let raw_jsonb = jsonb::RawJsonb::new(bytes);
558                                // Try to deserialize directly to i64
559                                match jsonb::from_raw_jsonb::<i64>(&raw_jsonb) {
560                                    Ok(int_val) => builder.append_value(int_val),
561                                    Err(e) => {
562                                        return Err(Error::invalid_input_source(format!(
563                                        "Failed to deserialize JSONB to i64 at index {}: {}",
564                                        i, e
565                                    )
566                                    .into()));
567                                    }
568                                }
569                            } else {
570                                builder.append_null();
571                            }
572                        }
573                        Arc::new(builder.finish())
574                    }
575                    DataType::Float64 => {
576                        let mut builder =
577                            arrow_array::builder::Float64Builder::with_capacity(binary_array.len());
578                        for i in 0..binary_array.len() {
579                            if binary_array.is_null(i) {
580                                builder.append_null();
581                            } else if let Some(bytes) = binary_array.value(i).into() {
582                                let raw_jsonb = jsonb::RawJsonb::new(bytes);
583                                // Try to deserialize directly to f64 (serde handles int->float conversion)
584                                match jsonb::from_raw_jsonb::<f64>(&raw_jsonb) {
585                                    Ok(float_val) => builder.append_value(float_val),
586                                    Err(e) => {
587                                        return Err(Error::invalid_input_source(format!(
588                                        "Failed to deserialize JSONB to f64 at index {}: {}",
589                                        i, e
590                                    )
591                                    .into()));
592                                    }
593                                }
594                            } else {
595                                builder.append_null();
596                            }
597                        }
598                        Arc::new(builder.finish())
599                    }
600                    DataType::Utf8 => {
601                        let mut builder = arrow_array::builder::StringBuilder::with_capacity(
602                            binary_array.len(),
603                            1024,
604                        );
605                        for i in 0..binary_array.len() {
606                            if binary_array.is_null(i) {
607                                builder.append_null();
608                            } else if let Some(bytes) = binary_array.value(i).into() {
609                                let raw_jsonb = jsonb::RawJsonb::new(bytes);
610                                // Try to deserialize to String, or use to_string() for any type
611                                match jsonb::from_raw_jsonb::<String>(&raw_jsonb) {
612                                    Ok(str_val) => builder.append_value(&str_val),
613                                    Err(_) => {
614                                        // For non-string types, convert to string representation
615                                        builder.append_value(raw_jsonb.to_string());
616                                    }
617                                }
618                            } else {
619                                builder.append_null();
620                            }
621                        }
622                        Arc::new(builder.finish())
623                    }
624                    DataType::LargeBinary => {
625                        // Keep as binary for array/object types
626                        value_array.clone()
627                    }
628                    _ => {
629                        return Err(Error::invalid_input_source(
630                            format!("Unsupported target type: {:?}", target_type).into(),
631                        ));
632                    }
633                };
634
635            // Get row_id column
636            let row_id_column = batch
637                .column_by_name(ROW_ID)
638                .ok_or_else(|| Error::invalid_input_source("Missing row_id column".into()))?
639                .clone();
640
641            // Create new batch with converted values
642            let new_schema = Arc::new(Schema::new(vec![
643                ArrowField::new(VALUE_COLUMN_NAME, target_type.clone(), true),
644                ArrowField::new(ROW_ID, DataType::UInt64, false),
645            ]));
646
647            let new_batch =
648                RecordBatch::try_new(new_schema.clone(), vec![converted_array, row_id_column])?;
649
650            converted_batches.push(new_batch);
651        }
652
653        // Create stream from converted batches
654        let schema = converted_batches
655            .first()
656            .map(|b| b.schema())
657            .ok_or_else(|| Error::invalid_input_source("No batches to convert".into()))?;
658
659        Ok(Box::pin(RecordBatchStreamAdapter::new(
660            schema,
661            futures::stream::iter(converted_batches.into_iter().map(Ok)),
662        )))
663    }
664}
665
666#[async_trait]
667impl ScalarIndexPlugin for JsonIndexPlugin {
668    fn name(&self) -> &str {
669        "Json"
670    }
671
672    fn new_training_request(
673        &self,
674        params: &str,
675        field: &Field,
676    ) -> Result<Box<dyn TrainingRequest>> {
677        if !matches!(field.data_type(), DataType::Binary | DataType::LargeBinary) {
678            return Err(Error::invalid_input_source(
679                "A JSON index can only be created on a Binary or LargeBinary field.".into(),
680            ));
681        }
682
683        // Initially use Utf8, will be refined during training with type inference
684        let target_type = DataType::Utf8;
685
686        let params = serde_json::from_str::<JsonIndexParameters>(params)?;
687        let registry = self.registry()?;
688        let target_plugin = registry.get_plugin_by_name(&params.target_index_type)?;
689        let target_request = target_plugin.new_training_request(
690            params.target_index_parameters.as_deref().unwrap_or("{}"),
691            &Field::new("", target_type, true),
692        )?;
693
694        Ok(Box::new(JsonTrainingRequest::new(params, target_request)))
695    }
696
697    fn provides_exact_answer(&self) -> bool {
698        // TODO: Need to lookup target plugin via details to figure this out correctly
699        true
700    }
701
702    fn attach_registry(&self, registry: Arc<IndexPluginRegistry>) {
703        let mut reg_ref = self.registry.lock().unwrap();
704        *reg_ref = Some(registry);
705    }
706
707    fn version(&self) -> u32 {
708        JSON_INDEX_VERSION
709    }
710
711    fn new_query_parser(
712        &self,
713        index_name: String,
714        index_details: &prost_types::Any,
715    ) -> Option<Box<dyn ScalarQueryParser>> {
716        // TODO: Allow return Result here
717        let registry = self.registry().unwrap();
718        let json_details =
719            crate::pb::JsonIndexDetails::decode(index_details.value.as_slice()).unwrap();
720        let target_details = json_details.target_details.as_ref().expect_ok().unwrap();
721        let target_plugin = registry.get_plugin_by_details(target_details).unwrap();
722        // TODO: Use something like ${index_name}_${path} for the index name?  Don't have access to path here tho
723        let target_parser = target_plugin.new_query_parser(index_name, index_details)?;
724        Some(Box::new(JsonQueryParser::new(
725            json_details.path.clone(),
726            target_parser,
727        )) as Box<dyn ScalarQueryParser>)
728    }
729
730    async fn train_index(
731        &self,
732        data: SendableRecordBatchStream,
733        index_store: &dyn IndexStore,
734        request: Box<dyn TrainingRequest>,
735        fragment_ids: Option<Vec<u32>>,
736        progress: Arc<dyn crate::progress::IndexBuildProgress>,
737    ) -> Result<CreatedIndex> {
738        let request = (request as Box<dyn std::any::Any>)
739            .downcast::<JsonTrainingRequest>()
740            .unwrap();
741        let path = request.parameters.path.clone();
742
743        // Extract JSON with type information
744        let (data_stream, inferred_type) =
745            Self::extract_json_with_type_info(data, path.clone()).await?;
746
747        // Convert the stream to properly typed values based on inferred type
748        let converted_stream =
749            Self::convert_stream_by_type(data_stream, inferred_type.clone()).await?;
750
751        // Update the target request with inferred type
752        let registry = self.registry()?;
753        let target_plugin = registry.get_plugin_by_name(&request.parameters.target_index_type)?;
754
755        // Create a new training request with the inferred type
756        let target_request = target_plugin.new_training_request(
757            request
758                .parameters
759                .target_index_parameters
760                .as_deref()
761                .unwrap_or("{}"),
762            &Field::new("", inferred_type, true),
763        )?;
764
765        let target_index = target_plugin
766            .train_index(
767                converted_stream,
768                index_store,
769                target_request,
770                fragment_ids,
771                progress,
772            )
773            .await?;
774
775        let index_details = crate::pb::JsonIndexDetails {
776            path,
777            target_details: Some(target_index.index_details),
778        };
779        Ok(CreatedIndex {
780            index_details: prost_types::Any::from_msg(&index_details)?,
781            index_version: JSON_INDEX_VERSION,
782        })
783    }
784
785    async fn load_index(
786        &self,
787        index_store: Arc<dyn IndexStore>,
788        index_details: &prost_types::Any,
789        frag_reuse_index: Option<Arc<FragReuseIndex>>,
790        cache: &LanceCache,
791    ) -> Result<Arc<dyn ScalarIndex>> {
792        let registry = self.registry().unwrap();
793        let json_details = crate::pb::JsonIndexDetails::decode(index_details.value.as_slice())?;
794        let target_details = json_details.target_details.as_ref().expect_ok()?;
795        let target_plugin = registry.get_plugin_by_details(target_details).unwrap();
796        let target_index = target_plugin
797            .load_index(index_store, target_details, frag_reuse_index, cache)
798            .await?;
799        Ok(Arc::new(JsonIndex::new(target_index, json_details.path)))
800    }
801
802    fn details_as_json(&self, details: &prost_types::Any) -> Result<serde_json::Value> {
803        let registry = self.registry().unwrap();
804        let json_details = crate::pb::JsonIndexDetails::decode(details.value.as_slice())?;
805        let target_details = json_details.target_details.as_ref().expect_ok()?;
806        let target_plugin = registry.get_plugin_by_details(target_details).unwrap();
807        let target_details_json = target_plugin.details_as_json(target_details)?;
808        Ok(serde_json::json!({
809            "path": json_details.path,
810            "target_details": target_details_json,
811        }))
812    }
813}
814
815#[cfg(test)]
816mod tests {
817    use super::*;
818    use arrow_array::{ArrayRef, RecordBatch};
819    use arrow_schema::{DataType, Field, Schema};
820    use std::sync::Arc;
821
822    // Note: The old test_detect_json_value_type test has been removed as we now use
823    // JSONB's inherent type information instead of string-based type detection
824
825    #[tokio::test]
826    async fn test_json_extract_with_type_info() {
827        use arrow_array::{LargeBinaryArray, UInt64Array};
828        use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
829        use futures::stream;
830
831        // Create test JSONB data
832        let json_data = vec![
833            r#"{"name": "Alice", "age": 30, "active": true}"#,
834            r#"{"name": "Bob", "age": 25, "active": false}"#,
835            r#"{"name": "Charlie", "age": 35, "active": true}"#,
836        ];
837
838        // Convert JSON strings to JSONB binary format
839        let mut jsonb_values = Vec::new();
840        for json_str in &json_data {
841            let owned_jsonb: jsonb::OwnedJsonb = json_str.parse().unwrap();
842            jsonb_values.push(Some(owned_jsonb.to_vec()));
843        }
844
845        // Create test batch with JSONB data
846        let schema = Arc::new(Schema::new(vec![
847            Field::new(VALUE_COLUMN_NAME, DataType::LargeBinary, true),
848            Field::new(ROW_ID, DataType::UInt64, false),
849        ]));
850
851        let jsonb_array = LargeBinaryArray::from(
852            jsonb_values
853                .iter()
854                .map(|v| v.as_deref())
855                .collect::<Vec<_>>(),
856        );
857        let row_ids = UInt64Array::from(vec![1, 2, 3]);
858
859        let batch = RecordBatch::try_new(
860            schema.clone(),
861            vec![
862                Arc::new(jsonb_array) as ArrayRef,
863                Arc::new(row_ids) as ArrayRef,
864            ],
865        )
866        .unwrap();
867
868        let stream = Box::pin(RecordBatchStreamAdapter::new(
869            schema.clone(),
870            stream::iter(vec![Ok(batch)]),
871        )) as SendableRecordBatchStream;
872
873        // Test type inference for integer field
874        let (_result_stream, inferred_type) =
875            JsonIndexPlugin::extract_json_with_type_info(stream, "$.age".to_string())
876                .await
877                .unwrap();
878
879        assert_eq!(inferred_type, DataType::Int64);
880
881        // Create new test stream for boolean field
882        let batch2 = RecordBatch::try_new(
883            schema.clone(),
884            vec![
885                Arc::new(LargeBinaryArray::from(vec![
886                    json_data[0]
887                        .parse::<jsonb::OwnedJsonb>()
888                        .ok()
889                        .map(|j| j.to_vec())
890                        .as_deref(),
891                    json_data[1]
892                        .parse::<jsonb::OwnedJsonb>()
893                        .ok()
894                        .map(|j| j.to_vec())
895                        .as_deref(),
896                    json_data[2]
897                        .parse::<jsonb::OwnedJsonb>()
898                        .ok()
899                        .map(|j| j.to_vec())
900                        .as_deref(),
901                ])) as ArrayRef,
902                Arc::new(UInt64Array::from(vec![1, 2, 3])) as ArrayRef,
903            ],
904        )
905        .unwrap();
906
907        let stream2 = Box::pin(RecordBatchStreamAdapter::new(
908            schema.clone(),
909            stream::iter(vec![Ok(batch2)]),
910        )) as SendableRecordBatchStream;
911
912        // Test type inference for boolean field
913        let (_, inferred_type) =
914            JsonIndexPlugin::extract_json_with_type_info(stream2, "$.active".to_string())
915                .await
916                .unwrap();
917
918        assert_eq!(inferred_type, DataType::Boolean);
919
920        // Create test stream for string field
921        let batch3 = RecordBatch::try_new(
922            schema.clone(),
923            vec![
924                Arc::new(LargeBinaryArray::from(vec![
925                    json_data[0]
926                        .parse::<jsonb::OwnedJsonb>()
927                        .ok()
928                        .map(|j| j.to_vec())
929                        .as_deref(),
930                    json_data[1]
931                        .parse::<jsonb::OwnedJsonb>()
932                        .ok()
933                        .map(|j| j.to_vec())
934                        .as_deref(),
935                    json_data[2]
936                        .parse::<jsonb::OwnedJsonb>()
937                        .ok()
938                        .map(|j| j.to_vec())
939                        .as_deref(),
940                ])) as ArrayRef,
941                Arc::new(UInt64Array::from(vec![1, 2, 3])) as ArrayRef,
942            ],
943        )
944        .unwrap();
945
946        let stream3 = Box::pin(RecordBatchStreamAdapter::new(
947            schema,
948            stream::iter(vec![Ok(batch3)]),
949        )) as SendableRecordBatchStream;
950
951        // Test type inference for string field
952        let (_, inferred_type) =
953            JsonIndexPlugin::extract_json_with_type_info(stream3, "$.name".to_string())
954                .await
955                .unwrap();
956
957        assert_eq!(inferred_type, DataType::Utf8);
958    }
959}
lance_index/scalar/json.rs

lance_index/scalar/
json.rs
