vortex_json/
arrow.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Arrow import and export support for the JSON extension dtype.
5
6use arrow_array::ArrayRef as ArrowArrayRef;
7use arrow_schema::Field;
8use arrow_schema::extension::ExtensionType;
9use arrow_schema::extension::Json as ArrowJson;
10use vortex_array::ArrayRef;
11use vortex_array::EmptyMetadata;
12use vortex_array::ExecutionCtx;
13use vortex_array::IntoArray;
14use vortex_array::arrays::ExtensionArray;
15use vortex_array::arrays::extension::ExtensionArrayExt;
16use vortex_array::arrow::ArrowExport;
17use vortex_array::arrow::ArrowExportVTable;
18use vortex_array::arrow::ArrowImport;
19use vortex_array::arrow::ArrowImportVTable;
20use vortex_array::arrow::ArrowSession;
21use vortex_array::arrow::ArrowSessionExt;
22use vortex_array::arrow::FromArrowArray;
23use vortex_array::dtype::DType;
24use vortex_array::dtype::extension::ExtDType;
25use vortex_array::dtype::extension::ExtVTable;
26use vortex_error::VortexExpect;
27use vortex_error::VortexResult;
28use vortex_session::registry::CachedId;
29use vortex_session::registry::Id;
30
31use crate::Json;
32
33/// Arrow's canonical JSON extension name cached as a registry id.
34static ARROW_JSON: CachedId = CachedId::new(ArrowJson::NAME);
35
36/// Returns whether an Arrow field contains valid canonical JSON extension metadata.
37fn has_valid_json_extension(field: &Field) -> bool {
38    field.extension_type_name() == Some(ArrowJson::NAME)
39        && ArrowJson::try_new_from_field_metadata(field.data_type(), field.metadata()).is_ok()
40}
41
42impl ArrowExportVTable for Json {
43    fn arrow_ext_id(&self) -> Id {
44        *ARROW_JSON
45    }
46
47    fn vortex_id(&self) -> Id {
48        Json.id()
49    }
50
51    fn to_arrow_field(
52        &self,
53        name: &str,
54        dtype: &DType,
55        session: &ArrowSession,
56    ) -> VortexResult<Option<Field>> {
57        let DType::Extension(ext_dtype) = dtype else {
58            return Ok(None);
59        };
60        if !ext_dtype.is::<Json>() {
61            return Ok(None);
62        }
63
64        let mut field = session.to_arrow_field(name, ext_dtype.storage_dtype())?;
65        field
66            .try_with_extension_type(ArrowJson::default())
67            .vortex_expect("Utf8 is a valid storage type for Arrow JSON");
68        Ok(Some(field))
69    }
70
71    fn execute_arrow(
72        &self,
73        array: ArrayRef,
74        target: &Field,
75        ctx: &mut ExecutionCtx,
76    ) -> VortexResult<ArrowExport> {
77        let is_json = array
78            .dtype()
79            .as_extension_opt()
80            .map(|ext_dtype| ext_dtype.is::<Json>())
81            .unwrap_or(false);
82        if !is_json {
83            return Ok(ArrowExport::Unsupported(array));
84        }
85
86        ArrowJson::try_new_from_field_metadata(target.data_type(), target.metadata())?;
87
88        let executed = array.execute::<ExtensionArray>(ctx)?;
89        let storage = executed.storage_array().clone();
90        let storage_field = Field::new(
91            String::new(),
92            target.data_type().clone(),
93            target.is_nullable(),
94        );
95        let session = ctx.session().clone();
96
97        let storage = session
98            .arrow()
99            .execute_arrow(storage, Some(&storage_field), ctx)?;
100
101        Ok(ArrowExport::Exported(storage))
102    }
103}
104
105impl ArrowImportVTable for Json {
106    fn arrow_ext_id(&self) -> Id {
107        *ARROW_JSON
108    }
109
110    fn from_arrow_field(&self, field: &Field) -> VortexResult<Option<DType>> {
111        if !has_valid_json_extension(field) {
112            return Ok(None);
113        }
114
115        Ok(Some(DType::Extension(
116            ExtDType::<Json>::try_new(EmptyMetadata, DType::Utf8(field.is_nullable().into()))?
117                .erased(),
118        )))
119    }
120
121    fn from_arrow_array(
122        &self,
123        array: ArrowArrayRef,
124        field: &Field,
125        dtype: &DType,
126    ) -> VortexResult<ArrowImport> {
127        let DType::Extension(ext_dtype) = dtype else {
128            return Ok(ArrowImport::Unsupported(array));
129        };
130        if !ext_dtype.is::<Json>() || !has_valid_json_extension(field) {
131            return Ok(ArrowImport::Unsupported(array));
132        }
133
134        let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?;
135        Ok(ArrowImport::Imported(
136            ExtensionArray::new(ext_dtype.clone(), storage).into_array(),
137        ))
138    }
139}
140
141#[cfg(test)]
142mod 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/arrow.rs

vortex_json/
arrow.rs
