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

risingwave_connector/sink/file_sink/
opendal_sink.rs
