mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-09-26 23:25:33 +00:00
feat: export logical tables from Metric physical scans (#9159)
* feat: add physical Metric table exporter Signed-off-by: jeremyhi <fengjiachun@gmail.com> * fix: drain Metric export writes before cancellation cleanup Signed-off-by: jeremyhi <fengjiachun@gmail.com> * perf: construct Metric export error context lazily Signed-off-by: jeremyhi <fengjiachun@gmail.com> * refactor: name the logical table export entry point Signed-off-by: jeremyhi <fengjiachun@gmail.com> * refactor: share Parquet writer for logical table exports Signed-off-by: jeremyhi <fengjiachun@gmail.com> * fix: preserve Parquet destinations and cancellation boundaries Signed-off-by: jeremyhi <fengjiachun@gmail.com> * refactor: clarify logical table export field names Signed-off-by: jeremyhi <fengjiachun@gmail.com> * refactor: clarify logical table export helper responsibilities Signed-off-by: jeremyhi <fengjiachun@gmail.com> * fix: validate logical export membership by table ID Signed-off-by: jeremyhi <fengjiachun@gmail.com> * test: simplify logical table export coverage and strengthen assertions Signed-off-by: jeremyhi <fengjiachun@gmail.com> --------- Signed-off-by: jeremyhi <fengjiachun@gmail.com>
This commit is contained in:
jeremyhi
Generated
+1
@@ -15229,6 +15229,7 @@ dependencies = [
|
||||
"tokio",
|
||||
"tokio-postgres",
|
||||
"tokio-stream",
|
||||
"tokio-util",
|
||||
"tonic 0.14.2",
|
||||
"tower 0.5.2",
|
||||
"url",
|
||||
|
||||
@@ -47,3 +47,4 @@ url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
common-test-util.workspace = true
|
||||
object-store = { workspace = true, features = ["testing"] }
|
||||
|
||||
@@ -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 { .. }
|
||||
|
||||
@@ -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<SharedBuffer> {
|
||||
/// 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<usize> {
|
||||
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;
|
||||
|
||||
@@ -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<ArrowWriter<Vec<u8>>>,
|
||||
sink: Writer,
|
||||
store: ObjectStore,
|
||||
path: String,
|
||||
limits: Option<ParquetWriterLimits>,
|
||||
}
|
||||
|
||||
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<ParquetWriterLimits>,
|
||||
) -> Result<Self> {
|
||||
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<u8>) -> 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<Bytes> {
|
||||
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::<Vec<_>>();
|
||||
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::<std::result::Result<Vec<_>, _>>()
|
||||
.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::<usize>()
|
||||
);
|
||||
let reader = read(&store, path).await;
|
||||
assert_eq!(reader.schema().fields(), batch.schema().fields());
|
||||
let actual = reader
|
||||
.build()
|
||||
.unwrap()
|
||||
.collect::<std::result::Result<Vec<_>, _>>()
|
||||
.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<tokio::sync::Notify>,
|
||||
release: Arc<tokio::sync::Notify>,
|
||||
closed: Arc<std::sync::atomic::AtomicBool>,
|
||||
}
|
||||
|
||||
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<object_store::layers::mock::Metadata> {
|
||||
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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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(),
|
||||
}
|
||||
|
||||
@@ -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<usize>,
|
||||
logical_tables: BTreeMap<TableId, LogicalTableProjection>,
|
||||
}
|
||||
|
||||
struct LogicalTableProjection {
|
||||
name: String,
|
||||
schema: SchemaRef,
|
||||
projection: Vec<usize>,
|
||||
}
|
||||
|
||||
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<Self> {
|
||||
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::<Result<Vec<_>>>()?;
|
||||
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::<Vec<_>>();
|
||||
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::<Vec<_>>();
|
||||
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<LogicalPlan> {
|
||||
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::<Vec<_>>();
|
||||
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<String, String>,
|
||||
time_range: Option<&TimestampRange>,
|
||||
limits: LogicalTableExportLimits,
|
||||
cancellation: &CancellationToken,
|
||||
query_ctx: QueryContextRef,
|
||||
) -> Result<LogicalTableExportSummary> {
|
||||
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<LogicalTableExportSummary> {
|
||||
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<ActiveWriter>,
|
||||
) -> Result<LogicalTableExportSummary> {
|
||||
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::<UInt32Array>()
|
||||
.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<Self> {
|
||||
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<ActiveWriter>,
|
||||
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::<Result<Vec<_>>>()?;
|
||||
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<usize> {
|
||||
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::<i32>().value(row).len(),
|
||||
DataType::LargeUtf8 => array.as_string::<i64>().value(row).len(),
|
||||
DataType::Binary => array.as_binary::<i32>().value(row).len(),
|
||||
DataType::LargeBinary => array.as_binary::<i64>().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::<i32>();
|
||||
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<usize> {
|
||||
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;
|
||||
@@ -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<Field>, 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<Option<u32>>, hosts: Vec<Option<&str>>) -> RecordBatch {
|
||||
use arrow::array::TimestampMillisecondArray;
|
||||
let rows = ids.len();
|
||||
let host: DictionaryArray<UInt32Type> = 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<RecordBatch>) -> 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<RecordBatch>) {
|
||||
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::<std::result::Result<_, _>>()
|
||||
.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::<i32>()
|
||||
.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::<UInt32Type>::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<object_store::layers::mock::oio::Writer>,
|
||||
started: Arc<tokio::sync::Notify>,
|
||||
release: Arc<tokio::sync::Notify>,
|
||||
}
|
||||
|
||||
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<object_store::layers::mock::Metadata> {
|
||||
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<object_store::layers::mock::Metadata> {
|
||||
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 { .. })
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
|
||||
@@ -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<Arc<TempDir>>,
|
||||
meta_selector: Option<SelectorRef>,
|
||||
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<DatanodeOptions>,
|
||||
@@ -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
|
||||
|
||||
@@ -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<Instance>, name: &str) -> table::TableRef {
|
||||
instance
|
||||
.catalog_manager()
|
||||
.table("greptime", "public", name, Some(&QueryContext::arc()))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn values(instance: &Arc<Instance>, query: &str) -> Vec<Vec<datatypes::value::Value>> {
|
||||
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::<Vec<_>>())
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn roundtrip(instance: &Arc<Instance>) {
|
||||
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;
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user