diff --git a/src/config/src/utils/parquet.rs b/src/config/src/utils/parquet.rs index c4bda5da00..821b2a9a44 100644 --- a/src/config/src/utils/parquet.rs +++ b/src/config/src/utils/parquet.rs @@ -34,6 +34,7 @@ use parquet::{ basic::{Compression, Encoding}, file::{metadata::KeyValue, properties::WriterProperties}, }; +use serde::{Deserialize, Serialize}; use vortex::{ VortexSessionDefault, array::{ArrayRef, VortexSessionExecute}, @@ -44,7 +45,31 @@ use vortex::{ session::VortexSession, }; -use crate::{FileFormat, config::*, ider, meta::stream::FileMeta}; +use crate::{FileFormat, config::*, ider, meta::stream::FileMeta, utils::json}; + +/// Key of the vortex metadata segment carrying the o2 [`FileMeta`]. +pub const VORTEX_FILE_META_KEY: &str = "o2_file_meta"; + +/// Same four fields [`new_parquet_writer`] writes into the parquet footer. +#[derive(Debug, Default, Serialize, Deserialize)] +#[serde(default)] +struct VortexFileMeta { + min_ts: i64, + max_ts: i64, + records: i64, + original_size: i64, +} + +/// Encode `metadata` for the [`VORTEX_FILE_META_KEY`] segment. +pub fn encode_vortex_file_meta(metadata: &FileMeta) -> Vec { + json::to_vec(&VortexFileMeta { + min_ts: metadata.min_ts, + max_ts: metadata.max_ts, + records: metadata.records, + original_size: metadata.original_size, + }) + .expect("file meta is always serializable") +} pub fn new_parquet_writer<'a>( buf: &'a mut Vec, @@ -496,6 +521,31 @@ mod tests { assert_eq!(read_metadata.original_size, metadata.original_size); } + #[test] + fn test_encode_decode_vortex_file_meta() { + let metadata = FileMeta { + min_ts: -1, + max_ts: i64::MAX, + records: 7, + original_size: 8, + compressed_size: 9, + index_size: 10, + bloom_ver: 11, + flattened: true, + }; + let decoded: VortexFileMeta = + json::from_slice(&encode_vortex_file_meta(&metadata)).unwrap(); + assert_eq!(decoded.min_ts, metadata.min_ts); + assert_eq!(decoded.max_ts, metadata.max_ts); + assert_eq!(decoded.records, metadata.records); + assert_eq!(decoded.original_size, metadata.original_size); + + // unknown and missing keys are tolerated + let decoded: VortexFileMeta = json::from_slice(br#"{"min_ts":5,"future":true}"#).unwrap(); + assert_eq!(decoded.min_ts, 5); + assert_eq!(decoded.max_ts, 0); + } + #[test] fn test_parse_file_key_columns() { let key = "files/default/logs/olympics/2022/10/03/10/6982652937134804993_1.parquet"; diff --git a/src/search/src/datafusion/merge/downsampling.rs b/src/search/src/datafusion/merge/downsampling.rs index 4587b84edc..90b0e61a22 100644 --- a/src/search/src/datafusion/merge/downsampling.rs +++ b/src/search/src/datafusion/merge/downsampling.rs @@ -190,6 +190,7 @@ async fn write_downsampled_parquet( // Finalize last file if it has data if file_meta.records > 0 { file_meta.min_ts = last_min_ts; + append_metadata(&mut writer, &file_meta)?; writer.close().await?; bufs.push(buf); file_metas.push(file_meta); @@ -403,6 +404,40 @@ mod tests { .unwrap() } + /// Every downsampled file must carry its own meta in the parquet footer, + /// including the last one. + #[tokio::test] + async fn test_write_downsampled_parquet_writes_metadata_for_every_file() { + let schema = create_test_schema(); + let batch = create_test_record_batch(); + + // one batch per file + let mut cfg = config::Config::default(); + cfg.compact.max_file_size = 1; + + let (tx, rx) = tokio::sync::mpsc::channel::(2); + tx.send(batch.clone()).await.unwrap(); + tx.send(batch).await.unwrap(); + drop(tx); + + let (bufs, file_metas) = + write_downsampled_parquet(rx, &schema, &[], &FileMeta::default(), &cfg) + .await + .unwrap(); + + assert_eq!(bufs.len(), 2); + assert_eq!(file_metas.len(), 2); + for (buf, file_meta) in bufs.into_iter().zip(file_metas) { + let footer = config::utils::parquet::read_metadata_from_bytes(&bytes::Bytes::from(buf)) + .await + .unwrap(); + assert_eq!(footer.min_ts, file_meta.min_ts); + assert_eq!(footer.max_ts, file_meta.max_ts); + assert_eq!(footer.records, file_meta.records); + assert_eq!(footer.original_size, file_meta.original_size); + } + } + #[test] fn test_get_max_timestamp() { let batch = create_test_record_batch(); diff --git a/src/search/src/datafusion/merge/mod.rs b/src/search/src/datafusion/merge/mod.rs index 3feb9b772c..20679064f2 100644 --- a/src/search/src/datafusion/merge/mod.rs +++ b/src/search/src/datafusion/merge/mod.rs @@ -19,7 +19,10 @@ use arrow::array::RecordBatch; use config::{ FileFormat, TIMESTAMP_COL_NAME, get_config, meta::stream::{FileMeta, StreamType}, - utils::{parquet::new_parquet_writer, util::DISTINCT_STREAM_PREFIX}, + utils::{ + parquet::{VORTEX_FILE_META_KEY, encode_vortex_file_meta, new_parquet_writer}, + util::DISTINCT_STREAM_PREFIX, + }, }; use datafusion::{ arrow::datatypes::Schema, @@ -182,7 +185,7 @@ pub async fn merge_parquet_files( ) .await? } - FileFormat::Vortex => write_vortex(schema, rx, read_task).await?, + FileFormat::Vortex => write_vortex(schema, &metadata, rx, read_task).await?, }; log::debug!( @@ -254,16 +257,21 @@ async fn write_parquet( async fn write_vortex( schema: Arc, + metadata: &FileMeta, mut rx: tokio::sync::mpsc::Receiver, read_task: tokio::task::JoinHandle>, ) -> Result> { + // metadata segments belong to the write options, they can't be appended at + // close time like parquet's, so `records` may drift from the rows written + let file_meta = encode_vortex_file_meta(metadata); let writer_task = VORTEX_RUNTIME.spawn_blocking(move || { VORTEX_RUNTIME.block_on(async move { let mut buf = Vec::new(); let session = VortexSession::default().with_tokio(); let dtype = DType::from_arrow(schema.as_ref()); - let write_options = - VortexWriteOptions::new(session.clone()).with_strategy(vortex_write_strategy()); + let write_options = VortexWriteOptions::new(session.clone()) + .with_strategy(vortex_write_strategy()) + .with_metadata_segment(VORTEX_FILE_META_KEY, file_meta); let mut writer = write_options.writer(&mut buf, dtype); while let Some(batch) = rx.recv().await { @@ -319,6 +327,7 @@ mod tests { use std::sync::Arc; use arrow_schema::{DataType, Field, Schema}; + use vortex::file::OpenOptionsSessionExt; use super::*; @@ -372,4 +381,52 @@ mod tests { // The exact behavior depends on implementation details assert!(result.is_ok() || result.is_err()); } + + #[tokio::test] + async fn test_write_vortex_carries_file_meta() { + use arrow::array::{Int64Array, StringArray}; + + let schema = create_test_schema(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![100, 200, 300])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + Arc::new(Int64Array::from(vec![1, 2, 3])), + ], + ) + .unwrap(); + + let metadata = FileMeta { + min_ts: 100, + max_ts: 300, + records: 3, + original_size: 1024, + ..Default::default() + }; + + let (tx, rx) = tokio::sync::mpsc::channel::(2); + tx.send(batch).await.unwrap(); + drop(tx); + let read_task = tokio::task::spawn(async { Ok(()) }); + + let buf = write_vortex(schema, &metadata, rx, read_task) + .await + .unwrap(); + + let session = VortexSession::default().with_tokio(); + let vxf = session + .open_options() + .include_metadata() + .open_buffer(vortex::buffer::Buffer::from(buf)) + .unwrap(); + let segment = vxf.metadata_segment(VORTEX_FILE_META_KEY).unwrap(); + let file_meta: config::utils::json::Value = + config::utils::json::from_slice(segment.as_slice()).unwrap(); + assert_eq!(file_meta["min_ts"], metadata.min_ts); + assert_eq!(file_meta["max_ts"], metadata.max_ts); + assert_eq!(file_meta["records"], metadata.records); + assert_eq!(file_meta["original_size"], metadata.original_size); + assert_eq!(vxf.row_count(), 3); + } }