使用 duckdb::arrow 实现表格输出的 DuckDB CLI 代码

试图让DeepSeek编写输出列名、并支持各种数据类型的代码,总是不成功,在duckdb-rs主页看到它的示例代码支持arrow表格,把此示例提交给DeepSeek, 并让他删除语法高亮代码,就能正常处理各种查询了。如下所示。

use std::{
    error::Error,
    io::{self, BufRead},
    time::Instant,
};

use duckdb::{params, Connection, arrow::record_batch::RecordBatch};
use duckdb::arrow::util::pretty::print_batches;

fn main() -> Result<(), Box<dyn Error>> {
    let conn = Connection::open_in_memory()?;

    // 检查是否是管道输入
    if atty::isnt(atty::Stream::Stdin) {
        // 管道模式
        let stdin = io::stdin();
        
        for line in stdin.lock().lines() {
            let sql = line?;
            if sql.trim().is_empty() {
                continue;
            }
            
            let start_time = Instant::now();
            let mut stmt = conn.prepare(&sql)?;
            
            // 使用 Arrow 格式查询并打印结果
            let rbs: Vec<RecordBatch> = stmt.query_arrow([])?.collect();
            print_batches(&rbs)?;
            
            eprintln!("Query executed in {:?}", start_time.elapsed());
        }
    } else {
        // 交互模式
        println!("DuckDB CLI (Arrow Output)");
        println!("Enter SQL commands or .exit to quit");
        let mut rl = rustyline::DefaultEditor::new()?;
        //let mut rl = rustyline::Editor::<SQLHelper, rustyline::history::DefaultHistory>::new()?;
        
        loop {
            match rl.readline("duckdb> ") {
                Ok(line) => {
                    let line = line.trim();
                    if line.eq_ignore_ascii_case(".exit") {
                        break;
                    }
                    
                    if line.is_empty() {
                        continue;
                    }
                    
                    // 执行并计时
                    let start_time = Instant::now();
                    let mut stmt = conn.prepare(line)?;
                    
                    // 使用 Arrow 格式查询并打印结果
                    let rbs: Vec<RecordBatch> = stmt.query_arrow([])?.collect();
                    print_batches(&rbs)?;
                    
                    println!("\nQuery executed in {:?}", start_time.elapsed());
                    
                    rl.add_history_entry(line);
                }
                Err(_) => break,
            }
        }
    }
    
    Ok(())
}

输出

DuckDB CLI (Arrow Output)
Enter SQL commands or .exit to quit
duckdb> select 1 a;
+---+
| a |
+---+
| 1 |
+---+

Query executed in 30.619746ms
duckdb> select * from range(3);
+-------+
| range |
+-------+
| 0     |
| 1     |
| 2     |
+-------+

Query executed in 23.698313ms
duckdb> create table t (a int,b varchar);
++
++

Query executed in 8.113396ms
duckdb> insert into t select 1,'a';
+-------+
| Count |
+-------+
| 1     |
+-------+

Query executed in 31.439682ms
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值