diff --git a/Cargo.lock b/Cargo.lock index 9fe8082..dbf8f9f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -70,6 +70,21 @@ dependencies = [ "memchr", ] +[[package]] +name = "alloc-no-stdlib" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" + +[[package]] +name = "alloc-stdlib" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195" +dependencies = [ + "alloc-no-stdlib", +] + [[package]] name = "android_system_properties" version = "0.1.5" @@ -387,12 +402,39 @@ dependencies = [ "serde_core", ] +[[package]] +name = "brotli" +version = "8.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor", +] + +[[package]] +name = "brotli-decompressor" +version = "5.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] + [[package]] name = "bumpalo" version = "3.19.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5dd9dc738b7a8311c7ade152424974d8115f2cdad61e8dab8dac9f2362298510" +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "bytes" version = "1.11.0" @@ -406,6 +448,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "755d2fce177175ffca841e9a06afdb2c4ab0f593d53b4dee48147dfaade85932" dependencies = [ "find-msvc-tools", + "jobserver", + "libc", "shlex", ] @@ -606,6 +650,7 @@ dependencies = [ "clap", "comfy-table", "nu-ansi-term", + "parquet", "reedline", "syntect", "tempfile", @@ -715,6 +760,7 @@ checksum = "b375d6465b98090a5f25b1c7703f3859783755aa9a80433b36e0379a3ec2f369" dependencies = [ "crc32fast", "miniz_oxide", + "zlib-rs", ] [[package]] @@ -742,10 +788,21 @@ checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", "libc", - "r-efi", + "r-efi 5.3.0", "wasip2", ] +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", +] + [[package]] name = "half" version = "2.7.1" @@ -804,6 +861,12 @@ dependencies = [ "hashbrown", ] +[[package]] +name = "integer-encoding" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" + [[package]] name = "is_terminal_polyfill" version = "1.70.2" @@ -825,6 +888,16 @@ version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" +[[package]] +name = "jobserver" +version = "0.1.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" +dependencies = [ + "getrandom 0.4.3", + "libc", +] + [[package]] name = "js-sys" version = "0.3.85" @@ -947,6 +1020,15 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "lz4_flex" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ef0d4ed8669f8f8826eb00dc878084aa8f253506c4fd5e8f58f5bce72ddb97e" +dependencies = [ + "twox-hash", +] + [[package]] name = "memchr" version = "2.7.6" @@ -1062,6 +1144,15 @@ dependencies = [ "pkg-config", ] +[[package]] +name = "ordered-float" +version = "2.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c" +dependencies = [ + "num-traits", +] + [[package]] name = "parking_lot" version = "0.12.5" @@ -1085,6 +1176,45 @@ dependencies = [ "windows-link", ] +[[package]] +name = "parquet" +version = "58.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d3f9f2205199603564127932b89695f52b62322f541d0fc7179d57c2e1c9877" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-ipc", + "arrow-schema", + "arrow-select", + "base64", + "brotli", + "bytes", + "chrono", + "flate2", + "half", + "hashbrown", + "lz4_flex", + "num-bigint", + "num-integer", + "num-traits", + "paste", + "seq-macro", + "simdutf8", + "snap", + "thrift", + "twox-hash", + "zstd", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "path-slash" version = "0.2.1" @@ -1167,6 +1297,12 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + [[package]] name = "redox_syscall" version = "0.5.18" @@ -1280,6 +1416,12 @@ version = "1.0.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" +[[package]] +name = "seq-macro" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" + [[package]] name = "serde" version = "1.0.228" @@ -1393,6 +1535,12 @@ version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +[[package]] +name = "snap" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b" + [[package]] name = "strip-ansi-escapes" version = "0.2.1" @@ -1520,6 +1668,17 @@ dependencies = [ "syn", ] +[[package]] +name = "thrift" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e54bc85fc7faa8bc175c4bab5b92ba8d9a3ce893d0e9f42cc455c8ab16a9e09" +dependencies = [ + "byteorder", + "integer-encoding", + "ordered-float", +] + [[package]] name = "time" version = "0.3.46" @@ -1597,6 +1756,12 @@ version = "1.1.1+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "756daf9b1013ebe47a8776667b466417e2d4c5679d441c26230efd9ef78692db" +[[package]] +name = "twox-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" + [[package]] name = "unicase" version = "2.9.0" @@ -1942,8 +2107,42 @@ dependencies = [ "syn", ] +[[package]] +name = "zlib-rs" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40990edd51aae2c2b6907af74ffb635029d5788228222c4bb811e9351c0caad3" + [[package]] name = "zmij" version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94f63c051f4fe3c1509da62131a678643c5b6fbdc9273b2b79d4378ebda003d2" + +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] diff --git a/Cargo.toml b/Cargo.toml index 96fa9ab..a7550c4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ arrow-schema = "58.1.0" clap = { version = "4.6.1", features = ["derive"] } comfy-table = "7.2.2" nu-ansi-term = "0.50.3" +parquet = "58.1.0" reedline = "0.47.0" syntect = "5.3.0" terminal-colorsaurus = "1.0.3" diff --git a/README.md b/README.md index 849c770..7ed3457 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ A command-line tool for querying databases via [ADBC](https://arrow.apache.org/a - **Interactive SQL shell** - Execute SQL queries with command history and intuitive navigation - **Syntax highlighting** - SQL queries highlighted for improved readability - **Formatted output** - Results displayed in clean, aligned tables with dynamic column width -- **File export** - Export query results to JSON, CSV, or Arrow IPC files +- **File export** - Export query results to JSON, JSON lines, CSV, Arrow IPC, Arrow IPC stream, or Parquet files - **Fast and lightweight** - Built in Rust for high performance and minimal resource usage ## Installation @@ -101,8 +101,11 @@ Execute a query and output the result to a file: ```sh databow --driver duckdb --query "SELECT 42 AS the_answer" --output result.json +databow --driver duckdb --query "SELECT 42 AS the_answer" --output result.jsonl databow --driver duckdb --query "SELECT 42 AS the_answer" --output result.csv databow --driver duckdb --query "SELECT 42 AS the_answer" --output result.arrow +databow --driver duckdb --query "SELECT 42 AS the_answer" --output result.arrows +databow --driver duckdb --query "SELECT 42 AS the_answer" --output result.parquet ``` ## Reference diff --git a/docs/index.md b/docs/index.md index 06e5a75..81aa504 100644 --- a/docs/index.md +++ b/docs/index.md @@ -19,5 +19,5 @@ databow is a command-line tool for querying databases. - **Interactive SQL shell** - Execute SQL queries with command history and intuitive navigation - **Syntax highlighting** - SQL queries highlighted for improved readability - **Formatted output** - Results displayed in clean, aligned tables with dynamic column width -- **File export** - Export query results to JSON, CSV, or Arrow IPC files +- **File export** - Export query results to JSON, JSON lines, CSV, Arrow IPC, Arrow IPC stream, or Parquet files - **Fast and lightweight** - Built in Rust for high performance and minimal resource usage diff --git a/docs/reference.md b/docs/reference.md index 25c203a..7c553e5 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -108,11 +108,14 @@ databow --driver duckdb --query "SELECT 42 AS the_answer" --output result.json The output format is inferred from the file extension: -| Extension | Format | -|-----------------|-----------| -| `.json` | JSON | -| `.csv` | CSV | -| `.arrow`, `.ipc`| Arrow IPC | +| Extension | Format | +|-----------------|------------------| +| `.json` | JSON | +| `.jsonl` | JSON lines | +| `.csv` | CSV | +| `.arrow`, `.ipc`| Arrow IPC (file) | +| `.arrows` | Arrow IPC stream | +| `.parquet` | Parquet | ## --help diff --git a/docs/tutorial.md b/docs/tutorial.md index 59ec6ae..6bd479f 100644 --- a/docs/tutorial.md +++ b/docs/tutorial.md @@ -143,7 +143,7 @@ $ databow --profile warehouse --file query.sql └───────────────────┘ ``` -Instead of printing query results to stdout, the [`--output` argument](/reference/#-output) can be used to write results to JSON, CSV, or Arrow IPC files: +Instead of printing query results to stdout, the [`--output` argument](/reference/#-output) can be used to write results to JSON, JSON lines, CSV, Arrow IPC, Arrow IPC stream, or Parquet files: ```console $ databow --profile warehouse --query "SELECT * FROM penguins" --output penguins.csv diff --git a/src/output.rs b/src/output.rs index 690335a..6b72066 100644 --- a/src/output.rs +++ b/src/output.rs @@ -2,26 +2,33 @@ // SPDX-License-Identifier: Apache-2.0 use arrow::csv::writer::Writer as CsvWriter; -use arrow::ipc::writer::FileWriter as IpcWriter; -use arrow::json::writer::{JsonArray, Writer as JsonWriter}; +use arrow::ipc::writer::{FileWriter as IpcWriter, StreamWriter as IpcStreamWriter}; +use arrow::json::writer::{JsonArray, LineDelimited, Writer as JsonWriter}; use arrow_array::RecordBatch; use arrow_schema::ArrowError; +use parquet::arrow::ArrowWriter; use std::fs::File; use std::path::Path; #[derive(Debug, Clone, Copy, PartialEq)] pub enum OutputFormat { Json, + Jsonl, Csv, Arrow, + ArrowStream, + Parquet, } impl OutputFormat { pub fn from_path(path: &Path) -> Result { match path.extension().and_then(|ext| ext.to_str()) { Some("json") => Ok(OutputFormat::Json), + Some("jsonl") => Ok(OutputFormat::Jsonl), Some("csv") => Ok(OutputFormat::Csv), Some("arrow" | "ipc") => Ok(OutputFormat::Arrow), + Some("arrows") => Ok(OutputFormat::ArrowStream), + Some("parquet") => Ok(OutputFormat::Parquet), Some(ext) => Err(format!("Unsupported file extension: '.{ext}'")), None => Err("Cannot infer format: no file extension".to_string()), } @@ -38,8 +45,11 @@ pub fn write_batches_to_file(batches: &[RecordBatch], path: &Path) -> Result<(), match format { OutputFormat::Json => write_json(batches, file), + OutputFormat::Jsonl => write_jsonl(batches, file), OutputFormat::Csv => write_csv(batches, file), OutputFormat::Arrow => write_arrow_ipc(batches, file), + OutputFormat::ArrowStream => write_arrow_stream(batches, file), + OutputFormat::Parquet => write_parquet(batches, file), } } @@ -53,6 +63,16 @@ fn write_json(batches: &[RecordBatch], file: File) -> Result<(), ArrowError> { Ok(()) } +fn write_jsonl(batches: &[RecordBatch], file: File) -> Result<(), ArrowError> { + let mut writer = JsonWriter::<_, LineDelimited>::new(file); + for batch in batches { + writer.write(batch)?; + } + writer.finish()?; + + Ok(()) +} + fn write_csv(batches: &[RecordBatch], file: File) -> Result<(), ArrowError> { let mut writer = CsvWriter::new(file); for batch in batches { @@ -76,6 +96,39 @@ fn write_arrow_ipc(batches: &[RecordBatch], file: File) -> Result<(), ArrowError Ok(()) } +fn write_arrow_stream(batches: &[RecordBatch], file: File) -> Result<(), ArrowError> { + if batches.is_empty() { + return Ok(()); + } + let schema = batches[0].schema(); + let mut writer = IpcStreamWriter::try_new(file, &schema)?; + for batch in batches { + writer.write(batch)?; + } + writer.finish()?; + + Ok(()) +} + +fn write_parquet(batches: &[RecordBatch], file: File) -> Result<(), ArrowError> { + if batches.is_empty() { + return Ok(()); + } + let schema = batches[0].schema(); + let mut writer = ArrowWriter::try_new(file, schema, None) + .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; + for batch in batches { + writer + .write(batch) + .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; + } + writer + .close() + .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; + + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -243,6 +296,118 @@ mod tests { assert!(!path.exists()); } + #[test] + fn test_output_format_from_path_arrows() { + let path = Path::new("output.arrows"); + assert_eq!( + OutputFormat::from_path(path).unwrap(), + OutputFormat::ArrowStream + ); + } + + #[test] + fn test_output_format_from_path_parquet() { + let path = Path::new("output.parquet"); + assert_eq!( + OutputFormat::from_path(path).unwrap(), + OutputFormat::Parquet + ); + } + + #[test] + fn test_output_format_from_path_jsonl() { + let path = Path::new("output.jsonl"); + assert_eq!(OutputFormat::from_path(path).unwrap(), OutputFormat::Jsonl); + } + + #[test] + fn test_write_arrow_stream() { + let dir = tempdir().unwrap(); + let path = dir.path().join("output.arrows"); + let batch = create_test_batch(); + + write_batches_to_file(std::slice::from_ref(&batch), &path).unwrap(); + + // Verify by reading it back as an IPC stream + let file = File::open(&path).unwrap(); + let reader = arrow::ipc::reader::StreamReader::try_new(file, None).unwrap(); + let read_batches: Vec = reader.map(|r| r.unwrap()).collect(); + + assert_eq!(read_batches.len(), 1); + assert_eq!(read_batches[0].num_rows(), batch.num_rows()); + assert_eq!(read_batches[0].num_columns(), batch.num_columns()); + } + + #[test] + fn test_write_arrow_stream_empty_batches_direct() { + let dir = tempdir().unwrap(); + let path = dir.path().join("output.arrows"); + let file = File::create(&path).unwrap(); + + let result = write_arrow_stream(&[], file); + + assert!(result.is_ok()); + } + + #[test] + fn test_write_parquet() { + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + + let dir = tempdir().unwrap(); + let path = dir.path().join("output.parquet"); + let batch = create_test_batch(); + + write_batches_to_file(std::slice::from_ref(&batch), &path).unwrap(); + + // Verify by reading it back + let file = File::open(&path).unwrap(); + let reader = ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap() + .build() + .unwrap(); + let read_batches: Vec = reader.map(|r| r.unwrap()).collect(); + + let total_rows: usize = read_batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, batch.num_rows()); + assert_eq!(read_batches[0].num_columns(), batch.num_columns()); + } + + #[test] + fn test_write_parquet_empty_batches_direct() { + let dir = tempdir().unwrap(); + let path = dir.path().join("output.parquet"); + let file = File::create(&path).unwrap(); + + let result = write_parquet(&[], file); + + assert!(result.is_ok()); + } + + #[test] + fn test_write_jsonl() { + let dir = tempdir().unwrap(); + let path = dir.path().join("output.jsonl"); + let batch = create_test_batch(); + + write_batches_to_file(std::slice::from_ref(&batch), &path).unwrap(); + + let mut file = File::open(&path).unwrap(); + let mut contents = String::new(); + file.read_to_string(&mut contents).unwrap(); + + // JSON lines: one object per line, not a wrapping array + assert!(!contents.trim_start().starts_with('[')); + let lines: Vec<&str> = contents.lines().filter(|l| !l.is_empty()).collect(); + assert_eq!(lines.len(), 3); + for line in &lines { + assert!(line.trim_start().starts_with('{')); + assert!(line.trim_end().ends_with('}')); + } + assert!(contents.contains("Alice")); + assert!(contents.contains("Bob")); + assert!(contents.contains("Charlie")); + } + #[test] fn test_write_multiple_batches() { let dir = tempdir().unwrap();