Skip to content

Commit a9ded61

Browse files
AdamGSAdamGS
committed
CR
Signed-off-by: Adam Gutglick <adam@spiraldb.com>
1 parent 5430d3a commit a9ded61

2 files changed

Lines changed: 106 additions & 101 deletions

File tree

‎vortex-json/src/arrow.rs‎

Lines changed: 106 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
//! Arrow import and export support for the JSON extension dtype.
55
66
use arrow_array::ArrayRef as ArrowArrayRef;
7-
use arrow_schema::DataType;
87
use arrow_schema::Field;
98
use arrow_schema::extension::ExtensionType;
109
use arrow_schema::extension::Json as ArrowJson;
@@ -22,7 +21,6 @@ use vortex_array::arrow::ArrowSession;
2221
use vortex_array::arrow::ArrowSessionExt;
2322
use vortex_array::arrow::FromArrowArray;
2423
use vortex_array::dtype::DType;
25-
use vortex_array::dtype::arrow::FromArrowType;
2624
use vortex_array::dtype::extension::ExtDType;
2725
use vortex_array::dtype::extension::ExtVTable;
2826
use vortex_error::VortexExpect;
@@ -54,7 +52,7 @@ impl ArrowExportVTable for Json {
5452
&self,
5553
name: &str,
5654
dtype: &DType,
57-
_session: &ArrowSession,
55+
session: &ArrowSession,
5856
) -> VortexResult<Option<Field>> {
5957
let DType::Extension(ext_dtype) = dtype else {
6058
return Ok(None);
@@ -63,7 +61,7 @@ impl ArrowExportVTable for Json {
6361
return Ok(None);
6462
}
6563

66-
let mut field = Field::new(name, DataType::Utf8, dtype.is_nullable());
64+
let mut field = session.to_arrow_field(name, ext_dtype.storage_dtype())?;
6765
field
6866
.try_with_extension_type(ArrowJson::default())
6967
.vortex_expect("Utf8 is a valid storage type for Arrow JSON");
@@ -92,14 +90,15 @@ impl ArrowExportVTable for Json {
9290
let storage_field = Field::new(
9391
String::new(),
9492
target.data_type().clone(),
95-
storage.dtype().is_nullable(),
93+
target.is_nullable(),
9694
);
9795
let session = ctx.session().clone();
98-
Ok(ArrowExport::Exported(session.arrow().execute_arrow(
99-
storage,
100-
Some(&storage_field),
101-
ctx,
102-
)?))
96+
97+
let storage = session
98+
.arrow()
99+
.execute_arrow(storage, Some(&storage_field), ctx)?;
100+
101+
Ok(ArrowExport::Exported(storage))
103102
}
104103
}
105104

@@ -113,9 +112,9 @@ impl ArrowImportVTable for Json {
113112
return Ok(None);
114113
}
115114

116-
let storage_dtype = DType::from_arrow(field);
117115
Ok(Some(DType::Extension(
118-
ExtDType::<Json>::try_new(EmptyMetadata, storage_dtype)?.erased(),
116+
ExtDType::<Json>::try_new(EmptyMetadata, DType::Utf8(field.is_nullable().into()))?
117+
.erased(),
119118
)))
120119
}
121120

@@ -138,3 +137,98 @@ impl ArrowImportVTable for Json {
138137
))
139138
}
140139
}
140+
141+
#[cfg(test)]
142+
mod tests {
143+
144+
use std::sync::Arc;
145+
146+
use arrow_array::Array;
147+
use arrow_array::ArrayRef as ArrowArrayRef;
148+
use arrow_array::StringArray;
149+
use arrow_array::cast::AsArray;
150+
use arrow_schema::DataType;
151+
use arrow_schema::Field;
152+
use arrow_schema::extension::ExtensionType;
153+
use arrow_schema::extension::Json as ArrowJson;
154+
use vortex_array::EmptyMetadata;
155+
use vortex_array::IntoArray;
156+
use vortex_array::VortexSessionExecute;
157+
use vortex_array::arrays::ExtensionArray;
158+
use vortex_array::arrays::VarBinArray;
159+
use vortex_array::arrow::ArrowSessionExt;
160+
use vortex_array::dtype::Nullability;
161+
use vortex_array::dtype::extension::ExtDType;
162+
use vortex_error::VortexExpect;
163+
use vortex_error::VortexResult;
164+
use vortex_session::VortexSession;
165+
166+
use crate::Json;
167+
use crate::initialize;
168+
169+
/// Export a JSON extension array to Arrow's canonical JSON extension.
170+
#[test]
171+
fn exports_json_extension_array_as_arrow_json() -> VortexResult<()> {
172+
let session = VortexSession::empty();
173+
initialize(&session);
174+
175+
let storage = VarBinArray::from_iter(
176+
[Some("{\"id\":1}"), Some("{\"id\":2}")],
177+
vortex_array::dtype::DType::Utf8(Nullability::NonNullable),
178+
)
179+
.into_array();
180+
let ext_dtype = ExtDType::<Json>::try_new(EmptyMetadata, storage.dtype().clone())?.erased();
181+
182+
dbg!(&ext_dtype);
183+
let array = ExtensionArray::new(ext_dtype, storage).into_array();
184+
185+
let field = session.arrow().to_arrow_field("data", array.dtype())?;
186+
assert_eq!(field.extension_type_name(), Some(ArrowJson::NAME));
187+
ArrowJson::try_new_from_field_metadata(field.data_type(), field.metadata())?;
188+
189+
dbg!(&field);
190+
191+
let exported = session.arrow().execute_arrow(
192+
array,
193+
Some(&field),
194+
&mut session.create_execution_ctx(),
195+
)?;
196+
197+
assert!(exported.data_type().is_string());
198+
199+
dbg!(exported.data_type());
200+
201+
let strings = exported.as_string_view();
202+
assert_eq!(strings.value(0), "{\"id\":1}");
203+
assert_eq!(strings.value(1), "{\"id\":2}");
204+
Ok(())
205+
}
206+
207+
/// Import Arrow's canonical JSON extension as a Vortex JSON extension array.
208+
#[test]
209+
fn imports_arrow_json_extension_array_as_vortex_json() -> VortexResult<()> {
210+
let session = VortexSession::empty();
211+
initialize(&session);
212+
213+
let mut field = Field::new("data", DataType::Utf8, false);
214+
field.try_with_extension_type(ArrowJson::default())?;
215+
let array = Arc::new(StringArray::from(vec!["{\"id\":1}", "{\"id\":2}"])) as ArrowArrayRef;
216+
217+
let imported = session.arrow().from_arrow_array(array, &field)?;
218+
let ext_dtype = imported
219+
.dtype()
220+
.as_extension_opt()
221+
.vortex_expect("expected JSON extension dtype");
222+
assert!(ext_dtype.is::<Json>());
223+
224+
let exported = session.arrow().execute_arrow(
225+
imported,
226+
Some(&field),
227+
&mut session.create_execution_ctx(),
228+
)?;
229+
let strings = exported.as_string::<i32>();
230+
assert_eq!(strings.value(0), "{\"id\":1}");
231+
assert_eq!(strings.value(1), "{\"id\":2}");
232+
Ok(())
233+
}
234+
}

‎vortex-json/src/lib.rs‎

Lines changed: 0 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -25,92 +25,3 @@ pub fn initialize(session: &VortexSession) {
2525
session.arrow().register_exporter(Arc::new(Json));
2626
session.arrow().register_importer(Arc::new(Json));
2727
}
28-
29-
#[cfg(test)]
30-
mod tests {
31-
//! Tests for JSON extension Arrow export.
32-
33-
use std::sync::Arc;
34-
35-
use arrow_array::Array;
36-
use arrow_array::ArrayRef as ArrowArrayRef;
37-
use arrow_array::StringArray;
38-
use arrow_array::cast::AsArray;
39-
use arrow_schema::DataType;
40-
use arrow_schema::Field;
41-
use arrow_schema::extension::ExtensionType;
42-
use arrow_schema::extension::Json as ArrowJson;
43-
use vortex_array::EmptyMetadata;
44-
use vortex_array::IntoArray;
45-
use vortex_array::VortexSessionExecute;
46-
use vortex_array::arrays::ExtensionArray;
47-
use vortex_array::arrays::VarBinArray;
48-
use vortex_array::arrow::ArrowSessionExt;
49-
use vortex_array::dtype::Nullability;
50-
use vortex_array::dtype::extension::ExtDType;
51-
use vortex_error::VortexExpect;
52-
use vortex_error::VortexResult;
53-
use vortex_session::VortexSession;
54-
55-
use crate::Json;
56-
use crate::initialize;
57-
58-
/// Export a JSON extension array to Arrow's canonical JSON extension.
59-
#[test]
60-
fn exports_json_extension_array_as_arrow_json() -> VortexResult<()> {
61-
let session = VortexSession::empty();
62-
initialize(&session);
63-
64-
let storage = VarBinArray::from_iter(
65-
[Some("{\"id\":1}"), Some("{\"id\":2}")],
66-
vortex_array::dtype::DType::Utf8(Nullability::NonNullable),
67-
)
68-
.into_array();
69-
let ext_dtype = ExtDType::<Json>::try_new(EmptyMetadata, storage.dtype().clone())?.erased();
70-
let array = ExtensionArray::new(ext_dtype, storage).into_array();
71-
72-
let field = session.arrow().to_arrow_field("data", array.dtype())?;
73-
assert_eq!(field.extension_type_name(), Some(ArrowJson::NAME));
74-
ArrowJson::try_new_from_field_metadata(field.data_type(), field.metadata())?;
75-
76-
let exported = session.arrow().execute_arrow(
77-
array,
78-
Some(&field),
79-
&mut session.create_execution_ctx(),
80-
)?;
81-
assert_eq!(exported.data_type(), &DataType::Utf8);
82-
83-
let strings = exported.as_string::<i32>();
84-
assert_eq!(strings.value(0), "{\"id\":1}");
85-
assert_eq!(strings.value(1), "{\"id\":2}");
86-
Ok(())
87-
}
88-
89-
/// Import Arrow's canonical JSON extension as a Vortex JSON extension array.
90-
#[test]
91-
fn imports_arrow_json_extension_array_as_vortex_json() -> VortexResult<()> {
92-
let session = VortexSession::empty();
93-
initialize(&session);
94-
95-
let mut field = Field::new("data", DataType::Utf8, false);
96-
field.try_with_extension_type(ArrowJson::default())?;
97-
let array = Arc::new(StringArray::from(vec!["{\"id\":1}", "{\"id\":2}"])) as ArrowArrayRef;
98-
99-
let imported = session.arrow().from_arrow_array(array, &field)?;
100-
let ext_dtype = imported
101-
.dtype()
102-
.as_extension_opt()
103-
.vortex_expect("expected JSON extension dtype");
104-
assert!(ext_dtype.is::<Json>());
105-
106-
let exported = session.arrow().execute_arrow(
107-
imported,
108-
Some(&field),
109-
&mut session.create_execution_ctx(),
110-
)?;
111-
let strings = exported.as_string::<i32>();
112-
assert_eq!(strings.value(0), "{\"id\":1}");
113-
assert_eq!(strings.value(1), "{\"id\":2}");
114-
Ok(())
115-
}
116-
}

0 commit comments

Comments
 (0)