O meu livro técnico favorito é o Designing Data-Intensive Applications do Kleppmann. Apesar de ser relativamente antigo, como ele trata mais do aspecto teórico que prático, continua sendo uma ótima leitura para engenheiros de dados e profissionais que precisam lidar com sistemas de larga escala.

Em uma entrevista recente com o autor – quando perguntado sobre dicas para um aspirante a engenheiros de dados – concordei plenamente com esse conselho:

Learn just enough abot the internals of the tools you are using, so that you have a reasonable mental model of what’s going on there. You don’t need to be able to, like, modify of the Kafka yourself, but I think having just enough of an idea of the internals that […] if the the performance goes bad, you have a way of visualizing in your head what’s going on. […]. It’s incredibly valuable to just have a bit of a mental model and not just treat it as a black box.

Uma das primeiras vezes, que entendi o valor de construir esses modelo mentais, foi quando aprendi sobre índices de banco de dados. Passei a entender melhor o plano de execução, quando a criação de um índice fazia sentido e otimizei consultas que outras pessoas não estavam conseguindo.

A solução mais comum para índices são as árvores B, esse tipo de estrutura é utilizada por diversos tipos de banco de dados: desde o SQLite, passando pelo Oracle Database e até mesmo em soluções distribuídas como o Dynamo DB.

Sendo uma estrutura tão prevalente em soluções de dados, acho importante que todo engenheiro de dados tenha uma noção de como ela funciona, construir esse modelo mental que Kleppmann comentou. Como eu queria brincar um pouco com a linguagem Rust, achei que era um bom exercício implementar uma árvore B na linguagem e escrever um pouco sobre.

O que é uma árvore B?

As árvores B são uma generalização das árvores binárias: ao invés de ter uma chave por nó, são k chaves por nós. Ela foi criada na década de 70 para lidar com dados persistentes, em um época em que a memória era escassa e a forma mais comum de armazenamento durável era o disco rígido. Atualmente, temos fartura de memória e armazenamento em SSD, novas estruturas como as LSM Trees foram desenvolvidas pensando nessa nova realidade.

Em soluções focadas em performance e escalabilidade, como o Cassandra por exemplo, faz muito sentido usar estruturas como as LSM Tree. Entretanto, algo que eu aprendi após anos trabalhando com “Big Data”: soluções com uso intensivo de memória têm performance excepcional, mas podem ficar inutilizáveis em cenários de escassez da mesma.

