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

risingwave_connector/sink/file_sink/
opendal_sink.rs
