Skip to main content

hyperdb_mcp/
ingest_arrow.rs

1// Copyright (c) 2026, Salesforce, Inc. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0 OR MIT
3
4//! Ingest Parquet and Arrow IPC files into Hyper tables.
5//!
6//! These are "Tier 1" (exact schema) ingest paths: types come directly from
7//! the file's embedded schema, so there is no guessing.
8//!
9//! # Parquet path: native hyperd read
10//!
11//! For Parquet files we let hyperd read the file directly via the
12//! `external(path, FORMAT => 'parquet')` table function. The MCP issues a
13//! single `CREATE TABLE ... AS SELECT * FROM external(...)` (replace mode)
14//! or `INSERT INTO ... SELECT * FROM external(...)` (append mode). This
15//! pushes the entire read into hyperd's C++ Parquet reader — no Rust-side
16//! Arrow decode, no per-row INSERTs, and native support for the full set of
17//! compressions and encodings hyperd understands (including Snappy, which
18//! isn't compiled into the Rust `parquet` crate in this build).
19//!
20//! Benchmarked at ~3.5M rows/sec on TPCH 1GB lineitem (6M rows / 360 MB
21//! in ~1.7s) — roughly three orders of magnitude faster than the previous
22//! row-by-row path.
23//!
24//! We still open the Parquet footer in Rust to infer column types — that's
25//! cheap (footer-only metadata read) and lets us report the schema in
26//! [`IngestResult`] and apply user-supplied schema overrides by wrapping
27//! each overridden column in an explicit `::TYPE` cast inside the SELECT
28//! projection.
29//!
30//! # Arrow IPC path: binary COPY via `ArrowInserter`
31//!
32//! For Arrow IPC (Feather) files we use hyperdb-api's `ArrowInserter` /
33//! `AsyncArrowInserter`, which push `RecordBatch`es straight through
34//! Hyper's binary COPY protocol as an Arrow IPC Stream. No text encoding,
35//! no per-row SQL — the batches we decode from the IPC File format are
36//! re-framed as IPC Stream messages and sent wholesale to the server.
37//!
38//! User-supplied schema overrides are applied by pre-creating the target
39//! table with the overridden column types and letting Hyper coerce the
40//! incoming values on insert. This avoids a second Rust-side transform
41//! pass over the data.
42
43use crate::engine::Engine;
44use crate::error::{ErrorCode, McpError};
45use crate::ingest::{IngestOptions, IngestResult};
46use crate::schema::ColumnSchema;
47use crate::stats::{IngestStats, StatsTimer};
48use arrow::datatypes::{DataType, Schema as ArrowSchema};
49use arrow::record_batch::RecordBatch;
50use hyperdb_api::AsyncConnection;
51use std::path::Path;
52
53/// Map an Arrow [`DataType`] to the corresponding Hyper SQL type name.
54///
55/// Unsigned integers are promoted to the next wider signed type (e.g.
56/// `UInt32` → `BIGINT`) because Hyper has no unsigned integer types.
57/// Unrecognized types fall back to `TEXT`.
58fn arrow_type_to_hyper(dt: &DataType) -> String {
59    match dt {
60        DataType::Boolean => "BOOL".into(),
61        DataType::Int8 | DataType::Int16 => "SMALLINT".into(),
62        DataType::Int32 | DataType::UInt16 => "INT".into(),
63        DataType::Int64 | DataType::UInt32 => "BIGINT".into(),
64        DataType::UInt64 => "BIGINT".into(),
65        DataType::Float16 | DataType::Float32 => "DOUBLE PRECISION".into(),
66        DataType::Float64 => "DOUBLE PRECISION".into(),
67        DataType::Utf8 | DataType::LargeUtf8 => "TEXT".into(),
68        DataType::Binary | DataType::LargeBinary => "BYTEA".into(),
69        DataType::Date32 | DataType::Date64 => "DATE".into(),
70        DataType::Time32(_) | DataType::Time64(_) => "TIME".into(),
71        DataType::Timestamp(_, None) => "TIMESTAMP".into(),
72        DataType::Timestamp(_, Some(_)) => "TIMESTAMPTZ".into(),
73        DataType::Decimal128(p, s) | DataType::Decimal256(p, s) => {
74            format!("NUMERIC({p}, {s})")
75        }
76        _ => "TEXT".into(),
77    }
78}
79
80/// Convert an Arrow schema to our internal [`ColumnSchema`] representation,
81/// preserving nullability from the Arrow field metadata.
82#[must_use]
83pub fn arrow_schema_to_columns(schema: &ArrowSchema) -> Vec<ColumnSchema> {
84    schema
85        .fields()
86        .iter()
87        .map(|f| ColumnSchema {
88            name: f.name().clone(),
89            hyper_type: arrow_type_to_hyper(f.data_type()),
90            nullable: f.is_nullable(),
91        })
92        .collect()
93}
94
95/// Serialize a slice of Arrow `RecordBatch`es to a single Arrow IPC Stream
96/// byte buffer. Used to feed the async inserter, which accepts raw IPC
97/// Stream bytes (not `RecordBatch` objects). A well-formed stream is
98/// schema message + one or more batch messages, in that order.
99fn record_batches_to_ipc_stream(batches: &[RecordBatch]) -> Result<Vec<u8>, McpError> {
100    let schema = batches
101        .first()
102        .map(arrow::array::RecordBatch::schema)
103        .ok_or_else(|| McpError::new(ErrorCode::EmptyData, "Arrow IPC file has no batches"))?;
104    let mut buf = Vec::new();
105    {
106        let mut writer =
107            arrow::ipc::writer::StreamWriter::try_new(&mut buf, &schema).map_err(|e| {
108                McpError::new(
109                    ErrorCode::InternalError,
110                    format!("Failed to create Arrow IPC StreamWriter: {e}"),
111                )
112            })?;
113        for batch in batches {
114            writer.write(batch).map_err(|e| {
115                McpError::new(
116                    ErrorCode::InternalError,
117                    format!("Failed to write Arrow batch: {e}"),
118                )
119            })?;
120        }
121        writer.finish().map_err(|e| {
122            McpError::new(
123                ErrorCode::InternalError,
124                format!("Failed to finish Arrow IPC stream: {e}"),
125            )
126        })?;
127    }
128    Ok(buf)
129}
130
131/// Read only the Parquet footer and return the inferred column list.
132///
133/// Hyperd will re-read the file when ingesting, so this pass is just for
134/// reporting the schema back to the caller and resolving schema overrides
135/// against real column names. It's cheap — the parquet crate only loads
136/// footer metadata, not row groups.
137fn infer_parquet_schema(path: &str) -> Result<Vec<ColumnSchema>, McpError> {
138    let file = std::fs::File::open(path)
139        .map_err(|e| McpError::new(ErrorCode::FileNotFound, format!("Cannot open file: {e}")))?;
140    let reader = parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder::try_new(file)
141        .map_err(|e| {
142            McpError::new(
143                ErrorCode::UnsupportedFormat,
144                format!("Invalid Parquet file: {e}"),
145            )
146        })?;
147    let arrow_schema = reader.schema();
148    Ok(arrow_schema_to_columns(arrow_schema))
149}
150
151/// Build the projection list for `SELECT <projection> FROM external(...)`.
152///
153/// When a user supplies a schema override, the overridden columns are
154/// wrapped in a `::TYPE AS name` cast so the target table receives the
155/// narrowed/widened type. Non-overridden columns pass through by name.
156/// Column names are quoted to tolerate case, spaces, and reserved words.
157fn parquet_select_projection(inferred: &[ColumnSchema], final_columns: &[ColumnSchema]) -> String {
158    inferred
159        .iter()
160        .zip(final_columns.iter())
161        .map(|(orig, col)| {
162            let quoted = format!("\"{}\"", col.name.replace('"', "\"\""));
163            if orig.hyper_type == col.hyper_type {
164                quoted
165            } else {
166                format!("{quoted}::{ty} AS {quoted}", ty = col.hyper_type)
167            }
168        })
169        .collect::<Vec<_>>()
170        .join(", ")
171}
172
173/// Build the SQL statement hyperd will run to ingest the Parquet file.
174///
175/// - Replace mode: `CREATE TABLE "t" AS SELECT <proj> FROM external(...)`.
176/// - Append mode: `INSERT INTO "t" SELECT <proj> FROM external(...)`.
177///
178/// The caller is responsible for issuing a preceding `DROP TABLE IF EXISTS`
179/// in replace mode — hyperd's `CREATE TABLE AS` fails rather than replacing.
180fn build_parquet_ingest_sql(
181    table: &str,
182    path: &str,
183    projection: &str,
184    is_replace: bool,
185    target_db: Option<&str>,
186) -> String {
187    let quoted_table = match target_db {
188        Some(db) => {
189            let esc_db = db.replace('"', "\"\"");
190            let esc_tbl = table.replace('"', "\"\"");
191            format!("\"{esc_db}\".\"public\".\"{esc_tbl}\"")
192        }
193        None => format!("\"{}\"", table.replace('"', "\"\"")),
194    };
195    let quoted_path = hyperdb_api::escape_string_literal(path);
196    if is_replace {
197        format!(
198            "CREATE TABLE {quoted_table} AS SELECT {projection} FROM external({quoted_path}, FORMAT => 'parquet')"
199        )
200    } else {
201        format!(
202            "INSERT INTO {quoted_table} SELECT {projection} FROM external({quoted_path}, FORMAT => 'parquet')"
203        )
204    }
205}
206
207/// Resolve the absolute ingest path, verifying the file exists.
208///
209/// Hyperd reads the Parquet file directly, so the path it receives must be
210/// unambiguous — we canonicalize it here rather than passing whatever the
211/// caller provided. A non-existent file fails fast with a clear error
212/// before we ever reach the server.
213fn resolve_parquet_path(path: &str) -> Result<(String, u64), McpError> {
214    let file_path = Path::new(path);
215    if !file_path.exists() {
216        return Err(McpError::new(
217            ErrorCode::FileNotFound,
218            format!("File not found: {path}"),
219        ));
220    }
221    let absolute = std::fs::canonicalize(file_path)
222        .map_err(|e| {
223            McpError::new(
224                ErrorCode::FileNotFound,
225                format!("Cannot resolve path {path}: {e}"),
226            )
227        })?
228        .to_string_lossy()
229        .into_owned();
230    let file_size = std::fs::metadata(file_path).map_or(0, |m| m.len());
231    Ok((absolute, file_size))
232}
233
234/// Count rows in a newly-created or freshly-appended table. Used for
235/// reporting after `CREATE TABLE AS`, which doesn't surface an affected
236/// row count the way `INSERT` does.
237fn count_rows_sync(engine: &Engine, table: &str) -> Result<u64, McpError> {
238    let quoted = format!("\"{}\"", table.replace('"', "\"\""));
239    let sql = format!("SELECT COUNT(*) FROM {quoted}");
240    let rows = engine.execute_query_to_json(&sql)?;
241    rows.first()
242        .and_then(|r| r.get("count"))
243        .and_then(serde_json::Value::as_u64)
244        .ok_or_else(|| {
245            McpError::new(
246                ErrorCode::InternalError,
247                "Could not read row count after parquet ingest",
248            )
249        })
250}
251
252/// Async twin of `count_rows_sync`.
253async fn count_rows_async(conn: &AsyncConnection, table: &str) -> Result<u64, McpError> {
254    let quoted = format!("\"{}\"", table.replace('"', "\"\""));
255    let sql = format!("SELECT COUNT(*) FROM {quoted}");
256    let row = conn.fetch_one(&sql).await.map_err(McpError::from)?;
257    let count: i64 = row.get(0).ok_or_else(|| {
258        McpError::new(
259            ErrorCode::InternalError,
260            "Could not read row count after parquet ingest",
261        )
262    })?;
263    // A negative row count from COUNT(*) indicates a catalog inconsistency;
264    // fall back to 0 rather than wrapping into a huge u64.
265    Ok(u64::try_from(count).unwrap_or(0))
266}
267
268/// Ingest a Parquet file into a Hyper table using hyperd's native parquet
269/// reader via `CREATE TABLE ... AS SELECT * FROM external(...)`.
270///
271/// This is one SQL round-trip; hyperd opens the file itself and writes the
272/// data straight into the target table. No Rust-side Arrow decode, no
273/// per-row INSERTs.
274///
275/// # Errors
276///
277/// - Propagates errors from `resolve_parquet_path` when `path` is
278///   missing or not canonicalizable.
279/// - Propagates errors from `infer_parquet_schema` (malformed footer)
280///   or [`crate::schema::apply_schema_override`].
281/// - Propagates any transaction error from the `CREATE TABLE AS
282///   SELECT` / `INSERT INTO ... SELECT` statement against hyperd.
283/// - Propagates any error from the post-ingest `COUNT(*)` in
284///   `count_rows_sync` when running in replace mode.
285pub fn ingest_parquet_file(
286    engine: &Engine,
287    path: &str,
288    opts: &IngestOptions,
289) -> Result<IngestResult, McpError> {
290    if opts.mode == "merge" {
291        return crate::ingest::merge_via_temp_table(engine, opts, |tmp_opts| {
292            ingest_parquet_file(engine, path, tmp_opts)
293        });
294    }
295    let timer = StatsTimer::start();
296    let (absolute_path, file_size) = resolve_parquet_path(path)?;
297
298    let inferred = infer_parquet_schema(&absolute_path)?;
299    let final_columns = match &opts.schema_override {
300        Some(s) => crate::schema::apply_schema_override(inferred.clone(), s)?,
301        None => inferred.clone(),
302    };
303    let projection = parquet_select_projection(&inferred, &final_columns);
304
305    let is_replace = opts.mode != "append";
306    let sql = build_parquet_ingest_sql(
307        &opts.table,
308        &absolute_path,
309        &projection,
310        is_replace,
311        opts.target_db.as_deref(),
312    );
313
314    // Issue the ingest DDL/DML. `CREATE TABLE ... AS SELECT` auto-commits
315    // (Hyper treats all DDL that way), so wrapping it in `execute_in_transaction`
316    // no longer buys us rollback — but it still gives us a clean error
317    // path that runs the transaction prelude + drops if needed.
318    let affected = engine.execute_in_transaction(|engine| {
319        if is_replace {
320            let qualified = crate::ingest::qualified_table(opts);
321            engine.execute_command(&format!("DROP TABLE IF EXISTS {qualified}"))?;
322        }
323        engine.execute_command(&sql)
324    })?;
325
326    // Row count: `CREATE TABLE AS` reports 0 affected, so for replace mode
327    // we have to COUNT(*). Critically, the COUNT must run *outside* the
328    // transaction above — when issued immediately after CTAS on the same
329    // connection inside a transaction we've observed the count coming back
330    // truncated (the low 17 bits of the true value). Running it as a fresh
331    // statement avoids whatever wire-state quirk is at play. For append
332    // mode `INSERT INTO ... SELECT` does report affected rows, so we use
333    // those directly.
334    let row_count = if is_replace {
335        count_rows_sync(engine, &opts.table)?
336    } else {
337        affected
338    };
339
340    let elapsed = timer.elapsed_ms();
341    let stats = IngestStats {
342        operation: "load_file".into(),
343        rows: row_count,
344        elapsed_ms: elapsed,
345        bytes_read: file_size,
346        bytes_stored: 0,
347        schema_inference_ms: Some(0),
348        table: opts.table.clone(),
349        file_format: Some("parquet".into()),
350        warning: None,
351        schema_changed: false,
352    };
353
354    Ok(IngestResult {
355        rows: row_count,
356        schema: final_columns,
357        stats,
358    })
359}
360
361/// Async twin of [`ingest_parquet_file`]. Uses an owned [`AsyncConnection`]
362/// so it can run on a pooled connection from the watcher without touching
363/// the engine's primary sync connection. Footer metadata is read on the
364/// blocking pool (tiny read, but file I/O).
365///
366/// # Errors
367///
368/// - Returns [`ErrorCode::InternalError`] if the spawn-blocking task
369///   panics (surfaced as a join error).
370/// - Propagates the same errors as [`ingest_parquet_file`]:
371///   path resolution, schema inference, transaction/ingest, and row
372///   counting failures.
373pub async fn ingest_parquet_file_async(
374    conn: &mut AsyncConnection,
375    path: &str,
376    opts: &IngestOptions,
377) -> Result<IngestResult, McpError> {
378    let timer = StatsTimer::start();
379    let (absolute_path, file_size) = resolve_parquet_path(path)?;
380
381    let path_for_infer = absolute_path.clone();
382    let override_owned = opts.schema_override.clone();
383    let (inferred, final_columns): (Vec<ColumnSchema>, Vec<ColumnSchema>) =
384        tokio::task::spawn_blocking(move || -> Result<_, McpError> {
385            let inferred = infer_parquet_schema(&path_for_infer)?;
386            let final_columns = match &override_owned {
387                Some(s) => crate::schema::apply_schema_override(inferred.clone(), s)?,
388                None => inferred.clone(),
389            };
390            Ok((inferred, final_columns))
391        })
392        .await
393        .map_err(|e| McpError::new(ErrorCode::InternalError, format!("Task join error: {e}")))??;
394
395    let projection = parquet_select_projection(&inferred, &final_columns);
396    let is_replace = opts.mode != "append";
397    let sql = build_parquet_ingest_sql(
398        &opts.table,
399        &absolute_path,
400        &projection,
401        is_replace,
402        opts.target_db.as_deref(),
403    );
404
405    let txn = conn.transaction().await.map_err(McpError::from)?;
406    let result: Result<u64, McpError> = async {
407        if is_replace {
408            let qualified = crate::ingest::qualified_table(opts);
409            txn.execute_command(&format!("DROP TABLE IF EXISTS {qualified}"))
410                .await
411                .map_err(McpError::from)?;
412        }
413        txn.execute_command(&sql).await.map_err(McpError::from)
414    }
415    .await;
416
417    let affected = match result {
418        Ok(n) => {
419            txn.commit().await.map_err(McpError::from)?;
420            n
421        }
422        Err(e) => {
423            if let Err(rb) = txn.rollback().await {
424                tracing::warn!("rollback after error failed: {}", rb);
425            }
426            return Err(e);
427        }
428    };
429
430    // The transaction was committed above (or the function returned on
431    // error). This COUNT(*) runs on the bare connection — outside any
432    // transaction — to mirror the sync path's source-of-truth behavior
433    // and avoid any post-CTAS wire-state quirks that the sync path
434    // documents.
435    let row_count = if is_replace {
436        count_rows_async(conn, &opts.table).await?
437    } else {
438        affected
439    };
440
441    let elapsed = timer.elapsed_ms();
442    let stats = IngestStats {
443        operation: "load_file".into(),
444        rows: row_count,
445        elapsed_ms: elapsed,
446        bytes_read: file_size,
447        bytes_stored: 0,
448        schema_inference_ms: Some(0),
449        table: opts.table.clone(),
450        file_format: Some("parquet".into()),
451        warning: None,
452        schema_changed: false,
453    };
454
455    Ok(IngestResult {
456        rows: row_count,
457        schema: final_columns,
458        stats,
459    })
460}
461
462/// Reject schema overrides on the Arrow IPC ingest path.
463///
464/// The binary COPY Arrow protocol is schema-strict: [`hyperdb_api::ArrowInserter`]
465/// validates that each incoming column's Arrow type matches the target
466/// column's Hyper type exactly, and Hyper rejects any mismatch with
467/// SQLSTATE 42804. Since Arrow IPC files carry exact typed schemas by
468/// construction, an override would almost always be a user error — so we
469/// surface it with a clear message rather than failing later at Hyper's
470/// type check. If a genuine override is ever needed here, add a
471/// Rust-side `arrow::compute::cast` pass over the batches.
472fn reject_ipc_schema_override(opts: &IngestOptions) -> Result<(), McpError> {
473    if opts.schema_override.is_some() {
474        return Err(McpError::new(
475            ErrorCode::SchemaMismatch,
476            "Schema overrides are not supported for Arrow IPC files. \
477             The embedded Arrow schema is authoritative on this path \
478             because the binary COPY protocol requires an exact type \
479             match between the file and the target table.",
480        ));
481    }
482    Ok(())
483}
484
485/// Magic prefix written at the start of an Arrow IPC **File** format
486/// (Feather v2). See the Arrow spec §IPC File Format. The Stream format
487/// does not start with this marker.
488const ARROW_IPC_FILE_MAGIC: &[u8] = b"ARROW1";
489
490/// Open an Arrow IPC file (either File/Feather format or Stream format)
491/// and return the inferred column schema plus the decoded batches. Used
492/// by both the sync and async ingest paths.
493///
494/// We auto-detect the sub-format by peeking the first 6 bytes: the File
495/// format opens with `ARROW1`, the Stream format does not. This matters
496/// because hyperdb-mcp's own export path emits raw IPC Stream bytes —
497/// the same shape Hyper's wire protocol speaks — so round-tripping
498/// `export → load_file` would otherwise fail on the Stream side with
499/// "Arrow file does not contain correct footer". Accepting both also
500/// matches what external producers emit: pyarrow's `new_file` writes
501/// File format, `new_stream` writes Stream format, and both are common.
502fn read_arrow_ipc_file(path: &str) -> Result<(Vec<ColumnSchema>, Vec<RecordBatch>), McpError> {
503    use std::io::Read;
504
505    let mut file = std::fs::File::open(path)
506        .map_err(|e| McpError::new(ErrorCode::FileNotFound, format!("Cannot open file: {e}")))?;
507
508    let mut magic = [0u8; 6];
509    let read = file.read(&mut magic).map_err(|e| {
510        McpError::new(
511            ErrorCode::UnsupportedFormat,
512            format!("Cannot read Arrow IPC header: {e}"),
513        )
514    })?;
515    // Rewind so whichever reader we pick sees the full file.
516    use std::io::Seek;
517    file.rewind().map_err(|e| {
518        McpError::new(
519            ErrorCode::InternalError,
520            format!("Cannot rewind file handle: {e}"),
521        )
522    })?;
523
524    let is_file_format = read == 6 && magic == ARROW_IPC_FILE_MAGIC;
525
526    if is_file_format {
527        let reader = arrow::ipc::reader::FileReader::try_new(file, None).map_err(|e| {
528            McpError::new(
529                ErrorCode::UnsupportedFormat,
530                format!("Invalid Arrow IPC file: {e}"),
531            )
532        })?;
533        let inferred = arrow_schema_to_columns(&reader.schema());
534        let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>().map_err(|e| {
535            McpError::new(
536                ErrorCode::InternalError,
537                format!("Arrow IPC read error: {e}"),
538            )
539        })?;
540        Ok((inferred, batches))
541    } else {
542        let reader = arrow::ipc::reader::StreamReader::try_new(file, None).map_err(|e| {
543            McpError::new(
544                ErrorCode::UnsupportedFormat,
545                format!("Invalid Arrow IPC stream: {e}"),
546            )
547        })?;
548        let inferred = arrow_schema_to_columns(&reader.schema());
549        let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>().map_err(|e| {
550            McpError::new(
551                ErrorCode::InternalError,
552                format!("Arrow IPC read error: {e}"),
553            )
554        })?;
555        Ok((inferred, batches))
556    }
557}
558
559/// Ingest an Arrow IPC (Feather) file into a Hyper table via the binary
560/// COPY protocol.
561///
562/// The file is read with `FileReader` to get `RecordBatch`es, then streamed
563/// through [`hyperdb_api::ArrowInserter`] — which re-frames each batch as an
564/// Arrow IPC Stream message and pushes it over COPY. No per-row SQL, no
565/// text encoding.
566///
567/// Schema overrides are rejected on this path — see
568/// `reject_ipc_schema_override` for the reasoning.
569///
570/// # Errors
571///
572/// - Propagates [`ErrorCode::InvalidArgument`] from
573///   `reject_ipc_schema_override` when `opts.schema_override` is set.
574/// - Returns [`ErrorCode::FileNotFound`] if `path` does not exist.
575/// - Returns [`ErrorCode::UnsupportedFormat`] or
576///   [`ErrorCode::InternalError`] when the file cannot be parsed as an
577///   Arrow IPC file or stream (see `read_arrow_ipc_file`).
578/// - Propagates transaction errors from [`Engine::create_table`] and
579///   from [`hyperdb_api::ArrowInserter`] operations (COPY setup, batch
580///   insert, or execute).
581pub fn ingest_arrow_ipc_file(
582    engine: &Engine,
583    path: &str,
584    opts: &IngestOptions,
585) -> Result<IngestResult, McpError> {
586    if opts.mode == "merge" {
587        return crate::ingest::merge_via_temp_table(engine, opts, |tmp_opts| {
588            ingest_arrow_ipc_file(engine, path, tmp_opts)
589        });
590    }
591    let timer = StatsTimer::start();
592    reject_ipc_schema_override(opts)?;
593
594    let file_path = Path::new(path);
595    if !file_path.exists() {
596        return Err(McpError::new(
597            ErrorCode::FileNotFound,
598            format!("File not found: {path}"),
599        ));
600    }
601    let file_size = std::fs::metadata(file_path).map_or(0, |m| m.len());
602
603    let (columns, batches) = read_arrow_ipc_file(path)?;
604
605    let is_replace = opts.mode != "append";
606    // Arrow IPC uses the binary COPY protocol which resolves table names via
607    // the search path. When targeting a non-primary database, temporarily
608    // redirect the search path for the duration of the transaction.
609    let _search_guard = if let Some(ref db) = opts.target_db {
610        Some(engine.scoped_search_path(db)?)
611    } else {
612        None
613    };
614    let row_count = engine.execute_in_transaction(|engine| {
615        engine.create_table_in(&opts.table, &columns, is_replace, opts.target_db.as_deref())?;
616
617        // Stream RecordBatches through the binary COPY protocol. Each
618        // batch is written to an IPC Stream segment internally — no
619        // text encoding, no per-row SQL.
620        let mut inserter =
621            hyperdb_api::ArrowInserter::from_table(engine.connection(), opts.table.as_str())
622                .map_err(McpError::from)?;
623        inserter
624            .insert_batches(batches.iter())
625            .map_err(McpError::from)?;
626        inserter.execute().map_err(McpError::from)
627    })?;
628
629    let elapsed = timer.elapsed_ms();
630    let stats = IngestStats {
631        operation: "load_file".into(),
632        rows: row_count,
633        elapsed_ms: elapsed,
634        bytes_read: file_size,
635        bytes_stored: 0,
636        schema_inference_ms: Some(0),
637        table: opts.table.clone(),
638        file_format: Some("arrow_ipc".into()),
639        warning: None,
640        schema_changed: false,
641    };
642
643    Ok(IngestResult {
644        rows: row_count,
645        schema: columns,
646        stats,
647    })
648}
649
650/// Async twin of [`ingest_arrow_ipc_file`]. File I/O and Arrow decode run
651/// on the blocking pool; the COPY stream is driven on the async connection.
652/// [`hyperdb_api::AsyncArrowInserter`] accepts raw IPC Stream bytes, so we
653/// serialize the decoded batches into one buffer and send it in a single
654/// `insert_data` call.
655///
656/// # Errors
657///
658/// - Propagates [`ErrorCode::InvalidArgument`] from
659///   `reject_ipc_schema_override` when a schema override is supplied.
660/// - Returns [`ErrorCode::FileNotFound`] if `path` does not exist.
661/// - Returns [`ErrorCode::InternalError`] if the spawn-blocking Arrow
662///   decode task panics (join error).
663/// - Propagates Arrow decode errors from `read_arrow_ipc_file` /
664///   `record_batches_to_ipc_stream`.
665/// - Propagates transaction errors from the async `CREATE TABLE` and
666///   from [`hyperdb_api::AsyncArrowInserter`] operations.
667pub async fn ingest_arrow_ipc_file_async(
668    conn: &mut AsyncConnection,
669    path: &str,
670    opts: &IngestOptions,
671) -> Result<IngestResult, McpError> {
672    let timer = StatsTimer::start();
673    reject_ipc_schema_override(opts)?;
674
675    let file_path = Path::new(path);
676    if !file_path.exists() {
677        return Err(McpError::new(
678            ErrorCode::FileNotFound,
679            format!("File not found: {path}"),
680        ));
681    }
682    let file_size = std::fs::metadata(file_path).map_or(0, |m| m.len());
683
684    let path_owned = path.to_string();
685    let (columns, ipc_stream): (Vec<ColumnSchema>, Vec<u8>) =
686        tokio::task::spawn_blocking(move || -> Result<_, McpError> {
687            let (columns, batches) = read_arrow_ipc_file(&path_owned)?;
688            let ipc_stream = record_batches_to_ipc_stream(&batches)?;
689            Ok((columns, ipc_stream))
690        })
691        .await
692        .map_err(|e| McpError::new(ErrorCode::InternalError, format!("Task join error: {e}")))??;
693
694    let is_replace = opts.mode != "append";
695    let table_def = crate::schema::build_table_def(&opts.table, &columns)?;
696
697    let txn = conn.transaction().await.map_err(McpError::from)?;
698    let result: Result<u64, McpError> = async {
699        crate::ingest::create_table_async(
700            txn.connection(),
701            &opts.table,
702            &columns,
703            is_replace,
704            opts.target_db.as_deref(),
705        )
706        .await?;
707
708        // AsyncArrowInserter resolves table names via the connection's
709        // search path. When targeting a non-primary DB, qualify the
710        // TableDefinition with that database. (Same trick the sync
711        // Arrow IPC path uses via scoped_search_path.)
712        let mut inserter = hyperdb_api::AsyncArrowInserter::new(txn.connection(), &table_def)
713            .map_err(McpError::from)?;
714        inserter
715            .insert_data(&ipc_stream)
716            .await
717            .map_err(McpError::from)?;
718        inserter.execute().await.map_err(McpError::from)
719    }
720    .await;
721
722    let row_count = match result {
723        Ok(n) => {
724            txn.commit().await.map_err(McpError::from)?;
725            n
726        }
727        Err(e) => {
728            if let Err(rb) = txn.rollback().await {
729                tracing::warn!("rollback after error failed: {}", rb);
730            }
731            return Err(e);
732        }
733    };
734
735    let elapsed = timer.elapsed_ms();
736    let stats = IngestStats {
737        operation: "load_file".into(),
738        rows: row_count,
739        elapsed_ms: elapsed,
740        bytes_read: file_size,
741        bytes_stored: 0,
742        schema_inference_ms: Some(0),
743        table: opts.table.clone(),
744        file_format: Some("arrow_ipc".into()),
745        warning: None,
746        schema_changed: false,
747    };
748
749    Ok(IngestResult {
750        rows: row_count,
751        schema: columns,
752        stats,
753    })
754}