chore: write o2 file meta into vortex files (#13553)
Refs #13398 ## What Parquet files carry `min_ts` / `max_ts` / `records` / `original_size` as footer key/values (`new_parquet_writer` on write, `FileMeta: From<&[KeyValue]>` on read). Vortex files carried nothing. Vortex 0.81.0 added user metadata segments (vortex-data/vortex#5813), which our pinned rev has, so the merge path now writes the same four fields into an `o2_file_meta` segment. - `config::utils::parquet`: `VORTEX_FILE_META_KEY` + `encode_vortex_file_meta` - `merge::write_vortex`: attaches the segment (the only place vortex files are produced) While in the same code, a related gap on the parquet side: `write_downsampled_parquet` calls `append_metadata` when it rolls over to a new file, but closed the **final** file without it, so the last file of every downsampling run shipped with an empty footer (`min_ts` / `max_ts` / `records` / `original_size` all read back as 0). Fixed here too. file_list was unaffected — the correct meta is returned via `MergeParquetResult::Multiple.file_metas` — the file itself just was not self-describing. ## Notes - Vortex metadata segments belong to the write options, so they must be known before the first row is written and cannot be appended at close time the way `append_metadata` does for parquet. `records` may therefore drift from the rows actually written — a reader that needs the exact count should take it from the footer `row_count()`, which is always accurate. - **Write side only.** There is no consumer for the read side today: `read_metadata_from_file` only ever sees WAL `.par` files (parquet), `read_metadata_from_bytes` has no callers, and the only reader of storage-file metadata is the `load_file_list_from_s3` recovery CLI via `storage::get_file_meta`, which is still parquet-only (unchanged, already failed on `.vortex` before this PR). Making that path format-aware needs a vortex dependency in `infra` and is better done when it is actually wanted. - For the same "must be known up front" reason, `write_downsampled_vortex` cannot carry the segment: it splits output by accumulated size and only knows each file's meta once that file is closed. Would need a different split strategy (or buffering batches per output file) — not attempted here. ## Test - `test_encode_decode_vortex_file_meta` — payload shape, unknown/missing keys tolerated - `test_write_vortex_carries_file_meta` — drives the real `write_vortex`, reopens the buffer with `include_metadata()` and reads the segment back - `test_write_downsampled_parquet_writes_metadata_for_every_file` — 1-byte file-size limit so each batch becomes its own file, then compares every buffer's footer against the returned `FileMeta`. Verified it fails without the downsampling fix (footer `min_ts` 0 vs expected 3000). This one is behind the `enterprise` feature, so it was run against the real enterprise workspace, not the OSS stub.
This commit is contained in:
GitHub
parent
6fb5eecdc5
commit
6856e0cfc4
|
|
@ -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<u8> {
|
||||
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<u8>,
|
||||
|
|
@ -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";
|
||||
|
|
|
|||
|
|
@ -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::<RecordBatch>(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();
|
||||
|
|
|
|||
|
|
@ -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<Schema>,
|
||||
metadata: &FileMeta,
|
||||
mut rx: tokio::sync::mpsc::Receiver<RecordBatch>,
|
||||
read_task: tokio::task::JoinHandle<Result<()>>,
|
||||
) -> Result<Vec<u8>> {
|
||||
// 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::<RecordBatch>(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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue