diff --git a/rust/lancedb/src/table/add_data.rs b/rust/lancedb/src/table/add_data.rs index 0bfec9d68..c82b85971 100644 --- a/rust/lancedb/src/table/add_data.rs +++ b/rust/lancedb/src/table/add_data.rs @@ -1019,14 +1019,23 @@ mod tests { #[case::utf8(DataType::Utf8)] #[case::large_utf8(DataType::LargeUtf8)] #[tokio::test] - async fn test_add_arrow_json_into_lance_json_table(#[case] input_type: DataType) { + async fn test_add_arrow_json_into_lance_json_table( + #[case] input_type: DataType, + #[values(false, true)] table_has_serialized_metadata: bool, + #[values(false, true)] input_has_serialized_metadata: bool, + ) { use arrow_array::{Array, cast::AsArray}; use lance_arrow::ARROW_EXT_NAME_KEY; use lance_arrow::json::{ARROW_JSON_EXT_NAME, JSON_EXT_NAME}; // Build a table whose "data" column is lance.json (LargeBinary + // ARROW:extension:name = "lance.json"). - let lance_json_field = lance_arrow::json::json_field("data", true); + let mut lance_json_field = lance_arrow::json::json_field("data", true); + if table_has_serialized_metadata { + let mut metadata = lance_json_field.metadata().clone(); + metadata.insert(lance_arrow::ARROW_EXT_META_KEY.to_string(), String::new()); + lance_json_field = lance_json_field.with_metadata(metadata); + } let table_schema = Arc::new(Schema::new(vec![lance_json_field])); let db = connect("memory://").execute().await.unwrap(); @@ -1050,10 +1059,13 @@ mod tests { // Build an arrow.json input field (Utf8/LargeUtf8 + arrow.json extension). // This is what PyArrow produces for pa.json_() arrays. - let arrow_json_metadata = std::collections::HashMap::from([( + let mut arrow_json_metadata = std::collections::HashMap::from([( ARROW_EXT_NAME_KEY.to_string(), ARROW_JSON_EXT_NAME.to_string(), )]); + if input_has_serialized_metadata { + arrow_json_metadata.insert(lance_arrow::ARROW_EXT_META_KEY.to_string(), String::new()); + } let arrow_json_field = Field::new("data", input_type.clone(), true).with_metadata(arrow_json_metadata); let arrow_json_schema = Arc::new(Schema::new(vec![arrow_json_field])); diff --git a/rust/lancedb/src/table/datafusion/cast.rs b/rust/lancedb/src/table/datafusion/cast.rs index 0df1cb4ba..7dcf3f0c5 100644 --- a/rust/lancedb/src/table/datafusion/cast.rs +++ b/rust/lancedb/src/table/datafusion/cast.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -use std::collections::HashMap; use std::sync::{Arc, LazyLock}; use arrow_array::StructArray; @@ -59,7 +58,7 @@ fn build_field_exprs( get_input_expr: &dyn Fn(usize) -> Arc, ) -> Result, FieldRef)>> { let config = Arc::new(ConfigOptions::default()); - let mut result = Vec::new(); + let mut result: Vec<(Arc, FieldRef)> = Vec::new(); for table_field in table_fields { let Some(input_idx) = input_fields @@ -72,14 +71,7 @@ fn build_field_exprs( let input_field = &input_fields[input_idx]; let input_expr = get_input_expr(input_idx); - // PyArrow's pa.json_() is already labelled arrow.json, which is what lance-core wants - // to see, so pass it straight through. - if is_json_field(table_field) && is_arrow_json_field(input_field) { - result.push((input_expr, Arc::clone(input_field) as FieldRef)); - continue; - } - - // Anything else destined for a json column needs its JSON leaves labelled; see + // JSON inputs need their leaves labelled with the stored metadata; see // `json_write_target`. Structs are excluded because the recursion below rebuilds them // field by field, which also handles reordered and partial input. if !matches!(table_field.data_type(), DataType::Struct(_)) @@ -299,11 +291,18 @@ fn arrow_json_storage_type(input: &DataType) -> Option { } } -fn arrow_json_field(name: &str, storage: DataType, nullable: bool) -> Field { - Field::new(name, storage, nullable).with_metadata(HashMap::from([( +fn arrow_json_field(table_field: &Field, storage: DataType) -> Field { + // Empty extension metadata and absent metadata are distinct to the file + // writer. Keep the stored field's metadata while requesting JSONB encoding. + let mut metadata = table_field.metadata().clone(); + metadata.insert( ARROW_EXT_NAME_KEY.to_string(), ARROW_JSON_EXT_NAME.to_string(), - )])) + ); + table_field + .clone() + .with_data_type(storage) + .with_metadata(metadata) } /// Rewrite `table_field` so that every lance.json leaf the input supplies as text becomes an @@ -320,11 +319,7 @@ fn json_write_target(input_field: &Field, table_field: &Field) -> Option } else { arrow_json_storage_type(input_field.data_type())? }; - return Some(arrow_json_field( - table_field.name(), - storage, - table_field.is_nullable(), - )); + return Some(arrow_json_field(table_field, storage)); } if !has_json_fields(table_field) { @@ -1007,7 +1002,10 @@ mod tests { )]); let input_item = if input_labelled { - Arc::new(arrow_json_field("item", item_type, true)) + Arc::new(arrow_json_field( + &Field::new("item", item_type.clone(), true), + item_type, + )) } else { Arc::new(Field::new("item", item_type, true)) };