liboxen 0.52.0

Oxen is a fast, unstructured data version control, to help version large machine learning datasets written in Rust.
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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

use duckdb::arrow::array::RecordBatch;
use polars::frame::DataFrame;

use crate::constants::TABLE_NAME;
use crate::core::db::data_frames::DataFrameError;
use crate::core::db::data_frames::workspace_df_db::schema_without_oxen_cols;
use crate::model::Schema;
use crate::model::data_frame::schema::DataType;
use crate::view::data_frames::columns::{ColumnToDelete, ColumnToUpdate, NewColumn};

use super::df_db;

pub fn add_column(
    conn: &duckdb::Connection,
    new_column: &NewColumn,
) -> Result<DataFrame, DataFrameError> {
    let table_schema = schema_without_oxen_cols(conn, TABLE_NAME)?;

    if table_schema.has_column(&new_column.name) {
        return Err(DataFrameError::ColumnNameAlreadyExists(
            new_column.name.clone(),
        ));
    }

    let inserted_df = polar_insert_column(conn, TABLE_NAME, new_column)?;
    Ok(inserted_df)
}

pub fn delete_column(
    conn: &duckdb::Connection,
    column_to_delete: &ColumnToDelete,
) -> Result<DataFrame, DataFrameError> {
    let table_schema = schema_without_oxen_cols(conn, TABLE_NAME)?;

    if !table_schema.has_column(&column_to_delete.name) {
        return Err(DataFrameError::ColumnNameNotFound(
            column_to_delete.name.clone(),
        ));
    }

    let inserted_df = polar_delete_column(conn, TABLE_NAME, column_to_delete)?;
    Ok(inserted_df)
}

pub fn update_column(
    conn: &duckdb::Connection,
    column_to_update: &ColumnToUpdate,
    table_schema: &Schema,
) -> Result<DataFrame, DataFrameError> {
    if !table_schema.has_column(&column_to_update.name) {
        return Err(DataFrameError::ColumnNameNotFound(
            column_to_update.name.clone(),
        ));
    }

    let inserted_df = polar_update_column(conn, TABLE_NAME, column_to_update)?;
    Ok(inserted_df)
}

pub fn polar_insert_column(
    conn: &duckdb::Connection,
    table_name: &str,
    new_column: &NewColumn,
) -> Result<DataFrame, DataFrameError> {
    let data_type = DataType::from_string(&new_column.data_type).to_sql();
    let sql = format!(
        "ALTER TABLE {} ADD COLUMN {} {}",
        table_name, new_column.name, data_type
    );
    conn.execute(&sql, [])?;

    let sql_query = format!("SELECT * FROM {table_name}");
    let result_set: Vec<RecordBatch> = conn.prepare(&sql_query)?.query_arrow([])?.collect();

    df_db::record_batches_to_polars_df(result_set)
}

pub fn polar_delete_column(
    conn: &duckdb::Connection,
    table_name: &str,
    column_to_delete: &ColumnToDelete,
) -> Result<DataFrame, DataFrameError> {
    // Corrected to DROP COLUMN instead of ADD COLUMN
    let sql = format!(
        "ALTER TABLE {} DROP COLUMN {}",
        table_name, column_to_delete.name
    );
    conn.execute(&sql, [])?;

    let sql_query = format!("SELECT * FROM {table_name}");
    let result_set: Vec<RecordBatch> = conn.prepare(&sql_query)?.query_arrow([])?.collect();

    df_db::record_batches_to_polars_df(result_set)
}

pub fn polar_update_column(
    conn: &duckdb::Connection,
    table_name: &str,
    column_to_update: &ColumnToUpdate,
) -> Result<DataFrame, DataFrameError> {
    let sql_commands = {
        let mut sql_commands = Vec::new();

        if let Some(ref new_data_type) = column_to_update.new_data_type {
            let data_type = DataType::from_string(new_data_type).to_sql();

            let update_type_sql = format!(
                "ALTER TABLE {} ALTER COLUMN {} TYPE {}",
                table_name, column_to_update.name, data_type
            );
            sql_commands.push(update_type_sql);
        }

        if let Some(ref new_name) = column_to_update.new_name {
            let rename_sql = format!(
                "ALTER TABLE {} RENAME COLUMN {} TO {}",
                table_name, column_to_update.name, new_name
            );
            sql_commands.push(rename_sql);
        }
        sql_commands
    };

    for sql in sql_commands {
        conn.execute(&sql, [])?;
    }

    let sql_query = format!("SELECT * FROM {table_name}");
    let result_set: Vec<RecordBatch> = conn.prepare(&sql_query)?.query_arrow([])?.collect();

    df_db::record_batches_to_polars_df(result_set)
}