use pandrs::column::{Column, ColumnTrait, StringColumn};
use pandrs::optimized::dataframe::OptimizedDataFrame;
use std::time::Instant;
fn main() {
println!("=== Test of Extended Features of Optimized DataFrame ===\n");
create_test_csv(100000);
let start = Instant::now();
let df = OptimizedDataFrame::from_csv("test_data.csv", true).unwrap();
let duration = start.elapsed();
println!("CSV Read Time: {duration:?}");
println!(
"Row Count: {}, Column Count: {}",
df.row_count(),
df.column_count()
);
let start = Instant::now();
let melted_df = df
.melt(
&["id"],
Some(&["name", "age", "score"]),
Some("variable"),
Some("value"),
)
.unwrap();
let duration = start.elapsed();
println!("\nMelt Operation Time: {duration:?}");
println!(
"Post-Transformation Row Count: {}, Column Count: {}",
melted_df.row_count(),
melted_df.column_count()
);
let start = Instant::now();
let applied_df = df
.apply(
|col| {
if col.column_type() == pandrs::column::ColumnType::String {
if let Some(str_col) = col.as_string() {
let mut new_data = Vec::with_capacity(str_col.len());
for i in 0..str_col.len() {
if let Ok(Some(val)) = str_col.get(i) {
new_data.push(val.to_uppercase());
} else {
new_data.push(String::new());
}
}
Ok(Column::String(StringColumn::new(new_data)))
} else {
Ok(col.column().clone())
}
} else {
Ok(col.column().clone())
}
},
Some(&["name"]),
)
.unwrap();
let duration = start.elapsed();
println!("\nApply Operation Time: {duration:?}");
println!(
"Post-Processing Row Count: {}, Column Count: {}",
applied_df.row_count(),
applied_df.column_count()
);
let start = Instant::now();
applied_df.to_csv("test_output.csv", true).unwrap();
let duration = start.elapsed();
println!("\nCSV Write Time: {duration:?}");
println!("\n=== Test of Optimized DataFrame Features Complete ===");
}
fn create_test_csv(rows: usize) {
use std::fs::File;
use std::io::{BufWriter, Write};
println!("Creating Test CSV File ({rows} rows)...");
let file = File::create("test_data.csv").unwrap();
let mut writer = BufWriter::new(file);
writeln!(writer, "id,name,age,score,category").unwrap();
let categories = ["A", "B", "C", "D", "E"];
let names = [
"Alice", "Bob", "Charlie", "David", "Emma", "Frank", "Grace", "Hannah", "Ian", "Julia",
];
for i in 0..rows {
let name = names[i % names.len()];
let age = 20 + (i % 50);
let score = (i % 100) as f64 / 10.0;
let category = categories[i % categories.len()];
writeln!(writer, "{i},{name},{age},{score},{category}").unwrap();
}
println!("Test CSV File Creation Complete");
}