crates/graphar-falkordb/src/bulk.rs
//! Native `GRAPH.BULK` fast-lane — pure-Rust encoding of FalkorDB's binary
//! bulk-insert protocol (the wire the `falkordb_bulk_loader` speaks) straight
//! from Arrow `RecordBatch`es, with **no CSV, no Python, no shell** (L9).
//!
//! `GRAPH.BULK` bypasses the Cypher parser/planner entirely: instead of
//! `UNWIND … CREATE` (parse + plan + execute per batch) it ships a **typed,
//! columnar** node/relation payload that the server ingests directly — the fast
//! lane for INITIAL / empty-graph population. `knut` is in the perfect position
//! to emit it because it already holds typed Arrow columns, exactly the shape
//! the bulk format wants. The `UNWIND` path stays the incremental/upsert
//! fallback (see [`crate::loader`]).
//!
//! ## Wire format (mirrors FalkorDB `src/bulk_insert/bulk_insert.c`)
//!
//! The command is
//! `GRAPH.BULK <key> [BEGIN] <node_count> <relation_count> <blob…>` where the
//! blobs are **all node blobs first, then all relation blobs**, each a separate
//! RESP bulk-string argument. Node global ids are assigned in blob order
//! (first node of the first blob = id `0`), which is exactly how a relation
//! blob's `uint64` endpoints reference them.
//!
//! **Node blob** = `label\0` · `u32 prop_count` · (`propname\0`)* ·
//! then, per node, `prop_count` typed property values in header order.
//!
//! **Relation blob** = `reltype\0` · `u32 prop_count` · (`propname\0`)* ·
//! then, per relation, `u64 src_id` · `u64 dst_id` · `prop_count` values.
//!
//! **Typed value** = 1 type byte then the payload:
//! `NULL`(0) → nothing; `BOOL`(1) → 1 byte; `DOUBLE`(2) → f64-LE; `STRING`(3) →
//! bytes + `\0`; `LONG`(4) → i64-LE; `ARRAY`(5) → `u64-LE len` + that many
//! typed values. Integers/floats are little-endian (the server reads them with
//! a raw host-order pointer cast; knut targets are little-endian).
//!
//! ## Parity with the UNWIND path
//!
//! Each node carries a `_gar_id` property = its absolute 0-based row index
//! **within its node group** — identical to what `insert_vertices` writes —
//! UNLESS the batches already carry a `_gar_id` column, in which case that column
//! is used verbatim (no duplicate; lets a caller like `knut-popviz` store a
//! *global*, per-group-offset `_gar_id` and stay byte-identical to its UNWIND
//! load). See [`encode_node_blob`]. Edge
//! endpoints are given as per-endpoint-label `_gar_id`s (the `src`/`dst`
//! columns, exactly like `insert_edges`) and translated here to global node ids
//! by adding the destination group's base offset. So a `GRAPH.BULK` load
//! produces the byte-identical graph (same labels, same `_gar_id`, same
//! properties, same relationships) a per-group `insert_vertices` +
//! `insert_edges` sequence would — see the container-gated parity test.
use std::collections::HashMap;
use arrow_array::{
Array, BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array, LargeListArray,
LargeStringArray, ListArray, RecordBatch, StringArray,
};
use arrow_schema::DataType;
use crate::cypher::temporal_iso;
use crate::error::{FalkorError, Result};
/// FalkorDB bulk property type bytes (`bulk_insert.c` `TYPE` enum). Public so
/// the DB-free encode tests can assert the exact wire bytes.
pub mod ty {
pub const NULL: u8 = 0;
pub const BOOL: u8 = 1;
pub const DOUBLE: u8 = 2;
pub const STRING: u8 = 3;
pub const LONG: u8 = 4;
pub const ARRAY: u8 = 5;
}
/// A group of same-label nodes to bulk-load. `_gar_id` is assigned as the
/// absolute 0-based row index across `batches` (matching `insert_vertices`), and
/// the group's nodes occupy a contiguous global-id range in the order groups are
/// passed to [`build_bulk_payload`].
pub struct BulkNodeGroup<'a> {
pub label: &'a str,
pub batches: &'a [RecordBatch],
}
/// A group of same-type relationships. `batches` carry Int64 `src`/`dst` columns
/// holding the endpoints' per-label `_gar_id`s (exactly like `insert_edges`);
/// every other column is an edge property. Endpoints are translated to global
/// node ids using the node groups' base offsets.
pub struct BulkEdgeGroup<'a> {
pub src_label: &'a str,
pub edge_label: &'a str,
pub dst_label: &'a str,
pub batches: &'a [RecordBatch],
}
/// The encoded `GRAPH.BULK` payload: the two header counts and the ordered blobs
/// (node blobs first, then relation blobs) ready to be sent as command args.
pub struct BulkPayload {
pub node_count: u64,
pub relation_count: u64,
/// Node blobs then relation blobs, each one RESP bulk-string argument.
pub blobs: Vec<Vec<u8>>,
/// Per-blob entity count, parallel to `blobs` (nodes for a node blob,
/// relations for a relation blob). Lets the streamed/appendable sender emit
/// each blob in its own `GRAPH.BULK` call with the right per-call counts
/// while the server keeps assigning global node ids in blob order.
pub blob_counts: Vec<u64>,
/// How many leading entries of `blobs` are **node** blobs (the rest are
/// relation blobs). All node blobs must be sent before any relation blob.
pub node_blob_len: usize,
}
// --------------------------------------------------------------------------
// Value encoding — the typed-cell primitives (pure, connection-free, tested
// byte-for-byte without a DB).
// --------------------------------------------------------------------------
/// Append a `LONG` (type byte + i64 little-endian).
fn put_long(buf: &mut Vec<u8>, v: i64) {
buf.push(ty::LONG);
buf.extend_from_slice(&v.to_le_bytes());
}
/// Append a `DOUBLE` (type byte + f64 little-endian).
fn put_double(buf: &mut Vec<u8>, v: f64) {
buf.push(ty::DOUBLE);
buf.extend_from_slice(&v.to_le_bytes());
}
/// Append a `STRING` (type byte + UTF-8 bytes + NUL terminator).
fn put_string(buf: &mut Vec<u8>, s: &str) {
buf.push(ty::STRING);
buf.extend_from_slice(s.as_bytes());
buf.push(0);
}
/// Append a NUL-terminated header string (no type byte — used for label /
/// reltype / property names).
fn put_cstr(buf: &mut Vec<u8>, s: &str) {
buf.extend_from_slice(s.as_bytes());
buf.push(0);
}
/// Encode one Arrow cell as a typed bulk value into `buf`. Mirrors
/// [`crate::cypher::arrow_value_to_cypher`]'s type dispatch so the stored value
/// is byte-identical in meaning to the UNWIND path: Int→LONG, Float→DOUBLE,
/// Utf8→STRING, Bool→BOOL, Date/Timestamp→STRING (shared ISO formatting),
/// List→ARRAY, null (and non-finite floats, which Cypher renders as `null`)→NULL.
fn encode_cell(buf: &mut Vec<u8>, array: &dyn Array, row: usize) -> Result<()> {
if array.is_null(row) {
buf.push(ty::NULL);
return Ok(());
}
match array.data_type() {
DataType::Boolean => {
let v = array
.as_any()
.downcast_ref::<BooleanArray>()
.unwrap()
.value(row);
buf.push(ty::BOOL);
buf.push(u8::from(v));
}
DataType::Int32 => {
let v = array
.as_any()
.downcast_ref::<Int32Array>()
.unwrap()
.value(row);
put_long(buf, v as i64);
}
DataType::Int64 => {
let v = array
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.value(row);
put_long(buf, v);
}
// Non-finite floats have no Cypher literal (the UNWIND path stores
// `null`); match that exactly so the two graphs agree.
DataType::Float32 => {
let v = array
.as_any()
.downcast_ref::<Float32Array>()
.unwrap()
.value(row);
if v.is_finite() {
put_double(buf, v as f64);
} else {
buf.push(ty::NULL);
}
}
DataType::Float64 => {
let v = array
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.value(row);
if v.is_finite() {
put_double(buf, v);
} else {
buf.push(ty::NULL);
}
}
DataType::Utf8 => {
let v = array
.as_any()
.downcast_ref::<StringArray>()
.unwrap()
.value(row);
put_string(buf, v);
}
DataType::LargeUtf8 => {
let v = array
.as_any()
.downcast_ref::<LargeStringArray>()
.unwrap()
.value(row);
put_string(buf, v);
}
DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _) => {
match temporal_iso(array, row) {
Some(s) => put_string(buf, &s),
None => buf.push(ty::NULL),
}
}
DataType::List(_) => {
let list = array.as_any().downcast_ref::<ListArray>().unwrap();
encode_array(buf, list.value(row).as_ref())?;
}
DataType::LargeList(_) => {
let list = array.as_any().downcast_ref::<LargeListArray>().unwrap();
encode_array(buf, list.value(row).as_ref())?;
}
other => return Err(FalkorError::UnsupportedType(format!("{other:?}"))),
}
Ok(())
}
/// Encode a whole Arrow array as a bulk `ARRAY` value: type byte, `u64` element
/// count, then each element as a typed value.
fn encode_array(buf: &mut Vec<u8>, child: &dyn Array) -> Result<()> {
buf.push(ty::ARRAY);
buf.extend_from_slice(&(child.len() as u64).to_le_bytes());
for i in 0..child.len() {
encode_cell(buf, child, i)?;
}
Ok(())
}
// --------------------------------------------------------------------------
// Blob size estimation — presize the blob Vecs from the Arrow data so encode
// never grows from a 64-byte stub through O(log n) reallocs. Estimates are
// UPPER BOUNDS (never under), exact for the fixed-width + string types.
// --------------------------------------------------------------------------
/// Upper-bound encoded-byte estimate for every value of `array` as typed bulk
/// cells. Exact for Bool (2/row), Int (9/row), finite Float (9/row; a NULL cell
/// is 1 ≤ 9) and tight for strings (offset-range bytes + 2/row); a measured
/// upper bound for temporals (ISO strings ≤ 40 bytes) and lists (9/row header +
/// the child's own estimate). Unknown types get 1/row (they error in
/// `encode_cell` before the size could matter).
fn estimate_value_bytes(array: &dyn Array) -> usize {
let n = array.len();
match array.data_type() {
DataType::Boolean => 2 * n,
DataType::Int32 | DataType::Int64 | DataType::Float32 | DataType::Float64 => 9 * n,
DataType::Utf8 => {
let a = array.as_any().downcast_ref::<StringArray>().unwrap();
let off = a.offsets();
(off[n] - off[0]) as usize + 2 * n
}
DataType::LargeUtf8 => {
let a = array.as_any().downcast_ref::<LargeStringArray>().unwrap();
let off = a.offsets();
(off[n] - off[0]) as usize + 2 * n
}
DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _) => 40 * n,
DataType::List(_) => {
let list = array.as_any().downcast_ref::<ListArray>().unwrap();
9 * n + estimate_value_bytes(list.values().as_ref())
}
DataType::LargeList(_) => {
let list = array.as_any().downcast_ref::<LargeListArray>().unwrap();
9 * n + estimate_value_bytes(list.values().as_ref())
}
_ => n,
}
}
/// Upper-bound byte size of one node blob (header + every row's typed cells),
/// `synth_gar` adding the synthesized `_gar_id` header name + 9 bytes per row.
fn estimate_node_blob_bytes(label: &str, batches: &[RecordBatch], synth_gar: bool) -> usize {
let props = prop_names(batches, &[]);
let mut est = label.len() + 1 + 4 + props.iter().map(|p| p.len() + 1).sum::<usize>();
if synth_gar {
est += "_gar_id".len() + 1;
}
for batch in batches {
if synth_gar {
est += 9 * batch.num_rows();
}
for name in &props {
if let Some(c) = batch.column_by_name(name) {
est += estimate_value_bytes(c.as_ref());
}
}
}
est
}
/// Upper-bound byte size of one relation blob: header + 16 endpoint bytes per
/// row + the property cells.
fn estimate_edge_blob_bytes(group: &BulkEdgeGroup<'_>) -> usize {
let props = prop_names(group.batches, &["src", "dst"]);
let mut est =
group.edge_label.len() + 1 + 4 + props.iter().map(|p| p.len() + 1).sum::<usize>();
for batch in group.batches {
est += 16 * batch.num_rows();
for name in &props {
if let Some(c) = batch.column_by_name(name) {
est += estimate_value_bytes(c.as_ref());
}
}
}
est
}
// --------------------------------------------------------------------------
// Blob encoding.
// --------------------------------------------------------------------------
/// The property column names of a group's batches, in schema order. All batches
/// in a group share the schema (they come from one label / reltype), so batch 0
/// defines it. `schema_ref()` borrows from the batch (not a temporary), keeping
/// the names alive for `'a`.
fn prop_names<'a>(batches: &'a [RecordBatch], skip: &[&str]) -> Vec<&'a str> {
batches[0]
.schema_ref()
.fields()
.iter()
.map(|f| f.name().as_str())
.filter(|n| !skip.contains(n))
.collect()
}
/// Encode one label's node blob and return `(blob, node_count)`.
///
/// **`_gar_id` handling** — two modes, auto-detected from the schema:
/// * If the batches **already carry a `_gar_id` column** (as `knut-popviz`'s
/// vertex batches do — a *global*, offset-based id), it is used **verbatim**
/// as the node's `_gar_id` property: this both avoids emitting a **duplicate**
/// `_gar_id` and stores the caller's global value, so a multi-label load whose
/// `_gar_id`s are offset per group stays byte-identical to the UNWIND path
/// (which writes exactly that column). The header/property order is the batch
/// schema order.
/// * Otherwise `_gar_id` is **synthesized** as the absolute 0-based row index
/// across the group and prepended to every batch column — the original scheme,
/// identical to what a `_gar_id`-less `insert_vertices` writes.
///
/// Either way, per-endpoint node ids used by [`encode_edge_blob`] are the group
/// **position** (0-based), whose global base offset is added in
/// [`build_bulk_payload`]; a caller that supplies global `_gar_id`s must pass the
/// node groups in the same cumulative order so position + base == `_gar_id`.
/// `gar_base` is the absolute 0-based row index of `batches[0][0]` **within its
/// node group** — the offset a synthesized `_gar_id` starts at, so a label split
/// across several blobs still assigns `_gar_id = group-relative row index`
/// (`0` for the whole-group / single-blob case). Ignored when the batches already
/// carry a `_gar_id` column (that column is used verbatim).
#[cfg(test)]
fn encode_node_blob(label: &str, batches: &[RecordBatch], gar_base: i64) -> Result<(Vec<u8>, u64)> {
encode_node_blob_ex(label, batches, gar_base, true)
}
/// [`encode_node_blob`], with `synth_gar = false` writing the batch's columns
/// and nothing else — no synthesized `_gar_id` (see [`build_bulk_payload_verbatim`]).
fn encode_node_blob_ex(label: &str, batches: &[RecordBatch], gar_base: i64, synth_gar: bool) -> Result<(Vec<u8>, u64)> {
if batches.is_empty() {
return Err(FalkorError::EmptyBatches);
}
// "has the column, or must not synthesize one": either way the header is
// just the batch's columns.
let has_gar_col = !synth_gar
|| batches[0]
.schema_ref()
.column_with_name("_gar_id")
.is_some();
// `prop_names` keeps a `_gar_id` column (it isn't in the skip list), so in
// the supplied-column mode it flows through as an ordinary property.
let props = prop_names(batches, &[]);
// Presize from the Arrow data: the estimate is an upper bound, so the vec
// never reallocates during encode (was `with_capacity(64)` + O(log n) growth).
let mut blob = Vec::with_capacity(estimate_node_blob_bytes(label, batches, !has_gar_col));
put_cstr(&mut blob, label);
if has_gar_col {
// `_gar_id` is already among the columns — header is just the columns.
blob.extend_from_slice(&(props.len() as u32).to_le_bytes());
for p in &props {
put_cstr(&mut blob, p);
}
} else {
// Header prop_count = synthesized _gar_id + every batch column.
blob.extend_from_slice(&((props.len() as u32) + 1).to_le_bytes());
put_cstr(&mut blob, "_gar_id");
for p in &props {
put_cstr(&mut blob, p);
}
}
let mut count: u64 = 0;
for batch in batches {
// Resolve the property columns once per batch (header order).
let cols: Vec<&dyn Array> = props
.iter()
.map(|name| {
batch
.column_by_name(name)
.map(|c| c.as_ref())
.ok_or_else(|| FalkorError::ColumnNotFound((*name).into()))
})
.collect::<Result<_>>()?;
for row in 0..batch.num_rows() {
if !has_gar_col {
put_long(&mut blob, gar_base + count as i64); // synthesized _gar_id
}
for col in &cols {
encode_cell(&mut blob, *col, row)?;
}
count += 1;
}
}
Ok((blob, count))
}
/// Encode one reltype's relation blob and return `(blob, relation_count)`.
/// `resolve` maps a `(label, per_label_gar_id)` to a global node id (adds the
/// endpoint group's base offset). Property columns are every column except
/// `src`/`dst`.
fn encode_edge_blob(
group: &BulkEdgeGroup<'_>,
resolve: impl Fn(&str, i64) -> Result<u64>,
) -> Result<(Vec<u8>, u64)> {
if group.batches.is_empty() {
return Err(FalkorError::EmptyBatches);
}
let props = prop_names(group.batches, &["src", "dst"]);
// Presize from the Arrow data (upper bound — no realloc during encode).
let mut blob = Vec::with_capacity(estimate_edge_blob_bytes(group));
put_cstr(&mut blob, group.edge_label);
blob.extend_from_slice(&(props.len() as u32).to_le_bytes());
for p in &props {
put_cstr(&mut blob, p);
}
let mut count: u64 = 0;
for batch in group.batches {
let src = batch
.column_by_name("src")
.ok_or_else(|| FalkorError::ColumnNotFound("src".into()))?;
let dst = batch
.column_by_name("dst")
.ok_or_else(|| FalkorError::ColumnNotFound("dst".into()))?;
let src = src
.as_any()
.downcast_ref::<Int64Array>()
.ok_or_else(|| FalkorError::UnsupportedType("edge `src` must be Int64".into()))?;
let dst = dst
.as_any()
.downcast_ref::<Int64Array>()
.ok_or_else(|| FalkorError::UnsupportedType("edge `dst` must be Int64".into()))?;
let cols: Vec<&dyn Array> = props
.iter()
.map(|name| {
batch
.column_by_name(name)
.map(|c| c.as_ref())
.ok_or_else(|| FalkorError::ColumnNotFound((*name).into()))
})
.collect::<Result<_>>()?;
for row in 0..batch.num_rows() {
let s = resolve(group.src_label, src.value(row))?;
let d = resolve(group.dst_label, dst.value(row))?;
blob.extend_from_slice(&s.to_le_bytes());
blob.extend_from_slice(&d.to_le_bytes());
for col in &cols {
encode_cell(&mut blob, *col, row)?;
}
count += 1;
}
}
Ok((blob, count))
}
/// Max rows encoded into a single `GRAPH.BULK` blob. Two problems this bounds:
///
/// 1. **Broken pipe at scale.** One blob per label meant the dominant label
/// (a 60M-row `TaxRecord` at full demo scale) became a **single** `GRAPH.BULK`
/// command whose payload sailed past FalkorDB's `proto-max-bulk-len` (512MB
/// default) — the server closed the connection mid-write (`Broken pipe`).
/// Splitting a label into ≤`MAX_ROWS_PER_BLOB` blobs keeps every command well
/// under the limit.
/// 2. **Single-core encode.** `build_bulk_payload` used to encode each label's
/// whole blob in a serial `for` loop, so that 60M-row label pinned **one** core
/// for the entire encode. Chunking yields many independent units, which now fan
/// across all cores via gatling.
///
/// ~2M rows is a few hundred MB at most for these schemas — comfortably under
/// 512MB with headroom, and small enough that a big label yields dozens of units
/// for the pool. It is a size cap, not a semantic one: blobs stay ordered, so the
/// server still assigns global ids exactly as the one-blob-per-label path did
/// (see [`build_bulk_payload`]).
pub const MAX_ROWS_PER_BLOB: usize = 2_000_000;
/// Split `batches` into consecutive `[lo, hi)` index windows whose cumulative row
/// count stays under `budget`, never splitting an individual batch (a lone
/// over-budget batch is its own window). Order preserved, so concatenating the
/// windows' encoded blobs reproduces the one-blob encoding's node order.
fn chunk_windows_with(batches: &[RecordBatch], budget: usize) -> Vec<std::ops::Range<usize>> {
let budget = budget.max(1);
let mut out = Vec::new();
let (mut lo, mut rows) = (0usize, 0usize);
for (i, b) in batches.iter().enumerate() {
if rows > 0 && rows + b.num_rows() > budget {
out.push(lo..i);
lo = i;
rows = 0;
}
rows += b.num_rows();
}
if lo < batches.len() {
out.push(lo..batches.len());
}
out
}
/// Build the full `GRAPH.BULK` payload from node + edge groups — the pure,
/// connection-free core (unit-tested byte-for-byte without a DB). Node groups
/// are laid out in order (their global-id ranges are contiguous and ascending),
/// each label's base offset recorded, and edge endpoints (per-label `_gar_id`s)
/// translated to global node ids. A label may appear in at most one node group.
///
/// Each group is split into ≤[`MAX_ROWS_PER_BLOB`] blobs (see [`chunk_windows`])
/// and the blobs are encoded **in parallel** across cores (one gatling unit per
/// blob; LAW #0: gatling, never rayon), then laid down in order. Because the
/// blobs stay ordered — every chunk of label `li` before label `li+1`, chunks in
/// row order — the server assigns the same contiguous global ids it would for one
/// blob per label, so the produced graph is identical; only the wire is split.
pub fn build_bulk_payload(
nodes: &[BulkNodeGroup<'_>],
edges: &[BulkEdgeGroup<'_>],
) -> Result<BulkPayload> {
build_bulk_payload_with(nodes, edges, MAX_ROWS_PER_BLOB)
}
/// [`build_bulk_payload`] with an explicit per-blob row budget — the impl, so the
/// chunk-splitting can be exercised on tiny data (the public entry pins the
/// production [`MAX_ROWS_PER_BLOB`]).
pub(crate) fn build_bulk_payload_with(
nodes: &[BulkNodeGroup<'_>],
edges: &[BulkEdgeGroup<'_>],
max_rows: usize,
) -> Result<BulkPayload> {
build_bulk_payload_impl(nodes, edges, max_rows, true)
}
/// A **nodes-only** `GRAPH.BULK` payload whose nodes carry exactly their batch
/// columns — no synthesized `_gar_id`. For a caller whose nodes are keyed on
/// their own id property (bifröst's first load of an Iceberg table) and whose
/// later writes are `MERGE`s on that key: a `_gar_id` on the bulk-loaded nodes
/// only would make them differ from every node merged afterwards. No edges —
/// bulk edges address nodes by `_gar_id` position.
pub fn build_bulk_payload_verbatim(nodes: &[BulkNodeGroup<'_>]) -> Result<BulkPayload> {
build_bulk_payload_impl(nodes, &[], MAX_ROWS_PER_BLOB, false)
}
/// An **edges-only** `GRAPH.BULK` payload whose `src`/`dst` columns hold
/// **absolute FalkorDB node ids** (`ID(n)`) of nodes already in the graph — for
/// appending relationships between nodes loaded earlier, where the per-label
/// `_gar_id` positions [`build_bulk_payload`] translates do not exist. The ids are
/// the caller's to resolve (`MATCH (n:L) RETURN n.key, ID(n)`); every other column
/// is an edge property. Blobs encoded in parallel (gatling), in order.
pub fn build_bulk_payload_edges_absolute(edges: &[BulkEdgeGroup<'_>]) -> Result<BulkPayload> {
use znippy_zoomies::gatling_forkjoin::gatling_for_each;
let mut plan: Vec<(usize, std::ops::Range<usize>)> = Vec::new();
for (gi, g) in edges.iter().enumerate() {
for w in chunk_windows_with(g.batches, MAX_ROWS_PER_BLOB) {
plan.push((gi, w));
}
}
if plan.is_empty() {
return Err(FalkorError::EmptyBatches);
}
let absolute = |label: &str, id: i64| -> Result<u64> {
u64::try_from(id).map_err(|_| FalkorError::UnsupportedType(format!("negative node id {id} for label {label}")))
};
let encoded: Vec<Result<(Vec<u8>, u64)>> = gatling_for_each(plan.len(), 0, |k| {
let (gi, w) = &plan[k];
let g = &edges[*gi];
let windowed = BulkEdgeGroup {
src_label: g.src_label,
edge_label: g.edge_label,
dst_label: g.dst_label,
batches: &g.batches[w.clone()],
};
encode_edge_blob(&windowed, absolute)
});
let mut blobs = Vec::with_capacity(encoded.len());
let mut blob_counts = Vec::with_capacity(encoded.len());
let mut relation_count = 0;
for r in encoded {
let (blob, n) = r?;
relation_count += n;
blobs.push(blob);
blob_counts.push(n);
}
Ok(BulkPayload { node_count: 0, relation_count, blobs, blob_counts, node_blob_len: 0 })
}
fn build_bulk_payload_impl(
nodes: &[BulkNodeGroup<'_>],
edges: &[BulkEdgeGroup<'_>],
max_rows: usize,
synth_gar: bool,
) -> Result<BulkPayload> {
use znippy_zoomies::gatling_forkjoin::gatling_for_each;
if nodes.is_empty() {
return Err(FalkorError::EmptyBatches);
}
// Label bases from full-label row totals (cheap; no encode), so edge resolve
// is exactly what the one-blob path computed. Also the ORDERED node-blob plan:
// (group index, batch-window) per chunk, group order then row order.
let mut base: HashMap<&str, (u64, u64)> = HashMap::new();
// (group index, batch-window, gar_base) per chunk — the group-relative row
// offset the chunk's first node sits at, so a split label keeps assigning
// contiguous `_gar_id`s.
let mut node_plan: Vec<(usize, std::ops::Range<usize>, i64)> = Vec::new();
let mut node_count: u64 = 0;
for (gi, g) in nodes.iter().enumerate() {
if base.contains_key(g.label) {
return Err(FalkorError::UnsupportedType(format!(
"duplicate node label in bulk load: {}",
g.label
)));
}
let total: u64 = g.batches.iter().map(|b| b.num_rows() as u64).sum();
base.insert(g.label, (node_count, total));
node_count += total;
let mut gar_base: i64 = 0; // group-relative offset of the current window
for w in chunk_windows_with(g.batches, max_rows) {
let rows: i64 = g.batches[w.clone()]
.iter()
.map(|b| b.num_rows() as i64)
.sum();
node_plan.push((gi, w, gar_base));
gar_base += rows;
}
}
// Encode every node blob in parallel, results index-ordered (== plan order).
let node_blobs: Vec<Result<(Vec<u8>, u64)>> = gatling_for_each(node_plan.len(), 0, |k| {
let entry = &node_plan[k];
let g = &nodes[entry.0];
encode_node_blob_ex(g.label, &g.batches[entry.1.clone()], entry.2, synth_gar)
});
let resolve = |label: &str, gar: i64| -> Result<u64> {
let (b, n) = base.get(label).ok_or_else(|| {
FalkorError::UnsupportedType(format!("edge endpoint label has no nodes: {label}"))
})?;
if gar < 0 || (gar as u64) >= *n {
return Err(FalkorError::UnsupportedType(format!(
"edge endpoint _gar_id {gar} out of range for label {label} ({n} nodes)"
)));
}
Ok(b + gar as u64)
};
// Ordered edge-blob plan: (group index, batch-window), then encode in parallel.
let mut edge_plan: Vec<(usize, std::ops::Range<usize>)> = Vec::new();
for (gi, g) in edges.iter().enumerate() {
for w in chunk_windows_with(g.batches, max_rows) {
edge_plan.push((gi, w));
}
}
let edge_blobs: Vec<Result<(Vec<u8>, u64)>> = gatling_for_each(edge_plan.len(), 0, |k| {
let entry = &edge_plan[k];
let g = &edges[entry.0];
let windowed = BulkEdgeGroup {
src_label: g.src_label,
edge_label: g.edge_label,
dst_label: g.dst_label,
batches: &g.batches[entry.1.clone()],
};
encode_edge_blob(&windowed, resolve)
});
// Assemble in order, propagating the first encode error.
let node_blob_len = node_blobs.len();
let mut blobs: Vec<Vec<u8>> = Vec::with_capacity(node_blobs.len() + edge_blobs.len());
let mut blob_counts: Vec<u64> = Vec::with_capacity(node_blobs.len() + edge_blobs.len());
for r in node_blobs {
let (blob, n) = r?;
blobs.push(blob);
blob_counts.push(n);
}
let mut relation_count: u64 = 0;
for r in edge_blobs {
let (blob, n) = r?;
relation_count += n;
blobs.push(blob);
blob_counts.push(n);
}
Ok(BulkPayload {
node_count,
relation_count,
blobs,
blob_counts,
node_blob_len,
})
}
#[cfg(test)]
mod tests {
use super::*;
use arrow_array::{Int64Array, StringArray};
use arrow_schema::{Field, Schema};
use std::sync::Arc;
fn people(names: &[&str], ages: &[i64]) -> RecordBatch {
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("name", DataType::Utf8, false),
Field::new("age", DataType::Int64, false),
])),
vec![
Arc::new(StringArray::from(names.to_vec())),
Arc::new(Int64Array::from(ages.to_vec())),
],
)
.unwrap()
}
fn edges(src: &[i64], dst: &[i64]) -> RecordBatch {
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("src", DataType::Int64, false),
Field::new("dst", DataType::Int64, false),
])),
vec![
Arc::new(Int64Array::from(src.to_vec())),
Arc::new(Int64Array::from(dst.to_vec())),
],
)
.unwrap()
}
/// The primitive value encoders must lay down the EXACT documented bytes.
/// RED-when-broken: a wrong type byte, endianness, or terminator flips these.
#[test]
fn typed_values_encode_exact_bytes() {
let mut b = Vec::new();
put_long(&mut b, 1);
assert_eq!(
b,
vec![ty::LONG, 1, 0, 0, 0, 0, 0, 0, 0],
"LONG = 4 + i64-LE"
);
let mut b = Vec::new();
put_double(&mut b, 1.5);
assert_eq!(b[0], ty::DOUBLE);
assert_eq!(&b[1..], &1.5f64.to_le_bytes(), "DOUBLE = 2 + f64-LE");
let mut b = Vec::new();
put_string(&mut b, "ab");
assert_eq!(
b,
vec![ty::STRING, b'a', b'b', 0],
"STRING = 3 + bytes + NUL"
);
}
/// `encode_cell` maps each Arrow type to the right bulk type byte, and both
/// nulls and non-finite floats collapse to `NULL` — matching the UNWIND
/// path's `null` rendering, the invariant the parity gate depends on.
#[test]
fn encode_cell_types_and_null_collapse() {
let bools = BooleanArray::from(vec![Some(true), None]);
let mut b = Vec::new();
encode_cell(&mut b, &bools, 0).unwrap();
assert_eq!(b, vec![ty::BOOL, 1]);
let mut b = Vec::new();
encode_cell(&mut b, &bools, 1).unwrap();
assert_eq!(b, vec![ty::NULL], "null bool → NULL");
let f = Float64Array::from(vec![f64::NAN, f64::INFINITY, 2.0]);
let mut b = Vec::new();
encode_cell(&mut b, &f, 0).unwrap();
assert_eq!(b, vec![ty::NULL], "NaN → NULL (Cypher parity)");
let mut b = Vec::new();
encode_cell(&mut b, &f, 1).unwrap();
assert_eq!(b, vec![ty::NULL], "+Inf → NULL");
let mut b = Vec::new();
encode_cell(&mut b, &f, 2).unwrap();
assert_eq!(b[0], ty::DOUBLE);
}
/// A node blob is `label\0` · `u32 prop_count` · prop names · then per node
/// `_gar_id` + columns. `_gar_id` is the absolute index across batches — the
/// same scheme `insert_vertices` writes, so the produced nodes match. Exact
/// byte assert (RED-when-broken on any header/layout drift).
#[test]
fn node_blob_layout_and_gar_ids() {
let (blob, n) = encode_node_blob(
"Person",
&[people(&["al"], &[30]), people(&["bo"], &[41])],
0,
)
.unwrap();
assert_eq!(n, 2, "two nodes across two batches");
let mut want = Vec::new();
put_cstr(&mut want, "Person");
want.extend_from_slice(&3u32.to_le_bytes()); // _gar_id + name + age
put_cstr(&mut want, "_gar_id");
put_cstr(&mut want, "name");
put_cstr(&mut want, "age");
// node 0: _gar_id 0, name 'al', age 30
put_long(&mut want, 0);
put_string(&mut want, "al");
put_long(&mut want, 30);
// node 1: _gar_id 1 (absolute across batches), name 'bo', age 41
put_long(&mut want, 1);
put_string(&mut want, "bo");
put_long(&mut want, 41);
assert_eq!(blob, want);
}
/// A batch that ALREADY carries a `_gar_id` column must NOT get a second,
/// synthesized one (the dedup), and its supplied — possibly global/offset —
/// values must be stored verbatim. This is the exact shape `knut-popviz`
/// feeds (its vertex batches carry a global, offset-based `_gar_id`), so this
/// is what keeps a multi-label popviz bulk load byte-identical to its UNWIND
/// load. RED-when-broken: re-introduce the synthetic `_gar_id` and the header
/// grows a duplicate `_gar_id` and the stored value flips to 0-based.
#[test]
fn node_blob_uses_supplied_gar_id_column_without_duplicating() {
// A popviz-shaped batch: `_gar_id` (global, offset 100) + `name`.
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("_gar_id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
])),
vec![
Arc::new(Int64Array::from(vec![100, 101])),
Arc::new(StringArray::from(vec!["al", "bo"])),
],
)
.unwrap();
let (blob, n) = encode_node_blob("Person", &[batch], 0).unwrap();
assert_eq!(n, 2);
// Header carries EXACTLY ONE `_gar_id` (dedup), in schema order.
let mut want = Vec::new();
put_cstr(&mut want, "Person");
want.extend_from_slice(&2u32.to_le_bytes()); // _gar_id + name — NOT 3
put_cstr(&mut want, "_gar_id");
put_cstr(&mut want, "name");
// node 0: the SUPPLIED global _gar_id 100 (not a synthesized 0), name 'al'
put_long(&mut want, 100);
put_string(&mut want, "al");
// node 1: supplied global _gar_id 101, name 'bo'
put_long(&mut want, 101);
put_string(&mut want, "bo");
assert_eq!(blob, want);
}
/// End-to-end payload: node global ids are contiguous across groups and edge
/// endpoints are translated from per-label `_gar_id` to global id by adding
/// the destination group's base. RED-when-broken: drop the base offset and
/// the `KNOWS` endpoints point at the wrong (or out-of-range) nodes.
#[test]
fn payload_counts_and_edge_endpoint_translation() {
let pbatch = people(&["a", "b"], &[1, 2]);
let people_g = BulkNodeGroup {
label: "Person",
batches: std::slice::from_ref(&pbatch),
};
let city_batch = people(&["X", "Y", "Z"], &[9, 9, 9]);
let city_g = BulkNodeGroup {
label: "City",
batches: std::slice::from_ref(&city_batch),
};
// Person(gar 0) -LIVES-> City(gar 2): global src 0, global dst = 2 (base) + 2 = 4.
let e = edges(&[0], &[2]);
let edge_g = BulkEdgeGroup {
src_label: "Person",
edge_label: "LIVES",
dst_label: "City",
batches: std::slice::from_ref(&e),
};
let p = build_bulk_payload(&[people_g, city_g], &[edge_g]).unwrap();
assert_eq!(p.node_count, 5, "2 Person + 3 City");
assert_eq!(p.relation_count, 1);
assert_eq!(p.blobs.len(), 3, "2 node blobs + 1 relation blob");
// The relation blob's endpoints: reltype header then u64 src=0, dst=4.
let rel = p.blobs.last().unwrap();
let mut want = Vec::new();
put_cstr(&mut want, "LIVES");
want.extend_from_slice(&0u32.to_le_bytes()); // no edge props
want.extend_from_slice(&0u64.to_le_bytes()); // global src
want.extend_from_slice(&4u64.to_le_bytes()); // global dst = City base(2) + gar(2)
assert_eq!(rel, &want);
}
#[test]
fn chunk_windows_split_on_the_row_budget_never_within_a_batch() {
let one = |n: usize| people(&vec!["x"; n], &vec![1i64; n]);
// three 1-row batches, budget 2 → [0,1] then [2].
let three = [one(1), one(1), one(1)];
assert_eq!(chunk_windows_with(&three, 2), vec![0..2, 2..3]);
// a generous budget keeps them in one window.
assert_eq!(chunk_windows_with(&three, 100), vec![0..3]);
// a single over-budget batch is its own window (never split mid-batch).
let big = [one(5)];
assert_eq!(chunk_windows_with(&big, 2), vec![0..1]);
}
#[test]
fn gar_base_continues_the_synthesized_gar_id_across_a_split() {
// No `_gar_id` column ⇒ it is synthesized as gar_base + row. A chunk that
// starts at group offset 7 must store 7, not restart at 0.
let b0 = encode_node_blob("P", &[people(&["a"], &[1])], 0).unwrap().0;
let b7 = encode_node_blob("P", &[people(&["a"], &[1])], 7).unwrap().0;
assert_eq!(
b0.len(),
b7.len(),
"gar_base changes only the _gar_id value"
);
let long = |v: i64| {
let mut o = vec![ty::LONG];
o.extend_from_slice(&v.to_le_bytes());
o
};
assert!(
b7.windows(9).any(|w| w == long(7).as_slice()),
"chunk stored _gar_id 7"
);
assert!(
b0.windows(9).any(|w| w == long(0).as_slice()),
"whole stored _gar_id 0"
);
}
#[test]
fn chunking_a_label_matches_the_single_blob_graph() {
// 3 Person rows across 3 batches; an edge whose dst (gar 2) lands in the
// SECOND chunk — so this also proves endpoints resolve across the split.
let b = [
people(&["a"], &[1]),
people(&["b"], &[2]),
people(&["c"], &[3]),
];
let e = edges(&[0], &[2]);
let build = |max_rows: usize| {
build_bulk_payload_with(
&[BulkNodeGroup {
label: "Person",
batches: &b,
}],
&[BulkEdgeGroup {
src_label: "Person",
edge_label: "KNOWS",
dst_label: "Person",
batches: std::slice::from_ref(&e),
}],
max_rows,
)
.unwrap()
};
let chunked = build(2);
let whole = build(usize::MAX);
assert_eq!(chunked.node_blob_len, 2, "3 rows / budget 2 → 2 node blobs");
assert_eq!(whole.node_blob_len, 1, "no budget → one blob for the label");
assert_eq!(&chunked.blob_counts[..chunked.node_blob_len], &[2, 1]);
// Same graph either way: counts identical …
assert_eq!(chunked.node_count, whole.node_count);
assert_eq!(chunked.node_count, 3);
assert_eq!(chunked.relation_count, whole.relation_count);
// … and the relation blob is byte-identical (dst resolves to global id 2
// whether or not the nodes were split — endpoints use blob-order ids).
assert_eq!(
chunked.blobs.last().unwrap(),
whole.blobs.last().unwrap(),
"edges resolve identically across the chunk boundary",
);
}
/// Out-of-range and unknown-label endpoints are loud errors, never a silent
/// dangling edge (the bulk format can't MATCH-fail like UNWIND does).
#[test]
fn edge_endpoint_out_of_range_errors() {
let one = people(&["a"], &[1]);
let pg = BulkNodeGroup {
label: "Person",
batches: std::slice::from_ref(&one),
};
let e = edges(&[0], &[5]); // dst 5 ≥ 1 node
let eg = BulkEdgeGroup {
src_label: "Person",
edge_label: "KNOWS",
dst_label: "Person",
batches: std::slice::from_ref(&e),
};
assert!(
build_bulk_payload(&[pg], &[eg]).is_err(),
"dst out of range must error"
);
let pg = BulkNodeGroup {
label: "Person",
batches: std::slice::from_ref(&one),
};
let e2 = edges(&[0], &[0]);
let eg = BulkEdgeGroup {
src_label: "Ghost",
edge_label: "KNOWS",
dst_label: "Person",
batches: std::slice::from_ref(&e2),
};
assert!(
build_bulk_payload(&[pg], &[eg]).is_err(),
"unknown src label must error"
);
}
/// Empty node set and a duplicate node label are rejected (the id-range
/// bookkeeping assumes one contiguous group per label).
#[test]
fn empty_and_duplicate_label_rejected() {
assert!(matches!(
build_bulk_payload(&[], &[]),
Err(FalkorError::EmptyBatches)
));
let b = people(&["a"], &[1]);
let g1 = BulkNodeGroup {
label: "Person",
batches: std::slice::from_ref(&b),
};
let g2 = BulkNodeGroup {
label: "Person",
batches: std::slice::from_ref(&b),
};
assert!(
build_bulk_payload(&[g1, g2], &[]).is_err(),
"duplicate label must error"
);
}
#[test]
fn a_verbatim_payload_writes_the_columns_and_no_gar_id() {
let batch = people(&["al", "bo"], &[30, 41]);
let p = build_bulk_payload_verbatim(&[BulkNodeGroup { label: "Person", batches: std::slice::from_ref(&batch) }]).unwrap();
assert_eq!((p.node_count, p.relation_count, p.blobs.len()), (2, 0, 1));
let mut want = Vec::new();
put_cstr(&mut want, "Person");
want.extend_from_slice(&2u32.to_le_bytes());
put_cstr(&mut want, "name");
put_cstr(&mut want, "age");
put_string(&mut want, "al");
put_long(&mut want, 30);
put_string(&mut want, "bo");
put_long(&mut want, 41);
assert_eq!(p.blobs[0], want);
// The default payload still synthesizes it.
let d = build_bulk_payload(&[BulkNodeGroup { label: "Person", batches: std::slice::from_ref(&batch) }], &[]).unwrap();
assert!(d.blobs[0].windows(8).any(|w| w == b"_gar_id\0"));
}
/// Blob `Vec`s are presized from the Arrow data instead of growing from a
/// 64-byte stub: for every supported column type the estimate is an **upper
/// bound** (so `Vec::with_capacity(estimate)` never reallocates — at most
/// one growth would only ever come from an estimate miss, which this test
/// forbids), and for the exact-size types (ints, floats, bools, strings) it
/// is tight. RED before the presizing lands: `estimate_node_blob_bytes` /
/// `estimate_edge_blob_bytes` do not exist.
#[test]
fn blob_estimates_are_upper_bounds_and_capacity_holds() {
use arrow_array::{BooleanArray, Float64Array};
// A fixture with every encode arm: ints, strings (varied lengths),
// floats incl. NaN (→ NULL), bools, nulls, and a list column.
let list = {
let mut lb = arrow_array::builder::ListBuilder::new(
arrow_array::builder::Int64Builder::new(),
);
lb.values().append_value(1);
lb.values().append_value(2);
lb.append(true);
lb.append(false); // null list
lb.values().append_value(3);
lb.append(true);
Arc::new(lb.finish())
};
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("name", DataType::Utf8, true),
Field::new("age", DataType::Int64, true),
Field::new("score", DataType::Float64, true),
Field::new("ok", DataType::Boolean, true),
Field::new("tags", DataType::List(Arc::new(Field::new("item", DataType::Int64, true))), true),
])),
vec![
Arc::new(StringArray::from(vec![Some("a-longer-name"), None, Some("x")])),
Arc::new(Int64Array::from(vec![Some(1), Some(2), None])),
Arc::new(Float64Array::from(vec![Some(1.5), Some(f64::NAN), None])),
Arc::new(BooleanArray::from(vec![Some(true), None, Some(false)])),
list,
],
)
.unwrap();
// Node blob: estimate is an upper bound and the vec never grew.
let est = estimate_node_blob_bytes("Person", std::slice::from_ref(&batch), true);
let (blob, n) = encode_node_blob("Person", std::slice::from_ref(&batch), 0).unwrap();
assert_eq!(n, 3);
assert!(
est >= blob.len(),
"node estimate {est} must cover encoded len {}",
blob.len()
);
assert!(
blob.capacity() >= est,
"blob was allocated below its estimate — presizing not applied"
);
// Edge blob likewise (src/dst skipped from props, 16 bytes/row endpoints).
let e = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("src", DataType::Int64, false),
Field::new("dst", DataType::Int64, false),
Field::new("w", DataType::Utf8, true),
])),
vec![
Arc::new(Int64Array::from(vec![0, 1])),
Arc::new(Int64Array::from(vec![1, 0])),
Arc::new(StringArray::from(vec![Some("heavy"), None])),
],
)
.unwrap();
let g = BulkEdgeGroup {
src_label: "Person",
edge_label: "KNOWS",
dst_label: "Person",
batches: std::slice::from_ref(&e),
};
let est = estimate_edge_blob_bytes(&g);
let (blob, n) = encode_edge_blob(&g, |_, id| Ok(id as u64)).unwrap();
assert_eq!(n, 2);
assert!(
est >= blob.len(),
"edge estimate {est} must cover encoded len {}",
blob.len()
);
assert!(blob.capacity() >= est);
}
/// ZERO-LOSS: presizing must not change a single encoded byte. The expected
/// bytes are laid down with the `put_*` primitives into a plain `Vec::new()`
/// — the old path's growth-from-nothing allocation — and the presized
/// encoder must produce exactly them, across nulls, NaN, strings, bools,
/// ints, floats and a synthesized `_gar_id`.
#[test]
fn presized_node_blob_bytes_identical_to_unsized_oracle() {
use arrow_array::{BooleanArray, Float64Array};
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("name", DataType::Utf8, true),
Field::new("score", DataType::Float64, true),
Field::new("ok", DataType::Boolean, true),
])),
vec![
Arc::new(StringArray::from(vec![Some("al"), None])),
Arc::new(Float64Array::from(vec![Some(0.25), Some(f64::NAN)])),
Arc::new(BooleanArray::from(vec![None, Some(false)])),
],
)
.unwrap();
let (blob, _) = encode_node_blob("P", std::slice::from_ref(&batch), 5).unwrap();
// Oracle: unsized Vec, same primitives, same order.
let mut want = Vec::new();
put_cstr(&mut want, "P");
want.extend_from_slice(&4u32.to_le_bytes());
put_cstr(&mut want, "_gar_id");
put_cstr(&mut want, "name");
put_cstr(&mut want, "score");
put_cstr(&mut want, "ok");
// row 0: gar 5, "al", 0.25, null bool
put_long(&mut want, 5);
put_string(&mut want, "al");
put_double(&mut want, 0.25);
want.push(ty::NULL);
// row 1: gar 6, null string, NaN→NULL, false
put_long(&mut want, 6);
want.push(ty::NULL);
want.push(ty::NULL);
want.push(ty::BOOL);
want.push(0);
assert_eq!(blob, want);
}
#[test]
fn an_absolute_edges_payload_writes_the_node_ids_verbatim() {
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("src", DataType::Int64, false),
Field::new("dst", DataType::Int64, false),
Field::new("w", DataType::Int64, false),
])),
vec![
Arc::new(Int64Array::from(vec![7, 9])),
Arc::new(Int64Array::from(vec![1000, 3])),
Arc::new(Int64Array::from(vec![5, 6])),
],
)
.unwrap();
let g = BulkEdgeGroup { src_label: "A", edge_label: "T", dst_label: "B", batches: std::slice::from_ref(&batch) };
let p = build_bulk_payload_edges_absolute(&[g]).unwrap();
assert_eq!((p.node_count, p.relation_count, p.node_blob_len, p.blobs.len()), (0, 2, 0, 1));
let mut want = Vec::new();
put_cstr(&mut want, "T");
want.extend_from_slice(&1u32.to_le_bytes());
put_cstr(&mut want, "w");
for (s, d, w) in [(7u64, 1000u64, 5i64), (9, 3, 6)] {
want.extend_from_slice(&s.to_le_bytes());
want.extend_from_slice(&d.to_le_bytes());
put_long(&mut want, w);
}
assert_eq!(p.blobs[0], want, "ids as given - no label base offset added");
}
}