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