 b/tests/data/test_parquet.py
+from itertools import islice
+from pathlib import Path
+import pyarrow.dataset
+import pyarrow.fs
+import pytest
+from confit.utils.random import set_seed
+from typing_extensions import Literal
+import edsnlp
+from edsnlp.data.converters import get_dict2doc_converter, get_doc2dict_converter
+from edsnlp.utils.collections import dl_to_ld
+def assert_doc_read(doc):
+    assert doc._.note_id == "subfolder/doc-1"
+    assert doc._.context_var == "test"
+    attrs = ("etat", "assertion")
+    spans_and_attributes = {
+        "__ents__": sorted(
+            [
+                (e.start, e.end, e.text, tuple(getattr(e._, key) for key in attrs))
+                for e in doc.ents
+            ]
+        ),
+        **{
+            name: sorted(
+                [
+                    (e.start, e.end, e.text, tuple(getattr(e._, key) for key in attrs))
+                    for e in doc.spans[name]
+                ]
+            )
+            for name in doc.spans
+        },
+    }
+    assert spans_and_attributes == {
+        "__ents__": [
+            (6, 7, "douleurs", (None, None)),
+            (7, 11, "dans le bras droit", (None, None)),
+            (17, 21, "problème \nde locomotion", (None, "absent")),
+            (25, 26, "AVC", ("passé", "non-associé")),
+            (35, 36, "rhume", ("présent", "hypothétique")),
+            (45, 46, "rhume", ("présent", "hypothétique")),
+            (51, 52, "Douleurs", (None, None)),
+            (52, 56, "dans le bras droit", (None, None)),
+            (68, 69, "anomalie", (None, "absent")),
+        ],
+        "anatomie": [
+            (9, 11, "bras droit", (None, None)),
+            (54, 56, "bras droit", (None, None)),
+        ],
+        "localisation": [
+            (7, 11, "dans le bras droit", (None, None)),
+            (52, 56, "dans le bras droit", (None, None)),
+        ],
+        "pathologie": [
+            (17, 21, "problème \nde locomotion", (None, "absent")),
+            (25, 26, "AVC", ("passé", "non-associé")),
+            (35, 36, "rhume", ("présent", "hypothétique")),
+            (45, 46, "rhume", ("présent", "hypothétique")),
+        ],
+        "sosy": [
+            (6, 7, "douleurs", (None, None)),
+            (51, 52, "Douleurs", (None, None)),
+            (68, 69, "anomalie", (None, "absent")),
+        ],
+    }
+GOLD_OMOP = {
+    "entities": [
+        {
+            "assertion": None,
+            "end_char": 38,
+            "etat": "test",
+            "lexical_variant": "douleurs",
+            "note_nlp_id": 0,
+            "note_nlp_source_value": "sosy",
+            "start_char": 30,
+        },
+        {
+            "assertion": None,
+            "end_char": 57,
+            "etat": None,
+            "lexical_variant": "dans le bras droit",
+            "note_nlp_id": 1,
+            "note_nlp_source_value": "localisation",
+            "start_char": 39,
+        },
+        {
+            "assertion": None,
+            "end_char": 57,
+            "etat": None,
+            "lexical_variant": "bras droit",
+            "note_nlp_id": 2,
+            "note_nlp_source_value": "anatomie",
+            "start_char": 47,
+        },
+        {
+            "assertion": "absent",
+            "end_char": 98,
+            "etat": None,
+            "lexical_variant": "problème \nde locomotion",
+            "note_nlp_id": 3,
+            "note_nlp_source_value": "pathologie",
+            "start_char": 75,
+        },
+        {
+            "assertion": "non-associé",
+            "end_char": 117,
+            "etat": "passé",
+            "lexical_variant": "AVC",
+            "note_nlp_id": 4,
+            "note_nlp_source_value": "pathologie",
+            "start_char": 114,
+        },
+        {
+            "assertion": "hypothétique",
+            "end_char": 164,
+            "etat": "présent",
+            "lexical_variant": "rhume",
+            "note_nlp_id": 5,
+            "note_nlp_source_value": "pathologie",
+            "start_char": 159,
+        },
+        {
+            "assertion": "hypothétique",
+            "end_char": 296,
+            "etat": "présent",
+            "lexical_variant": "rhume",
+            "note_nlp_id": 6,
+            "note_nlp_source_value": "pathologie",
+            "start_char": 291,
+        },
+        {
+            "assertion": None,
+            "end_char": 314,
+            "etat": None,
+            "lexical_variant": "Douleurs",
+            "note_nlp_id": 7,
+            "note_nlp_source_value": "sosy",
+            "start_char": 306,
+        },
+        {
+            "assertion": None,
+            "end_char": 333,
+            "etat": None,
+            "lexical_variant": "dans le bras droit",
+            "note_nlp_id": 8,
+            "note_nlp_source_value": "localisation",
+            "start_char": 315,
+        },
+        {
+            "assertion": None,
+            "end_char": 333,
+            "etat": None,
+            "lexical_variant": "bras droit",
+            "note_nlp_id": 9,
+            "note_nlp_source_value": "anatomie",
+            "start_char": 323,
+        },
+        {
+            "assertion": "absent",
+            "end_char": 386,
+            "etat": None,
+            "lexical_variant": "anomalie",
+            "note_nlp_id": 10,
+            "note_nlp_source_value": "sosy",
+            "start_char": 378,
+        },
+    ],
+    "note_id": "subfolder/doc-1",
+    "context_var": "test",
+    "note_text": "Le patient est admis pour des douleurs dans le bras droit, mais "
+    "n'a pas de problème \n"
+    "de locomotion. \n"
+    "Historique d'AVC dans la famille. pourrait être un cas de "
+    "rhume.\n"
+    "NBNbWbWbNbWbNBNbNbWbWbNBNbWbNbNbWbNBNbWbNbNBWbWbNbNbNBWbNbWbNbWBNbNbWbNbNBNbWb"
+    "WbNbWBNbNbWbNBNbWbWbNb\n"
+    "Pourrait être un cas de rhume.\n"
+    "Motif :\n"
+    "Douleurs dans le bras droit.\n"
+    "ANTÉCÉDENTS\n"
+    "Le patient est déjà venu\n"
+    "Pas d'anomalie détectée.\n",
+}
+def assert_doc_write_omop(exported_obj):
+    assert exported_obj == GOLD_OMOP
+def assert_doc_write_ents(exported_objs):
+    in_converter, kwargs = get_dict2doc_converter(
+        "omop",
+        dict(
+            span_attributes=["etat", "assertion"],
+            doc_attributes=["context_var"],
+        ),
+    )
+    doc = in_converter(GOLD_OMOP, **kwargs)
+    out_converter, kwargs = get_doc2dict_converter(
+        "ents",
+        dict(
+            span_attributes=["etat", "assertion"],
+            doc_attributes=["context_var"],
+            span_getter=["ents", "sosy", "localisation", "anatomie", "pathologie"],
+        ),
+    )
+    GOLD_ENTS = out_converter(doc, **kwargs)
+    assert exported_objs == GOLD_ENTS
+def test_read_write_in_worker(blank_nlp, tmpdir):
+    input_dir = Path(__file__).parent.parent.resolve() / "resources" / "docs.parquet"
+    output_dir = Path(tmpdir)
+    edsnlp.data.read_parquet(
+        input_dir,
+        converter="omop",
+        span_attributes=["etat", "assertion"],
+        doc_attributes=["context_var"],
+        read_in_worker=True,
+    ).write_parquet(
+        output_dir / "docs.parquet",
+        converter="omop",
+        doc_attributes=["context_var"],
+        span_attributes=["etat", "assertion"],
+        span_getter=["ents", "sosy", "localisation", "anatomie", "pathologie"],
+        write_in_worker=True,
+    )
+    # fmt: off
+    assert (
+            list(dl_to_ld(pyarrow.dataset.dataset(output_dir / "docs.parquet").to_table().to_pydict()))  # noqa: E501
+            == list(dl_to_ld(pyarrow.dataset.dataset(input_dir).to_table().to_pydict()))
+    )
+    # fmt: on
+def test_read_to_parquet(blank_nlp, tmpdir, run_in_test_dir):
+    output_dir = Path(tmpdir)
+    fs = pyarrow.fs.LocalFileSystem()
+    doc = list(
+        edsnlp.data.read_parquet(
+            "../resources/docs.parquet",
+            converter="omop",
+            span_attributes=["etat", "assertion"],
+            doc_attributes=["context_var"],
+            filesystem=fs,
+        )
+    )[0]
+    assert_doc_read(doc)
+    doc.ents[0]._.etat = "test"
+    edsnlp.data.write_parquet(
+        [doc],
+        output_dir,
+        converter="omop",
+        doc_attributes=["context_var"],
+        span_attributes=["etat", "assertion"],
+        span_getter=["ents", "sosy", "localisation", "anatomie", "pathologie"],
+    )
+    assert_doc_write_omop(
+        next(dl_to_ld(pyarrow.dataset.dataset(output_dir).to_table().to_pydict()))
+    )
+    with pytest.raises(FileExistsError):
+        edsnlp.data.write_parquet(
+            [doc],
+            output_dir,
+            converter="omop",
+            doc_attributes=["context_var"],
+            span_attributes=["etat", "assertion"],
+            span_getter=["ents", "sosy", "localisation", "anatomie", "pathologie"],
+        )
+    edsnlp.data.write_parquet(
+        [doc],
+        output_dir,
+        converter="omop",
+        doc_attributes=["context_var"],
+        span_attributes=["etat", "assertion"],
+        span_getter=["ents", "sosy", "localisation", "anatomie", "pathologie"],
+        overwrite=True,
+    )
+def test_read_to_parquet_ents(blank_nlp, tmpdir):
+    input_dir = Path(__file__).parent.parent.resolve() / "resources" / "docs.parquet"
+    output_dir = Path(tmpdir)
+    fs = pyarrow.fs.LocalFileSystem()
+    doc = list(
+        edsnlp.data.read_parquet(
+            input_dir,
+            converter="omop",
+            span_attributes=["etat", "assertion"],
+            doc_attributes=["context_var"],
+            filesystem=fs,
+        )
+    )[0]
+    assert_doc_read(doc)
+    doc.ents[0]._.etat = "test"
+    edsnlp.data.write_parquet(
+        [doc],
+        output_dir,
+        converter="ents",
+        doc_attributes=["context_var"],
+        span_attributes=["etat", "assertion"],
+        span_getter=["ents", "sosy", "localisation", "anatomie", "pathologie"],
+        num_rows_per_file=1024,  # deprecated but test for backward compatibility
+    )
+    assert_doc_write_ents(
+        list(dl_to_ld(pyarrow.dataset.dataset(output_dir).to_table().to_pydict()))
+    )
+    with pytest.raises(FileExistsError):
+        edsnlp.data.write_parquet(
+            [doc],
+            output_dir,
+            converter="ents",
+            doc_attributes=["context_var"],
+            span_attributes=["etat", "assertion"],
+            span_getter=["ents", "sosy", "localisation", "anatomie", "pathologie"],
+        )
+@pytest.mark.parametrize("num_cpu_workers", [0, 2])
+@pytest.mark.parametrize("shuffle", ["dataset", "fragment"])
+@pytest.mark.parametrize("shuffle_reader", [False, None])
+def test_read_shuffle_loop(
+    num_cpu_workers: int,
+    shuffle: Literal["dataset", "fragment"],
+    shuffle_reader: bool,
+):
+    input_dir = Path(__file__).parent.parent.resolve() / "resources" / "docs.parquet"
+    notes = (
+        edsnlp.data.read_parquet(input_dir, loop=True)
+        .shuffle(batch_by=shuffle, seed=42, shuffle_reader=shuffle_reader)
+        .map(lambda x: x["note_id"])
+        .set_processing(num_cpu_workers=num_cpu_workers)
+    )
+    notes = list(islice(notes, 6))
+    if not (num_cpu_workers > 1 and not shuffle_reader):
+        assert notes == [
+            "subfolder/doc-2",
+            "subfolder/doc-1",
+            "subfolder/doc-3",
+            "subfolder/doc-3",
+            "subfolder/doc-2",
+            "subfolder/doc-1",
+        ]
+@pytest.mark.parametrize("num_cpu_workers", [0, 2])
+@pytest.mark.parametrize("work_unit", ["record", "fragment"])
+@pytest.mark.parametrize("shuffle", [False, "dataset", "fragment"])
+def test_read_work_unit(
+    num_cpu_workers,
+    work_unit: Literal["record", "fragment"],
+    shuffle: Literal[False, "dataset", "fragment"],
+):
+    if shuffle == "dataset" and work_unit == "fragment":
+        pytest.skip("Dataset-level shuffle is not supported with fragment work unit")
+    input_dir = Path(__file__).parent.parent.resolve() / "resources" / "docs.parquet"
+    set_seed(42)
+    stream = edsnlp.data.read_parquet(
+        input_dir, work_unit=work_unit, shuffle=shuffle
+    ).set_processing(
+        num_cpu_workers=num_cpu_workers,
+    )
+    stream = stream.map_batches(
+        lambda b: "|".join(sorted([x["note_id"] for x in b])), batch_size=1000
+    )
+    if work_unit == "fragment" and num_cpu_workers == 2 or num_cpu_workers == 0:
+        assert list(stream) == ["subfolder/doc-1|subfolder/doc-2|subfolder/doc-3"]
+    else:
+        assert list(stream) == ["subfolder/doc-1|subfolder/doc-3", "subfolder/doc-2"]
+@pytest.mark.parametrize(
+    "num_cpu_workers,write_in_worker",
+    [
+        (0, False),
+        (2, True),
+        (2, False),
+    ],
+)
+def test_write_parquet_fragment(tmpdir, num_cpu_workers, write_in_worker):
+    input_dir = Path(__file__).parent.parent.resolve() / "resources" / "docs.parquet"
+    output_dir = Path(tmpdir)
+    notes = edsnlp.data.read_parquet(input_dir, converter="omop")
+    notes = notes.map_batches(
+        lambda b: [y for x in b for y in (x,) * (4 if "doc-2" in x._.note_id else 1)]
+    )
+    notes = notes.set_processing(
+        num_cpu_workers=num_cpu_workers,
+        deterministic=True,  # by default
+    )
+    notes.write_parquet(
+        output_dir,
+        batch_size="fragment",
+        converter=lambda x: {"note_id": x._.note_id},
+        write_in_worker=write_in_worker,
+    )
+    input_ds = pyarrow.dataset.dataset(input_dir)
+    inputs = [o["note_id"] for o in dl_to_ld(input_ds.to_table().to_pydict())]
+    assert len(list(input_ds.get_fragments())) == 1
+    assert inputs == ["subfolder/doc-1", "subfolder/doc-2", "subfolder/doc-3"]
+    out_ds = pyarrow.dataset.dataset(output_dir)
+    outs = [o["note_id"] for o in dl_to_ld(out_ds.to_table().to_pydict())]
+    if write_in_worker and num_cpu_workers == 2:
+        # Depending on the order in which the 2 workers produces the batches
+        assert len(list(out_ds.get_fragments())) == 2
+        assert outs == [
+            "subfolder/doc-1",
+            "subfolder/doc-3",
+            "subfolder/doc-2",
+            "subfolder/doc-2",
+            "subfolder/doc-2",
+            "subfolder/doc-2",
+        ] or outs == [
+            "subfolder/doc-2",
+            "subfolder/doc-2",
+            "subfolder/doc-2",
+            "subfolder/doc-2",
+            "subfolder/doc-1",
+            "subfolder/doc-3",
+        ]
+    elif not write_in_worker and num_cpu_workers == 2:
+        assert len(list(out_ds.get_fragments())) == 1
+        assert outs == [
+            "subfolder/doc-1",
+            "subfolder/doc-2",
+            "subfolder/doc-3",
+            "subfolder/doc-2",
+            "subfolder/doc-2",
+            "subfolder/doc-2",
+        ]
+    else:  # simple case
+        assert len(list(out_ds.get_fragments())) == 1
+        assert outs == [
+            "subfolder/doc-1",
+            "subfolder/doc-2",
+            "subfolder/doc-2",
+            "subfolder/doc-2",
+            "subfolder/doc-2",
+            "subfolder/doc-3",
+        ]