risingwave_connector/sink/iceberg/
pk_index_file.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//! Data file naming for pk-index sinks that do not use the legacy file path layout.
16//!
17//! Every data file written by the sink writer or the pk-index compaction is placed at
18//! `{data_location}/{uuid}.parquet`, so the PK index only needs to store the 16-byte UUID and
19//! can rebuild the full path from it.
20
21use anyhow::anyhow;
22use iceberg::spec::{DataFileFormat, PartitionKey, TableMetadata};
23use iceberg::writer::file_writer::location_generator::{
24    DefaultLocationGenerator, FileNameGenerator, LocationGenerator,
25};
26use uuid::Uuid;
27
28use crate::sink::{Result, SinkError};
29
30const PK_INDEX_DATA_FILE_SUFFIX: &str = ".parquet";
31
32/// Maps between data file UUIDs and data file paths of a pk-index sink table.
33#[derive(Clone, Debug, PartialEq, Eq)]
34pub struct PkIndexDataFileLocator {
35    data_location: String,
36}
37
38impl PkIndexDataFileLocator {
39    pub fn new(table_metadata: &TableMetadata) -> Result<Self> {
40        // Resolve the data location exactly as the default generator does, so the iceberg-compaction
41        // writer, which places files with `DefaultLocationGenerator`, agrees with this locator.
42        let location = DefaultLocationGenerator::new(table_metadata)
43            .map_err(|err| SinkError::Iceberg(anyhow!(err)))?
44            .generate_location(None, "");
45        Ok(Self::with_data_location(location))
46    }
47
48    pub fn with_data_location(data_location: impl Into<String>) -> Self {
49        let mut data_location = data_location.into();
50        while data_location.ends_with('/') {
51            data_location.pop();
52        }
53        Self { data_location }
54    }
55
56    pub fn file_path(&self, uuid: Uuid) -> String {
57        format!(
58            "{}/{}{}",
59            self.data_location,
60            uuid.hyphenated(),
61            PK_INDEX_DATA_FILE_SUFFIX
62        )
63    }
64
65    pub fn file_path_from_bytes(&self, uuid: &[u8]) -> Result<String> {
66        let uuid = Uuid::from_slice(uuid).map_err(|err| {
67            SinkError::Iceberg(anyhow!(err).context("invalid pk-index data file uuid"))
68        })?;
69        Ok(self.file_path(uuid))
70    }
71
72    /// Extracts the UUID from a data file path. Fails if the path is not exactly the one
73    /// [`Self::file_path`] would produce, since the full path could not be rebuilt otherwise.
74    pub fn file_uuid(&self, path: &str) -> Result<Uuid> {
75        path.strip_prefix(self.data_location.as_str())
76            .and_then(|name| name.strip_prefix('/'))
77            .and_then(|name| name.strip_suffix(PK_INDEX_DATA_FILE_SUFFIX))
78            .and_then(|uuid| Uuid::try_parse(uuid).ok())
79            .filter(|uuid| self.file_path(*uuid) == path)
80            .ok_or_else(|| {
81                SinkError::Iceberg(anyhow!(
82                    "data file path {path} does not match the pk-index layout `{}/{{uuid}}{}`",
83                    self.data_location,
84                    PK_INDEX_DATA_FILE_SUFFIX
85                ))
86            })
87    }
88}
89
90/// Places every data file directly under the data location, without partition directories.
91impl LocationGenerator for PkIndexDataFileLocator {
92    fn generate_location(&self, _partition_key: Option<&PartitionKey>, file_name: &str) -> String {
93        format!("{}/{}", self.data_location, file_name)
94    }
95}
96
97/// Names every data file `{uuid}.parquet` with a fresh UUID v7.
98#[derive(Clone, Debug, Default)]
99pub struct PkIndexDataFileNameGenerator;
100
101impl FileNameGenerator for PkIndexDataFileNameGenerator {
102    fn generate_file_name(&self) -> String {
103        debug_assert_eq!(
104            PK_INDEX_DATA_FILE_SUFFIX,
105            format!(".{}", DataFileFormat::Parquet)
106        );
107        format!(
108            "{}{}",
109            Uuid::now_v7().hyphenated(),
110            PK_INDEX_DATA_FILE_SUFFIX
111        )
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    #[test]
120    fn test_file_path_round_trip() {
121        let locator = PkIndexDataFileLocator::with_data_location("s3://bucket/db/t/data/");
122        let file_name = PkIndexDataFileNameGenerator.generate_file_name();
123        let path = locator.generate_location(None, &file_name);
124        assert!(path.starts_with("s3://bucket/db/t/data/"));
125        assert!(path.ends_with(".parquet"));
126
127        let uuid = locator.file_uuid(&path).unwrap();
128        assert_eq!(uuid.get_version_num(), 7);
129        assert_eq!(locator.file_path(uuid), path);
130        assert_eq!(locator.file_path_from_bytes(uuid.as_bytes()).unwrap(), path);
131    }
132
133    #[test]
134    fn test_file_uuid_rejects_other_layouts() {
135        let locator = PkIndexDataFileLocator::with_data_location("s3://bucket/t/data");
136        let uuid = Uuid::now_v7();
137        for path in [
138            format!("s3://bucket/t/data/p=1/{uuid}.parquet"),
139            format!("s3://bucket/other/data/{uuid}.parquet"),
140            format!("s3://bucket/t/data/{}.parquet", uuid.simple()),
141            format!("s3://bucket/t/data/{uuid}.orc"),
142            "s3://bucket/t/data/1-00000-0193.parquet".to_owned(),
143        ] {
144            assert!(locator.file_uuid(&path).is_err(), "{path}");
145        }
146        assert!(locator.file_path_from_bytes(&[1, 2, 3]).is_err());
147    }
148}
risingwave_connector/sink/iceberg/pk_index_file.rs

risingwave_connector/sink/iceberg/
pk_index_file.rs
