feat: Add NDJSON source to new streaming engine #21562

nameexhaustion · 2025-03-03T08:32:17Z

Adds a NDJSON source to the new streaming engine, complete with slice / row_index support.

Benchmarks

Test system - Apple M3 Pro 11-core, 36G memory

Note that decreases in performance are only observed on the scale of milliseconds in micro benchmarks.

Full scan

dataset	time (s, mem-engine)	time (s, new-streaming)	speedup
iris.jsonl (150 rows, 5 cols)	0.00006	0.0001	0.6x
rank_317.jsonl (514,749 rows, 3 cols)	0.34	0.3	1.1x
657.xz_s-2302B.jsonl (269,716,216 rows, 3 cols)	5.7	5.7	1.0x

Dataset sources

iris.jsonl - https://j.mp/iriscsv (converted to jsonl)
rank_317.jsonl - https://huggingface.co/datasets/BAAI/IndustryCorpus/blob/main/IndustryCorpus/en/tech/rank_317.jsonl.gz (decompressed before scan)
657.xz_s-2302B.jsonl - https://huggingface.co/datasets/dchasap/spec_cpu_branch_traces/blob/main/657.xz_s/657.xz_s-2302B.jsonl

Slice benchmarks

Dataset: 657.xz_s-2302B.jsonl (269,716,216 rows, 3 cols)

dataset	time (s, mem-engine)	time (s, new-streaming)	speedup
.head(100)	0.000099	0.00024	0.41x
.head(1_000_000)	0.022	0.026	0.85x
.slice(134_858_108) (skips first half of file)	5.6	4.5	1.2x
.slice(268_716_216) (last 1M rows)	5.5	3.9	1.4x
.tail(1_000_000)	5.4	0.026	210.0x
.with_row_index().tail(1_000_000)	5.7	0.74	7.7x
.tail(100)	5.5	0.00028	20000.0x
.with_row_index().tail(100)	5.7	0.68	8.4x

Slice optimizations (single file):

For positive non-zero slice offsets, adds an optimization to skip parsing lines until the starting point, which can reduce memory usage compared to the in-mem engine
For negative slice offsets, adds an optimization to read the file backwards, and also skips lines until the starting point
- If a row index is also requested, json lines are parsed until the point at which the slice ends, after which the source transitions into row-counting mode to calculate the row index offset:

_

Fixes Add NDJSON source for new-streaming engine #21356

codecov · 2025-03-03T09:08:41Z

Codecov Report

Attention: Patch coverage is 61.50794% with 388 lines in your changes missing coverage. Please review.

Project coverage is 79.52%. Comparing base (d556155) to head (8575bbb).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...src/nodes/io_sources/ndjson/negative_slice_pass.rs	0.00%	186 Missing ⚠️
...s/polars-stream/src/nodes/io_sources/ndjson/mod.rs	74.81%	102 Missing ⚠️
.../nodes/io_sources/ndjson/line_batch_distributor.rs	67.66%	43 Missing ⚠️
...rc/nodes/io_sources/ndjson/line_batch_processor.rs	69.81%	32 Missing ⚠️
...rc/nodes/io_sources/ndjson/row_index_limit_pass.rs	80.35%	22 Missing ⚠️
crates/polars-io/src/ndjson/core.rs	95.45%	1 Missing ⚠️
...stream/src/nodes/io_sources/ndjson/chunk_reader.rs	96.00%	1 Missing ⚠️
crates/polars-utils/src/idx_mapper.rs	94.73%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #21562      +/-   ##
==========================================
- Coverage   79.59%   79.52%   -0.08%     
==========================================
  Files        1591     1598       +7     
  Lines      229525   230516     +991     
  Branches     2632     2632              
==========================================
+ Hits       182699   183319     +620     
- Misses      46219    46590     +371     
  Partials      607      607

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

coastalwhite · 2025-03-03T09:33:53Z

Ritchie, please wait with merging. I want to quickly review this.

nameexhaustion · 2025-03-03T09:37:24Z

This is how the components can be connected (to help with review) -

coastalwhite

It looks mostly fine, I mostly did a scan as it is a rather large PR. Some points here and there.

I really do think the oneshot channel change should be reverted, as a deadlock is much harder to debug than an incorrect result. We can quite easily add a debug assertion for the correctness of the return value if that is your concern.

coastalwhite · 2025-03-03T09:38:21Z

crates/polars-stream/src/nodes/io_sources/ndjson/mod.rs

+    async fn unrestricted_row_count(&mut self) -> PolarsResult<IdxSize> {
+        let mem_slice = self.scan_source_bytes()?;
+
+        let num_rows = ndjson::count_rows_par(&mem_slice, None);


We cannot call into rayon here

I think this should be safe? The rayon threads should never call back into the async executor.

The new-streaming CSV source also goes into rayon underneath in csv::read::count_rows_from_slice() -

polars/crates/polars-io/src/csv/read/parser.rs

Line 107 in d556155

let count_result: PolarsResult<usize> = POOL.install(|| iter.sum());

But let me know if we want to remove this

I was not aware that the CSV source does that. It should not.

Alright, I have updated this PR to count without rayon. I will update the one for CSV later in a separate PR

crates/polars-stream/src/nodes/io_sources/ndjson/mod.rs

py-polars/tests/unit/io/test_lazy_json.py

coastalwhite · 2025-03-03T10:03:12Z

crates/polars-stream/src/nodes/io_sources/ndjson/negative_slice_pass.rs

+        r > l
+    }));
+
+    accumulate_dataframes_vertical_unchecked(acc_morsels.drain(..).rev().map(|(_, df)| df))


this should use the owned variant

should already be owned -

polars/crates/polars-core/src/utils/mod.rs

Lines 781 to 783 in d556155

/// This takes ownership of the DataFrame so that drop is called earlier.

/// Does not check if schema is correct

pub fn accumulate_dataframes_vertical_unchecked<I>(dfs: I) -> DataFrame

unless you meant the Vec, but that one is drain as we reuse it

Ah sorry. For some reason I thought we had separate variants for that.

crates/polars-stream/src/nodes/io_sources/ndjson/mod.rs

coastalwhite · 2025-03-03T10:06:19Z

crates/polars-stream/src/nodes/io_sources/ndjson/mod.rs

+            if let Ok(v) = std::env::var("POLARS_FORCE_NDJSON_CHUNK_SIZE")
+                .map(|x| x.parse::<usize>().expect("integer"))
+            {
+                v
+            } else {
+                chunk_size
+            }


let v = std::env::var("POLARS_FORCE_NDJSON_CHUNK_SIZE") .map_or(chunk_size, |x| x.parse::<usize>().expect("expected `POLARS_FORCE_NDJSON_CHUNK_SIZE` to be an integer")) .max(1);

added

but I'll allow the user to set 0 😄

nameexhaustion · 2025-03-03T12:33:47Z

.github/workflows/benchmark.yml

@@ -73,7 +73,7 @@ jobs:

      - name: Install Polars release build
        run: |
-          pip install target/wheels/polars*.whl
+          pip install --force-reinstall target/wheels/polars*.whl


fix benchmark workflow failed due to having an older version installed for some reason -

https://github.com/pola-rs/polars/actions/runs/13630270606/job/38096297749

codspeed-hq · 2025-03-03T12:46:00Z

CodSpeed Performance Report

Merging #21562 will improve performances by 28.04%

_{Comparing nameexhaustion:ndjson-source (8575bbb) with main (d556155)}

Summary

⚡ 1 improvements
✅ 40 untouched benchmarks

Benchmarks breakdown

	Benchmark	`BASE`	`HEAD`	Change
⚡	`test_to_numpy_series_with_nulls`	438.4 µs	342.4 µs	+28.04%

c

ae43bf4

