mirror of
https://github.com/bitcoinresearchkit/brk.git
synced 2026-04-28 16:49:58 -07:00
query: init
This commit is contained in:
32
crates/brk_query/src/format.rs
Normal file
32
crates/brk_query/src/format.rs
Normal file
@@ -0,0 +1,32 @@
|
||||
use clap::ValueEnum;
|
||||
use color_eyre::eyre::eyre;
|
||||
use serde::Deserialize;
|
||||
|
||||
#[allow(clippy::upper_case_acronyms)]
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum, Deserialize)]
|
||||
pub enum Format {
|
||||
JSON,
|
||||
CSV,
|
||||
TSV,
|
||||
}
|
||||
|
||||
impl TryFrom<Option<String>> for Format {
|
||||
type Error = color_eyre::Report;
|
||||
fn try_from(value: Option<String>) -> Result<Self, Self::Error> {
|
||||
if let Some(value) = value {
|
||||
let value = value.to_lowercase();
|
||||
let value = value.as_str();
|
||||
if value == "csv" {
|
||||
Ok(Self::CSV)
|
||||
} else if value == "tsv" {
|
||||
Ok(Self::TSV)
|
||||
} else if value == "json" {
|
||||
Ok(Self::JSON)
|
||||
} else {
|
||||
Err(eyre!("Fail"))
|
||||
}
|
||||
} else {
|
||||
Err(eyre!("Fail"))
|
||||
}
|
||||
}
|
||||
}
|
||||
86
crates/brk_query/src/index.rs
Normal file
86
crates/brk_query/src/index.rs
Normal file
@@ -0,0 +1,86 @@
|
||||
use std::fmt::{self, Debug};
|
||||
|
||||
use color_eyre::eyre::eyre;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
|
||||
pub enum Index {
|
||||
Addressindex,
|
||||
Dateindex,
|
||||
Height,
|
||||
P2PK33index,
|
||||
P2PK65index,
|
||||
P2PKHindex,
|
||||
P2SHindex,
|
||||
P2TRindex,
|
||||
P2WPKHindex,
|
||||
P2WSHindex,
|
||||
Txindex,
|
||||
Txinindex,
|
||||
Txoutindex,
|
||||
}
|
||||
|
||||
impl Index {
|
||||
pub fn all() -> [Self; 13] {
|
||||
[
|
||||
Self::Addressindex,
|
||||
Self::Dateindex,
|
||||
Self::Height,
|
||||
Self::P2PK33index,
|
||||
Self::P2PK65index,
|
||||
Self::P2PKHindex,
|
||||
Self::P2SHindex,
|
||||
Self::P2TRindex,
|
||||
Self::P2WPKHindex,
|
||||
Self::P2WSHindex,
|
||||
Self::Txindex,
|
||||
Self::Txinindex,
|
||||
Self::Txoutindex,
|
||||
]
|
||||
}
|
||||
|
||||
pub fn ids(&self) -> &[&str] {
|
||||
match self {
|
||||
Self::Dateindex => &["d", "date", "dateindex"],
|
||||
Self::Height => &["h", "height"],
|
||||
Self::Txindex => &["txi", "txindex"],
|
||||
Self::Txinindex => &["txini", "txinindex"],
|
||||
Self::Txoutindex => &["txouti", "txoutindex"],
|
||||
Self::Addressindex => &["addri", "addressindex"],
|
||||
Self::P2PK33index => &["p2pk33i", "p2pk33index"],
|
||||
Self::P2PK65index => &["p2pk65i", "p2pk65index"],
|
||||
Self::P2PKHindex => &["p2pkhi", "p2pkhindex"],
|
||||
Self::P2SHindex => &["p2shi", "p2shindex"],
|
||||
Self::P2TRindex => &["p2tri", "p2trindex"],
|
||||
Self::P2WPKHindex => &["p2wpkhi", "p2wpkhindex"],
|
||||
Self::P2WSHindex => &["p2wshi", "p2wshindex"],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<&str> for Index {
|
||||
type Error = color_eyre::Report;
|
||||
fn try_from(value: &str) -> Result<Self, Self::Error> {
|
||||
Ok(match value {
|
||||
v if (Self::Dateindex).ids().contains(&v) => Self::Dateindex,
|
||||
v if (Self::Height).ids().contains(&v) => Self::Height,
|
||||
v if (Self::Txindex).ids().contains(&v) => Self::Txindex,
|
||||
v if (Self::Txinindex).ids().contains(&v) => Self::Txinindex,
|
||||
v if (Self::Txoutindex).ids().contains(&v) => Self::Txoutindex,
|
||||
v if (Self::Addressindex).ids().contains(&v) => Self::Addressindex,
|
||||
v if (Self::P2PK33index).ids().contains(&v) => Self::P2PK33index,
|
||||
v if (Self::P2PK65index).ids().contains(&v) => Self::P2PK65index,
|
||||
v if (Self::P2PKHindex).ids().contains(&v) => Self::P2PKHindex,
|
||||
v if (Self::P2SHindex).ids().contains(&v) => Self::P2SHindex,
|
||||
v if (Self::P2TRindex).ids().contains(&v) => Self::P2TRindex,
|
||||
v if (Self::P2WPKHindex).ids().contains(&v) => Self::P2WPKHindex,
|
||||
v if (Self::P2WSHindex).ids().contains(&v) => Self::P2WSHindex,
|
||||
_ => return Err(eyre!("Bad index")),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for Index {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
Debug::fmt(self, f)
|
||||
}
|
||||
}
|
||||
@@ -2,3 +2,159 @@
|
||||
#![doc = "\n## Example\n\n```rust"]
|
||||
#![doc = include_str!("main.rs")]
|
||||
#![doc = "```"]
|
||||
|
||||
mod format;
|
||||
mod index;
|
||||
mod params;
|
||||
mod tree;
|
||||
|
||||
use std::fmt;
|
||||
|
||||
use brk_computer::Computer;
|
||||
use brk_indexer::Indexer;
|
||||
pub use format::Format;
|
||||
pub use index::Index;
|
||||
pub use params::Params;
|
||||
use serde::Serialize;
|
||||
use tree::VecIdToIndexToVec;
|
||||
|
||||
pub struct Query<'a> {
|
||||
pub vecid_to_index_to_vec: VecIdToIndexToVec<'a>,
|
||||
indexer: &'a Indexer,
|
||||
computer: &'a Computer,
|
||||
}
|
||||
|
||||
impl<'a> Query<'a> {
|
||||
pub fn build(indexer: &'a Indexer, computer: &'a Computer) -> Self {
|
||||
let mut vecs = VecIdToIndexToVec::default();
|
||||
|
||||
indexer.vecs.as_any_vecs().into_iter().for_each(|vec| vecs.insert(vec));
|
||||
computer.vecs.as_any_vecs().into_iter().for_each(|vec| vecs.insert(vec));
|
||||
|
||||
Self {
|
||||
vecid_to_index_to_vec: vecs,
|
||||
indexer,
|
||||
computer,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn search(
|
||||
&self,
|
||||
index: Index,
|
||||
values: &[&str],
|
||||
from: Option<i64>,
|
||||
to: Option<i64>,
|
||||
format: Option<Format>,
|
||||
) -> color_eyre::Result<QueryResponse> {
|
||||
let ids = values
|
||||
.iter()
|
||||
.map(|s| {
|
||||
(
|
||||
s.to_owned(),
|
||||
self.vecid_to_index_to_vec.get(&s.to_lowercase().replace("_", "-")),
|
||||
)
|
||||
})
|
||||
.filter(|(_, opt)| opt.is_some())
|
||||
.map(|(id, vec)| (id, vec.unwrap()))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
if ids.is_empty() {
|
||||
return Ok(QueryResponse::default(format));
|
||||
}
|
||||
|
||||
let mut values = ids
|
||||
.iter()
|
||||
.flat_map(|(_, i_to_v)| i_to_v.get(&index))
|
||||
.map(|vec| -> brk_vec::Result<Vec<serde_json::Value>> { vec.collect_range_values(from, to) })
|
||||
.collect::<brk_vec::Result<Vec<_>>>()?;
|
||||
|
||||
if values.is_empty() {
|
||||
return Ok(QueryResponse::default(format));
|
||||
}
|
||||
|
||||
let ids_last_i = ids.len() - 1;
|
||||
|
||||
Ok(match format {
|
||||
Some(Format::CSV) | Some(Format::TSV) => {
|
||||
let delimiter = if format == Some(Format::CSV) { ',' } else { '\t' };
|
||||
|
||||
let mut text = ids
|
||||
.into_iter()
|
||||
.map(|(id, _)| id)
|
||||
.collect::<Vec<_>>()
|
||||
.join(&delimiter.to_string());
|
||||
|
||||
text.push('\n');
|
||||
|
||||
let values_len = values.first().unwrap().len();
|
||||
|
||||
(0..values_len).for_each(|i| {
|
||||
let mut line = "".to_string();
|
||||
values.iter().enumerate().for_each(|(id_i, v)| {
|
||||
line += &v.get(i).unwrap().to_string();
|
||||
if id_i == ids_last_i {
|
||||
line.push('\n');
|
||||
} else {
|
||||
line.push(delimiter);
|
||||
}
|
||||
});
|
||||
text += &line;
|
||||
});
|
||||
|
||||
if format == Some(Format::CSV) {
|
||||
QueryResponse::CSV(text)
|
||||
} else {
|
||||
QueryResponse::TSV(text)
|
||||
}
|
||||
}
|
||||
Some(Format::JSON) | None => {
|
||||
if values.len() == 1 {
|
||||
let mut values = values.pop().unwrap();
|
||||
if values.len() == 1 {
|
||||
let value = values.pop().unwrap();
|
||||
QueryResponse::Json(Value::Single(value))
|
||||
} else {
|
||||
QueryResponse::Json(Value::List(values))
|
||||
}
|
||||
} else {
|
||||
QueryResponse::Json(Value::Matrix(values))
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum QueryResponse {
|
||||
Json(Value),
|
||||
CSV(String),
|
||||
TSV(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum Value {
|
||||
Matrix(Vec<Vec<serde_json::Value>>),
|
||||
List(Vec<serde_json::Value>),
|
||||
Single(serde_json::Value),
|
||||
}
|
||||
|
||||
impl QueryResponse {
|
||||
fn default(format: Option<Format>) -> Self {
|
||||
match format {
|
||||
Some(Format::CSV) => QueryResponse::CSV("".to_string()),
|
||||
Some(Format::TSV) => QueryResponse::TSV("".to_string()),
|
||||
_ => QueryResponse::Json(Value::Single(serde_json::Value::Null)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for QueryResponse {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
Self::Json(value) => write!(f, "{}", serde_json::to_string_pretty(value).unwrap()),
|
||||
Self::CSV(string) => write!(f, "{}", string),
|
||||
Self::TSV(string) => write!(f, "{}", string),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,3 +1,23 @@
|
||||
fn main() {
|
||||
println!("Hello, world!");
|
||||
use std::path::Path;
|
||||
|
||||
use brk_computer::Computer;
|
||||
use brk_indexer::Indexer;
|
||||
use brk_query::{Index, Query};
|
||||
|
||||
pub fn main() -> color_eyre::Result<()> {
|
||||
color_eyre::install()?;
|
||||
|
||||
let outputs_dir = Path::new("../../_outputs");
|
||||
|
||||
let indexer = Indexer::import(&outputs_dir.join("indexed"))?;
|
||||
|
||||
let computer = Computer::import(&outputs_dir.join("computed"))?;
|
||||
|
||||
let query = Query::build(&indexer, &computer);
|
||||
|
||||
dbg!(query.search(Index::Height, &["date"], Some(-1), None, None)?);
|
||||
dbg!(query.search(Index::Height, &["date"], Some(-10), None, None)?);
|
||||
dbg!(query.search(Index::Height, &["date", "timestamp"], Some(-10), None, None)?);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
18
crates/brk_query/src/params.rs
Normal file
18
crates/brk_query/src/params.rs
Normal file
@@ -0,0 +1,18 @@
|
||||
use clap::Parser;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::Format;
|
||||
|
||||
#[derive(Debug, Deserialize, Parser)]
|
||||
pub struct Params {
|
||||
#[clap(short, long)]
|
||||
pub index: String,
|
||||
#[clap(short, long, value_delimiter = ' ', num_args = 1..)]
|
||||
pub values: Vec<String>,
|
||||
#[clap(short, long, allow_hyphen_values = true)]
|
||||
pub from: Option<i64>,
|
||||
#[clap(short, long, allow_hyphen_values = true)]
|
||||
pub to: Option<i64>,
|
||||
#[clap(long)]
|
||||
pub format: Option<Format>,
|
||||
}
|
||||
38
crates/brk_query/src/tree.rs
Normal file
38
crates/brk_query/src/tree.rs
Normal file
@@ -0,0 +1,38 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use brk_vec::AnyStorableVec;
|
||||
use derive_deref::{Deref, DerefMut};
|
||||
|
||||
use super::index::Index;
|
||||
|
||||
#[derive(Default, Deref, DerefMut)]
|
||||
pub struct VecIdToIndexToVec<'a>(BTreeMap<String, IndexToVec<'a>>);
|
||||
|
||||
impl<'a> VecIdToIndexToVec<'a> {
|
||||
// Not the most performant or type safe but only built once so that's okay
|
||||
pub fn insert(&mut self, vec: &'a dyn AnyStorableVec) {
|
||||
let file_name = vec.file_name();
|
||||
let split = file_name.split("_to_").collect::<Vec<_>>();
|
||||
if split.len() != 2 {
|
||||
panic!();
|
||||
}
|
||||
let str = vec.index_type_to_string().split("::").last().unwrap().to_lowercase();
|
||||
let index = Index::try_from(str.as_str())
|
||||
.inspect_err(|_| {
|
||||
dbg!(str);
|
||||
})
|
||||
.unwrap();
|
||||
if split[0] != index.to_string().to_lowercase() {
|
||||
dbg!(split[0], index.to_string());
|
||||
panic!();
|
||||
}
|
||||
let key = split[1].to_string().replace("_", "-");
|
||||
let prev = self.entry(key).or_default().insert(index, vec);
|
||||
if prev.is_some() {
|
||||
panic!()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default, Deref, DerefMut)]
|
||||
pub struct IndexToVec<'a>(BTreeMap<Index, &'a dyn AnyStorableVec>);
|
||||
Reference in New Issue
Block a user