Apesar de ser uma estrutura com mais de 50 anos, desenvolvida em um cenário diferente do atual, segue sendo popular em novas soluções de dados por vários motivos:

  • as ávores B não geram pressão em memória;
  • são flexíveis como uma árvore binária (e.g. possibilidade de chaves parcias; suporte a múltiplos operadores de busca(>, < e =);
  • dado é armazenado de forma ordenada;
  • desempenho suficiente para muitos casos de uso.

Não faz muito sentido eu fazer (mais) uma explicação detalhada de como elas funcionam, porque existem infinitos materias sobre o assunto nos mais diversos formatos: aulas online, vídeos do Akita, blog posts e livros de algoritmos. É mais importante que o leitor procure esses materiais para entender sobre a estrutura, do que se preocupar em ler o restante desse post: eu fiz para meu próprio entretenimento, não como algo útil ou didático necessariamente.

Para fazer a implementação em Rust, revi o assunto no famoso “Cormen” e fiz praticamente uma cópia do proposto no livro. Implementei apenas a parte de inserção e busca, usando a estrutura para indexar arquivos no formato csv.

A estrutura da estrutura

Ao implementar uma estrutura de dados, normalmente começo imaginando as sub-estruturas que vou precisar e só depois penso na manipulação. Para essa implementação da árvore B, utilizei três sub-estruturas: BTree, Node e Key.

BTree

A BTree é a estrutura pública, que será utilizada pelas outras aplicações, contendo três campos:

  • root é o nó raíz por onde serão iniciadas as buscas;
  • order indica a quantidade de nós máximos que uma árvore pode conter;
  • path é o diretório em que os arquivos da árvore serão armazenados.
pub struct BTree {
    root: Node,
    order: usize,
    path: String,
}

Node

O Node é a estrutura mais importante, que armazena e organiza as chaves de forma que as buscas possam ser feita em log(N). É composta pelos seguintes campos:

  • key é um vetor contendo as chaves ordenadas;
  • children é um vetor contendo os filhos de cada chave;
  • leaf indica se o nó é uma folha, informação importante para tratar casos especiais de inclusão;
  • filename é o nome do arquivo em que esse nó é armazenado.
pub struct Node {
    pub keys: Vec<Key>,
    pub children: Vec<String>,
    pub leaf: bool,
    pub filename: String,
}

Key

A estrutura Key é a mais simples, contendo apenas dois campos:

  • value é o valor da chave propriamente dito, optei por deixar como String por simplicidade de implementação. Em uma solução produtiva, seria interessante ter algo mais flexível como um vetor de bytes ou um campo com tipo genérico.

  • position é uma tupla com a posição da chave dentro do arquivo csv, o primeiro campo é o offset no arquivo e a segunda é a quantidade de bytes para ser lido a partir do offset.

pub struct Key {
    pub value: String,
    pub position: (u64, u64),
}

Criando a árvore

A busca em árvore costuma ser algo simples, a complexidade é maior para criação e manutenção da mesma. Optei por fazer apenas a inserção, que é o mínimo necessário para criar uma árvore funcional.

A inclusão na árvore é feito pela função insert, associada ao objeto Btree. A lógica de inclusão é feita pelo objeto Node, mas antes existe um desvio para o cenário em que o nó raiz está cheio.

pub fn is_full(&self, order: usize) -> bool {
    self.keys.len() == 2 * order - 1
}

pub fn insert(&mut self, key: Key) {
    if self.root.is_full(self.order) {
        let mut new_root = Node::empty(self.order, false, &self.path);
        new_root.children.push(self.root.clone().filename);
        new_root.split(0, self.order, &self.path);
        self.root = new_root;
    }

    self.root.insert(key, self.order, &self.path);
    self.save();
}

Se o nó da raiz estiver cheio, um nó vazio é criado e a raiz vira filho desse novo nó. Após tratar esse caso específico, a inserção do registro é feito pela função insert na estrutura Node:

fn find_position(&self, key: &Key) -> usize {
    let mut idx = 0;

    for (i, iter_key) in self.keys.iter().enumerate() {
        idx = i;
        if iter_key.value > key.value {
            break;
        }
    }

    if idx + 1 == self.keys.len() {
        if key.value > self.keys[idx].value {
            idx += 1;
        }
    }

    idx
}

fn add_key(&mut self, idx: usize, key: Key) {
    if self.keys.len() == 0 {
        self.keys.push(key);
    } else {
        self.keys.insert(idx, key);
    }
    self.save();
}

pub fn insert(&mut self, key: Key, order: usize, path: &str) {
    if self.leaf {
        self.add_key(self.find_position(&key), key.clone());
    } else {
        let mut idx = self.find_position(&key);

        if Node::load(&self.children[idx]).is_full(order) {
            self.split(idx, order, path);
            idx = self.find_position(&key);
        }

        Node::load(&self.children[idx]).insert(key, order, path);
    }
}

A função add_key assume que o nó tem espaço e não precisa ser separado em múltiplos. Se for um nó folha, essa função é chamada diretamente na posição definida por find_position. Caso contrário, o nó filho em que a chave será inserida é procurada pela mesma função find_position, esse nó filho é separado em dois caso esteja cheio.

Para separar um nó, a chave central do nó é usada como referência. No caso abaixo, a chave S vai para o nó superior, os filhos à esquerda permanecem no nó original e os filhos à direita são movidos para um novo nó.

... N W ...... N S W ...P Q R S T U VP Q RT U V

A função split faz esse processo de rotação das chaves, que é realizado antes da inclusão em um nó cheio:

pub fn split(&mut self, pivot: usize, order: usize, path: &str) {
    let left = &mut Node::load(&self.children[pivot]);
    let key = left.keys[order - 1].clone();

    let right = Node {
        keys: left.keys[order..left.keys.len()].to_owned(),
        children: match left.leaf {
            true => Vec::with_capacity(2 * order),
            false => left.children[order..left.children.len()].to_owned(),
        },
        leaf: left.leaf,
        filename: Node::filename(path),
    };

    right.save();

    left.keys.resize(order - 1, Key::create("", (0, 0)));

    if !left.leaf {
        left.children.resize(order, String::from(""));
    }

    left.save();

    self.keys.insert(pivot, key);
    self.children.insert(pivot + 1, right.filename);

    self.save()
}

Em várias momentos, as estruturas são persistidas e lidas do armazenamento, usando as funções savee load respectivamente. Essas funções serializam e desserializam o objeto Node no formato json, usando a biblioteca serde_json.

pub fn save(&self) {
    let mut file = File::create(&self.filename).unwrap();
    file.write_all(serde_json::to_string(self).unwrap().as_bytes())
        .unwrap();
}

pub fn load(filename: &str) -> Node {
    serde_json::from_slice(&fs::read(filename).unwrap()).unwrap()
}

Cada nó é salvo em um arquivo com nome aleatório, gerado pela função filename:

fn filename(path: &str) -> String {
    format!("{}/{}.json", path, Uuid::new_v4().to_string())
}

O formato json não é o ideal em termos de performance, mas é legível por humanos e prático para ser utilizado em um código para estudos. Abaixo, um exemplo de nó salvo em disco, gerado pelos testes unitários:

{
    "keys": [
        {
            "value": "9b35db6e-49d9-4d91-b5f6-973b8a9111f4",
            "position": [
                0,
                0
            ]
        },
        {
            "value": "d47fbd93-cb1a-4e01-82b2-4325aa8cb1a3",
            "position": [
                0,
                0
            ]
        }
    ],
    "children": [
        "btree_test_search/d1740cf9-8fa9-471d-9ba6-99ff1462f5c6.json",
        "btree_test_search/9ba9bd01-6b3e-4441-a70b-253c86424cd3.json",
        "btree_test_search/7248b1e6-88f6-4f8e-a07b-0952347adc47.json"
    ],
    "leaf": false,
    "filename": "btree_test_search/1c3df73f-f321-4b3a-a545-9917ad2ecb4e.json"
}

Buscando na árvore

A busca é feito de forma recursiva, como é comum em estruturas de árvore. A função pública search da Btree tem apenas o parâmetro value com o valor da chave a ser procurado, a função search_tree é a que de fato realiza a busca recursiva.

fn search_tree(node: &Node, value: String) -> Option<Key> {
    for (i, key) in node.keys.iter().enumerate() {
        if key.value == value {
            return Some(key.clone());
        } else if key.value > value {
            if node.leaf {
                return None;
            } else {
                return BTree::search_tree(&Node::load(&node.children[i]), value);
            }
        }
    }

    if node.leaf {
        None
    } else {
        BTree::search_tree(&Node::load(&node.children[node.keys.len()]), value)
    }
}

pub fn search(&self, value: &str) -> Option<Key> {
    BTree::search_tree(&self.root, value.to_string())
}

Com a busca e a inserção criadas, já temos o mínimo para fazer o indexador.

Indexando os CSVs

Os arquivos CSVs são muito utilizados para lidar com dados tabulares. A organização física desses arquivos é parecida com a de uma tabela em um banco relacional, sendo orientado a linhas e com marcadores utilizados para indetificar início e fim dos campos e registros.

A função index_file recebe um arquivo e uma árvore, itera linha-a-linha no arquivo e armazena três informações na árvore: o valor da chave, a posição da linha no arquivo e a quantidade de bytes por linha.

pub fn index_file(file: &File, tree: &mut BTree) {
    let mut reader = BufReader::new(file);

    let mut buf = String::new();
    let mut offset: u64 = 0;

    loop {
        buf.clear();

        let size: u64 = reader
            .read_line(&mut buf)
            .expect("reading from cursor shouldn't fail")
            .try_into()
            .unwrap();

        if size == 0 {
            break;
        }

        let key_value = match get_key(0, &buf) {
            None => {
                offset += size;
                continue;
            }
            Some(value) => value,
        };

        tree.insert(Key::create(key_value, (offset, size)));
        offset += size;
    }
}

Para recuperar as informações do arquivo, a função recebe o arquivo e uma tupla contendo a posição da linha e seu tamanho.

pub fn read_line(file: &mut File, position: (u64, u64)) -> Result<String, Box<dyn error::Error>> {
    let (start, offset) = position;
    file.seek(SeekFrom::Start(start))?;

    let mut read_buf = vec![0; offset.try_into().unwrap()];
    file.read_exact(&mut read_buf)?;

    match String::from_utf8(read_buf) {
        Err(e) => Err(Box::new(e)),
        Ok(line) => Ok(line),
    }
}

E a performance?

Para testar o indexador, gerei um arquivo CSVs com 10.000.000 de registros e aproximadamente 500MB de tamanho. Abaixo, o script para indexar esse arquivo usando nós de ordem 1.000:

fn main() -> std::io::Result<()> {

    let filename = "/home/gdarruda/Projects/sandbox/clients.csv";
    let mut file = File::open(filename)?;
    let mut tree = index::btree::BTree::create(1000, "/home/gdarruda/btree_files");
    csv::index_file(&file, &mut tree);

    Ok(())
}

Para indexar esse arquivo, foram necessárias 3:17 horas, mas não passou de 3MB de consumo total de memória. Um tempo alto para indexar, teríamos várias oportunidades de otimização – melhorar serialização; não forçar escrita em disco para cada inclusão; otimizar os tamanhos dos nós; usar cache – mas esse não é o ponto do exercício proposto.

Olhando para o desempenho da busca – considerando o pior caso, que é procurar e não encontrar um registro – procurando por 1.000 chaves aleatórias não existentes, o tempo médio de busca foi de 474µs com desvio de 26µs.

fn main() -> std::io::Result<()> {

    let filename = "/home/gdarruda/Projects/sandbox/clients.csv";
    let mut file = File::open(filename)?;
    let tree = index::btree::BTree::load("/home/gdarruda/btree_files");

    for uuid in (0..1000).map(|_| Uuid::new_v4().to_string()) {

        let now = SystemTime::now();

        match tree.search(&uuid) {
            None => {},
            Some(key) => {
                match csv::read_line(&mut file, key.position) {
                    Err(e) => {println!("Error: {}", e)},
                    Ok(line) => {println!("Found line: {}", line)}
                }
            }
        };

        match now.elapsed() {
            Ok(elapsed) => {
                println!("{}", elapsed.as_micros());
            }
            Err(e) => {
                println!("Error: {e:?}");
            }
        }

    }

    Ok(())
}

Pode-se pensar em otimizações para essa parte da busca também, mas essa implementação simplória já demonstra a utilidade de uma estrutura como essa para grandes base de dados: boa performance de busca, baixíssimo consumo de memória.

Por que fazer isso?

Não sou a favor que os programadores fiquem recriando banco de dados e frameworks: apesar de ser uma ótima forma de aprender profundamente sobre um tema, provavelmente não é a forma mais eficiente. De qualquer forma, acho fundamental ter uma noção intuitiva de como as coisas funcionam e não focar apenas nas APIs das ferramentas.

Mais importante que entender os detalhes da minha implementação, eu gostaria que os programadores conseguissem responder perguntas sobre índices basedos em árvore:

  • Se vou ler 50% da tabela, um índice me ajuda? E se for 75%? E se for 5%?
  • Índices fazem sentido para armazenamento colunar, como um arquivo parquet por exemplo?
  • Quais as vantagens e desvantagens de colocar múltiplas colunas em um único índice?
  • Em uma coluna com poucos valores distintos, faz mais sentido usar ela como partição ou criar um índice? Por quê?

Obviamente, todas essas perguntas têm vários “depende” em uma situação real, mas é possível ter uma noção das respostas pelo conhecimento teórico. São perguntas que realmente aparecem no dia-a-dia, escolher a melhor solução de armazenamento costumar ser uma porta de sentido único.

Quando eu comecei a carreira, o padrão era usar bancos relacionais, no máximo a discussão era qual deles escolher. Hoje em dia, pode-se recorrer a uma miríade de soluções especializadas para cada caso de uso: desde armazenar dados como arquivos em storage, passando por bancos relacionais e diversos tipos de bancos NoSQL especializados (e.g. grafos, chave-valor, documento, MPP).

As soluções especializadas podem ser mais escaláveis e nem sempre usar Postgres para tudo é a melhor opção, mas ferramentas especializadas podem ser inutilizáveis quando adaptadas para cenários fora de seu caso de uso. Aprender sobre todas as soluções de todos provedores de cloud é inviável, focar em conhecer os conceitos das ferramentas é melhor que ler infinitos artigos no Medium comparando as soluções em termos de “prós e contras”.