Streaming large JSON, XML and CSV files

"""The same 200,000 orders as a JSON array, JSON Lines, XML and CSV: per-city totals computed by reading each file whole
and by streaming it, with the peak memory of each way, as tracemalloc counts it."""
import csv
import json
import os
import random
import tracemalloc
import xml.etree.ElementTree as ET
from collections import Counter

CITIES = ["Pune", "Kochi", "Jaipur", "Indore", "Surat", "Mysuru"]
N = 200_000


def orders():
    rng = random.Random(42)
    for i in range(N):
        yield {"id": f"o{i:06d}", "city": rng.choice(CITIES), "paise": rng.randrange(1_000, 500_000), "items": rng.randint(1, 5)}


# Write the four files once; each holds the same orders.
with open("orders.json", "w", encoding="utf-8") as f:
    json.dump(list(orders()), f)
with open("orders.jsonl", "w", encoding="utf-8") as f:
    for o in orders():
        f.write(json.dumps(o) + "\n")
with open("orders.xml", "w", encoding="utf-8") as f:
    f.write("<orders>")
    for o in orders():
        f.write(f'<order id="{o["id"]}"><city>{o["city"]}</city><paise>{o["paise"]}</paise><items>{o["items"]}</items></order>')
    f.write("</orders>\n")
with open("orders.csv", "w", encoding="utf-8", newline="") as f:
    w = csv.DictWriter(f, ["id", "city", "paise", "items"])
    w.writeheader()
    w.writerows(orders())


def json_whole():
    totals = Counter()
    with open("orders.json", encoding="utf-8") as f:
        for o in json.load(f):
            totals[o["city"]] += o["paise"]
    return totals


def json_lines():
    totals = Counter()
    with open("orders.jsonl", encoding="utf-8") as f:
        for line in f:
            o = json.loads(line)
            totals[o["city"]] += o["paise"]
    return totals


def xml_whole():
    totals = Counter()
    for order in ET.parse("orders.xml").getroot():
        totals[order.findtext("city")] += int(order.findtext("paise"))
    return totals


def xml_stream():
    totals = Counter()
    root = None
    for event, el in ET.iterparse("orders.xml", events=("start", "end")):
        if root is None:
            root = el
        elif event == "end" and el.tag == "order":
            totals[el.findtext("city")] += int(el.findtext("paise"))
            root.clear()  # drop the finished orders, or the tree keeps every one of them
    return totals


def csv_stream():
    totals = Counter()
    with open("orders.csv", encoding="utf-8", newline="") as f:
        for row in csv.DictReader(f):
            totals[row["city"]] += int(row["paise"])
    return totals


results = []
print(f"{'reader':<34}{'file MB':>8}{'peak MB':>9}")
for name, path, read in [
    ("json.load, the whole array", "orders.json", json_whole),
    ("JSON Lines, one line at a time", "orders.jsonl", json_lines),
    ("ElementTree.parse, the whole tree", "orders.xml", xml_whole),
    ("iterparse with root.clear()", "orders.xml", xml_stream),
    ("csv.DictReader, one row at a time", "orders.csv", csv_stream),
]:
    tracemalloc.start()
    totals = read()
    peak = tracemalloc.get_traced_memory()[1]
    tracemalloc.stop()
    results.append(totals)
    print(f"{name:<34}{os.path.getsize(path) / 1e6:>8.1f}{peak / 1e6:>9.1f}")

print("same totals from all five:", all(r == results[0] for r in results))
print("Pune:", results[0]["Pune"], "paise")

reader                             file MB  peak MB
json.load, the whole array            13.0     76.0
JSON Lines, one line at a time        12.8      0.2
ElementTree.parse, the whole tree     16.6    137.4
iterparse with root.clear()           16.6      0.4
csv.DictReader, one row at a time      4.8      0.2
same totals from all five: True
Pune: 8340826118 paise

"""One JSON array of 200,000 orders, read item by item from 64 KB chunks with the standard library's raw_decode, against
json.load of the same file: the same totals, and the peak memory of each."""
import json
import random
import tracemalloc
from collections import Counter

CITIES = ["Pune", "Kochi", "Jaipur", "Indore", "Surat", "Mysuru"]
rng = random.Random(42)
with open("orders.json", "w", encoding="utf-8") as f:
    json.dump([{"id": f"o{i:06d}", "city": rng.choice(CITIES), "paise": rng.randrange(1_000, 500_000)} for i in range(200_000)], f)


def iter_array(f, chunk_size=65_536):
    """Yield the items of the JSON array in the text file f, reading chunk_size characters at a time."""
    decoder = json.JSONDecoder()
    buf, pos, eof = "", 0, False

    def more():
        nonlocal buf, pos, eof
        piece = f.read(chunk_size)
        eof = piece == ""
        buf, pos = buf[pos:] + piece, 0  # keep only what is not yet read

    def skip_space():
        nonlocal pos
        while True:
            while pos < len(buf) and buf[pos] in " \t\r\n":
                pos += 1
            if pos < len(buf) or eof:
                return
            more()

    more()
    skip_space()
    if buf[pos:pos + 1] != "[":
        raise ValueError("not a JSON array")
    pos += 1
    first = True
    while True:
        skip_space()
        if buf[pos:pos + 1] == "]":
            return
        if not first:
            if buf[pos:pos + 1] != ",":
                raise ValueError(f"expected , or ] at character {pos}")
            pos += 1
            skip_space()
        while True:
            try:
                item, end = decoder.raw_decode(buf, pos)
                if end < len(buf) or eof:  # a value at the very end of the buffer may go on in the next chunk
                    break
            except json.JSONDecodeError:
                if eof:
                    raise
            more()
        pos, first = end, False
        yield item


def totals_of(items):
    totals = Counter()
    for o in items:
        totals[o["city"]] += o["paise"]
    return totals


for name, read in [("json.load", lambda f: totals_of(json.load(f))), ("iter_array, 64 KB chunks", lambda f: totals_of(iter_array(f)))]:
    with open("orders.json", encoding="utf-8") as f:
        tracemalloc.start()
        totals = read(f)
        peak = tracemalloc.get_traced_memory()[1]
        tracemalloc.stop()
    print(f"{name:<26} peak {peak / 1e6:5.1f} MB, Pune {totals['Pune']} paise")

json.load                  peak  73.5 MB, Pune 8297547192 paise
iter_array, 64 KB chunks   peak   0.3 MB, Pune 8297547192 paise

// One JSON array of 50,000 orders, fed to a small incremental reader in 1,000-byte chunks, as a network stream would
// deliver it. The reader cuts out each item by counting brackets outside strings, then parses the item alone.
const CITIES = ["Pune", "Kochi", "Jaipur", "Indore", "Surat", "Mysuru"];
let seed = 42;
const next = () => (seed = (seed * 1103515245 + 12345) % 2 ** 31); // a fixed sequence, so every run is the same

const orders = Array.from({ length: 50_000 }, (_, i) => ({ id: `o${i}`, city: CITIES[Math.floor((next() / 2 ** 31) * 6)], note: i % 997 === 0 ? 'says "hi" [ok]' : "" }));
const bytes = new TextEncoder().encode(JSON.stringify(orders));

function* chunks(size) {
  for (let i = 0; i < bytes.length; i += size) yield bytes.subarray(i, i + size);
}

/** Yields each item of a JSON array whose bytes arrive in chunks. */
function* items(source) {
  const decoder = new TextDecoder(); // stream: true below keeps a character split between chunks
  let depth = 0, inString = false, escaped = false, item = "", split = 0;
  for (const chunk of source) {
    const text = decoder.decode(chunk, { stream: true });
    for (const c of text) {
      if (depth >= 2 || (depth === 1 && c !== "," && c !== "]" && !/\s/.test(c))) item += c;
      if (inString) {
        if (escaped) escaped = false;
        else if (c === "\\") escaped = true;
        else if (c === '"') inString = false;
      } else if (c === '"') inString = true;
      else if (c === "{" || c === "[") depth++;
      else if (c === "}" || c === "]") {
        depth--;
        if (depth === 1 && item) {
          yield JSON.parse(item);
          item = "";
        }
      }
    }
    if (item) split++; // this chunk ended in the middle of an item
  }
  console.log(`chunks that ended inside an item: ${split}`);
}

