diff --git a/Cargo.lock b/Cargo.lock index 3d0781a9031..fd9eeaefdd7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -15229,6 +15229,7 @@ dependencies = [ "tokio", "tokio-postgres", "tokio-stream", + "tokio-util", "tonic 0.14.2", "tower 0.5.2", "url", diff --git a/src/common/datasource/Cargo.toml b/src/common/datasource/Cargo.toml index a601dae1d90..d1ca609c606 100644 --- a/src/common/datasource/Cargo.toml +++ b/src/common/datasource/Cargo.toml @@ -47,3 +47,4 @@ url.workspace = true [dev-dependencies] common-test-util.workspace = true +object-store = { workspace = true, features = ["testing"] } diff --git a/src/common/datasource/src/error.rs b/src/common/datasource/src/error.rs index dbb0dec417c..1cc04d5a3de 100644 --- a/src/common/datasource/src/error.rs +++ b/src/common/datasource/src/error.rs @@ -27,6 +27,15 @@ use url::ParseError; #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display("Parquet write cancelled"))] + ParquetWriteCancelled {}, + + #[snafu(display("Parquet writer limits must be positive"))] + InvalidParquetWriterLimits {}, + + #[snafu(display("Parquet writer resource limit exceeded: {reason}"))] + ParquetWriterResource { reason: String }, + #[snafu(display("Unsupported compression type: {}", compression_type))] UnsupportedCompressionType { compression_type: String, @@ -275,6 +284,9 @@ impl ErrorExt for Error { fn status_code(&self) -> StatusCode { use Error::*; match self { + ParquetWriteCancelled {} => StatusCode::Cancelled, + InvalidParquetWriterLimits {} => StatusCode::InvalidArguments, + ParquetWriterResource { .. } => StatusCode::Suspended, BuildBackend { .. } | ListObjects { .. } | ReadObject { .. } diff --git a/src/common/datasource/src/file_format/parquet.rs b/src/common/datasource/src/file_format/parquet.rs index 43b281077e0..0032de6b31e 100644 --- a/src/common/datasource/src/file_format/parquet.rs +++ b/src/common/datasource/src/file_format/parquet.rs @@ -29,22 +29,17 @@ use datafusion::parquet::file::metadata::{ use datafusion::physical_plan::SendableRecordBatchStream; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_datasource::PartitionedFile; -use datatypes::schema::SchemaRef; use futures::StreamExt; use futures::future::BoxFuture; use object_store::{FuturesAsyncReader, ObjectStore}; -use parquet::arrow::AsyncArrowWriter; use parquet::arrow::arrow_reader::ArrowReaderOptions; -use parquet::basic::{Compression, Encoding, ZstdLevel}; -use parquet::file::properties::{WriterProperties, WriterPropertiesBuilder}; -use parquet::schema::types::ColumnPath; use snafu::ResultExt; -use tokio_util::compat::{Compat, FuturesAsyncReadCompatExt, FuturesAsyncWriteCompatExt}; +use tokio_util::compat::{Compat, FuturesAsyncReadCompatExt}; -use crate::DEFAULT_WRITE_BUFFER_SIZE; use crate::buffered_writer::{ArrowWriterCloser, DfRecordBatchEncoder}; -use crate::error::{self, Result, WriteObjectSnafu, WriteParquetSnafu}; +use crate::error::{self, Result}; use crate::file_format::FileFormat; +use crate::parquet_writer::ParquetFileWriter; use crate::share_buffer::SharedBuffer; #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] @@ -212,61 +207,23 @@ impl ArrowWriterCloser for ArrowWriter { /// Returns number of rows written. pub async fn stream_to_parquet( mut stream: SendableRecordBatchStream, - schema: datatypes::schema::SchemaRef, store: ObjectStore, path: &str, concurrency: usize, ) -> Result { - 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), - schema, - ) - .build(); - let inner_writer = store - .writer_with(path) - .concurrent(concurrency) - .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) - .await - .map(|w| w.into_futures_async_write().compat_write()) - .context(WriteObjectSnafu { path })?; - - let mut writer = AsyncArrowWriter::try_new(inner_writer, stream.schema(), Some(write_props)) - .context(WriteParquetSnafu { path })?; + let mut writer = + ParquetFileWriter::open(stream.schema(), store, path, concurrency, None).await?; let mut rows_written = 0; - while let Some(batch) = stream.next().await { let batch = batch.context(error::ReadRecordBatchSnafu)?; - writer - .write(&batch) - .await - .context(WriteParquetSnafu { path })?; - rows_written += batch.num_rows(); + let rows = batch.num_rows(); + writer.write(batch, None).await?; + rows_written += rows; } - writer.close().await.context(WriteParquetSnafu { path })?; + writer.finish(None).await?; Ok(rows_written) } -/// Customizes per-column properties. -fn column_wise_config( - mut props: WriterPropertiesBuilder, - schema: SchemaRef, -) -> WriterPropertiesBuilder { - // Disable dictionary for timestamp column, since for increasing timestamp column, - // the dictionary pages will be larger than data pages. - for col in schema.column_schemas() { - if col.data_type.is_timestamp() { - let path = ColumnPath::new(vec![col.name.clone()]); - props = props - .set_column_dictionary_enabled(path.clone(), false) - .set_column_encoding(path, Encoding::DELTA_BINARY_PACKED) - } - } - props -} - #[cfg(test)] mod tests { use common_test_util::find_workspace_path; diff --git a/src/common/datasource/src/parquet_writer.rs b/src/common/datasource/src/parquet_writer.rs index 088d17fb831..ae1fcec04a9 100644 --- a/src/common/datasource/src/parquet_writer.rs +++ b/src/common/datasource/src/parquet_writer.rs @@ -12,11 +12,215 @@ // See the License for the specific language governing permissions and // limitations under the License. +use arrow::datatypes::{DataType, SchemaRef}; +use arrow::record_batch::RecordBatch; use bytes::Bytes; use futures::future::BoxFuture; -use object_store::Writer; +use object_store::{ObjectStore, Writer}; +use parquet::arrow::ArrowWriter; use parquet::arrow::async_writer::AsyncFileWriter; +use parquet::basic::{Compression, Encoding, ZstdLevel}; use parquet::errors::ParquetError; +use parquet::file::properties::WriterProperties; +use parquet::schema::types::ColumnPath; +use snafu::{IntoError, ResultExt, ensure}; +use tokio_util::sync::CancellationToken; + +use crate::DEFAULT_WRITE_BUFFER_SIZE; +use crate::error::{self, Result}; + +/// Limits for one Parquet file. Flush thresholds are not hard memory caps. +#[derive(Clone, Copy, Debug)] +pub struct ParquetWriterLimits { + /// Maximum rows in a row group. + pub row_group_rows: usize, + /// Flush when encoder memory or encoded size reaches this threshold. + pub flush_threshold_bytes: usize, + /// Maximum row groups retained in the file footer. + pub max_row_groups: usize, +} + +/// Encodes batches on a blocking runtime and writes one object-store file. +/// Callers must await each operation before aborting; dropping an in-flight +/// operation can leave a filesystem worker running after cleanup. +pub struct ParquetFileWriter { + encoder: Option>>, + sink: Writer, + store: ObjectStore, + path: String, + limits: Option, +} + +impl ParquetFileWriter { + /// Open a file using COPY's encoding settings. Destination ownership and + /// overwrite policy belong to the caller. None preserves Parquet's defaults. + pub async fn open( + schema: SchemaRef, + store: ObjectStore, + path: &str, + concurrency: usize, + limits: Option, + ) -> Result { + let mut props = WriterProperties::builder() + .set_compression(Compression::ZSTD(ZstdLevel::default())) + .set_statistics_truncate_length(None) + .set_column_index_truncate_length(None); + if let Some(limits) = limits { + ensure!( + limits.row_group_rows > 0 + && limits.flush_threshold_bytes > 0 + && limits.max_row_groups > 0, + error::InvalidParquetWriterLimitsSnafu + ); + props = props + .set_max_row_group_row_count(Some(limits.row_group_rows)) + .set_max_row_group_bytes(None); + } + for field in schema.fields() { + if matches!(field.data_type(), DataType::Timestamp(_, _)) { + let column = ColumnPath::new(vec![field.name().clone()]); + props = props + .set_column_dictionary_enabled(column.clone(), false) + .set_column_encoding(column, Encoding::DELTA_BINARY_PACKED); + } + } + let encoder = ArrowWriter::try_new(Vec::new(), schema, Some(props.build())) + .context(error::WriteParquetSnafu { path })?; + let sink = store + .writer_with(path) + .concurrent(concurrency) + .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) + .await + .context(error::WriteObjectSnafu { path })?; + Ok(Self { + encoder: Some(encoder), + sink, + store, + path: path.to_owned(), + limits, + }) + } + + /// Write a batch, enforcing file limits across batch and row-group boundaries. + pub async fn write( + &mut self, + batch: RecordBatch, + cancellation: Option<&CancellationToken>, + ) -> Result<()> { + let mut offset = 0; + while offset < batch.num_rows() { + check_cancelled(cancellation)?; + let mut encoder = self.encoder.take().ok_or_else(|| { + error::WriteParquetSnafu { path: &self.path } + .into_error(ParquetError::General("Parquet writer is closed".into())) + })?; + let len = self.limits.map_or(batch.num_rows() - offset, |limits| { + (limits.row_group_rows - encoder.in_progress_rows()).min(batch.num_rows() - offset) + }); + let slice = batch.slice(offset, len); + let limits = self.limits; + let path = self.path.clone(); + let (encoder, bytes) = common_runtime::spawn_blocking_global(move || { + if let Some(limits) = limits { + ensure!( + encoder.flushed_row_groups().len() < limits.max_row_groups, + error::ParquetWriterResourceSnafu { + reason: "Parquet row-group metadata budget exceeded" + } + ); + } + encoder + .write(&slice) + .context(error::WriteParquetSnafu { path: &path })?; + if limits.is_some_and(|limits| { + encoder.memory_size() >= limits.flush_threshold_bytes + || encoder.in_progress_size() >= limits.flush_threshold_bytes + }) { + encoder + .flush() + .context(error::WriteParquetSnafu { path: &path })?; + } + // Draining preserves ArrowWriter's cumulative file offsets. + let bytes = std::mem::take(encoder.inner_mut()); + Ok::<_, error::Error>((encoder, bytes)) + }) + .await + .context(error::JoinHandleSnafu)??; + self.encoder = Some(encoder); + check_cancelled(cancellation)?; + self.write_bytes(bytes).await?; + check_cancelled(cancellation)?; + offset += len; + } + Ok(()) + } + + async fn write_bytes(&mut self, bytes: Vec) -> Result<()> { + if !bytes.is_empty() { + self.sink + .write(bytes) + .await + .context(error::WriteObjectSnafu { path: &self.path })?; + } + Ok(()) + } + + /// Write the footer and close the file. Retains the handle for cleanup on error. + pub async fn finish(&mut self, cancellation: Option<&CancellationToken>) -> Result<()> { + check_cancelled(cancellation)?; + let mut encoder = self.encoder.take().ok_or_else(|| { + error::WriteParquetSnafu { path: &self.path } + .into_error(ParquetError::General("Parquet writer is closed".into())) + })?; + let path = self.path.clone(); + let bytes = common_runtime::spawn_blocking_global(move || { + encoder + .finish() + .context(error::WriteParquetSnafu { path: &path })?; + Ok::<_, error::Error>(std::mem::take(encoder.inner_mut())) + }) + .await + .context(error::JoinHandleSnafu)??; + self.write_bytes(bytes).await?; + check_cancelled(cancellation)?; + self.sink + .close() + .await + .context(error::WriteObjectSnafu { path: &self.path })?; + check_cancelled(cancellation)?; + Ok(()) + } + + /// Abort an exclusively owned file after all in-flight operations have completed. + /// If the backend cannot abort, deletion assumes this attempt owns the path. + pub async fn abort(mut self) -> Result<()> { + let result = self.sink.abort().await; + if result + .as_ref() + .is_err_and(|e| e.kind() == object_store::ErrorKind::Unsupported) + { + let store = self.store.clone(); + let path = self.path.clone(); + // Secure filesystem writers require dropping the handle before deletion. + drop(self); + store + .delete(&path) + .await + .context(error::WriteObjectSnafu { path })?; + } else { + result.context(error::WriteObjectSnafu { path: &self.path })?; + } + Ok(()) + } +} + +fn check_cancelled(cancellation: Option<&CancellationToken>) -> Result<()> { + ensure!( + cancellation.is_none_or(|token| !token.is_cancelled()), + error::ParquetWriteCancelledSnafu + ); + Ok(()) +} /// Bridges opendal [Writer] with parquet [AsyncFileWriter]. pub struct AsyncWriter { @@ -50,3 +254,244 @@ impl AsyncFileWriter for AsyncWriter { }) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{Int64Array, TimestampMillisecondArray}; + use common_error::ext::ErrorExt; + use common_error::status_code::StatusCode; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + + use super::*; + use crate::file_format::parquet::stream_to_parquet; + + fn batch() -> RecordBatch { + RecordBatch::try_from_iter([ + ( + "value", + Arc::new(Int64Array::from(vec![Some(1), None, Some(3), Some(4)])) + as arrow::array::ArrayRef, + ), + ( + "ts", + Arc::new(TimestampMillisecondArray::from(vec![1, 2, 3, 4])), + ), + ]) + .unwrap() + } + + async fn read(store: &ObjectStore, path: &str) -> ParquetRecordBatchReaderBuilder { + ParquetRecordBatchReaderBuilder::try_new(store.read(path).await.unwrap().to_bytes()) + .unwrap() + } + + #[tokio::test] + async fn row_and_byte_limits_split_batches() { + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let batch = batch(); + for (row_group_rows, flush_threshold_bytes, split) in [(2, usize::MAX, 3), (100, 1, 2)] { + let mut writer = ParquetFileWriter::open( + batch.schema(), + store.clone(), + "groups.parquet", + 1, + Some(ParquetWriterLimits { + row_group_rows, + flush_threshold_bytes, + max_row_groups: 2, + }), + ) + .await + .unwrap(); + writer.write(batch.slice(0, split), None).await.unwrap(); + writer + .write(batch.slice(split, batch.num_rows() - split), None) + .await + .unwrap(); + writer.finish(None).await.unwrap(); + let reader = read(&store, "groups.parquet").await; + assert_eq!(reader.metadata().num_row_groups(), 2); + for group in reader.metadata().row_groups() { + assert_eq!(group.num_rows(), 2); + let encodings = group.column(1).encodings().collect::>(); + assert!(encodings.contains(&Encoding::DELTA_BINARY_PACKED)); + assert!(!encodings.contains(&Encoding::RLE_DICTIONARY)); + assert_eq!( + group.column(0).compression(), + Compression::ZSTD(ZstdLevel::default()) + ); + } + let actual = reader + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!( + arrow::compute::concat_batches(&batch.schema(), &actual).unwrap(), + batch + ); + } + } + + #[tokio::test] + async fn oversized_batch_stops_at_footer_limit_and_can_abort() { + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let batch = batch(); + let mut writer = ParquetFileWriter::open( + batch.schema(), + store.clone(), + "limited.parquet", + 1, + Some(ParquetWriterLimits { + row_group_rows: 2, + flush_threshold_bytes: usize::MAX, + max_row_groups: 1, + }), + ) + .await + .unwrap(); + let err = writer.write(batch, None).await.unwrap_err(); + assert!(matches!(err, error::Error::ParquetWriterResource { .. })); + assert_eq!(err.status_code(), StatusCode::Suspended); + writer.abort().await.unwrap(); + assert!(!store.exists("limited.parquet").await.unwrap()); + } + + #[tokio::test] + async fn copy_stream_preserves_values_and_empty_schema() { + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let batch = batch(); + for (path, batches) in [ + ("copy.parquet", vec![batch.clone()]), + ("empty.parquet", vec![]), + ] { + let stream = RecordBatchStreamAdapter::new( + batch.schema(), + futures::stream::iter(batches.clone().into_iter().map(Ok)), + ); + assert_eq!( + stream_to_parquet(Box::pin(stream), store.clone(), path, 2) + .await + .unwrap(), + batches.iter().map(RecordBatch::num_rows).sum::() + ); + let reader = read(&store, path).await; + assert_eq!(reader.schema().fields(), batch.schema().fields()); + let actual = reader + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(actual, batches); + } + } + + #[tokio::test] + async fn failed_copy_preserves_untouched_securefs_destination() { + let directory = common_test_util::temp_dir::create_temp_dir("copy_existing"); + let path = directory.path().join("existing.parquet"); + std::fs::write(&path, b"original bytes").unwrap(); + let access = crate::object_store::LocalFileAccess::sandboxed(directory.path()).unwrap(); + let store = crate::object_store::build_backend_for_write( + &format!("{}/", directory.path().display()), + &Default::default(), + &access, + ) + .await + .unwrap(); + let stream = RecordBatchStreamAdapter::new( + batch().schema(), + futures::stream::iter(vec![Err(datafusion::error::DataFusionError::Execution( + "injected input failure".into(), + ))]), + ); + let err = stream_to_parquet(Box::pin(stream), store, "existing.parquet", 1) + .await + .unwrap_err(); + assert!(matches!(err, error::Error::ReadRecordBatch { .. })); + assert_eq!(std::fs::read(path).unwrap(), b"original bytes"); + } + + struct PausedFooterWriter { + inner: object_store::layers::mock::oio::Writer, + paused: bool, + started: Arc, + release: Arc, + closed: Arc, + } + + impl object_store::layers::mock::oio::Write for PausedFooterWriter { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + if !self.paused { + self.paused = true; + self.started.notify_one(); + self.release.notified().await; + } + self.inner.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + self.closed.store(true, std::sync::atomic::Ordering::SeqCst); + self.inner.close().await + } + async fn abort(&mut self) -> object_store::Result<()> { + self.inner.abort().await + } + } + + #[tokio::test] + async fn cancellation_during_footer_write_waits_then_aborts_before_close() { + use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory}; + let started = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let closed = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let factory: MockWriterFactory = Arc::new({ + let (started, release, closed) = (started.clone(), release.clone(), closed.clone()); + move |_, _, inner| { + Box::new(PausedFooterWriter { + inner, + paused: false, + started: started.clone(), + release: release.clone(), + closed: closed.clone(), + }) + } + }); + let store = ObjectStore::new(object_store::services::Memory::default()) + .unwrap() + .layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let mut writer = + ParquetFileWriter::open(batch().schema(), store.clone(), "cancel.parquet", 1, None) + .await + .unwrap(); + // Force footer bytes through the sink before its close operation. + writer.sink = store.writer_with("cancel.parquet").chunk(1).await.unwrap(); + let cancellation = CancellationToken::new(); + let result = { + let finish = writer.finish(Some(&cancellation)); + tokio::pin!(finish); + tokio::select! { + result = &mut finish => panic!("finished before footer write: {result:?}"), + _ = started.notified() => {}, + } + cancellation.cancel(); + assert!(futures::poll!(&mut finish).is_pending()); + release.notify_one(); + finish.await + }; + assert!(matches!( + result, + Err(error::Error::ParquetWriteCancelled {}) + )); + assert!(!closed.load(std::sync::atomic::Ordering::SeqCst)); + writer.abort().await.unwrap(); + assert!(!store.exists("cancel.parquet").await.unwrap()); + } +} diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index 4303e831002..4ed4927b41f 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -89,5 +89,6 @@ axum.workspace = true catalog = { workspace = true, features = ["testing"] } common-meta = { workspace = true, features = ["testing"] } common-test-util.workspace = true +object-store = { workspace = true, features = ["testing"] } path-slash = "0.2" url.workspace = true diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index ae680be33ba..3c1a6ae3604 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -85,6 +85,15 @@ pub enum Error { source: common_meta::error::Error, }, + #[snafu(display("Invalid logical table export: {reason}"))] + InvalidLogicalTableExport { reason: String }, + + #[snafu(display("Logical table export resource limit exceeded: {reason}"))] + LogicalTableExportResource { reason: String }, + + #[snafu(display("Logical table export cancelled"))] + LogicalTableExportCancelled {}, + #[snafu(display("Unexpected, violated: {}", violated))] Unexpected { violated: String, @@ -1071,6 +1080,9 @@ impl ErrorExt for Error { Error::InvalidTimeIndexType { .. } | Error::InvalidTimezone { .. } => { StatusCode::InvalidArguments } + Error::InvalidLogicalTableExport { .. } => StatusCode::InvalidArguments, + Error::LogicalTableExportResource { .. } => StatusCode::Suspended, + Error::LogicalTableExportCancelled { .. } => StatusCode::Cancelled, Error::InvalidProcessId { .. } => StatusCode::InvalidArguments, Error::ProcessManagerMissing { .. } => StatusCode::Unexpected, Error::TimestampFormatNotSupported { .. } => StatusCode::InvalidArguments, diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 150020bc9f3..69d387de206 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -22,6 +22,7 @@ mod cursor; pub mod ddl; mod describe; mod dml; +pub mod export_logical_tables; mod kill; pub mod semantic_graph; mod set; diff --git a/src/operator/src/statement/copy_table_to.rs b/src/operator/src/statement/copy_table_to.rs index 4cf9448d8b6..faa06b88216 100644 --- a/src/operator/src/statement/copy_table_to.rs +++ b/src/operator/src/statement/copy_table_to.rs @@ -25,8 +25,8 @@ use common_datasource::object_store::build_backend_for_write_with_path; use common_query::Output; use common_recordbatch::adapter::DfRecordBatchStreamAdapter; use common_recordbatch::{ - RecordBatchStream, SendableRecordBatchMapper, SendableRecordBatchStream, - map_json_type_to_string, map_json_type_to_string_schema, + SendableRecordBatchMapper, SendableRecordBatchStream, map_json_type_to_string, + map_json_type_to_string_schema, }; use common_telemetry::{debug, tracing}; use datafusion::datasource::DefaultTableSource; @@ -85,18 +85,14 @@ impl StatementExecutor { ) .await .context(error::WriteStreamToFileSnafu { path }), - Format::Parquet(_) => { - let schema = stream.schema(); - stream_to_parquet( - Box::pin(DfRecordBatchStreamAdapter::new(stream)), - schema, - object_store, - path, - WRITE_CONCURRENCY, - ) - .await - .context(error::WriteStreamToFileSnafu { path }) - } + Format::Parquet(_) => stream_to_parquet( + Box::pin(DfRecordBatchStreamAdapter::new(stream)), + object_store, + path, + WRITE_CONCURRENCY, + ) + .await + .context(error::WriteStreamToFileSnafu { path }), _ => error::UnsupportedFormatSnafu { format: format.clone(), } diff --git a/src/operator/src/statement/export_logical_tables.rs b/src/operator/src/statement/export_logical_tables.rs new file mode 100644 index 00000000000..7c38af14ba0 --- /dev/null +++ b/src/operator/src/statement/export_logical_tables.rs @@ -0,0 +1,653 @@ +// Copyright 2023 Greptime Team +// +// 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. + +//! Internal physical-table export into logical-table Parquet files. +//! +//! Callers must authorize the captured tables and destination before invoking +//! this component, and keep the source stable. Each destination must be owned +//! exclusively by the export attempt. Completion and cleanup of closed files +//! belong to the caller's chunk protocol, not individual file completion. + +use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::sync::Arc; + +use arrow::array::{Array, AsArray, UInt32Array}; +use arrow::compute::cast; +use arrow::datatypes::{DataType, SchemaRef}; +use arrow::downcast_dictionary_array; +use arrow::record_batch::RecordBatch; +use common_datasource::object_store::build_backend_for_write; +use common_datasource::parquet_writer::{ParquetFileWriter, ParquetWriterLimits}; +use common_meta::key::table_route::{TableRouteManager, TableRouteValue}; +use common_query::OutputData; +use common_recordbatch::SendableRecordBatchStream; +use common_time::range::TimestampRange; +use datafusion::datasource::DefaultTableSource; +use datafusion_common::TableReference as DfTableReference; +use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, col}; +use futures::StreamExt; +use object_store::ObjectStore; +use session::context::QueryContextRef; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; +use store_api::metric_engine_consts::{ + DATA_SCHEMA_TABLE_ID_COLUMN_NAME as TABLE_ID, DATA_SCHEMA_TSID_COLUMN_NAME as TSID, + LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY, +}; +use table::TableRef; +use table::metadata::TableId; +use table::table::adapter::DfTableProviderAdapter; +use tokio_util::sync::CancellationToken; + +use crate::error::{self, InvalidLogicalTableExportSnafu, LogicalTableExportResourceSnafu, Result}; +use crate::statement::StatementExecutor; + +/// Export preprocessing and per-file limits. Query memory and spill remain +/// governed by the query engine. +/// Batch checks happen after allocation; flush thresholds are not hard RSS caps. +#[derive(Clone, Copy, Debug)] +pub struct LogicalTableExportLimits { + /// Maximum retained scan batch before routing. + pub input_batch_bytes: usize, + /// Maximum estimated expanded logical values in one conversion slice. + pub conversion_bytes: usize, + /// Limits owned by the single-file Parquet writer. + pub writer: ParquetWriterLimits, +} + +impl Default for LogicalTableExportLimits { + fn default() -> Self { + Self { + input_batch_bytes: 64 * 1024 * 1024, + conversion_bytes: 1024 * 1024, + writer: ParquetWriterLimits { + row_group_rows: 8192, + flush_threshold_bytes: 8 * 1024 * 1024, + max_row_groups: 4096, + }, + } + } +} + +impl LogicalTableExportLimits { + fn validate(self) -> Result<()> { + ensure!( + self.input_batch_bytes > 0 && self.conversion_bytes > 0, + InvalidLogicalTableExportSnafu { + reason: "export limits must be positive" + } + ); + Ok(()) + } +} + +/// Metadata captured for one physical table and its selected logical tables. +/// Construct one unit per physical table, schema and time chunk. Capturing these +/// references supplies no snapshot isolation or locking guarantee. +pub struct LogicalTableExport { + physical_table: TableRef, + scan_projection: Vec, + logical_tables: BTreeMap, +} + +struct LogicalTableProjection { + name: String, + schema: SchemaRef, + projection: Vec, +} + +impl LogicalTableExport { + /// Capture schemas from selected Metric table references. + /// Export validates their physical-table association against table routes. + pub fn try_new(physical: TableRef, tables: &[TableRef]) -> Result { + let physical_info = physical.table_info(); + ensure!( + physical_info.meta.engine == METRIC_ENGINE_NAME + && physical_info + .meta + .options + .extra_options + .contains_key(PHYSICAL_TABLE_METADATA_KEY) + && !tables.is_empty(), + InvalidLogicalTableExportSnafu { + reason: "expected a Metric physical table and selected logical tables" + } + ); + let physical_schema = physical.schema(); + let id_index = physical_schema.column_index_by_name(TABLE_ID).context( + InvalidLogicalTableExportSnafu { + reason: "physical schema has no __table_id", + }, + )?; + let mut scan_projection = BTreeSet::from([id_index]); + let mut logical_tables = BTreeMap::new(); + for table in tables { + let info = table.table_info(); + let name = &info.name; + ensure!( + !name.contains('/') && !name.contains('\\'), + InvalidLogicalTableExportSnafu { + reason: "logical table names must not contain path separators" + } + ); + ensure!( + info.catalog_name == physical_info.catalog_name + && info.schema_name == physical_info.schema_name + && info.meta.engine == METRIC_ENGINE_NAME + && info + .meta + .options + .extra_options + .contains_key(LOGICAL_TABLE_METADATA_KEY), + InvalidLogicalTableExportSnafu { + reason: format!( + "{name} is not a Metric logical table in the physical table schema" + ) + } + ); + let schema = table.schema().arrow_schema().clone(); + let indices = schema + .fields() + .iter() + .map(|field| { + ensure!( + field.name() != TABLE_ID && field.name() != TSID, + InvalidLogicalTableExportSnafu { + reason: "logical schema contains internal Metric columns" + } + ); + let index = physical_schema + .column_index_by_name(field.name()) + .with_context(|| InvalidLogicalTableExportSnafu { + reason: format!("physical schema lacks {}", field.name()), + })?; + ensure!( + physical_schema.arrow_schema().field(index).data_type() + == field.data_type(), + InvalidLogicalTableExportSnafu { + reason: format!("physical/logical type mismatch for {}", field.name()) + } + ); + Ok(index) + }) + .collect::>>()?; + scan_projection.extend(indices.iter().copied()); + ensure!( + logical_tables + .insert( + info.table_id(), + LogicalTableProjection { + name: name.clone(), + schema, + projection: indices, + } + ) + .is_none(), + InvalidLogicalTableExportSnafu { + reason: "duplicate logical table" + } + ); + } + let scan_projection = scan_projection.into_iter().collect::>(); + for file in logical_tables.values_mut() { + for index in &mut file.projection { + *index = scan_projection.binary_search(index).map_err(|_| { + error::UnexpectedSnafu { + violated: "logical column missing from physical projection", + } + .build() + })?; + } + } + Ok(Self { + physical_table: physical, + scan_projection, + logical_tables, + }) + } + + async fn validate_table_routes(&self, manager: &TableRouteManager) -> Result<()> { + let table_ids = self.logical_tables.keys().copied().collect::>(); + let routes = manager + .table_route_storage() + .batch_get(&table_ids) + .await + .context(error::TableMetadataManagerSnafu)?; + let physical_table_id = self.physical_table.table_info().table_id(); + for (table_id, route) in table_ids.into_iter().zip(routes) { + ensure!( + matches!(route, Some(TableRouteValue::Logical(route)) if route.physical_table_id() == physical_table_id), + InvalidLogicalTableExportSnafu { + reason: format!( + "logical table {table_id} does not belong to physical table {physical_table_id}" + ) + } + ); + } + Ok(()) + } + + fn build_plan(&self, time_range: Option<&TimestampRange>) -> Result { + let info = self.physical_table.table_info(); + let filters = self + .physical_table + .schema() + .timestamp_column() + .and_then(|column| { + common_query::logical_plan::build_filter_from_timestamp(&column.name, time_range) + }) + .into_iter() + .collect::>(); + let source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(self.physical_table.clone()), + ))); + let mut builder = LogicalPlanBuilder::scan_with_filters( + DfTableReference::full( + info.catalog_name.clone(), + info.schema_name.clone(), + info.name.clone(), + ), + source, + Some(self.scan_projection.clone()), + filters.clone(), + ) + .context(error::BuildDfLogicalPlanSnafu)?; + for filter in filters { + builder = builder + .filter(filter) + .context(error::BuildDfLogicalPlanSnafu)?; + } + builder + .sort(vec![col(TABLE_ID).sort(true, false)]) + .context(error::BuildDfLogicalPlanSnafu)? + .build() + .context(error::BuildDfLogicalPlanSnafu) + } +} + +/// Counts returned only after every selected logical file has been closed. +#[derive(Debug, Default, PartialEq, Eq)] +pub struct LogicalTableExportSummary { + pub rows: usize, + pub skipped_rows: usize, + pub files: usize, +} + +impl StatementExecutor { + /// Export selected logical tables to a fresh directory using one physical query. + /// The caller bounds concurrent units and owns retry/cleanup of this directory. + /// Cancellation waits for in-flight I/O before aborting the active upload. + /// Closed files remain an incomplete chunk until the caller publishes its + /// completion state. + #[allow(clippy::too_many_arguments)] + pub async fn export_logical_tables( + &self, + unit: &LogicalTableExport, + directory: &str, + connection: &HashMap, + time_range: Option<&TimestampRange>, + limits: LogicalTableExportLimits, + cancellation: &CancellationToken, + query_ctx: QueryContextRef, + ) -> Result { + limits.validate()?; + let (store, stream) = tokio::select! { + biased; + _ = cancellation.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(), + result = async { + unit.validate_table_routes(self.table_metadata_manager.table_route_manager()).await?; + let store = build_backend_for_write(&format!("{}/", directory.trim_end_matches('/')), connection, &self.local_file_access) + .await.context(error::BuildBackendSnafu)?; + let output = self.query_engine.execute(unit.build_plan(time_range)?, query_ctx) + .await.context(error::ExecLogicalPlanSnafu)?; + let stream = match output.data { + OutputData::Stream(stream) => stream, + OutputData::RecordBatches(batches) => batches.as_stream(), + _ => return error::UnexpectedSnafu { violated: "expected physical query rows" }.fail(), + }; + Ok((store, stream)) + } => result?, + }; + export_stream(unit, stream, &store, limits, cancellation).await + } +} + +async fn export_stream( + unit: &LogicalTableExport, + stream: SendableRecordBatchStream, + store: &ObjectStore, + limits: LogicalTableExportLimits, + cancellation: &CancellationToken, +) -> Result { + let mut active = None; + let result = write_tables(unit, stream, store, limits, cancellation, &mut active).await; + if result.is_err() + && let Some(writer) = active + && let Err(cleanup_error) = writer + .writer + .abort() + .await + .map_err(|error| map_writer_error(error, &writer.path)) + { + common_telemetry::warn!(cleanup_error; "Failed to clean up incomplete Metric export file"); + } + result +} + +fn check_cancelled(cancellation: &CancellationToken) -> Result<()> { + ensure!( + !cancellation.is_cancelled(), + error::LogicalTableExportCancelledSnafu + ); + Ok(()) +} + +async fn write_tables( + unit: &LogicalTableExport, + mut stream: SendableRecordBatchStream, + store: &ObjectStore, + limits: LogicalTableExportLimits, + cancellation: &CancellationToken, + active: &mut Option, +) -> Result { + let id_index = + stream + .schema() + .column_index_by_name(TABLE_ID) + .context(error::UnexpectedSnafu { + violated: "physical query omitted __table_id", + })?; + let mut summary = LogicalTableExportSummary::default(); + let mut previous = None; + let mut written = BTreeSet::new(); + loop { + let batch = tokio::select! { + biased; + _ = cancellation.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(), + batch = stream.next() => batch, + }; + let Some(batch) = batch else { + break; + }; + let batch = batch + .context(error::BuildRecordBatchSnafu)? + .into_df_record_batch(); + ensure!( + batch.get_array_memory_size() <= limits.input_batch_bytes, + LogicalTableExportResourceSnafu { + reason: "scan batch exceeds input byte budget" + } + ); + let ids = batch + .column(id_index) + .as_any() + .downcast_ref::() + .context(error::UnexpectedSnafu { + violated: "__table_id must be UInt32", + })?; + ensure!( + ids.null_count() == 0, + error::UnexpectedSnafu { + violated: "null __table_id" + } + ); + let mut start = 0; + while start < batch.num_rows() { + check_cancelled(cancellation)?; + let id = ids.value(start); + ensure!( + previous.is_none_or(|last| last <= id), + error::UnexpectedSnafu { + violated: "physical query is not ordered by __table_id" + } + ); + previous = Some(id); + let mut end = start + 1; + while end < batch.num_rows() && ids.value(end) == id { + end += 1; + } + if active.as_ref().is_some_and(|writer| writer.table_id != id) { + finish_active(active, cancellation).await?; + } + check_cancelled(cancellation)?; + if let Some(file) = unit.logical_tables.get(&id) { + if active.is_none() { + *active = Some(ActiveWriter::open(id, file, store, limits).await?); + written.insert(id); + summary.files += 1; + } + let projected = batch + .project(&file.projection) + .context(error::ProjectSchemaSnafu)?; + let writer = active.as_mut().context(error::UnexpectedSnafu { + violated: "missing logical writer", + })?; + let mut offset = start; + while offset < end { + let (expanded, consumed) = expand_bounded_slice( + projected.clone(), + file.schema.clone(), + offset, + end, + limits.conversion_bytes, + ) + .await?; + check_cancelled(cancellation)?; + // Do not drop an in-flight file operation before cleanup. + writer + .writer + .write(expanded, Some(cancellation)) + .await + .map_err(|error| map_writer_error(error, &writer.path))?; + check_cancelled(cancellation)?; + offset += consumed; + summary.rows += consumed; + } + } else { + summary.skipped_rows += end - start; + } + start = end; + } + } + finish_active(active, cancellation).await?; + for (&id, file) in &unit.logical_tables { + if !written.contains(&id) { + check_cancelled(cancellation)?; + *active = Some(ActiveWriter::open(id, file, store, limits).await?); + finish_active(active, cancellation).await?; + summary.files += 1; + } + } + check_cancelled(cancellation)?; + Ok(summary) +} + +struct ActiveWriter { + table_id: u32, + path: String, + writer: ParquetFileWriter, +} + +impl ActiveWriter { + async fn open( + id: u32, + table: &LogicalTableProjection, + store: &ObjectStore, + limits: LogicalTableExportLimits, + ) -> Result { + let path = format!("{}.parquet", table.name); + ensure!( + !store + .exists(&path) + .await + .context(error::ReadObjectSnafu { path: &path })?, + InvalidLogicalTableExportSnafu { + reason: format!("output already exists: {path}") + } + ); + let writer = ParquetFileWriter::open( + table.schema.clone(), + store.clone(), + &path, + 1, + Some(limits.writer), + ) + .await + .map_err(|error| map_writer_error(error, &path))?; + Ok(Self { + table_id: id, + path, + writer, + }) + } +} + +async fn finish_active( + active: &mut Option, + cancellation: &CancellationToken, +) -> Result<()> { + if let Some(writer) = active.as_mut() { + writer + .writer + .finish(Some(cancellation)) + .await + .map_err(|error| map_writer_error(error, &writer.path))?; + check_cancelled(cancellation)?; + *active = None; + } + Ok(()) +} + +async fn expand_bounded_slice( + batch: RecordBatch, + schema: SchemaRef, + start: usize, + end: usize, + budget: usize, +) -> Result<(RecordBatch, usize)> { + common_runtime::spawn_blocking_global(move || { + let len = rows_within_budget(&batch, start, end, budget)?; + let slice = batch.slice(start, len); + let arrays = slice + .columns() + .iter() + .zip(schema.fields()) + .map(|(array, field)| cast(array, field.data_type()).context(error::ComputeArrowSnafu)) + .collect::>>()?; + let expanded = RecordBatch::try_new(schema, arrays).context(error::ComputeArrowSnafu)?; + Ok((expanded, len)) + }) + .await + .context(error::JoinTaskSnafu)? +} + +fn map_writer_error(source: common_datasource::error::Error, path: &str) -> error::Error { + match source { + common_datasource::error::Error::ParquetWriteCancelled {} => { + error::LogicalTableExportCancelledSnafu.build() + } + common_datasource::error::Error::InvalidParquetWriterLimits {} => { + InvalidLogicalTableExportSnafu { + reason: "Parquet writer limits must be positive", + } + .build() + } + common_datasource::error::Error::ParquetWriterResource { reason } => { + LogicalTableExportResourceSnafu { reason }.build() + } + source => error::WriteStreamToFileSnafu { path }.into_error(source), + } +} + +// Count only selected logical values before dictionary expansion, including nested +// histogram lists/structs. Offset and validity overhead is charged per value. +fn estimate_value_size(array: &dyn Array, row: usize) -> Result { + if array.is_null(row) { + return Ok(32); + } + let bytes = match array.data_type() { + DataType::Boolean => 1, + DataType::Null => 0, + DataType::Utf8 => array.as_string::().value(row).len(), + DataType::LargeUtf8 => array.as_string::().value(row).len(), + DataType::Binary => array.as_binary::().value(row).len(), + DataType::LargeBinary => array.as_binary::().value(row).len(), + DataType::Struct(_) => { + array + .as_struct() + .columns() + .iter() + .try_fold(0usize, |sum, child| { + Ok::<_, error::Error>( + sum.saturating_add(estimate_value_size(child.as_ref(), row)?), + ) + })? + } + DataType::List(_) => { + let list = array.as_list::(); + let offsets = list.value_offsets(); + (offsets[row] as usize..offsets[row + 1] as usize).try_fold(0usize, |sum, index| { + Ok::<_, error::Error>( + sum.saturating_add(estimate_value_size(list.values().as_ref(), index)?), + ) + })? + } + DataType::Dictionary(_, _) => { + downcast_dictionary_array! { + array => { + match array.key(row) { + Some(index) => estimate_value_size(array.values().as_ref(), index)?, + None => 0, + } + }, + _ => return error::UnexpectedSnafu { violated: "invalid dictionary array" }.fail(), + } + } + other => other + .primitive_width() + .with_context(|| InvalidLogicalTableExportSnafu { + reason: format!("unsupported Metric Parquet type: {other}"), + })?, + }; + Ok(bytes.saturating_add(16)) +} + +fn rows_within_budget( + batch: &RecordBatch, + start: usize, + end: usize, + budget: usize, +) -> Result { + let mut bytes = 0usize; + let mut row = start; + while row < end { + let row_bytes = batch.columns().iter().try_fold(0usize, |sum, array| { + Ok::<_, error::Error>(sum.saturating_add(estimate_value_size(array.as_ref(), row)?)) + })?; + if row_bytes > budget.saturating_sub(bytes) { + break; + } + bytes += row_bytes; + row += 1; + } + ensure!( + row > start, + LogicalTableExportResourceSnafu { + reason: "one expanded logical row exceeds conversion byte budget" + } + ); + Ok(row - start) +} + +#[cfg(test)] +mod tests; diff --git a/src/operator/src/statement/export_logical_tables/tests.rs b/src/operator/src/statement/export_logical_tables/tests.rs new file mode 100644 index 00000000000..a85418fa2e6 --- /dev/null +++ b/src/operator/src/statement/export_logical_tables/tests.rs @@ -0,0 +1,578 @@ +// Copyright 2023 Greptime Team +// +// 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 arrow::array::{ + ArrayRef, DictionaryArray, Float64Array, Int32Array, ListArray, StringArray, StructArray, +}; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::{Field, Schema, UInt32Type}; +use bytes::Bytes; +use common_recordbatch::{RecordBatch as GreptimeRecordBatch, RecordBatches}; +use datafusion::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use table::test_util::EmptyTable; +use table::test_util::table_info::test_table_info; + +use super::*; + +fn table(id: u32, name: &str, fields: Vec, physical: bool) -> TableRef { + let schema = Arc::new(datatypes::schema::Schema::try_from(Schema::new(fields)).unwrap()); + let mut info = test_table_info(id, name, "public", "greptime", schema); + info.meta.engine = METRIC_ENGINE_NAME.into(); + let (key, value) = if physical { + (PHYSICAL_TABLE_METADATA_KEY, "") + } else { + (LOGICAL_TABLE_METADATA_KEY, "phy") + }; + info.meta + .options + .extra_options + .insert(key.into(), value.into()); + EmptyTable::from_table_info(&info) +} + +fn export_limits() -> LogicalTableExportLimits { + let mut limits = LogicalTableExportLimits::default(); + limits.writer.row_group_rows = 1; + limits +} + +fn unit() -> LogicalTableExport { + let ts = Field::new( + "ts", + DataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + false, + ); + let host = Field::new("host", DataType::Utf8, true); + let a = Field::new("a", DataType::Float64, true); + let b = Field::new("b", DataType::Float64, true); + let physical = table( + 1024, + "phy", + vec![ + Field::new(TABLE_ID, DataType::UInt32, false), + ts.clone(), + host.clone(), + a.clone(), + b.clone(), + ], + true, + ); + let tables = vec![ + table( + 1025, + "cpu.v1", + vec![host.clone(), a.clone(), ts.clone()], + false, + ), + table(1026, "empty", vec![ts.clone(), a], false), + table(1027, "requests", vec![ts, b, host], false), + ]; + LogicalTableExport::try_new(physical, &tables).unwrap() +} + +fn batch(ids: Vec>, hosts: Vec>) -> RecordBatch { + use arrow::array::TimestampMillisecondArray; + let rows = ids.len(); + let host: DictionaryArray = hosts.into_iter().collect(); + RecordBatch::try_from_iter_with_nullable([ + (TABLE_ID, Arc::new(UInt32Array::from(ids)) as ArrayRef, true), + ( + "ts", + Arc::new(TimestampMillisecondArray::from(vec![100; rows])), + false, + ), + ("host", Arc::new(host), true), + ( + "a", + Arc::new(Float64Array::from(vec![Some(1.5); rows])), + true, + ), + ("b", Arc::new(Float64Array::from(vec![None; rows])), true), + ]) + .unwrap() +} + +fn stream(batches: Vec) -> SendableRecordBatchStream { + let schema = Arc::new(datatypes::schema::Schema::try_from(batches[0].schema()).unwrap()); + let batches = batches + .into_iter() + .map(|batch| GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch)) + .collect(); + RecordBatches::try_new(schema, batches).unwrap().as_stream() +} + +async fn read(store: &ObjectStore, path: &str) -> (SchemaRef, Vec) { + let bytes = store.read(path).await.unwrap().to_bytes(); + let builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap(); + let schema = builder.schema().clone(); + ( + schema, + builder + .build() + .unwrap() + .collect::>() + .unwrap(), + ) +} + +#[tokio::test] +async fn routes_across_batches_and_writes_empty_files() { + let unit = unit(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let batches = vec![ + batch(vec![Some(1025), Some(1025)], vec![Some(""), None]), + batch( + vec![Some(1025), Some(1027), Some(1028)], + vec![Some("a"), Some("b"), Some("unselected")], + ), + ]; + let result = write_tables( + &unit, + stream(batches), + &store, + export_limits(), + &CancellationToken::new(), + &mut None, + ) + .await + .unwrap(); + assert_eq!( + result, + LogicalTableExportSummary { + rows: 4, + skipped_rows: 1, + files: 3 + } + ); + let (schema, cpu) = read(&store, "cpu.v1.parquet").await; + assert_eq!(schema.fields(), unit.logical_tables[&1025].schema.fields()); + let values: Vec<_> = cpu + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_string::() + .iter() + .map(|s| s.map(str::to_owned)) + }) + .collect(); + assert_eq!(values, vec![Some("".into()), None, Some("a".into())]); + let (schema, empty) = read(&store, "empty.parquet").await; + assert_eq!(schema.fields(), unit.logical_tables[&1026].schema.fields()); + assert!(empty.is_empty()); + let (schema, requests) = read(&store, "requests.parquet").await; + assert_eq!(schema.fields(), unit.logical_tables[&1027].schema.fields()); + assert_eq!(requests[0].column(1).null_count(), 1); +} + +#[tokio::test] +async fn rejects_invalid_order_ids_and_resource_exhaustion() { + let unit = unit(); + for (batches, limits, message) in [ + ( + vec![ + batch(vec![Some(1027)], vec![Some("a")]), + batch(vec![Some(1025)], vec![Some("b")]), + ], + LogicalTableExportLimits::default(), + "not ordered", + ), + ( + vec![batch(vec![None], vec![None])], + LogicalTableExportLimits::default(), + "null __table_id", + ), + ( + vec![batch(vec![Some(1025)], vec![Some("a")])], + LogicalTableExportLimits { + input_batch_bytes: 1, + ..Default::default() + }, + "input byte budget", + ), + ( + vec![batch(vec![Some(1025)], vec![Some("too big")])], + LogicalTableExportLimits { + conversion_bytes: 1, + ..Default::default() + }, + "one expanded logical row", + ), + ] { + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let mut active = None; + let err = write_tables( + &unit, + stream(batches), + &store, + limits, + &CancellationToken::new(), + &mut active, + ) + .await + .unwrap_err(); + assert!(err.to_string().contains(message), "{err}"); + if let Some(writer) = active { + writer.writer.abort().await.unwrap(); + } + } +} + +#[tokio::test] +async fn existing_outputs_are_not_overwritten() { + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + store.write("cpu.v1.parquet", "keep").await.unwrap(); + let err = write_tables( + &unit(), + stream(vec![batch(vec![Some(1025)], vec![None])]), + &store, + LogicalTableExportLimits::default(), + &CancellationToken::new(), + &mut None, + ) + .await + .unwrap_err(); + assert!(err.to_string().contains("already exists")); + assert_eq!( + store.read("cpu.v1.parquet").await.unwrap().to_bytes(), + Bytes::from_static(b"keep") + ); +} + +#[test] +fn dictionary_and_nested_histogram_values_are_bounded_before_expansion() { + let dictionary = DictionaryArray::::try_new( + UInt32Array::from(vec![0, 0, 0]), + Arc::new(StringArray::from(vec!["x".repeat(4096)])), + ) + .unwrap(); + let nested = ListArray::new( + Arc::new(Field::new("item", DataType::Int32, true)), + OffsetBuffer::new(vec![0i32, 2, 3, 5].into()), + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])), + None, + ); + let histogram = StructArray::from(vec![( + Arc::new(Field::new("buckets", nested.data_type().clone(), true)), + Arc::new(nested) as ArrayRef, + )]); + let batch = RecordBatch::try_from_iter([ + ("tag", Arc::new(dictionary) as ArrayRef), + ("histogram", Arc::new(histogram)), + ]) + .unwrap(); + assert_eq!(rows_within_budget(&batch, 0, 3, 4300).unwrap(), 1); + // The dictionary and container overhead fit; the nested list elements do not. + assert!(rows_within_budget(&batch, 0, 3, 4180).is_err()); + assert_eq!(rows_within_budget(&batch, 0, 3, 15000).unwrap(), 3); +} + +#[test] +fn validates_selected_schemas_and_projects_only_selected_columns() { + let unit = unit(); + let selected = table( + 1025, + "only", + vec![Field::new("a", DataType::Float64, true)], + false, + ); + let one = + LogicalTableExport::try_new(unit.physical_table.clone(), std::slice::from_ref(&selected)) + .unwrap(); + assert_eq!(one.scan_projection, vec![0, 3]); + assert_eq!(one.logical_tables[&1025].projection, vec![1]); + assert!( + LogicalTableExport::try_new(unit.physical_table.clone(), &[selected.clone(), selected]) + .is_err() + ); + let unsafe_name = table(1030, "a/b", vec![], false); + assert!(LogicalTableExport::try_new(unit.physical_table.clone(), &[unsafe_name]).is_err()); + let wrong_type = table( + 1030, + "wrong", + vec![Field::new("a", DataType::Int32, true)], + false, + ); + assert!(LogicalTableExport::try_new(unit.physical_table, &[wrong_type]).is_err()); +} + +#[tokio::test] +async fn cancellation_drops_input_and_aborts_active_upload() { + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let unit = unit(); + let batch = batch(vec![Some(1025)], vec![Some("a")]); + let schema = batch.schema(); + let token = CancellationToken::new(); + let trigger = token.clone(); + let batches = async_stream::stream! { + yield Ok(batch); + // Resuming the source proves the previous batch reached the writer. + trigger.cancel(); + std::future::pending::<()>().await; + }; + let df_stream = + datafusion::physical_plan::stream::RecordBatchStreamAdapter::new(schema, batches); + let stream = + common_recordbatch::adapter::RecordBatchStreamAdapter::try_new(Box::pin(df_stream)) + .unwrap(); + let result = export_stream(&unit, Box::pin(stream), &store, export_limits(), &token).await; + assert!(matches!( + result, + Err(error::Error::LogicalTableExportCancelled { .. }) + )); + assert!(!store.exists("cpu.v1.parquet").await.unwrap()); + assert!(!store.exists("empty.parquet").await.unwrap()); +} + +#[tokio::test] +async fn native_histogram_parquet_roundtrip() { + fn sample(data_type: &DataType) -> ArrayRef { + match data_type { + DataType::Struct(fields) => Arc::new(StructArray::new( + fields.clone(), + fields.iter().map(|f| sample(f.data_type())).collect(), + None, + )), + DataType::List(field) => Arc::new(ListArray::new( + field.clone(), + OffsetBuffer::new(vec![0i32, 1, 2].into()), + sample(field.data_type()), + None, + )), + DataType::Int32 => Arc::new(Int32Array::from(vec![Some(1), None])), + DataType::Int64 => Arc::new(arrow::array::Int64Array::from(vec![Some(2), None])), + DataType::Float64 => Arc::new(Float64Array::from(vec![Some(3.5), None])), + DataType::Timestamp(_, _) => { + Arc::new(arrow::array::TimestampMillisecondArray::from(vec![ + Some(4), + None, + ])) + } + other => panic!("unexpected histogram field {other}"), + } + } + let data_type = common_query::native_histogram::native_histogram_arrow_type(); + let field = Field::new("histogram", data_type.clone(), true); + let unit = LogicalTableExport::try_new( + table( + 1024, + "phy", + vec![Field::new(TABLE_ID, DataType::UInt32, false), field.clone()], + true, + ), + &[table(1025, "histogram", vec![field], false)], + ) + .unwrap(); + let batch = RecordBatch::try_from_iter_with_nullable([ + ( + TABLE_ID, + Arc::new(UInt32Array::from(vec![1025; 2])) as ArrayRef, + false, + ), + ("histogram", sample(&data_type), true), + ]) + .unwrap(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let summary = export_stream( + &unit, + stream(vec![batch.clone()]), + &store, + export_limits(), + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(summary.rows, 2); + let (schema, batches) = read(&store, "histogram.parquet").await; + assert_eq!(schema.fields(), unit.logical_tables[&1025].schema.fields()); + assert_eq!( + arrow::compute::concat_batches(&schema, &batches).unwrap(), + batch.project(&[1]).unwrap() + ); +} + +struct PausedFileWriter { + inner: Option, + started: Arc, + release: Arc, +} + +impl object_store::layers::mock::oio::Write for PausedFileWriter { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + let mut inner = self.inner.take().unwrap(); + let started = self.started.clone(); + let release = self.release.clone(); + // Like SecureFs's blocking open, this task survives a dropped I/O future. + let (inner, result) = tokio::spawn(async move { + started.notify_one(); + release.notified().await; + let result = inner.write(bytes).await; + (inner, result) + }) + .await + .unwrap(); + self.inner = Some(inner); + result + } + + async fn close(&mut self) -> object_store::Result { + self.inner.as_mut().unwrap().close().await + } + + async fn abort(&mut self) -> object_store::Result<()> { + match self.inner.as_mut() { + Some(inner) => inner.abort().await, + None => Err(object_store::Error::new( + object_store::ErrorKind::Unsupported, + "open is pending", + )), + } + } +} + +#[tokio::test] +async fn cancellation_waits_for_file_creation_before_cleanup() { + use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory}; + let directory = common_test_util::temp_dir::create_temp_dir("metric_export_pending_open"); + let access = + common_datasource::object_store::LocalFileAccess::sandboxed(directory.path()).unwrap(); + let store = build_backend_for_write( + &format!("{}/", directory.path().display()), + &HashMap::new(), + &access, + ) + .await + .unwrap(); + let started = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let factory: MockWriterFactory = Arc::new({ + let started = started.clone(); + let release = release.clone(); + move |_, _, inner| { + Box::new(PausedFileWriter { + inner: Some(inner), + started: started.clone(), + release: release.clone(), + }) + } + }); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let cancellation = CancellationToken::new(); + let unit = unit(); + let export = export_stream( + &unit, + stream(vec![batch(vec![Some(1025)], vec![Some("a")])]), + &store, + LogicalTableExportLimits::default(), + &cancellation, + ); + tokio::pin!(export); + tokio::select! { + result = &mut export => panic!("export completed before the file open: {result:?}"), + _ = started.notified() => {}, + } + cancellation.cancel(); + assert!(futures::poll!(&mut export).is_pending()); + release.notify_one(); + let result = export.await; + assert!(matches!( + result, + Err(error::Error::LogicalTableExportCancelled { .. }) + )); + assert!(!store.exists("cpu.v1.parquet").await.unwrap()); +} + +struct FailedAbortWriter(object_store::layers::mock::oio::Writer); + +impl object_store::layers::mock::oio::Write for FailedAbortWriter { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.0.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + self.0.close().await + } + async fn abort(&mut self) -> object_store::Result<()> { + Err(object_store::Error::new( + object_store::ErrorKind::PermissionDenied, + "injected abort failure", + )) + } +} + +#[tokio::test] +async fn cleanup_failure_preserves_resource_error() { + use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory}; + let factory: MockWriterFactory = Arc::new(|_, _, inner| Box::new(FailedAbortWriter(inner))); + let store = ObjectStore::new(object_store::services::Memory::default()) + .unwrap() + .layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let mut limits = export_limits(); + limits.writer.max_row_groups = 1; + let result = export_stream( + &unit(), + stream(vec![batch(vec![Some(1025); 2], vec![Some("a"); 2])]), + &store, + limits, + &CancellationToken::new(), + ) + .await; + assert!(matches!( + result, + Err(error::Error::LogicalTableExportResource { .. }) + )); +} + +#[tokio::test] +async fn validates_membership_by_table_route() { + use common_meta::kv_backend::TxnService; + use common_meta::kv_backend::memory::MemoryKvBackend; + + let unit = unit(); + for later_physical_id in [None, Some(2048), Some(1024)] { + let kv = Arc::new(MemoryKvBackend::default()); + let manager = TableRouteManager::new(kv.clone()); + for (&table_id, physical_id) in + unit.logical_tables + .keys() + .zip([Some(1024), later_physical_id, Some(1024)]) + { + if let Some(physical_id) = physical_id { + let (txn, _) = manager + .table_route_storage() + .build_create_txn(table_id, &TableRouteValue::logical(physical_id)) + .unwrap(); + assert!(kv.txn(txn).await.unwrap().succeeded); + } + } + let result = unit.validate_table_routes(&manager).await; + if later_physical_id == Some(1024) { + result.unwrap(); + } else { + assert!(matches!( + result, + Err(error::Error::InvalidLogicalTableExport { .. }) + )); + } + } +} diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index a2d6c70e394..d1b0774af9e 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -126,6 +126,7 @@ session = { workspace = true, features = ["testing"] } similar-asserts.workspace = true store-api.workspace = true tokio-postgres = { workspace = true } +tokio-util.workspace = true url = "2.3" urlencoding = "2.1" yaml-rust = "0.4" diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index e40c707fe98..3e53b6104a0 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -35,6 +35,7 @@ use client::Client; use client::client_manager::NodeClients; use cmd::frontend::create_heartbeat_task; use common_base::Plugins; +use common_datasource::object_store::LocalFileAccess; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_meta::DatanodeId; use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder}; @@ -172,6 +173,7 @@ pub struct GreptimeDbClusterBuilder { frontend_auto_create_table: bool, shared_home_dir: Option>, meta_selector: Option, + local_file_access: LocalFileAccess, } impl GreptimeDbClusterBuilder { @@ -206,6 +208,7 @@ impl GreptimeDbClusterBuilder { frontend_auto_create_table: true, shared_home_dir: None, meta_selector: None, + local_file_access: LocalFileAccess::default(), } } @@ -269,6 +272,12 @@ impl GreptimeDbClusterBuilder { self } + /// Configure the frontend COPY sandbox for filesystem integration tests. + pub fn with_local_file_access(mut self, access: LocalFileAccess) -> Self { + self.local_file_access = access; + self + } + pub async fn build_with( &self, datanode_options: Vec, @@ -512,6 +521,7 @@ impl GreptimeDbClusterBuilder { meta_client.clone(), Arc::new(ProcessManager::new(fe_opts.grpc.server_addr.clone(), None)), ) + .with_local_file_access(self.local_file_access.clone()) .with_local_cache_invalidator(cache_registry) .try_build() .await diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs new file mode 100644 index 00000000000..4aeb745b885 --- /dev/null +++ b/tests-integration/tests/export_logical_tables.rs @@ -0,0 +1,198 @@ +// Copyright 2023 Greptime Team +// +// 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 std::sync::Arc; + +use common_query::OutputData; +use common_time::Timestamp; +use common_time::range::TimestampRange; +use frontend::instance::Instance; +use operator::statement::export_logical_tables::{LogicalTableExport, LogicalTableExportLimits}; +use session::context::QueryContext; +use tests_integration::cluster::GreptimeDbClusterBuilder; +use tests_integration::standalone::GreptimeDbStandaloneBuilder; +use tests_integration::test_util::execute_sql as sql; +use tokio_util::sync::CancellationToken; + +async fn table(instance: &Arc, name: &str) -> table::TableRef { + instance + .catalog_manager() + .table("greptime", "public", name, Some(&QueryContext::arc())) + .await + .unwrap() + .unwrap() +} + +async fn values(instance: &Arc, query: &str) -> Vec> { + let output = sql(instance, query).await; + let batches = match output.data { + OutputData::Stream(stream) => common_recordbatch::util::collect_batches(stream) + .await + .unwrap(), + OutputData::RecordBatches(batches) => batches, + _ => panic!("expected query rows"), + }; + batches + .iter() + .flat_map(|b| { + let columns = datatypes::vectors::Helper::try_into_vectors(b.columns()).unwrap(); + (0..b.num_rows()) + .map(|row| columns.iter().map(|col| col.get(row)).collect::>()) + .collect::>() + }) + .collect() +} + +async fn roundtrip(instance: &Arc) { + let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); + for (physical, encoding) in [("phy", "dense"), ("other_phy", "sparse")] { + sql(instance, &format!("CREATE TABLE {physical} (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY) PARTITION ON COLUMNS (host) (host < 'm', host >= 'm') ENGINE=metric WITH (physical_metric_table='', primary_key_encoding='{encoding}')")).await; + for (suffix, extra, key) in [ + ("cpu.v1", "zone_tag STRING,", ", zone_tag"), + ("requests", "service_tag STRING,", ", service_tag"), + ("empty", "", ""), + ] { + let name = format!("{physical}_{suffix}"); + sql(instance, &format!("CREATE TABLE \"{name}\" (host STRING, {extra} val DOUBLE, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host{key})) ENGINE=metric WITH (on_physical_table='{physical}')")).await; + if suffix != "empty" { + sql(instance, &format!("INSERT INTO \"{name}\" (host,val,ts) VALUES ('a',1,1),('a',NULL,2),('z',3,3),('z',4,4)")).await; + } + } + // This live but unselected table models rows outside the routing whitelist. + sql(instance, &format!("CREATE TABLE {physical}_excluded (host STRING, huge_tag STRING, val DOUBLE, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host, huge_tag)) ENGINE=metric WITH (on_physical_table='{physical}')")).await; + sql( + instance, + &format!("INSERT INTO {physical}_excluded (host, huge_tag, val, ts) VALUES ('z','ignore',9,2)"), + ) + .await; + let names = ["cpu.v1", "requests", "empty"].map(|suffix| format!("{physical}_{suffix}")); + let tables = vec![ + table(instance, &names[0]).await, + table(instance, &names[1]).await, + table(instance, &names[2]).await, + ]; + let renamed = format!("renamed_{physical}"); + sql( + instance, + &format!("ALTER TABLE {physical} RENAME {renamed}"), + ) + .await; + let unit = LogicalTableExport::try_new(table(instance, &renamed).await, &tables).unwrap(); + let range = + TimestampRange::new(Timestamp::new_millisecond(2), Timestamp::new_millisecond(4)) + .unwrap(); + sql(instance, &format!("CREATE TABLE target_{physical} (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY) ENGINE=metric WITH (physical_metric_table='')")).await; + let mut limits = LogicalTableExportLimits::default(); + limits.writer.row_group_rows = 1; + for partitions in [1, 2, 4] { + let directory = destination.path().join(format!("{physical}_{partitions}")); + let mut ctx = QueryContext::with("greptime", "public"); + ctx.set_extension( + query::datafusion::QUERY_PARALLELISM_HINT, + partitions.to_string(), + ); + let summary = instance + .statement_executor() + .export_logical_tables( + &unit, + directory.to_str().unwrap(), + &Default::default(), + Some(&range), + limits, + &CancellationToken::new(), + Arc::new(ctx), + ) + .await + .unwrap(); + assert_eq!(summary.rows, 4); + assert_eq!(summary.files, 3); + assert_eq!(summary.skipped_rows, 1); + for (index, name) in names.iter().enumerate() { + let restored = format!("restore_{physical}_{partitions}_{index}"); + let (extra, key) = [ + ("zone_tag STRING,", ", zone_tag"), + ("service_tag STRING,", ", service_tag"), + ("", ""), + ][index]; + sql(instance, &format!("CREATE TABLE {restored} (host STRING, {extra} val DOUBLE, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host{key})) ENGINE=metric WITH (on_physical_table='target_{physical}')")).await; + sql( + instance, + &format!( + "COPY {restored} FROM '{}/{}.parquet' WITH (FORMAT='parquet')", + directory.display(), + name + ), + ) + .await; + let expected = values( + instance, + &format!("SELECT * FROM \"{name}\" WHERE ts >= 2 AND ts < 4 ORDER BY host, ts"), + ) + .await; + let actual = values( + instance, + &format!("SELECT * FROM {restored} ORDER BY host, ts"), + ) + .await; + assert_eq!(actual, expected); + } + } + let cancellation = CancellationToken::new(); + cancellation.cancel(); + let directory = destination.path().join(format!("cancel_{physical}")); + let result = instance + .statement_executor() + .export_logical_tables( + &unit, + directory.to_str().unwrap(), + &Default::default(), + None, + LogicalTableExportLimits::default(), + &cancellation, + QueryContext::arc(), + ) + .await; + assert!(matches!( + result, + Err(operator::error::Error::LogicalTableExportCancelled { .. }) + )); + assert!(!directory.exists()); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn physical_export_standalone_roundtrip() { + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new("physical_export") + .build() + .await; + roundtrip(standalone.fe_instance()).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn physical_export_distributed_roundtrip() { + common_telemetry::init_default_ut_logging(); + let cluster = GreptimeDbClusterBuilder::new("physical_export") + .await + .with_datanodes(2) + .with_local_file_access( + common_datasource::object_store::LocalFileAccess::sandboxed( + common_test_util::find_workspace_path("."), + ) + .unwrap(), + ) + .build(false) + .await; + roundtrip(cluster.fe_instance()).await; +} diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 28a81155724..b61b6cd3e19 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -36,6 +36,7 @@ mod repartition; mod repartition_event; #[macro_use] mod repartition_expr_version; +mod export_logical_tables; mod mysql; mod reconciliation_event; mod view_ddl_event;