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

operator/statement/
copy_table_to.rs