github-actions bot added enhancement New feature or an improvement of an existing feature python Related to Python Polars rust Related to Rust Polars labels Mar 3, 2025

nameexhaustion added 3 commits March 3, 2025 19:40

c

257b19d

clippy

e26c7b2

c

2236043

nameexhaustion marked this pull request as ready for review March 3, 2025 09:11

nameexhaustion requested review from ritchie46, c-peters, alexander-beedie, MarcoGorelli, reswqa and orlp as code owners March 3, 2025 09:11

coastalwhite requested changes Mar 3, 2025

View reviewed changes

nameexhaustion marked this pull request as draft March 3, 2025 10:59

nameexhaustion added 10 commits March 3, 2025 22:02

env var parsing

948c5f9

decompress once

eb1f359

add slice(-3, 999) to test

7ba2703

more comments

34a99a2

drop acc_morsels

ed71f3f

Merge remote-tracking branch 'origin/main' into ndjson-source

a09be21

slice->pre_slice

3b6caaa

count without rayon

3a7527b

benchmark.yml pip force-reinstall

d7d5956

todo for parallel rowcount

8575bbb

nameexhaustion commented Mar 3, 2025

View reviewed changes

nameexhaustion marked this pull request as ready for review March 3, 2025 12:45

coastalwhite approved these changes Mar 3, 2025

View reviewed changes

ritchie46 merged commit d922a4b into pola-rs:main Mar 3, 2025
29 checks passed

nameexhaustion mentioned this pull request Mar 4, 2025

fix: Fix CSV count with comment prefix skipped empty lines #21577

Merged

anath2 pushed a commit to anath2/polars that referenced this pull request Mar 5, 2025

feat: Add NDJSON source to new streaming engine (pola-rs#21562)

9080f84

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Add NDJSON source to new streaming engine #21562

feat: Add NDJSON source to new streaming engine #21562

nameexhaustion commented Mar 3, 2025 •

edited

Loading

codecov bot commented Mar 3, 2025 •

edited

Loading

coastalwhite commented Mar 3, 2025

nameexhaustion commented Mar 3, 2025

coastalwhite left a comment

coastalwhite Mar 3, 2025

nameexhaustion Mar 3, 2025

coastalwhite Mar 3, 2025

nameexhaustion Mar 3, 2025

coastalwhite Mar 3, 2025

nameexhaustion Mar 3, 2025 •

edited

Loading

coastalwhite Mar 3, 2025

coastalwhite Mar 3, 2025

nameexhaustion Mar 3, 2025 •

edited

Loading

nameexhaustion Mar 3, 2025 •

edited

Loading

codspeed-hq bot commented Mar 3, 2025

	/// This takes ownership of the DataFrame so that drop is called earlier.
	/// Does not check if schema is correct
	pub fn accumulate_dataframes_vertical_unchecked<I>(dfs: I) -> DataFrame

feat: Add NDJSON source to new streaming engine #21562

feat: Add NDJSON source to new streaming engine #21562

Conversation

nameexhaustion commented Mar 3, 2025 • edited Loading

Benchmarks

Full scan

Slice benchmarks

codecov bot commented Mar 3, 2025 • edited Loading

Codecov Report

coastalwhite commented Mar 3, 2025

nameexhaustion commented Mar 3, 2025

coastalwhite left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nameexhaustion Mar 3, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nameexhaustion Mar 3, 2025 • edited Loading

Choose a reason for hiding this comment

nameexhaustion Mar 3, 2025 • edited Loading

Choose a reason for hiding this comment

codspeed-hq bot commented Mar 3, 2025

Merging #21562 will improve performances by 28.04%

Summary

Benchmarks breakdown

nameexhaustion commented Mar 3, 2025 •

edited

Loading

codecov bot commented Mar 3, 2025 •

edited

Loading

nameexhaustion Mar 3, 2025 •

edited

Loading

nameexhaustion Mar 3, 2025 •

edited

Loading

nameexhaustion Mar 3, 2025 •

edited

Loading