ckb_rocksdb/ops/
multi_get.rs

1// Copyright 2019 Tyler Neely
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
16use crate::{ffi, ColumnFamily, DBPinnableSlice, DBVector};
17use libc::c_char;
18use std::ptr;
19
20use crate::{handle::Handle, Error, ReadOptions};
21
22pub type CFAndKey<'a> = (&'a ColumnFamily, Box<[u8]>);
23
24pub trait MultiGet<R> {
25    fn multi_get_full<K, I>(
26        &self,
27        keys: I,
28        readopts: Option<&R>,
29    ) -> Vec<Result<Option<DBVector>, Error>>
30    where
31        K: AsRef<[u8]>,
32        I: IntoIterator<Item = K>;
33
34    fn multi_get<K, I>(&self, keys: I) -> Vec<Result<Option<DBVector>, Error>>
35    where
36        K: AsRef<[u8]>,
37        I: IntoIterator<Item = K>,
38    {
39        self.multi_get_full(keys, None)
40    }
41
42    fn multi_get_opt<K, I>(&self, keys: I, readopts: &R) -> Vec<Result<Option<DBVector>, Error>>
43    where
44        K: AsRef<[u8]>,
45        I: IntoIterator<Item = K>,
46    {
47        self.multi_get_full(keys, Some(readopts))
48    }
49}
50
51pub trait MultiGetCF<R> {
52    fn multi_get_cf_full<'a, K, I>(
53        &self,
54        keys_cf: I,
55        readopts: Option<&R>,
56    ) -> Vec<Result<Option<DBVector>, Error>>
57    where
58        K: AsRef<[u8]>,
59        I: IntoIterator<Item = (&'a ColumnFamily, K)>;
60
61    fn multi_get_cf<'a, K, I>(&self, keys_cf: I) -> Vec<Result<Option<DBVector>, Error>>
62    where
63        K: AsRef<[u8]>,
64        I: IntoIterator<Item = (&'a ColumnFamily, K)>,
65    {
66        self.multi_get_cf_full(keys_cf, None)
67    }
68
69    fn multi_get_cf_opt<'a, K, I>(
70        &self,
71        keys_cf: I,
72        readopts: &R,
73    ) -> Vec<Result<Option<DBVector>, Error>>
74    where
75        K: AsRef<[u8]>,
76        I: IntoIterator<Item = (&'a ColumnFamily, K)>,
77    {
78        self.multi_get_cf_full(keys_cf, Some(readopts))
79    }
80}
81
82impl<T> MultiGet<ReadOptions> for T
83where
84    T: Handle<ffi::rocksdb_t> + super::Read,
85{
86    /// Return the values associated with the given keys using read options.
87    fn multi_get_full<K, I>(
88        &self,
89        keys: I,
90        readopts: Option<&ReadOptions>,
91    ) -> Vec<Result<Option<DBVector>, Error>>
92    where
93        K: AsRef<[u8]>,
94        I: IntoIterator<Item = K>,
95    {
96        let mut default_readopts = None;
97        let ro_handle = match ReadOptions::input_or_default(readopts, &mut default_readopts) {
98            Ok(ro) => ro,
99            Err(e) => {
100                let key_count = keys.into_iter().count();
101
102                return vec![e; key_count]
103                    .iter()
104                    .map(|e| Err(e.to_owned()))
105                    .collect();
106            }
107        };
108
109        let (keys, keys_sizes): (Vec<Box<[u8]>>, Vec<_>) = keys
110            .into_iter()
111            .map(|k| (Box::from(k.as_ref()), k.as_ref().len()))
112            .unzip();
113        let ptr_keys: Vec<_> = keys.iter().map(|k| k.as_ptr() as *const c_char).collect();
114        let mut values = vec![ptr::null_mut(); keys.len()];
115        let mut values_sizes = vec![0_usize; keys.len()];
116        let mut errors = vec![ptr::null_mut(); keys.len()];
117        unsafe {
118            ffi::rocksdb_multi_get(
119                self.handle(),
120                ro_handle,
121                ptr_keys.len(),
122                ptr_keys.as_ptr(),
123                keys_sizes.as_ptr(),
124                values.as_mut_ptr(),
125                values_sizes.as_mut_ptr(),
126                errors.as_mut_ptr(),
127            );
128        }
129        convert_values(values, values_sizes, errors)
130    }
131}
132
133impl<T> MultiGetCF<ReadOptions> for T
134where
135    T: Handle<ffi::rocksdb_t> + super::Read,
136{
137    /// Return the values associated with the given keys and column families using read options.
138    fn multi_get_cf_full<'a, K, I>(
139        &self,
140        keys: I,
141        readopts: Option<&ReadOptions>,
142    ) -> Vec<Result<Option<DBVector>, Error>>
143    where
144        K: AsRef<[u8]>,
145        I: IntoIterator<Item = (&'a ColumnFamily, K)>,
146    {
147        let mut default_readopts = None;
148        let ro_handle = match ReadOptions::input_or_default(readopts, &mut default_readopts) {
149            Ok(ro) => ro,
150            Err(e) => {
151                let key_count = keys.into_iter().count();
152                return vec![e; key_count]
153                    .iter()
154                    .map(|e| Err(e.to_owned()))
155                    .collect();
156            }
157        };
158
159        let (cfs_and_keys, keys_sizes): (Vec<CFAndKey>, Vec<_>) = keys
160            .into_iter()
161            .map(|(cf, key)| ((cf, Box::from(key.as_ref())), key.as_ref().len()))
162            .unzip();
163        let ptr_keys: Vec<_> = cfs_and_keys
164            .iter()
165            .map(|(_, k)| k.as_ptr() as *const c_char)
166            .collect();
167        let ptr_cfs: Vec<_> = cfs_and_keys
168            .iter()
169            .map(|(c, _)| c.inner as *const _)
170            .collect();
171
172        let mut values = vec![ptr::null_mut(); ptr_keys.len()];
173        let mut values_sizes = vec![0_usize; ptr_keys.len()];
174        let mut errors = vec![ptr::null_mut(); ptr_keys.len()];
175        unsafe {
176            ffi::rocksdb_multi_get_cf(
177                self.handle(),
178                ro_handle,
179                ptr_cfs.as_ptr(),
180                ptr_keys.len(),
181                ptr_keys.as_ptr(),
182                keys_sizes.as_ptr(),
183                values.as_mut_ptr(),
184                values_sizes.as_mut_ptr(),
185                errors.as_mut_ptr(),
186            );
187        }
188
189        convert_values(values, values_sizes, errors)
190    }
191}
192
193pub trait BatchedMultiGetCF<R> {
194    fn batched_multi_get_cf_full<'a, K, I>(
195        &self,
196        cf: &ColumnFamily,
197        keys: I,
198        sorted_input: bool,
199        readopts: Option<&R>,
200    ) -> Vec<Result<Option<DBPinnableSlice>, Error>>
201    where
202        K: AsRef<[u8]> + 'a + ?Sized,
203        I: IntoIterator<Item = &'a K>;
204
205    fn batched_multi_get_cf<'a, K, I>(
206        &self,
207        cf: &ColumnFamily,
208        keys: I,
209        sorted_input: bool,
210    ) -> Vec<Result<Option<DBPinnableSlice>, Error>>
211    where
212        K: AsRef<[u8]> + 'a + ?Sized,
213        I: IntoIterator<Item = &'a K>,
214    {
215        self.batched_multi_get_cf_full(cf, keys, sorted_input, None)
216    }
217
218    fn batched_multi_get_cf_opt<'a, K, I>(
219        &self,
220        cf: &ColumnFamily,
221        keys: I,
222        sorted_input: bool,
223        readopts: &R,
224    ) -> Vec<Result<Option<DBPinnableSlice>, Error>>
225    where
226        K: AsRef<[u8]> + 'a + ?Sized,
227        I: IntoIterator<Item = &'a K>,
228    {
229        self.batched_multi_get_cf_full(cf, keys, sorted_input, Some(readopts))
230    }
231}
232
233impl<T> BatchedMultiGetCF<ReadOptions> for T
234where
235    T: Handle<ffi::rocksdb_t> + super::Read,
236{
237    fn batched_multi_get_cf_full<'a, K, I>(
238        &self,
239        cf: &ColumnFamily,
240        keys: I,
241        sorted_input: bool,
242        readopts: Option<&ReadOptions>,
243    ) -> Vec<Result<Option<DBPinnableSlice>, Error>>
244    where
245        K: AsRef<[u8]> + 'a + ?Sized,
246        I: IntoIterator<Item = &'a K>,
247    {
248        let mut default_readopts = None;
249        let ro_handle = match ReadOptions::input_or_default(readopts, &mut default_readopts) {
250            Ok(ro) => ro,
251            Err(e) => {
252                let key_count = keys.into_iter().count();
253
254                return vec![e; key_count]
255                    .iter()
256                    .map(|e| Err(e.to_owned()))
257                    .collect();
258            }
259        };
260
261        let (ptr_keys, keys_sizes): (Vec<_>, Vec<_>) = keys
262            .into_iter()
263            .map(|k| {
264                let k = k.as_ref();
265                (k.as_ptr() as *const c_char, k.len())
266            })
267            .unzip();
268
269        let mut pinned_values = vec![ptr::null_mut(); ptr_keys.len()];
270        let mut errors = vec![ptr::null_mut(); ptr_keys.len()];
271
272        unsafe {
273            ffi::rocksdb_batched_multi_get_cf(
274                self.handle(),
275                ro_handle,
276                cf.inner,
277                ptr_keys.len(),
278                ptr_keys.as_ptr(),
279                keys_sizes.as_ptr(),
280                pinned_values.as_mut_ptr(),
281                errors.as_mut_ptr(),
282                sorted_input,
283            );
284            pinned_values
285                .into_iter()
286                .zip(errors.into_iter())
287                .map(|(v, e)| {
288                    if e.is_null() {
289                        if v.is_null() {
290                            Ok(None)
291                        } else {
292                            Ok(Some(DBPinnableSlice::from_c(v)))
293                        }
294                    } else {
295                        Err(Error::new(crate::ffi_util::error_message(e)))
296                    }
297                })
298                .collect()
299        }
300    }
301}
302
303pub fn convert_values(
304    values: Vec<*mut c_char>,
305    values_sizes: Vec<usize>,
306    errors: Vec<*mut c_char>,
307) -> Vec<Result<Option<DBVector>, Error>> {
308    values
309        .into_iter()
310        .zip(values_sizes.into_iter())
311        .zip(errors.into_iter())
312        .map(|((v, s), e)| {
313            if e.is_null() {
314                if v.is_null() {
315                    return Ok(None);
316                }
317                unsafe { Ok(Some(DBVector::from_c(v as *mut u8, s))) }
318            } else {
319                Err(Error::new(crate::ffi_util::error_message(e)))
320            }
321        })
322        .collect()
323}
ckb_rocksdb/ops/multi_get.rs

ckb_rocksdb/ops/
multi_get.rs
