Simplifies RocksDB SST API

pull/171/head
Tpt 3 years ago
parent 367a1b4585
commit 9f414c13fd
  1. 49
      lib/src/storage/backend/rocksdb.rs
  2. 29
      lib/src/storage/mod.rs

@ -138,11 +138,7 @@ impl Db {
in_memory: bool, in_memory: bool,
for_bulk_load: bool, for_bulk_load: bool,
) -> Result<DbHandler> { ) -> Result<DbHandler> {
let c_path = CString::new( let c_path = path_to_cstring(path)?;
path.to_str()
.ok_or_else(|| invalid_input_error("The DB path is not valid UTF-8"))?,
)
.map_err(invalid_input_error)?;
unsafe { unsafe {
let options = rocksdb_options_create(); let options = rocksdb_options_create();
@ -500,31 +496,23 @@ impl Db {
pub fn new_sst_file(&self) -> Result<SstFileWriter> { pub fn new_sst_file(&self) -> Result<SstFileWriter> {
unsafe { unsafe {
let path = self.0.path.join(random::<u128>().to_string());
let writer = rocksdb_sstfilewriter_create(self.0.env_options, self.0.options); let writer = rocksdb_sstfilewriter_create(self.0.env_options, self.0.options);
let cpath = CString::new( ffi_result!(rocksdb_sstfilewriter_open(
self.0 writer,
.path path_to_cstring(&path)?.as_ptr()
.join(random::<u128>().to_string()) ))?;
.to_string_lossy() Ok(SstFileWriter { writer, path })
.to_string(),
)
.map_err(invalid_input_error)?;
ffi_result!(rocksdb_sstfilewriter_open(writer, cpath.as_ptr()))?;
Ok(SstFileWriter { writer, cpath })
} }
} }
pub fn write_stt_files( pub fn write_stt_files(&self, ssts_for_cf: Vec<(&ColumnFamily, PathBuf)>) -> Result<()> {
&self, for (cf, path) in ssts_for_cf {
writers_with_cf: Vec<(&ColumnFamily, SstFileWriter)>,
) -> Result<()> {
for (cf, writer) in writers_with_cf {
unsafe { unsafe {
ffi_result!(rocksdb_sstfilewriter_finish(writer.writer))?;
ffi_result!(rocksdb_transactiondb_ingest_external_file_cf( ffi_result!(rocksdb_transactiondb_ingest_external_file_cf(
self.0.db, self.0.db,
cf.0, cf.0,
&writer.cpath.as_ptr(), &path_to_cstring(&path)?.as_ptr(),
1, 1,
self.0.ingest_external_file_options self.0.ingest_external_file_options
))?; ))?;
@ -764,7 +752,7 @@ impl Iter {
pub struct SstFileWriter { pub struct SstFileWriter {
writer: *mut rocksdb_sstfilewriter_t, writer: *mut rocksdb_sstfilewriter_t,
cpath: CString, path: PathBuf,
} }
impl Drop for SstFileWriter { impl Drop for SstFileWriter {
@ -803,6 +791,13 @@ impl SstFileWriter {
)) ))
} }
} }
pub fn finish(self) -> Result<PathBuf> {
unsafe {
ffi_result!(rocksdb_sstfilewriter_finish(self.writer))?;
}
Ok(self.path.clone())
}
} }
fn convert_error(ptr: *const c_char) -> Error { fn convert_error(ptr: *const c_char) -> Error {
@ -950,3 +945,11 @@ unsafe extern "C" fn compactionfilter_name(filter: *mut c_void) -> *const c_char
let filter = &*(filter as *const CompactionFilter); let filter = &*(filter as *const CompactionFilter);
filter.name.as_ptr() filter.name.as_ptr()
} }
fn path_to_cstring(path: &Path) -> Result<CString> {
CString::new(
path.to_str()
.ok_or_else(|| invalid_input_error("The DB path is not valid UTF-8"))?,
)
.map_err(invalid_input_error)
}

@ -9,8 +9,6 @@ use crate::storage::binary_encoder::{
use crate::storage::numeric_encoder::{ use crate::storage::numeric_encoder::{
insert_term, remove_term, EncodedQuad, EncodedTerm, StrHash, StrLookup, insert_term, remove_term, EncodedQuad, EncodedTerm, StrHash, StrLookup,
}; };
#[cfg(not(target_arch = "wasm32"))]
use backend::SstFileWriter;
use backend::{ use backend::{
ColumnFamily, ColumnFamilyDefinition, CompactionAction, CompactionFilter, Db, Iter, ColumnFamily, ColumnFamilyDefinition, CompactionAction, CompactionFilter, Db, Iter,
MergeOperator, WriteBatchWithIndex, MergeOperator, WriteBatchWithIndex,
@ -23,6 +21,7 @@ use std::io::Result;
use std::mem::take; use std::mem::take;
#[cfg(not(target_arch = "wasm32"))] #[cfg(not(target_arch = "wasm32"))]
use std::path::Path; use std::path::Path;
use std::path::PathBuf;
mod backend; mod backend;
mod binary_encoder; mod binary_encoder;
@ -1114,13 +1113,13 @@ impl BulkLoader {
id2str_sst.merge(&k, &buffer)?; id2str_sst.merge(&k, &buffer)?;
buffer.clear(); buffer.clear();
} }
to_load.push((&self.storage.id2str_cf, id2str_sst)); to_load.push((&self.storage.id2str_cf, id2str_sst.finish()?));
} }
if !self.triples.is_empty() { if !self.triples.is_empty() {
to_load.push(( to_load.push((
&self.storage.dspo_cf, &self.storage.dspo_cf,
self.add_keys_to_sst( self.build_sst_for_keys(
self.triples.iter().map(|quad| { self.triples.iter().map(|quad| {
encode_term_triple(&quad.subject, &quad.predicate, &quad.object) encode_term_triple(&quad.subject, &quad.predicate, &quad.object)
}), }),
@ -1128,7 +1127,7 @@ impl BulkLoader {
)); ));
to_load.push(( to_load.push((
&self.storage.dpos_cf, &self.storage.dpos_cf,
self.add_keys_to_sst( self.build_sst_for_keys(
self.triples.iter().map(|quad| { self.triples.iter().map(|quad| {
encode_term_triple(&quad.predicate, &quad.object, &quad.subject) encode_term_triple(&quad.predicate, &quad.object, &quad.subject)
}), }),
@ -1136,7 +1135,7 @@ impl BulkLoader {
)); ));
to_load.push(( to_load.push((
&self.storage.dosp_cf, &self.storage.dosp_cf,
self.add_keys_to_sst( self.build_sst_for_keys(
self.triples.iter().map(|quad| { self.triples.iter().map(|quad| {
encode_term_triple(&quad.object, &quad.subject, &quad.predicate) encode_term_triple(&quad.object, &quad.subject, &quad.predicate)
}), }),
@ -1149,12 +1148,12 @@ impl BulkLoader {
let quads = take(&mut self.graphs); let quads = take(&mut self.graphs);
to_load.push(( to_load.push((
&self.storage.graphs_cf, &self.storage.graphs_cf,
self.add_keys_to_sst(quads.into_iter().map(|g| encode_term(&g)))?, self.build_sst_for_keys(quads.into_iter().map(|g| encode_term(&g)))?,
)); ));
to_load.push(( to_load.push((
&self.storage.gspo_cf, &self.storage.gspo_cf,
self.add_keys_to_sst(self.quads.iter().map(|quad| { self.build_sst_for_keys(self.quads.iter().map(|quad| {
encode_term_quad( encode_term_quad(
&quad.graph_name, &quad.graph_name,
&quad.subject, &quad.subject,
@ -1165,7 +1164,7 @@ impl BulkLoader {
)); ));
to_load.push(( to_load.push((
&self.storage.gpos_cf, &self.storage.gpos_cf,
self.add_keys_to_sst(self.quads.iter().map(|quad| { self.build_sst_for_keys(self.quads.iter().map(|quad| {
encode_term_quad( encode_term_quad(
&quad.graph_name, &quad.graph_name,
&quad.object, &quad.object,
@ -1176,7 +1175,7 @@ impl BulkLoader {
)); ));
to_load.push(( to_load.push((
&self.storage.gosp_cf, &self.storage.gosp_cf,
self.add_keys_to_sst(self.quads.iter().map(|quad| { self.build_sst_for_keys(self.quads.iter().map(|quad| {
encode_term_quad( encode_term_quad(
&quad.graph_name, &quad.graph_name,
&quad.object, &quad.object,
@ -1187,7 +1186,7 @@ impl BulkLoader {
)); ));
to_load.push(( to_load.push((
&self.storage.spog_cf, &self.storage.spog_cf,
self.add_keys_to_sst(self.quads.iter().map(|quad| { self.build_sst_for_keys(self.quads.iter().map(|quad| {
encode_term_quad( encode_term_quad(
&quad.subject, &quad.subject,
&quad.predicate, &quad.predicate,
@ -1198,7 +1197,7 @@ impl BulkLoader {
)); ));
to_load.push(( to_load.push((
&self.storage.posg_cf, &self.storage.posg_cf,
self.add_keys_to_sst(self.quads.iter().map(|quad| { self.build_sst_for_keys(self.quads.iter().map(|quad| {
encode_term_quad( encode_term_quad(
&quad.object, &quad.object,
&quad.subject, &quad.subject,
@ -1209,7 +1208,7 @@ impl BulkLoader {
)); ));
to_load.push(( to_load.push((
&self.storage.ospg_cf, &self.storage.ospg_cf,
self.add_keys_to_sst(self.quads.iter().map(|quad| { self.build_sst_for_keys(self.quads.iter().map(|quad| {
encode_term_quad( encode_term_quad(
&quad.object, &quad.object,
&quad.subject, &quad.subject,
@ -1240,14 +1239,14 @@ impl BulkLoader {
) )
} }
fn add_keys_to_sst(&self, values: impl Iterator<Item = Vec<u8>>) -> Result<SstFileWriter> { fn build_sst_for_keys(&self, values: impl Iterator<Item = Vec<u8>>) -> Result<PathBuf> {
let mut values = values.collect::<Vec<_>>(); let mut values = values.collect::<Vec<_>>();
values.sort_unstable(); values.sort_unstable();
let mut sst = self.storage.db.new_sst_file()?; let mut sst = self.storage.db.new_sst_file()?;
for t in values { for t in values {
sst.insert_empty(&t)?; sst.insert_empty(&t)?;
} }
Ok(sst) sst.finish()
} }
} }

Loading…
Cancel
Save