crates/graphar-flight/src/export.rs
//! Graph → lakehouse: export FalkorDB query results back to GraphAr files.
//!
//! The reverse of the ingest bridge. [`FalkorDbLoader`](graphar_falkordb) writes
//! Arrow into FalkorDB; this reads a Cypher query's result rows back out as
//! Arrow and writes them as a chunked GraphAr file (parquet/csv/orc/json/
//! arrow-ipc) — so a graph computed in FalkorDB can land back in the lakehouse
//! (and from there into Iceberg, object storage, etc.). Both directions of the
//! columnar ↔ graph bridge now exist.
use std::path::{Path, PathBuf};
use graphar::FileType;
use crate::error::Result;
use crate::falkor::FalkorExecutor;
/// Outcome of an [`export_cypher`] / [`export_cypher_chunked`] call.
#[derive(Debug, Clone, Default)]
pub struct ExportReport {
/// Rows written.
pub rows: usize,
/// Columns in the result schema.
pub columns: usize,
/// Chunk files written, in order (`chunk0.<ext>`, `chunk1.<ext>`, …).
///
/// Empty for [`export_cypher`] (single-file, single-path export); populated
/// by [`export_cypher_chunked`]. Additive — single-file callers can keep
/// ignoring it.
pub chunks: Vec<PathBuf>,
}
/// Run `cypher` against the executor's graph and write the result rows as one
/// GraphAr chunk file at `path` in `file_type`. Returns what was written.
///
/// The schema is inferred from FalkorDB's response (auto-string mode), so any
/// `RETURN` shape works without pre-registration. For very large results this
/// writes a single chunk; chunked/streamed export is a future extension.
///
/// ```no_run
/// # async fn demo() -> graphar_flight::Result<()> {
/// use graphar_flight::{FalkorExecutor, export::export_cypher};
/// use graphar::FileType;
///
/// let mut exec = FalkorExecutor::connect("redis://127.0.0.1:6379", "social").await?;
/// let report = export_cypher(
/// &mut exec,
/// "MATCH (n:Person) RETURN n._gar_id AS id, n.name AS name",
/// "/tmp/people.parquet",
/// FileType::Parquet,
/// ).await?;
/// println!("exported {} rows", report.rows);
/// # Ok(()) }
/// ```
pub async fn export_cypher(
executor: &mut FalkorExecutor,
cypher: &str,
path: impl AsRef<Path>,
file_type: FileType,
) -> Result<ExportReport> {
let batch = executor.query_auto(cypher).await?;
let rows = batch.num_rows();
let columns = batch.num_columns();
graphar::io::write_chunk(path, &[batch], &file_type)?;
Ok(ExportReport {
rows,
columns,
chunks: Vec::new(),
})
}
/// Run `cypher` and write the result as **multiple** GraphAr chunk files —
/// `chunk0.<ext>`, `chunk1.<ext>`, … — under `dir`, each holding at most
/// `chunk_rows` rows. This is the *chunked* export for huge results: instead of
/// one monolithic file it lands the result the way GraphAr naturally tiles a
/// vertex/edge collection (one file per row-range), so a `VertexReader` (or a
/// plain `read_chunk` per file) can stream the chunks back without ever holding
/// the whole result in memory.
///
/// Returns an [`ExportReport`] whose `chunks` lists the files written in order.
///
/// ## Streaming honesty — why this is slice-and-write, not pull-stream
///
/// FalkorDB's `GRAPH.QUERY` is a single RESP request/response: the server
/// serializes the **entire** result set into one reply with no server-side
/// cursor, and [`FalkorExecutor::query_auto`] therefore materializes it as one
/// in-memory [`RecordBatch`](arrow_array::RecordBatch). There is no API to pull
/// the result in pages, so we cannot bound the *read* side below the full
/// result. (FalkorDB ≥ 4 exposes `GRAPH.QUERY … TIMEOUT`/result-set limits but
/// no streaming cursor over RESP.)
///
/// What this *does* bound is the **write** side and downstream consumption: the
/// fetched batch is zero-copy-sliced into `chunk_rows`-sized record batches
/// (`RecordBatch::slice` shares the underlying Arrow buffers — no row copy), and
/// each slice is written and dropped before the next, so peak *encoder* memory
/// (parquet row-group buffers, CSV/JSON string buffers, etc.) is bounded by
/// `chunk_rows` rather than the whole result, and each output file is bounded
/// too. Reading back with `read_chunk` per file then streams chunk-by-chunk.
///
/// When the upstream is the *lakehouse* rather than FalkorDB, the Iceberg export
/// path ([`export_cypher_to_iceberg`], feature `skade`) already appends in
/// batches; this function is the GraphAr-file equivalent.
///
/// `chunk_rows` must be non-zero; `0` is treated as `1`. An empty result still
/// writes a single empty `chunk0.<ext>` so consumers always find a chunk file.
///
/// ```no_run
/// # async fn demo() -> graphar_flight::Result<()> {
/// use graphar_flight::{FalkorExecutor, export::export_cypher_chunked};
/// use graphar::FileType;
///
/// let mut exec = FalkorExecutor::connect("redis://127.0.0.1:6379", "social").await?;
/// let report = export_cypher_chunked(
/// &mut exec,
/// "MATCH (n:Person) RETURN n._gar_id AS id, n.name AS name",
/// "/tmp/people_chunks",
/// FileType::Parquet,
/// 100_000, // ≤ 100k rows per chunk file
/// ).await?;
/// println!("exported {} rows across {} chunks", report.rows, report.chunks.len());
/// # Ok(()) }
/// ```
pub async fn export_cypher_chunked(
executor: &mut FalkorExecutor,
cypher: &str,
dir: impl AsRef<Path>,
file_type: FileType,
chunk_rows: usize,
) -> Result<ExportReport> {
let batch = executor.query_auto(cypher).await?;
write_batch_chunked(&batch, dir, &file_type, chunk_rows)
}
/// Slice one in-memory [`RecordBatch`](arrow_array::RecordBatch) into
/// `chunk_rows`-sized record batches and write each as `chunk{i}.<ext>` under
/// `dir`. Split out from [`export_cypher_chunked`] so the chunking can be
/// exercised (round-tripped) without a FalkorDB connection.
///
/// The slicing is zero-copy ([`RecordBatch::slice`](arrow_array::RecordBatch::slice)
/// shares Arrow buffers); peak extra memory is one encoder's worth of
/// `chunk_rows` rows at a time.
pub fn write_batch_chunked(
batch: &arrow_array::RecordBatch,
dir: impl AsRef<Path>,
file_type: &FileType,
chunk_rows: usize,
) -> Result<ExportReport> {
// `graphar::io::write_chunk` creates each path's parent directory, so `dir`
// itself is created on the first write — no explicit `create_dir_all` (and
// no need for `FlightSqlError: From<io::Error>`).
let dir = dir.as_ref();
let ext = file_type.extension();
let total = batch.num_rows();
let columns = batch.num_columns();
let step = chunk_rows.max(1);
let mut chunks = Vec::new();
if total == 0 {
// Always emit one (empty) chunk so consumers find a chunk0 file.
let path = dir.join(format!("chunk0.{ext}"));
graphar::io::write_chunk(&path, std::slice::from_ref(batch), file_type)?;
chunks.push(path);
return Ok(ExportReport {
rows: 0,
columns,
chunks,
});
}
let mut offset = 0;
let mut idx = 0;
while offset < total {
let len = step.min(total - offset);
// Zero-copy view of rows [offset, offset+len) — shares the source
// buffers, so no row data is duplicated to form the slice.
let slice = batch.slice(offset, len);
let path = dir.join(format!("chunk{idx}.{ext}"));
graphar::io::write_chunk(&path, &[slice], file_type)?;
chunks.push(path);
offset += len;
idx += 1;
}
Ok(ExportReport {
rows: total,
columns,
chunks,
})
}
// ── Graph → Iceberg (feature `skade`) ───────────────────────────────────────
/// Export a Cypher query's results **straight into an Iceberg table** (the
/// graph → lakehouse, Iceberg edition). Runs `cypher`, then appends the rows to
/// `table_name` in the [`skade`] embedded Iceberg warehouse rooted at
/// `warehouse_dir`, creating the table on first write.
///
/// knut and skade now both build on arrow 59, so the result is handed to skade
/// directly ([`to_skade_batch`] is an identity clone) — no serialization, no
/// temp file, no FFI. (The parquet-bridge variant
/// [`append_arrow_to_iceberg_via_parquet`] is retained only for the comparison
/// bench.)
///
/// ```no_run
/// # #[cfg(feature = "skade")]
/// # async fn demo() -> graphar_flight::Result<()> {
/// use graphar_flight::{FalkorExecutor, export::export_cypher_to_iceberg};
/// let mut exec = FalkorExecutor::connect("redis://127.0.0.1:6379", "social").await?;
/// let report = export_cypher_to_iceberg(
/// &mut exec,
/// "MATCH (n:Person) RETURN n._gar_id AS id, n.name AS name",
/// "/tmp/knut_warehouse",
/// "people",
/// ).await?;
/// # let _ = report; Ok(()) }
/// ```
#[cfg(feature = "skade")]
pub async fn export_cypher_to_iceberg(
executor: &mut FalkorExecutor,
cypher: &str,
warehouse_dir: impl AsRef<Path>,
table_name: &str,
) -> Result<ExportReport> {
let batch = executor.query_auto(cypher).await?;
let rows = batch.num_rows();
let columns = batch.num_columns();
append_arrow_to_iceberg(warehouse_dir.as_ref(), table_name, batch).await?;
Ok(ExportReport {
rows,
columns,
chunks: Vec::new(),
})
}
/// Append one arrow-59 `RecordBatch` to an Iceberg table via skade.
///
/// knut and skade share arrow 59, so the batch is handed to skade directly
/// ([`to_skade_batch`] is an identity clone — an Arc bump, no buffer copy, no
/// FFI). Exposed (feature-gated) so the bridge can be exercised without a
/// FalkorDB connection.
#[cfg(feature = "skade")]
pub async fn append_arrow_to_iceberg(
warehouse_dir: &Path,
table_name: &str,
batch: arrow_array::RecordBatch,
) -> Result<usize> {
let rows = batch.num_rows();
if rows == 0 {
return Ok(0);
}
let skade_batch = to_skade_batch(&batch)?;
append_skade_batch(warehouse_dir, table_name, &skade_batch).await?;
Ok(rows)
}
/// Same as [`append_arrow_to_iceberg`] but bridges through a temp **parquet**
/// file instead of the direct identity hand-off. Kept so the `knut.skade_iceberg_export`
/// (zero-copy) and `knut.skade_iceberg_export_parquet` benchers can quantify the
/// difference; the zero-copy path is the default everywhere else.
#[cfg(feature = "skade")]
pub async fn append_arrow_to_iceberg_via_parquet(
warehouse_dir: &Path,
table_name: &str,
batch: arrow_array::RecordBatch,
) -> Result<usize> {
use crate::error::FlightSqlError;
let rows = batch.num_rows();
let tmp = tempfile::Builder::new()
.prefix("knut-iceberg-")
.suffix(".parquet")
.tempfile()
.map_err(|e| FlightSqlError::Skade(format!("temp file: {e}")))?;
graphar::io::write_chunk(tmp.path(), &[batch], &FileType::Parquet)?;
let skade_batches = read_parquet_as_skade(tmp.path())?;
drop(tmp);
let Some(first) = skade_batches.first() else {
return Ok(0);
};
let schema = skade::arrow_array::RecordBatch::schema(first);
let wh = skade::Warehouse::open(warehouse_dir)
.await
.map_err(|e| FlightSqlError::Skade(e.to_string()))?;
let mut table = wh
.table_or_create(table_name, &schema)
.await
.map_err(|e| FlightSqlError::Skade(e.to_string()))?;
table
.append(&skade_batches)
.await
.map_err(|e| FlightSqlError::Skade(e.to_string()))?;
Ok(rows)
}
/// Open the warehouse and append one arrow-59 batch to `table_name`.
#[cfg(feature = "skade")]
async fn append_skade_batch(
warehouse_dir: &Path,
table_name: &str,
skade_batch: &skade::arrow_array::RecordBatch,
) -> Result<()> {
use crate::error::FlightSqlError;
let schema = skade::arrow_array::RecordBatch::schema(skade_batch);
let wh = skade::Warehouse::open(warehouse_dir)
.await
.map_err(|e| FlightSqlError::Skade(e.to_string()))?;
let mut table = wh
.table_or_create(table_name, &schema)
.await
.map_err(|e| FlightSqlError::Skade(e.to_string()))?;
table
.append(std::slice::from_ref(skade_batch))
.await
.map_err(|e| FlightSqlError::Skade(e.to_string()))?;
Ok(())
}
/// arrow → skade record-batch bridge — now an **identity clone**.
///
/// knut and skade both build on Apache Arrow 59, so
/// `skade::arrow_array::RecordBatch` *is* `arrow_array::RecordBatch` (the same
/// crate, same major). The old C-Data-Interface FFI transmute bridge
/// (arrow 56 → arrow 57) is therefore gone: a plain `clone()` (Arc bumps, no
/// buffer copy) hands skade the batch directly. Kept as a named seam so the
/// export path and its round-trip test read the same.
#[cfg(feature = "skade")]
pub fn to_skade_batch(batch: &arrow_array::RecordBatch) -> Result<skade::arrow_array::RecordBatch> {
Ok(batch.clone())
}
/// Read a parquet file as skade's (arrow-59) record batches.
#[cfg(feature = "skade")]
fn read_parquet_as_skade(path: &Path) -> Result<Vec<skade::arrow_array::RecordBatch>> {
use crate::error::FlightSqlError;
use skade::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
let file =
std::fs::File::open(path).map_err(|e| FlightSqlError::Skade(format!("open: {e}")))?;
let reader = ParquetRecordBatchReaderBuilder::try_new(file)
.map_err(|e| FlightSqlError::Skade(format!("parquet: {e}")))?
.build()
.map_err(|e| FlightSqlError::Skade(format!("parquet: {e}")))?;
let mut out = Vec::new();
for b in reader {
out.push(b.map_err(|e| FlightSqlError::Skade(format!("parquet read: {e}")))?);
}
Ok(out)
}
#[cfg(test)]
mod chunk_tests {
use super::*;
use std::sync::Arc;
use arrow_array::{Array, Int64Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema};
/// A 7-row batch (id 0..6, name "n0".."n6").
fn sample_batch() -> RecordBatch {
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
])),
vec![
Arc::new(Int64Array::from((0..7).collect::<Vec<i64>>())),
Arc::new(StringArray::from(
(0..7).map(|i| format!("n{i}")).collect::<Vec<_>>(),
)),
],
)
.unwrap()
}
#[test]
fn chunked_export_round_trips_and_files_are_bounded() {
// 7 rows, chunk_rows = 3 → ceil(7/3) = 3 files: 3 + 3 + 1 rows.
let batch = sample_batch();
let dir = tempfile::tempdir().unwrap();
let report = write_batch_chunked(&batch, dir.path(), &FileType::Parquet, 3).unwrap();
assert_eq!(report.rows, 7);
assert_eq!(report.columns, 2);
assert_eq!(report.chunks.len(), 3, "ceil(7/3) chunk files");
// Each file exists, is named chunk{i}.parquet, and holds ≤ chunk_rows rows.
let mut total = 0usize;
let mut all_ids = Vec::new();
let mut all_names = Vec::new();
for (i, path) in report.chunks.iter().enumerate() {
assert_eq!(
path.file_name().unwrap().to_str().unwrap(),
format!("chunk{i}.parquet"),
);
let batches = graphar::io::read_chunk(path, &FileType::Parquet).unwrap();
let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
assert!(rows <= 3, "chunk {i} bounded by chunk_rows (got {rows})");
total += rows;
for b in &batches {
let ids = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let names = b
.column_by_name("name")
.unwrap()
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
for r in 0..b.num_rows() {
all_ids.push(ids.value(r));
all_names.push(names.value(r).to_string());
}
}
}
// Round trip: every input row came back exactly once, in order.
assert_eq!(total, 7, "total rows read back == input rows");
assert_eq!(all_ids, (0..7).collect::<Vec<i64>>());
assert_eq!(
all_names,
(0..7).map(|i| format!("n{i}")).collect::<Vec<String>>(),
);
}
#[test]
fn chunk_rows_at_least_total_writes_single_file() {
let batch = sample_batch();
let dir = tempfile::tempdir().unwrap();
let report = write_batch_chunked(&batch, dir.path(), &FileType::Parquet, 100).unwrap();
assert_eq!(report.chunks.len(), 1, "one chunk when chunk_rows >= total");
assert_eq!(report.rows, 7);
}
#[test]
fn empty_result_writes_one_empty_chunk_preserving_schema() {
let empty = RecordBatch::new_empty(sample_batch().schema());
let dir = tempfile::tempdir().unwrap();
let report = write_batch_chunked(&empty, dir.path(), &FileType::Parquet, 3).unwrap();
assert_eq!(report.rows, 0);
assert_eq!(report.chunks.len(), 1, "empty result still emits chunk0");
let batches = graphar::io::read_chunk(&report.chunks[0], &FileType::Parquet).unwrap();
let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(rows, 0);
// The file exists and is readable (an empty parquet yields zero record
// batches on read-back; the point is the chunk0 file is always present
// so a consumer never hits a missing-file error for an empty result).
assert!(report.chunks[0].exists());
}
#[test]
fn chunk_rows_zero_is_treated_as_one() {
let batch = sample_batch();
let dir = tempfile::tempdir().unwrap();
let report = write_batch_chunked(&batch, dir.path(), &FileType::Parquet, 0).unwrap();
assert_eq!(report.chunks.len(), 7, "chunk_rows 0 → 1 row per file");
}
}
#[cfg(all(test, feature = "skade"))]
mod skade_tests {
use super::*;
use std::sync::Arc;
#[test]
fn ffi_bridge_preserves_values_and_nulls() {
// arrow-59 batch with three types incl. a null → skade via the now
// identity bridge (knut + skade share arrow 59) → assert every value
// (and the null) survives the hand-off. Reads the result back through
// skade's typed arrays, which are the same arrow-59 types.
use arrow_array::{Float64Array, Int64Array, StringArray};
use arrow_schema::{DataType, Field, Schema};
let ws_batch = arrow_array::RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("score", DataType::Float64, false),
Field::new("name", DataType::Utf8, true),
])),
vec![
Arc::new(Int64Array::from(vec![10, 20, 30])),
Arc::new(Float64Array::from(vec![1.5, 2.5, 3.5])),
Arc::new(StringArray::from(vec![Some("a"), None, Some("c")])),
],
)
.unwrap();
let skade_batch = to_skade_batch(&ws_batch).unwrap();
assert_eq!(skade_batch.num_rows(), 3);
assert_eq!(skade_batch.num_columns(), 3);
let ids = skade_batch
.column(0)
.as_any()
.downcast_ref::<skade::arrow_array::Int64Array>()
.unwrap();
assert_eq!(ids.values(), &[10, 20, 30]);
let scores = skade_batch
.column(1)
.as_any()
.downcast_ref::<skade::arrow_array::Float64Array>()
.unwrap();
assert_eq!(scores.values(), &[1.5, 2.5, 3.5]);
use skade::arrow_array::Array;
let names = skade_batch
.column(2)
.as_any()
.downcast_ref::<skade::arrow_array::StringArray>()
.unwrap();
assert_eq!(names.value(0), "a");
assert!(names.is_null(1), "the null survived the bridge");
assert_eq!(names.value(2), "c");
}
#[tokio::test]
async fn append_and_read_back_via_iceberg() {
// Build an arrow-59 batch (no FalkorDB), append to a fresh warehouse,
// then read it back through skade — exercises the (now identity-clone)
// arrow bridge and the skade write/read path with no container.
use arrow_array::{Int64Array, StringArray};
use arrow_schema::{DataType, Field, Schema};
let batch = arrow_array::RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
])),
vec![
Arc::new(Int64Array::from(vec![1, 2, 3])),
Arc::new(StringArray::from(vec!["a", "b", "c"])),
],
)
.unwrap();
let dir = tempfile::tempdir().unwrap();
let n = append_arrow_to_iceberg(dir.path(), "people", batch)
.await
.unwrap();
assert_eq!(n, 3);
// Read it back from the Iceberg table.
let wh = skade::Warehouse::open(dir.path()).await.unwrap();
let table = wh.table("people").await.unwrap();
let count = table.count().await.unwrap();
assert_eq!(count, 3, "3 rows landed in the Iceberg table");
}
}