Line data Source code
1 : use std::{sync::Arc, time::SystemTime};
2 :
3 : use anyhow::Context;
4 : use bytes::{buf::Writer, BufMut, BytesMut};
5 : use chrono::{Datelike, Timelike};
6 : use futures::{Stream, StreamExt};
7 : use parquet::{
8 : basic::Compression,
9 : file::{
10 : metadata::RowGroupMetaDataPtr,
11 : properties::{WriterProperties, WriterPropertiesPtr, DEFAULT_PAGE_SIZE},
12 : writer::SerializedFileWriter,
13 : },
14 : record::RecordWriter,
15 : };
16 : use remote_storage::{GenericRemoteStorage, RemotePath, TimeoutOrCancel};
17 : use tokio::{sync::mpsc, time};
18 : use tokio_util::sync::CancellationToken;
19 : use tracing::{debug, info, Span};
20 : use utils::backoff;
21 :
22 : use crate::config::{remote_storage_from_toml, OptRemoteStorageConfig};
23 :
24 : use super::{RequestMonitoring, LOG_CHAN};
25 :
26 12 : #[derive(clap::Args, Clone, Debug)]
27 : pub struct ParquetUploadArgs {
28 : /// Storage location to upload the parquet files to.
29 : /// Encoded as toml (same format as pageservers), eg
30 : /// `{bucket_name='the-bucket',bucket_region='us-east-1',prefix_in_bucket='proxy',endpoint='http://minio:9000'}`
31 : #[clap(long, default_value = "{}", value_parser = remote_storage_from_toml)]
32 0 : parquet_upload_remote_storage: OptRemoteStorageConfig,
33 :
34 : /// How many rows to include in a row group
35 6 : #[clap(long, default_value_t = 8192)]
36 0 : parquet_upload_row_group_size: usize,
37 :
38 : /// How large each column page should be in bytes
39 6 : #[clap(long, default_value_t = DEFAULT_PAGE_SIZE)]
40 0 : parquet_upload_page_size: usize,
41 :
42 : /// How large the total parquet file should be in bytes
43 6 : #[clap(long, default_value_t = 100_000_000)]
44 0 : parquet_upload_size: i64,
45 :
46 : /// How long to wait before forcing a file upload
47 : #[clap(long, default_value = "20m", value_parser = humantime::parse_duration)]
48 0 : parquet_upload_maximum_duration: tokio::time::Duration,
49 :
50 : /// What level of compression to use
51 6 : #[clap(long, default_value_t = Compression::UNCOMPRESSED)]
52 0 : parquet_upload_compression: Compression,
53 : }
54 :
55 : // Occasional network issues and such can cause remote operations to fail, and
56 : // that's expected. If a upload fails, we log it at info-level, and retry.
57 : // But after FAILED_UPLOAD_WARN_THRESHOLD retries, we start to log it at WARN
58 : // level instead, as repeated failures can mean a more serious problem. If it
59 : // fails more than FAILED_UPLOAD_RETRIES times, we give up
60 : pub const FAILED_UPLOAD_WARN_THRESHOLD: u32 = 3;
61 : pub const FAILED_UPLOAD_MAX_RETRIES: u32 = 10;
62 :
63 : // the parquet crate leaves a lot to be desired...
64 : // what follows is an attempt to write parquet files with minimal allocs.
65 : // complication: parquet is a columnar format, while we want to write in as rows.
66 : // design:
67 : // * we batch up to 1024 rows, then flush them into a 'row group'
68 : // * after each rowgroup write, we check the length of the file and upload to s3 if large enough
69 :
70 6688010 : #[derive(parquet_derive::ParquetRecordWriter)]
71 : pub struct RequestData {
72 : region: &'static str,
73 : protocol: &'static str,
74 : /// Must be UTC. The derive macro doesn't like the timezones
75 : timestamp: chrono::NaiveDateTime,
76 : session_id: uuid::Uuid,
77 : peer_addr: String,
78 : username: Option<String>,
79 : application_name: Option<String>,
80 : endpoint_id: Option<String>,
81 : database: Option<String>,
82 : project: Option<String>,
83 : branch: Option<String>,
84 : auth_method: Option<&'static str>,
85 : error: Option<&'static str>,
86 : /// Success is counted if we form a HTTP response with sql rows inside
87 : /// Or if we make it to proxy_pass
88 : success: bool,
89 : /// Indicates if the cplane started the new compute node for this request.
90 : cold_start_info: &'static str,
91 : /// Tracks time from session start (HTTP request/libpq TCP handshake)
92 : /// Through to success/failure
93 : duration_us: u64,
94 : }
95 :
96 : impl From<&RequestMonitoring> for RequestData {
97 0 : fn from(value: &RequestMonitoring) -> Self {
98 0 : Self {
99 0 : session_id: value.session_id,
100 0 : peer_addr: value.peer_addr.to_string(),
101 0 : timestamp: value.first_packet.naive_utc(),
102 0 : username: value.user.as_deref().map(String::from),
103 0 : application_name: value.application.as_deref().map(String::from),
104 0 : endpoint_id: value.endpoint_id.as_deref().map(String::from),
105 0 : database: value.dbname.as_deref().map(String::from),
106 0 : project: value.project.as_deref().map(String::from),
107 0 : branch: value.branch.as_deref().map(String::from),
108 0 : auth_method: value.auth_method.as_ref().map(|x| match x {
109 0 : super::AuthMethod::Web => "web",
110 0 : super::AuthMethod::ScramSha256 => "scram_sha_256",
111 0 : super::AuthMethod::ScramSha256Plus => "scram_sha_256_plus",
112 0 : super::AuthMethod::Cleartext => "cleartext",
113 0 : }),
114 0 : protocol: value.protocol.as_str(),
115 0 : region: value.region,
116 0 : error: value.error_kind.as_ref().map(|e| e.to_metric_label()),
117 0 : success: value.success,
118 0 : cold_start_info: value.cold_start_info.as_str(),
119 0 : duration_us: SystemTime::from(value.first_packet)
120 0 : .elapsed()
121 0 : .unwrap_or_default()
122 0 : .as_micros() as u64, // 584 millenia... good enough
123 0 : }
124 0 : }
125 : }
126 :
127 : /// Parquet request context worker
128 : ///
129 : /// It listened on a channel for all completed requests, extracts the data and writes it into a parquet file,
130 : /// then uploads a completed batch to S3
131 0 : pub async fn worker(
132 0 : cancellation_token: CancellationToken,
133 0 : config: ParquetUploadArgs,
134 0 : ) -> anyhow::Result<()> {
135 0 : let Some(remote_storage_config) = config.parquet_upload_remote_storage else {
136 0 : tracing::warn!("parquet request upload: no s3 bucket configured");
137 0 : return Ok(());
138 : };
139 :
140 0 : let (tx, mut rx) = mpsc::unbounded_channel();
141 0 : LOG_CHAN.set(tx.downgrade()).unwrap();
142 0 :
143 0 : // setup row stream that will close on cancellation
144 0 : tokio::spawn(async move {
145 0 : cancellation_token.cancelled().await;
146 : // dropping this sender will cause the channel to close only once
147 : // all the remaining inflight requests have been completed.
148 0 : drop(tx);
149 0 : });
150 0 : let rx = futures::stream::poll_fn(move |cx| rx.poll_recv(cx));
151 0 : let rx = rx.map(RequestData::from);
152 :
153 0 : let storage =
154 0 : GenericRemoteStorage::from_config(&remote_storage_config).context("remote storage init")?;
155 :
156 0 : let properties = WriterProperties::builder()
157 0 : .set_data_page_size_limit(config.parquet_upload_page_size)
158 0 : .set_compression(config.parquet_upload_compression);
159 0 :
160 0 : let parquet_config = ParquetConfig {
161 0 : propeties: Arc::new(properties.build()),
162 0 : rows_per_group: config.parquet_upload_row_group_size,
163 0 : file_size: config.parquet_upload_size,
164 0 : max_duration: config.parquet_upload_maximum_duration,
165 0 :
166 0 : #[cfg(any(test, feature = "testing"))]
167 0 : test_remote_failures: 0,
168 0 : };
169 0 :
170 0 : worker_inner(storage, rx, parquet_config).await
171 0 : }
172 :
173 : struct ParquetConfig {
174 : propeties: WriterPropertiesPtr,
175 : rows_per_group: usize,
176 : file_size: i64,
177 :
178 : max_duration: tokio::time::Duration,
179 :
180 : #[cfg(any(test, feature = "testing"))]
181 : test_remote_failures: u64,
182 : }
183 :
184 10 : async fn worker_inner(
185 10 : storage: GenericRemoteStorage,
186 10 : rx: impl Stream<Item = RequestData>,
187 10 : config: ParquetConfig,
188 10 : ) -> anyhow::Result<()> {
189 : #[cfg(any(test, feature = "testing"))]
190 10 : let storage = if config.test_remote_failures > 0 {
191 4 : GenericRemoteStorage::unreliable_wrapper(storage, config.test_remote_failures)
192 : } else {
193 6 : storage
194 : };
195 :
196 10 : let mut rx = std::pin::pin!(rx);
197 10 :
198 10 : let mut rows = Vec::with_capacity(config.rows_per_group);
199 :
200 10 : let schema = rows.as_slice().schema()?;
201 10 : let buffer = BytesMut::new();
202 10 : let w = buffer.writer();
203 10 : let mut w = SerializedFileWriter::new(w, schema.clone(), config.propeties.clone())?;
204 :
205 10 : let mut last_upload = time::Instant::now();
206 10 :
207 10 : let mut len = 0;
208 418010 : while let Some(row) = rx.next().await {
209 418000 : rows.push(row);
210 418000 : let force = last_upload.elapsed() > config.max_duration;
211 418000 : if rows.len() == config.rows_per_group || force {
212 210 : let rg_meta;
213 210 : (rows, w, rg_meta) = flush_rows(rows, w).await?;
214 210 : len += rg_meta.compressed_size();
215 417790 : }
216 418000 : if len > config.file_size || force {
217 56 : last_upload = time::Instant::now();
218 233 : let file = upload_parquet(w, len, &storage).await?;
219 56 : w = SerializedFileWriter::new(file, schema.clone(), config.propeties.clone())?;
220 56 : len = 0;
221 417944 : }
222 : }
223 :
224 10 : if !rows.is_empty() {
225 2 : let rg_meta;
226 2 : (_, w, rg_meta) = flush_rows(rows, w).await?;
227 2 : len += rg_meta.compressed_size();
228 8 : }
229 :
230 10 : if !w.flushed_row_groups().is_empty() {
231 24 : let _: Writer<BytesMut> = upload_parquet(w, len, &storage).await?;
232 4 : }
233 :
234 10 : Ok(())
235 10 : }
236 :
237 212 : async fn flush_rows<W>(
238 212 : rows: Vec<RequestData>,
239 212 : mut w: SerializedFileWriter<W>,
240 212 : ) -> anyhow::Result<(
241 212 : Vec<RequestData>,
242 212 : SerializedFileWriter<W>,
243 212 : RowGroupMetaDataPtr,
244 212 : )>
245 212 : where
246 212 : W: std::io::Write + Send + 'static,
247 212 : {
248 212 : let span = Span::current();
249 212 : let (mut rows, w, rg_meta) = tokio::task::spawn_blocking(move || {
250 212 : let _enter = span.enter();
251 :
252 212 : let mut rg = w.next_row_group()?;
253 212 : rows.as_slice().write_to_row_group(&mut rg)?;
254 212 : let rg_meta = rg.close()?;
255 :
256 212 : let size = rg_meta.compressed_size();
257 212 : let compression = rg_meta.compressed_size() as f64 / rg_meta.total_byte_size() as f64;
258 212 :
259 212 : debug!(size, compression, "flushed row group to parquet file");
260 :
261 212 : Ok::<_, parquet::errors::ParquetError>((rows, w, rg_meta))
262 212 : })
263 212 : .await
264 212 : .unwrap()?;
265 :
266 212 : rows.clear();
267 212 : Ok((rows, w, rg_meta))
268 212 : }
269 :
270 62 : async fn upload_parquet(
271 62 : w: SerializedFileWriter<Writer<BytesMut>>,
272 62 : len: i64,
273 62 : storage: &GenericRemoteStorage,
274 62 : ) -> anyhow::Result<Writer<BytesMut>> {
275 62 : let len_uncompressed = w
276 62 : .flushed_row_groups()
277 62 : .iter()
278 212 : .map(|rg| rg.total_byte_size())
279 62 : .sum::<i64>();
280 :
281 : // I don't know how compute intensive this is, although it probably isn't much... better be safe than sorry.
282 : // finish method only available on the fork: https://github.com/apache/arrow-rs/issues/5253
283 62 : let (writer, metadata) = tokio::task::spawn_blocking(move || w.finish())
284 62 : .await
285 62 : .unwrap()?;
286 :
287 62 : let mut buffer = writer.into_inner();
288 62 : let data = buffer.split().freeze();
289 62 :
290 62 : let compression = len as f64 / len_uncompressed as f64;
291 62 : let size = data.len();
292 62 : let now = chrono::Utc::now();
293 62 : let id = uuid::Uuid::new_v7(uuid::Timestamp::from_unix(
294 62 : uuid::NoContext,
295 62 : // we won't be running this in 1970. this cast is ok
296 62 : now.timestamp() as u64,
297 62 : now.timestamp_subsec_nanos(),
298 62 : ));
299 62 :
300 62 : info!(
301 0 : %id,
302 0 : rows = metadata.num_rows,
303 0 : size, compression, "uploading request parquet file"
304 0 : );
305 :
306 62 : let year = now.year();
307 62 : let month = now.month();
308 62 : let day = now.day();
309 62 : let hour = now.hour();
310 : // segment files by time for S3 performance
311 62 : let path = RemotePath::from_string(&format!(
312 62 : "{year:04}/{month:02}/{day:02}/{hour:02}/requests_{id}.parquet"
313 62 : ))?;
314 62 : let cancel = CancellationToken::new();
315 62 : backoff::retry(
316 86 : || async {
317 86 : let stream = futures::stream::once(futures::future::ready(Ok(data.clone())));
318 86 : storage
319 86 : .upload(stream, data.len(), &path, None, &cancel)
320 195 : .await
321 172 : },
322 62 : TimeoutOrCancel::caused_by_cancel,
323 62 : FAILED_UPLOAD_WARN_THRESHOLD,
324 62 : FAILED_UPLOAD_MAX_RETRIES,
325 62 : "request_data_upload",
326 62 : // we don't want cancellation to interrupt here, so we make a dummy cancel token
327 62 : &cancel,
328 62 : )
329 195 : .await
330 62 : .ok_or_else(|| anyhow::Error::new(TimeoutOrCancel::Cancel))
331 62 : .and_then(|x| x)
332 62 : .context("request_data_upload")?;
333 :
334 62 : Ok(buffer.writer())
335 62 : }
336 :
337 : #[cfg(test)]
338 : mod tests {
339 : use std::{net::Ipv4Addr, num::NonZeroUsize, sync::Arc};
340 :
341 : use camino::Utf8Path;
342 : use clap::Parser;
343 : use futures::{Stream, StreamExt};
344 : use itertools::Itertools;
345 : use parquet::{
346 : basic::{Compression, ZstdLevel},
347 : file::{
348 : properties::{WriterProperties, DEFAULT_PAGE_SIZE},
349 : reader::FileReader,
350 : serialized_reader::SerializedFileReader,
351 : },
352 : };
353 : use rand::{rngs::StdRng, Rng, SeedableRng};
354 : use remote_storage::{
355 : GenericRemoteStorage, RemoteStorageConfig, RemoteStorageKind, S3Config,
356 : DEFAULT_MAX_KEYS_PER_LIST_RESPONSE, DEFAULT_REMOTE_STORAGE_S3_CONCURRENCY_LIMIT,
357 : };
358 : use tokio::{sync::mpsc, time};
359 : use walkdir::WalkDir;
360 :
361 : use super::{worker_inner, ParquetConfig, ParquetUploadArgs, RequestData};
362 :
363 4 : #[derive(Parser)]
364 : struct ProxyCliArgs {
365 : #[clap(flatten)]
366 : parquet_upload: ParquetUploadArgs,
367 : }
368 :
369 : #[test]
370 2 : fn default_parser() {
371 2 : let ProxyCliArgs { parquet_upload } = ProxyCliArgs::parse_from(["proxy"]);
372 2 : assert_eq!(parquet_upload.parquet_upload_remote_storage, None);
373 2 : assert_eq!(parquet_upload.parquet_upload_row_group_size, 8192);
374 2 : assert_eq!(parquet_upload.parquet_upload_page_size, DEFAULT_PAGE_SIZE);
375 2 : assert_eq!(parquet_upload.parquet_upload_size, 100_000_000);
376 2 : assert_eq!(
377 2 : parquet_upload.parquet_upload_maximum_duration,
378 2 : time::Duration::from_secs(20 * 60)
379 2 : );
380 2 : assert_eq!(
381 2 : parquet_upload.parquet_upload_compression,
382 2 : Compression::UNCOMPRESSED
383 2 : );
384 2 : }
385 :
386 : #[test]
387 2 : fn full_parser() {
388 2 : let ProxyCliArgs { parquet_upload } = ProxyCliArgs::parse_from([
389 2 : "proxy",
390 2 : "--parquet-upload-remote-storage",
391 2 : "{bucket_name='default',prefix_in_bucket='proxy/',bucket_region='us-east-1',endpoint='http://minio:9000'}",
392 2 : "--parquet-upload-row-group-size",
393 2 : "100",
394 2 : "--parquet-upload-page-size",
395 2 : "10000",
396 2 : "--parquet-upload-size",
397 2 : "10000000",
398 2 : "--parquet-upload-maximum-duration",
399 2 : "10m",
400 2 : "--parquet-upload-compression",
401 2 : "zstd(5)",
402 2 : ]);
403 2 : assert_eq!(
404 2 : parquet_upload.parquet_upload_remote_storage,
405 2 : Some(RemoteStorageConfig {
406 2 : storage: RemoteStorageKind::AwsS3(S3Config {
407 2 : bucket_name: "default".into(),
408 2 : bucket_region: "us-east-1".into(),
409 2 : prefix_in_bucket: Some("proxy/".into()),
410 2 : endpoint: Some("http://minio:9000".into()),
411 2 : concurrency_limit: NonZeroUsize::new(
412 2 : DEFAULT_REMOTE_STORAGE_S3_CONCURRENCY_LIMIT
413 2 : )
414 2 : .unwrap(),
415 2 : max_keys_per_list_response: DEFAULT_MAX_KEYS_PER_LIST_RESPONSE,
416 2 : }),
417 2 : timeout: RemoteStorageConfig::DEFAULT_TIMEOUT,
418 2 : })
419 2 : );
420 2 : assert_eq!(parquet_upload.parquet_upload_row_group_size, 100);
421 2 : assert_eq!(parquet_upload.parquet_upload_page_size, 10000);
422 2 : assert_eq!(parquet_upload.parquet_upload_size, 10_000_000);
423 2 : assert_eq!(
424 2 : parquet_upload.parquet_upload_maximum_duration,
425 2 : time::Duration::from_secs(10 * 60)
426 2 : );
427 2 : assert_eq!(
428 2 : parquet_upload.parquet_upload_compression,
429 2 : Compression::ZSTD(ZstdLevel::try_new(5).unwrap())
430 2 : );
431 2 : }
432 :
433 418000 : fn generate_request_data(rng: &mut impl Rng) -> RequestData {
434 418000 : RequestData {
435 418000 : session_id: uuid::Builder::from_random_bytes(rng.gen()).into_uuid(),
436 418000 : peer_addr: Ipv4Addr::from(rng.gen::<[u8; 4]>()).to_string(),
437 418000 : timestamp: chrono::NaiveDateTime::from_timestamp_millis(
438 418000 : rng.gen_range(1703862754..1803862754),
439 418000 : )
440 418000 : .unwrap(),
441 418000 : application_name: Some("test".to_owned()),
442 418000 : username: Some(hex::encode(rng.gen::<[u8; 4]>())),
443 418000 : endpoint_id: Some(hex::encode(rng.gen::<[u8; 16]>())),
444 418000 : database: Some(hex::encode(rng.gen::<[u8; 16]>())),
445 418000 : project: Some(hex::encode(rng.gen::<[u8; 16]>())),
446 418000 : branch: Some(hex::encode(rng.gen::<[u8; 16]>())),
447 418000 : auth_method: None,
448 418000 : protocol: ["tcp", "ws", "http"][rng.gen_range(0..3)],
449 418000 : region: "us-east-1",
450 418000 : error: None,
451 418000 : success: rng.gen(),
452 418000 : cold_start_info: "no",
453 418000 : duration_us: rng.gen_range(0..30_000_000),
454 418000 : }
455 418000 : }
456 :
457 14 : fn random_stream(len: usize) -> impl Stream<Item = RequestData> + Unpin {
458 14 : let mut rng = StdRng::from_seed([0x39; 32]);
459 14 : futures::stream::iter(
460 418000 : std::iter::repeat_with(move || generate_request_data(&mut rng)).take(len),
461 14 : )
462 14 : }
463 :
464 10 : async fn run_test(
465 10 : tmpdir: &Utf8Path,
466 10 : config: ParquetConfig,
467 10 : rx: impl Stream<Item = RequestData>,
468 10 : ) -> Vec<(u64, usize, i64)> {
469 10 : let remote_storage_config = RemoteStorageConfig {
470 10 : storage: RemoteStorageKind::LocalFs(tmpdir.to_path_buf()),
471 10 : timeout: std::time::Duration::from_secs(120),
472 10 : };
473 10 : let storage = GenericRemoteStorage::from_config(&remote_storage_config).unwrap();
474 10 :
475 609 : worker_inner(storage, rx, config).await.unwrap();
476 10 :
477 10 : let mut files = WalkDir::new(tmpdir.as_std_path())
478 10 : .into_iter()
479 112 : .filter_map(|entry| entry.ok())
480 112 : .filter(|entry| entry.file_type().is_file())
481 62 : .map(|entry| entry.path().to_path_buf())
482 10 : .collect_vec();
483 10 : files.sort();
484 10 :
485 10 : files
486 10 : .into_iter()
487 62 : .map(|path| std::fs::File::open(tmpdir.as_std_path().join(path)).unwrap())
488 62 : .map(|file| {
489 62 : (
490 62 : file.metadata().unwrap(),
491 62 : SerializedFileReader::new(file).unwrap().metadata().clone(),
492 62 : )
493 62 : })
494 62 : .map(|(file_meta, parquet_meta)| {
495 62 : (
496 62 : file_meta.len(),
497 62 : parquet_meta.num_row_groups(),
498 62 : parquet_meta.file_metadata().num_rows(),
499 62 : )
500 62 : })
501 10 : .collect()
502 10 : }
503 :
504 : #[tokio::test]
505 2 : async fn verify_parquet_no_compression() {
506 2 : let tmpdir = camino_tempfile::tempdir().unwrap();
507 2 :
508 2 : let config = ParquetConfig {
509 2 : propeties: Arc::new(WriterProperties::new()),
510 2 : rows_per_group: 2_000,
511 2 : file_size: 1_000_000,
512 2 : max_duration: time::Duration::from_secs(20 * 60),
513 2 : test_remote_failures: 0,
514 2 : };
515 2 :
516 2 : let rx = random_stream(50_000);
517 123 : let file_stats = run_test(tmpdir.path(), config, rx).await;
518 2 :
519 2 : assert_eq!(
520 2 : file_stats,
521 2 : [
522 2 : (1314385, 3, 6000),
523 2 : (1314378, 3, 6000),
524 2 : (1314438, 3, 6000),
525 2 : (1314395, 3, 6000),
526 2 : (1314525, 3, 6000),
527 2 : (1314367, 3, 6000),
528 2 : (1314159, 3, 6000),
529 2 : (1314395, 3, 6000),
530 2 : (438352, 1, 2000)
531 2 : ]
532 2 : );
533 2 :
534 2 : tmpdir.close().unwrap();
535 2 : }
536 :
537 : #[tokio::test]
538 2 : async fn verify_parquet_min_compression() {
539 2 : let tmpdir = camino_tempfile::tempdir().unwrap();
540 2 :
541 2 : let config = ParquetConfig {
542 2 : propeties: Arc::new(
543 2 : WriterProperties::builder()
544 2 : .set_compression(parquet::basic::Compression::ZSTD(ZstdLevel::default()))
545 2 : .build(),
546 2 : ),
547 2 : rows_per_group: 2_000,
548 2 : file_size: 1_000_000,
549 2 : max_duration: time::Duration::from_secs(20 * 60),
550 2 : test_remote_failures: 0,
551 2 : };
552 2 :
553 2 : let rx = random_stream(50_000);
554 92 : let file_stats = run_test(tmpdir.path(), config, rx).await;
555 2 :
556 2 : // with compression, there are fewer files with more rows per file
557 2 : assert_eq!(
558 2 : file_stats,
559 2 : [
560 2 : (1220633, 5, 10000),
561 2 : (1226783, 5, 10000),
562 2 : (1228577, 5, 10000),
563 2 : (1227939, 5, 10000),
564 2 : (1219217, 5, 10000)
565 2 : ]
566 2 : );
567 2 :
568 2 : tmpdir.close().unwrap();
569 2 : }
570 :
571 : #[tokio::test]
572 2 : async fn verify_parquet_strong_compression() {
573 2 : let tmpdir = camino_tempfile::tempdir().unwrap();
574 2 :
575 2 : let config = ParquetConfig {
576 2 : propeties: Arc::new(
577 2 : WriterProperties::builder()
578 2 : .set_compression(parquet::basic::Compression::ZSTD(
579 2 : ZstdLevel::try_new(10).unwrap(),
580 2 : ))
581 2 : .build(),
582 2 : ),
583 2 : rows_per_group: 2_000,
584 2 : file_size: 1_000_000,
585 2 : max_duration: time::Duration::from_secs(20 * 60),
586 2 : test_remote_failures: 0,
587 2 : };
588 2 :
589 2 : let rx = random_stream(50_000);
590 92 : let file_stats = run_test(tmpdir.path(), config, rx).await;
591 2 :
592 2 : // with strong compression, the files are smaller
593 2 : assert_eq!(
594 2 : file_stats,
595 2 : [
596 2 : (1206280, 5, 10000),
597 2 : (1206011, 5, 10000),
598 2 : (1206304, 5, 10000),
599 2 : (1206292, 5, 10000),
600 2 : (1206547, 5, 10000)
601 2 : ]
602 2 : );
603 2 :
604 2 : tmpdir.close().unwrap();
605 2 : }
606 :
607 : #[tokio::test]
608 2 : async fn verify_parquet_unreliable_upload() {
609 2 : let tmpdir = camino_tempfile::tempdir().unwrap();
610 2 :
611 2 : let config = ParquetConfig {
612 2 : propeties: Arc::new(WriterProperties::new()),
613 2 : rows_per_group: 2_000,
614 2 : file_size: 1_000_000,
615 2 : max_duration: time::Duration::from_secs(20 * 60),
616 2 : test_remote_failures: 2,
617 2 : };
618 2 :
619 2 : let rx = random_stream(50_000);
620 124 : let file_stats = run_test(tmpdir.path(), config, rx).await;
621 2 :
622 2 : assert_eq!(
623 2 : file_stats,
624 2 : [
625 2 : (1314385, 3, 6000),
626 2 : (1314378, 3, 6000),
627 2 : (1314438, 3, 6000),
628 2 : (1314395, 3, 6000),
629 2 : (1314525, 3, 6000),
630 2 : (1314367, 3, 6000),
631 2 : (1314159, 3, 6000),
632 2 : (1314395, 3, 6000),
633 2 : (438352, 1, 2000)
634 2 : ]
635 2 : );
636 2 :
637 2 : tmpdir.close().unwrap();
638 2 : }
639 :
640 : #[tokio::test(start_paused = true)]
641 2 : async fn verify_parquet_regular_upload() {
642 2 : let tmpdir = camino_tempfile::tempdir().unwrap();
643 2 :
644 2 : let config = ParquetConfig {
645 2 : propeties: Arc::new(WriterProperties::new()),
646 2 : rows_per_group: 2_000,
647 2 : file_size: 1_000_000,
648 2 : max_duration: time::Duration::from_secs(60),
649 2 : test_remote_failures: 2,
650 2 : };
651 2 :
652 2 : let (tx, mut rx) = mpsc::unbounded_channel();
653 2 :
654 2 : tokio::spawn(async move {
655 8 : for _ in 0..3 {
656 6 : let mut s = random_stream(3000);
657 18006 : while let Some(r) = s.next().await {
658 18000 : tx.send(r).unwrap();
659 18000 : }
660 6 : time::sleep(time::Duration::from_secs(70)).await
661 2 : }
662 2 : });
663 2 :
664 18142 : let rx = futures::stream::poll_fn(move |cx| rx.poll_recv(cx));
665 178 : let file_stats = run_test(tmpdir.path(), config, rx).await;
666 2 :
667 2 : // files are smaller than the size threshold, but they took too long to fill so were flushed early
668 2 : assert_eq!(
669 2 : file_stats,
670 2 : [(658823, 2, 3001), (658537, 2, 3000), (658333, 2, 2999)]
671 2 : );
672 2 :
673 2 : tmpdir.close().unwrap();
674 2 : }
675 : }
|
