From 222bfb74aac04d05c4bc2701bffc13bbb0655588 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Wed, 5 Aug 2026 23:40:53 +0800 Subject: [PATCH] feat: add Parquet table name metadata Signed-off-by: jeremyhi --- src/common/datasource/src/error.rs | 10 +- .../datasource/src/file_format/parquet.rs | 351 +++++++++++++++++- src/operator/src/statement/copy_query_to.rs | 2 +- src/operator/src/statement/copy_table_to.rs | 52 ++- 4 files changed, 403 insertions(+), 12 deletions(-) diff --git a/src/common/datasource/src/error.rs b/src/common/datasource/src/error.rs index dbb0dec417c..c334c3e9b07 100644 --- a/src/common/datasource/src/error.rs +++ b/src/common/datasource/src/error.rs @@ -260,6 +260,13 @@ pub enum Error { error: parquet::errors::ParquetError, }, + #[snafu(display("Invalid Parquet table name metadata: {}", reason))] + InvalidParquetTableNameMetadata { + reason: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(transparent)] DataFusion { #[snafu(implicit)] @@ -298,7 +305,8 @@ impl ErrorExt for Error { | ReadParquetSnafu { .. } | ParquetToSchema { .. } | ParseFormat { .. } - | MergeSchema { .. } => StatusCode::InvalidArguments, + | MergeSchema { .. } + | InvalidParquetTableNameMetadata { .. } => StatusCode::InvalidArguments, JoinHandle { .. } | ReadRecordBatch { .. } diff --git a/src/common/datasource/src/file_format/parquet.rs b/src/common/datasource/src/file_format/parquet.rs index 43b281077e0..7b27eba665c 100644 --- a/src/common/datasource/src/file_format/parquet.rs +++ b/src/common/datasource/src/file_format/parquet.rs @@ -24,7 +24,7 @@ use datafusion::parquet::arrow::async_reader::AsyncFileReader; use datafusion::parquet::arrow::{ArrowWriter, parquet_to_arrow_schema}; use datafusion::parquet::errors::{ParquetError, Result as ParquetResult}; use datafusion::parquet::file::metadata::{ - PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader, + KeyValue, PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader, }; use datafusion::physical_plan::SendableRecordBatchStream; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; @@ -43,10 +43,82 @@ use tokio_util::compat::{Compat, FuturesAsyncReadCompatExt, FuturesAsyncWriteCom use crate::DEFAULT_WRITE_BUFFER_SIZE; use crate::buffered_writer::{ArrowWriterCloser, DfRecordBatchEncoder}; -use crate::error::{self, Result, WriteObjectSnafu, WriteParquetSnafu}; +use crate::error::{ + self, InvalidParquetTableNameMetadataSnafu, Result, WriteObjectSnafu, WriteParquetSnafu, +}; use crate::file_format::FileFormat; use crate::share_buffer::SharedBuffer; +/// Parquet file metadata key for the source GreptimeDB table name. +/// +/// The value is an opaque, unqualified GreptimeDB table name. +pub const PARQUET_TABLE_NAME_KEY: &str = "greptime.table_name"; + +/// Extracts the table name identity from one Parquet file's key/value metadata. +/// +/// Missing metadata or a missing key represents an unknown table identity. +pub fn parquet_table_name_from_metadata(metadata: Option<&[KeyValue]>) -> Result> { + let mut table_name = None; + for item in metadata.unwrap_or_default() { + if item.key != PARQUET_TABLE_NAME_KEY { + continue; + } + + let Some(value) = item.value.as_deref() else { + return InvalidParquetTableNameMetadataSnafu { + reason: format!("metadata key '{PARQUET_TABLE_NAME_KEY}' has no value"), + } + .fail(); + }; + if value.is_empty() { + return InvalidParquetTableNameMetadataSnafu { + reason: format!("metadata key '{PARQUET_TABLE_NAME_KEY}' has an empty value"), + } + .fail(); + } + if table_name.is_some_and(|table_name| table_name != value) { + return InvalidParquetTableNameMetadataSnafu { + reason: format!("metadata key '{PARQUET_TABLE_NAME_KEY}' has conflicting values"), + } + .fail(); + } + table_name = Some(value); + } + + Ok(table_name) +} + +/// Extracts one consistent table identity from a batch of Parquet file metadata. +pub fn parquet_table_name_from_metadata_batch<'a>( + metadata: impl IntoIterator>, +) -> Result> { + let mut table_name = None; + let mut saw_missing = false; + for metadata in metadata { + match parquet_table_name_from_metadata(metadata)? { + Some(value) => { + if saw_missing || table_name.is_some_and(|table_name| table_name != value) { + return InvalidParquetTableNameMetadataSnafu { + reason: "Parquet files have inconsistent table names".to_string(), + } + .fail(); + } + table_name = Some(value); + } + None => { + if table_name.is_some() { + return InvalidParquetTableNameMetadataSnafu { + reason: "Parquet files have inconsistent table names".to_string(), + } + .fail(); + } + saw_missing = true; + } + } + } + Ok(table_name) +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub struct ParquetFormat {} @@ -211,17 +283,42 @@ impl ArrowWriterCloser for ArrowWriter { /// /// Returns number of rows written. pub async fn stream_to_parquet( - mut stream: SendableRecordBatchStream, + stream: SendableRecordBatchStream, schema: datatypes::schema::SchemaRef, store: ObjectStore, path: &str, concurrency: usize, ) -> Result { + stream_to_parquet_with_metadata(stream, schema, store, path, concurrency, Vec::new()).await +} + +/// Output the stream to a parquet file with custom key/value footer metadata. +/// +/// Returns number of rows written. +pub async fn stream_to_parquet_with_metadata( + mut stream: SendableRecordBatchStream, + schema: datatypes::schema::SchemaRef, + store: ObjectStore, + path: &str, + concurrency: usize, + metadata: Vec<(String, String)>, +) -> Result { + let metadata = if metadata.is_empty() { + None + } else { + Some( + metadata + .into_iter() + .map(|(key, value)| KeyValue::new(key, value)) + .collect(), + ) + }; let write_props = column_wise_config( WriterProperties::builder() .set_compression(Compression::ZSTD(ZstdLevel::default())) .set_statistics_truncate_length(None) - .set_column_index_truncate_length(None), + .set_column_index_truncate_length(None) + .set_key_value_metadata(metadata), schema, ) .build(); @@ -269,10 +366,19 @@ fn column_wise_config( #[cfg(test)] mod tests { + use std::sync::Arc; + + use common_recordbatch::adapter::DfRecordBatchStreamAdapter; + use common_recordbatch::{RecordBatch, RecordBatches}; use common_test_util::find_workspace_path; + use datafusion::parquet::arrow::async_reader::AsyncFileReader; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::{ColumnSchema, Schema as GreptimeSchema}; + use datatypes::vectors::{UInt32Vector, VectorRef}; + use parquet::file::metadata::KeyValue; use super::*; - use crate::test_util::{format_schema, test_store}; + use crate::test_util::{format_schema, test_store, test_tmp_store}; fn test_data_root() -> String { find_workspace_path("/src/common/datasource/tests/parquet") @@ -280,6 +386,241 @@ mod tests { .to_string() } + fn test_record_batch_stream() -> (SendableRecordBatchStream, datatypes::schema::SchemaRef) { + let schema = Arc::new(GreptimeSchema::new(vec![ColumnSchema::new( + "number", + ConcreteDataType::uint32_datatype(), + false, + )])); + let columns: Vec = vec![Arc::new(UInt32Vector::from_slice([1, 2, 3]))]; + let batch = RecordBatch::new(schema.clone(), columns).unwrap(); + let batches = RecordBatches::try_new(schema.clone(), vec![batch]).unwrap(); + ( + Box::pin(DfRecordBatchStreamAdapter::new(batches.as_stream())), + schema, + ) + } + + async fn written_key_values(store: ObjectStore, path: &str) -> Option> { + let mut reader = LazyParquetFileReader::new(store, path.to_string(), None); + reader + .get_metadata(None) + .await + .unwrap() + .file_metadata() + .key_value_metadata() + .cloned() + } + + #[tokio::test] + async fn legacy_parquet_writer_does_not_write_table_name_metadata() { + let (store, dir) = test_tmp_store("legacy_parquet_writer_metadata"); + let path = format!("{}/output.parquet", dir.path().display()); + let (stream, schema) = test_record_batch_stream(); + + assert_eq!( + 3, + stream_to_parquet(stream, schema, store.clone(), &path, 1) + .await + .unwrap() + ); + + let user_key_values = written_key_values(store, &path) + .await + .unwrap_or_default() + .into_iter() + .filter(|item| item.key != "ARROW:schema") + .collect::>(); + assert!(user_key_values.is_empty(), "{user_key_values:?}"); + } + + #[tokio::test] + async fn parquet_writer_writes_table_name_metadata_to_footer() { + let (store, dir) = test_tmp_store("parquet_writer_table_name_metadata"); + let path = format!("{}/output.parquet", dir.path().display()); + let (stream, schema) = test_record_batch_stream(); + + assert_eq!( + 3, + stream_to_parquet_with_metadata( + stream, + schema, + store.clone(), + &path, + 1, + vec![(PARQUET_TABLE_NAME_KEY.to_string(), "my_table".to_string(),)], + ) + .await + .unwrap() + ); + + let table_names = written_key_values(store, &path) + .await + .unwrap_or_default() + .into_iter() + .filter(|item| item.key == PARQUET_TABLE_NAME_KEY) + .collect::>(); + assert_eq!(1, table_names.len()); + assert_eq!(Some("my_table"), table_names[0].value.as_deref()); + } + + #[test] + fn missing_parquet_table_name_is_unknown() { + let metadata = [KeyValue::new("other".to_string(), "value".to_string())]; + + assert_eq!( + None, + parquet_table_name_from_metadata(Some(&metadata)).unwrap() + ); + } + + #[test] + fn present_parquet_table_name_is_returned_as_opaque_value() { + let metadata = [KeyValue::new( + PARQUET_TABLE_NAME_KEY.to_string(), + "schema/table.with.dots".to_string(), + )]; + + assert_eq!( + Some("schema/table.with.dots"), + parquet_table_name_from_metadata(Some(&metadata)).unwrap() + ); + } + + #[test] + fn parquet_table_name_without_value_is_invalid() { + let metadata = [KeyValue::new( + PARQUET_TABLE_NAME_KEY.to_string(), + None::, + )]; + + assert!(parquet_table_name_from_metadata(Some(&metadata)).is_err()); + } + + #[test] + fn empty_parquet_table_name_is_invalid() { + let metadata = [KeyValue::new( + PARQUET_TABLE_NAME_KEY.to_string(), + String::new(), + )]; + + assert!(parquet_table_name_from_metadata(Some(&metadata)).is_err()); + } + + #[test] + fn conflicting_parquet_table_names_are_invalid() { + let metadata = [ + KeyValue::new(PARQUET_TABLE_NAME_KEY.to_string(), "first".to_string()), + KeyValue::new(PARQUET_TABLE_NAME_KEY.to_string(), "second".to_string()), + ]; + + assert!(parquet_table_name_from_metadata(Some(&metadata)).is_err()); + } + + #[test] + fn identical_duplicate_parquet_table_names_are_accepted() { + let metadata = [ + KeyValue::new(PARQUET_TABLE_NAME_KEY.to_string(), "same".to_string()), + KeyValue::new(PARQUET_TABLE_NAME_KEY.to_string(), "same".to_string()), + ]; + + assert_eq!( + Some("same"), + parquet_table_name_from_metadata(Some(&metadata)).unwrap() + ); + } + + #[test] + fn empty_parquet_metadata_batch_is_unknown() { + assert_eq!( + None, + parquet_table_name_from_metadata_batch(std::iter::empty()).unwrap() + ); + } + + #[test] + fn all_missing_parquet_metadata_batch_is_unknown() { + assert_eq!( + None, + parquet_table_name_from_metadata_batch([None, None]).unwrap() + ); + } + + #[test] + fn same_parquet_table_name_across_batch_is_returned() { + let first = [KeyValue::new( + PARQUET_TABLE_NAME_KEY.to_string(), + "same".to_string(), + )]; + let second = [KeyValue::new( + PARQUET_TABLE_NAME_KEY.to_string(), + "same".to_string(), + )]; + + assert_eq!( + Some("same"), + parquet_table_name_from_metadata_batch([ + Some(first.as_slice()), + Some(second.as_slice()) + ]) + .unwrap() + ); + } + + #[test] + fn different_parquet_table_names_across_batch_are_invalid() { + let first = [KeyValue::new( + PARQUET_TABLE_NAME_KEY.to_string(), + "first".to_string(), + )]; + let second = [KeyValue::new( + PARQUET_TABLE_NAME_KEY.to_string(), + "second".to_string(), + )]; + + assert!( + parquet_table_name_from_metadata_batch([ + Some(first.as_slice()), + Some(second.as_slice()) + ]) + .is_err() + ); + } + + #[test] + fn invalid_parquet_table_name_in_batch_is_propagated() { + let invalid = [KeyValue::new( + PARQUET_TABLE_NAME_KEY.to_string(), + None::, + )]; + + assert!(parquet_table_name_from_metadata_batch([Some(invalid.as_slice())]).is_err()); + } + + #[test] + fn mixed_present_and_missing_parquet_table_names_are_invalid() { + let present = [KeyValue::new( + PARQUET_TABLE_NAME_KEY.to_string(), + "table".to_string(), + )]; + let missing = [KeyValue::new("other".to_string(), "value".to_string())]; + + assert!( + parquet_table_name_from_metadata_batch([ + Some(present.as_slice()), + Some(missing.as_slice()), + ]) + .is_err() + ); + assert!( + parquet_table_name_from_metadata_batch([ + Some(missing.as_slice()), + Some(present.as_slice()), + ]) + .is_err() + ); + } + #[tokio::test] async fn infer_schema_basic() { let json = ParquetFormat::default(); diff --git a/src/operator/src/statement/copy_query_to.rs b/src/operator/src/statement/copy_query_to.rs index 98fc92e23eb..625bcd7f390 100644 --- a/src/operator/src/statement/copy_query_to.rs +++ b/src/operator/src/statement/copy_query_to.rs @@ -37,7 +37,7 @@ impl StatementExecutor { let format = Format::try_from(with).context(error::ParseFileFormatSnafu)?; debug!("Copy query to location: {location}"); - self.copy_to_file(&format, query_output, location, connection) + self.copy_to_file(&format, None, query_output, location, connection) .await } } diff --git a/src/operator/src/statement/copy_table_to.rs b/src/operator/src/statement/copy_table_to.rs index 4cf9448d8b6..a9344629d4a 100644 --- a/src/operator/src/statement/copy_table_to.rs +++ b/src/operator/src/statement/copy_table_to.rs @@ -20,7 +20,9 @@ use common_base::readable_size::ReadableSize; use common_datasource::file_format::Format; use common_datasource::file_format::csv::stream_to_csv; use common_datasource::file_format::json::stream_to_json; -use common_datasource::file_format::parquet::stream_to_parquet; +use common_datasource::file_format::parquet::{ + PARQUET_TABLE_NAME_KEY, stream_to_parquet_with_metadata, +}; use common_datasource::object_store::build_backend_for_write_with_path; use common_query::Output; use common_recordbatch::adapter::DfRecordBatchStreamAdapter; @@ -49,11 +51,21 @@ const WRITE_BUFFER_THRESHOLD: ReadableSize = ReadableSize::mb(8); /// Default number of concurrent write, it only works on object store backend(e.g., S3). const WRITE_CONCURRENCY: usize = 8; +fn parquet_metadata(format: &Format, table_name: Option<&str>) -> Vec<(String, String)> { + match (format, table_name) { + (Format::Parquet(_), Some(table_name)) => { + vec![(PARQUET_TABLE_NAME_KEY.to_string(), table_name.to_string())] + } + _ => Vec::new(), + } +} + impl StatementExecutor { async fn stream_to_file( &self, stream: SendableRecordBatchStream, format: &Format, + parquet_table_name: Option<&str>, object_store: ObjectStore, path: &str, ) -> Result { @@ -87,12 +99,13 @@ impl StatementExecutor { .context(error::WriteStreamToFileSnafu { path }), Format::Parquet(_) => { let schema = stream.schema(); - stream_to_parquet( + stream_to_parquet_with_metadata( Box::pin(DfRecordBatchStreamAdapter::new(stream)), schema, object_store, path, WRITE_CONCURRENCY, + parquet_metadata(format, parquet_table_name), ) .await .context(error::WriteStreamToFileSnafu { path }) @@ -157,13 +170,14 @@ impl StatementExecutor { } = &req; debug!("Copy table: {table_id} to location: {location}"); - self.copy_to_file(&format, output, location, connection) + self.copy_to_file(&format, Some(&req.table_name), output, location, connection) .await } pub(crate) async fn copy_to_file( &self, format: &Format, + parquet_table_name: Option<&str>, output: Output, location: &str, connection: &HashMap, @@ -184,7 +198,35 @@ impl StatementExecutor { let filename = backend.object_path.context(error::UnexpectedSnafu { violated: format!("Expected filename, path: {location}"), })?; - self.stream_to_file(stream, format, backend.object_store, &filename) - .await + self.stream_to_file( + stream, + format, + parquet_table_name, + backend.object_store, + &filename, + ) + .await + } +} + +#[cfg(test)] +mod tests { + use common_datasource::file_format::Format; + use common_datasource::file_format::json::JsonFormat; + use common_datasource::file_format::parquet::{PARQUET_TABLE_NAME_KEY, ParquetFormat}; + + use super::parquet_metadata; + + #[test] + fn parquet_metadata_routes_only_table_parquet_exports() { + let parquet = Format::Parquet(ParquetFormat::default()); + let json = Format::Json(JsonFormat::default()); + + assert_eq!( + vec![(PARQUET_TABLE_NAME_KEY.to_string(), "my_table".to_string())], + parquet_metadata(&parquet, Some("my_table")) + ); + assert!(parquet_metadata(&parquet, None).is_empty()); + assert!(parquet_metadata(&json, Some("my_table")).is_empty()); } }