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 client::OutputData;
19use common_base::readable_size::ReadableSize;
20use common_datasource::file_format::Format;
21use common_datasource::file_format::csv::stream_to_csv;
22use common_datasource::file_format::json::stream_to_json;
23use common_datasource::file_format::parquet::stream_to_parquet;
24use common_datasource::object_store::build_backend_for_write_with_path;
25use common_query::Output;
26use common_recordbatch::adapter::DfRecordBatchStreamAdapter;
27use common_recordbatch::{
28    RecordBatchStream, SendableRecordBatchMapper, SendableRecordBatchStream,
29    map_json_type_to_string, map_json_type_to_string_schema,
30};
31use common_telemetry::{debug, tracing};
32use datafusion::datasource::DefaultTableSource;
33use datafusion_common::TableReference as DfTableReference;
34use datafusion_expr::LogicalPlanBuilder;
35use object_store::ObjectStore;
36use session::context::QueryContextRef;
37use snafu::{OptionExt, ResultExt};
38use table::requests::CopyTableRequest;
39use table::table::adapter::DfTableProviderAdapter;
40use table::table_reference::TableReference;
41
42use crate::error::{self, BuildDfLogicalPlanSnafu, ExecLogicalPlanSnafu, Result};
43use crate::statement::StatementExecutor;
44
45// The buffer size should be greater than 5MB (minimum multipart upload size).
46/// Buffer size to flush data to object stores.
47const WRITE_BUFFER_THRESHOLD: ReadableSize = ReadableSize::mb(8);
48
49/// Default number of concurrent write, it only works on object store backend(e.g., S3).
50const WRITE_CONCURRENCY: usize = 8;
51
52impl StatementExecutor {
53    async fn stream_to_file(
54        &self,
55        stream: SendableRecordBatchStream,
56        format: &Format,
57        object_store: ObjectStore,
58        path: &str,
59    ) -> Result<usize> {
60        let threshold = WRITE_BUFFER_THRESHOLD.as_bytes() as usize;
61
62        let stream = Box::pin(SendableRecordBatchMapper::new(
63            stream,
64            map_json_type_to_string,
65            map_json_type_to_string_schema,
66        ));
67        match format {
68            Format::Csv(format) => stream_to_csv(
69                Box::pin(DfRecordBatchStreamAdapter::new(stream)),
70                object_store,
71                path,
72                threshold,
73                WRITE_CONCURRENCY,
74                format,
75            )
76            .await
77            .context(error::WriteStreamToFileSnafu { path }),
78            Format::Json(format) => stream_to_json(
79                Box::pin(DfRecordBatchStreamAdapter::new(stream)),
80                object_store,
81                path,
82                threshold,
83                WRITE_CONCURRENCY,
84                format,
85            )
86            .await
87            .context(error::WriteStreamToFileSnafu { path }),
88            Format::Parquet(_) => {
89                let schema = stream.schema();
90                stream_to_parquet(
91                    Box::pin(DfRecordBatchStreamAdapter::new(stream)),
92                    schema,
93                    object_store,
94                    path,
95                    WRITE_CONCURRENCY,
96                )
97                .await
98                .context(error::WriteStreamToFileSnafu { path })
99            }
100            _ => error::UnsupportedFormatSnafu {
101                format: format.clone(),
102            }
103            .fail(),
104        }
105    }
106
107    #[tracing::instrument(skip_all)]
108    pub(crate) async fn copy_table_to(
109        &self,
110        req: CopyTableRequest,
111        query_ctx: QueryContextRef,
112    ) -> Result<usize> {
113        let table_ref = TableReference::full(&req.catalog_name, &req.schema_name, &req.table_name);
114        let table = self.get_table(&table_ref).await?;
115        let table_id = table.table_info().table_id();
116        let format = Format::try_from(&req.with).context(error::ParseFileFormatSnafu)?;
117
118        let df_table_ref = DfTableReference::from(table_ref);
119
120        let filters = table
121            .schema()
122            .timestamp_column()
123            .and_then(|c| {
124                common_query::logical_plan::build_filter_from_timestamp(
125                    &c.name,
126                    req.timestamp_range.as_ref(),
127                )
128            })
129            .into_iter()
130            .collect::<Vec<_>>();
131
132        let table_provider = Arc::new(DfTableProviderAdapter::new(table));
133        let table_source = Arc::new(DefaultTableSource::new(table_provider));
134
135        let mut builder = LogicalPlanBuilder::scan_with_filters(
136            df_table_ref,
137            table_source,
138            None,
139            filters.clone(),
140        )
141        .context(BuildDfLogicalPlanSnafu)?;
142        for f in filters {
143            builder = builder.filter(f).context(BuildDfLogicalPlanSnafu)?;
144        }
145        let plan = builder.build().context(BuildDfLogicalPlanSnafu)?;
146
147        let output = self
148            .query_engine
149            .execute(plan, query_ctx)
150            .await
151            .context(ExecLogicalPlanSnafu)?;
152
153        let CopyTableRequest {
154            location,
155            connection,
156            ..
157        } = &req;
158
159        debug!("Copy table: {table_id} to location: {location}");
160        self.copy_to_file(&format, output, location, connection)
161            .await
162    }
163
164    pub(crate) async fn copy_to_file(
165        &self,
166        format: &Format,
167        output: Output,
168        location: &str,
169        connection: &HashMap<String, String>,
170    ) -> Result<usize> {
171        let output = output
172            .map_dictionary_to_values()
173            .context(error::BuildRecordBatchSnafu)?;
174        let stream = match output.data {
175            OutputData::Stream(stream) => stream,
176            OutputData::RecordBatches(record_batches) => record_batches.as_stream(),
177            _ => unreachable!(),
178        };
179
180        let backend =
181            build_backend_for_write_with_path(location, connection, &self.local_file_access)
182                .await
183                .context(error::BuildBackendSnafu)?;
184        let filename = backend.object_path.context(error::UnexpectedSnafu {
185            violated: format!("Expected filename, path: {location}"),
186        })?;
187        self.stream_to_file(stream, format, backend.object_store, &filename)
188            .await
189    }
190}
operator/statement/copy_table_to.rs

operator/statement/
copy_table_to.rs
