From 33ea956c7b36c3a86c03879d639106cd2e1b7438 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= Date: Mon, 24 Sep 2018 16:53:33 +0200 Subject: [PATCH] feat: Add a way to index from a csv file --- Cargo.lock | 44 +++++++++++ Cargo.toml | 1 + raptor-http/Cargo.toml | 2 - raptor-http/src/main.rs | 6 ++ raptor-indexer-csv/.gitignore | 3 + raptor-indexer-csv/Cargo.toml | 17 +++++ raptor-indexer-csv/src/main.rs | 136 +++++++++++++++++++++++++++++++++ raptor-indexer/Cargo.toml | 2 - raptor-search/Cargo.toml | 2 - raptor/Cargo.toml | 2 - 10 files changed, 207 insertions(+), 8 deletions(-) create mode 100644 raptor-indexer-csv/.gitignore create mode 100644 raptor-indexer-csv/Cargo.toml create mode 100644 raptor-indexer-csv/src/main.rs diff --git a/Cargo.lock b/Cargo.lock index 73adf77ce..04f194895 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -177,6 +177,23 @@ name = "crossbeam-utils" version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "csv" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "csv-core 0.1.4 (registry+https://github.com/rust-lang/crates.io-index)", + "serde 1.0.75 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] +name = "csv-core" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "memchr 2.1.0 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "digest" version = "0.7.5" @@ -452,6 +469,16 @@ name = "matches" version = "0.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "memchr" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "cfg-if 0.1.5 (registry+https://github.com/rust-lang/crates.io-index)", + "libc 0.2.43 (registry+https://github.com/rust-lang/crates.io-index)", + "version_check 0.1.4 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "memmap" version = "0.6.2" @@ -723,6 +750,20 @@ dependencies = [ "unidecode 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "raptor-indexer-csv" +version = "0.1.0" +dependencies = [ + "csv 1.0.1 (registry+https://github.com/rust-lang/crates.io-index)", + "moby-name-gen 0.1.0 (registry+https://github.com/rust-lang/crates.io-index)", + "raptor 0.1.0", + "rocksdb 0.3.0 (git+https://github.com/pingcap/rust-rocksdb.git)", + "serde 1.0.75 (registry+https://github.com/rust-lang/crates.io-index)", + "serde_derive 1.0.75 (registry+https://github.com/rust-lang/crates.io-index)", + "structopt 0.2.10 (registry+https://github.com/rust-lang/crates.io-index)", + "unidecode 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "raptor-search" version = "0.1.0" @@ -1327,6 +1368,8 @@ dependencies = [ "checksum crossbeam-deque 0.6.1 (registry+https://github.com/rust-lang/crates.io-index)" = "3486aefc4c0487b9cb52372c97df0a48b8c249514af1ee99703bf70d2f2ceda1" "checksum crossbeam-epoch 0.5.2 (registry+https://github.com/rust-lang/crates.io-index)" = "30fecfcac6abfef8771151f8be4abc9e4edc112c2bcb233314cafde2680536e9" "checksum crossbeam-utils 0.5.0 (registry+https://github.com/rust-lang/crates.io-index)" = "677d453a17e8bd2b913fa38e8b9cf04bcdbb5be790aa294f2389661d72036015" +"checksum csv 1.0.1 (registry+https://github.com/rust-lang/crates.io-index)" = "0bbb6d75aae072248e381715437855a69595e5a97d3abbf748fe7a95e65d77fa" +"checksum csv-core 0.1.4 (registry+https://github.com/rust-lang/crates.io-index)" = "4dd8e6d86f7ba48b4276ef1317edc8cc36167546d8972feb4a2b5fec0b374105" "checksum digest 0.7.5 (registry+https://github.com/rust-lang/crates.io-index)" = "5b29c278aa8fd30796bd977169e8004b4aa88cdcd2f32a6eb22bc2d5d38df94a" "checksum dtoa 0.4.3 (registry+https://github.com/rust-lang/crates.io-index)" = "6d301140eb411af13d3115f9a562c85cc6b541ade9dfa314132244aaee7489dd" "checksum elapsed 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)" = "6f4e5af126dafd0741c2ad62d47f68b28602550102e5f0dd45c8a97fc8b49c29" @@ -1361,6 +1404,7 @@ dependencies = [ "checksum log 0.4.5 (registry+https://github.com/rust-lang/crates.io-index)" = "d4fcce5fa49cc693c312001daf1d13411c4a5283796bac1084299ea3e567113f" "checksum lz4-sys 1.8.0 (git+https://github.com/busyjay/lz4-rs.git?branch=adjust-build)" = "" "checksum matches 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)" = "7ffc5c5338469d4d3ea17d269fa8ea3512ad247247c30bd2df69e68309ed0a08" +"checksum memchr 2.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "4b3629fe9fdbff6daa6c33b90f7c08355c1aca05a3d01fa8063b822fcf185f3b" "checksum memmap 0.6.2 (registry+https://github.com/rust-lang/crates.io-index)" = "e2ffa2c986de11a9df78620c01eeaaf27d94d3ff02bf81bfcca953102dd0c6ff" "checksum memoffset 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "0f9dc261e2b62d7a622bf416ea3c5245cdd5d9a7fcc428c0d06804dfce1775b3" "checksum mime 0.3.9 (registry+https://github.com/rust-lang/crates.io-index)" = "4b082692d3f6cf41b453af73839ce3dfc212c4411cbb2441dff80a716e38bd79" diff --git a/Cargo.toml b/Cargo.toml index fcfd400e4..6c69591b9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,6 +2,7 @@ members = [ "raptor", "raptor-indexer", + "raptor-indexer-csv", "raptor-search", "raptor-http", ] diff --git a/raptor-http/Cargo.toml b/raptor-http/Cargo.toml index b78f6bdb8..017731ede 100644 --- a/raptor-http/Cargo.toml +++ b/raptor-http/Cargo.toml @@ -1,5 +1,3 @@ -cargo-features = ["edition"] - [package] edition = "2018" name = "raptor-http" diff --git a/raptor-http/src/main.rs b/raptor-http/src/main.rs index 323219b65..9a71cce24 100644 --- a/raptor-http/src/main.rs +++ b/raptor-http/src/main.rs @@ -37,6 +37,7 @@ struct Document<'a> { id: u64, title: &'a str, description: &'a str, + image: &'a str, } type CommonWords = HashSet; @@ -79,10 +80,15 @@ where M: AsRef, let description = database.as_ref().get(description_key.as_bytes()).unwrap().unwrap(); let description = unsafe { from_utf8_unchecked(&description) }; + let image_key = format!("{}-image", document.document_id); + let image = database.as_ref().get(image_key.as_bytes()).unwrap().unwrap(); + let image = unsafe { from_utf8_unchecked(&image) }; + let document = Document { id: document.document_id, title: title, description: description, + image: image, }; if !first { write!(&mut body, ",")? } diff --git a/raptor-indexer-csv/.gitignore b/raptor-indexer-csv/.gitignore new file mode 100644 index 000000000..70e3cae73 --- /dev/null +++ b/raptor-indexer-csv/.gitignore @@ -0,0 +1,3 @@ + +/target +**/*.rs.bk diff --git a/raptor-indexer-csv/Cargo.toml b/raptor-indexer-csv/Cargo.toml new file mode 100644 index 000000000..34a56316f --- /dev/null +++ b/raptor-indexer-csv/Cargo.toml @@ -0,0 +1,17 @@ +[package] +edition = "2018" +name = "raptor-indexer-csv" +version = "0.1.0" +authors = ["Kerollmops "] + +[dependencies] +raptor = { path = "../raptor" } +serde = "1.0" +structopt = "0.2" +serde_derive = "1.0" +csv = "1.0" +unidecode = "0.3" +moby-name-gen = "0.1" + +[dependencies.rocksdb] +git = "https://github.com/pingcap/rust-rocksdb.git" diff --git a/raptor-indexer-csv/src/main.rs b/raptor-indexer-csv/src/main.rs new file mode 100644 index 000000000..b26dfcd9f --- /dev/null +++ b/raptor-indexer-csv/src/main.rs @@ -0,0 +1,136 @@ +// TODO make the raptor binary expose multiple subcommand +// make only one binary + +#[macro_use] extern crate serde_derive; + +use std::path::{Path, PathBuf}; +use std::collections::{HashSet, BTreeMap}; +use std::io::{self, BufReader, BufRead}; +use std::fs::File; +use std::iter; + +use csv::ReaderBuilder; +use structopt::StructOpt; +use raptor::{MetadataBuilder, DocIndex}; +use rocksdb::{SstFileWriter, EnvOptions, ColumnFamilyOptions}; +use unidecode::unidecode; + +#[derive(Debug, StructOpt)] +#[structopt(name = "raptor-indexer-csv", about = "A Raptor binary to index csv stored products.")] +struct Opt { + /// The stop word file, each word must be separated by a newline. + #[structopt(long = "stop-words", parse(from_os_str))] + stop_words: PathBuf, + + /// The csv file to index. + #[structopt(parse(from_os_str))] + products: PathBuf, +} + +#[derive(Debug, Deserialize)] +struct Product { + #[serde(rename = "_unit_id")] + id: u64, + #[serde(rename = "product_title")] + title: String, + #[serde(rename = "product_image")] + image: String, + #[serde(rename = "product_description")] + description: String, +} + +type CommonWords = HashSet; + +fn common_words

(path: P) -> io::Result +where P: AsRef, +{ + let file = File::open(path)?; + let file = BufReader::new(file); + let mut set = HashSet::new(); + for line in file.lines().filter_map(|l| l.ok()) { + for word in line.split_whitespace() { + set.insert(word.to_owned()); + } + } + Ok(set) +} + +fn main() { + let opt = Opt::from_args(); + + let common_words = common_words(opt.stop_words).expect("reading stop words"); + + // TODO add a subcommand to pack these files in a tar.xxx archive + let random_name = moby_name_gen::random_name(); + let map_file = format!("{}.map", random_name); + let idx_file = format!("{}.idx", random_name); + let sst_file = format!("{}.sst", random_name); + + let env_options = EnvOptions::new(); + let cf_options = ColumnFamilyOptions::new(); + let mut sst_file_writer = SstFileWriter::new(env_options, cf_options); + sst_file_writer.open(&sst_file).expect("open the sst file"); + + let map = File::create(&map_file).unwrap(); + let indexes = File::create(&idx_file).unwrap(); + let mut builder = MetadataBuilder::new(map, indexes); + let mut fields = BTreeMap::new(); + + let mut rdr = ReaderBuilder::new().from_path(opt.products).expect("reading product file"); + let mut errors = 0; + + for result in rdr.deserialize() { + let product: Product = match result { + Ok(product) => product, + Err(e) => { eprintln!("{:?}", e); errors += 1; continue }, + }; + + { + let title = iter::repeat(0).zip(product.title.split_whitespace()).filter(|&(_, w)| !common_words.contains(w)).enumerate(); + let description = iter::repeat(1).zip(product.description.split_whitespace()).filter(|&(_, w)| !common_words.contains(w)).enumerate(); + + let words = title.chain(description); + for (i, (attr, word)) in words { + let doc_index = DocIndex { + document: product.id, + attribute: attr, + attribute_index: i as u32, + }; + // insert the exact representation + let word_lower = word.to_lowercase(); + + // and the unidecoded lowercased version + let word_unidecoded = unidecode(word).to_lowercase(); + if word_lower != word_unidecoded { + builder.insert(word_unidecoded, doc_index); + } + + builder.insert(word_lower, doc_index); + } + } + + // TODO simplify this by using functions and + // use the MetadataBuilder internal BTreeMap ? + let key = format!("{}-title", product.id); + let value = product.title; + fields.insert(key, value); + + let key = format!("{}-description", product.id); + let value = product.description; + fields.insert(key, value); + + let key = format!("{}-image", product.id); + let value = product.image; + fields.insert(key, value); + } + + for (key, value) in fields { + sst_file_writer.put(key.as_bytes(), value.as_bytes()).unwrap(); + } + let _sst_file_info = sst_file_writer.finish().unwrap(); + + builder.finish().unwrap(); + + println!("Found {} errorneous lines", errors); + println!("Succesfully created {:?} dump.", random_name); +} diff --git a/raptor-indexer/Cargo.toml b/raptor-indexer/Cargo.toml index bb8da0325..eed5086b6 100644 --- a/raptor-indexer/Cargo.toml +++ b/raptor-indexer/Cargo.toml @@ -1,5 +1,3 @@ -cargo-features = ["edition"] - [package] edition = "2018" name = "raptor-indexer" diff --git a/raptor-search/Cargo.toml b/raptor-search/Cargo.toml index fb16a0ee2..b18105b72 100644 --- a/raptor-search/Cargo.toml +++ b/raptor-search/Cargo.toml @@ -1,5 +1,3 @@ -cargo-features = ["edition"] - [package] edition = "2018" name = "raptor-search" diff --git a/raptor/Cargo.toml b/raptor/Cargo.toml index 74c2b4551..1e4ed8549 100644 --- a/raptor/Cargo.toml +++ b/raptor/Cargo.toml @@ -1,5 +1,3 @@ -cargo-features = ["edition"] - [package] edition = "2018" name = "raptor"