swellow 0.2.0

SQL-first migration CLI
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83

use crate::db::error::{EngineError, EngineErrorKind};
use arrow::array::{Array, StringArray};
use arrow::record_batch::RecordBatch;


/// Retrieves and downcasts a column from an Arrow RecordBatch.
pub fn get_column<'a, T: Array + 'static>(
    batch: &'a arrow::record_batch::RecordBatch,
    column_index: usize,
) -> Result<&'a T, EngineError> {
    let array = batch
        .columns()
        .get(column_index)
        .ok_or_else(|| EngineError {
            kind: EngineErrorKind::ColumnIndexOutOfBounds {
                column_index,
                num_columns: batch.num_columns(),
            },
        })?;

    let schema = batch.schema();
    let field = schema
        .fields()
        .get(column_index)
        .ok_or_else(|| EngineError {
            kind: EngineErrorKind::ColumnIndexOutOfBounds {
                column_index,
                num_columns: batch.schema().fields().len(),
            },
        })?;

    array
        .as_any()
        .downcast_ref::<T>()
        .ok_or_else(|| EngineError {
            kind: EngineErrorKind::ColumnTypeMismatch {
                column_index,
                expected: std::any::type_name::<T>(),
                found: field.data_type().clone(),
            },
        })
}


/// Retrieves and downcasts a column from an Arrow RecordBatch by column name.
pub fn get_column_by_name<'a, T: Array + 'static>(
    batch: &'a arrow::record_batch::RecordBatch,
    column_name: &str,
) -> Result<&'a T, EngineError> {
    // Find the column index by name
    let column_index = batch
        .schema()
        .index_of(column_name)
        .map_err(|_| EngineError {
            kind: EngineErrorKind::ColumnNotFound {
                column_name: column_name.to_string(),
            },
        })?;

    // Get column by index
    get_column::<T>(batch, column_index)
}


/// Get the first string value from a batch's column
pub fn get_first_string(
    batch: &RecordBatch,
    column_name: &str,
) -> Result<String, EngineError> {
    let array = get_column_by_name::<StringArray>(
        batch, column_name
    )?;

    if array.is_valid(0) {
        Ok(array.value(0).to_string())
    } else {
        Err(EngineError {
            kind: EngineErrorKind::InvalidSchema {
                stderr: format!("Column {} is empty", column_name),
            }
        })
    }
}