const totals = {};
let count = 0;
for (const o of items(chunks(1_000))) {
  totals[o.city] = (totals[o.city] ?? 0) + 1;
  count++;
}
const whole = {};
for (const o of JSON.parse(new TextDecoder().decode(bytes))) whole[o.city] = (whole[o.city] ?? 0) + 1;
console.log(`${bytes.length} bytes, ${count} items read one at a time`);
console.log("orders per city:", JSON.stringify(totals));
console.log("the same as JSON.parse of the whole text:", JSON.stringify(totals) === JSON.stringify(whole));

chunks that ended inside an item: 1943
2055924 bytes, 50000 items read one at a time
orders per city: {"Indore":7950,"Mysuru":8265,"Surat":8172,"Pune":8575,"Kochi":8461,"Jaipur":8577}
the same as JSON.parse of the whole text: true

// Two JSON Lines records whose bytes arrive in two chunks, cut inside the first record, between the bytes of its
// rupee sign (E2 82 B9 in UTF-8). Three ways to decode the chunks, then a reader that also keeps the unfinished line.
const text = '{"city":"Pune","price":"₹1,250"}\n{"city":"Kochi","price":"₹980"}\n';
const bytes = new TextEncoder().encode(text);
const at = new TextEncoder().encode(text.slice(0, text.indexOf("₹"))).length; // where the sign's three bytes start
const parts = [bytes.subarray(0, at + 2), bytes.subarray(at + 2)]; // cut after two of them
console.log(`chunk sizes: ${parts.map((p) => p.length).join(" + ")} bytes; the first ends with byte ${parts[0].at(-1).toString(16)}`);

// 1. A new decoder for every chunk: each half of the sign becomes U+FFFD, without an error.
const separately = parts.map((p) => new TextDecoder().decode(p)).join("");
console.log("decoded one chunk at a time:", JSON.stringify(separately.split("\n")[0]));

// 2. The same, with fatal: true: the damage becomes an error instead.
try {
  new TextDecoder("utf-8", { fatal: true }).decode(parts[0]);
} catch (e) {
  console.log(`with fatal: true: ${e.name}: ${e.message}`);
}

// 3. One decoder for the whole stream, with stream: true, keeps the unfinished bytes for the next chunk.
const decoder = new TextDecoder();
const streamed = parts.map((p) => decoder.decode(p, { stream: true })).join("") + decoder.decode();
console.log("one decoder, stream: true:  ", JSON.stringify(streamed.split("\n")[0]));

// 4. A web stream, as fetch() hands a body over: bytes through TextDecoderStream, then whole lines only.
const body = new ReadableStream({
  start(controller) {
    for (const p of parts) controller.enqueue(p);
    controller.close();
  },
});
let pending = "";
const records = [];
for await (const piece of body.pipeThrough(new TextDecoderStream())) {
  const lines = (pending + piece).split("\n");
  pending = lines.pop(); // the last piece has no newline yet: keep it for the next chunk
  for (const line of lines) if (line) records.push(JSON.parse(line));
}
console.log("records read through TextDecoderStream:", JSON.stringify(records));

chunk sizes: 26 + 43 bytes; the first ends with byte 82
decoded one chunk at a time: "{\"city\":\"Pune\",\"price\":\"��1,250\"}"
with fatal: true: TypeError: The encoded data was not valid for encoding utf-8
one decoder, stream: true:   "{\"city\":\"Pune\",\"price\":\"₹1,250\"}"
records read through TextDecoderStream: [{"city":"Pune","price":"₹1,250"},{"city":"Kochi","price":"₹980"}]

"""The same two chunks in Python: bytes.decode per chunk, errors="replace", and an incremental decoder."""
import codecs
import json

text = '{"city":"Pune","price":"₹1,250"}\n{"city":"Kochi","price":"₹980"}\n'
data = text.encode("utf-8")
at = len(text[: text.index("₹")].encode("utf-8"))
parts = [data[: at + 2], data[at + 2 :]]

try:
    "".join(p.decode("utf-8") for p in parts)
except UnicodeDecodeError as e:
    print("decode per chunk:", e)

replaced = "".join(p.decode("utf-8", errors="replace") for p in parts)
print('errors="replace":', json.loads(replaced.splitlines()[0])["price"])

decoder = codecs.getincrementaldecoder("utf-8")()
pending, records = "", []
for p in parts:
    lines = (pending + decoder.decode(p)).split("\n")
    pending = lines.pop()  # an unfinished line waits for the next chunk
    records += [json.loads(line) for line in lines if line]
decoder.decode(b"", final=True)  # raises if the stream ended inside a character
print("incremental decoder:", [r["price"] for r in records])

decode per chunk: 'utf-8' codec can't decode bytes in position 24-25: unexpected end of data
errors="replace": ��1,250
incremental decoder: ['₹1,250', '₹980']

// A fast producer of 20,000 JSON Lines records and a slow consumer that finishes one write per turn of the event
// loop, as a busy disk or socket does. How many bytes wait in memory when write()'s answer is ignored, and when it is
// respected by hand or by pipeline().
import { once } from "node:events";
import { Readable, Writable, getDefaultHighWaterMark } from "node:stream";
import { pipeline } from "node:stream/promises";

const LINE = '{"city":"Pune","paise":125000}\n';
const COUNT = 20_000;
const HWM = 16_384;

function slowSink() {
  return new Writable({ highWaterMark: HWM, write: (_chunk, _encoding, done) => setImmediate(done) });
}

async function ignoring() {
  const sink = slowSink();
  let most = 0, refusals = 0;
  for (let i = 0; i < COUNT; i++) {
    if (!sink.write(LINE)) refusals++; // the answer is false, and the loop writes anyway
    most = Math.max(most, sink.writableLength);
  }
  sink.end();
  await once(sink, "finish");
  return { most, refusals };
}

async function respecting() {
  const sink = slowSink();
  let most = 0, refusals = 0;
  for (let i = 0; i < COUNT; i++) {
    const ok = sink.write(LINE);
    most = Math.max(most, sink.writableLength);
    if (!ok) {
      refusals++;
      await once(sink, "drain"); // wait until the buffer has emptied
    }
  }
  sink.end();
  await once(sink, "finish");
  return { most, refusals };
}

async function piped() {
  const sink = slowSink();
  let most = 0;
  async function* lines() {
    for (let i = 0; i < COUNT; i++) {
      most = Math.max(most, sink.writableLength);
      yield LINE;
    }
  }
  await pipeline(Readable.from(lines()), sink);
  return { most, refusals: "handled by pipeline" };
}

console.log(`Node.js ${process.versions.node}: default highWaterMark ${getDefaultHighWaterMark(false)} bytes; this sink's is ${HWM}`);
console.log(`${COUNT} records, ${COUNT * LINE.length} bytes in all`);
for (const [name, run] of [["write() answer ignored", ignoring], ["write() answer respected", respecting], ["pipeline()", piped]]) {
  const { most, refusals } = await run();
  console.log(`${name.padEnd(26)} most bytes waiting: ${String(most).padStart(6)}   write() said false: ${refusals}`);
}

Node.js 24.21.0: default highWaterMark 65536 bytes; this sink's is 16384
20000 records, 620000 bytes in all
write() answer ignored     most bytes waiting: 620000   write() said false: 19472
write() answer respected   most bytes waiting:  16399   write() said false: 37
pipeline()                 most bytes waiting:  16399   write() said false: handled by pipeline