risingwave_connector/sink/file_sink/
opendal_sink.rs

1// Copyright 2024 RisingWave Labs
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::{BTreeMap, HashMap};
16use std::fmt::Write;
17use std::marker::PhantomData;
18use std::sync::Arc;
19use std::time::{Duration, SystemTime, UNIX_EPOCH};
20
21use anyhow::anyhow;
22use bytes::BytesMut;
23use chrono::{TimeZone, Utc};
24use opendal::{FuturesAsyncWriter, Operator, Writer as OpendalWriter};
25use parquet::arrow::AsyncArrowWriter;
26use parquet::basic::Compression;
27use parquet::file::properties::WriterProperties;
28use risingwave_common::array::arrow::IcebergArrowConvert;
29use risingwave_common::array::arrow::arrow_schema_iceberg::{self, SchemaRef};
30use risingwave_common::array::{Op, StreamChunk};
31use risingwave_common::catalog::Schema;
32use risingwave_pb::id::ExecutorId;
33use serde::Deserialize;
34use serde_json::Value;
35use serde_with::{DisplayFromStr, serde_as};
36use strum_macros::{Display, EnumString};
37use tokio_util::compat::{Compat, FuturesAsyncWriteCompatExt};
38use uuid::Uuid;
39use with_options::WithOptions;
40
41use crate::enforce_secret::EnforceSecret;
42use crate::sink::catalog::SinkEncode;
43use crate::sink::encoder::{
44    JsonEncoder, JsonbHandlingMode, RowEncoder, TimeHandlingMode, TimestampHandlingMode,
45    TimestamptzHandlingMode,
46};
47use crate::sink::file_sink::batching_log_sink::BatchingLogSinker;
48use crate::sink::{Result, Sink, SinkError, SinkFormatDesc, SinkParam, UnknownFields};
49use crate::source::TryFromBTreeMap;
50use crate::with_options::WithOptions;
51
52pub const DEFAULT_ROLLOVER_SECONDS: usize = 10;
53pub const DEFAULT_MAX_ROW_COUNR: usize = 10240;
54
55pub fn default_rollover_seconds() -> usize {
56    DEFAULT_ROLLOVER_SECONDS
57}
58pub fn default_max_row_count() -> usize {
59    DEFAULT_MAX_ROW_COUNR
60}
61/// The `FileSink` struct represents a file sink that uses the `OpendalSinkBackend` trait for its backend implementation.
62///
63/// # Type Parameters
64///
65/// - S: The type parameter S represents the concrete implementation of the `OpendalSinkBackend` trait used by this file sink.
66#[derive(Debug, Clone)]
67pub struct FileSink<S: OpendalSinkBackend> {
68    pub(crate) op: Operator,
69    /// The path to the file where the sink writes data.
70    pub(crate) path: String,
71    /// The schema describing the structure of the data being written to the file sink.
72    pub(crate) schema: Schema,
73    pub(crate) is_append_only: bool,
74    /// The batching strategy for sinking data to files.
75    pub(crate) batching_strategy: BatchingStrategy,
76
77    /// The description of the sink's format.
78    pub(crate) format_desc: SinkFormatDesc,
79    pub(crate) engine_type: EngineType,
80    pub(crate) unknown_fields: HashMap<String, String>,
81    pub(crate) _marker: PhantomData<S>,
82}
83
84impl<S: OpendalSinkBackend> EnforceSecret for FileSink<S> {}
85
86/// The `OpendalSinkBackend` trait unifies the behavior of various sink backends
87/// implemented through `OpenDAL`(`<https://github.com/apache/opendal>`).
88///
89/// # Type Parameters
90///
91/// - Properties: Represents the necessary parameters for establishing a backend.
92///
93/// # Constants
94///
95/// - `SINK_NAME`: A static string representing the name of the sink.
96///
97/// # Functions
98///
99/// - `from_btreemap`: Automatically parse the required parameters from the input create sink statement.
100/// - `new_operator`: Creates a new operator using the provided backend properties.
101/// - `get_path`: Returns the path of the sink file specified by the user's create sink statement.
102pub trait OpendalSinkBackend: Send + Sync + 'static + Clone + PartialEq {
103    type Properties: TryFromBTreeMap + UnknownFields + Send + Sync + Clone + WithOptions;
104    const SINK_NAME: &'static str;
105
106    fn from_btreemap(btree_map: BTreeMap<String, String>) -> Result<Self::Properties>;
107    fn new_operator(properties: Self::Properties) -> Result<Operator>;
108    fn get_path(properties: Self::Properties) -> String;
109    fn get_engine_type() -> EngineType;
110    fn get_batching_strategy(properties: Self::Properties) -> BatchingStrategy;
111}
112
113#[derive(Clone, Debug)]
114pub enum EngineType {
115    Gcs,
116    S3,
117    Fs,
118    Azblob,
119    Webhdfs,
120    Snowflake,
121}
122
123impl<S: OpendalSinkBackend> Sink for FileSink<S> {
124    type LogSinker = BatchingLogSinker;
125
126    const SINK_NAME: &'static str = S::SINK_NAME;
127
128    fn validate_unknown_fields(&self) -> Result<()> {
129        crate::sink::validate_sink_unknown_fields(&self.unknown_fields)
130    }
131
132    async fn validate(&self) -> Result<()> {
133        if matches!(self.engine_type, EngineType::Snowflake) {
134            risingwave_common::license::Feature::SnowflakeSink
135                .check_available()
136                .map_err(|e| anyhow::anyhow!(e))?;
137        }
138        if !self.is_append_only {
139            return Err(SinkError::Config(anyhow!(
140                "File sink only supports append-only mode at present. \
141                    Please change the query to append-only, and specify it \
142                    explicitly after the `FORMAT ... ENCODE ...` statement. \
143                    For example, `FORMAT xxx ENCODE xxx(force_append_only='true')`"
144            )));
145        }
146
147        if self.format_desc.encode != SinkEncode::Parquet
148            && self.format_desc.encode != SinkEncode::Json
149        {
150            return Err(SinkError::Config(anyhow!(
151                "File sink only supports `PARQUET` and `JSON` encode at present."
152            )));
153        }
154
155        match self.op.list(&self.path).await {
156            Ok(_) => Ok(()),
157            Err(e) => Err(anyhow!(e).into()),
158        }
159    }
160
161    async fn new_log_sinker(
162        &self,
163        writer_param: crate::sink::SinkWriterParam,
164    ) -> Result<Self::LogSinker> {
165        let writer = OpenDalSinkWriter::new(
166            self.op.clone(),
167            &self.path,
168            self.schema.clone(),
169            writer_param.executor_id,
170            &self.format_desc,
171            self.engine_type.clone(),
172            self.batching_strategy.clone(),
173        )?;
174        Ok(BatchingLogSinker::new(writer))
175    }
176}
177
178impl<S: OpendalSinkBackend> TryFrom<SinkParam> for FileSink<S> {
179    type Error = SinkError;
180
181    fn try_from(param: SinkParam) -> std::result::Result<Self, Self::Error> {
182        let schema = param.schema();
183        let config = S::from_btreemap(param.properties)?;
184        let unknown_fields = crate::sink::UnknownFields::unknown_fields(&config);
185        let path = S::get_path(config.clone());
186        let op = S::new_operator(config.clone())?;
187        let batching_strategy = S::get_batching_strategy(config);
188        let engine_type = S::get_engine_type();
189        let format_desc = match param.format_desc {
190            Some(desc) => desc,
191            None => {
192                if let EngineType::Snowflake = engine_type {
193                    SinkFormatDesc::plain_json_for_snowflake_only()
194                } else {
195                    return Err(SinkError::Config(anyhow!("missing FORMAT ... ENCODE ...")));
196                }
197            }
198        };
199        Ok(Self {
200            op,
201            path,
202            schema,
203            is_append_only: param.sink_type.is_append_only(),
204            batching_strategy,
205            format_desc,
206            engine_type,
207            unknown_fields,
208            _marker: PhantomData,
209        })
210    }
211}
212
213pub struct OpenDalSinkWriter {
214    schema: SchemaRef,
215    operator: Operator,
216    sink_writer: Option<FileWriterEnum>,
217    write_path: String,
218    executor_id: ExecutorId,
219    unique_writer_id: Uuid,
220    encode_type: SinkEncode,
221    row_encoder: JsonEncoder,
222    engine_type: EngineType,
223    pub(crate) batching_strategy: BatchingStrategy,
224    current_bached_row_num: usize,
225    created_time: SystemTime,
226    file_seq: u64,
227}
228
229/// The `FileWriterEnum` enum represents different types of file writers used for various sink
230/// implementations.
231///
232/// # Variants
233///
234/// - `ParquetFileWriter`: Represents a Parquet file writer using the `AsyncArrowWriter<W>`
235///   for writing data to a Parquet file. It accepts an implementation of W: `AsyncWrite` + `Unpin` + `Send`
236///   as the underlying writer. In this case, the `OpendalWriter` serves as the underlying writer.
237/// - `FileWriter`: Represents a basic `OpenDAL` writer, for writing files in encodes other than parquet.
238///
239/// The choice of writer used during the actual writing process depends on the encode type of the sink.
240enum FileWriterEnum {
241    ParquetFileWriter(AsyncArrowWriter<Compat<FuturesAsyncWriter>>),
242    FileWriter(OpendalWriter),
243}
244
245/// Public interface exposed to `BatchingLogSinker`, used to write chunk and commit files.
246impl OpenDalSinkWriter {
247    /// This method writes a chunk.
248    pub async fn write_batch(&mut self, chunk: StreamChunk) -> Result<()> {
249        if self.sink_writer.is_none() {
250            assert_eq!(self.current_bached_row_num, 0);
251            self.create_sink_writer().await?;
252        };
253        self.append_only(chunk).await?;
254        Ok(())
255    }
256
257    /// This method close current writer, finish writing a file and returns whether the commit is successful.
258    pub async fn commit(&mut self) -> Result<bool> {
259        if let Some(sink_writer) = self.sink_writer.take() {
260            match sink_writer {
261                FileWriterEnum::ParquetFileWriter(w) => {
262                    let bytes_written = w.bytes_written();
263                    if bytes_written > 0 {
264                        w.close().await?;
265                        tracing::debug!(
266                            "writer {} (executor_id: {}, created_time: {}) finish write file, bytes_written: {}",
267                            self.unique_writer_id,
268                            self.executor_id,
269                            self.created_time
270                                .duration_since(UNIX_EPOCH)
271                                .expect("Time went backwards")
272                                .as_secs(),
273                            bytes_written
274                        );
275                    }
276                }
277                FileWriterEnum::FileWriter(mut w) => {
278                    w.close().await?;
279                }
280            };
281            self.current_bached_row_num = 0;
282            return Ok(true);
283        }
284        Ok(false)
285    }
286
287    /// Returns whether there is pending data (i.e., an active writer that has not been committed yet).
288    pub fn has_pending_data(&self) -> bool {
289        self.sink_writer.is_some()
290    }
291
292    // Try commit if the batching condition is met.
293    pub async fn try_commit(&mut self) -> Result<bool> {
294        if self.can_commit() {
295            return self.commit().await;
296        }
297        Ok(false)
298    }
299}
300
301/// Private methods related to batching.
302impl OpenDalSinkWriter {
303    /// Method for judging whether batch condition is met.
304    fn can_commit(&self) -> bool {
305        self.duration_seconds_since_writer_created() >= self.batching_strategy.rollover_seconds
306            || self.current_bached_row_num >= self.batching_strategy.max_row_count
307    }
308
309    fn path_partition_prefix(&self, duration: &Duration) -> String {
310        let datetime = Utc
311            .timestamp_opt(duration.as_secs() as i64, 0)
312            .single()
313            .expect("Failed to convert timestamp to DateTime<Utc>")
314            .with_timezone(&Utc);
315        let path_partition_prefix = self
316            .batching_strategy
317            .path_partition_prefix
318            .as_ref()
319            .unwrap_or(&PathPartitionPrefix::None);
320        match path_partition_prefix {
321            PathPartitionPrefix::None => "".to_owned(),
322            PathPartitionPrefix::Day => datetime.format("%Y-%m-%d/").to_string(),
323            PathPartitionPrefix::Month => datetime.format("/%Y-%m/").to_string(),
324            PathPartitionPrefix::Hour => datetime.format("/%Y-%m-%d %H:00/").to_string(),
325        }
326    }
327
328    fn duration_seconds_since_writer_created(&self) -> usize {
329        let now = SystemTime::now();
330        now.duration_since(self.created_time)
331            .expect("Time went backwards")
332            .as_secs() as usize
333    }
334
335    // Method for writing chunk and update related batching condition.
336    async fn append_only(&mut self, chunk: StreamChunk) -> Result<()> {
337        match self
338            .sink_writer
339            .as_mut()
340            .ok_or_else(|| SinkError::File("Sink writer is not created.".to_owned()))?
341        {
342            FileWriterEnum::ParquetFileWriter(w) => {
343                let batch =
344                    IcebergArrowConvert.to_record_batch(self.schema.clone(), chunk.data_chunk())?;
345                let batch_row_nums = batch.num_rows();
346                w.write(&batch).await?;
347                self.current_bached_row_num += batch_row_nums;
348            }
349            FileWriterEnum::FileWriter(w) => {
350                let mut chunk_buf = BytesMut::new();
351                let batch_row_nums = chunk.data_chunk().capacity();
352                // write the json representations of the row(s) in current chunk to `chunk_buf`
353                for (op, row) in chunk.rows() {
354                    assert_eq!(op, Op::Insert, "expect all `op(s)` to be `Op::Insert`");
355                    // to prevent temporary string allocation,
356                    // so we directly write to `chunk_buf` implicitly via `write_fmt`.
357                    writeln!(
358                        chunk_buf,
359                        "{}",
360                        Value::Object(self.row_encoder.encode(row)?)
361                    )
362                    .unwrap(); // write to a `BytesMut` should never fail
363                }
364                w.write(chunk_buf.freeze()).await?;
365                self.current_bached_row_num += batch_row_nums;
366            }
367        }
368        Ok(())
369    }
370}
371
372/// Init methods.
373impl OpenDalSinkWriter {
374    pub fn new(
375        operator: Operator,
376        write_path: &str,
377        rw_schema: Schema,
378        executor_id: ExecutorId,
379        format_desc: &SinkFormatDesc,
380        engine_type: EngineType,
381        batching_strategy: BatchingStrategy,
382    ) -> Result<Self> {
383        let arrow_schema = convert_rw_schema_to_arrow_schema(rw_schema.clone())?;
384        let jsonb_handling_mode = JsonbHandlingMode::from_options(&format_desc.options)?;
385        let row_encoder = JsonEncoder::new(
386            rw_schema,
387            None,
388            crate::sink::encoder::DateHandlingMode::String,
389            TimestampHandlingMode::String,
390            TimestamptzHandlingMode::UtcString,
391            TimeHandlingMode::String,
392            jsonb_handling_mode,
393        );
394        Ok(Self {
395            schema: Arc::new(arrow_schema),
396            write_path: write_path.to_owned(),
397            operator,
398            sink_writer: None,
399            executor_id,
400            unique_writer_id: Uuid::now_v7(),
401            encode_type: format_desc.encode.clone(),
402            row_encoder,
403            engine_type,
404            batching_strategy,
405            current_bached_row_num: 0,
406            created_time: SystemTime::now(),
407            file_seq: 0,
408        })
409    }
410
411    async fn create_object_writer(&mut self) -> Result<OpendalWriter> {
412        // Todo: specify more file suffixes based on encode_type.
413        let suffix = match self.encode_type {
414            SinkEncode::Parquet => "parquet",
415            SinkEncode::Json => "json",
416            _ => unimplemented!(),
417        };
418
419        let create_time = self
420            .created_time
421            .duration_since(UNIX_EPOCH)
422            .expect("Time went backwards");
423
424        // With batching in place, the file writing process is decoupled from checkpoints.
425        // The current file naming convention is as follows:
426        // 1. A subdirectory is defined based on `path_partition_prefix` (e.g., by day、hour or month or none.).
427        // 2. The file name includes a unique UUID (v7, which contains timestamp) and the creation time in seconds since the UNIX epoch.
428        // If the engine type is `Fs`, the path is automatically handled, and the filename does not include a path prefix.
429        // 3. For the Snowflake Sink, the `write_path` parameter can be empty.
430        // When the `write_path` is not specified, the data will be written to the root of the specified bucket.
431        let object_name = {
432            let base_path = match self.engine_type {
433                EngineType::Fs => "".to_owned(),
434                EngineType::Snowflake if self.write_path.is_empty() => "".to_owned(),
435                _ => format!("{}/", self.write_path),
436            };
437            let current_file_seq = self.file_seq;
438            self.file_seq = self.file_seq.checked_add(1).expect("file seq overflow");
439
440            format!(
441                "{}{}{}_{}_{}.{}",
442                base_path,
443                self.path_partition_prefix(&create_time),
444                self.unique_writer_id,
445                create_time.as_secs(),
446                current_file_seq,
447                suffix,
448            )
449        };
450        Ok(self
451            .operator
452            .writer_with(&object_name)
453            .concurrent(8)
454            .await?)
455    }
456
457    async fn create_sink_writer(&mut self) -> Result<()> {
458        // Update the `created_time` to the current time when creating a new writer.
459        self.created_time = SystemTime::now();
460
461        let object_writer = self.create_object_writer().await?;
462        match self.encode_type {
463            SinkEncode::Parquet => {
464                let props = WriterProperties::builder().set_compression(Compression::SNAPPY);
465                let parquet_writer: tokio_util::compat::Compat<opendal::FuturesAsyncWriter> =
466                    object_writer.into_futures_async_write().compat_write();
467                self.sink_writer = Some(FileWriterEnum::ParquetFileWriter(
468                    AsyncArrowWriter::try_new(
469                        parquet_writer,
470                        self.schema.clone(),
471                        Some(props.build()),
472                    )?,
473                ));
474            }
475            _ => {
476                self.sink_writer = Some(FileWriterEnum::FileWriter(object_writer));
477            }
478        }
479        self.current_bached_row_num = 0;
480
481        Ok(())
482    }
483}
484
485fn convert_rw_schema_to_arrow_schema(
486    rw_schema: risingwave_common::catalog::Schema,
487) -> anyhow::Result<arrow_schema_iceberg::Schema> {
488    let mut schema_fields = HashMap::new();
489    rw_schema.fields.iter().for_each(|field| {
490        let res = schema_fields.insert(&field.name, &field.data_type);
491        // This assert is to make sure there is no duplicate field name in the schema.
492        assert!(res.is_none())
493    });
494    let mut arrow_fields = vec![];
495    for rw_field in &rw_schema.fields {
496        let arrow_field = IcebergArrowConvert
497            .to_arrow_field(&rw_field.name.clone(), &rw_field.data_type.clone())?;
498
499        arrow_fields.push(arrow_field);
500    }
501
502    Ok(arrow_schema_iceberg::Schema::new(arrow_fields))
503}
504
505/// `BatchingStrategy` represents the strategy for batching data before writing to files.
506///
507/// This struct contains settings that control how data is collected and
508/// partitioned based on specified criteria:
509///
510/// - `max_row_count`: Optional maximum number of rows to accumulate before writing.
511/// - `rollover_seconds`: Optional time interval (in seconds) to trigger a write,
512///   regardless of the number of accumulated rows.
513/// - `path_partition_prefix`: Specifies how files are organized into directories
514///   based on creation time (e.g., by day, month, or hour).
515
516#[serde_as]
517#[derive(Default, Deserialize, Debug, Clone, WithOptions)]
518pub struct BatchingStrategy {
519    #[serde(default = "default_max_row_count")]
520    #[serde_as(as = "DisplayFromStr")]
521    pub max_row_count: usize,
522    #[serde(default = "default_rollover_seconds")]
523    #[serde_as(as = "DisplayFromStr")]
524    pub rollover_seconds: usize,
525    #[serde(default)]
526    #[serde_as(as = "Option<DisplayFromStr>")]
527    pub path_partition_prefix: Option<PathPartitionPrefix>,
528}
529
530/// `PathPartitionPrefix` defines the granularity of file partitions based on creation time.
531///
532/// Each variant specifies how files are organized into directories:
533/// - `None`: No partitioning.
534/// - `Day`: Files are written in a directory for each day.
535/// - `Month`: Files are written in a directory for each month.
536/// - `Hour`: Files are written in a directory for each hour.
537#[derive(Default, Debug, Clone, PartialEq, Display, Deserialize, EnumString)]
538#[strum(serialize_all = "snake_case")]
539pub enum PathPartitionPrefix {
540    #[default]
541    None = 0,
542    #[serde(alias = "day")]
543    Day = 1,
544    #[serde(alias = "month")]
545    Month = 2,
546    #[serde(alias = "hour")]
547    Hour = 3,
548}
risingwave_connector/sink/file_sink/opendal_sink.rs

risingwave_connector/sink/file_sink/
opendal_sink.rs
