Skip to content

refactor: merge JsonConvertingReader into SchemaAdapter #7678

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xloya wants to merge 1 commit into
base: main
Choose a base branch
from
+61 −87
Open

refactor: merge JsonConvertingReader into SchemaAdapter #7678

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 19 additions & 71 deletions rust/lance/src/dataset/fragment.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ use datafusion::logical_expr::Expr;
use datafusion::scalar::ScalarValue;
use futures::future::{BoxFuture, try_join_all};
use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt, join, stream};
use lance_arrow::json::{convert_json_columns, has_json_fields, is_arrow_json_field};
use lance_arrow::{RecordBatchExt, SchemaExt};
use lance_core::datatypes::{OnMissing, OnTypeMismatch, SchemaCompareOptions};
use lance_core::utils::address::RowAddress;
Expand Down Expand Up @@ -65,6 +64,7 @@ use super::updater::Updater;
use super::{NewColumnTransform, WriteParams, schema_evolution};
use crate::dataset::Dataset;
use crate::dataset::fragment::session::FragmentSession;
use crate::dataset::utils::SchemaAdapter;
use crate::io::deletion::read_dataset_deletion_file;

/// Result of [`FileFragment::update_columns_with_offsets`]: updated fragment metadata, modified field ids,
Expand Down Expand Up @@ -1907,17 +1907,10 @@ impl FileFragment {
)
.await?;
// Hash join: rows matched on the right-hand stream rewrite columns; track physical offsets via `_rowaddr`.
// Convert Arrow JSON columns (Utf8) to Lance JSON (LargeBinary) in the right stream
// so they match the physical storage format read from the fragment's left batch.
let right_stream: Box<dyn RecordBatchReader + Send> = if right_schema
.fields()
.iter()
.any(|f| is_arrow_json_field(f) || has_json_fields(f))
{
Box::new(JsonConvertingReader::new(right_stream))
} else {
right_stream
};
// Convert the right stream from its logical form (Arrow JSON, view types)
// to the physical form stored on disk so it matches the fragment's left batch.
let right_stream =
SchemaAdapter::new(right_schema.clone()).to_physical_reader(right_stream);
let joiner = Arc::new(HashJoiner::try_new(right_stream, right_on).await?);
let mut matched_offsets = RoaringBitmap::new();
let frag_id_u32 = u32::try_from(self.metadata.id).map_err(|_| {
Expand Down Expand Up @@ -2958,59 +2951,6 @@ impl FragmentReader {
}
}

/// A wrapper around a `RecordBatchReader` that converts Arrow JSON columns
/// (Utf8/LargeUtf8 with `arrow.json` extension) to Lance JSON columns
/// (LargeBinary with `lance.json` extension / JSONB format).
///
/// This is needed when user-provided data contains Arrow JSON fields but the
/// dataset stores them in Lance's JSONB binary format.
struct JsonConvertingReader {
inner: Box<dyn RecordBatchReader + Send>,
schema: arrow_schema::SchemaRef,
}

impl JsonConvertingReader {
fn new(inner: Box<dyn RecordBatchReader + Send>) -> Self {
use lance_arrow::json::arrow_json_to_lance_json;

// Build the converted schema (Arrow JSON fields → Lance JSON fields)
let orig_schema = inner.schema();
let new_fields: Vec<arrow_schema::FieldRef> = orig_schema
.fields()
.iter()
.map(|f| {
if is_arrow_json_field(f) || has_json_fields(f) {
Arc::new(arrow_json_to_lance_json(f))
} else {
Arc::clone(f)
}
})
.collect();
let schema = Arc::new(arrow_schema::Schema::new_with_metadata(
new_fields,
orig_schema.metadata().clone(),
));

Self { inner, schema }
}
}

impl Iterator for JsonConvertingReader {
type Item = std::result::Result<RecordBatch, arrow_schema::ArrowError>;

fn next(&mut self) -> Option<Self::Item> {
self.inner
.next()
.map(|result| result.and_then(|batch| convert_json_columns(&batch)))
}
}

impl RecordBatchReader for JsonConvertingReader {
fn schema(&self) -> arrow_schema::SchemaRef {
self.schema.clone()
}
}

#[cfg(test)]
mod tests {
use arrow_arith::numeric::mul;
Expand Down Expand Up @@ -4488,13 +4428,16 @@ mod tests {

#[tokio::test]
async fn test_update_columns_with_json_extension_type() {
use arrow_array::UInt64Array;
use arrow_array::{StringViewArray, UInt64Array};
use lance_arrow::ARROW_EXT_NAME_KEY;
use lance_arrow::json::ARROW_JSON_EXT_NAME;
use lance_core::ROW_ID;
use std::collections::HashMap;

// Create a dataset with an Arrow JSON extension column
// Create a dataset with an Arrow JSON extension column and a Utf8View
// column. Both are logical types that Lance stores in a different
// physical form (JSONB LargeBinary / Utf8), so the update path must
// convert the right-hand stream to match.
let test_dir = TempStrDir::default();
let mut json_metadata = HashMap::new();
json_metadata.insert(
Expand All @@ -4505,6 +4448,7 @@ mod tests {
ArrowField::new("id", DataType::Int64, false),
ArrowField::new("name", DataType::Utf8, true),
ArrowField::new("meta", DataType::Utf8, true).with_metadata(json_metadata.clone()),
ArrowField::new("desc", DataType::Utf8View, true),
]));
let batch = RecordBatch::try_new(
schema.clone(),
Expand All @@ -4518,6 +4462,7 @@ mod tests {
r#"{"x":4}"#,
r#"{"x":5}"#,
])),
Arc::new(StringViewArray::from(vec!["d1", "d2", "d3", "d4", "d5"])),
],
)
.unwrap();
Expand All @@ -4526,11 +4471,12 @@ mod tests {
.await
.unwrap();

// Build the right stream with Arrow JSON column (Utf8 + arrow.json extension)
// Only update rows with row_id 1 and 3
// Build the right stream with an Arrow JSON column (Utf8 + arrow.json
// extension) and a Utf8View column. Only update rows with row_id 1 and 3.
let update_schema = Arc::new(ArrowSchema::new(vec![
ArrowField::new(ROW_ID, DataType::UInt64, false),
ArrowField::new("meta", DataType::Utf8, true).with_metadata(json_metadata),
ArrowField::new("desc", DataType::Utf8View, true),
]));
let update_batch = RecordBatch::try_new(
update_schema.clone(),
Expand All @@ -4540,6 +4486,7 @@ mod tests {
r#"{"updated":true,"id":2}"#,
r#"{"updated":true,"id":4}"#,
])),
Arc::new(StringViewArray::from(vec!["d2-new", "d4-new"])),
],
)
.unwrap();
Expand All @@ -4548,9 +4495,10 @@ mod tests {
update_schema,
));

