graphar_flight/

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 58, so the result is handed to skade
/// directly ([`bridge_56_to_57`] 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-58 `RecordBatch` to an Iceberg table via skade.
///
/// knut and skade share arrow 58, so the batch is handed to skade directly
/// ([`bridge_56_to_57`] 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 batch57 = bridge_56_to_57(&batch)?;
    append_batch57(warehouse_dir, table_name, &batch57).await?;
    Ok(rows)
}

/// Same as [`append_arrow_to_iceberg`] but bridges through a temp **parquet**
/// file instead of the C Data Interface. 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 batches57 = read_parquet_as_skade(tmp.path())?;
    drop(tmp);
    let Some(first) = batches57.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(&batches57)
        .await
        .map_err(|e| FlightSqlError::Skade(e.to_string()))?;
    Ok(rows)
}

/// Open the warehouse and append one arrow-57 batch to `table_name`.
#[cfg(feature = "skade")]
async fn append_batch57(
    warehouse_dir: &Path,
    table_name: &str,
    batch57: &skade::arrow_array::RecordBatch,
) -> Result<()> {
    use crate::error::FlightSqlError;
    let schema = skade::arrow_array::RecordBatch::schema(batch57);
    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(batch57))
        .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 58, 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 bridge_56_to_57(
    batch: &arrow_array::RecordBatch,
) -> Result<skade::arrow_array::RecordBatch> {
    Ok(batch.clone())
}

/// Read a parquet file as skade's (arrow-57) 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-58 batch with three types incl. a null → skade via the now
        // identity bridge (knut + skade share arrow 58) → assert every value
        // (and the null) survives the hand-off. Reads the result back through
        // skade's typed arrays, which are the same arrow-58 types.
        use arrow_array::{Float64Array, Int64Array, StringArray};
        use arrow_schema::{DataType, Field, Schema};

        let batch56 = 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 batch57 = bridge_56_to_57(&batch56).unwrap();
        assert_eq!(batch57.num_rows(), 3);
        assert_eq!(batch57.num_columns(), 3);

        let ids = batch57
            .column(0)
            .as_any()
            .downcast_ref::<skade::arrow_array::Int64Array>()
            .unwrap();
        assert_eq!(ids.values(), &[10, 20, 30]);

        let scores = batch57
            .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 = batch57
            .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-58 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");
    }
}