operator/statement/
import_packed.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::sync::Arc;
17
18use common_datasource::file_format::Format;
19use common_datasource::file_format::parquet::packed_reader::{
20    PackReadWindows, PackedParquetReader, WINDOW_SIZE,
21};
22use common_datasource::object_store::{BuiltBackend, build_backend_with_path};
23use common_datasource::packed_snapshot::{ObjectKind, PACK_INDEX_FILE, PackIndex};
24use common_error::ext::BoxedError;
25use common_query::Output;
26use futures::stream::FuturesUnordered;
27use futures::{FutureExt, StreamExt};
28use session::context::QueryContextRef;
29use snafu::ResultExt;
30use table::TableRef;
31use table::requests::CopyDatabaseRequest;
32use tokio_util::compat::FuturesAsyncReadCompatExt;
33use tokio_util::sync::CancellationToken;
34
35use crate::error::{self, Result};
36use crate::statement::StatementExecutor;
37use crate::statement::database_copy::validate_database_directory;
38
39/// All table identities are resolved before the frontend authorizes any write.
40pub struct PreparedPackedImport {
41    backend: BuiltBackend,
42    index: PackIndex,
43    parallelism: usize,
44    pub tables: Vec<TableRef>,
45}
46
47impl StatementExecutor {
48    pub async fn prepare_packed_import(
49        &self,
50        req: CopyDatabaseRequest,
51        ctx: &QueryContextRef,
52    ) -> Result<PreparedPackedImport> {
53        validate_database_directory(&req.location)?;
54        if req.with.get("metric_data_layout").map(String::as_str) != Some("packed")
55            || !matches!(
56                Format::try_from(&req.with).context(error::ParseFileFormatSnafu)?,
57                Format::Parquet(_)
58            )
59            || req
60                .with
61                .get("continue_on_error")
62                .is_some_and(|v| v != "false")
63        {
64            return error::InvalidCopyParameterSnafu {
65                key: "metric_data_layout",
66                value: "packed requires parquet and continue_on_error=false",
67            }
68            .fail();
69        }
70        let backend =
71            build_backend_with_path(&req.location, &req.connection, &self.local_file_access)
72                .await
73                .context(error::BuildBackendSnafu)?;
74        if !backend
75            .is_file(PACK_INDEX_FILE)
76            .await
77            .context(error::ReadObjectSnafu {
78                path: PACK_INDEX_FILE,
79            })?
80        {
81            return error::InvalidCopyParameterSnafu {
82                key: "index",
83                value: "not a regular file",
84            }
85            .fail();
86        }
87        let bytes = backend
88            .object_store
89            .read(PACK_INDEX_FILE)
90            .await
91            .context(error::ReadObjectSnafu {
92                path: PACK_INDEX_FILE,
93            })?
94            .to_bytes();
95        let mut index: PackIndex = serde_json::from_slice(&bytes).map_err(|e| {
96            error::InvalidCopyParameterSnafu {
97                key: "pack_index",
98                value: e.to_string(),
99            }
100            .build()
101        })?;
102        index
103            .validate()
104            .map_err(BoxedError::new)
105            .context(error::ExternalSnafu)?;
106        index
107            .tables
108            .sort_by(|a, b| (&a.object, a.offset).cmp(&(&b.object, b.offset)));
109        let names: Vec<_> = index.tables.iter().map(|t| t.table_name.clone()).collect();
110        let mut tables = self
111            .capture_database_export_tables(&req, Some(&names), ctx)
112            .await?;
113        if tables.len() != names.len() {
114            return error::InvalidCopyParameterSnafu {
115                key: "index",
116                value: "physical tables and views are not data targets",
117            }
118            .fail();
119        }
120        let positions: HashMap<_, _> = names
121            .iter()
122            .enumerate()
123            .map(|(i, n)| (n.as_str(), i))
124            .collect();
125        tables.sort_by_key(|table| positions.get(table.table_info().name.as_str()).copied());
126        for object in &index.objects {
127            let meta = backend
128                .object_store
129                .stat(&object.path)
130                .await
131                .context(error::ReadObjectSnafu { path: &object.path })?;
132            if !backend
133                .is_file_with_mode(&object.path, meta.mode())
134                .await
135                .context(error::ReadObjectSnafu { path: &object.path })?
136            {
137                return error::InvalidCopyParameterSnafu {
138                    key: "object",
139                    value: &object.path,
140                }
141                .fail();
142            }
143            if meta.content_length() != object.length {
144                return error::InvalidCopyParameterSnafu {
145                    key: "object_length",
146                    value: &object.path,
147                }
148                .fail();
149            }
150        }
151        Ok(PreparedPackedImport {
152            backend,
153            index,
154            tables,
155            parallelism: crate::statement::database_copy::parse_parallelism_from_option_map(
156                &req.with,
157            ),
158        })
159    }
160
161    pub async fn import_packed(
162        &self,
163        plan: PreparedPackedImport,
164        cancellation: &CancellationToken,
165        ctx: QueryContextRef,
166    ) -> Result<Output> {
167        let windows = PackReadWindows::new(plan.backend.object_store.clone());
168        let mut pending = PendingPackedInserts::new(plan.parallelism);
169        let objects: HashMap<_, _> = plan
170            .index
171            .objects
172            .iter()
173            .map(|o| (o.path.as_str(), o))
174            .collect();
175        let result: Result<()> = async {
176            for (entry, table) in plan.index.tables.iter().zip(plan.tables) {
177                if cancellation.is_cancelled() {
178                    return error::PackedImportCancelledSnafu.fail();
179                }
180                let object = objects.get(entry.object.as_str()).ok_or_else(|| {
181                    error::InvalidCopyParameterSnafu {
182                        key: "object",
183                        value: &entry.object,
184                    }
185                    .build()
186                })?;
187                match object.kind {
188                    ObjectKind::Pack if entry.length <= WINDOW_SIZE as u64 => {
189                        let reader = PackedParquetReader::new(
190                            Arc::clone(&windows),
191                            object.path.clone(),
192                            object.length,
193                            entry.offset,
194                            entry.length,
195                        )
196                        .context(error::ReadParquetMetadataSnafu)?;
197                        self.copy_indexed_parquet(
198                            reader,
199                            table,
200                            entry.row_count,
201                            &mut pending,
202                            cancellation,
203                            ctx.clone(),
204                        )
205                        .await?
206                    }
207                    ObjectKind::Pack | ObjectKind::Parquet => {
208                        // The virtual file keeps footer/column offsets relative to this
209                        // stream. Parquet decoder allocations are separate from the cache.
210                        let reader = plan
211                            .backend
212                            .object_store
213                            .reader_with(&object.path)
214                            .chunk(256 * 1024)
215                            .await
216                            .context(error::ReadObjectSnafu { path: &object.path })?
217                            .into_futures_async_read(entry.offset..entry.offset + entry.length)
218                            .await
219                            .context(error::ReadObjectSnafu { path: &object.path })?
220                            .compat();
221                        self.copy_indexed_parquet(
222                            reader,
223                            table,
224                            entry.row_count,
225                            &mut pending,
226                            cancellation,
227                            ctx.clone(),
228                        )
229                        .await?
230                    }
231                };
232            }
233            Ok(())
234        }
235        .await;
236        let drained = pending.drain().await;
237        common_telemetry::debug!(
238            peak_pending_bytes = pending.peak_pending_bytes,
239            peak_batch_bytes = pending.peak_batch_bytes,
240            "Packed import insertion payloads; codec and compressed windows accounted separately"
241        );
242        result?;
243        drained?;
244        Ok(Output::new(
245            common_query::OutputData::AffectedRows(pending.rows),
246            common_query::OutputMeta::new_with_cost(pending.cost),
247        ))
248    }
249}
250
251/// Retains all started writes until completion, including after the first error.
252pub(crate) struct PendingPackedInserts {
253    tasks: FuturesUnordered<tokio::task::JoinHandle<(usize, Result<Output>)>>,
254    bytes: usize,
255    parallelism: usize,
256    rows: usize,
257    cost: usize,
258    peak_pending_bytes: usize,
259    peak_batch_bytes: usize,
260}
261
262impl PendingPackedInserts {
263    fn new(parallelism: usize) -> Self {
264        Self {
265            tasks: FuturesUnordered::new(),
266            bytes: 0,
267            parallelism,
268            rows: 0,
269            cost: 0,
270            peak_pending_bytes: 0,
271            peak_batch_bytes: 0,
272        }
273    }
274
275    /// Backpressure before decoding, while allowing one current batch outside the queue.
276    pub(crate) async fn before_decode(&mut self) -> Result<()> {
277        while let Some(Some(output)) = self.tasks.next().now_or_never() {
278            self.complete(output)?;
279        }
280        while self.tasks.len() >= self.parallelism || self.bytes >= 32 * 1024 * 1024 {
281            self.complete_one().await?;
282        }
283        Ok(())
284    }
285
286    pub(crate) async fn admit(
287        &mut self,
288        bytes: usize,
289        insert: impl std::future::Future<Output = Result<Output>> + Send + 'static,
290        cancellation: &CancellationToken,
291    ) -> Result<()> {
292        while let Some(Some(output)) = self.tasks.next().now_or_never() {
293            self.complete(output)?;
294        }
295        self.peak_batch_bytes = self.peak_batch_bytes.max(bytes);
296        const BUDGET: usize = 32 * 1024 * 1024;
297        if bytes > BUDGET {
298            self.drain().await?;
299            if cancellation.is_cancelled() {
300                return error::PackedImportCancelledSnafu.fail();
301            }
302            let (rows, cost) = insert.await?.extract_rows_and_cost();
303            self.rows += rows;
304            self.cost += cost;
305            return Ok(());
306        }
307        while self.tasks.len() >= self.parallelism || self.bytes + bytes > BUDGET {
308            self.complete_one().await?;
309        }
310        if cancellation.is_cancelled() {
311            return error::PackedImportCancelledSnafu.fail();
312        }
313        self.bytes += bytes;
314        self.peak_pending_bytes = self.peak_pending_bytes.max(self.bytes);
315        self.tasks.push(common_runtime::spawn_query(
316            async move { (bytes, insert.await) },
317        ));
318        Ok(())
319    }
320
321    async fn complete_one(&mut self) -> Result<()> {
322        if let Some(output) = self.tasks.next().await {
323            self.complete(output)?;
324        }
325        Ok(())
326    }
327
328    fn complete(
329        &mut self,
330        output: std::result::Result<(usize, Result<Output>), common_runtime::JoinError>,
331    ) -> Result<()> {
332        let (bytes, output) = output.context(error::JoinTaskSnafu)?;
333        self.bytes -= bytes;
334        let (rows, cost) = output?.extract_rows_and_cost();
335        self.rows += rows;
336        self.cost += cost;
337        Ok(())
338    }
339
340    async fn drain(&mut self) -> Result<()> {
341        let mut first = None;
342        while !self.tasks.is_empty() {
343            if let Err(error) = self.complete_one().await {
344                first.get_or_insert(error);
345            }
346        }
347        first.map_or(Ok(()), Err)
348    }
349}
350
351#[cfg(test)]
352mod tests {
353    use std::sync::atomic::{AtomicUsize, Ordering};
354
355    use super::*;
356
357    #[tokio::test]
358    async fn oversized_batch_drains_and_runs_exclusively() {
359        let cancellation = CancellationToken::new();
360        let mut queue = PendingPackedInserts::new(4);
361        let completed = Arc::new(AtomicUsize::new(0));
362        for _ in 0..2 {
363            let completed = completed.clone();
364            queue
365                .admit(
366                    1024,
367                    async move {
368                        completed.fetch_add(1, Ordering::SeqCst);
369                        Ok(Output::new_with_affected_rows(1))
370                    },
371                    &cancellation,
372                )
373                .await
374                .unwrap();
375        }
376        let observed = completed.clone();
377        queue
378            .admit(
379                40 * 1024 * 1024,
380                async move {
381                    assert_eq!(observed.load(Ordering::SeqCst), 2);
382                    Ok(Output::new_with_affected_rows(3))
383                },
384                &cancellation,
385            )
386            .await
387            .unwrap();
388        assert_eq!(queue.rows, 5);
389        assert!(queue.tasks.is_empty());
390        assert!(queue.peak_pending_bytes <= 32 * 1024 * 1024);
391        assert_eq!(queue.peak_batch_bytes, 40 * 1024 * 1024);
392    }
393
394    #[tokio::test]
395    async fn concurrent_insert_queue_drains_after_error_and_cancel() {
396        let cancellation = CancellationToken::new();
397        let mut queue = PendingPackedInserts::new(2);
398        let barrier = Arc::new(tokio::sync::Barrier::new(2));
399        let completed = Arc::new(AtomicUsize::new(0));
400        for fail in [true, false] {
401            let barrier = barrier.clone();
402            let completed = completed.clone();
403            queue
404                .admit(
405                    16 * 1024 * 1024,
406                    async move {
407                        barrier.wait().await;
408                        completed.fetch_add(1, Ordering::SeqCst);
409                        if fail {
410                            error::InvalidCopyParameterSnafu {
411                                key: "test",
412                                value: "failure",
413                            }
414                            .fail()
415                        } else {
416                            Ok(Output::new_with_affected_rows(1))
417                        }
418                    },
419                    &cancellation,
420                )
421                .await
422                .unwrap();
423        }
424        cancellation.cancel();
425        assert!(
426            queue
427                .admit(1, async { panic!("must not start") }, &cancellation)
428                .await
429                .is_err()
430        );
431        let _ = queue.drain().await;
432        assert_eq!(completed.load(Ordering::SeqCst), 2);
433        assert_eq!(queue.bytes, 0);
434        assert!(queue.tasks.is_empty());
435    }
436}
operator/statement/import_packed.rs

operator/statement/
import_packed.rs
