Merge #561

561: Enriched documents batch reader r=curquiza a=Kerollmops

~This PR is based on #555 and must be rebased on main after it has been merged to ease the review.~
This PR contains the work in #555 and can be merged on main as soon as reviewed and approved.

- [x] Create an `EnrichedDocumentsBatchReader` that contains the external documents id.
- [x] Extract the primary key name and make it accessible in the `EnrichedDocumentsBatchReader`.
- [x] Use the external id from the `EnrichedDocumentsBatchReader` in the `Transform::read_documents`.
- [x] Remove the `update_primary_key` from the _transform.rs_ file.
- [x] Really generate the auto-generated documents ids.
- [x] Insert the (auto-generated) document ids in the document while processing it in `Transform::read_documents`.

Co-authored-by: Kerollmops <clement@meilisearch.com>

benchmarks/benches/indexing.rs

@@ -132,12 +132,13 @@ fn indexing_songs_default(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::SMOL_SONGS, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -169,12 +170,13 @@ fn reindexing_songs_default(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::SMOL_SONGS, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -184,12 +186,13 @@ fn reindexing_songs_default(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::SMOL_SONGS, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -223,11 +226,12 @@ fn deleting_songs_in_batches_default(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
                 let indexing_config = IndexDocumentsConfig::default();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
                 let documents = utils::documents_from(datasets_paths::SMOL_SONGS, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -279,11 +283,12 @@ fn indexing_songs_in_three_batches_default(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
                 let indexing_config = IndexDocumentsConfig::default();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
                 let documents = utils::documents_from(datasets_paths::SMOL_SONGS_1_2, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
 
                 wtxn.commit().unwrap();
@@ -294,19 +299,21 @@ fn indexing_songs_in_three_batches_default(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
                 let documents = utils::documents_from(datasets_paths::SMOL_SONGS_3_4, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
 
                 let indexing_config = IndexDocumentsConfig::default();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
                 let documents = utils::documents_from(datasets_paths::SMOL_SONGS_4_4, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
 
                 wtxn.commit().unwrap();
@@ -339,13 +346,14 @@ fn indexing_songs_without_faceted_numbers(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::SMOL_SONGS, "csv");
 
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -377,12 +385,13 @@ fn indexing_songs_without_faceted_fields(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::SMOL_SONGS, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -415,12 +424,13 @@ fn indexing_wiki(c: &mut Criterion) {
                 let indexing_config =
                     IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::SMOL_WIKI_ARTICLES, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -452,12 +462,13 @@ fn reindexing_wiki(c: &mut Criterion) {
                 let indexing_config =
                     IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::SMOL_WIKI_ARTICLES, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -468,12 +479,13 @@ fn reindexing_wiki(c: &mut Criterion) {
                 let indexing_config =
                     IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::SMOL_WIKI_ARTICLES, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -507,11 +519,12 @@ fn deleting_wiki_in_batches_default(c: &mut Criterion) {
                 let mut wtxn = index.write_txn().unwrap();
                 let indexing_config =
                     IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
                 let documents = utils::documents_from(datasets_paths::SMOL_WIKI_ARTICLES, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -564,12 +577,13 @@ fn indexing_wiki_in_three_batches(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config =
                     IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
                 let documents =
                     utils::documents_from(datasets_paths::SMOL_WIKI_ARTICLES_1_2, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
 
                 wtxn.commit().unwrap();
@@ -581,24 +595,26 @@ fn indexing_wiki_in_three_batches(c: &mut Criterion) {
                 let indexing_config =
                     IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents =
                     utils::documents_from(datasets_paths::SMOL_WIKI_ARTICLES_3_4, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
 
                 let indexing_config =
                     IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents =
                     utils::documents_from(datasets_paths::SMOL_WIKI_ARTICLES_4_4, "csv");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
 
                 wtxn.commit().unwrap();
@@ -631,12 +647,13 @@ fn indexing_movies_default(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::MOVIES, "json");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -667,12 +684,13 @@ fn reindexing_movies_default(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::MOVIES, "json");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -682,12 +700,13 @@ fn reindexing_movies_default(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::MOVIES, "json");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -720,11 +739,12 @@ fn deleting_movies_in_batches_default(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
                 let indexing_config = IndexDocumentsConfig::default();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
                 let documents = utils::documents_from(datasets_paths::MOVIES, "json");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -775,12 +795,13 @@ fn indexing_movies_in_three_batches(c: &mut Criterion) {
                 // as we don't care about the time it takes.
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::MOVIES_1_2, "json");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
 
                 wtxn.commit().unwrap();
@@ -791,21 +812,23 @@ fn indexing_movies_in_three_batches(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::MOVIES_3_4, "json");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
 
                 let indexing_config = IndexDocumentsConfig::default();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::MOVIES_4_4, "json");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
 
                 wtxn.commit().unwrap();
@@ -861,12 +884,13 @@ fn indexing_nested_movies_default(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::NESTED_MOVIES, "json");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -922,11 +946,12 @@ fn deleting_nested_movies_in_batches_default(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
                 let indexing_config = IndexDocumentsConfig::default();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
                 let documents = utils::documents_from(datasets_paths::NESTED_MOVIES, "json");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -984,12 +1009,13 @@ fn indexing_nested_movies_without_faceted_fields(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::NESTED_MOVIES, "json");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 
@@ -1021,12 +1047,13 @@ fn indexing_geo(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::SMOL_ALL_COUNTRIES, "jsonl");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
 
                 wtxn.commit().unwrap();
@@ -1058,12 +1085,13 @@ fn reindexing_geo(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::SMOL_ALL_COUNTRIES, "jsonl");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
 
                 wtxn.commit().unwrap();
@@ -1074,12 +1102,13 @@ fn reindexing_geo(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let indexing_config = IndexDocumentsConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
 
                 let documents = utils::documents_from(datasets_paths::SMOL_ALL_COUNTRIES, "jsonl");
-                builder.add_documents(documents).unwrap();
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
 
                 wtxn.commit().unwrap();
@@ -1113,11 +1142,12 @@ fn deleting_geo_in_batches_default(c: &mut Criterion) {
                 let config = IndexerConfig::default();
                 let mut wtxn = index.write_txn().unwrap();
                 let indexing_config = IndexDocumentsConfig::default();
-                let mut builder =
+                let builder =
                     IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ())
                         .unwrap();
-                let documents = utils::documents_from(datasets_paths::SMOL_ALL_COUNTRIES, "json");
-                builder.add_documents(documents).unwrap();
+                let documents = utils::documents_from(datasets_paths::SMOL_ALL_COUNTRIES, "jsonl");
+                let (builder, user_error) = builder.add_documents(documents).unwrap();
+                user_error.unwrap();
                 builder.execute().unwrap();
                 wtxn.commit().unwrap();
 

benchmarks/benches/utils.rs

@@ -7,12 +7,12 @@ use std::path::Path;
 
 use criterion::BenchmarkId;
 use heed::EnvOpenOptions;
-use milli::documents::DocumentBatchReader;
+use milli::documents::{DocumentsBatchBuilder, DocumentsBatchReader};
 use milli::update::{
     IndexDocuments, IndexDocumentsConfig, IndexDocumentsMethod, IndexerConfig, Settings,
 };
-use milli::{Filter, Index};
-use serde_json::{Map, Value};
+use milli::{Filter, Index, Object};
+use serde_json::Value;
 
 pub struct Conf<'a> {
     /// where we are going to create our database.mmdb directory
@@ -96,12 +96,10 @@ pub fn base_setup(conf: &Conf) -> Index {
         update_method: IndexDocumentsMethod::ReplaceDocuments,
         ..Default::default()
     };
-    let mut builder =
-        IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
+    let builder = IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
     let documents = documents_from(conf.dataset, conf.dataset_format);
-
-    builder.add_documents(documents).unwrap();
-
+    let (builder, user_error) = builder.add_documents(documents).unwrap();
+    user_error.unwrap();
     builder.execute().unwrap();
     wtxn.commit().unwrap();
 
@@ -140,7 +138,7 @@ pub fn run_benches(c: &mut criterion::Criterion, confs: &[Conf]) {
     }
 }
 
-pub fn documents_from(filename: &str, filetype: &str) -> DocumentBatchReader<impl BufRead + Seek> {
+pub fn documents_from(filename: &str, filetype: &str) -> DocumentsBatchReader<impl BufRead + Seek> {
     let reader =
         File::open(filename).expect(&format!("could not find the dataset in: {}", filename));
     let reader = BufReader::new(reader);
@@ -150,39 +148,35 @@ pub fn documents_from(filename: &str, filetype: &str) -> DocumentBatchReader<imp
         "jsonl" => documents_from_jsonl(reader).unwrap(),
         otherwise => panic!("invalid update format {:?}", otherwise),
     };
-    DocumentBatchReader::from_reader(Cursor::new(documents)).unwrap()
+    DocumentsBatchReader::from_reader(Cursor::new(documents)).unwrap()
 }
 
-fn documents_from_jsonl(mut reader: impl BufRead) -> anyhow::Result<Vec<u8>> {
-    let mut writer = Cursor::new(Vec::new());
-    let mut documents = milli::documents::DocumentBatchBuilder::new(&mut writer)?;
+fn documents_from_jsonl(reader: impl BufRead) -> anyhow::Result<Vec<u8>> {
+    let mut documents = DocumentsBatchBuilder::new(Vec::new());
 
-    let mut buf = String::new();
-
-    while reader.read_line(&mut buf)? > 0 {
-        documents.extend_from_json(&mut buf.as_bytes())?;
-        buf.clear();
+    for result in serde_json::Deserializer::from_reader(reader).into_iter::<Object>() {
+        let object = result?;
+        documents.append_json_object(&object)?;
     }
-    documents.finish()?;
 
-    Ok(writer.into_inner())
+    documents.into_inner().map_err(Into::into)
 }
 
 fn documents_from_json(reader: impl BufRead) -> anyhow::Result<Vec<u8>> {
-    let mut writer = Cursor::new(Vec::new());
-    let mut documents = milli::documents::DocumentBatchBuilder::new(&mut writer)?;
+    let mut documents = DocumentsBatchBuilder::new(Vec::new());
 
-    documents.extend_from_json(reader)?;
-    documents.finish()?;
+    documents.append_json_array(reader)?;
 
-    Ok(writer.into_inner())
+    documents.into_inner().map_err(Into::into)
 }
 
 fn documents_from_csv(reader: impl BufRead) -> anyhow::Result<Vec<u8>> {
-    let mut writer = Cursor::new(Vec::new());
-    milli::documents::DocumentBatchBuilder::from_csv(reader, &mut writer)?.finish()?;
+    let csv = csv::Reader::from_reader(reader);
 
-    Ok(writer.into_inner())
+    let mut documents = DocumentsBatchBuilder::new(Vec::new());
+    documents.append_csv(csv)?;
+
+    documents.into_inner().map_err(Into::into)
 }
 
 enum AllowedType {
@@ -222,14 +216,14 @@ impl<R: Read> CSVDocumentDeserializer<R> {
 }
 
 impl<R: Read> Iterator for CSVDocumentDeserializer<R> {
-    type Item = anyhow::Result<Map<String, Value>>;
+    type Item = anyhow::Result<Object>;
 
     fn next(&mut self) -> Option<Self::Item> {
         let csv_document = self.documents.next()?;
 
         match csv_document {
             Ok(csv_document) => {
-                let mut document = Map::new();
+                let mut document = Object::new();
 
                 for ((field_name, field_type), value) in
                     self.headers.iter().zip(csv_document.into_iter())

cli/src/main.rs

@@ -8,12 +8,12 @@ use std::time::Instant;
 use byte_unit::Byte;
 use eyre::Result;
 use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
+use milli::documents::{DocumentsBatchBuilder, DocumentsBatchReader};
 use milli::update::UpdateIndexingStep::{
     ComputeIdsAndMergeDocuments, IndexDocuments, MergeDataIntoFinalDatabase, RemapDocumentAddition,
 };
 use milli::update::{self, IndexDocumentsConfig, IndexDocumentsMethod, IndexerConfig};
-use milli::Index;
-use serde_json::{Map, Value};
+use milli::{Index, Object};
 use structopt::StructOpt;
 
 #[cfg(target_os = "linux")]
@@ -225,9 +225,9 @@ impl Performer for DocumentAddition {
             DocumentAdditionFormat::Jsonl => documents_from_jsonl(reader)?,
         };
 
-        let reader = milli::documents::DocumentBatchReader::from_reader(Cursor::new(documents))?;
+        let reader = DocumentsBatchReader::from_reader(Cursor::new(documents))?;
 
-        println!("Adding {} documents to the index.", reader.len());
+        println!("Adding {} documents to the index.", reader.documents_count());
 
         let mut txn = index.write_txn()?;
         let config = milli::update::IndexerConfig { log_every_n: Some(100), ..Default::default() };
@@ -255,7 +255,7 @@ impl Performer for DocumentAddition {
             let bar = progesses.add(bar);
             bars.push(bar);
         }
-        let mut addition = milli::update::IndexDocuments::new(
+        let addition = milli::update::IndexDocuments::new(
             &mut txn,
             &index,
             &config,
@@ -263,7 +263,10 @@ impl Performer for DocumentAddition {
             |step| indexing_callback(step, &bars),
         )
         .unwrap();
-        addition.add_documents(reader)?;
+        let (addition, user_error) = addition.add_documents(reader)?;
+        if let Err(error) = user_error {
+            return Err(error.into());
+        }
 
         std::thread::spawn(move || {
             progesses.join().unwrap();
@@ -321,35 +324,32 @@ fn indexing_callback(step: milli::update::UpdateIndexingStep, bars: &[ProgressBa
 }
 
 fn documents_from_jsonl(reader: impl Read) -> Result<Vec<u8>> {
-    let mut writer = Cursor::new(Vec::new());
-    let mut documents = milli::documents::DocumentBatchBuilder::new(&mut writer)?;
+    let mut documents = DocumentsBatchBuilder::new(Vec::new());
+    let reader = BufReader::new(reader);
 
-    let mut buf = String::new();
-    let mut reader = BufReader::new(reader);
-
-    while reader.read_line(&mut buf)? > 0 {
-        documents.extend_from_json(&mut buf.as_bytes())?;
+    for result in serde_json::Deserializer::from_reader(reader).into_iter::<Object>() {
+        let object = result?;
+        documents.append_json_object(&object)?;
     }
-    documents.finish()?;
 
-    Ok(writer.into_inner())
+    documents.into_inner().map_err(Into::into)
 }
 
 fn documents_from_json(reader: impl Read) -> Result<Vec<u8>> {
-    let mut writer = Cursor::new(Vec::new());
-    let mut documents = milli::documents::DocumentBatchBuilder::new(&mut writer)?;
+    let mut documents = DocumentsBatchBuilder::new(Vec::new());
 
-    documents.extend_from_json(reader)?;
-    documents.finish()?;
+    documents.append_json_array(reader)?;
 
-    Ok(writer.into_inner())
+    documents.into_inner().map_err(Into::into)
 }
 
 fn documents_from_csv(reader: impl Read) -> Result<Vec<u8>> {
-    let mut writer = Cursor::new(Vec::new());
-    milli::documents::DocumentBatchBuilder::from_csv(reader, &mut writer)?.finish()?;
+    let csv = csv::Reader::from_reader(reader);
 
-    Ok(writer.into_inner())
+    let mut documents = DocumentsBatchBuilder::new(Vec::new());
+    documents.append_csv(csv)?;
+
+    documents.into_inner().map_err(Into::into)
 }
 
 #[derive(Debug, StructOpt)]
@@ -423,7 +423,7 @@ impl Search {
         filter: &Option<String>,
         offset: &Option<usize>,
         limit: &Option<usize>,
-    ) -> Result<Vec<Map<String, Value>>> {
+    ) -> Result<Vec<Object>> {
         let txn = index.read_txn()?;
         let mut search = index.search(&txn);
 

http-ui/src/main.rs

@@ -3,7 +3,7 @@ mod update_store;
 use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 use std::fmt::Display;
 use std::fs::{create_dir_all, File};
-use std::io::{BufRead, BufReader, Cursor, Read};
+use std::io::{BufReader, Cursor, Read};
 use std::net::SocketAddr;
 use std::num::{NonZeroU32, NonZeroUsize};
 use std::path::PathBuf;
@@ -18,7 +18,7 @@ use either::Either;
 use flate2::read::GzDecoder;
 use futures::{stream, FutureExt, StreamExt};
 use heed::EnvOpenOptions;
-use milli::documents::DocumentBatchReader;
+use milli::documents::{DocumentsBatchBuilder, DocumentsBatchReader};
 use milli::tokenizer::TokenizerBuilder;
 use milli::update::UpdateIndexingStep::*;
 use milli::update::{
@@ -26,11 +26,11 @@ use milli::update::{
 };
 use milli::{
     obkv_to_json, CompressionType, Filter as MilliFilter, FilterCondition, FormatOptions, Index,
-    MatcherBuilder, SearchResult, SortError,
+    MatcherBuilder, Object, SearchResult, SortError,
 };
 use once_cell::sync::OnceCell;
 use serde::{Deserialize, Serialize};
-use serde_json::{Map, Value};
+use serde_json::Value;
 use structopt::StructOpt;
 use tokio::fs::File as TFile;
 use tokio::io::AsyncWriteExt;
@@ -169,11 +169,7 @@ impl<'s, A: AsRef<[u8]>> Highlighter<'s, A> {
         }
     }
 
-    fn highlight_record(
-        &self,
-        object: &mut Map<String, Value>,
-        attributes_to_highlight: &HashSet<String>,
-    ) {
+    fn highlight_record(&self, object: &mut Object, attributes_to_highlight: &HashSet<String>) {
         // TODO do we need to create a string for element that are not and needs to be highlight?
         for (key, value) in object.iter_mut() {
             if attributes_to_highlight.contains(key) {
@@ -378,7 +374,7 @@ async fn main() -> anyhow::Result<()> {
                         });
                     };
 
-                    let mut builder = milli::update::IndexDocuments::new(
+                    let builder = milli::update::IndexDocuments::new(
                         &mut wtxn,
                         &index_cloned,
                         GLOBAL_CONFIG.get().unwrap(),
@@ -399,10 +395,10 @@ async fn main() -> anyhow::Result<()> {
                         otherwise => panic!("invalid update format {:?}", otherwise),
                     };
 
-                    let documents = DocumentBatchReader::from_reader(Cursor::new(documents))?;
-
-                    builder.add_documents(documents)?;
+                    let documents = DocumentsBatchReader::from_reader(Cursor::new(documents))?;
 
+                    let (builder, user_error) = builder.add_documents(documents)?;
+                    let _count = user_error?;
                     let result = builder.execute();
 
                     match result {
@@ -708,7 +704,7 @@ async fn main() -> anyhow::Result<()> {
     #[derive(Debug, Serialize)]
     #[serde(rename_all = "camelCase")]
     struct Answer {
-        documents: Vec<Map<String, Value>>,
+        documents: Vec<Object>,
         number_of_candidates: u64,
         facets: BTreeMap<String, BTreeMap<String, u64>>,
     }
@@ -1032,35 +1028,33 @@ async fn main() -> anyhow::Result<()> {
     Ok(())
 }
 
-fn documents_from_jsonl(reader: impl io::Read) -> anyhow::Result<Vec<u8>> {
-    let mut writer = Cursor::new(Vec::new());
-    let mut documents = milli::documents::DocumentBatchBuilder::new(&mut writer)?;
+fn documents_from_jsonl(reader: impl Read) -> anyhow::Result<Vec<u8>> {
+    let mut documents = DocumentsBatchBuilder::new(Vec::new());
+    let reader = BufReader::new(reader);
 
-    for result in BufReader::new(reader).lines() {
-        let line = result?;
-        documents.extend_from_json(Cursor::new(line))?;
+    for result in serde_json::Deserializer::from_reader(reader).into_iter::<Object>() {
+        let object = result?;
+        documents.append_json_object(&object)?;
     }
 
-    documents.finish()?;
-
-    Ok(writer.into_inner())
+    documents.into_inner().map_err(Into::into)
 }
 
-fn documents_from_json(reader: impl io::Read) -> anyhow::Result<Vec<u8>> {
-    let mut writer = Cursor::new(Vec::new());
-    let mut documents = milli::documents::DocumentBatchBuilder::new(&mut writer)?;
+fn documents_from_json(reader: impl Read) -> anyhow::Result<Vec<u8>> {
+    let mut documents = DocumentsBatchBuilder::new(Vec::new());
 
-    documents.extend_from_json(reader)?;
-    documents.finish()?;
+    documents.append_json_array(reader)?;
 
-    Ok(writer.into_inner())
+    documents.into_inner().map_err(Into::into)
 }
 
-fn documents_from_csv(reader: impl io::Read) -> anyhow::Result<Vec<u8>> {
-    let mut writer = Cursor::new(Vec::new());
-    milli::documents::DocumentBatchBuilder::from_csv(reader, &mut writer)?.finish()?;
+fn documents_from_csv(reader: impl Read) -> anyhow::Result<Vec<u8>> {
+    let csv = csv::Reader::from_reader(reader);
 
-    Ok(writer.into_inner())
+    let mut documents = DocumentsBatchBuilder::new(Vec::new());
+    documents.append_csv(csv)?;
+
+    documents.into_inner().map_err(Into::into)
 }
 
 #[cfg(test)]

milli/Cargo.toml

@@ -17,7 +17,7 @@ flatten-serde-json = { path = "../flatten-serde-json" }
 fst = "0.4.7"
 fxhash = "0.2.1"
 geoutils = "0.4.1"
-grenad = { version = "0.4.1", default-features = false, features = ["tempfile"] }
+grenad = { version = "0.4.2", default-features = false, features = ["tempfile"] }
 heed = { git = "https://github.com/meilisearch/heed", tag = "v0.12.1", default-features = false, features = ["lmdb", "sync-read-txn"] }
 json-depth-checker = { path = "../json-depth-checker" }
 levenshtein_automata = { version = "0.2.1", features = ["fst_automaton"] }

milli/fuzz/.gitignore

@@ -1,2 +1,5 @@
+Cargo.lock
+target/
+
 /corpus/
 /artifacts/

milli/fuzz/fuzz_targets/indexing.rs

@@ -7,10 +7,10 @@ use anyhow::{bail, Result};
 use arbitrary_json::ArbitraryValue;
 use heed::EnvOpenOptions;
 use libfuzzer_sys::fuzz_target;
-use milli::documents::{DocumentBatchBuilder, DocumentBatchReader};
+use milli::documents::{DocumentsBatchBuilder, DocumentsBatchReader};
 use milli::update::{IndexDocuments, IndexDocumentsConfig, IndexerConfig, Settings};
 use milli::Index;
-use serde_json::Value;
+use serde_json::{Map, Value};
 
 #[cfg(target_os = "linux")]
 #[global_allocator]
@@ -19,21 +19,26 @@ static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
 /// reads json from input and write an obkv batch to writer.
 pub fn read_json(input: impl Read, writer: impl Write + Seek) -> Result<usize> {
     let writer = BufWriter::new(writer);
-    let mut builder = DocumentBatchBuilder::new(writer)?;
-    builder.extend_from_json(input)?;
+    let mut builder = DocumentsBatchBuilder::new(writer);
 
-    if builder.len() == 0 {
+    let values: Vec<Object> = serde_json::from_reader(input)?;
+    if builder.documents_count() == 0 {
         bail!("Empty payload");
     }
 
-    let count = builder.finish()?;
+    for object in values {
+        builder.append_json_object(&object)?;
+    }
 
-    Ok(count)
+    let count = builder.documents_count();
+    let vector = builder.into_inner()?;
+
+    Ok(count as usize)
 }
 
 fn index_documents(
     index: &mut milli::Index,
-    documents: DocumentBatchReader<Cursor<Vec<u8>>>,
+    documents: DocumentsBatchReader<Cursor<Vec<u8>>>,
 ) -> Result<()> {
     let config = IndexerConfig::default();
     let mut wtxn = index.write_txn()?;
@@ -98,7 +103,7 @@ fuzz_target!(|batches: Vec<Vec<ArbitraryValue>>| {
             // We ignore all malformed documents
             if let Ok(_) = read_json(json.as_bytes(), &mut documents) {
                 documents.rewind().unwrap();
-                let documents = DocumentBatchReader::from_reader(documents).unwrap();
+                let documents = DocumentsBatchReader::from_reader(documents).unwrap();
                 // A lot of errors can come out of milli and we don't know which ones are normal or not
                 // so we are only going to look for the unexpected panics.
                 let _ = index_documents(&mut index, documents);

milli/src/documents/builder.rs

@@ -1,157 +1,170 @@
-use std::collections::BTreeMap;
-use std::io;
-use std::io::{Cursor, Write};
+use std::io::{self, Write};
 
-use byteorder::{BigEndian, WriteBytesExt};
-use serde::Deserializer;
-use serde_json::Value;
+use grenad::{CompressionType, WriterBuilder};
+use serde::de::Deserializer;
+use serde_json::{to_writer, Value};
 
-use super::serde_impl::DocumentVisitor;
-use super::{ByteCounter, DocumentsBatchIndex, DocumentsMetadata, Error};
-use crate::FieldId;
+use super::{DocumentsBatchIndex, Error, DOCUMENTS_BATCH_INDEX_KEY};
+use crate::documents::serde_impl::DocumentVisitor;
+use crate::Object;
 
 /// The `DocumentsBatchBuilder` provides a way to build a documents batch in the intermediary
 /// format used by milli.
 ///
-/// The writer used by the DocumentBatchBuilder can be read using a `DocumentBatchReader` to
-/// iterate over the documents.
+/// The writer used by the `DocumentsBatchBuilder` can be read using a `DocumentsBatchReader`
+/// to iterate over the documents.
 ///
 /// ## example:
 /// ```
-/// use milli::documents::DocumentBatchBuilder;
 /// use serde_json::json;
-/// use std::io::Cursor;
+/// use milli::documents::DocumentsBatchBuilder;
 ///
-/// let json = r##"{"id": 1, "name": "foo"}"##;
-/// let mut writer = Cursor::new(Vec::new());
-/// let mut builder = DocumentBatchBuilder::new(&mut writer).unwrap();
-/// builder.extend_from_json(&mut json.as_bytes()).unwrap();
-/// builder.finish().unwrap();
+/// let json = json!({ "id": 1, "name": "foo" });
+///
+/// let mut builder = DocumentsBatchBuilder::new(Vec::new());
+/// builder.append_json_object(json.as_object().unwrap()).unwrap();
+/// let _vector = builder.into_inner().unwrap();
 /// ```
-pub struct DocumentBatchBuilder<W> {
-    inner: ByteCounter<W>,
-    index: DocumentsBatchIndex,
+pub struct DocumentsBatchBuilder<W> {
+    /// The inner grenad writer, the last value must always be the `DocumentsBatchIndex`.
+    writer: grenad::Writer<W>,
+    /// A map that creates the relation between field ids and field names.
+    fields_index: DocumentsBatchIndex,
+    /// The number of documents that were added to this builder,
+    /// it doesn't take the primary key of the documents into account at this point.
+    documents_count: u32,
+
+    /// A buffer to store a temporary obkv buffer and avoid reallocating.
     obkv_buffer: Vec<u8>,
+    /// A buffer to serialize the values and avoid reallocating,
+    /// serialized values are stored in an obkv.
     value_buffer: Vec<u8>,
-    values: BTreeMap<FieldId, Value>,
-    count: usize,
 }
 
-impl<W: io::Write + io::Seek> DocumentBatchBuilder<W> {
-    pub fn new(writer: W) -> Result<Self, Error> {
-        let index = DocumentsBatchIndex::default();
-        let mut writer = ByteCounter::new(writer);
-        // add space to write the offset of the metadata at the end of the writer
-        writer.write_u64::<BigEndian>(0)?;
-
-        Ok(Self {
-            inner: writer,
-            index,
+impl<W: Write> DocumentsBatchBuilder<W> {
+    pub fn new(writer: W) -> DocumentsBatchBuilder<W> {
+        DocumentsBatchBuilder {
+            writer: WriterBuilder::new().compression_type(CompressionType::None).build(writer),
+            fields_index: DocumentsBatchIndex::default(),
+            documents_count: 0,
             obkv_buffer: Vec::new(),
             value_buffer: Vec::new(),
-            values: BTreeMap::new(),
-            count: 0,
-        })
+        }
     }
 
-    /// Returns the number of documents that have been written to the builder.
-    pub fn len(&self) -> usize {
-        self.count
+    /// Returns the number of documents inserted into this builder.
+    pub fn documents_count(&self) -> u32 {
+        self.documents_count
     }
 
-    /// This method must be called after the document addition is terminated. It will put the
-    /// metadata at the end of the file, and write the metadata offset at the beginning on the
-    /// file.
-    pub fn finish(self) -> Result<usize, Error> {
-        let Self { inner: ByteCounter { mut writer, count: offset }, index, count, .. } = self;
+    /// Appends a new JSON object into the batch and updates the `DocumentsBatchIndex` accordingly.
+    pub fn append_json_object(&mut self, object: &Object) -> io::Result<()> {
+        // Make sure that we insert the fields ids in order as the obkv writer has this requirement.
+        let mut fields_ids: Vec<_> = object.keys().map(|k| self.fields_index.insert(&k)).collect();
+        fields_ids.sort_unstable();
 
-        let meta = DocumentsMetadata { count, index };
-
-        bincode::serialize_into(&mut writer, &meta)?;
-
-        writer.seek(io::SeekFrom::Start(0))?;
-        writer.write_u64::<BigEndian>(offset as u64)?;
-
-        writer.flush()?;
-
-        Ok(count)
+        self.obkv_buffer.clear();
+        let mut writer = obkv::KvWriter::new(&mut self.obkv_buffer);
+        for field_id in fields_ids {
+            let key = self.fields_index.name(field_id).unwrap();
+            self.value_buffer.clear();
+            to_writer(&mut self.value_buffer, &object[key])?;
+            writer.insert(field_id, &self.value_buffer)?;
         }
 
-    /// Extends the builder with json documents from a reader.
-    pub fn extend_from_json<R: io::Read>(&mut self, reader: R) -> Result<(), Error> {
+        let internal_id = self.documents_count.to_be_bytes();
+        let document_bytes = writer.into_inner()?;
+        self.writer.insert(internal_id, &document_bytes)?;
+        self.documents_count += 1;
+
+        Ok(())
+    }
+
+    /// Appends a new JSON array of objects into the batch and updates the `DocumentsBatchIndex` accordingly.
+    pub fn append_json_array<R: io::Read>(&mut self, reader: R) -> Result<(), Error> {
         let mut de = serde_json::Deserializer::from_reader(reader);
-
-        let mut visitor = DocumentVisitor {
-            inner: &mut self.inner,
-            index: &mut self.index,
-            obkv_buffer: &mut self.obkv_buffer,
-            value_buffer: &mut self.value_buffer,
-            values: &mut self.values,
-            count: &mut self.count,
-        };
-
-        de.deserialize_any(&mut visitor).map_err(Error::JsonError)?
+        let mut visitor = DocumentVisitor::new(self);
+        de.deserialize_any(&mut visitor)?
     }
 
-    /// Creates a builder from a reader of CSV documents.
-    ///
-    /// Since all fields in a csv documents are guaranteed to be ordered, we are able to perform
-    /// optimisations, and extending from another CSV is not allowed.
-    pub fn from_csv<R: io::Read>(reader: R, writer: W) -> Result<Self, Error> {
-        let mut this = Self::new(writer)?;
-        // Ensure that this is the first and only addition made with this builder
-        debug_assert!(this.index.is_empty());
-
-        let mut records = csv::Reader::from_reader(reader);
-
-        let headers = records
+    /// Appends a new CSV file into the batch and updates the `DocumentsBatchIndex` accordingly.
+    pub fn append_csv<R: io::Read>(&mut self, mut reader: csv::Reader<R>) -> Result<(), Error> {
+        // Make sure that we insert the fields ids in order as the obkv writer has this requirement.
+        let mut typed_fields_ids: Vec<_> = reader
             .headers()?
             .into_iter()
             .map(parse_csv_header)
-            .map(|(k, t)| (this.index.insert(&k), t))
-            .collect::<BTreeMap<_, _>>();
+            .map(|(k, t)| (self.fields_index.insert(k), t))
+            .enumerate()
+            .collect();
+        // Make sure that we insert the fields ids in order as the obkv writer has this requirement.
+        typed_fields_ids.sort_unstable_by_key(|(_, (fid, _))| *fid);
 
-        for (i, record) in records.into_records().enumerate() {
-            let record = record?;
-            this.obkv_buffer.clear();
-            let mut writer = obkv::KvWriter::new(&mut this.obkv_buffer);
-            for (value, (fid, ty)) in record.into_iter().zip(headers.iter()) {
-                let value = match ty {
+        let mut record = csv::StringRecord::new();
+        let mut line = 0;
+        while reader.read_record(&mut record)? {
+            // We increment here and not at the end of the while loop to take
+            // the header offset into account.
+            line += 1;
+
+            self.obkv_buffer.clear();
+            let mut writer = obkv::KvWriter::new(&mut self.obkv_buffer);
+
+            for (i, (field_id, type_)) in typed_fields_ids.iter() {
+                self.value_buffer.clear();
+
+                let value = &record[*i];
+                match type_ {
                     AllowedType::Number => {
                         if value.trim().is_empty() {
-                            Value::Null
+                            to_writer(&mut self.value_buffer, &Value::Null)?;
                         } else {
-                            value.trim().parse::<f64>().map(Value::from).map_err(|error| {
-                                Error::ParseFloat {
-                                    error,
-                                    // +1 for the header offset.
-                                    line: i + 1,
-                                    value: value.to_string(),
+                            match value.trim().parse::<f64>() {
+                                Ok(float) => {
+                                    to_writer(&mut self.value_buffer, &float)?;
+                                }
+                                Err(error) => {
+                                    return Err(Error::ParseFloat {
+                                        error,
+                                        line,
+                                        value: value.to_string(),
+                                    });
+                                }
                             }
-                            })?
                         }
                     }
                     AllowedType::String => {
                         if value.is_empty() {
-                            Value::Null
+                            to_writer(&mut self.value_buffer, &Value::Null)?;
                         } else {
-                            Value::String(value.to_string())
+                            to_writer(&mut self.value_buffer, value)?;
                         }
                     }
-                };
-
-                this.value_buffer.clear();
-                serde_json::to_writer(Cursor::new(&mut this.value_buffer), &value)?;
-                writer.insert(*fid, &this.value_buffer)?;
                 }
 
-            this.inner.write_u32::<BigEndian>(this.obkv_buffer.len() as u32)?;
-            this.inner.write_all(&this.obkv_buffer)?;
-
-            this.count += 1;
+                // We insert into the obkv writer the value buffer that has been filled just above.
+                writer.insert(*field_id, &self.value_buffer)?;
             }
 
-        Ok(this)
+            let internal_id = self.documents_count.to_be_bytes();
+            let document_bytes = writer.into_inner()?;
+            self.writer.insert(internal_id, &document_bytes)?;
+            self.documents_count += 1;
+        }
+
+        Ok(())
+    }
+
+    /// Flushes the content on disk and stores the final version of the `DocumentsBatchIndex`.
+    pub fn into_inner(mut self) -> io::Result<W> {
+        let DocumentsBatchBuilder { mut writer, fields_index, .. } = self;
+
+        // We serialize and insert the `DocumentsBatchIndex` as the last key of the grenad writer.
+        self.value_buffer.clear();
+        to_writer(&mut self.value_buffer, &fields_index)?;
+        writer.insert(DOCUMENTS_BATCH_INDEX_KEY, &self.value_buffer)?;
+
+        writer.into_inner()
     }
 }
 
@@ -161,16 +174,16 @@ enum AllowedType {
     Number,
 }
 
-fn parse_csv_header(header: &str) -> (String, AllowedType) {
+fn parse_csv_header(header: &str) -> (&str, AllowedType) {
     // if there are several separators we only split on the last one.
     match header.rsplit_once(':') {
         Some((field_name, field_type)) => match field_type {
-            "string" => (field_name.to_string(), AllowedType::String),
-            "number" => (field_name.to_string(), AllowedType::Number),
+            "string" => (field_name, AllowedType::String),
+            "number" => (field_name, AllowedType::Number),
             // if the pattern isn't reconized, we keep the whole field.
-            _otherwise => (header.to_string(), AllowedType::String),
+            _otherwise => (header, AllowedType::String),
         },
-        None => (header.to_string(), AllowedType::String),
+        None => (header, AllowedType::String),
     }
 }
 
@@ -178,35 +191,20 @@ fn parse_csv_header(header: &str) -> (String, AllowedType) {
 mod test {
     use std::io::Cursor;
 
-    use serde_json::{json, Map};
+    use serde_json::json;
 
     use super::*;
-    use crate::documents::DocumentBatchReader;
-
-    fn obkv_to_value(obkv: &obkv::KvReader<FieldId>, index: &DocumentsBatchIndex) -> Value {
-        let mut map = Map::new();
-
-        for (fid, value) in obkv.iter() {
-            let field_name = index.name(fid).unwrap().clone();
-            let value: Value = serde_json::from_slice(value).unwrap();
-
-            map.insert(field_name, value);
-        }
-
-        Value::Object(map)
-    }
+    use crate::documents::{obkv_to_object, DocumentsBatchReader};
 
     #[test]
     fn add_single_documents_json() {
-        let mut cursor = Cursor::new(Vec::new());
-        let mut builder = DocumentBatchBuilder::new(&mut cursor).unwrap();
-
         let json = serde_json::json!({
             "id": 1,
             "field": "hello!",
         });
 
-        builder.extend_from_json(Cursor::new(serde_json::to_vec(&json).unwrap())).unwrap();
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_json_object(json.as_object().unwrap()).unwrap();
 
         let json = serde_json::json!({
             "blabla": false,
@@ -214,100 +212,64 @@ mod test {
             "id": 1,
         });
 
-        builder.extend_from_json(Cursor::new(serde_json::to_vec(&json).unwrap())).unwrap();
+        builder.append_json_object(json.as_object().unwrap()).unwrap();
 
-        assert_eq!(builder.len(), 2);
+        assert_eq!(builder.documents_count(), 2);
+        let vector = builder.into_inner().unwrap();
 
-        builder.finish().unwrap();
-
-        cursor.set_position(0);
-
-        let mut reader = DocumentBatchReader::from_reader(cursor).unwrap();
-
-        let (index, document) = reader.next_document_with_index().unwrap().unwrap();
+        let (mut cursor, index) = DocumentsBatchReader::from_reader(Cursor::new(vector))
+            .unwrap()
+            .into_cursor_and_fields_index();
         assert_eq!(index.len(), 3);
+
+        let document = cursor.next_document().unwrap().unwrap();
         assert_eq!(document.iter().count(), 2);
 
-        let (index, document) = reader.next_document_with_index().unwrap().unwrap();
-        assert_eq!(index.len(), 3);
+        let document = cursor.next_document().unwrap().unwrap();
         assert_eq!(document.iter().count(), 3);
 
-        assert!(reader.next_document_with_index().unwrap().is_none());
-    }
-
-    #[test]
-    fn add_documents_seq_json() {
-        let mut cursor = Cursor::new(Vec::new());
-        let mut builder = DocumentBatchBuilder::new(&mut cursor).unwrap();
-
-        let json = serde_json::json!([{
-            "id": 1,
-            "field": "hello!",
-        },{
-            "blabla": false,
-            "field": "hello!",
-            "id": 1,
-        }
-        ]);
-
-        builder.extend_from_json(Cursor::new(serde_json::to_vec(&json).unwrap())).unwrap();
-
-        assert_eq!(builder.len(), 2);
-
-        builder.finish().unwrap();
-
-        cursor.set_position(0);
-
-        let mut reader = DocumentBatchReader::from_reader(cursor).unwrap();
-
-        let (index, document) = reader.next_document_with_index().unwrap().unwrap();
-        assert_eq!(index.len(), 3);
-        assert_eq!(document.iter().count(), 2);
-
-        let (index, document) = reader.next_document_with_index().unwrap().unwrap();
-        assert_eq!(index.len(), 3);
-        assert_eq!(document.iter().count(), 3);
-
-        assert!(reader.next_document_with_index().unwrap().is_none());
+        assert!(cursor.next_document().unwrap().is_none());
     }
 
     #[test]
     fn add_documents_csv() {
-        let mut cursor = Cursor::new(Vec::new());
+        let csv_content = "id:number,field:string\n1,hello!\n2,blabla";
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content));
 
-        let csv = "id:number,field:string\n1,hello!\n2,blabla";
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_csv(csv).unwrap();
+        assert_eq!(builder.documents_count(), 2);
+        let vector = builder.into_inner().unwrap();
 
-        let builder =
-            DocumentBatchBuilder::from_csv(Cursor::new(csv.as_bytes()), &mut cursor).unwrap();
-        builder.finish().unwrap();
-
-        cursor.set_position(0);
-
-        let mut reader = DocumentBatchReader::from_reader(cursor).unwrap();
-
-        let (index, document) = reader.next_document_with_index().unwrap().unwrap();
+        let (mut cursor, index) = DocumentsBatchReader::from_reader(Cursor::new(vector))
+            .unwrap()
+            .into_cursor_and_fields_index();
         assert_eq!(index.len(), 2);
+
+        let document = cursor.next_document().unwrap().unwrap();
         assert_eq!(document.iter().count(), 2);
 
-        let (_index, document) = reader.next_document_with_index().unwrap().unwrap();
+        let document = cursor.next_document().unwrap().unwrap();
         assert_eq!(document.iter().count(), 2);
 
-        assert!(reader.next_document_with_index().unwrap().is_none());
+        assert!(cursor.next_document().unwrap().is_none());
     }
 
     #[test]
     fn simple_csv_document() {
-        let documents = r#"city,country,pop
+        let csv_content = r#"city,country,pop
 "Boston","United States","4628910""#;
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content));
 
-        let mut buf = Vec::new();
-        DocumentBatchBuilder::from_csv(documents.as_bytes(), Cursor::new(&mut buf))
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_csv(csv).unwrap();
+        let vector = builder.into_inner().unwrap();
+
+        let (mut cursor, index) = DocumentsBatchReader::from_reader(Cursor::new(vector))
             .unwrap()
-            .finish()
-            .unwrap();
-        let mut reader = DocumentBatchReader::from_reader(Cursor::new(buf)).unwrap();
-        let (index, doc) = reader.next_document_with_index().unwrap().unwrap();
-        let val = obkv_to_value(&doc, index);
+            .into_cursor_and_fields_index();
+        let doc = cursor.next_document().unwrap().unwrap();
+        let val = obkv_to_object(&doc, &index).map(Value::from).unwrap();
 
         assert_eq!(
             val,
@@ -318,22 +280,25 @@ mod test {
             })
         );
 
-        assert!(reader.next_document_with_index().unwrap().is_none());
+        assert!(cursor.next_document().unwrap().is_none());
     }
 
     #[test]
     fn coma_in_field() {
-        let documents = r#"city,country,pop
+        let csv_content = r#"city,country,pop
 "Boston","United, States","4628910""#;
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content));
 
-        let mut buf = Vec::new();
-        DocumentBatchBuilder::from_csv(documents.as_bytes(), Cursor::new(&mut buf))
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_csv(csv).unwrap();
+        let vector = builder.into_inner().unwrap();
+
+        let (mut cursor, index) = DocumentsBatchReader::from_reader(Cursor::new(vector))
             .unwrap()
-            .finish()
-            .unwrap();
-        let mut reader = DocumentBatchReader::from_reader(Cursor::new(buf)).unwrap();
-        let (index, doc) = reader.next_document_with_index().unwrap().unwrap();
-        let val = obkv_to_value(&doc, index);
+            .into_cursor_and_fields_index();
+
+        let doc = cursor.next_document().unwrap().unwrap();
+        let val = obkv_to_object(&doc, &index).map(Value::from).unwrap();
 
         assert_eq!(
             val,
@@ -347,17 +312,20 @@ mod test {
 
     #[test]
     fn quote_in_field() {
-        let documents = r#"city,country,pop
+        let csv_content = r#"city,country,pop
 "Boston","United"" States","4628910""#;
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content));
 
-        let mut buf = Vec::new();
-        DocumentBatchBuilder::from_csv(documents.as_bytes(), Cursor::new(&mut buf))
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_csv(csv).unwrap();
+        let vector = builder.into_inner().unwrap();
+
+        let (mut cursor, index) = DocumentsBatchReader::from_reader(Cursor::new(vector))
             .unwrap()
-            .finish()
-            .unwrap();
-        let mut reader = DocumentBatchReader::from_reader(Cursor::new(buf)).unwrap();
-        let (index, doc) = reader.next_document_with_index().unwrap().unwrap();
-        let val = obkv_to_value(&doc, index);
+            .into_cursor_and_fields_index();
+
+        let doc = cursor.next_document().unwrap().unwrap();
+        let val = obkv_to_object(&doc, &index).map(Value::from).unwrap();
 
         assert_eq!(
             val,
@@ -371,17 +339,20 @@ mod test {
 
     #[test]
     fn integer_in_field() {
-        let documents = r#"city,country,pop:number
+        let csv_content = r#"city,country,pop:number
 "Boston","United States","4628910""#;
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content));
 
-        let mut buf = Vec::new();
-        DocumentBatchBuilder::from_csv(documents.as_bytes(), Cursor::new(&mut buf))
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_csv(csv).unwrap();
+        let vector = builder.into_inner().unwrap();
+
+        let (mut cursor, index) = DocumentsBatchReader::from_reader(Cursor::new(vector))
             .unwrap()
-            .finish()
-            .unwrap();
-        let mut reader = DocumentBatchReader::from_reader(Cursor::new(buf)).unwrap();
-        let (index, doc) = reader.next_document_with_index().unwrap().unwrap();
-        let val = obkv_to_value(&doc, index);
+            .into_cursor_and_fields_index();
+
+        let doc = cursor.next_document().unwrap().unwrap();
+        let val = obkv_to_object(&doc, &index).map(Value::from).unwrap();
 
         assert_eq!(
             val,
@@ -395,17 +366,20 @@ mod test {
 
     #[test]
     fn float_in_field() {
-        let documents = r#"city,country,pop:number
+        let csv_content = r#"city,country,pop:number
 "Boston","United States","4628910.01""#;
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content));
 
-        let mut buf = Vec::new();
-        DocumentBatchBuilder::from_csv(documents.as_bytes(), Cursor::new(&mut buf))
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_csv(csv).unwrap();
+        let vector = builder.into_inner().unwrap();
+
+        let (mut cursor, index) = DocumentsBatchReader::from_reader(Cursor::new(vector))
             .unwrap()
-            .finish()
-            .unwrap();
-        let mut reader = DocumentBatchReader::from_reader(Cursor::new(buf)).unwrap();
-        let (index, doc) = reader.next_document_with_index().unwrap().unwrap();
-        let val = obkv_to_value(&doc, index);
+            .into_cursor_and_fields_index();
+
+        let doc = cursor.next_document().unwrap().unwrap();
+        let val = obkv_to_object(&doc, &index).map(Value::from).unwrap();
 
         assert_eq!(
             val,
@@ -419,17 +393,20 @@ mod test {
 
     #[test]
     fn several_colon_in_header() {
-        let documents = r#"city:love:string,country:state,pop
+        let csv_content = r#"city:love:string,country:state,pop
 "Boston","United States","4628910""#;
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content));
 
-        let mut buf = Vec::new();
-        DocumentBatchBuilder::from_csv(documents.as_bytes(), Cursor::new(&mut buf))
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_csv(csv).unwrap();
+        let vector = builder.into_inner().unwrap();
+
+        let (mut cursor, index) = DocumentsBatchReader::from_reader(Cursor::new(vector))
             .unwrap()
-            .finish()
-            .unwrap();
-        let mut reader = DocumentBatchReader::from_reader(Cursor::new(buf)).unwrap();
-        let (index, doc) = reader.next_document_with_index().unwrap().unwrap();
-        let val = obkv_to_value(&doc, index);
+            .into_cursor_and_fields_index();
+
+        let doc = cursor.next_document().unwrap().unwrap();
+        let val = obkv_to_object(&doc, &index).map(Value::from).unwrap();
 
         assert_eq!(
             val,
@@ -443,17 +420,20 @@ mod test {
 
     #[test]
     fn ending_by_colon_in_header() {
-        let documents = r#"city:,country,pop
+        let csv_content = r#"city:,country,pop
 "Boston","United States","4628910""#;
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content));
 
-        let mut buf = Vec::new();
-        DocumentBatchBuilder::from_csv(documents.as_bytes(), Cursor::new(&mut buf))
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_csv(csv).unwrap();
+        let vector = builder.into_inner().unwrap();
+
+        let (mut cursor, index) = DocumentsBatchReader::from_reader(Cursor::new(vector))
             .unwrap()
-            .finish()
-            .unwrap();
-        let mut reader = DocumentBatchReader::from_reader(Cursor::new(buf)).unwrap();
-        let (index, doc) = reader.next_document_with_index().unwrap().unwrap();
-        let val = obkv_to_value(&doc, index);
+            .into_cursor_and_fields_index();
+
+        let doc = cursor.next_document().unwrap().unwrap();
+        let val = obkv_to_object(&doc, &index).map(Value::from).unwrap();
 
         assert_eq!(
             val,
@@ -467,17 +447,20 @@ mod test {
 
     #[test]
     fn starting_by_colon_in_header() {
-        let documents = r#":city,country,pop
+        let csv_content = r#":city,country,pop
 "Boston","United States","4628910""#;
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content));
 
-        let mut buf = Vec::new();
-        DocumentBatchBuilder::from_csv(documents.as_bytes(), Cursor::new(&mut buf))
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_csv(csv).unwrap();
+        let vector = builder.into_inner().unwrap();
+
+        let (mut cursor, index) = DocumentsBatchReader::from_reader(Cursor::new(vector))
             .unwrap()
-            .finish()
-            .unwrap();
-        let mut reader = DocumentBatchReader::from_reader(Cursor::new(buf)).unwrap();
-        let (index, doc) = reader.next_document_with_index().unwrap().unwrap();
-        let val = obkv_to_value(&doc, index);
+            .into_cursor_and_fields_index();
+
+        let doc = cursor.next_document().unwrap().unwrap();
+        let val = obkv_to_object(&doc, &index).map(Value::from).unwrap();
 
         assert_eq!(
             val,
@@ -492,32 +475,37 @@ mod test {
     #[ignore]
     #[test]
     fn starting_by_colon_in_header2() {
-        let documents = r#":string,country,pop
+        let csv_content = r#":string,country,pop
 "Boston","United States","4628910""#;
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content));
 
-        let mut buf = Vec::new();
-        DocumentBatchBuilder::from_csv(documents.as_bytes(), Cursor::new(&mut buf))
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_csv(csv).unwrap();
+        let vector = builder.into_inner().unwrap();
+
+        let (mut cursor, _) = DocumentsBatchReader::from_reader(Cursor::new(vector))
             .unwrap()
-            .finish()
-            .unwrap();
-        let mut reader = DocumentBatchReader::from_reader(Cursor::new(buf)).unwrap();
+            .into_cursor_and_fields_index();
 
-        assert!(reader.next_document_with_index().is_err());
+        assert!(cursor.next_document().is_err());
     }
 
     #[test]
     fn double_colon_in_header() {
-        let documents = r#"city::string,country,pop
+        let csv_content = r#"city::string,country,pop
 "Boston","United States","4628910""#;
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content));
 
-        let mut buf = Vec::new();
-        DocumentBatchBuilder::from_csv(documents.as_bytes(), Cursor::new(&mut buf))
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_csv(csv).unwrap();
+        let vector = builder.into_inner().unwrap();
+
+        let (mut cursor, index) = DocumentsBatchReader::from_reader(Cursor::new(vector))
             .unwrap()
-            .finish()
-            .unwrap();
-        let mut reader = DocumentBatchReader::from_reader(Cursor::new(buf)).unwrap();
-        let (index, doc) = reader.next_document_with_index().unwrap().unwrap();
-        let val = obkv_to_value(&doc, index);
+            .into_cursor_and_fields_index();
+
+        let doc = cursor.next_document().unwrap().unwrap();
+        let val = obkv_to_object(&doc, &index).map(Value::from).unwrap();
 
         assert_eq!(
             val,
@@ -531,34 +519,32 @@ mod test {
 
     #[test]
     fn bad_type_in_header() {
-        let documents = r#"city,country:number,pop
+        let csv_content = r#"city,country:number,pop
 "Boston","United States","4628910""#;
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content));
 
-        let mut buf = Vec::new();
-        assert!(
-            DocumentBatchBuilder::from_csv(documents.as_bytes(), Cursor::new(&mut buf)).is_err()
-        );
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        assert!(builder.append_csv(csv).is_err());
     }
 
     #[test]
     fn bad_column_count1() {
-        let documents = r#"city,country,pop
-"Boston","United States","4628910", "too much""#;
+        let csv_content = r#"city,country,pop
+"Boston","United States","4628910", "too much
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content"#;
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content));
 
-        let mut buf = Vec::new();
-        assert!(
-            DocumentBatchBuilder::from_csv(documents.as_bytes(), Cursor::new(&mut buf)).is_err()
-        );
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        assert!(builder.append_csv(csv).is_err());
     }
 
     #[test]
     fn bad_column_count2() {
-        let documents = r#"city,country,pop
+        let csv_content = r#"city,country,pop
 "Boston","United States""#;
+        let csv = csv::Reader::from_reader(Cursor::new(csv_content));
 
-        let mut buf = Vec::new();
-        assert!(
-            DocumentBatchBuilder::from_csv(documents.as_bytes(), Cursor::new(&mut buf)).is_err()
-        );
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        assert!(builder.append_csv(csv).is_err());
     }
 }

milli/src/documents/enriched.rs

@@ -0,0 +1,109 @@
+use std::fs::File;
+use std::{io, str};
+
+use obkv::KvReader;
+
+use super::{
+    DocumentsBatchCursor, DocumentsBatchCursorError, DocumentsBatchIndex, DocumentsBatchReader,
+    Error,
+};
+use crate::update::DocumentId;
+use crate::FieldId;
+
+/// The `EnrichedDocumentsBatchReader` provides a way to iterate over documents that have
+/// been created with a `DocumentsBatchWriter` and, for the enriched data,
+/// a simple `grenad::Reader<File>`.
+///
+/// The documents are returned in the form of `obkv::Reader` where each field is identified with a
+/// `FieldId`. The mapping between the field ids and the field names is done thanks to the index.
+pub struct EnrichedDocumentsBatchReader<R> {
+    documents: DocumentsBatchReader<R>,
+    primary_key: String,
+    external_ids: grenad::ReaderCursor<File>,
+}
+
+impl<R: io::Read + io::Seek> EnrichedDocumentsBatchReader<R> {
+    pub fn new(
+        documents: DocumentsBatchReader<R>,
+        primary_key: String,
+        external_ids: grenad::Reader<File>,
+    ) -> Result<Self, Error> {
+        if documents.documents_count() as u64 == external_ids.len() {
+            Ok(EnrichedDocumentsBatchReader {
+                documents,
+                primary_key,
+                external_ids: external_ids.into_cursor()?,
+            })
+        } else {
+            Err(Error::InvalidEnrichedData)
+        }
+    }
+
+    pub fn documents_count(&self) -> u32 {
+        self.documents.documents_count()
+    }
+
+    pub fn primary_key(&self) -> &str {
+        &self.primary_key
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.documents.is_empty()
+    }
+
+    pub fn documents_batch_index(&self) -> &DocumentsBatchIndex {
+        self.documents.documents_batch_index()
+    }
+
+    /// This method returns a forward cursor over the enriched documents.
+    pub fn into_cursor_and_fields_index(
+        self,
+    ) -> (EnrichedDocumentsBatchCursor<R>, DocumentsBatchIndex) {
+        let EnrichedDocumentsBatchReader { documents, primary_key, mut external_ids } = self;
+        let (documents, fields_index) = documents.into_cursor_and_fields_index();
+        external_ids.reset();
+        (EnrichedDocumentsBatchCursor { documents, primary_key, external_ids }, fields_index)
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct EnrichedDocument<'a> {
+    pub document: KvReader<'a, FieldId>,
+    pub document_id: DocumentId,
+}
+
+pub struct EnrichedDocumentsBatchCursor<R> {
+    documents: DocumentsBatchCursor<R>,
+    primary_key: String,
+    external_ids: grenad::ReaderCursor<File>,
+}
+
+impl<R> EnrichedDocumentsBatchCursor<R> {
+    pub fn primary_key(&self) -> &str {
+        &self.primary_key
+    }
+    /// Resets the cursor to be able to read from the start again.
+    pub fn reset(&mut self) {
+        self.documents.reset();
+        self.external_ids.reset();
+    }
+}
+
+impl<R: io::Read + io::Seek> EnrichedDocumentsBatchCursor<R> {
+    /// Returns the next document, starting from the first one. Subsequent calls to
+    /// `next_document` advance the document reader until all the documents have been read.
+    pub fn next_enriched_document(
+        &mut self,
+    ) -> Result<Option<EnrichedDocument>, DocumentsBatchCursorError> {
+        let document = self.documents.next_document()?;
+        let document_id = match self.external_ids.move_on_next()? {
+            Some((_, bytes)) => serde_json::from_slice(bytes).map(Some)?,
+            None => None,
+        };
+
+        match document.zip(document_id) {
+            Some((document, document_id)) => Ok(Some(EnrichedDocument { document, document_id })),
+            None => Ok(None),
+        }
+    }
+}

milli/src/documents/mod.rs

@@ -1,24 +1,41 @@
 mod builder;
-/// The documents module defines an intermediary document format that milli uses for indexation, and
-/// provides an API to easily build and read such documents.
-///
-/// The `DocumentBatchBuilder` interface allows to write batches of documents to a writer, that can
-/// later be read by milli using the `DocumentBatchReader` interface.
+mod enriched;
 mod reader;
 mod serde_impl;
 
 use std::fmt::{self, Debug};
 use std::io;
+use std::str::Utf8Error;
 
 use bimap::BiHashMap;
-pub use builder::DocumentBatchBuilder;
-pub use reader::DocumentBatchReader;
+pub use builder::DocumentsBatchBuilder;
+pub use enriched::{EnrichedDocument, EnrichedDocumentsBatchCursor, EnrichedDocumentsBatchReader};
+use obkv::KvReader;
+pub use reader::{DocumentsBatchCursor, DocumentsBatchCursorError, DocumentsBatchReader};
 use serde::{Deserialize, Serialize};
 
-use crate::FieldId;
+use crate::error::{FieldIdMapMissingEntry, InternalError};
+use crate::{FieldId, Object, Result};
+
+/// The key that is used to store the `DocumentsBatchIndex` datastructure,
+/// it is the absolute last key of the list.
+const DOCUMENTS_BATCH_INDEX_KEY: [u8; 8] = u64::MAX.to_be_bytes();
+
+/// Helper function to convert an obkv reader into a JSON object.
+pub fn obkv_to_object(obkv: &KvReader<FieldId>, index: &DocumentsBatchIndex) -> Result<Object> {
+    obkv.iter()
+        .map(|(field_id, value)| {
+            let field_name = index.name(field_id).ok_or_else(|| {
+                FieldIdMapMissingEntry::FieldId { field_id, process: "obkv_to_object" }
+            })?;
+            let value = serde_json::from_slice(value).map_err(InternalError::SerdeJson)?;
+            Ok((field_name.to_string(), value))
+        })
+        .collect()
+}
 
 /// A bidirectional map that links field ids to their name in a document batch.
-#[derive(Default, Debug, Serialize, Deserialize)]
+#[derive(Default, Clone, Debug, Serialize, Deserialize)]
 pub struct DocumentsBatchIndex(pub BiHashMap<FieldId, String>);
 
 impl DocumentsBatchIndex {
@@ -46,15 +63,16 @@ impl DocumentsBatchIndex {
         self.0.iter()
     }
 
-    pub fn name(&self, id: FieldId) -> Option<&String> {
-        self.0.get_by_left(&id)
+    pub fn name(&self, id: FieldId) -> Option<&str> {
+        self.0.get_by_left(&id).map(AsRef::as_ref)
     }
 
-    pub fn recreate_json(
-        &self,
-        document: &obkv::KvReaderU16,
-    ) -> Result<serde_json::Map<String, serde_json::Value>, crate::Error> {
-        let mut map = serde_json::Map::new();
+    pub fn id(&self, name: &str) -> Option<FieldId> {
+        self.0.get_by_right(name).cloned()
+    }
+
+    pub fn recreate_json(&self, document: &obkv::KvReaderU16) -> Result<Object> {
+        let mut map = Object::new();
 
         for (k, v) in document.iter() {
             // TODO: TAMO: update the error type
@@ -69,50 +87,22 @@ impl DocumentsBatchIndex {
     }
 }
 
-#[derive(Debug, Serialize, Deserialize)]
-struct DocumentsMetadata {
-    count: usize,
-    index: DocumentsBatchIndex,
-}
-
-pub struct ByteCounter<W> {
-    count: usize,
-    writer: W,
-}
-
-impl<W> ByteCounter<W> {
-    fn new(writer: W) -> Self {
-        Self { count: 0, writer }
-    }
-}
-
-impl<W: io::Write> io::Write for ByteCounter<W> {
-    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        let count = self.writer.write(buf)?;
-        self.count += count;
-        Ok(count)
-    }
-
-    fn flush(&mut self) -> io::Result<()> {
-        self.writer.flush()
-    }
-}
-
 #[derive(Debug)]
 pub enum Error {
     ParseFloat { error: std::num::ParseFloatError, line: usize, value: String },
     InvalidDocumentFormat,
-    Custom(String),
-    JsonError(serde_json::Error),
-    CsvError(csv::Error),
-    Serialize(bincode::Error),
+    InvalidEnrichedData,
+    InvalidUtf8(Utf8Error),
+    Csv(csv::Error),
+    Json(serde_json::Error),
+    Serialize(serde_json::Error),
+    Grenad(grenad::Error),
     Io(io::Error),
-    DocumentTooLarge,
 }
 
 impl From<csv::Error> for Error {
     fn from(e: csv::Error) -> Self {
-        Self::CsvError(e)
+        Self::Csv(e)
     }
 }
 
@@ -122,15 +112,21 @@ impl From<io::Error> for Error {
     }
 }
 
-impl From<bincode::Error> for Error {
-    fn from(other: bincode::Error) -> Self {
-        Self::Serialize(other)
+impl From<serde_json::Error> for Error {
+    fn from(other: serde_json::Error) -> Self {
+        Self::Json(other)
     }
 }
 
-impl From<serde_json::Error> for Error {
-    fn from(other: serde_json::Error) -> Self {
-        Self::JsonError(other)
+impl From<grenad::Error> for Error {
+    fn from(other: grenad::Error) -> Self {
+        Self::Grenad(other)
+    }
+}
+
+impl From<Utf8Error> for Error {
+    fn from(other: Utf8Error) -> Self {
+        Self::InvalidUtf8(other)
     }
 }
 
@@ -140,13 +136,16 @@ impl fmt::Display for Error {
             Error::ParseFloat { error, line, value } => {
                 write!(f, "Error parsing number {:?} at line {}: {}", value, line, error)
             }
-            Error::Custom(s) => write!(f, "Unexpected serialization error: {}", s),
-            Error::InvalidDocumentFormat => f.write_str("Invalid document addition format."),
-            Error::JsonError(err) => write!(f, "Couldn't serialize document value: {}", err),
+            Error::InvalidDocumentFormat => {
+                f.write_str("Invalid document addition format, missing the documents batch index.")
+            }
+            Error::InvalidEnrichedData => f.write_str("Invalid enriched data."),
+            Error::InvalidUtf8(e) => write!(f, "{}", e),
             Error::Io(e) => write!(f, "{}", e),
-            Error::DocumentTooLarge => f.write_str("Provided document is too large (>2Gib)"),
             Error::Serialize(e) => write!(f, "{}", e),
-            Error::CsvError(e) => write!(f, "{}", e),
+            Error::Grenad(e) => write!(f, "{}", e),
+            Error::Csv(e) => write!(f, "{}", e),
+            Error::Json(e) => write!(f, "{}", e),
         }
     }
 }
@@ -158,15 +157,25 @@ impl std::error::Error for Error {}
 macro_rules! documents {
     ($data:tt) => {{
         let documents = serde_json::json!($data);
-        let mut writer = std::io::Cursor::new(Vec::new());
-        let mut builder = crate::documents::DocumentBatchBuilder::new(&mut writer).unwrap();
-        let documents = serde_json::to_vec(&documents).unwrap();
-        builder.extend_from_json(std::io::Cursor::new(documents)).unwrap();
-        builder.finish().unwrap();
+        let documents = match documents {
+            object @ serde_json::Value::Object(_) => vec![object],
+            serde_json::Value::Array(objects) => objects,
+            invalid => {
+                panic!("an array of objects must be specified, {:#?} is not an array", invalid)
+            }
+        };
 
-        writer.set_position(0);
+        let mut builder = crate::documents::DocumentsBatchBuilder::new(Vec::new());
+        for document in documents {
+            let object = match document {
+                serde_json::Value::Object(object) => object,
+                invalid => panic!("an object must be specified, {:#?} is not an object", invalid),
+            };
+            builder.append_json_object(&object).unwrap();
+        }
 
-        crate::documents::DocumentBatchReader::from_reader(writer).unwrap()
+        let vector = builder.into_inner().unwrap();
+        crate::documents::DocumentsBatchReader::from_reader(std::io::Cursor::new(vector)).unwrap()
     }};
 }
 
@@ -180,7 +189,7 @@ mod test {
 
     #[test]
     fn create_documents_no_errors() {
-        let json = json!({
+        let value = json!({
             "number": 1,
             "string": "this is a field",
             "array": ["an", "array"],
@@ -190,26 +199,18 @@ mod test {
             "bool": true
         });
 
-        let json = serde_json::to_vec(&json).unwrap();
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_json_object(value.as_object().unwrap()).unwrap();
+        let vector = builder.into_inner().unwrap();
 
-        let mut v = Vec::new();
-        let mut cursor = io::Cursor::new(&mut v);
+        let (mut documents, index) = DocumentsBatchReader::from_reader(Cursor::new(vector))
+            .unwrap()
+            .into_cursor_and_fields_index();
 
-        let mut builder = DocumentBatchBuilder::new(&mut cursor).unwrap();
-
-        builder.extend_from_json(Cursor::new(json)).unwrap();
-
-        builder.finish().unwrap();
-
-        let mut documents =
-            DocumentBatchReader::from_reader(io::Cursor::new(cursor.into_inner())).unwrap();
-
-        assert_eq!(documents.index().iter().count(), 5);
-
-        let reader = documents.next_document_with_index().unwrap().unwrap();
-
-        assert_eq!(reader.1.iter().count(), 5);
-        assert!(documents.next_document_with_index().unwrap().is_none());
+        assert_eq!(index.iter().count(), 5);
+        let reader = documents.next_document().unwrap().unwrap();
+        assert_eq!(reader.iter().count(), 5);
+        assert!(documents.next_document().unwrap().is_none());
     }
 
     #[test]
@@ -221,101 +222,56 @@ mod test {
             "toto": false,
         });
 
-        let doc1 = serde_json::to_vec(&doc1).unwrap();
-        let doc2 = serde_json::to_vec(&doc2).unwrap();
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_json_object(doc1.as_object().unwrap()).unwrap();
+        builder.append_json_object(doc2.as_object().unwrap()).unwrap();
+        let vector = builder.into_inner().unwrap();
 
-        let mut v = Vec::new();
-        let mut cursor = io::Cursor::new(&mut v);
-
-        let mut builder = DocumentBatchBuilder::new(&mut cursor).unwrap();
-
-        builder.extend_from_json(Cursor::new(doc1)).unwrap();
-        builder.extend_from_json(Cursor::new(doc2)).unwrap();
-
-        builder.finish().unwrap();
-
-        let mut documents =
-            DocumentBatchReader::from_reader(io::Cursor::new(cursor.into_inner())).unwrap();
-
-        assert_eq!(documents.index().iter().count(), 2);
-
-        let reader = documents.next_document_with_index().unwrap().unwrap();
-
-        assert_eq!(reader.1.iter().count(), 1);
-        assert!(documents.next_document_with_index().unwrap().is_some());
-        assert!(documents.next_document_with_index().unwrap().is_none());
-    }
-
-    #[test]
-    fn add_documents_array() {
-        let docs = json!([
-            { "toto": false },
-            { "tata": "hello" },
-        ]);
-
-        let docs = serde_json::to_vec(&docs).unwrap();
-
-        let mut v = Vec::new();
-        let mut cursor = io::Cursor::new(&mut v);
-
-        let mut builder = DocumentBatchBuilder::new(&mut cursor).unwrap();
-
-        builder.extend_from_json(Cursor::new(docs)).unwrap();
-
-        builder.finish().unwrap();
-
-        let mut documents =
-            DocumentBatchReader::from_reader(io::Cursor::new(cursor.into_inner())).unwrap();
-
-        assert_eq!(documents.index().iter().count(), 2);
-
-        let reader = documents.next_document_with_index().unwrap().unwrap();
-
-        assert_eq!(reader.1.iter().count(), 1);
-        assert!(documents.next_document_with_index().unwrap().is_some());
-        assert!(documents.next_document_with_index().unwrap().is_none());
-    }
-
-    #[test]
-    fn add_invalid_document_format() {
-        let mut v = Vec::new();
-        let mut cursor = io::Cursor::new(&mut v);
-
-        let mut builder = DocumentBatchBuilder::new(&mut cursor).unwrap();
-
-        let docs = json!([[
-            { "toto": false },
-            { "tata": "hello" },
-        ]]);
-
-        let docs = serde_json::to_vec(&docs).unwrap();
-        assert!(builder.extend_from_json(Cursor::new(docs)).is_err());
-
-        let docs = json!("hello");
-        let docs = serde_json::to_vec(&docs).unwrap();
-
-        assert!(builder.extend_from_json(Cursor::new(docs)).is_err());
+        let (mut documents, index) = DocumentsBatchReader::from_reader(io::Cursor::new(vector))
+            .unwrap()
+            .into_cursor_and_fields_index();
+        assert_eq!(index.iter().count(), 2);
+        let reader = documents.next_document().unwrap().unwrap();
+        assert_eq!(reader.iter().count(), 1);
+        assert!(documents.next_document().unwrap().is_some());
+        assert!(documents.next_document().unwrap().is_none());
     }
 
     #[test]
     fn test_nested() {
-        let mut docs = documents!([{
+        let docs_reader = documents!([{
             "hello": {
                 "toto": ["hello"]
             }
         }]);
 
-        let (_index, doc) = docs.next_document_with_index().unwrap().unwrap();
-
+        let (mut cursor, _) = docs_reader.into_cursor_and_fields_index();
+        let doc = cursor.next_document().unwrap().unwrap();
         let nested: Value = serde_json::from_slice(doc.get(0).unwrap()).unwrap();
         assert_eq!(nested, json!({ "toto": ["hello"] }));
     }
 
     #[test]
-    fn out_of_order_fields() {
+    fn out_of_order_json_fields() {
         let _documents = documents!([
             {"id": 1,"b": 0},
             {"id": 2,"a": 0,"b": 0},
         ]);
     }
+
+    #[test]
+    fn out_of_order_csv_fields() {
+        let csv1_content = "id:number,b\n1,0";
+        let csv1 = csv::Reader::from_reader(Cursor::new(csv1_content));
+
+        let csv2_content = "id:number,a,b\n2,0,0";
+        let csv2 = csv::Reader::from_reader(Cursor::new(csv2_content));
+
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_csv(csv1).unwrap();
+        builder.append_csv(csv2).unwrap();
+        let vector = builder.into_inner().unwrap();
+
+        DocumentsBatchReader::from_reader(Cursor::new(vector)).unwrap();
+    }
 }

milli/src/documents/reader.rs

@@ -1,11 +1,9 @@
-use std::io;
-use std::io::{BufReader, Read};
-use std::mem::size_of;
+use std::convert::TryInto;
+use std::{error, fmt, io};
 
-use byteorder::{BigEndian, ReadBytesExt};
 use obkv::KvReader;
 
-use super::{DocumentsBatchIndex, DocumentsMetadata, Error};
+use super::{DocumentsBatchIndex, Error, DOCUMENTS_BATCH_INDEX_KEY};
 use crate::FieldId;
 
 /// The `DocumentsBatchReader` provides a way to iterate over documents that have been created with
@@ -13,63 +11,106 @@ use crate::FieldId;
 ///
 /// The documents are returned in the form of `obkv::Reader` where each field is identified with a
 /// `FieldId`. The mapping between the field ids and the field names is done thanks to the index.
-pub struct DocumentBatchReader<R> {
-    reader: BufReader<R>,
-    metadata: DocumentsMetadata,
-    buffer: Vec<u8>,
-    seen_documents: usize,
+pub struct DocumentsBatchReader<R> {
+    cursor: grenad::ReaderCursor<R>,
+    fields_index: DocumentsBatchIndex,
+}
+
+impl<R: io::Read + io::Seek> DocumentsBatchReader<R> {
+    pub fn new(cursor: DocumentsBatchCursor<R>, fields_index: DocumentsBatchIndex) -> Self {
+        Self { cursor: cursor.cursor, fields_index }
     }
 
-impl<R: io::Read + io::Seek> DocumentBatchReader<R> {
     /// Construct a `DocumentsReader` from a reader.
     ///
-    /// It first retrieves the index, then moves to the first document. Subsequent calls to
-    /// `next_document` advance the document reader until all the documents have been read.
-    pub fn from_reader(mut reader: R) -> Result<Self, Error> {
-        let mut buffer = Vec::new();
+    /// It first retrieves the index, then moves to the first document. Use the `into_cursor`
+    /// method to iterator over the documents, from the first to the last.
+    pub fn from_reader(reader: R) -> Result<Self, Error> {
+        let reader = grenad::Reader::new(reader)?;
+        let mut cursor = reader.into_cursor()?;
 
-        let meta_offset = reader.read_u64::<BigEndian>()?;
-        reader.seek(io::SeekFrom::Start(meta_offset))?;
-        reader.read_to_end(&mut buffer)?;
-        let metadata: DocumentsMetadata = bincode::deserialize(&buffer)?;
+        let fields_index = match cursor.move_on_key_equal_to(DOCUMENTS_BATCH_INDEX_KEY)? {
+            Some((_, value)) => serde_json::from_slice(value).map_err(Error::Serialize)?,
+            None => return Err(Error::InvalidDocumentFormat),
+        };
 
-        reader.seek(io::SeekFrom::Start(size_of::<u64>() as u64))?;
-        buffer.clear();
-
-        let reader = BufReader::new(reader);
-
-        Ok(Self { reader, metadata, buffer, seen_documents: 0 })
+        Ok(DocumentsBatchReader { cursor, fields_index })
     }
 
-    /// Returns the next document in the reader, and wraps it in an `obkv::KvReader`, along with a
-    /// reference to the addition index.
-    pub fn next_document_with_index<'a>(
-        &'a mut self,
-    ) -> io::Result<Option<(&'a DocumentsBatchIndex, KvReader<'a, FieldId>)>> {
-        if self.seen_documents < self.metadata.count {
-            let doc_len = self.reader.read_u32::<BigEndian>()?;
-            self.buffer.resize(doc_len as usize, 0);
-            self.reader.read_exact(&mut self.buffer)?;
-            self.seen_documents += 1;
-
-            let reader = KvReader::new(&self.buffer);
-            Ok(Some((&self.metadata.index, reader)))
-        } else {
-            Ok(None)
-        }
-    }
-
-    /// Return the fields index for the documents batch.
-    pub fn index(&self) -> &DocumentsBatchIndex {
-        &self.metadata.index
-    }
-
-    /// Returns the number of documents in the reader.
-    pub fn len(&self) -> usize {
-        self.metadata.count
+    pub fn documents_count(&self) -> u32 {
+        self.cursor.len().saturating_sub(1).try_into().expect("Invalid number of documents")
     }
 
     pub fn is_empty(&self) -> bool {
-        self.len() == 0
+        self.cursor.len().saturating_sub(1) == 0
+    }
+
+    pub fn documents_batch_index(&self) -> &DocumentsBatchIndex {
+        &self.fields_index
+    }
+
+    /// This method returns a forward cursor over the documents.
+    pub fn into_cursor_and_fields_index(self) -> (DocumentsBatchCursor<R>, DocumentsBatchIndex) {
+        let DocumentsBatchReader { cursor, fields_index } = self;
+        let mut cursor = DocumentsBatchCursor { cursor };
+        cursor.reset();
+        (cursor, fields_index)
+    }
+}
+
+/// A forward cursor over the documents in a `DocumentsBatchReader`.
+pub struct DocumentsBatchCursor<R> {
+    cursor: grenad::ReaderCursor<R>,
+}
+
+impl<R> DocumentsBatchCursor<R> {
+    /// Resets the cursor to be able to read from the start again.
+    pub fn reset(&mut self) {
+        self.cursor.reset();
+    }
+}
+
+impl<R: io::Read + io::Seek> DocumentsBatchCursor<R> {
+    /// Returns the next document, starting from the first one. Subsequent calls to
+    /// `next_document` advance the document reader until all the documents have been read.
+    pub fn next_document(
+        &mut self,
+    ) -> Result<Option<KvReader<FieldId>>, DocumentsBatchCursorError> {
+        match self.cursor.move_on_next()? {
+            Some((key, value)) if key != DOCUMENTS_BATCH_INDEX_KEY => {
+                Ok(Some(KvReader::new(value)))
+            }
+            _otherwise => Ok(None),
+        }
+    }
+}
+
+/// The possible error thrown by the `DocumentsBatchCursor` when iterating on the documents.
+#[derive(Debug)]
+pub enum DocumentsBatchCursorError {
+    Grenad(grenad::Error),
+    SerdeJson(serde_json::Error),
+}
+
+impl From<grenad::Error> for DocumentsBatchCursorError {
+    fn from(error: grenad::Error) -> DocumentsBatchCursorError {
+        DocumentsBatchCursorError::Grenad(error)
+    }
+}
+
+impl From<serde_json::Error> for DocumentsBatchCursorError {
+    fn from(error: serde_json::Error) -> DocumentsBatchCursorError {
+        DocumentsBatchCursorError::SerdeJson(error)
+    }
+}
+
+impl error::Error for DocumentsBatchCursorError {}
+
+impl fmt::Display for DocumentsBatchCursorError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            DocumentsBatchCursorError::Grenad(e) => e.fmt(f),
+            DocumentsBatchCursorError::SerdeJson(e) => e.fmt(f),
+        }
     }
 }

milli/src/documents/serde_impl.rs

@@ -1,14 +1,11 @@
-use std::collections::BTreeMap;
 use std::fmt;
-use std::io::{Cursor, Write};
+use std::io::Write;
 
-use byteorder::WriteBytesExt;
 use serde::de::{DeserializeSeed, MapAccess, SeqAccess, Visitor};
-use serde::Deserialize;
-use serde_json::Value;
 
-use super::{ByteCounter, DocumentsBatchIndex, Error};
-use crate::FieldId;
+use super::Error;
+use crate::documents::DocumentsBatchBuilder;
+use crate::Object;
 
 macro_rules! tri {
     ($e:expr) => {
@@ -19,54 +16,15 @@ macro_rules! tri {
     };
 }
 
-struct FieldIdResolver<'a>(&'a mut DocumentsBatchIndex);
-
-impl<'a, 'de> DeserializeSeed<'de> for FieldIdResolver<'a> {
-    type Value = FieldId;
-
-    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
-    where
-        D: serde::Deserializer<'de>,
-    {
-        deserializer.deserialize_str(self)
-    }
-}
-
-impl<'a, 'de> Visitor<'de> for FieldIdResolver<'a> {
-    type Value = FieldId;
-
-    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
-    where
-        E: serde::de::Error,
-    {
-        Ok(self.0.insert(v))
-    }
-
-    fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "a string")
-    }
-}
-
-struct ValueDeserializer;
-
-impl<'de> DeserializeSeed<'de> for ValueDeserializer {
-    type Value = serde_json::Value;
-
-    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
-    where
-        D: serde::Deserializer<'de>,
-    {
-        serde_json::Value::deserialize(deserializer)
-    }
-}
-
 pub struct DocumentVisitor<'a, W> {
-    pub inner: &'a mut ByteCounter<W>,
-    pub index: &'a mut DocumentsBatchIndex,
-    pub obkv_buffer: &'a mut Vec<u8>,
-    pub value_buffer: &'a mut Vec<u8>,
-    pub values: &'a mut BTreeMap<FieldId, Value>,
-    pub count: &'a mut usize,
+    inner: &'a mut DocumentsBatchBuilder<W>,
+    object: Object,
+}
+
+impl<'a, W> DocumentVisitor<'a, W> {
+    pub fn new(inner: &'a mut DocumentsBatchBuilder<W>) -> Self {
+        DocumentVisitor { inner, object: Object::new() }
+    }
 }
 
 impl<'a, 'de, W: Write> Visitor<'de> for &mut DocumentVisitor<'a, W> {
@@ -88,28 +46,12 @@ impl<'a, 'de, W: Write> Visitor<'de> for &mut DocumentVisitor<'a, W> {
     where
         A: MapAccess<'de>,
     {
-        while let Some((key, value)) =
-            map.next_entry_seed(FieldIdResolver(&mut *self.index), ValueDeserializer)?
-        {
-            self.values.insert(key, value);
+        self.object.clear();
+        while let Some((key, value)) = map.next_entry()? {
+            self.object.insert(key, value);
         }
 
-        self.obkv_buffer.clear();
-        let mut obkv = obkv::KvWriter::new(Cursor::new(&mut *self.obkv_buffer));
-        for (key, value) in self.values.iter() {
-            self.value_buffer.clear();
-            // This is guaranteed to work
-            tri!(serde_json::to_writer(Cursor::new(&mut *self.value_buffer), value));
-            tri!(obkv.insert(*key, &self.value_buffer));
-        }
-
-        let reader = tri!(obkv.into_inner()).into_inner();
-
-        tri!(self.inner.write_u32::<byteorder::BigEndian>(reader.len() as u32));
-        tri!(self.inner.write_all(reader));
-
-        *self.count += 1;
-        self.values.clear();
+        tri!(self.inner.append_json_object(&self.object));
 
         Ok(Ok(()))
     }

milli/src/error.rs

@@ -4,12 +4,11 @@ use std::{io, str};
 
 use heed::{Error as HeedError, MdbError};
 use rayon::ThreadPoolBuildError;
-use serde_json::{Map, Value};
+use serde_json::Value;
 use thiserror::Error;
 
-use crate::{CriterionError, DocumentId, FieldId, SortError};
-
-pub type Object = Map<String, Value>;
+use crate::documents::{self, DocumentsBatchCursorError};
+use crate::{CriterionError, DocumentId, FieldId, Object, SortError};
 
 pub fn is_reserved_keyword(keyword: &str) -> bool {
     ["_geo", "_geoDistance", "_geoPoint", "_geoRadius"].contains(&keyword)
@@ -37,6 +36,8 @@ pub enum InternalError {
     FieldIdMappingMissingEntry { key: FieldId },
     #[error(transparent)]
     Fst(#[from] fst::Error),
+    #[error(transparent)]
+    DocumentsError(#[from] documents::Error),
     #[error("Invalid compression type have been specified to grenad.")]
     GrenadInvalidCompressionType,
     #[error("Invalid grenad file with an invalid version format.")]
@@ -123,6 +124,8 @@ only composed of alphanumeric characters (a-z A-Z 0-9), hyphens (-) and undersco
     MaxDatabaseSizeReached,
     #[error("Document doesn't have a `{}` attribute: `{}`.", .primary_key, serde_json::to_string(.document).unwrap())]
     MissingDocumentId { primary_key: String, document: Object },
+    #[error("Document have too many matching `{}` attribute: `{}`.", .primary_key, serde_json::to_string(.document).unwrap())]
+    TooManyDocumentIds { primary_key: String, document: Object },
     #[error("The primary key inference process failed because the engine did not find any fields containing `id` substring in their name. If your document identifier does not contain any `id` substring, you can set the primary key of the index.")]
     MissingPrimaryKey,
     #[error("There is no more space left on the device. Consider increasing the size of the disk/partition.")]
@@ -141,13 +144,19 @@ only composed of alphanumeric characters (a-z A-Z 0-9), hyphens (-) and undersco
 
 #[derive(Error, Debug)]
 pub enum GeoError {
+    #[error("The `_geo` field in the document with the id: `{document_id}` is not an object. Was expecting an object with the `_geo.lat` and `_geo.lng` fields but instead got `{value}`.")]
+    NotAnObject { document_id: Value, value: Value },
+    #[error("Could not find latitude nor longitude in the document with the id: `{document_id}`. Was expecting `_geo.lat` and `_geo.lng` fields.")]
+    MissingLatitudeAndLongitude { document_id: Value },
     #[error("Could not find latitude in the document with the id: `{document_id}`. Was expecting a `_geo.lat` field.")]
     MissingLatitude { document_id: Value },
     #[error("Could not find longitude in the document with the id: `{document_id}`. Was expecting a `_geo.lng` field.")]
     MissingLongitude { document_id: Value },
-    #[error("Could not parse latitude in the document with the id: `{document_id}`. Was expecting a number but instead got `{value}`.")]
+    #[error("Could not parse latitude nor longitude in the document with the id: `{document_id}`. Was expecting finite numbers but instead got `{lat}` and `{lng}`.")]
+    BadLatitudeAndLongitude { document_id: Value, lat: Value, lng: Value },
+    #[error("Could not parse latitude in the document with the id: `{document_id}`. Was expecting a finite number but instead got `{value}`.")]
     BadLatitude { document_id: Value, value: Value },
-    #[error("Could not parse longitude in the document with the id: `{document_id}`. Was expecting a number but instead got `{value}`.")]
+    #[error("Could not parse longitude in the document with the id: `{document_id}`. Was expecting a finite number but instead got `{value}`.")]
     BadLongitude { document_id: Value, value: Value },
 }
 
@@ -178,6 +187,7 @@ macro_rules! error_from_sub_error {
 error_from_sub_error! {
     FieldIdMapMissingEntry => InternalError,
     fst::Error => InternalError,
+    documents::Error => InternalError,
     str::Utf8Error => InternalError,
     ThreadPoolBuildError => InternalError,
     SerializationError => InternalError,
@@ -203,6 +213,15 @@ where
     }
 }
 
+impl From<DocumentsBatchCursorError> for Error {
+    fn from(error: DocumentsBatchCursorError) -> Error {
+        match error {
+            DocumentsBatchCursorError::Grenad(e) => Error::from(e),
+            DocumentsBatchCursorError::SerdeJson(e) => Error::from(InternalError::from(e)),
+        }
+    }
+}
+
 impl From<Infallible> for Error {
     fn from(_error: Infallible) -> Error {
         unreachable!()

milli/src/index.rs

@@ -1212,10 +1212,11 @@ pub(crate) mod tests {
 
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1234,7 +1235,7 @@ pub(crate) mod tests {
         // we add all the documents a second time. we are supposed to get the same
         // field_distribution in the end
         let mut wtxn = index.write_txn().unwrap();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
         let content = documents!([
@@ -1242,7 +1243,8 @@ pub(crate) mod tests {
             { "id": 2, "name": "bob", "age": 20 },
             { "id": 2, "name": "bob", "age": 20 },
         ]);
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1265,10 +1267,11 @@ pub(crate) mod tests {
         ]);
 
         let mut wtxn = index.write_txn().unwrap();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1333,10 +1336,11 @@ pub(crate) mod tests {
 
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1390,10 +1394,11 @@ pub(crate) mod tests {
         ]);
 
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 

milli/src/lib.rs

@@ -20,7 +20,7 @@ use std::hash::BuildHasherDefault;
 pub use filter_parser::{Condition, FilterCondition};
 use fxhash::{FxHasher32, FxHasher64};
 pub use grenad::CompressionType;
-use serde_json::{Map, Value};
+use serde_json::Value;
 pub use {charabia as tokenizer, heed};
 
 pub use self::asc_desc::{AscDesc, AscDescError, Member, SortError};
@@ -43,20 +43,21 @@ pub use self::search::{
 
 pub type Result<T> = std::result::Result<T, error::Error>;
 
+pub type Attribute = u32;
+pub type BEU32 = heed::zerocopy::U32<heed::byteorder::BE>;
+pub type BEU64 = heed::zerocopy::U64<heed::byteorder::BE>;
+pub type DocumentId = u32;
 pub type FastMap4<K, V> = HashMap<K, V, BuildHasherDefault<FxHasher32>>;
 pub type FastMap8<K, V> = HashMap<K, V, BuildHasherDefault<FxHasher64>>;
+pub type FieldDistribution = BTreeMap<String, u64>;
+pub type FieldId = u16;
+pub type Object = serde_json::Map<String, serde_json::Value>;
+pub type Position = u32;
+pub type RelativePosition = u16;
 pub type SmallString32 = smallstr::SmallString<[u8; 32]>;
 pub type SmallVec16<T> = smallvec::SmallVec<[T; 16]>;
 pub type SmallVec32<T> = smallvec::SmallVec<[T; 32]>;
 pub type SmallVec8<T> = smallvec::SmallVec<[T; 8]>;
-pub type BEU32 = heed::zerocopy::U32<heed::byteorder::BE>;
-pub type BEU64 = heed::zerocopy::U64<heed::byteorder::BE>;
-pub type Attribute = u32;
-pub type DocumentId = u32;
-pub type FieldId = u16;
-pub type Position = u32;
-pub type RelativePosition = u16;
-pub type FieldDistribution = BTreeMap<String, u64>;
 
 /// A GeoPoint is a point in cartesian plan, called xyz_point in the code. Its metadata
 /// is a tuple composed of 1. the DocumentId of the associated document and 2. the original point
@@ -82,7 +83,7 @@ pub fn obkv_to_json(
     displayed_fields: &[FieldId],
     fields_ids_map: &FieldsIdsMap,
     obkv: obkv::KvReaderU16,
-) -> Result<Map<String, Value>> {
+) -> Result<Object> {
     displayed_fields
         .iter()
         .copied()

milli/src/search/distinct/mod.rs

@@ -35,7 +35,7 @@ mod test {
     use roaring::RoaringBitmap;
     use serde_json::{json, Value};
 
-    use crate::documents::{DocumentBatchBuilder, DocumentBatchReader};
+    use crate::documents::{DocumentsBatchBuilder, DocumentsBatchReader};
     use crate::index::tests::TempIndex;
     use crate::index::Index;
     use crate::update::{
@@ -43,14 +43,11 @@ mod test {
     };
     use crate::{DocumentId, FieldId, BEU32};
 
-    static JSON: Lazy<Vec<u8>> = Lazy::new(generate_documents);
-
-    fn generate_documents() -> Vec<u8> {
+    static JSON: Lazy<Vec<u8>> = Lazy::new(|| {
         let mut rng = rand::thread_rng();
         let num_docs = rng.gen_range(10..30);
 
-        let mut cursor = Cursor::new(Vec::new());
-        let mut builder = DocumentBatchBuilder::new(&mut cursor).unwrap();
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
         let txts = ["Toto", "Titi", "Tata"];
         let cats = (1..10).map(|i| i.to_string()).collect::<Vec<_>>();
         let cat_ints = (1..10).collect::<Vec<_>>();
@@ -63,7 +60,7 @@ mod test {
             let mut sample_ints = cat_ints.clone();
             sample_ints.shuffle(&mut rng);
 
-            let doc = json!({
+            let json = json!({
                 "id": i,
                 "txt": txt,
                 "cat-int": rng.gen_range(0..3),
@@ -71,13 +68,16 @@ mod test {
                 "cat-ints": sample_ints[..(rng.gen_range(0..3))],
             });
 
-            let doc = Cursor::new(serde_json::to_vec(&doc).unwrap());
-            builder.extend_from_json(doc).unwrap();
+            let object = match json {
+                Value::Object(object) => object,
+                _ => panic!(),
+            };
+
+            builder.append_json_object(&object).unwrap();
         }
 
-        builder.finish().unwrap();
-        cursor.into_inner()
-    }
+        builder.into_inner().unwrap()
+    });
 
     /// Returns a temporary index populated with random test documents, the FieldId for the
     /// distinct attribute, and the RoaringBitmap with the document ids.
@@ -97,20 +97,22 @@ mod test {
             update_method: IndexDocumentsMethod::ReplaceDocuments,
             ..Default::default()
         };
-        let mut addition =
+        let addition =
             IndexDocuments::new(&mut txn, &index, &config, indexing_config, |_| ()).unwrap();
 
         let reader =
-            crate::documents::DocumentBatchReader::from_reader(Cursor::new(&*JSON)).unwrap();
+            crate::documents::DocumentsBatchReader::from_reader(Cursor::new(JSON.as_slice()))
+                .unwrap();
 
-        addition.add_documents(reader).unwrap();
+        let (addition, user_error) = addition.add_documents(reader).unwrap();
+        user_error.unwrap();
         addition.execute().unwrap();
 
         let fields_map = index.fields_ids_map(&txn).unwrap();
         let fid = fields_map.id(&distinct).unwrap();
 
-        let documents = DocumentBatchReader::from_reader(Cursor::new(&*JSON)).unwrap();
-        let map = (0..documents.len() as u32).collect();
+        let documents = DocumentsBatchReader::from_reader(Cursor::new(JSON.as_slice())).unwrap();
+        let map = (0..documents.documents_count() as u32).collect();
 
         txn.commit().unwrap();
 

milli/src/search/facet/filter.rs

@@ -648,10 +648,11 @@ mod tests {
 
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         wtxn.commit().unwrap();

milli/src/update/clear_documents.rs

@@ -100,9 +100,10 @@ mod tests {
         ]);
         let indexing_config = IndexDocumentsConfig::default();
         let config = IndexerConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         // Clear all documents from the database.

milli/src/update/delete_documents.rs

@@ -657,13 +657,13 @@ mod tests {
     fn insert_documents<'t, R: std::io::Read + std::io::Seek>(
         wtxn: &mut RwTxn<'t, '_>,
         index: &'t Index,
-        documents: crate::documents::DocumentBatchReader<R>,
+        documents: crate::documents::DocumentsBatchReader<R>,
     ) {
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
-            IndexDocuments::new(wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        builder.add_documents(documents).unwrap();
+        let builder = IndexDocuments::new(wtxn, &index, &config, indexing_config, |_| ()).unwrap();
+        let (builder, user_error) = builder.add_documents(documents).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
     }
 
@@ -701,9 +701,10 @@ mod tests {
         ]);
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         // delete those documents, ids are synchronous therefore 0, 1, and 2.
@@ -736,9 +737,10 @@ mod tests {
 
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         // Delete not all of the documents but some of them.

milli/src/update/index_documents/enrich.rs

@@ -0,0 +1,365 @@
+use std::io::{Read, Seek};
+use std::result::Result as StdResult;
+use std::{fmt, iter};
+
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+
+use crate::documents::{DocumentsBatchIndex, DocumentsBatchReader, EnrichedDocumentsBatchReader};
+use crate::error::{GeoError, InternalError, UserError};
+use crate::update::index_documents::{obkv_to_object, writer_into_reader};
+use crate::{FieldId, Index, Object, Result};
+
+/// The symbol used to define levels in a nested primary key.
+const PRIMARY_KEY_SPLIT_SYMBOL: char = '.';
+
+/// The default primary that is used when not specified.
+const DEFAULT_PRIMARY_KEY: &str = "id";
+
+/// This function validates and enrich the documents by checking that:
+///  - we can infer a primary key,
+///  - all the documents id exist and are extracted,
+///  - the validity of them but also,
+///  - the validity of the `_geo` field depending on the settings.
+pub fn enrich_documents_batch<R: Read + Seek>(
+    rtxn: &heed::RoTxn,
+    index: &Index,
+    autogenerate_docids: bool,
+    reader: DocumentsBatchReader<R>,
+) -> Result<StdResult<EnrichedDocumentsBatchReader<R>, UserError>> {
+    let (mut cursor, mut documents_batch_index) = reader.into_cursor_and_fields_index();
+
+    let mut external_ids = tempfile::tempfile().map(grenad::Writer::new)?;
+    let mut uuid_buffer = [0; uuid::fmt::Hyphenated::LENGTH];
+
+    // The primary key *field id* that has already been set for this index or the one
+    // we will guess by searching for the first key that contains "id" as a substring.
+    let primary_key = match index.primary_key(rtxn)? {
+        Some(primary_key) if primary_key.contains(PRIMARY_KEY_SPLIT_SYMBOL) => {
+            PrimaryKey::nested(primary_key)
+        }
+        Some(primary_key) => match documents_batch_index.id(primary_key) {
+            Some(id) => PrimaryKey::flat(primary_key, id),
+            None if autogenerate_docids => {
+                PrimaryKey::flat(primary_key, documents_batch_index.insert(primary_key))
+            }
+            None => {
+                return match cursor.next_document()? {
+                    Some(first_document) => Ok(Err(UserError::MissingDocumentId {
+                        primary_key: primary_key.to_string(),
+                        document: obkv_to_object(&first_document, &documents_batch_index)?,
+                    })),
+                    None => Ok(Err(UserError::MissingPrimaryKey)),
+                };
+            }
+        },
+        None => {
+            let guessed = documents_batch_index
+                .iter()
+                .filter(|(_, name)| name.to_lowercase().contains(DEFAULT_PRIMARY_KEY))
+                .min_by_key(|(fid, _)| *fid);
+            match guessed {
+                Some((id, name)) => PrimaryKey::flat(name.as_str(), *id),
+                None if autogenerate_docids => PrimaryKey::flat(
+                    DEFAULT_PRIMARY_KEY,
+                    documents_batch_index.insert(DEFAULT_PRIMARY_KEY),
+                ),
+                None => return Ok(Err(UserError::MissingPrimaryKey)),
+            }
+        }
+    };
+
+    // If the settings specifies that a _geo field must be used therefore we must check the
+    // validity of it in all the documents of this batch and this is when we return `Some`.
+    let geo_field_id = match documents_batch_index.id("_geo") {
+        Some(geo_field_id) if index.sortable_fields(rtxn)?.contains("_geo") => Some(geo_field_id),
+        _otherwise => None,
+    };
+
+    let mut count = 0;
+    while let Some(document) = cursor.next_document()? {
+        let document_id = match fetch_or_generate_document_id(
+            &document,
+            &documents_batch_index,
+            primary_key,
+            autogenerate_docids,
+            &mut uuid_buffer,
+            count,
+        )? {
+            Ok(document_id) => document_id,
+            Err(user_error) => return Ok(Err(user_error)),
+        };
+
+        if let Some(geo_value) = geo_field_id.and_then(|fid| document.get(fid)) {
+            if let Err(user_error) = validate_geo_from_json(&document_id, geo_value)? {
+                return Ok(Err(UserError::from(user_error)));
+            }
+        }
+
+        let document_id = serde_json::to_vec(&document_id).map_err(InternalError::SerdeJson)?;
+        external_ids.insert(count.to_be_bytes(), document_id)?;
+
+        count += 1;
+    }
+
+    let external_ids = writer_into_reader(external_ids)?;
+    let primary_key_name = primary_key.name().to_string();
+    let reader = EnrichedDocumentsBatchReader::new(
+        DocumentsBatchReader::new(cursor, documents_batch_index),
+        primary_key_name,
+        external_ids,
+    )?;
+
+    Ok(Ok(reader))
+}
+
+/// Retrieve the document id after validating it, returning a `UserError`
+/// if the id is invalid or can't be guessed.
+fn fetch_or_generate_document_id(
+    document: &obkv::KvReader<FieldId>,
+    documents_batch_index: &DocumentsBatchIndex,
+    primary_key: PrimaryKey,
+    autogenerate_docids: bool,
+    uuid_buffer: &mut [u8; uuid::fmt::Hyphenated::LENGTH],
+    count: u32,
+) -> Result<StdResult<DocumentId, UserError>> {
+    match primary_key {
+        PrimaryKey::Flat { name: primary_key, field_id: primary_key_id } => {
+            match document.get(primary_key_id) {
+                Some(document_id_bytes) => {
+                    let document_id = serde_json::from_slice(document_id_bytes)
+                        .map_err(InternalError::SerdeJson)?;
+                    match validate_document_id_value(document_id)? {
+                        Ok(document_id) => Ok(Ok(DocumentId::retrieved(document_id))),
+                        Err(user_error) => Ok(Err(user_error)),
+                    }
+                }
+                None if autogenerate_docids => {
+                    let uuid = uuid::Uuid::new_v4().as_hyphenated().encode_lower(uuid_buffer);
+                    Ok(Ok(DocumentId::generated(uuid.to_string(), count)))
+                }
+                None => Ok(Err(UserError::MissingDocumentId {
+                    primary_key: primary_key.to_string(),
+                    document: obkv_to_object(&document, &documents_batch_index)?,
+                })),
+            }
+        }
+        nested @ PrimaryKey::Nested { .. } => {
+            let mut matching_documents_ids = Vec::new();
+            for (first_level_name, right) in nested.possible_level_names() {
+                if let Some(field_id) = documents_batch_index.id(first_level_name) {
+                    if let Some(value_bytes) = document.get(field_id) {
+                        let object = serde_json::from_slice(value_bytes)
+                            .map_err(InternalError::SerdeJson)?;
+                        fetch_matching_values(object, right, &mut matching_documents_ids);
+
+                        if matching_documents_ids.len() >= 2 {
+                            return Ok(Err(UserError::TooManyDocumentIds {
+                                primary_key: nested.name().to_string(),
+                                document: obkv_to_object(&document, &documents_batch_index)?,
+                            }));
+                        }
+                    }
+                }
+            }
+
+            match matching_documents_ids.pop() {
+                Some(document_id) => match validate_document_id_value(document_id)? {
+                    Ok(document_id) => Ok(Ok(DocumentId::retrieved(document_id))),
+                    Err(user_error) => Ok(Err(user_error)),
+                },
+                None => Ok(Err(UserError::MissingDocumentId {
+                    primary_key: nested.name().to_string(),
+                    document: obkv_to_object(&document, &documents_batch_index)?,
+                })),
+            }
+        }
+    }
+}
+
+/// A type that represent the type of primary key that has been set
+/// for this index, a classic flat one or a nested one.
+#[derive(Debug, Clone, Copy)]
+enum PrimaryKey<'a> {
+    Flat { name: &'a str, field_id: FieldId },
+    Nested { name: &'a str },
+}
+
+impl PrimaryKey<'_> {
+    fn flat(name: &str, field_id: FieldId) -> PrimaryKey {
+        PrimaryKey::Flat { name, field_id }
+    }
+
+    fn nested(name: &str) -> PrimaryKey {
+        PrimaryKey::Nested { name }
+    }
+
+    fn name(&self) -> &str {
+        match self {
+            PrimaryKey::Flat { name, .. } => name,
+            PrimaryKey::Nested { name } => name,
+        }
+    }
+
+    /// Returns an `Iterator` that gives all the possible fields names the primary key
+    /// can have depending of the first level name and deepnes of the objects.
+    fn possible_level_names(&self) -> impl Iterator<Item = (&str, &str)> + '_ {
+        let name = self.name();
+        name.match_indices(PRIMARY_KEY_SPLIT_SYMBOL)
+            .map(move |(i, _)| (&name[..i], &name[i + PRIMARY_KEY_SPLIT_SYMBOL.len_utf8()..]))
+            .chain(iter::once((name, "")))
+    }
+}
+
+/// A type that represents a document id that has been retrieved from a document or auto-generated.
+///
+/// In case the document id has been auto-generated, the document nth is kept to help
+/// users debug if there is an issue with the document itself.
+#[derive(Serialize, Deserialize, Clone)]
+pub enum DocumentId {
+    Retrieved { value: String },
+    Generated { value: String, document_nth: u32 },
+}
+
+impl DocumentId {
+    fn retrieved(value: String) -> DocumentId {
+        DocumentId::Retrieved { value }
+    }
+
+    fn generated(value: String, document_nth: u32) -> DocumentId {
+        DocumentId::Generated { value, document_nth }
+    }
+
+    fn debug(&self) -> String {
+        format!("{:?}", self)
+    }
+
+    pub fn is_generated(&self) -> bool {
+        matches!(self, DocumentId::Generated { .. })
+    }
+
+    pub fn value(&self) -> &str {
+        match self {
+            DocumentId::Retrieved { value } => value,
+            DocumentId::Generated { value, .. } => value,
+        }
+    }
+}
+
+impl fmt::Debug for DocumentId {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            DocumentId::Retrieved { value } => write!(f, "{:?}", value),
+            DocumentId::Generated { value, document_nth } => {
+                write!(f, "{{{:?}}} of the {}nth document", value, document_nth)
+            }
+        }
+    }
+}
+
+fn starts_with(selector: &str, key: &str) -> bool {
+    selector.strip_prefix(key).map_or(false, |tail| {
+        tail.chars().next().map(|c| c == PRIMARY_KEY_SPLIT_SYMBOL).unwrap_or(true)
+    })
+}
+
+pub fn fetch_matching_values(value: Value, selector: &str, output: &mut Vec<Value>) {
+    match value {
+        Value::Object(object) => fetch_matching_values_in_object(object, selector, "", output),
+        otherwise => output.push(otherwise),
+    }
+}
+
+pub fn fetch_matching_values_in_object(
+    object: Object,
+    selector: &str,
+    base_key: &str,
+    output: &mut Vec<Value>,
+) {
+    for (key, value) in object {
+        let base_key = if base_key.is_empty() {
+            key.to_string()
+        } else {
+            format!("{}{}{}", base_key, PRIMARY_KEY_SPLIT_SYMBOL, key)
+        };
+
+        if starts_with(selector, &base_key) {
+            match value {
+                Value::Object(object) => {
+                    fetch_matching_values_in_object(object, selector, &base_key, output)
+                }
+                value => output.push(value),
+            }
+        }
+    }
+}
+
+/// Returns a trimmed version of the document id or `None` if it is invalid.
+pub fn validate_document_id(document_id: &str) -> Option<&str> {
+    let document_id = document_id.trim();
+    if !document_id.is_empty()
+        && document_id.chars().all(|c| matches!(c, 'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_'))
+    {
+        Some(document_id)
+    } else {
+        None
+    }
+}
+
+/// Parses a Json encoded document id and validate it, returning a user error when it is one.
+pub fn validate_document_id_value(document_id: Value) -> Result<StdResult<String, UserError>> {
+    match document_id {
+        Value::String(string) => match validate_document_id(&string) {
+            Some(s) if s.len() == string.len() => Ok(Ok(string)),
+            Some(s) => Ok(Ok(s.to_string())),
+            None => Ok(Err(UserError::InvalidDocumentId { document_id: Value::String(string) })),
+        },
+        Value::Number(number) if number.is_i64() => Ok(Ok(number.to_string())),
+        content => Ok(Err(UserError::InvalidDocumentId { document_id: content.clone() })),
+    }
+}
+
+/// Try to extract an `f64` from a JSON `Value` and return the `Value`
+/// in the `Err` variant if it failed.
+pub fn extract_finite_float_from_value(value: Value) -> StdResult<f64, Value> {
+    let number = match value {
+        Value::Number(ref n) => match n.as_f64() {
+            Some(number) => number,
+            None => return Err(value),
+        },
+        Value::String(ref s) => match s.parse::<f64>() {
+            Ok(number) => number,
+            Err(_) => return Err(value),
+        },
+        value => return Err(value),
+    };
+
+    if number.is_finite() {
+        Ok(number)
+    } else {
+        Err(value)
+    }
+}
+
+pub fn validate_geo_from_json(id: &DocumentId, bytes: &[u8]) -> Result<StdResult<(), GeoError>> {
+    use GeoError::*;
+    let debug_id = || Value::from(id.debug());
+    match serde_json::from_slice(bytes).map_err(InternalError::SerdeJson)? {
+        Value::Object(mut object) => match (object.remove("lat"), object.remove("lng")) {
+            (Some(lat), Some(lng)) => {
+                match (extract_finite_float_from_value(lat), extract_finite_float_from_value(lng)) {
+                    (Ok(_), Ok(_)) => Ok(Ok(())),
+                    (Err(value), Ok(_)) => Ok(Err(BadLatitude { document_id: debug_id(), value })),
+                    (Ok(_), Err(value)) => Ok(Err(BadLongitude { document_id: debug_id(), value })),
+                    (Err(lat), Err(lng)) => {
+                        Ok(Err(BadLatitudeAndLongitude { document_id: debug_id(), lat, lng }))
+                    }
+                }
+            }
+            (None, Some(_)) => Ok(Err(MissingLatitude { document_id: debug_id() })),
+            (Some(_), None) => Ok(Err(MissingLongitude { document_id: debug_id() })),
+            (None, None) => Ok(Err(MissingLatitudeAndLongitude { document_id: debug_id() })),
+        },
+        value => Ok(Err(NotAnObject { document_id: debug_id(), value })),
+    }
+}

milli/src/update/index_documents/extract/extract_geo_points.rs

@@ -1,12 +1,12 @@
 use std::fs::File;
 use std::io;
-use std::result::Result as StdResult;
 
 use concat_arrays::concat_arrays;
 use serde_json::Value;
 
 use super::helpers::{create_writer, writer_into_reader, GrenadParameters};
 use crate::error::GeoError;
+use crate::update::index_documents::extract_finite_float_from_value;
 use crate::{FieldId, InternalError, Result};
 
 /// Extracts the geographical coordinates contained in each document under the `_geo` field.
@@ -29,9 +29,9 @@ pub fn extract_geo_points<R: io::Read + io::Seek>(
         let obkv = obkv::KvReader::new(value);
         // since we only needs the primary key when we throw an error we create this getter to
         // lazily get it when needed
-        let primary_key = || -> Value {
-            let primary_key = obkv.get(primary_key_id).unwrap();
-            serde_json::from_slice(primary_key).unwrap()
+        let document_id = || -> Value {
+            let document_id = obkv.get(primary_key_id).unwrap();
+            serde_json::from_slice(document_id).unwrap()
         };
 
         // first we get the two fields
@@ -40,32 +40,24 @@ pub fn extract_geo_points<R: io::Read + io::Seek>(
 
         if let Some((lat, lng)) = lat.zip(lng) {
             // then we extract the values
-            let lat = extract_float_from_value(
+            let lat = extract_finite_float_from_value(
                 serde_json::from_slice(lat).map_err(InternalError::SerdeJson)?,
             )
-            .map_err(|lat| GeoError::BadLatitude { document_id: primary_key(), value: lat })?;
+            .map_err(|lat| GeoError::BadLatitude { document_id: document_id(), value: lat })?;
 
-            let lng = extract_float_from_value(
+            let lng = extract_finite_float_from_value(
                 serde_json::from_slice(lng).map_err(InternalError::SerdeJson)?,
             )
-            .map_err(|lng| GeoError::BadLongitude { document_id: primary_key(), value: lng })?;
+            .map_err(|lng| GeoError::BadLongitude { document_id: document_id(), value: lng })?;
 
             let bytes: [u8; 16] = concat_arrays![lat.to_ne_bytes(), lng.to_ne_bytes()];
             writer.insert(docid_bytes, bytes)?;
         } else if lat.is_none() && lng.is_some() {
-            return Err(GeoError::MissingLatitude { document_id: primary_key() })?;
+            return Err(GeoError::MissingLatitude { document_id: document_id() })?;
         } else if lat.is_some() && lng.is_none() {
-            return Err(GeoError::MissingLongitude { document_id: primary_key() })?;
+            return Err(GeoError::MissingLongitude { document_id: document_id() })?;
         }
     }
 
     Ok(writer_into_reader(writer)?)
 }
-
-fn extract_float_from_value(value: Value) -> StdResult<f64, Value> {
-    match value {
-        Value::Number(ref n) => n.as_f64().ok_or(value),
-        Value::String(ref s) => s.parse::<f64>().map_err(|_| value),
-        value => Err(value),
-    }
-}

milli/src/update/index_documents/mod.rs

@@ -1,3 +1,4 @@
+mod enrich;
 mod extract;
 mod helpers;
 mod transform;
@@ -7,6 +8,7 @@ use std::collections::HashSet;
 use std::io::{Cursor, Read, Seek};
 use std::iter::FromIterator;
 use std::num::{NonZeroU32, NonZeroUsize};
+use std::result::Result as StdResult;
 
 use crossbeam_channel::{Receiver, Sender};
 use heed::types::Str;
@@ -17,6 +19,11 @@ use serde::{Deserialize, Serialize};
 use slice_group_by::GroupBy;
 use typed_chunk::{write_typed_chunk_into_index, TypedChunk};
 
+use self::enrich::enrich_documents_batch;
+pub use self::enrich::{
+    extract_finite_float_from_value, validate_document_id, validate_document_id_value,
+    validate_geo_from_json, DocumentId,
+};
 pub use self::helpers::{
     as_cloneable_grenad, create_sorter, create_writer, fst_stream_into_hashset,
     fst_stream_into_vec, merge_cbo_roaring_bitmaps, merge_roaring_bitmaps,
@@ -25,13 +32,14 @@ pub use self::helpers::{
 };
 use self::helpers::{grenad_obkv_into_chunks, GrenadParameters};
 pub use self::transform::{Transform, TransformOutput};
-use crate::documents::DocumentBatchReader;
+use crate::documents::{obkv_to_object, DocumentsBatchReader};
+use crate::error::UserError;
 pub use crate::update::index_documents::helpers::CursorClonableMmap;
 use crate::update::{
     self, Facets, IndexerConfig, UpdateIndexingStep, WordPrefixDocids,
     WordPrefixPairProximityDocids, WordPrefixPositionDocids, WordsPrefixesFst,
 };
-use crate::{Index, Result, RoaringBitmapCodec, UserError};
+use crate::{Index, Result, RoaringBitmapCodec};
 
 static MERGED_DATABASE_COUNT: usize = 7;
 static PREFIX_DATABASE_COUNT: usize = 5;
@@ -117,29 +125,42 @@ where
 
     /// Adds a batch of documents to the current builder.
     ///
-    /// Since the documents are progressively added to the writer, a failure will cause a stale
-    /// builder, and the builder must be discarded.
+    /// Since the documents are progressively added to the writer, a failure will cause only
+    /// return an error and not the `IndexDocuments` struct as it is invalid to use it afterward.
     ///
     /// Returns the number of documents added to the builder.
-    pub fn add_documents<R>(&mut self, reader: DocumentBatchReader<R>) -> Result<u64>
-    where
-        R: Read + Seek,
-    {
+    pub fn add_documents<R: Read + Seek>(
+        mut self,
+        reader: DocumentsBatchReader<R>,
+    ) -> Result<(Self, StdResult<u64, UserError>)> {
         // Early return when there is no document to add
         if reader.is_empty() {
-            return Ok(0);
+            return Ok((self, Ok(0)));
         }
 
+        // We check for user errors in this validator and if there is one, we can return
+        // the `IndexDocument` struct as it is valid to send more documents into it.
+        // However, if there is an internal error we throw it away!
+        let enriched_documents_reader = match enrich_documents_batch(
+            self.wtxn,
+            self.index,
+            self.config.autogenerate_docids,
+            reader,
+        )? {
+            Ok(reader) => reader,
+            Err(user_error) => return Ok((self, Err(user_error))),
+        };
+
         let indexed_documents = self
             .transform
             .as_mut()
             .expect("Invalid document addition state")
-            .read_documents(reader, self.wtxn, &self.progress)?
+            .read_documents(enriched_documents_reader, self.wtxn, &self.progress)?
             as u64;
 
         self.added_documents += indexed_documents;
 
-        Ok(indexed_documents)
+        Ok((self, Ok(indexed_documents)))
     }
 
     #[logging_timer::time("IndexDocuments::{}")]
@@ -590,9 +611,8 @@ mod tests {
     use maplit::hashset;
 
     use super::*;
-    use crate::documents::DocumentBatchBuilder;
+    use crate::documents::DocumentsBatchBuilder;
     use crate::update::DeleteDocuments;
-    use crate::HashMap;
 
     #[test]
     fn simple_document_replacement() {
@@ -611,10 +631,11 @@ mod tests {
 
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -627,10 +648,11 @@ mod tests {
         // Second we send 1 document with id 1, to erase the previous ones.
         let mut wtxn = index.write_txn().unwrap();
         let content = documents!([ { "id": 1, "name": "updated kevin" } ]);
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -647,9 +669,11 @@ mod tests {
             { "id": 2, "name": "updated kevina" },
             { "id": 3, "name": "updated benoit" }
         ]);
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
+        builder.execute().unwrap();
         wtxn.commit().unwrap();
 
         // Check that there is **always** 3 documents.
@@ -679,10 +703,11 @@ mod tests {
             update_method: IndexDocumentsMethod::UpdateDocuments,
             ..Default::default()
         };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -707,9 +732,10 @@ mod tests {
         // Second we send 1 document with id 1, to force it to be merged with the previous one.
         let mut wtxn = index.write_txn().unwrap();
         let content = documents!([ { "id": 1, "age": 25 } ]);
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -750,9 +776,10 @@ mod tests {
         ]);
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        assert!(builder.add_documents(content).is_err());
+        let (_builder, user_error) = builder.add_documents(content).unwrap();
+        assert!(user_error.is_err());
         wtxn.commit().unwrap();
 
         // Check that there is no document.
@@ -779,10 +806,11 @@ mod tests {
         let config = IndexerConfig::default();
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -799,9 +827,10 @@ mod tests {
         // Second we send 1 document with the generated uuid, to erase the previous ones.
         let mut wtxn = index.write_txn().unwrap();
         let content = documents!([ { "name": "updated kevin", "id": kevin_uuid } ]);
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -841,9 +870,10 @@ mod tests {
         ]);
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -858,9 +888,10 @@ mod tests {
         let content = documents!([ { "name": "new kevin" } ]);
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -883,9 +914,10 @@ mod tests {
         let content = documents!([]);
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -909,19 +941,21 @@ mod tests {
         let content = documents!([ { "id": "brume bleue", "name": "kevin" } ]);
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        assert!(builder.add_documents(content).is_err());
+        let (_builder, user_error) = builder.add_documents(content).unwrap();
+        assert!(user_error.is_err());
         wtxn.commit().unwrap();
 
         // First we send 1 document with a valid id.
         let mut wtxn = index.write_txn().unwrap();
         // There is a space in the document id.
         let content = documents!([ { "id": 32, "name": "kevin" } ]);
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -948,9 +982,10 @@ mod tests {
         ]);
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -994,9 +1029,10 @@ mod tests {
             update_method: IndexDocumentsMethod::ReplaceDocuments,
             ..Default::default()
         };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        builder.add_documents(documents).unwrap();
+        let (builder, user_error) = builder.add_documents(documents).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1005,7 +1041,7 @@ mod tests {
             update_method: IndexDocumentsMethod::UpdateDocuments,
             ..Default::default()
         };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
         let documents = documents!([
           {
@@ -1015,7 +1051,8 @@ mod tests {
           }
         ]);
 
-        builder.add_documents(documents).unwrap();
+        let (builder, user_error) = builder.add_documents(documents).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
     }
@@ -1042,9 +1079,10 @@ mod tests {
             update_method: IndexDocumentsMethod::ReplaceDocuments,
             ..Default::default()
         };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-        builder.add_documents(documents).unwrap();
+        let (builder, user_error) = builder.add_documents(documents).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1084,10 +1122,11 @@ mod tests {
           { "id": 2, "_geo": { "lng": "42" }, "_geo.lat": "31" },
           { "id": 3, "_geo.lat": 31, "_geo.lng": "42" },
         ]);
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(documents).unwrap();
+        let (builder, user_error) = builder.add_documents(documents).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1123,10 +1162,11 @@ mod tests {
         let documents = documents!([
           { "id": 0, "_geo": { "lng": 42 } }
         ]);
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(documents).unwrap();
+        let (builder, user_error) = builder.add_documents(documents).unwrap();
+        user_error.unwrap();
         let error = builder.execute().unwrap_err();
         assert_eq!(
             &error.to_string(),
@@ -1136,10 +1176,11 @@ mod tests {
         let documents = documents!([
           { "id": 0, "_geo": { "lat": 42 } }
         ]);
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(documents).unwrap();
+        let (builder, user_error) = builder.add_documents(documents).unwrap();
+        user_error.unwrap();
         let error = builder.execute().unwrap_err();
         assert_eq!(
             &error.to_string(),
@@ -1149,40 +1190,43 @@ mod tests {
         let documents = documents!([
           { "id": 0, "_geo": { "lat": "lol", "lng": 42 } }
         ]);
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(documents).unwrap();
+        let (builder, user_error) = builder.add_documents(documents).unwrap();
+        user_error.unwrap();
         let error = builder.execute().unwrap_err();
         assert_eq!(
             &error.to_string(),
-            r#"Could not parse latitude in the document with the id: `0`. Was expecting a number but instead got `"lol"`."#
+            r#"Could not parse latitude in the document with the id: `0`. Was expecting a finite number but instead got `"lol"`."#
         );
 
         let documents = documents!([
           { "id": 0, "_geo": { "lat": [12, 13], "lng": 42 } }
         ]);
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(documents).unwrap();
+        let (builder, user_error) = builder.add_documents(documents).unwrap();
+        user_error.unwrap();
         let error = builder.execute().unwrap_err();
         assert_eq!(
             &error.to_string(),
-            r#"Could not parse latitude in the document with the id: `0`. Was expecting a number but instead got `[12,13]`."#
+            r#"Could not parse latitude in the document with the id: `0`. Was expecting a finite number but instead got `[12,13]`."#
         );
 
         let documents = documents!([
           { "id": 0, "_geo": { "lat": 12, "lng": "hello" } }
         ]);
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(documents).unwrap();
+        let (builder, user_error) = builder.add_documents(documents).unwrap();
+        user_error.unwrap();
         let error = builder.execute().unwrap_err();
         assert_eq!(
             &error.to_string(),
-            r#"Could not parse longitude in the document with the id: `0`. Was expecting a number but instead got `"hello"`."#
+            r#"Could not parse longitude in the document with the id: `0`. Was expecting a finite number but instead got `"hello"`."#
         );
     }
 
@@ -1202,10 +1246,11 @@ mod tests {
         ]);
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         assert_eq!(index.primary_key(&wtxn).unwrap(), Some("objectId"));
@@ -1222,10 +1267,11 @@ mod tests {
             { "objectId": 30,  "title": "Hamlet", "_geo": { "lat": 12, "lng": 89 } }
         ]);
 
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         let external_documents_ids = index.external_documents_ids(&wtxn).unwrap();
         assert!(external_documents_ids.get("30").is_some());
@@ -1234,10 +1280,11 @@ mod tests {
             { "objectId": 30,  "title": "Hamlet", "_geo": { "lat": 12, "lng": 89 } }
         ]);
 
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         wtxn.commit().unwrap();
@@ -1252,28 +1299,25 @@ mod tests {
 
         let mut wtxn = index.write_txn().unwrap();
 
-        let mut big_object = HashMap::new();
-        big_object.insert(S("id"), "wow");
+        let mut big_object = serde_json::Map::new();
+        big_object.insert(S("id"), serde_json::Value::from("wow"));
         for i in 0..1000 {
             let key = i.to_string();
-            big_object.insert(key, "I am a text!");
+            big_object.insert(key, serde_json::Value::from("I am a text!"));
         }
 
-        let mut cursor = Cursor::new(Vec::new());
-
-        let mut builder = DocumentBatchBuilder::new(&mut cursor).unwrap();
-        let big_object = Cursor::new(serde_json::to_vec(&big_object).unwrap());
-        builder.extend_from_json(big_object).unwrap();
-        builder.finish().unwrap();
-        cursor.set_position(0);
-        let content = DocumentBatchReader::from_reader(cursor).unwrap();
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_json_object(&big_object).unwrap();
+        let vector = builder.into_inner().unwrap();
+        let content = DocumentsBatchReader::from_reader(Cursor::new(vector)).unwrap();
 
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         wtxn.commit().unwrap();
@@ -1288,30 +1332,27 @@ mod tests {
 
         let mut wtxn = index.write_txn().unwrap();
 
-        let mut big_object = HashMap::new();
-        big_object.insert(S("id"), "wow");
+        let mut big_object = serde_json::Map::new();
+        big_object.insert(S("id"), serde_json::Value::from("wow"));
         let content: String = (0..=u16::MAX)
             .into_iter()
             .map(|p| p.to_string())
             .reduce(|a, b| a + " " + b.as_ref())
             .unwrap();
-        big_object.insert("content".to_string(), &content);
+        big_object.insert("content".to_string(), serde_json::Value::from(content));
 
-        let mut cursor = Cursor::new(Vec::new());
-
-        let big_object = serde_json::to_string(&big_object).unwrap();
-        let mut builder = DocumentBatchBuilder::new(&mut cursor).unwrap();
-        builder.extend_from_json(&mut big_object.as_bytes()).unwrap();
-        builder.finish().unwrap();
-        cursor.set_position(0);
-        let content = DocumentBatchReader::from_reader(cursor).unwrap();
+        let mut builder = DocumentsBatchBuilder::new(Vec::new());
+        builder.append_json_object(&big_object).unwrap();
+        let vector = builder.into_inner().unwrap();
+        let content = DocumentsBatchReader::from_reader(Cursor::new(vector)).unwrap();
 
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         wtxn.commit().unwrap();
@@ -1366,10 +1407,11 @@ mod tests {
 
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         wtxn.commit().unwrap();
@@ -1419,10 +1461,11 @@ mod tests {
 
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         wtxn.commit().unwrap();
@@ -1551,10 +1594,11 @@ mod tests {
         ]);
 
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1583,6 +1627,58 @@ mod tests {
         assert_eq!(documents_ids, vec![3]);
     }
 
+    #[test]
+    fn retrieve_a_b_nested_document_id() {
+        let path = tempfile::tempdir().unwrap();
+        let mut options = EnvOpenOptions::new();
+        options.map_size(10 * 1024 * 1024); // 10 MB
+        let index = Index::new(options, &path).unwrap();
+        let config = IndexerConfig::default();
+
+        let mut wtxn = index.write_txn().unwrap();
+        let mut builder = update::Settings::new(&mut wtxn, &index, &config);
+        builder.set_primary_key("a.b".to_owned());
+        builder.execute(|_| ()).unwrap();
+
+        let content = documents!({ "a" : { "b" : { "c" :  1 }}});
+        let indexing_config = IndexDocumentsConfig::default();
+        let builder =
+            IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
+                .unwrap();
+        let (_builder, user_error) = builder.add_documents(content).unwrap();
+
+        // There must be an issue with the primary key no present in the given document
+        user_error.unwrap_err();
+    }
+
+    #[test]
+    fn retrieve_a_b_c_nested_document_id() {
+        let path = tempfile::tempdir().unwrap();
+        let mut options = EnvOpenOptions::new();
+        options.map_size(10 * 1024 * 1024); // 10 MB
+        let index = Index::new(options, &path).unwrap();
+        let config = IndexerConfig::default();
+
+        let mut wtxn = index.write_txn().unwrap();
+        let mut builder = update::Settings::new(&mut wtxn, &index, &config);
+        builder.set_primary_key("a.b.c".to_owned());
+        builder.execute(|_| ()).unwrap();
+
+        let content = documents!({ "a" : { "b" : { "c" :  1 }}});
+        let indexing_config = IndexDocumentsConfig::default();
+        let builder =
+            IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
+                .unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
+        builder.execute().unwrap();
+        wtxn.commit().unwrap();
+
+        let rtxn = index.read_txn().unwrap();
+        let external_documents_ids = index.external_documents_ids(&rtxn).unwrap();
+        assert!(external_documents_ids.get("1").is_some());
+    }
+
     #[test]
     fn test_facets_generation() {
         let path = tempfile::tempdir().unwrap();
@@ -1621,10 +1717,11 @@ mod tests {
         // index the documents
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         wtxn.commit().unwrap();
@@ -1713,10 +1810,11 @@ mod tests {
         let mut wtxn = index.write_txn().unwrap();
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1730,10 +1828,11 @@ mod tests {
         let mut wtxn = index.write_txn().unwrap();
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1752,10 +1851,11 @@ mod tests {
         let mut wtxn = index.write_txn().unwrap();
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1780,10 +1880,11 @@ mod tests {
         let mut wtxn = index.write_txn().unwrap();
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1825,10 +1926,11 @@ mod tests {
         let mut wtxn = index.write_txn().unwrap();
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
     }
@@ -1843,28 +1945,31 @@ mod tests {
 
         // Create 200 documents with a long text
         let content = {
-            let documents: Vec<_> = (0..200i32)
+            let documents_iter = (0..200i32)
                 .into_iter()
                 .map(|i| serde_json::json!({ "id": i, "script": script }))
-                .collect();
+                .filter_map(|json| match json {
+                    serde_json::Value::Object(object) => Some(object),
+                    _ => None,
+                });
 
-            let mut writer = std::io::Cursor::new(Vec::new());
-            let mut builder = crate::documents::DocumentBatchBuilder::new(&mut writer).unwrap();
-            let documents = serde_json::to_vec(&documents).unwrap();
-            builder.extend_from_json(std::io::Cursor::new(documents)).unwrap();
-            builder.finish().unwrap();
-            writer.set_position(0);
-            crate::documents::DocumentBatchReader::from_reader(writer).unwrap()
+            let mut builder = crate::documents::DocumentsBatchBuilder::new(Vec::new());
+            for object in documents_iter {
+                builder.append_json_object(&object).unwrap();
+            }
+            let vector = builder.into_inner().unwrap();
+            crate::documents::DocumentsBatchReader::from_reader(Cursor::new(vector)).unwrap()
         };
 
         // Index those 200 long documents
         let mut wtxn = index.write_txn().unwrap();
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         // Create one long document
@@ -1875,10 +1980,11 @@ mod tests {
         // Index this one long document
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         wtxn.commit().unwrap();
@@ -1892,7 +1998,7 @@ mod tests {
         let index = Index::new(options, tmp).unwrap();
         let mut wtxn = index.write_txn().unwrap();
         let indexer_config = IndexerConfig::default();
-        let mut builder = IndexDocuments::new(
+        let builder = IndexDocuments::new(
             &mut wtxn,
             &index,
             &indexer_config,
@@ -1921,8 +2027,10 @@ mod tests {
             "branch_id_number": 0
         }]};
 
-        builder.add_documents(doc1).unwrap();
-        builder.add_documents(doc2).unwrap();
+        let (builder, user_error) = builder.add_documents(doc1).unwrap();
+        user_error.unwrap();
+        let (builder, user_error) = builder.add_documents(doc2).unwrap();
+        user_error.unwrap();
 
         builder.execute().unwrap();
 
@@ -1931,4 +2039,51 @@ mod tests {
 
         assert_eq!(ids.len(), map.len());
     }
+
+    #[test]
+    fn primary_key_must_not_contain_floats() {
+        let tmp = tempfile::tempdir().unwrap();
+        let mut options = EnvOpenOptions::new();
+        options.map_size(4096 * 100);
+        let index = Index::new(options, tmp).unwrap();
+        let mut wtxn = index.write_txn().unwrap();
+        let indexer_config = IndexerConfig::default();
+        let builder = IndexDocuments::new(
+            &mut wtxn,
+            &index,
+            &indexer_config,
+            IndexDocumentsConfig::default(),
+            |_| (),
+        )
+        .unwrap();
+
+        let doc1 = documents! {[{
+            "id": -228142,
+            "title": "asdsad",
+        }]};
+
+        let doc2 = documents! {[{
+            "id": 228143.56,
+            "title": "something",
+        }]};
+
+        let doc3 = documents! {[{
+            "id": -228143.56,
+            "title": "something",
+        }]};
+
+        let doc4 = documents! {[{
+            "id": 2.0,
+            "title": "something",
+        }]};
+
+        let (builder, user_error) = builder.add_documents(doc1).unwrap();
+        user_error.unwrap();
+        let (builder, user_error) = builder.add_documents(doc2).unwrap();
+        assert!(user_error.is_err());
+        let (builder, user_error) = builder.add_documents(doc3).unwrap();
+        assert!(user_error.is_err());
+        let (_builder, user_error) = builder.add_documents(doc4).unwrap();
+        assert!(user_error.is_err());
+    }
 }

milli/src/update/index_documents/transform.rs

@@ -9,12 +9,12 @@ use heed::RoTxn;
 use itertools::Itertools;
 use obkv::{KvReader, KvWriter};
 use roaring::RoaringBitmap;
-use serde_json::{Map, Value};
+use serde_json::Value;
 use smartstring::SmartString;
 
 use super::helpers::{create_sorter, create_writer, keep_latest_obkv, merge_obkvs, MergeFn};
 use super::{IndexDocumentsMethod, IndexerConfig};
-use crate::documents::{DocumentBatchReader, DocumentsBatchIndex};
+use crate::documents::{DocumentsBatchIndex, EnrichedDocument, EnrichedDocumentsBatchReader};
 use crate::error::{Error, InternalError, UserError};
 use crate::index::db_name;
 use crate::update::{AvailableDocumentsIds, UpdateIndexingStep};
@@ -23,8 +23,6 @@ use crate::{
     Result, BEU32,
 };
 
-const DEFAULT_PRIMARY_KEY_NAME: &str = "id";
-
 pub struct TransformOutput {
     pub primary_key: String,
     pub fields_ids_map: FieldsIdsMap,
@@ -84,18 +82,6 @@ fn create_fields_mapping(
         .collect()
 }
 
-/// Look for a key containing the [DEFAULT_PRIMARY_KEY_NAME] in the fields.
-/// It doesn't look in the subfield because we don't want to enable the
-/// primary key inference on nested objects.
-fn find_primary_key(index: &DocumentsBatchIndex) -> Option<&str> {
-    index
-        .iter()
-        .sorted_by_key(|(k, _)| *k)
-        .map(|(_, v)| v)
-        .find(|v| v.to_lowercase().contains(DEFAULT_PRIMARY_KEY_NAME))
-        .map(String::as_str)
-}
-
 impl<'a, 'i> Transform<'a, 'i> {
     pub fn new(
         wtxn: &mut heed::RwTxn,
@@ -152,7 +138,7 @@ impl<'a, 'i> Transform<'a, 'i> {
 
     pub fn read_documents<R, F>(
         &mut self,
-        mut reader: DocumentBatchReader<R>,
+        reader: EnrichedDocumentsBatchReader<R>,
         wtxn: &mut heed::RwTxn,
         progress_callback: F,
     ) -> Result<usize>
@@ -160,33 +146,25 @@ impl<'a, 'i> Transform<'a, 'i> {
         R: Read + Seek,
         F: Fn(UpdateIndexingStep) + Sync,
     {
-        let fields_index = reader.index();
+        let (mut cursor, fields_index) = reader.into_cursor_and_fields_index();
+
         let external_documents_ids = self.index.external_documents_ids(wtxn)?;
 
-        let mapping = create_fields_mapping(&mut self.fields_ids_map, fields_index)?;
+        let mapping = create_fields_mapping(&mut self.fields_ids_map, &fields_index)?;
 
-        let alternative_name = self
-            .index
-            .primary_key(wtxn)?
-            .or_else(|| find_primary_key(fields_index))
-            .map(String::from);
+        let primary_key = cursor.primary_key().to_string();
+        let primary_key_id =
+            self.fields_ids_map.insert(&primary_key).ok_or(UserError::AttributeLimitReached)?;
 
-        let (primary_key_id, primary_key_name) = compute_primary_key_pair(
-            self.index.primary_key(wtxn)?,
-            &mut self.fields_ids_map,
-            alternative_name,
-            self.autogenerate_docids,
-        )?;
-
-        let primary_key_id_nested = primary_key_name.contains('.');
-
-        let mut flattened_document = None;
         let mut obkv_buffer = Vec::new();
-        let mut flattened_obkv_buffer = Vec::new();
         let mut documents_count = 0;
-        let mut external_id_buffer = Vec::new();
+        let mut docid_buffer: Vec<u8> = Vec::new();
         let mut field_buffer: Vec<(u16, Cow<[u8]>)> = Vec::new();
-        while let Some((addition_index, document)) = reader.next_document_with_index()? {
+        while let Some(enriched_document) = cursor.next_enriched_document()? {
+            let EnrichedDocument { document, document_id } = enriched_document;
+
+            // drop_and_reuse is called instead of .clear() to communicate to the compiler that field_buffer
+            // does not keep references from the cursor between loop iterations
             let mut field_buffer_cache = drop_and_reuse(field_buffer);
             if self.indexer_settings.log_every_n.map_or(false, |len| documents_count % len == 0) {
                 progress_callback(UpdateIndexingStep::RemapDocumentAddition {
@@ -194,52 +172,21 @@ impl<'a, 'i> Transform<'a, 'i> {
                 });
             }
 
+            // When the document id has been auto-generated by the `enrich_documents_batch`
+            // we must insert this document id into the remaped document.
+            let external_id = document_id.value();
+            if document_id.is_generated() {
+                serde_json::to_writer(&mut docid_buffer, external_id)
+                    .map_err(InternalError::SerdeJson)?;
+                field_buffer_cache.push((primary_key_id, Cow::from(&docid_buffer)));
+            }
+
             for (k, v) in document.iter() {
                 let mapped_id =
                     *mapping.get(&k).ok_or(InternalError::FieldIdMappingMissingEntry { key: k })?;
                 field_buffer_cache.push((mapped_id, Cow::from(v)));
             }
 
-            // We need to make sure that every document has a primary key. After we have remapped
-            // all the fields in the document, we try to find the primary key value. If we can find
-            // it, transform it into a string and validate it, and then update it in the
-            // document. If none is found, and we were told to generate missing document ids, then
-            // we create the missing field, and update the new document.
-            let mut uuid_buffer = [0; uuid::fmt::Hyphenated::LENGTH];
-            let external_id = if primary_key_id_nested {
-                let mut field_buffer_cache = field_buffer_cache.clone();
-                self.flatten_from_field_mapping(
-                    &mapping,
-                    &document,
-                    &mut flattened_obkv_buffer,
-                    &mut field_buffer_cache,
-                )?;
-                flattened_document = Some(&flattened_obkv_buffer);
-                let document = KvReader::new(&flattened_obkv_buffer);
-
-                update_primary_key(
-                    document,
-                    &addition_index,
-                    primary_key_id,
-                    &primary_key_name,
-                    &mut uuid_buffer,
-                    &mut field_buffer_cache,
-                    &mut external_id_buffer,
-                    self.autogenerate_docids,
-                )?
-            } else {
-                update_primary_key(
-                    document,
-                    &addition_index,
-                    primary_key_id,
-                    &primary_key_name,
-                    &mut uuid_buffer,
-                    &mut field_buffer_cache,
-                    &mut external_id_buffer,
-                    self.autogenerate_docids,
-                )?
-            };
-
             // Insertion in a obkv need to be done with keys ordered. For now they are ordered
             // according to the document addition key order, so we sort it according to the
             // fieldids map keys order.
@@ -294,18 +241,12 @@ impl<'a, 'i> Transform<'a, 'i> {
             }
 
             // We use the extracted/generated user id as the key for this document.
-            self.original_sorter.insert(&docid.to_be_bytes(), obkv_buffer.clone())?;
+            self.original_sorter.insert(&docid.to_be_bytes(), &obkv_buffer)?;
             documents_count += 1;
 
-            if let Some(flatten) = flattened_document {
-                self.flattened_sorter.insert(docid.to_be_bytes(), &flatten)?;
-            } else {
             match self.flatten_from_fields_ids_map(KvReader::new(&obkv_buffer))? {
                 Some(buffer) => self.flattened_sorter.insert(docid.to_be_bytes(), &buffer)?,
-                    None => {
-                        self.flattened_sorter.insert(docid.to_be_bytes(), obkv_buffer.clone())?
-                    }
-                }
+                None => self.flattened_sorter.insert(docid.to_be_bytes(), &obkv_buffer)?,
             }
 
             progress_callback(UpdateIndexingStep::RemapDocumentAddition {
@@ -313,7 +254,7 @@ impl<'a, 'i> Transform<'a, 'i> {
             });
 
             field_buffer = drop_and_reuse(field_buffer_cache);
-            external_id_buffer.clear();
+            docid_buffer.clear();
             obkv_buffer.clear();
         }
 
@@ -322,7 +263,7 @@ impl<'a, 'i> Transform<'a, 'i> {
         });
 
         self.index.put_fields_ids_map(wtxn, &self.fields_ids_map)?;
-        self.index.put_primary_key(wtxn, &primary_key_name)?;
+        self.index.put_primary_key(wtxn, &primary_key)?;
         self.documents_count += documents_count;
         // Now that we have a valid sorter that contains the user id and the obkv we
         // give it to the last transforming function which returns the TransformOutput.
@@ -384,61 +325,6 @@ impl<'a, 'i> Transform<'a, 'i> {
         Ok(Some(buffer))
     }
 
-    // Flatten a document from a field mapping generated by [create_fields_mapping]
-    fn flatten_from_field_mapping(
-        &mut self,
-        mapping: &HashMap<FieldId, FieldId>,
-        obkv: &KvReader<FieldId>,
-        output_buffer: &mut Vec<u8>,
-        field_buffer_cache: &mut Vec<(u16, Cow<[u8]>)>,
-    ) -> Result<()> {
-        // store the keys and values of the json + the original obkv
-        let mut key_value: Vec<(FieldId, Cow<[u8]>)> = Vec::new();
-
-        // if the primary_key is nested we need to flatten the document before being able to do anything
-        let mut doc = serde_json::Map::new();
-
-        // we recreate a json containing only the fields that needs to be flattened.
-        // all the raw values get inserted directly in the `key_value` vec.
-        for (key, value) in obkv.iter() {
-            if json_depth_checker::should_flatten_from_unchecked_slice(value) {
-                let key =
-                    mapping.get(&key).ok_or(InternalError::FieldIdMappingMissingEntry { key })?;
-                let key =
-                    self.fields_ids_map.name(*key).ok_or(FieldIdMapMissingEntry::FieldId {
-                        field_id: *key,
-                        process: "Flatten from field mapping.",
-                    })?;
-                let value = serde_json::from_slice::<serde_json::Value>(value)
-                    .map_err(InternalError::SerdeJson)?;
-                doc.insert(key.to_string(), value);
-            } else {
-                key_value.push((key, value.into()));
-            }
-        }
-
-        let flattened = flatten_serde_json::flatten(&doc);
-
-        // Once we have the flattened version we insert all the new generated fields_ids
-        // (if any) in the fields ids map and serialize the value.
-        for (key, value) in flattened.into_iter() {
-            let fid = self.fields_ids_map.insert(&key).ok_or(UserError::AttributeLimitReached)?;
-            let value = serde_json::to_vec(&value).map_err(InternalError::SerdeJson)?;
-            key_value.push((fid, value.clone().into()));
-
-            if field_buffer_cache.iter().find(|(id, _)| *id == fid).is_none() {
-                field_buffer_cache.push((fid, value.into()));
-            }
-        }
-
-        // we sort the key. If there was a conflict between the obkv and the new generated value the
-        // keys will be consecutive.
-        key_value.sort_unstable_by_key(|(key, _)| *key);
-
-        Self::create_obkv_from_key_value(&mut key_value, output_buffer)?;
-        Ok(())
-    }
-
     /// Generate an obkv from a slice of key / value sorted by key.
     fn create_obkv_from_key_value(
         key_value: &mut [(FieldId, Cow<[u8]>)],
@@ -744,50 +630,6 @@ impl<'a, 'i> Transform<'a, 'i> {
     }
 }
 
-/// Given an optional primary key and an optional alternative name, returns the (field_id, attr_name)
-/// for the primary key according to the following rules:
-/// - if primary_key is `Some`, returns the id and the name, else
-/// - if alternative_name is Some, adds alternative to the fields_ids_map, and returns the pair, else
-/// - if autogenerate_docids is true, insert the default id value in the field ids map ("id") and
-/// returns the pair, else
-/// - returns an error.
-fn compute_primary_key_pair(
-    primary_key: Option<&str>,
-    fields_ids_map: &mut FieldsIdsMap,
-    alternative_name: Option<String>,
-    autogenerate_docids: bool,
-) -> Result<(FieldId, String)> {
-    match primary_key {
-        Some(primary_key) => {
-            let id = fields_ids_map.insert(primary_key).ok_or(UserError::AttributeLimitReached)?;
-            Ok((id, primary_key.to_string()))
-        }
-        None => {
-            let name = match alternative_name {
-                Some(key) => key,
-                None => {
-                    if !autogenerate_docids {
-                        // If there is no primary key in the current document batch, we must
-                        // return an error and not automatically generate any document id.
-                        return Err(UserError::MissingPrimaryKey.into());
-                    }
-                    DEFAULT_PRIMARY_KEY_NAME.to_string()
-                }
-            };
-            let id = fields_ids_map.insert(&name).ok_or(UserError::AttributeLimitReached)?;
-            Ok((id, name))
-        }
-    }
-}
-
-fn validate_document_id(document_id: &str) -> Option<&str> {
-    let document_id = document_id.trim();
-    Some(document_id).filter(|id| {
-        !id.is_empty()
-            && id.chars().all(|c| matches!(c, 'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_'))
-    })
-}
-
 /// Drops all the value of type `U` in vec, and reuses the allocation to create a `Vec<T>`.
 ///
 /// The size and alignment of T and U must match.
@@ -799,63 +641,6 @@ fn drop_and_reuse<U, T>(mut vec: Vec<U>) -> Vec<T> {
     vec.into_iter().map(|_| unreachable!()).collect()
 }
 
-fn update_primary_key<'a>(
-    document: KvReader<'a, FieldId>,
-    addition_index: &DocumentsBatchIndex,
-    primary_key_id: FieldId,
-    primary_key_name: &str,
-    uuid_buffer: &'a mut [u8; uuid::fmt::Hyphenated::LENGTH],
-    field_buffer_cache: &mut Vec<(u16, Cow<'a, [u8]>)>,
-    mut external_id_buffer: &'a mut Vec<u8>,
-    autogenerate_docids: bool,
-) -> Result<Cow<'a, str>> {
-    match field_buffer_cache.iter_mut().find(|(id, _)| *id == primary_key_id) {
-        Some((_, bytes)) => {
-            let value = match serde_json::from_slice(bytes).map_err(InternalError::SerdeJson)? {
-                Value::String(string) => match validate_document_id(&string) {
-                    Some(s) if s.len() == string.len() => string,
-                    Some(s) => s.to_string(),
-                    None => {
-                        return Err(UserError::InvalidDocumentId {
-                            document_id: Value::String(string),
-                        }
-                        .into())
-                    }
-                },
-                Value::Number(number) => number.to_string(),
-                content => {
-                    return Err(UserError::InvalidDocumentId { document_id: content.clone() }.into())
-                }
-            };
-            serde_json::to_writer(external_id_buffer, &value).map_err(InternalError::SerdeJson)?;
-            Ok(Cow::Owned(value))
-        }
-        None if autogenerate_docids => {
-            let uuid = uuid::Uuid::new_v4().as_hyphenated().encode_lower(uuid_buffer);
-            serde_json::to_writer(&mut external_id_buffer, &uuid)
-                .map_err(InternalError::SerdeJson)?;
-            field_buffer_cache.push((primary_key_id, external_id_buffer.as_slice().into()));
-            Ok(Cow::Borrowed(&*uuid))
-        }
-        None => {
-            let mut json = Map::new();
-            for (key, value) in document.iter() {
-                let key = addition_index.name(key).cloned();
-                let value = serde_json::from_slice::<Value>(&value).ok();
-
-                if let Some((k, v)) = key.zip(value) {
-                    json.insert(k, v);
-                }
-            }
-
-            Err(UserError::MissingDocumentId {
-                primary_key: primary_key_name.to_string(),
-                document: json,
-            })?
-        }
-    }
-}
-
 impl TransformOutput {
     // find and insert the new field ids
     pub fn compute_real_facets(&self, rtxn: &RoTxn, index: &Index) -> Result<HashSet<String>> {
@@ -869,88 +654,3 @@ impl TransformOutput {
             .collect())
     }
 }
-
-#[cfg(test)]
-mod test {
-    use super::*;
-
-    mod compute_primary_key {
-        use big_s::S;
-
-        use super::{compute_primary_key_pair, FieldsIdsMap};
-
-        #[test]
-        fn should_return_primary_key_if_is_some() {
-            let mut fields_map = FieldsIdsMap::new();
-            fields_map.insert("toto").unwrap();
-            let result = compute_primary_key_pair(
-                Some("toto"),
-                &mut fields_map,
-                Some("tata".to_string()),
-                false,
-            );
-            assert_eq!(result.unwrap(), (0, "toto".to_string()));
-            assert_eq!(fields_map.len(), 1);
-
-            // and with nested fields
-            let mut fields_map = FieldsIdsMap::new();
-            fields_map.insert("toto.tata").unwrap();
-            let result = compute_primary_key_pair(
-                Some("toto.tata"),
-                &mut fields_map,
-                Some(S("titi")),
-                false,
-            );
-            assert_eq!(result.unwrap(), (0, "toto.tata".to_string()));
-            assert_eq!(fields_map.len(), 1);
-        }
-
-        #[test]
-        fn should_return_alternative_if_primary_is_none() {
-            let mut fields_map = FieldsIdsMap::new();
-            let result =
-                compute_primary_key_pair(None, &mut fields_map, Some("tata".to_string()), false);
-            assert_eq!(result.unwrap(), (0, S("tata")));
-            assert_eq!(fields_map.len(), 1);
-        }
-
-        #[test]
-        fn should_return_default_if_both_are_none() {
-            let mut fields_map = FieldsIdsMap::new();
-            let result = compute_primary_key_pair(None, &mut fields_map, None, true);
-            assert_eq!(result.unwrap(), (0, S("id")));
-            assert_eq!(fields_map.len(), 1);
-        }
-
-        #[test]
-        fn should_return_err_if_both_are_none_and_recompute_is_false() {
-            let mut fields_map = FieldsIdsMap::new();
-            let result = compute_primary_key_pair(None, &mut fields_map, None, false);
-            assert!(result.is_err());
-            assert_eq!(fields_map.len(), 0);
-        }
-    }
-
-    mod primary_key_inference {
-        use big_s::S;
-        use bimap::BiHashMap;
-
-        use crate::documents::DocumentsBatchIndex;
-        use crate::update::index_documents::transform::find_primary_key;
-
-        #[test]
-        fn primary_key_infered_on_first_field() {
-            // We run the test multiple times to change the order in which the fields are iterated upon.
-            for _ in 1..50 {
-                let mut map = BiHashMap::new();
-                map.insert(1, S("fakeId"));
-                map.insert(2, S("fakeId"));
-                map.insert(3, S("fakeId"));
-                map.insert(4, S("fakeId"));
-                map.insert(0, S("realId"));
-
-                assert_eq!(find_primary_key(&DocumentsBatchIndex(map)), Some("realId"));
-            }
-        }
-    }
-}

milli/src/update/mod.rs

@@ -3,7 +3,7 @@ pub use self::clear_documents::ClearDocuments;
 pub use self::delete_documents::{DeleteDocuments, DocumentDeletionResult};
 pub use self::facets::Facets;
 pub use self::index_documents::{
-    DocumentAdditionResult, IndexDocuments, IndexDocumentsConfig, IndexDocumentsMethod,
+    DocumentAdditionResult, DocumentId, IndexDocuments, IndexDocumentsConfig, IndexDocumentsMethod,
 };
 pub use self::indexer_config::IndexerConfig;
 pub use self::settings::{Setting, Settings};

milli/src/update/settings.rs

@@ -735,10 +735,11 @@ mod tests {
         ]);
         let config = IndexerConfig::default();
         let indexing_config = IndexDocumentsConfig::default();
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -798,10 +799,11 @@ mod tests {
         let config = IndexerConfig::default();
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -850,10 +852,11 @@ mod tests {
         let config = IndexerConfig::default();
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -880,10 +883,11 @@ mod tests {
         let config = IndexerConfig::default();
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         // In the same transaction we change the displayed fields to be only the age.
@@ -934,10 +938,11 @@ mod tests {
         ]);
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -974,10 +979,11 @@ mod tests {
 
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1016,10 +1022,11 @@ mod tests {
         ]);
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1067,10 +1074,11 @@ mod tests {
         ]);
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1110,10 +1118,11 @@ mod tests {
         ]);
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1142,10 +1151,11 @@ mod tests {
         let config = IndexerConfig::default();
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1172,10 +1182,11 @@ mod tests {
         let config = IndexerConfig::default();
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         // In the same transaction we provide some stop_words
@@ -1251,10 +1262,11 @@ mod tests {
         let config = IndexerConfig::default();
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
 
         // In the same transaction provide some synonyms
@@ -1389,10 +1401,11 @@ mod tests {
         ]);
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 
@@ -1452,10 +1465,11 @@ mod tests {
         ]);
         let indexing_config =
             IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-        let mut builder =
+        let builder =
             IndexDocuments::new(&mut wtxn, &index, &config, indexing_config.clone(), |_| ())
                 .unwrap();
-        builder.add_documents(content).unwrap();
+        let (builder, user_error) = builder.add_documents(content).unwrap();
+        user_error.unwrap();
         builder.execute().unwrap();
         wtxn.commit().unwrap();
 

milli/tests/search/facet_distribution.rs

@@ -3,9 +3,10 @@ use std::io::Cursor;
 use big_s::S;
 use heed::EnvOpenOptions;
 use maplit::hashset;
-use milli::documents::{DocumentBatchBuilder, DocumentBatchReader};
+use milli::documents::{DocumentsBatchBuilder, DocumentsBatchReader};
 use milli::update::{IndexDocuments, IndexDocumentsConfig, IndexerConfig, Settings};
-use milli::{FacetDistribution, Index};
+use milli::{FacetDistribution, Index, Object};
+use serde_json::Deserializer;
 
 #[test]
 fn test_facet_distribution_with_no_facet_values() {
@@ -28,38 +29,33 @@ fn test_facet_distribution_with_no_facet_values() {
     let config = IndexerConfig { max_memory: Some(10 * 1024 * 1024), ..Default::default() };
     let indexing_config = IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
 
-    let mut builder =
-        IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-    let mut cursor = Cursor::new(Vec::new());
-    let mut documents_builder = DocumentBatchBuilder::new(&mut cursor).unwrap();
+    let builder = IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
+    let mut documents_builder = DocumentsBatchBuilder::new(Vec::new());
     let reader = Cursor::new(
-        r#"[
-        {
+        r#"{
             "id": 123,
             "title": "What a week, hu...",
             "genres": [],
             "tags": ["blue"]
-        },
+        }
         {
             "id": 345,
             "title": "I am the pig!",
             "tags": ["red"]
-        }
-    ]"#,
+        }"#,
     );
 
-    for doc in serde_json::Deserializer::from_reader(reader).into_iter::<serde_json::Value>() {
-        let doc = Cursor::new(serde_json::to_vec(&doc.unwrap()).unwrap());
-        documents_builder.extend_from_json(doc).unwrap();
+    for result in Deserializer::from_reader(reader).into_iter::<Object>() {
+        let object = result.unwrap();
+        documents_builder.append_json_object(&object).unwrap();
     }
 
-    documents_builder.finish().unwrap();
-
-    cursor.set_position(0);
+    let vector = documents_builder.into_inner().unwrap();
 
     // index documents
-    let content = DocumentBatchReader::from_reader(cursor).unwrap();
-    builder.add_documents(content).unwrap();
+    let content = DocumentsBatchReader::from_reader(Cursor::new(vector)).unwrap();
+    let (builder, user_error) = builder.add_documents(content).unwrap();
+    user_error.unwrap();
     builder.execute().unwrap();
 
     wtxn.commit().unwrap();

milli/tests/search/mod.rs

@@ -6,10 +6,11 @@ use big_s::S;
 use either::{Either, Left, Right};
 use heed::EnvOpenOptions;
 use maplit::{hashmap, hashset};
-use milli::documents::{DocumentBatchBuilder, DocumentBatchReader};
+use milli::documents::{DocumentsBatchBuilder, DocumentsBatchReader};
 use milli::update::{IndexDocuments, IndexDocumentsConfig, IndexerConfig, Settings};
-use milli::{AscDesc, Criterion, DocumentId, Index, Member};
+use milli::{AscDesc, Criterion, DocumentId, Index, Member, Object};
 use serde::Deserialize;
+use serde_json::Deserializer;
 use slice_group_by::GroupBy;
 
 mod distinct;
@@ -60,24 +61,21 @@ pub fn setup_search_index_with_criteria(criteria: &[Criterion]) -> Index {
     let config = IndexerConfig { max_memory: Some(10 * 1024 * 1024), ..Default::default() };
     let indexing_config = IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
 
-    let mut builder =
-        IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
-    let mut cursor = Cursor::new(Vec::new());
-    let mut documents_builder = DocumentBatchBuilder::new(&mut cursor).unwrap();
+    let builder = IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
+    let mut documents_builder = DocumentsBatchBuilder::new(Vec::new());
     let reader = Cursor::new(CONTENT.as_bytes());
 
-    for doc in serde_json::Deserializer::from_reader(reader).into_iter::<serde_json::Value>() {
-        let doc = Cursor::new(serde_json::to_vec(&doc.unwrap()).unwrap());
-        documents_builder.extend_from_json(doc).unwrap();
+    for result in Deserializer::from_reader(reader).into_iter::<Object>() {
+        let object = result.unwrap();
+        documents_builder.append_json_object(&object).unwrap();
     }
 
-    documents_builder.finish().unwrap();
-
-    cursor.set_position(0);
+    let vector = documents_builder.into_inner().unwrap();
 
     // index documents
-    let content = DocumentBatchReader::from_reader(cursor).unwrap();
-    builder.add_documents(content).unwrap();
+    let content = DocumentsBatchReader::from_reader(Cursor::new(vector)).unwrap();
+    let (builder, user_error) = builder.add_documents(content).unwrap();
+    user_error.unwrap();
     builder.execute().unwrap();
 
     wtxn.commit().unwrap();

milli/tests/search/query_criteria.rs

@@ -5,7 +5,7 @@ use big_s::S;
 use heed::EnvOpenOptions;
 use itertools::Itertools;
 use maplit::hashset;
-use milli::documents::{DocumentBatchBuilder, DocumentBatchReader};
+use milli::documents::{DocumentsBatchBuilder, DocumentsBatchReader};
 use milli::update::{IndexDocuments, IndexDocumentsConfig, IndexerConfig, Settings};
 use milli::{AscDesc, Criterion, Index, Member, Search, SearchResult};
 use rand::Rng;
@@ -390,11 +390,9 @@ fn criteria_ascdesc() {
     // index documents
     let config = IndexerConfig { max_memory: Some(10 * 1024 * 1024), ..Default::default() };
     let indexing_config = IndexDocumentsConfig { autogenerate_docids: true, ..Default::default() };
-    let mut builder =
-        IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
+    let builder = IndexDocuments::new(&mut wtxn, &index, &config, indexing_config, |_| ()).unwrap();
 
-    let mut cursor = Cursor::new(Vec::new());
-    let mut batch_builder = DocumentBatchBuilder::new(&mut cursor).unwrap();
+    let mut batch_builder = DocumentsBatchBuilder::new(Vec::new());
 
     (0..ASC_DESC_CANDIDATES_THRESHOLD + 1).for_each(|_| {
         let mut rng = rand::thread_rng();
@@ -412,17 +410,19 @@ fn criteria_ascdesc() {
             "age": age,
         });
 
-        let json = Cursor::new(serde_json::to_vec(&json).unwrap());
-        batch_builder.extend_from_json(json).unwrap();
+        let object = match json {
+            serde_json::Value::Object(object) => object,
+            _ => panic!(),
+        };
+
+        batch_builder.append_json_object(&object).unwrap();
     });
 
-    batch_builder.finish().unwrap();
+    let vector = batch_builder.into_inner().unwrap();
 
-    cursor.set_position(0);
-
-    let reader = DocumentBatchReader::from_reader(cursor).unwrap();
-
-    builder.add_documents(reader).unwrap();
+    let reader = DocumentsBatchReader::from_reader(Cursor::new(vector)).unwrap();
+    let (builder, user_error) = builder.add_documents(reader).unwrap();
+    user_error.unwrap();
     builder.execute().unwrap();
 
     wtxn.commit().unwrap();

milli/tests/search/typo_tolerance.rs

@@ -106,35 +106,31 @@ fn test_typo_disabled_on_word() {
     options.map_size(4096 * 100);
     let index = Index::new(options, tmp.path()).unwrap();
 
-    let documents = json!([
-        {
+    let mut builder = milli::documents::DocumentsBatchBuilder::new(Vec::new());
+    let doc1 = json!({
         "id": 1usize,
         "data": "zealand",
-        },
-        {
+    });
+
+    let doc2 = json!({
         "id": 2usize,
         "data": "zearand",
-        },
-    ]);
+    });
 
-    let mut writer = std::io::Cursor::new(Vec::new());
-    let mut builder = milli::documents::DocumentBatchBuilder::new(&mut writer).unwrap();
-    let documents = serde_json::to_vec(&documents).unwrap();
-    builder.extend_from_json(std::io::Cursor::new(documents)).unwrap();
-    builder.finish().unwrap();
+    builder.append_json_object(doc1.as_object().unwrap()).unwrap();
+    builder.append_json_object(doc2.as_object().unwrap()).unwrap();
+    let vector = builder.into_inner().unwrap();
 
-    writer.set_position(0);
-
-    let documents = milli::documents::DocumentBatchReader::from_reader(writer).unwrap();
+    let documents =
+        milli::documents::DocumentsBatchReader::from_reader(std::io::Cursor::new(vector)).unwrap();
 
     let mut txn = index.write_txn().unwrap();
     let config = IndexerConfig::default();
     let indexing_config = IndexDocumentsConfig::default();
-    let mut builder =
-        IndexDocuments::new(&mut txn, &index, &config, indexing_config, |_| ()).unwrap();
-
-    builder.add_documents(documents).unwrap();
+    let builder = IndexDocuments::new(&mut txn, &index, &config, indexing_config, |_| ()).unwrap();
 
+    let (builder, user_error) = builder.add_documents(documents).unwrap();
+    user_error.unwrap();
     builder.execute().unwrap();
     txn.commit().unwrap();