paimon_datafusion/
blob_reader.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::{Arc, RwLock};
19
20use paimon::io::FileIO;
21
22#[derive(Clone, Debug)]
23struct BlobFileIO {
24    prefix: String,
25    file_io: FileIO,
26}
27
28/// Session-scoped registry of Paimon [`FileIO`] instances for BlobDescriptor reads.
29#[derive(Clone, Debug, Default)]
30pub struct BlobReaderRegistry {
31    readers: Arc<RwLock<Vec<BlobFileIO>>>,
32}
33
34fn allows_backslash_boundary(location: &str) -> bool {
35    let bytes = location.as_bytes();
36    location.starts_with("file:/")
37        || (bytes.len() >= 3
38            && bytes[0].is_ascii_alphabetic()
39            && bytes[1] == b':'
40            && matches!(bytes[2], b'\\' | b'/'))
41}
42
43fn matches_location(uri: &str, location: &str) -> bool {
44    uri == location
45        || uri.strip_prefix(location).is_some_and(|suffix| {
46            location.ends_with('/')
47                || suffix.starts_with('/')
48                || (allows_backslash_boundary(location)
49                    && (location.ends_with('\\') || suffix.starts_with('\\')))
50        })
51}
52
53impl BlobReaderRegistry {
54    pub fn register(&self, prefix: impl Into<String>, file_io: FileIO) {
55        let prefix = prefix.into();
56        let mut readers = self.readers.write().unwrap_or_else(|e| e.into_inner());
57        if let Some(existing) = readers.iter_mut().find(|reader| reader.prefix == prefix) {
58            existing.file_io = file_io;
59            return;
60        }
61        readers.push(BlobFileIO { prefix, file_io });
62    }
63
64    pub fn register_if_absent(&self, prefix: impl Into<String>, file_io: FileIO) {
65        let prefix = prefix.into();
66        let mut readers = self.readers.write().unwrap_or_else(|e| e.into_inner());
67        if readers.iter().any(|reader| reader.prefix == prefix) {
68            return;
69        }
70        readers.push(BlobFileIO { prefix, file_io });
71    }
72
73    pub fn resolve(&self, uri: &str) -> Option<FileIO> {
74        let readers = self.readers.read().unwrap_or_else(|e| e.into_inner());
75        readers
76            .iter()
77            .filter(|reader| matches_location(uri, &reader.prefix))
78            .max_by_key(|reader| reader.prefix.len())
79            .map(|reader| reader.file_io.clone())
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use std::fs;
86
87    use paimon::io::{FileIOBuilder, FileRead};
88    use paimon::spec::BlobDescriptor;
89
90    use super::*;
91
92    fn memory_file_io() -> FileIO {
93        FileIOBuilder::new("memory").build().unwrap()
94    }
95
96    async fn write_file(file_io: &FileIO, uri: &str, contents: &[u8]) {
97        file_io
98            .new_output(uri)
99            .unwrap()
100            .write(contents.to_vec().into())
101            .await
102            .unwrap();
103    }
104
105    async fn assert_resolves_to(registry: &BlobReaderRegistry, uri: &str, expected: &[u8]) {
106        let file_io = registry
107            .resolve(uri)
108            .unwrap_or_else(|| panic!("expected registered FileIO for {uri}"));
109        let bytes = file_io.new_input(uri).unwrap().read().await.unwrap();
110        assert_eq!(&bytes[..], expected);
111    }
112
113    #[test]
114    fn resolves_only_same_location_or_descendants() {
115        let cases = [
116            ("memory:/foo", "memory:/foo", true),
117            ("memory:/foo", "memory:/foo/blob.bin", true),
118            ("memory:/foo", "memory:/foo/nested/blob.bin", true),
119            ("memory:/foo", "memory:/foobar/private.bin", false),
120            ("s3://bucket", "s3://bucket-other/blob.bin", false),
121            ("s3://bucket/table", "gs://bucket/table/blob.bin", false),
122            ("memory:/foo/", "memory:/foo/", true),
123            ("memory:/foo/", "memory:/foo/blob.bin", true),
124            ("memory:/foo/", "memory:/foo", false),
125            (r"C:\warehouse\table", r"C:\warehouse\table\blob.bin", true),
126            (
127                r"C:\warehouse\table",
128                r"C:\warehouse\table-other\blob.bin",
129                false,
130            ),
131            (
132                "C:\\warehouse\\table\\",
133                r"C:\warehouse\table\blob.bin",
134                true,
135            ),
136            ("C:\\warehouse\\table\\", r"C:\warehouse\table", false),
137            (
138                r"file:/C:\warehouse\table",
139                r"file:/C:\warehouse\table\blob.bin",
140                true,
141            ),
142            (
143                r"file:/C:\warehouse\table",
144                r"file:/C:\warehouse\table-other\blob.bin",
145                false,
146            ),
147            ("s3://bucket/table", r"s3://bucket/table\blob.bin", false),
148        ];
149
150        for (location, uri, expected) in cases {
151            let registry = BlobReaderRegistry::default();
152            registry.register(location, memory_file_io());
153            assert_eq!(
154                registry.resolve(uri).is_some(),
155                expected,
156                "location={location}, uri={uri}"
157            );
158        }
159    }
160
161    #[tokio::test]
162    async fn resolves_with_longest_valid_location() {
163        let parent_file_io = memory_file_io();
164        let child_file_io = memory_file_io();
165        let registry = BlobReaderRegistry::default();
166        registry.register("memory:/warehouse/db", parent_file_io.clone());
167        registry.register("memory:/warehouse/db/foo", child_file_io.clone());
168
169        let child_uri = "memory:/warehouse/db/foo/blob.bin";
170        write_file(&parent_file_io, child_uri, b"parent").await;
171        write_file(&child_file_io, child_uri, b"child").await;
172        assert_resolves_to(&registry, child_uri, b"child").await;
173
174        let sibling_uri = "memory:/warehouse/db/foobar/private.bin";
175        write_file(&parent_file_io, sibling_uri, b"parent").await;
176        write_file(&child_file_io, sibling_uri, b"child").await;
177        assert_resolves_to(&registry, sibling_uri, b"parent").await;
178    }
179
180    #[tokio::test]
181    async fn preserves_registration_semantics() {
182        let uri = "memory:/warehouse/db/foo";
183        let initial_file_io = memory_file_io();
184        let replacement_file_io = memory_file_io();
185        write_file(&initial_file_io, uri, b"initial").await;
186        write_file(&replacement_file_io, uri, b"replacement").await;
187
188        let registry = BlobReaderRegistry::default();
189        registry.register(uri, initial_file_io.clone());
190        registry.register(uri, replacement_file_io.clone());
191        assert_resolves_to(&registry, uri, b"replacement").await;
192
193        let registry = BlobReaderRegistry::default();
194        registry.register(uri, initial_file_io);
195        registry.register_if_absent(uri, replacement_file_io);
196        assert_resolves_to(&registry, uri, b"initial").await;
197    }
198
199    #[tokio::test]
200    async fn resolves_file_blob_descriptor_with_file_io() {
201        let directory = tempfile::tempdir().unwrap();
202        let blob_path = directory.path().join("blob.bin");
203        fs::write(&blob_path, b"prefixpayloadsuffix").unwrap();
204
205        let descriptor = BlobDescriptor::new(
206            blob_path.to_string_lossy().to_string(),
207            6,
208            "payload".len() as i64,
209        );
210        let descriptor = BlobDescriptor::deserialize(&descriptor.serialize()).unwrap();
211
212        let registry = BlobReaderRegistry::default();
213        let file_io = FileIOBuilder::new("file").build().unwrap();
214        registry.register(directory.path().to_string_lossy().to_string(), file_io);
215
216        let resolved_file_io = registry
217            .resolve(descriptor.uri())
218            .expect("file blob descriptor should resolve to registered FileIO");
219        let input = resolved_file_io.new_input(descriptor.uri()).unwrap();
220        let reader = input.reader().await.unwrap();
221        let start = descriptor.offset() as u64;
222        let end = start + descriptor.length() as u64;
223        let bytes = reader.read(start..end).await.unwrap();
224
225        assert_eq!(&bytes[..], b"payload");
226    }
227}
paimon_datafusion/blob_reader.rs

paimon_datafusion/
blob_reader.rs
