risingwave_connector/sink/iceberg/
position_delete.rs

1// Copyright 2026 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
15//! Shared iceberg position-delete (Puffin deletion vector) helpers.
16
17use std::collections::HashMap;
18use std::sync::Arc;
19
20use anyhow::{Context, Result, anyhow, bail};
21use futures::StreamExt;
22use iceberg::arrow::schema_to_arrow_schema;
23use iceberg::delete_vector::DeleteVector;
24use iceberg::io::FileIO;
25use iceberg::puffin::{CompressionCodec, PuffinReader, PuffinWriter};
26use iceberg::spec::{DataContentType, DataFile, DataFileBuilder, DataFileFormat, PartitionKey};
27use iceberg::table::Table;
28use iceberg::writer::base_writer::position_delete_file_writer::POSITION_DELETE_SCHEMA;
29use iceberg::writer::file_writer::location_generator::{
30    DefaultFileNameGenerator, DefaultLocationGenerator, FileNameGenerator, LocationGenerator,
31};
32use iceberg::writer::file_writer::{
33    FileWriter, FileWriterBuilder, ParquetWriter, ParquetWriterBuilder,
34};
35use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask};
36use parquet::file::properties::WriterProperties;
37use risingwave_common::array::arrow::arrow_array_iceberg::{
38    Array, ArrayRef, Int64Array, RecordBatch, StringArray,
39};
40use risingwave_common::array::arrow::arrow_schema_iceberg::SchemaRef as ArrowSchemaRef;
41
42use crate::sink::iceberg::{IcebergConfig, PARQUET_CREATED_BY};
43use crate::source::iceberg::parquet_file_handler::ParquetFileReader;
44
45/// Puffin blob property for deletion vector cardinality.
46const DELETION_VECTOR_PROPERTY_CARDINALITY: &str = "cardinality";
47/// Puffin blob property for referenced data file path.
48const DELETION_VECTOR_PROPERTY_REFERENCED_DATA_FILE: &str = "referenced-data-file";
49
50/// Reads the deletion-vector positions of a single Puffin DV `DataFile`.
51pub async fn read_dv_positions_from_data_file(
52    file_io: &FileIO,
53    data_file: &DataFile,
54) -> Result<DeleteVector> {
55    let blob_offset = data_file.content_offset().with_context(|| {
56        format!(
57            "DV file {} missing content_offset for referenced data file {:?}",
58            data_file.file_path(),
59            data_file.referenced_data_file()
60        )
61    })?;
62    let blob_length = data_file.content_size_in_bytes().with_context(|| {
63        format!(
64            "DV file {} missing content_size_in_bytes for referenced data file {:?}",
65            data_file.file_path(),
66            data_file.referenced_data_file()
67        )
68    })?;
69
70    let input_file = file_io.new_input(data_file.file_path())?;
71    let puffin_reader = PuffinReader::new(input_file);
72    let file_metadata = puffin_reader.file_metadata().await?;
73    let blob_metadata = file_metadata
74        .blobs()
75        .iter()
76        .find(|blob| blob.offset() == blob_offset as u64 && blob.length() == blob_length as u64)
77        .with_context(|| {
78            format!(
79                "DV blob metadata not found in {} at offset={} length={}",
80                data_file.file_path(),
81                blob_offset,
82                blob_length
83            )
84        })?;
85    let blob = puffin_reader.blob(blob_metadata).await?;
86
87    let delete_vector = DeleteVector::from_puffin_blob(blob)?;
88    Ok(delete_vector)
89}
90
91/// Reads the positions stored in a V2 Parquet position-delete file into a [`DeleteVector`].
92///
93/// The file's schema is `(file_path, pos)`. Callers only invoke this after the entry's
94/// `referenced_data_file` already matched the target data file, and the files we write
95/// are file-scoped (every row shares one `file_path`), so the `file_path` column is
96/// redundant here: we project only the `pos` column and read every value.
97pub async fn read_parquet_position_deletes_from_file(
98    file_io: &FileIO,
99    delete_file: &DataFile,
100) -> Result<DeleteVector> {
101    let input_file = file_io.new_input(delete_file.file_path())?;
102    let metadata = input_file.metadata().await?;
103    let reader = input_file.reader().await?;
104    let parquet_reader = ParquetFileReader::new(metadata, reader);
105    let builder = ParquetRecordBatchStreamBuilder::new(parquet_reader).await?;
106    // Project only the `pos` leaf (column index 1) so the `file_path` column is never decoded.
107    let projection = ProjectionMask::leaves(builder.parquet_schema(), [1]);
108    let mut stream = builder.with_projection(projection).build()?;
109
110    let mut delete_vector = DeleteVector::default();
111    while let Some(batch) = stream.next().await {
112        let batch = batch?;
113        // Only the projected `pos` column is present in the batch.
114        let positions = batch.columns()[0]
115            .as_any()
116            .downcast_ref::<Int64Array>()
117            .context("position-delete pos column should be an Int64Array")?;
118        for pos in positions {
119            let pos = pos.with_context(|| {
120                format!(
121                    "null value in position-delete file {}",
122                    delete_file.file_path()
123                )
124            })?;
125            delete_vector.insert(pos as u64);
126        }
127    }
128
129    Ok(delete_vector)
130}
131
132/// Reads the deleted positions of a single position-delete `DataFile` regardless of on-disk format,
133pub async fn read_position_deletes_from_file(
134    file_io: &FileIO,
135    delete_file: &DataFile,
136) -> Result<DeleteVector> {
137    match delete_file.file_format() {
138        DataFileFormat::Puffin => read_dv_positions_from_data_file(file_io, delete_file).await,
139        DataFileFormat::Parquet => {
140            read_parquet_position_deletes_from_file(file_io, delete_file).await
141        }
142        other => bail!(
143            "position-delete file {} has unsupported format {:?}; expected Puffin or Parquet",
144            delete_file.file_path(),
145            other
146        ),
147    }
148}
149
150/// Writes `delete_vector` as a single Puffin deletion-vector blob referencing `data_file_path`,
151/// and returns its [`DataFile`] metadata (content `PositionDeletes`, format `Puffin`) with
152/// `referenced_data_file` set.
153pub async fn write_dv_puffin_file(
154    table: &Table,
155    location_generator: &DefaultLocationGenerator,
156    file_name_generator: &DefaultFileNameGenerator,
157    data_file_path: String,
158    delete_vector: &DeleteVector,
159    partition_key: Option<&PartitionKey>,
160) -> Result<DataFile> {
161    let file_name = file_name_generator.generate_file_name();
162    let location = location_generator.generate_location(partition_key, &file_name);
163    let output_file = table.file_io().new_output(&location)?;
164    let mut writer = PuffinWriter::new(&output_file, HashMap::new(), false).await?;
165
166    let cardinality = delete_vector.len();
167    let properties = HashMap::from([
168        (
169            DELETION_VECTOR_PROPERTY_CARDINALITY.to_owned(),
170            cardinality.to_string(),
171        ),
172        (
173            DELETION_VECTOR_PROPERTY_REFERENCED_DATA_FILE.to_owned(),
174            data_file_path.clone(),
175        ),
176    ]);
177    let blob = delete_vector.to_puffin_blob(properties)?;
178    writer.add(blob, CompressionCodec::None).await?;
179
180    let result = writer.close_with_metadata().await?;
181    let blob_metadata = result
182        .blobs_metadata
183        .first()
184        .context("blob metadata should be present")?;
185
186    let mut builder = DataFileBuilder::default();
187    builder
188        .content(DataContentType::PositionDeletes)
189        .file_path(location)
190        .file_format(DataFileFormat::Puffin)
191        .record_count(cardinality)
192        .file_size_in_bytes(result.file_size_in_bytes)
193        .referenced_data_file(Some(data_file_path))
194        .content_offset(Some(blob_metadata.offset() as i64))
195        .content_size_in_bytes(Some(blob_metadata.length() as i64));
196    if let Some(partition_key) = partition_key {
197        builder
198            .partition(partition_key.data().clone())
199            .partition_spec_id(partition_key.spec().spec_id());
200    }
201    builder
202        .build()
203        .context("Failed to build deletion vector file metadata")
204}
205
206/// How many positions to buffer before flushing one `(file_path, pos)` batch to the writer.
207const POSITION_DELETE_WRITE_CHUNK_SIZE: usize = 1024;
208
209/// Writes `delete_vector` as a single file-scoped Parquet position-delete file referencing
210/// `data_file_path`, and returns its [`DataFile`] metadata (content `PositionDeletes`, format
211/// `Parquet`) with `referenced_data_file` set.
212pub async fn write_parquet_position_delete_file(
213    table: &Table,
214    location_generator: &DefaultLocationGenerator,
215    file_name_generator: &DefaultFileNameGenerator,
216    config: &IcebergConfig,
217    data_file_path: String,
218    delete_vector: &DeleteVector,
219    partition_key: Option<&PartitionKey>,
220) -> Result<DataFile> {
221    let file_name = file_name_generator.generate_file_name();
222    let location = location_generator.generate_location(partition_key, &file_name);
223    let output_file = table.file_io().new_output(&location)?;
224
225    let parquet_writer_properties = WriterProperties::builder()
226        .set_compression(config.get_parquet_compression())
227        .set_max_row_group_bytes(config.write_parquet_max_row_group_bytes())
228        .set_created_by(PARQUET_CREATED_BY.to_owned())
229        .build();
230    let mut writer = ParquetWriterBuilder::new(
231        parquet_writer_properties,
232        POSITION_DELETE_SCHEMA.clone().into(),
233    )
234    .build(output_file)
235    .await?;
236
237    // The position-delete schema is `(file_path, pos)` with reserved field IDs; derive the matching
238    // Arrow schema so the written column field IDs line up.
239    let arrow_schema: ArrowSchemaRef = Arc::new(schema_to_arrow_schema(&POSITION_DELETE_SCHEMA)?);
240
241    let mut positions: Vec<i64> = Vec::with_capacity(POSITION_DELETE_WRITE_CHUNK_SIZE);
242    for pos in delete_vector.iter() {
243        positions.push(pos as i64);
244        if positions.len() == POSITION_DELETE_WRITE_CHUNK_SIZE {
245            write_position_delete_chunk(
246                &mut writer,
247                &arrow_schema,
248                &data_file_path,
249                std::mem::take(&mut positions),
250            )
251            .await?;
252            positions.reserve(POSITION_DELETE_WRITE_CHUNK_SIZE);
253        }
254    }
255    if !positions.is_empty() {
256        write_position_delete_chunk(&mut writer, &arrow_schema, &data_file_path, positions).await?;
257    }
258
259    let data_files = writer.close().await?;
260    // `close` will yield exactly one builder here.
261    let [mut builder] = data_files.try_into().map_err(|_| {
262        anyhow!("position-delete writer produced invalid file count for {data_file_path}")
263    })?;
264
265    // `ParquetWriter` builds the file as `DataContentType::Data` with an empty partition; override
266    // those for a file-scoped V2 position-delete file and attach `referenced_data_file`.
267    builder
268        .content(DataContentType::PositionDeletes)
269        .referenced_data_file(Some(data_file_path));
270    if let Some(partition_key) = partition_key {
271        builder
272            .partition(partition_key.data().clone())
273            .partition_spec_id(partition_key.spec().spec_id());
274    }
275    builder
276        .build()
277        .context("Failed to build position-delete file metadata")
278}
279
280/// Writes one chunk of `positions` as a `(file_path, pos)` batch into `writer`. Every row shares
281/// `data_file_path` because the delete file is file-scoped.
282async fn write_position_delete_chunk(
283    writer: &mut ParquetWriter,
284    arrow_schema: &ArrowSchemaRef,
285    data_file_path: &str,
286    positions: Vec<i64>,
287) -> Result<()> {
288    let path_column: ArrayRef = Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
289        data_file_path,
290        positions.len(),
291    )));
292    let pos_column: ArrayRef = Arc::new(Int64Array::from(positions));
293    let batch = RecordBatch::try_new(arrow_schema.clone(), vec![path_column, pos_column])
294        .map_err(|e| anyhow!(e))?;
295    writer.write(&batch).await?;
296    Ok(())
297}
risingwave_connector/sink/iceberg/position_delete.rs

risingwave_connector/sink/iceberg/
position_delete.rs
