diff --git a/quickwit/Cargo.lock b/quickwit/Cargo.lock index 4c929de6fc0..601594bdcda 100644 --- a/quickwit/Cargo.lock +++ b/quickwit/Cargo.lock @@ -7562,6 +7562,7 @@ dependencies = [ "futures", "http 1.5.0", "itertools 0.15.0", + "memchr", "mockall", "mrecordlog", "prost 0.14.4", diff --git a/quickwit/Cargo.toml b/quickwit/Cargo.toml index 1ade0c4c0e2..a58d5f3797f 100644 --- a/quickwit/Cargo.toml +++ b/quickwit/Cargo.toml @@ -126,6 +126,7 @@ libz-sys = "1.1" lru = "0.18" matches = "0.1" md5 = "0.8" +memchr = "2" metrics = "0.24" metrics-opentelemetry = "=0.24032.8" metrics-exporter-prometheus = { version = "0.18", default-features = false } diff --git a/quickwit/quickwit-ingest/Cargo.toml b/quickwit/quickwit-ingest/Cargo.toml index 4ecaeddf401..ed8ce9f496a 100644 --- a/quickwit/quickwit-ingest/Cargo.toml +++ b/quickwit/quickwit-ingest/Cargo.toml @@ -19,6 +19,7 @@ fail = { workspace = true, optional = true } futures = { workspace = true } http = { workspace = true } itertools = { workspace = true } +memchr = { workspace = true } mockall = { workspace = true, optional = true } mrecordlog = { workspace = true } prost = { workspace = true } diff --git a/quickwit/quickwit-ingest/benches/mrecord_bench.rs b/quickwit/quickwit-ingest/benches/mrecord_bench.rs index 228a5567877..5e750d3d914 100644 --- a/quickwit/quickwit-ingest/benches/mrecord_bench.rs +++ b/quickwit/quickwit-ingest/benches/mrecord_bench.rs @@ -28,7 +28,8 @@ use std::time::{Duration, Instant}; use bytes::{BufMut, Bytes}; use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; use mrecordlog::MultiRecordLog; -use quickwit_ingest::MRecord; +use quickwit_ingest::{DocBatchV2Builder, MRecord}; +use quickwit_proto::types::DocUidGenerator; /// Representative document sizes (bytes): a small log line up to a large structured event. const DOC_SIZES: [usize; 4] = [128, 1_024, 8_192, 65_536]; @@ -180,5 +181,44 @@ fn bench_append_batch(criterion: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_encode, bench_decode, bench_append_batch); +fn bench_doc_batch_builder(criterion: &mut Criterion) { + let mut group = criterion.benchmark_group("doc_batch_v2_builder"); + let docs: Vec = (0..10_000) + .map(|index| make_doc(BATCH_DOC_SIZES[index % BATCH_DOC_SIZES.len()])) + .collect(); + let total_bytes: usize = docs.iter().map(Bytes::len).sum(); + group.throughput(Throughput::Bytes(total_bytes as u64)); + + group.bench_function("default", |bencher| { + bencher.iter(|| { + let mut builder = DocBatchV2Builder::default(); + let mut doc_uid_generator = DocUidGenerator::default(); + for doc in &docs { + builder.add_doc(doc_uid_generator.next_doc_uid(), doc); + } + black_box(builder.build()) + }); + }); + + group.bench_function("preallocated", |bencher| { + bencher.iter(|| { + let mut builder = DocBatchV2Builder::with_capacity(total_bytes); + let mut doc_uid_generator = DocUidGenerator::default(); + for doc in &docs { + builder.add_doc(doc_uid_generator.next_doc_uid(), doc); + } + black_box(builder.build()) + }); + }); + + group.finish(); +} + +criterion_group!( + benches, + bench_encode, + bench_decode, + bench_append_batch, + bench_doc_batch_builder +); criterion_main!(benches); diff --git a/quickwit/quickwit-ingest/src/ingest_v2/mod.rs b/quickwit/quickwit-ingest/src/ingest_v2/mod.rs index 5a6c6e2c791..9b567b82fd7 100644 --- a/quickwit/quickwit-ingest/src/ingest_v2/mod.rs +++ b/quickwit/quickwit-ingest/src/ingest_v2/mod.rs @@ -153,6 +153,14 @@ pub struct DocBatchV2Builder { } impl DocBatchV2Builder { + pub fn with_capacity(doc_buffer_capacity: usize) -> Self { + Self { + doc_uids: Vec::new(), + doc_buffer: BytesMut::with_capacity(doc_buffer_capacity), + doc_lengths: Vec::new(), + } + } + /// Adds a document to the batch. pub fn add_doc(&mut self, doc_uid: DocUid, doc: &[u8]) { self.doc_uids.push(doc_uid); diff --git a/quickwit/quickwit-ingest/src/lib.rs b/quickwit/quickwit-ingest/src/lib.rs index 7529bed1ebc..f0c6d36d859 100644 --- a/quickwit/quickwit-ingest/src/lib.rs +++ b/quickwit/quickwit-ingest/src/lib.rs @@ -24,6 +24,7 @@ mod ingest_v2; mod memory_capacity; mod metrics; mod mrecordlog_async; +mod ndjson; mod notifications; mod position; mod queue; @@ -39,6 +40,7 @@ pub use ingest_api_service::{GetMemoryCapacity, GetPartitionId, IngestApiService pub use ingest_service::*; pub use ingest_v2::*; pub use memory_capacity::MemoryCapacity; +pub use ndjson::{doc_batch_v2_from_ndjson, split_ndjson_lines}; pub use position::Position; pub use queue::Queues; use quickwit_actors::{Mailbox, Universe}; diff --git a/quickwit/quickwit-ingest/src/ndjson.rs b/quickwit/quickwit-ingest/src/ndjson.rs new file mode 100644 index 00000000000..402e736a064 --- /dev/null +++ b/quickwit/quickwit-ingest/src/ndjson.rs @@ -0,0 +1,151 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use bytes::Bytes; +use quickwit_proto::ingest::DocBatchV2; +use quickwit_proto::types::DocUidGenerator; + +/// Iterates over the non-blank lines of an NDJSON body. +pub fn split_ndjson_lines(ndjson_body: &[u8]) -> impl Iterator { + let mut line_start = 0; + memchr::memchr_iter(b'\n', ndjson_body) + .chain(std::iter::once(ndjson_body.len())) + .map(move |line_end| { + let line = &ndjson_body[line_start..line_end]; + line_start = line_end + 1; + line + }) + .filter(|line| !is_empty_or_blank_line(line)) +} + +/// Builds a [`DocBatchV2`] whose doc buffer is the NDJSON body itself, without copying. The docs +/// must cover the whole buffer, so newlines and blank lines stay attached to neighboring docs. +pub fn doc_batch_v2_from_ndjson(ndjson_body: Bytes) -> Option { + let mut doc_uids = Vec::new(); + let mut doc_lengths = Vec::new(); + let mut doc_uid_generator = DocUidGenerator::default(); + let mut segment_start = 0usize; + let mut line_start = 0usize; + + for position in memchr::memchr_iter(b'\n', &ndjson_body) { + let line = &ndjson_body[line_start..position]; + if !is_empty_or_blank_line(line) { + doc_uids.push(doc_uid_generator.next_doc_uid()); + doc_lengths.push((position + 1 - segment_start) as u32); + segment_start = position + 1; + } + line_start = position + 1; + } + + let line = &ndjson_body[line_start..]; + if !is_empty_or_blank_line(line) { + doc_uids.push(doc_uid_generator.next_doc_uid()); + doc_lengths.push((ndjson_body.len() - segment_start) as u32); + segment_start = ndjson_body.len(); + } + + if doc_uids.is_empty() { + return None; + } + if segment_start < ndjson_body.len() { + let trailing_whitespace_len = ndjson_body.len() - segment_start; + let last_doc_len = doc_lengths + .last_mut() + .expect("doc lengths should not be empty"); + *last_doc_len += trailing_whitespace_len as u32; + } + + Some(DocBatchV2 { + doc_uids, + doc_buffer: ndjson_body, + doc_lengths, + }) +} + +#[inline] +fn is_empty_or_blank_line(line: &[u8]) -> bool { + line.is_empty() || line.iter().all(|ch| ch.is_ascii_whitespace()) +} + +#[cfg(test)] +mod tests { + use std::str; + + use super::*; + + #[test] + fn test_split_ndjson_lines() { + let test_cases = [ + // an empty line is inserted before the metadata action and the doc + (&b"\n{ \"create\" : { \"_index\" : \"my-index-1\", \"_id\" : \"1\"} }\n{\"id\": 1, \"message\": \"push\"}"[..], 2), + // a blank line is inserted before the metadata action and the doc + (&b" \n{ \"create\" : { \"_index\" : \"my-index-1\", \"_id\" : \"1\"} }\n{\"id\": 1, \"message\": \"push\"}"[..], 2), + // an empty line is inserted after the metadata action and before the doc + (&b"{ \"create\" : { \"_index\" : \"my-index-1\", \"_id\" : \"1\"} }\n\n{\"id\": 1, \"message\": \"push\"}"[..], 2), + // a blank line is inserted after the metadata action and before the doc + (&b"{ \"create\" : { \"_index\" : \"my-index-1\", \"_id\" : \"1\"} }\n \n{\"id\": 1, \"message\": \"push\"}"[..], 2), + ]; + for &(input, expected_count) in &test_cases { + assert_eq!(split_ndjson_lines(input).count(), expected_count); + } + let lines: Vec<&[u8]> = split_ndjson_lines(b"a\n\nbc\n \nd").collect(); + assert_eq!(lines, [&b"a"[..], b"bc", b"d"]); + + let lines: Vec<&[u8]> = split_ndjson_lines(b"a\n").collect(); + assert_eq!(lines, [&b"a"[..]]); + + assert_eq!(split_ndjson_lines(b"").count(), 0); + assert_eq!(split_ndjson_lines(b"\n \n\t").count(), 0); + } + + #[test] + fn test_doc_batch_v2_from_ndjson_zero_copy() { + let body = Bytes::from_static(b"\n {\"id\":1}\n\n{\"id\":2}\n \n"); + let doc_batch = doc_batch_v2_from_ndjson(body.clone()).unwrap(); + assert_eq!(doc_batch.num_docs(), 2); + assert_eq!(doc_batch.doc_buffer, body); + + let docs: Vec = doc_batch.docs().map(|(_doc_uid, doc)| doc).collect(); + assert_eq!(str::from_utf8(&docs[0]).unwrap(), "\n {\"id\":1}\n"); + assert_eq!(str::from_utf8(&docs[1]).unwrap(), "\n{\"id\":2}\n \n"); + } + + #[test] + fn test_doc_batch_v2_from_ndjson_no_trailing_newline() { + let body = Bytes::from_static(b"{\"id\":1}\n{\"id\":2}"); + let doc_batch = doc_batch_v2_from_ndjson(body).unwrap(); + let docs: Vec = doc_batch.docs().map(|(_doc_uid, doc)| doc).collect(); + assert_eq!(docs, [&b"{\"id\":1}\n"[..], b"{\"id\":2}"]); + } + + #[test] + fn test_doc_batch_v2_from_blank_ndjson() { + assert!(doc_batch_v2_from_ndjson(Bytes::new()).is_none()); + assert!(doc_batch_v2_from_ndjson(Bytes::from_static(b"\n \n\t")).is_none()); + } + + #[test] + fn test_doc_batch_v2_from_ndjson_matches_split_ndjson_lines() { + let body = Bytes::from_static(b" \n{\"a\":1}\n\n \t\n{\"b\":2}\n{\"c\":3}\n\n"); + let doc_batch = doc_batch_v2_from_ndjson(body.clone()).unwrap(); + let docs: Vec = doc_batch.docs().map(|(_doc_uid, doc)| doc).collect(); + let lines: Vec<&[u8]> = split_ndjson_lines(&body).collect(); + assert_eq!(docs.len(), lines.len()); + for (doc, line) in docs.iter().zip(&lines) { + assert_eq!(doc.trim_ascii(), *line); + } + let doc_uids: Vec<_> = doc_batch.docs().map(|(doc_uid, _doc)| doc_uid).collect(); + assert!(doc_uids.windows(2).all(|pair| pair[0] < pair[1])); + } +} diff --git a/quickwit/quickwit-serve/src/elasticsearch_api/bulk.rs b/quickwit/quickwit-serve/src/elasticsearch_api/bulk.rs index c39e3072355..a7574c7a369 100644 --- a/quickwit/quickwit-serve/src/elasticsearch_api/bulk.rs +++ b/quickwit/quickwit-serve/src/elasticsearch_api/bulk.rs @@ -18,6 +18,7 @@ use std::time::Instant; use bytesize::ByteSize; use quickwit_ingest::{ CommitType, DocBatchBuilder, IngestRequest, IngestService, IngestServiceClient, + split_ndjson_lines, }; use quickwit_proto::ingest::router::IngestRouterServiceClient; use quickwit_proto::types::IndexId; @@ -29,7 +30,6 @@ use crate::elasticsearch_api::filter::{elastic_bulk_filter, elastic_index_bulk_f use crate::elasticsearch_api::make_elastic_api_response; use crate::elasticsearch_api::model::{BulkAction, ElasticBulkOptions, ElasticsearchError}; use crate::format::extract_format_from_qs; -use crate::ingest_api::lines; use crate::rest::recover_fn; use crate::{Body, with_arg}; @@ -111,7 +111,7 @@ async fn elastic_ingest_bulk( } let now = Instant::now(); let mut doc_batch_builders = HashMap::new(); - let mut lines = lines(&body.content).enumerate(); + let mut lines = split_ndjson_lines(&body.content).enumerate(); while let Some((line_number, line)) = lines.next() { let action = serde_json::from_slice::(line).map_err(|error| { diff --git a/quickwit/quickwit-serve/src/elasticsearch_api/bulk_v2.rs b/quickwit/quickwit-serve/src/elasticsearch_api/bulk_v2.rs index 18159866f41..a101cfea321 100644 --- a/quickwit/quickwit-serve/src/elasticsearch_api/bulk_v2.rs +++ b/quickwit/quickwit-serve/src/elasticsearch_api/bulk_v2.rs @@ -17,7 +17,7 @@ use std::time::Instant; use quickwit_common::rate_limited_error; use quickwit_config::{INGEST_V2_SOURCE_ID, validate_identifier}; -use quickwit_ingest::IngestRequestV2Builder; +use quickwit_ingest::{IngestRequestV2Builder, split_ndjson_lines}; use quickwit_proto::ingest::CommitTypeV2; use quickwit_proto::ingest::router::{ IngestFailureReason, IngestResponseV2, IngestRouterService, IngestRouterServiceClient, @@ -29,7 +29,6 @@ use warp::hyper::StatusCode; use super::model::ElasticException; use crate::Body; use crate::elasticsearch_api::model::{BulkAction, ElasticBulkOptions, ElasticsearchError}; -use crate::ingest_api::lines; #[derive(Debug, Default, Serialize, Deserialize)] pub(crate) struct ElasticBulkResponse { @@ -88,7 +87,7 @@ pub(crate) async fn elastic_bulk_ingest_v2( ) -> Result { let now = Instant::now(); let mut ingest_request_builder = IngestRequestV2Builder::default(); - let mut lines = lines(&body.content).enumerate(); + let mut lines = split_ndjson_lines(&body.content).enumerate(); let mut per_subrequest_doc_handles: HashMap> = HashMap::new(); let mut action_count = 0; let mut invalid_index_id_items = Vec::new(); diff --git a/quickwit/quickwit-serve/src/ingest_api/mod.rs b/quickwit/quickwit-serve/src/ingest_api/mod.rs index 940313fbf6a..3207d7dcc6d 100644 --- a/quickwit/quickwit-serve/src/ingest_api/mod.rs +++ b/quickwit/quickwit-serve/src/ingest_api/mod.rs @@ -16,7 +16,7 @@ mod response; mod rest_handler; pub use response::{RestIngestResponse, RestParseFailure}; +pub(crate) use rest_handler::ingest_api_handlers; #[cfg(test)] pub(crate) use rest_handler::tests::setup_ingest_v1_service; pub use rest_handler::{IngestApi, IngestApiSchemas}; -pub(crate) use rest_handler::{ingest_api_handlers, lines}; diff --git a/quickwit/quickwit-serve/src/ingest_api/rest_handler.rs b/quickwit/quickwit-serve/src/ingest_api/rest_handler.rs index 9040c8fd700..ea91987547f 100644 --- a/quickwit/quickwit-serve/src/ingest_api/rest_handler.rs +++ b/quickwit/quickwit-serve/src/ingest_api/rest_handler.rs @@ -12,17 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. -use bytes::{Buf, Bytes}; +use bytes::Buf; use quickwit_config::{INGEST_V2_SOURCE_ID, IngestApiConfig, validate_identifier}; use quickwit_ingest::{ - CommitType, DocBatchBuilder, DocBatchV2Builder, FetchResponse, IngestRequest, IngestService, - IngestServiceClient, IngestServiceError, TailRequest, + CommitType, DocBatchBuilder, FetchResponse, IngestRequest, IngestService, IngestServiceClient, + IngestServiceError, TailRequest, doc_batch_v2_from_ndjson, split_ndjson_lines, }; use quickwit_proto::ingest::CommitTypeV2; use quickwit_proto::ingest::router::{ IngestRequestV2, IngestRouterService, IngestRouterServiceClient, IngestSubrequest, }; -use quickwit_proto::types::{DocUidGenerator, IndexId}; +use quickwit_proto::types::IndexId; use serde::Deserialize; use warp::{Filter, Rejection}; @@ -176,7 +176,7 @@ async fn ingest_v1( // The size of the body should be an upper bound of the size of the batch. The removal of the // end of line character for each doc compensates the addition of the `DocCommand` header. let mut doc_batch_builder = DocBatchBuilder::with_capacity(index_id, body.content.remaining()); - for line in lines(&body.content) { + for line in split_ndjson_lines(&body.content) { doc_batch_builder.ingest_doc(line); } let ingest_req = IngestRequest { @@ -193,14 +193,7 @@ async fn ingest_v2( ingest_options: IngestOptions, ingest_router: IngestRouterServiceClient, ) -> Result { - let mut doc_batch_builder = DocBatchV2Builder::default(); - let mut doc_uid_generator = DocUidGenerator::default(); - - for doc in lines(&body.content) { - doc_batch_builder.add_doc(doc_uid_generator.next_doc_uid(), doc); - } - drop(body); - let doc_batch_opt = doc_batch_builder.build(); + let doc_batch_opt = doc_batch_v2_from_ndjson(body.content); let Some(doc_batch) = doc_batch_opt else { let response = RestIngestResponse::default(); @@ -274,22 +267,10 @@ async fn tail_endpoint( Ok(fetch_response) } -pub(crate) fn lines(body: &Bytes) -> impl Iterator { - body.split(|byte| byte == &b'\n') - .filter(|line| !is_empty_or_blank_line(line)) -} - -#[inline] -fn is_empty_or_blank_line(line: &[u8]) -> bool { - line.is_empty() || line.iter().all(|ch| ch.is_ascii_whitespace()) -} - #[cfg(test)] pub(crate) mod tests { - use std::str; use std::time::Duration; - use bytes::Bytes; use quickwit_actors::{Mailbox, Universe}; use quickwit_config::IngestApiConfig; use quickwit_ingest::{ @@ -299,25 +280,6 @@ pub(crate) mod tests { use quickwit_proto::ingest::router::IngestRouterServiceClient; use super::{RestIngestResponse, ingest_api_handlers}; - use crate::ingest_api::lines; - - #[test] - fn test_process_lines() { - let test_cases = [ - // an empty line is inserted before the metadata action and the doc - (&b"\n{ \"create\" : { \"_index\" : \"my-index-1\", \"_id\" : \"1\"} }\n{\"id\": 1, \"message\": \"push\"}"[..], 2), - // a blank line is inserted before the metadata action and the doc - (&b" \n{ \"create\" : { \"_index\" : \"my-index-1\", \"_id\" : \"1\"} }\n{\"id\": 1, \"message\": \"push\"}"[..], 2), - // an empty line is inserted after the metadata action and before the doc - (&b"{ \"create\" : { \"_index\" : \"my-index-1\", \"_id\" : \"1\"} }\n\n{\"id\": 1, \"message\": \"push\"}"[..], 2), - // a blank line is inserted after the metadata action and before the doc - (&b"{ \"create\" : { \"_index\" : \"my-index-1\", \"_id\" : \"1\"} }\n \n{\"id\": 1, \"message\": \"push\"}"[..], 2), - ]; - - for &(input, expected_count) in &test_cases { - assert_eq!(lines(&Bytes::from(input)).count(), expected_count); - } - } pub(crate) async fn setup_ingest_v1_service( queues: &[&str], diff --git a/quickwit/rest-api-tests/scenarii/aggregations/0002-doc-len.yaml b/quickwit/rest-api-tests/scenarii/aggregations/0002-doc-len.yaml index 08b437015b9..3a89f925e75 100644 --- a/quickwit/rest-api-tests/scenarii/aggregations/0002-doc-len.yaml +++ b/quickwit/rest-api-tests/scenarii/aggregations/0002-doc-len.yaml @@ -12,7 +12,8 @@ json: expected: aggregations: doc_len: - value: 952.0 + # Native ingest retains the NDJSON separators: 952 JSON bytes + 10 newlines. + value: 962.0 --- # Test doc len isn't shown when querying documents method: [GET]