operator/statement/
copy_table_to.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 arrow::array::{Array, ArrayRef, AsArray, UInt64Array, make_array};
19use arrow::datatypes::DataType;
20use client::OutputData;
21use common_base::readable_size::ReadableSize;
22use common_datasource::file_format::Format;
23use common_datasource::file_format::csv::stream_to_csv;
24use common_datasource::file_format::json::stream_to_json;
25use common_datasource::file_format::parquet::stream_to_parquet;
26use common_datasource::object_store::build_backend_for_write_with_path;
27use common_datasource::packed_writer::PackedTableWriter;
28use common_datasource::parquet_writer::ParquetFileWriter;
29use common_query::Output;
30use common_recordbatch::adapter::DfRecordBatchStreamAdapter;
31use common_recordbatch::{
32    SendableRecordBatchMapper, SendableRecordBatchStream, map_json_type_to_string,
33    map_json_type_to_string_schema,
34};
35use common_telemetry::{debug, tracing};
36use datafusion::datasource::DefaultTableSource;
37use datafusion_common::TableReference as DfTableReference;
38use datafusion_expr::LogicalPlanBuilder;
39use futures::StreamExt;
40use object_store::ObjectStore;
41use session::context::QueryContextRef;
42use snafu::{OptionExt, ResultExt, ensure};
43use table::TableRef;
44use table::requests::CopyTableRequest;
45use table::table::adapter::DfTableProviderAdapter;
46use table::table_reference::TableReference;
47use tokio_util::sync::CancellationToken;
48
49use crate::error::{self, BuildDfLogicalPlanSnafu, ExecLogicalPlanSnafu, Result};
50use crate::statement::StatementExecutor;
51use crate::statement::export_logical_tables::writers::ExportWriteBudget;
52use crate::statement::export_logical_tables::{
53    expand_export_batch, map_writer_error, rows_within_budget,
54};
55
56// The buffer size should be greater than 5MB (minimum multipart upload size).
57/// Buffer size to flush data to object stores.
58const WRITE_BUFFER_THRESHOLD: ReadableSize = ReadableSize::mb(8);
59
60/// Default number of concurrent write, it only works on object store backend(e.g., S3).
61const WRITE_CONCURRENCY: usize = 8;
62
63impl StatementExecutor {
64    async fn stream_to_file(
65        &self,
66        stream: SendableRecordBatchStream,
67        format: &Format,
68        object_store: ObjectStore,
69        path: &str,
70    ) -> Result<usize> {
71        let threshold = WRITE_BUFFER_THRESHOLD.as_bytes() as usize;
72
73        let stream = Box::pin(SendableRecordBatchMapper::new(
74            stream,
75            map_json_type_to_string,
76            map_json_type_to_string_schema,
77        ));
78        match format {
79            Format::Csv(format) => stream_to_csv(
80                Box::pin(DfRecordBatchStreamAdapter::new(stream)),
81                object_store,
82                path,
83                threshold,
84                WRITE_CONCURRENCY,
85                format,
86            )
87            .await
88            .context(error::WriteStreamToFileSnafu { path }),
89            Format::Json(format) => stream_to_json(
90                Box::pin(DfRecordBatchStreamAdapter::new(stream)),
91                object_store,
92                path,
93                threshold,
94                WRITE_CONCURRENCY,
95                format,
96            )
97            .await
98            .context(error::WriteStreamToFileSnafu { path }),
99            Format::Parquet(_) => stream_to_parquet(
100                Box::pin(DfRecordBatchStreamAdapter::new(stream)),
101                object_store,
102                path,
103                WRITE_CONCURRENCY,
104            )
105            .await
106            .context(error::WriteStreamToFileSnafu { path }),
107            _ => error::UnsupportedFormatSnafu {
108                format: format.clone(),
109            }
110            .fail(),
111        }
112    }
113
114    #[tracing::instrument(skip_all)]
115    pub(crate) async fn copy_table_to(
116        &self,
117        req: CopyTableRequest,
118        query_ctx: QueryContextRef,
119    ) -> Result<usize> {
120        let table_ref = TableReference::full(&req.catalog_name, &req.schema_name, &req.table_name);
121        let table = self.get_table(&table_ref).await?;
122        self.copy_captured_table_to(table, req, query_ctx).await
123    }
124
125    pub(crate) async fn copy_captured_table_to(
126        &self,
127        table: TableRef,
128        req: CopyTableRequest,
129        query_ctx: QueryContextRef,
130    ) -> Result<usize> {
131        self.copy_captured_table_to_managed(table, req, query_ctx, None, None)
132            .await
133    }
134
135    pub(crate) async fn copy_captured_table_to_managed(
136        &self,
137        table: TableRef,
138        req: CopyTableRequest,
139        query_ctx: QueryContextRef,
140        managed: Option<(&ExportWriteBudget, &CancellationToken)>,
141        packed: Option<PackedTableWriter>,
142    ) -> Result<usize> {
143        let info = table.table_info();
144        let table_ref = TableReference::full(&info.catalog_name, &info.schema_name, &info.name);
145        let table_id = info.table_id();
146        let format = Format::try_from(&req.with).context(error::ParseFileFormatSnafu)?;
147
148        let df_table_ref = DfTableReference::from(table_ref);
149
150        let filters = table
151            .schema()
152            .timestamp_column()
153            .and_then(|c| {
154                common_query::logical_plan::build_filter_from_timestamp(
155                    &c.name,
156                    req.timestamp_range.as_ref(),
157                )
158            })
159            .into_iter()
160            .collect::<Vec<_>>();
161
162        let table_provider = Arc::new(DfTableProviderAdapter::new(table));
163        let table_source = Arc::new(DefaultTableSource::new(table_provider));
164
165        let mut builder = LogicalPlanBuilder::scan_with_filters(
166            df_table_ref,
167            table_source,
168            None,
169            filters.clone(),
170        )
171        .context(BuildDfLogicalPlanSnafu)?;
172        for f in filters {
173            builder = builder.filter(f).context(BuildDfLogicalPlanSnafu)?;
174        }
175        let plan = builder.build().context(BuildDfLogicalPlanSnafu)?;
176
177        let output = self
178            .query_engine
179            .execute(plan, query_ctx)
180            .await
181            .context(ExecLogicalPlanSnafu)?;
182
183        let CopyTableRequest {
184            location,
185            connection,
186            ..
187        } = &req;
188
189        debug!("Copy table: {table_id} to location: {location}");
190        self.copy_to_file_managed(&format, output, location, connection, managed, packed)
191            .await
192    }
193
194    pub(crate) async fn copy_to_file(
195        &self,
196        format: &Format,
197        output: Output,
198        location: &str,
199        connection: &HashMap<String, String>,
200    ) -> Result<usize> {
201        self.copy_to_file_managed(format, output, location, connection, None, None)
202            .await
203    }
204
205    async fn copy_to_file_managed(
206        &self,
207        format: &Format,
208        output: Output,
209        location: &str,
210        connection: &HashMap<String, String>,
211        managed: Option<(&ExportWriteBudget, &CancellationToken)>,
212        packed: Option<PackedTableWriter>,
213    ) -> Result<usize> {
214        let output = if managed.is_none() {
215            output
216                .map_dictionary_to_values()
217                .context(error::BuildRecordBatchSnafu)?
218        } else {
219            output
220        };
221        let stream = match output.data {
222            OutputData::Stream(stream) => stream,
223            OutputData::RecordBatches(record_batches) => record_batches.as_stream(),
224            _ => unreachable!(),
225        };
226
227        let backend =
228            build_backend_for_write_with_path(location, connection, &self.local_file_access)
229                .await
230                .context(error::BuildBackendSnafu)?;
231        let filename = backend.object_path.context(error::UnexpectedSnafu {
232            violated: format!("Expected filename, path: {location}"),
233        })?;
234        if let Some((budget, token)) = managed {
235            stream_to_managed_parquet_with_packed(
236                stream,
237                backend.object_store,
238                &filename,
239                budget,
240                token,
241                packed,
242            )
243            .await
244        } else {
245            self.stream_to_file(stream, format, backend.object_store, &filename)
246                .await
247        }
248    }
249}
250
251#[cfg(test)]
252pub(crate) async fn stream_to_managed_parquet(
253    stream: SendableRecordBatchStream,
254    store: ObjectStore,
255    path: &str,
256    budget: &ExportWriteBudget,
257    token: &CancellationToken,
258) -> Result<usize> {
259    stream_to_managed_parquet_with_packed(stream, store, path, budget, token, None).await
260}
261
262async fn stream_to_managed_parquet_with_packed(
263    mut stream: SendableRecordBatchStream,
264    store: ObjectStore,
265    path: &str,
266    budget: &ExportWriteBudget,
267    token: &CancellationToken,
268    packed: Option<PackedTableWriter>,
269) -> Result<usize> {
270    use common_recordbatch::{RecordBatch, map_dictionary_to_values_schema};
271    let original = stream.schema();
272    let (expanded_schema, expand) = map_dictionary_to_values_schema(original.clone());
273    let json_columns = expanded_schema
274        .column_schemas()
275        .iter()
276        .enumerate()
277        .filter_map(|(index, column)| {
278            (column.data_type.is_json() && !column.data_type.is_json2()).then_some(index)
279        })
280        .collect::<Vec<_>>();
281    let (mapped_schema, json) = map_json_type_to_string_schema(expanded_schema.clone());
282    let output_schema = if json {
283        mapped_schema
284    } else {
285        expanded_schema.clone()
286    };
287    let packed_destination = packed.is_some();
288    let mut writer = if let Some(packed) = packed {
289        ParquetFileWriter::open_packed(
290            output_schema.arrow_schema().clone(),
291            store,
292            path,
293            Some(
294                crate::statement::export_logical_tables::LogicalTableExportLimits::default().writer,
295            ),
296            packed,
297        )
298    } else {
299        ParquetFileWriter::open(output_schema.arrow_schema().clone(), store, path, 1, None).await
300    }
301    .context(error::WriteStreamToFileSnafu { path })?;
302    let mut started = false;
303    let result = async {
304        let mut rows = 0;
305        loop {
306            let batch = tokio::select! {
307                biased;
308                _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(),
309                batch = stream.next() => batch,
310            };
311            let Some(batch) = batch else { break };
312            let batch = batch
313                .context(error::BuildRecordBatchSnafu)?
314                .into_df_record_batch();
315            let mut offset = 0;
316            while offset < batch.num_rows() {
317                // The scan batch remains query-owned. Bound the downstream slice and
318                // detach its buffers below if it would retain more than its reservation.
319                let (conversion, retained) =
320                    ExportWriteBudget::conversion_budget(0, batch.num_columns(), usize::MAX)?;
321                let input = batch.clone();
322                let json_columns = json_columns.clone();
323                let (len, estimated) = common_runtime::spawn_blocking_global(move || {
324                    rows_within_budget(&input, offset, input.num_rows(), conversion, &json_columns)
325                })
326                .await
327                .context(error::JoinTaskSnafu)??;
328                let reservation = retained.saturating_add(estimated.saturating_mul(4));
329                let permit = budget.reserve(reservation, token).await?;
330                let batch = batch.clone();
331                let (original, expanded_schema, output_schema) = (
332                    original.clone(),
333                    expanded_schema.clone(),
334                    output_schema.clone(),
335                );
336                let (batch, len) = common_runtime::spawn_blocking_global(move || {
337                    let mut batch = RecordBatch::from_df_record_batch(
338                        original.clone(),
339                        batch.slice(offset, len),
340                    );
341                    if expand {
342                        batch = RecordBatch::from_df_record_batch(
343                            expanded_schema.clone(),
344                            expand_export_batch(
345                                &batch.into_df_record_batch(),
346                                expanded_schema.arrow_schema().clone(),
347                            )?,
348                        );
349                    }
350                    if json {
351                        batch = map_json_type_to_string(batch, &expanded_schema, &output_schema)
352                            .context(error::BuildRecordBatchSnafu)?;
353                    }
354                    let mut batch = batch.into_df_record_batch();
355                    if batch.get_array_memory_size() > reservation {
356                        let indices = UInt64Array::from_iter_values(0..batch.num_rows() as u64);
357                        batch = arrow::compute::take_record_batch(&batch, &indices)
358                            .context(error::ComputeArrowSnafu)?;
359                        let arrays = batch
360                            .columns()
361                            .iter()
362                            .map(compact_view_buffers)
363                            .collect::<Result<Vec<_>>>()?;
364                        batch = arrow::record_batch::RecordBatch::try_new(batch.schema(), arrays)
365                            .context(error::ComputeArrowSnafu)?;
366                    }
367                    ensure!(
368                        batch.get_array_memory_size() <= reservation,
369                        error::LogicalTableExportResourceSnafu {
370                            reason: "converted backing buffers exceed reservation"
371                        }
372                    );
373                    Ok::<_, error::Error>((batch, len))
374                })
375                .await
376                .context(error::JoinTaskSnafu)??;
377                started = true;
378                let write = writer.write(batch, Some(token)).await;
379                drop(permit);
380                write.map_err(|error| map_writer_error(error, path))?;
381                rows += len;
382                offset += len;
383            }
384        }
385        started = true;
386        writer
387            .finish(Some(token))
388            .await
389            .map_err(|error| map_writer_error(error, path))?;
390        Ok(rows)
391    }
392    .await;
393    if result.is_err() {
394        token.cancel();
395        if (started || packed_destination)
396            && let Err(error) = writer.abort().await
397        {
398            common_telemetry::warn!(error; "Failed to abort ordinary export file");
399        }
400    }
401    result
402}
403
404// Arrow take copies ordinary buffers but shares view data buffers, including
405// views nested inside lists or structs. Reclaim those unselected values too.
406fn compact_view_buffers(array: &ArrayRef) -> Result<ArrayRef> {
407    match array.data_type() {
408        DataType::Utf8View => Ok(Arc::new(array.as_string_view().gc())),
409        DataType::BinaryView => Ok(Arc::new(array.as_binary_view().gc())),
410        _ => {
411            let data = array.to_data();
412            if data.child_data().is_empty() {
413                return Ok(array.clone());
414            }
415            let children = data
416                .child_data()
417                .iter()
418                .map(|child| {
419                    compact_view_buffers(&make_array(child.clone())).map(|array| array.to_data())
420                })
421                .collect::<Result<Vec<_>>>()?;
422            Ok(make_array(
423                data.into_builder()
424                    .child_data(children)
425                    .build()
426                    .context(error::ComputeArrowSnafu)?,
427            ))
428        }
429    }
430}
431
432#[cfg(test)]
433mod tests {
434    use arrow::array::{
435        ArrayRef, BinaryArray, BinaryViewArray, DictionaryArray, Int32Array, StringArray,
436        StringViewArray,
437    };
438    use arrow::datatypes::Int32Type;
439    use common_recordbatch::{RecordBatch, RecordBatches};
440    use datafusion::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
441    use datatypes::prelude::ConcreteDataType;
442    use datatypes::schema::{ColumnSchema, Schema};
443
444    use super::*;
445
446    #[tokio::test]
447    async fn managed_copy_preserves_existing_file_before_sink_open() {
448        let temp_dir = common_test_util::temp_dir::create_temp_dir("managed_copy_existing_file");
449        let store = object_store::secure_fs::SecureFsRoot::open(temp_dir.path())
450            .unwrap()
451            .build_operator();
452        let path = "existing.parquet";
453        store.write(path, "original").await.unwrap();
454        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
455            "value",
456            ConcreteDataType::int32_datatype(),
457            false,
458        )]));
459        let batch = arrow::record_batch::RecordBatch::try_new(
460            schema.arrow_schema().clone(),
461            vec![Arc::new(Int32Array::from(vec![1]))],
462        )
463        .unwrap();
464        let batches = async_stream::stream! {
465            yield Ok(batch);
466            yield Err(datafusion::error::DataFusionError::Execution("source failed".into()));
467        };
468        let stream = datafusion::physical_plan::stream::RecordBatchStreamAdapter::new(
469            schema.arrow_schema().clone(),
470            batches,
471        );
472        let stream =
473            common_recordbatch::adapter::RecordBatchStreamAdapter::try_new(Box::pin(stream))
474                .unwrap();
475
476        let result = stream_to_managed_parquet(
477            Box::pin(stream),
478            store.clone(),
479            path,
480            &ExportWriteBudget::new(1),
481            &CancellationToken::new(),
482        )
483        .await;
484
485        assert!(result.is_err());
486        assert_eq!(
487            store.read(path).await.unwrap().to_bytes().as_ref(),
488            b"original"
489        );
490    }
491
492    #[tokio::test]
493    async fn managed_copy_rechunks_large_backing_buffers() {
494        let value = "x".repeat(9 * 1024);
495        for view in [false, true] {
496            let data_type = if view {
497                ConcreteDataType::utf8_view_datatype()
498            } else {
499                ConcreteDataType::string_datatype()
500            };
501            let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
502                "value", data_type, false,
503            )]));
504            let values = vec![value.as_str(); 8192];
505            let array: ArrayRef = if view {
506                Arc::new(StringViewArray::from(values))
507            } else {
508                Arc::new(StringArray::from(values))
509            };
510            let batch = arrow::record_batch::RecordBatch::try_new(
511                schema.arrow_schema().clone(),
512                vec![array],
513            )
514            .unwrap();
515            assert!(batch.get_array_memory_size() > 64 * 1024 * 1024);
516            // Also exercise a small slice that still references the large allocation.
517            for input in [batch.clone(), batch.slice(200, 3)] {
518                let rows = input.num_rows();
519                let batches = RecordBatches::try_new(
520                    schema.clone(),
521                    vec![RecordBatch::from_df_record_batch(schema.clone(), input)],
522                )
523                .unwrap();
524                let store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
525                let budget = ExportWriteBudget::new(1);
526                assert_eq!(
527                    stream_to_managed_parquet(
528                        batches.as_stream(),
529                        store.clone(),
530                        "large.parquet",
531                        &budget,
532                        &CancellationToken::new(),
533                    )
534                    .await
535                    .unwrap(),
536                    rows
537                );
538                let reader = ParquetRecordBatchReaderBuilder::try_new(
539                    store.read("large.parquet").await.unwrap().to_bytes(),
540                )
541                .unwrap()
542                .build()
543                .unwrap();
544                let mut actual_rows = 0;
545                for batch in reader {
546                    let batch = batch.unwrap();
547                    let strings = arrow::compute::cast(batch.column(0), &DataType::Utf8).unwrap();
548                    assert!(
549                        strings
550                            .as_string::<i32>()
551                            .iter()
552                            .all(|item| item == Some(value.as_str()))
553                    );
554                    actual_rows += batch.num_rows();
555                }
556                assert_eq!(actual_rows, rows);
557                assert_eq!(budget.available(), (1, 64 * 1024 * 1024));
558            }
559        }
560    }
561
562    #[tokio::test]
563    async fn managed_copy_preserves_dictionary_json_views_and_empty_schema() {
564        let schema = Arc::new(Schema::new(vec![
565            ColumnSchema::new(
566                "host",
567                ConcreteDataType::dictionary_datatype(
568                    ConcreteDataType::int32_datatype(),
569                    ConcreteDataType::string_datatype(),
570                ),
571                true,
572            ),
573            ColumnSchema::new("json", ConcreteDataType::json_datatype(), true),
574            ColumnSchema::new("text_view", ConcreteDataType::utf8_view_datatype(), true),
575            ColumnSchema::new(
576                "binary_view",
577                ConcreteDataType::binary_view_datatype(),
578                true,
579            ),
580        ]));
581        let dictionary = DictionaryArray::<Int32Type>::new(
582            Int32Array::from(vec![Some(0), None, Some(1), Some(1)]),
583            Arc::new(StringArray::from(vec!["tag", &"x".repeat(3 * 1024 * 1024)])),
584        );
585        let json = datatypes::types::parse_string_to_jsonb(&format!(
586            r#"{{"value":"escaped\ntext{}","n":123456789}}"#,
587            "y".repeat(160_000)
588        ))
589        .unwrap();
590        let arrays = vec![
591            Arc::new(dictionary) as ArrayRef,
592            Arc::new(BinaryArray::from(vec![
593                Some(json.as_slice()),
594                None,
595                None,
596                Some(json.as_slice()),
597            ])),
598            Arc::new(StringViewArray::from(vec![
599                Some("short"),
600                None,
601                Some("long view"),
602                Some("long view"),
603            ])),
604            Arc::new(BinaryViewArray::from(vec![
605                Some(&b"short"[..]),
606                None,
607                Some(&b"long view"[..]),
608                Some(&b"long view"[..]),
609            ])),
610        ];
611        let batch =
612            arrow::record_batch::RecordBatch::try_new(schema.arrow_schema().clone(), arrays)
613                .unwrap();
614        let batch = RecordBatch::from_df_record_batch(schema.clone(), batch);
615        for empty in [false, true] {
616            let batches = RecordBatches::try_new(
617                schema.clone(),
618                if empty { vec![] } else { vec![batch.clone()] },
619            )
620            .unwrap();
621            let store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
622            let budget = ExportWriteBudget::new(1);
623            let token = CancellationToken::new();
624            let managed = stream_to_managed_parquet(
625                batches.as_stream(),
626                store.clone(),
627                "managed.parquet",
628                &budget,
629                &token,
630            )
631            .await
632            .unwrap();
633            let output = Output::new_with_record_batches(batches)
634                .map_dictionary_to_values()
635                .unwrap();
636            let OutputData::RecordBatches(batches) = output.data else {
637                panic!("expected batches");
638            };
639            let stream = SendableRecordBatchMapper::new(
640                batches.as_stream(),
641                map_json_type_to_string,
642                map_json_type_to_string_schema,
643            );
644            let ordinary = stream_to_parquet(
645                Box::pin(DfRecordBatchStreamAdapter::new(Box::pin(stream))),
646                store.clone(),
647                "ordinary.parquet",
648                WRITE_CONCURRENCY,
649            )
650            .await
651            .unwrap();
652            assert_eq!(managed, ordinary);
653            let mut outputs = Vec::new();
654            for name in ["managed.parquet", "ordinary.parquet"] {
655                let reader = ParquetRecordBatchReaderBuilder::try_new(
656                    store.read(name).await.unwrap().to_bytes(),
657                )
658                .unwrap();
659                let schema = reader.schema().clone();
660                let values = reader
661                    .build()
662                    .unwrap()
663                    .collect::<std::result::Result<Vec<_>, _>>()
664                    .unwrap();
665                outputs.push((schema, values));
666            }
667            assert_eq!(outputs[0], outputs[1]);
668            assert_eq!(budget.available(), (1, 64 * 1024 * 1024));
669        }
670    }
671}
operator/statement/copy_table_to.rs

operator/statement/
copy_table_to.rs
