use pandrs::error::Result;
use pandrs::{Column, Int64Column, OptimizedDataFrame, StringColumn};
#[test]
#[allow(clippy::result_large_err)]
fn test_inner_join() -> Result<()> {
let mut left_df = OptimizedDataFrame::new();
let id_col = Int64Column::new(vec![1, 2, 3, 4]);
left_df.add_column("id", Column::Int64(id_col))?;
let name_col = StringColumn::new(vec![
"Alice".to_string(),
"Bob".to_string(),
"Charlie".to_string(),
"Dave".to_string(),
]);
left_df.add_column("name", Column::String(name_col))?;
let mut right_df = OptimizedDataFrame::new();
let id_col = Int64Column::new(vec![1, 2, 5, 6]);
right_df.add_column("id", Column::Int64(id_col))?;
let value_col = Int64Column::new(vec![100, 200, 500, 600]);
right_df.add_column("value", Column::Int64(value_col))?;
let joined = left_df.inner_join(&right_df, "id", "id")?;
assert_eq!(joined.row_count(), 2);
assert_eq!(joined.column_count(), 3);
Ok(())
}
#[test]
#[allow(clippy::result_large_err)]
fn test_left_join() -> Result<()> {
let mut left_df = OptimizedDataFrame::new();
let id_col = Int64Column::new(vec![1, 2, 3, 4]);
left_df.add_column("id", Column::Int64(id_col))?;
let name_col = StringColumn::new(vec![
"Alice".to_string(),
"Bob".to_string(),
"Charlie".to_string(),
"Dave".to_string(),
]);
left_df.add_column("name", Column::String(name_col))?;
let mut right_df = OptimizedDataFrame::new();
let id_col = Int64Column::new(vec![1, 2, 5, 6]);
right_df.add_column("id", Column::Int64(id_col))?;
let value_col = Int64Column::new(vec![100, 200, 500, 600]);
right_df.add_column("value", Column::Int64(value_col))?;
let joined = left_df.left_join(&right_df, "id", "id")?;
assert_eq!(joined.row_count(), 4);
assert_eq!(joined.column_count(), 3);
Ok(())
}
#[test]
#[allow(clippy::result_large_err)]
fn test_right_join() -> Result<()> {
let mut left_df = OptimizedDataFrame::new();
let id_col = Int64Column::new(vec![1, 2, 3, 4]);
left_df.add_column("id", Column::Int64(id_col))?;
let name_col = StringColumn::new(vec![
"Alice".to_string(),
"Bob".to_string(),
"Charlie".to_string(),
"Dave".to_string(),
]);
left_df.add_column("name", Column::String(name_col))?;
let mut right_df = OptimizedDataFrame::new();
let id_col = Int64Column::new(vec![1, 2, 5, 6]);
right_df.add_column("id", Column::Int64(id_col))?;
let value_col = Int64Column::new(vec![100, 200, 500, 600]);
right_df.add_column("value", Column::Int64(value_col))?;
let joined = left_df.right_join(&right_df, "id", "id")?;
assert_eq!(joined.row_count(), 4);
assert_eq!(joined.column_count(), 3);
Ok(())
}
#[test]
#[allow(clippy::result_large_err)]
fn test_outer_join() -> Result<()> {
let mut left_df = OptimizedDataFrame::new();
let id_col = Int64Column::new(vec![1, 2, 3, 4]);
left_df.add_column("id", Column::Int64(id_col))?;
let name_col = StringColumn::new(vec![
"Alice".to_string(),
"Bob".to_string(),
"Charlie".to_string(),
"Dave".to_string(),
]);
left_df.add_column("name", Column::String(name_col))?;
let mut right_df = OptimizedDataFrame::new();
let id_col = Int64Column::new(vec![1, 2, 5, 6]);
right_df.add_column("id", Column::Int64(id_col))?;
let value_col = Int64Column::new(vec![100, 200, 500, 600]);
right_df.add_column("value", Column::Int64(value_col))?;
let joined = left_df.outer_join(&right_df, "id", "id")?;
assert_eq!(joined.row_count(), 6);
assert_eq!(joined.column_count(), 3);
Ok(())
}
#[test]
#[allow(clippy::result_large_err)]
fn test_join_different_column_names() -> Result<()> {
let mut left_df = OptimizedDataFrame::new();
let left_id_col = Int64Column::new(vec![1, 2, 3, 4]);
left_df.add_column("left_id", Column::Int64(left_id_col))?;
let name_col = StringColumn::new(vec![
"Alice".to_string(),
"Bob".to_string(),
"Charlie".to_string(),
"Dave".to_string(),
]);
left_df.add_column("name", Column::String(name_col))?;
let mut right_df = OptimizedDataFrame::new();
let right_id_col = Int64Column::new(vec![1, 2, 5, 6]);
right_df.add_column("right_id", Column::Int64(right_id_col))?;
let value_col = Int64Column::new(vec![100, 200, 500, 600]);
right_df.add_column("value", Column::Int64(value_col))?;
let joined = left_df.inner_join(&right_df, "left_id", "right_id")?;
assert_eq!(joined.row_count(), 2);
assert_eq!(joined.column_count(), 3);
Ok(())
}
#[test]
#[allow(clippy::result_large_err)]
fn test_empty_join() -> Result<()> {
let empty_df = OptimizedDataFrame::new();
let result = empty_df.inner_join(&empty_df, "id", "id");
assert!(result.is_err());
let mut left_df = OptimizedDataFrame::new();
let id_col = Int64Column::new(vec![1, 2, 3, 4]);
left_df.add_column("id", Column::Int64(id_col))?;
let mut right_df = OptimizedDataFrame::new();
let id_col = Int64Column::new(vec![5, 6, 7, 8]);
right_df.add_column("id", Column::Int64(id_col))?;
let joined = left_df.inner_join(&right_df, "id", "id")?;
assert_eq!(joined.row_count(), 0);
Ok(())
}