// Perform update_columns - this should NOT fail with type mismatch
// Previously this would error with:
// Perform update_columns - this should NOT fail with a type mismatch.
// Previously the JSON column would error with:
// "It is not possible to interleave arrays of different data types (Utf8 and LargeBinary)"
// and the view column would likewise mismatch (Utf8View vs stored Utf8).
let mut fragment = dataset.get_fragment(0).unwrap();
let (updated_fragment, fields_modified) = fragment
.update_columns(right_stream, ROW_ID, ROW_ID)
Expand Down
58 changes: 42 additions & 16 deletions rust/lance/src/dataset/utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright The Lance Authors

use crate::Result;
use arrow_array::{ArrayRef, RecordBatch, UInt64Array};
use arrow_array::{ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader, UInt64Array};
use arrow_schema::{
DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef,
};
Expand Down Expand Up @@ -225,18 +225,12 @@ impl SchemaAdapter {
}
}

/// Convert a logical stream into a physical stream.
pub fn to_physical_stream(
&self,
stream: SendableRecordBatchStream,
) -> SendableRecordBatchStream {
if !self.requires_physical_conversion() {
return stream;
}

let arrow_schema = stream.schema();
let mut new_fields = Vec::with_capacity(arrow_schema.fields().len());
for field in arrow_schema.fields() {
/// Build the physical Arrow schema for `logical_schema`: Arrow JSON fields
/// become Lance JSON fields and view types are downcast to their classic
/// offset equivalents. Fields needing no conversion are left unchanged.
fn physical_schema(logical_schema: &ArrowSchemaRef) -> ArrowSchemaRef {
let mut new_fields = Vec::with_capacity(logical_schema.fields().len());
for field in logical_schema.fields() {
if has_arrow_json_fields(field) {
new_fields.push(Arc::new(arrow_json_to_lance_json(field)));
} else if let Some(phys) = physical_field(field) {
Expand All @@ -245,10 +239,42 @@ impl SchemaAdapter {
new_fields.push(Arc::clone(field));
}
}
let converted_schema = Arc::new(ArrowSchema::new_with_metadata(
Arc::new(ArrowSchema::new_with_metadata(
new_fields,
arrow_schema.metadata().clone(),
));
logical_schema.metadata().clone(),
))
}

/// Wrap a synchronous [`RecordBatchReader`] so each batch is converted from
/// its logical form to the physical form Lance stores on disk (Arrow JSON →
/// Lance JSON, view types → offset types). Returns the reader unchanged when
/// no field needs conversion.
pub fn to_physical_reader(
&self,
reader: Box<dyn RecordBatchReader + Send>,
) -> Box<dyn RecordBatchReader + Send> {
if !self.requires_physical_conversion() {
return reader;
}
let schema = Self::physical_schema(&reader.schema());
let converted = reader.map(|batch| {
let batch = batch?;
let batch = convert_json_columns(&batch)?;
downcast_view_columns(&batch)
});
Box::new(RecordBatchIterator::new(converted, schema))
}

/// Convert a logical stream into a physical stream.
pub fn to_physical_stream(
&self,
stream: SendableRecordBatchStream,
) -> SendableRecordBatchStream {
if !self.requires_physical_conversion() {
return stream;
}

let converted_schema = Self::physical_schema(&stream.schema());

let converted_stream = stream.map(move |batch_result| {
batch_result.and_then(|batch| {
Expand Down
Loading