polars_io/json/mod.rs
1//! # (De)serialize JSON files.
2//!
3//! ## Read JSON to a DataFrame
4//!
5//! ## Example
6//!
7//! ```
8//! use polars_core::prelude::*;
9//! use polars_io::prelude::*;
10//! use std::io::Cursor;
11//! use std::num::NonZeroUsize;
12//!
13//! let basic_json = r#"{"a":1, "b":2.0, "c":false, "d":"4"}
14//! {"a":-10, "b":-3.5, "c":true, "d":"4"}
15//! {"a":2, "b":0.6, "c":false, "d":"text"}
16//! {"a":1, "b":2.0, "c":false, "d":"4"}
17//! {"a":7, "b":-3.5, "c":true, "d":"4"}
18//! {"a":1, "b":0.6, "c":false, "d":"text"}
19//! {"a":1, "b":2.0, "c":false, "d":"4"}
20//! {"a":5, "b":-3.5, "c":true, "d":"4"}
21//! {"a":1, "b":0.6, "c":false, "d":"text"}
22//! {"a":1, "b":2.0, "c":false, "d":"4"}
23//! {"a":1, "b":-3.5, "c":true, "d":"4"}
24//! {"a":1, "b":0.6, "c":false, "d":"text"}"#;
25//! let file = Cursor::new(basic_json);
26//! let df = JsonReader::new(file)
27//! .with_json_format(JsonFormat::JsonLines)
28//! .infer_schema_len(NonZeroUsize::new(3))
29//! .with_batch_size(NonZeroUsize::new(3).unwrap())
30//! .finish()
31//! .unwrap();
32//!
33//! println!("{:?}", df);
34//! ```
35//! >>> Outputs:
36//!
37//! ```text
38//! +-----+--------+-------+--------+
39//! | a | b | c | d |
40//! | --- | --- | --- | --- |
41//! | i64 | f64 | bool | str |
42//! +=====+========+=======+========+
43//! | 1 | 2 | false | "4" |
44//! +-----+--------+-------+--------+
45//! | -10 | -3.5e0 | true | "4" |
46//! +-----+--------+-------+--------+
47//! | 2 | 0.6 | false | "text" |
48//! +-----+--------+-------+--------+
49//! | 1 | 2 | false | "4" |
50//! +-----+--------+-------+--------+
51//! | 7 | -3.5e0 | true | "4" |
52//! +-----+--------+-------+--------+
53//! | 1 | 0.6 | false | "text" |
54//! +-----+--------+-------+--------+
55//! | 1 | 2 | false | "4" |
56//! +-----+--------+-------+--------+
57//! | 5 | -3.5e0 | true | "4" |
58//! +-----+--------+-------+--------+
59//! | 1 | 0.6 | false | "text" |
60//! +-----+--------+-------+--------+
61//! | 1 | 2 | false | "4" |
62//! +-----+--------+-------+--------+
63//! ```
64//!
65pub(crate) mod infer;
66
67use std::io::Write;
68use std::num::NonZeroUsize;
69use std::ops::Deref;
70
71use arrow::array::LIST_VALUES_NAME;
72use arrow::legacy::conversion::chunk_to_struct;
73use polars_core::chunked_array::cast::CastOptions;
74use polars_core::error::to_compute_err;
75use polars_core::prelude::*;
76use polars_error::{PolarsResult, polars_bail};
77use polars_json::json::write::FallibleStreamingIterator;
78use simd_json::BorrowedValue;
79
80use crate::mmap::{MmapBytesReader, ReaderBytes};
81use crate::prelude::*;
82
83/// Reject dtypes that `polars-json` cannot serialize.
84pub fn ensure_json_writable(_dtype: &DataType) -> PolarsResult<()> {
85 #[cfg(feature = "object")]
86 polars_ensure!(
87 !_dtype.contains_objects(),
88 ComputeError: "cannot write 'Object' datatype to json"
89 );
90 #[cfg(feature = "dtype-map")]
91 polars_ensure!(
92 !_dtype.contains_map(),
93 ComputeError:
94 "cannot write 'Map' datatype to json\n\nConsider `Expr.map.entries` to write the entries as a list of structs."
95 );
96 Ok(())
97}
98
99/// The format to use to write the DataFrame to JSON: `Json` (a JSON array)
100/// or `JsonLines` (each row output on a separate line).
101///
102/// In either case, each row is serialized as a JSON object whose keys are the column names and
103/// whose values are the row's corresponding values.
104pub enum JsonFormat {
105 /// A single JSON array containing each DataFrame row as an object. The length of the array is the number of rows in
106 /// the DataFrame.
107 ///
108 /// Use this to create valid JSON that can be deserialized back into an array in one fell swoop.
109 Json,
110 /// Each DataFrame row is serialized as a JSON object on a separate line. The number of lines in the output is the
111 /// number of rows in the DataFrame.
112 ///
113 /// The [JSON Lines](https://jsonlines.org) format makes it easy to read records in a streaming fashion, one (line)
114 /// at a time. But the output in its entirety is not valid JSON; only the individual lines are.
115 ///
116 /// It is recommended to use the file extension `.jsonl` when saving as JSON Lines.
117 JsonLines,
118}
119
120/// Writes a DataFrame to JSON.
121///
122/// Under the hood, this uses [`arrow2::io::json`](https://docs.rs/arrow2/latest/arrow2/io/json/write/fn.write.html).
123/// `arrow2` generally serializes types that are not JSON primitives, such as Date and DateTime, as their
124/// `Display`-formatted versions. For instance, a (naive) DateTime column is formatted as the String `"yyyy-mm-dd
125/// HH:MM:SS"`. To control how non-primitive columns are serialized, convert them to String or another primitive type
126/// before serializing.
127#[must_use]
128pub struct JsonWriter<W: Write> {
129 /// File or Stream handler
130 buffer: W,
131 json_format: JsonFormat,
132}
133
134impl<W: Write> JsonWriter<W> {
135 pub fn with_json_format(mut self, format: JsonFormat) -> Self {
136 self.json_format = format;
137 self
138 }
139}
140
141impl<W> SerWriter<W> for JsonWriter<W>
142where
143 W: Write,
144{
145 /// Create a new `JsonWriter` writing to `buffer` with format `JsonFormat::JsonLines`. To specify a different
146 /// format, use e.g., [`JsonWriter::new(buffer).with_json_format(JsonFormat::Json)`](JsonWriter::with_json_format).
147 fn new(buffer: W) -> Self {
148 JsonWriter {
149 buffer,
150 json_format: JsonFormat::JsonLines,
151 }
152 }
153
154 fn finish(&mut self, df: &mut DataFrame) -> PolarsResult<()> {
155 df.align_chunks_par();
156 let fields = df
157 .columns()
158 .iter()
159 .map(|s| {
160 ensure_json_writable(s.dtype())?;
161 Ok(s.field().to_arrow(CompatLevel::newest()))
162 })
163 .collect::<PolarsResult<Vec<_>>>()?;
164 let batches = df
165 .iter_chunks(CompatLevel::newest(), false)
166 .map(|chunk| Ok(Box::new(chunk_to_struct(chunk, fields.clone())) as ArrayRef));
167
168 match self.json_format {
169 JsonFormat::JsonLines => {
170 let serializer = polars_json::ndjson::write::Serializer::new(batches, vec![]);
171 let writer =
172 polars_json::ndjson::write::FileWriter::new(&mut self.buffer, serializer);
173 writer.collect::<PolarsResult<()>>()?;
174 },
175 JsonFormat::Json => {
176 let serializer = polars_json::json::write::Serializer::new(batches, vec![]);
177 polars_json::json::write::write(&mut self.buffer, serializer)?;
178 },
179 }
180
181 Ok(())
182 }
183}
184
185pub struct BatchedWriter<W: Write> {
186 writer: W,
187}
188
189impl<W> BatchedWriter<W>
190where
191 W: Write,
192{
193 pub fn new(writer: W) -> Self {
194 BatchedWriter { writer }
195 }
196 /// Write a batch to the json writer.
197 ///
198 /// # Panics
199 /// The caller must ensure the chunks in the given [`DataFrame`] are aligned.
200 pub fn write_batch(&mut self, df: &DataFrame) -> PolarsResult<()> {
201 let fields = df
202 .columns()
203 .iter()
204 .map(|s| {
205 ensure_json_writable(s.dtype())?;
206 Ok(s.field().to_arrow(CompatLevel::newest()))
207 })
208 .collect::<PolarsResult<Vec<_>>>()?;
209 let chunks = df.iter_chunks(CompatLevel::newest(), false);
210 let batches =
211 chunks.map(|chunk| Ok(Box::new(chunk_to_struct(chunk, fields.clone())) as ArrayRef));
212 let mut serializer = polars_json::ndjson::write::Serializer::new(batches, vec![]);
213 while let Some(block) = serializer.next()? {
214 self.writer.write_all(block)?;
215 }
216 Ok(())
217 }
218}
219
220/// Reads JSON in one of the formats in [`JsonFormat`] into a DataFrame.
221#[must_use]
222pub struct JsonReader<'a, R>
223where
224 R: MmapBytesReader,
225{
226 reader: R,
227 rechunk: bool,
228 ignore_errors: bool,
229 infer_schema_len: Option<NonZeroUsize>,
230 batch_size: NonZeroUsize,
231 projection: Option<Vec<PlSmallStr>>,
232 schema: Option<SchemaRef>,
233 schema_overwrite: Option<&'a Schema>,
234 json_format: JsonFormat,
235}
236
237pub fn remove_bom(bytes: &[u8]) -> PolarsResult<&[u8]> {
238 if bytes.starts_with(&[0xEF, 0xBB, 0xBF]) {
239 // UTF-8 BOM
240 Ok(&bytes[3..])
241 } else if bytes.starts_with(&[0xFE, 0xFF]) || bytes.starts_with(&[0xFF, 0xFE]) {
242 // UTF-16 BOM
243 polars_bail!(ComputeError: "utf-16 not supported")
244 } else {
245 Ok(bytes)
246 }
247}
248impl<R> SerReader<R> for JsonReader<'_, R>
249where
250 R: MmapBytesReader,
251{
252 fn new(reader: R) -> Self {
253 JsonReader {
254 reader,
255 rechunk: true,
256 ignore_errors: false,
257 infer_schema_len: Some(NonZeroUsize::new(100).unwrap()),
258 batch_size: NonZeroUsize::new(8192).unwrap(),
259 projection: None,
260 schema: None,
261 schema_overwrite: None,
262 json_format: JsonFormat::Json,
263 }
264 }
265
266 fn set_rechunk(mut self, rechunk: bool) -> Self {
267 self.rechunk = rechunk;
268 self
269 }
270
271 /// Take the SerReader and return a parsed DataFrame.
272 ///
273 /// Because JSON values specify their types (number, string, etc), no upcasting or conversion is performed between
274 /// incompatible types in the input. In the event that a column contains mixed dtypes, is it unspecified whether an
275 /// error is returned or whether elements of incompatible dtypes are replaced with `null`.
276 fn finish(mut self) -> PolarsResult<DataFrame> {
277 let pre_rb: ReaderBytes = (&mut self.reader).into();
278 let bytes = remove_bom(pre_rb.deref())?;
279 let rb = ReaderBytes::Borrowed(bytes);
280 let out = match self.json_format {
281 JsonFormat::Json => {
282 polars_ensure!(!self.ignore_errors, InvalidOperation: "'ignore_errors' only supported in ndjson");
283 let mut bytes = rb.deref().to_vec();
284 let owned = &mut vec![];
285 #[expect(deprecated)] // JSON is not a row-format
286 compression::maybe_decompress_bytes(&bytes, owned)?;
287 // the easiest way to avoid ownership issues is by implicitly figuring out if
288 // decompression happened (owned is only populated on decompress), then pick which bytes to parse
289 let json_value = if owned.is_empty() {
290 simd_json::to_borrowed_value(&mut bytes).map_err(to_compute_err)?
291 } else {
292 simd_json::to_borrowed_value(owned).map_err(to_compute_err)?
293 };
294 if let BorrowedValue::Array(array) = &json_value {
295 if array.is_empty() & self.schema.is_none() & self.schema_overwrite.is_none() {
296 return Ok(DataFrame::empty());
297 }
298 }
299
300 let allow_extra_fields_in_struct = self.schema.is_some();
301
302 let mut schema = if let Some(schema) = self.schema {
303 Arc::unwrap_or_clone(schema)
304 } else {
305 // Infer.
306 let inner_dtype = if let BorrowedValue::Array(values) = &json_value {
307 infer::json_values_to_supertype(
308 values,
309 self.infer_schema_len
310 .unwrap_or(NonZeroUsize::new(usize::MAX).unwrap()),
311 )?
312 } else {
313 DataType::from_arrow_dtype(&polars_json::json::infer(&json_value)?)
314 };
315
316 let DataType::Struct(fields) = inner_dtype else {
317 polars_bail!(ComputeError: "can only deserialize json objects")
318 };
319
320 Schema::from_iter(fields)
321 };
322
323 if let Some(overwrite) = self.schema_overwrite {
324 overwrite_schema(&mut schema, overwrite)?;
325 }
326
327 let mut needs_cast = false;
328 let deserialize_schema = schema
329 .iter()
330 .map(|(name, dt)| {
331 Field::new(
332 name.clone(),
333 dt.clone().map_leaves(&mut |leaf_dt| {
334 // Deserialize enums and categoricals as strings first.
335 match leaf_dt {
336 #[cfg(feature = "dtype-categorical")]
337 DataType::Enum(..) | DataType::Categorical(..) => {
338 needs_cast = true;
339 DataType::String
340 },
341 leaf_dt => leaf_dt,
342 }
343 }),
344 )
345 })
346 .collect();
347
348 let arrow_dtype =
349 DataType::Struct(deserialize_schema).to_arrow(CompatLevel::newest());
350
351 let arrow_dtype = if let BorrowedValue::Array(_) = &json_value {
352 ArrowDataType::LargeList(Box::new(arrow::datatypes::Field::new(
353 LIST_VALUES_NAME,
354 arrow_dtype,
355 true,
356 )))
357 } else {
358 arrow_dtype
359 };
360
361 let arr = polars_json::json::deserialize(
362 &json_value,
363 arrow_dtype,
364 allow_extra_fields_in_struct,
365 )?;
366
367 let arr = arr.as_any().downcast_ref::<StructArray>().ok_or_else(
368 || polars_err!(ComputeError: "can only deserialize json objects"),
369 )?;
370
371 let mut df = DataFrame::try_from(arr.clone())?;
372
373 if df.width() == 0 && df.height() <= 1 {
374 // read_json("{}")
375 unsafe { df.set_height(0) };
376 }
377
378 if needs_cast {
379 for (col, dt) in unsafe { df.columns_mut() }
380 .iter_mut()
381 .zip(schema.iter_values())
382 {
383 *col = col.cast_with_options(
384 dt,
385 if self.ignore_errors {
386 CastOptions::NonStrict
387 } else {
388 CastOptions::Strict
389 },
390 )?;
391 }
392 }
393
394 df
395 },
396 JsonFormat::JsonLines => {
397 let mut json_reader = CoreJsonReader::new(
398 rb,
399 None,
400 self.schema,
401 self.schema_overwrite,
402 None,
403 1024, // sample size
404 NonZeroUsize::new(1 << 18).unwrap(),
405 false,
406 self.infer_schema_len,
407 self.ignore_errors,
408 None,
409 None,
410 None,
411 )?;
412 let mut df: DataFrame = json_reader.as_df()?;
413 if self.rechunk {
414 df.rechunk_mut_par();
415 }
416
417 df
418 },
419 };
420
421 // TODO! Ensure we don't materialize the columns we don't need
422 if let Some(proj) = self.projection.as_deref() {
423 out.select(proj.iter().cloned())
424 } else {
425 Ok(out)
426 }
427 }
428}
429
430impl<'a, R> JsonReader<'a, R>
431where
432 R: MmapBytesReader,
433{
434 /// Set the JSON file's schema
435 pub fn with_schema(mut self, schema: SchemaRef) -> Self {
436 self.schema = Some(schema);
437 self
438 }
439
440 /// Overwrite parts of the inferred schema.
441 pub fn with_schema_overwrite(mut self, schema: &'a Schema) -> Self {
442 self.schema_overwrite = Some(schema);
443 self
444 }
445
446 /// Set the JSON reader to infer the schema of the file. Currently, this is only used when reading from
447 /// [`JsonFormat::JsonLines`], as [`JsonFormat::Json`] reads in the entire array anyway.
448 ///
449 /// When using [`JsonFormat::JsonLines`], `max_records = None` will read the entire buffer in order to infer the
450 /// schema, `Some(1)` would look only at the first record, `Some(2)` the first two records, etc.
451 ///
452 /// It is an error to pass `max_records = Some(0)`, as a schema cannot be inferred from 0 records when deserializing
453 /// from JSON (unlike CSVs, there is no header row to inspect for column names).
454 pub fn infer_schema_len(mut self, max_records: Option<NonZeroUsize>) -> Self {
455 self.infer_schema_len = max_records;
456 self
457 }
458
459 /// Set the batch size (number of records to load at one time)
460 ///
461 /// This heavily influences loading time.
462 pub fn with_batch_size(mut self, batch_size: NonZeroUsize) -> Self {
463 self.batch_size = batch_size;
464 self
465 }
466
467 /// Set the reader's column projection: the names of the columns to keep after deserialization. If `None`, all
468 /// columns are kept.
469 ///
470 /// Setting `projection` to the columns you want to keep is more efficient than deserializing all of the columns and
471 /// then dropping the ones you don't want.
472 pub fn with_projection(mut self, projection: Option<Vec<PlSmallStr>>) -> Self {
473 self.projection = projection;
474 self
475 }
476
477 pub fn with_json_format(mut self, format: JsonFormat) -> Self {
478 self.json_format = format;
479 self
480 }
481
482 /// Return a `null` if an error occurs during parsing.
483 pub fn with_ignore_errors(mut self, ignore: bool) -> Self {
484 self.ignore_errors = ignore;
485 self
486 }
487}