use std::collections::HashMap;
use std::hash::Hash;
use polars_core::prelude::{
AnyValue, Categorical32Chunked, CategoricalPhysical, Column, DataFrame, DataType,
PlSeedableRandomStateQuality, SeedableFromU64SeedExt, Series,
};
use polars_dtype::categorical::Categories;
use crate::error::{RangeFrameError, Result};
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Side {
Left,
Right,
}
#[derive(Clone, Copy, Debug)]
struct PairKeyEntry {
group_id: u32,
side: Side,
row_idx: usize,
}
#[derive(Clone, Copy, Debug)]
struct SingleKeyEntry {
group_id: u32,
row_idx: usize,
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
enum KeyValue {
Bool(bool),
Int(i64),
UInt(u64),
Str(String),
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct RowKey(Vec<KeyValue>);
pub fn factorize(df: &DataFrame, match_by: &[String]) -> Result<Vec<u32>> {
if match_by.is_empty() {
return Ok(vec![0; df.height()]);
}
validate_match_by_columns(df, match_by)?;
if match_by.len() == 1 {
let column_name = &match_by[0];
if let Some(ids) = factorize_single_column(df.column(column_name)?, column_name)? {
return Ok(ids);
}
}
let mut keys = normalize_key_frame(df, match_by)?;
factorize_hashed_single(&mut keys).or_else(|_| factorize_single_fallback(df, match_by))
}
pub fn factorize_pair(
left: &DataFrame,
right: &DataFrame,
match_by: &[String],
) -> Result<(Vec<u32>, Vec<u32>)> {
factorize_pair_by(left, right, match_by, match_by)
}
pub fn factorize_pair_by(
left: &DataFrame,
right: &DataFrame,
left_match_by: &[String],
right_match_by: &[String],
) -> Result<(Vec<u32>, Vec<u32>)> {
if left_match_by.len() != right_match_by.len() {
return Err(RangeFrameError::Polars(
polars_core::prelude::PolarsError::ComputeError(
"left_match_by and right_match_by must have the same number of columns".into(),
),
));
}
if left_match_by.is_empty() {
return Ok((vec![0; left.height()], vec![0; right.height()]));
}
validate_match_by_columns(left, left_match_by)?;
validate_match_by_columns(right, right_match_by)?;
if left_match_by.len() == 1 {
let left_column_name = &left_match_by[0];
let right_column_name = &right_match_by[0];
if let Some(ids) = factorize_single_column_pair(
left.column(left_column_name)?,
right.column(right_column_name)?,
left_column_name,
)? {
return Ok(ids);
}
}
if let Some((mut left_keys, mut right_keys)) =
normalize_key_frames_by(left, right, left_match_by, right_match_by)?
{
return factorize_hashed_pair(&mut left_keys, &mut right_keys);
}
factorize_pair_fallback_by(left, right, left_match_by, right_match_by)
}
fn validate_match_by_columns(df: &DataFrame, match_by: &[String]) -> Result<()> {
for column in match_by {
if df.column(column).is_err() {
return Err(RangeFrameError::MissingMatchBy(column.clone()));
}
}
Ok(())
}
fn normalize_key_frame(df: &DataFrame, match_by: &[String]) -> Result<DataFrame> {
let mut columns = Vec::with_capacity(match_by.len());
for column_name in match_by {
let column = df.column(column_name)?;
ensure_no_nulls(column, column_name)?;
columns.push(normalize_column(column_name, column)?);
}
Ok(DataFrame::new_infer_height(columns)?)
}
fn normalize_key_frames_by(
left: &DataFrame,
right: &DataFrame,
left_match_by: &[String],
right_match_by: &[String],
) -> Result<Option<(DataFrame, DataFrame)>> {
let mut left_columns = Vec::with_capacity(left_match_by.len());
let mut right_columns = Vec::with_capacity(right_match_by.len());
for (left_column_name, right_column_name) in left_match_by.iter().zip(right_match_by) {
let left_column = left.column(left_column_name)?;
let right_column = right.column(right_column_name)?;
ensure_no_nulls(left_column, left_column_name)?;
ensure_no_nulls(right_column, right_column_name)?;
let Some((left_normalized, right_normalized)) =
normalize_column_pair(left_column_name, left_column, right_column)?
else {
return Ok(None);
};
left_columns.push(left_normalized);
right_columns.push(right_normalized);
}
Ok(Some((
DataFrame::new_infer_height(left_columns)?,
DataFrame::new_infer_height(right_columns)?,
)))
}
fn ensure_no_nulls(column: &Column, column_name: &str) -> Result<()> {
if column.null_count() > 0 {
return Err(RangeFrameError::NullKeyValue {
column: column_name.to_owned(),
});
}
Ok(())
}
fn normalize_column(column_name: &str, column: &Column) -> Result<Column> {
let series = column.as_materialized_series();
match column.dtype() {
DataType::Boolean => Ok(column.clone()),
dtype if is_signed_integer_dtype(dtype) => cast_column(series, DataType::Int64),
dtype if is_unsigned_integer_dtype(dtype) => cast_column(series, DataType::UInt64),
dtype if is_string_like_dtype(dtype) => cast_column(series, DataType::String),
dtype => Err(RangeFrameError::UnsupportedKeyDtype {
column: column_name.to_owned(),
dtype: dtype.to_string(),
}),
}
}
fn normalize_column_pair(
column_name: &str,
left: &Column,
right: &Column,
) -> Result<Option<(Column, Column)>> {
let left_series = left.as_materialized_series();
let right_series = right.as_materialized_series();
let pair = match (left.dtype(), right.dtype()) {
(DataType::Boolean, DataType::Boolean) => Some((left.clone(), right.clone())),
(left_dtype, right_dtype)
if is_signed_integer_dtype(left_dtype) && is_signed_integer_dtype(right_dtype) =>
{
Some((
cast_column(left_series, DataType::Int64)?,
cast_column(right_series, DataType::Int64)?,
))
}
(left_dtype, right_dtype)
if is_unsigned_integer_dtype(left_dtype) && is_unsigned_integer_dtype(right_dtype) =>
{
Some((
cast_column(left_series, DataType::UInt64)?,
cast_column(right_series, DataType::UInt64)?,
))
}
(left_dtype, right_dtype)
if is_string_like_dtype(left_dtype) && is_string_like_dtype(right_dtype) =>
{
Some((
cast_column(left_series, DataType::String)?,
cast_column(right_series, DataType::String)?,
))
}
_ => None,
};
if pair.is_none() {
let left_dtype = left.dtype();
let right_dtype = right.dtype();
if is_fast_key_dtype(left_dtype) && is_fast_key_dtype(right_dtype) {
return Ok(None);
}
return Err(RangeFrameError::UnsupportedKeyDtype {
column: column_name.to_owned(),
dtype: format!("{left_dtype} vs {right_dtype}"),
});
}
Ok(pair)
}
fn cast_column(series: &Series, dtype: DataType) -> Result<Column> {
Ok(Column::from(series.cast(&dtype)?))
}
fn cast_series(series: &Series, dtype: DataType) -> Result<Series> {
Ok(series.cast(&dtype)?)
}
fn factorize_single_column(column: &Column, column_name: &str) -> Result<Option<Vec<u32>>> {
ensure_no_nulls(column, column_name)?;
let series = column.as_materialized_series();
match column.dtype() {
DataType::Boolean => {
let values = series.bool().expect("validated boolean dtype");
Ok(Some(factorize_scalar_values_single(
values.into_no_null_iter(),
column.len(),
)))
}
dtype if is_signed_integer_dtype(dtype) => {
let cast = cast_series(series, DataType::Int64)?;
let values = cast.i64().expect("cast to Int64 succeeded");
Ok(Some(factorize_scalar_values_single(
values.into_no_null_iter(),
column.len(),
)))
}
dtype if is_unsigned_integer_dtype(dtype) => {
let cast = cast_series(series, DataType::UInt64)?;
let values = cast.u64().expect("cast to UInt64 succeeded");
Ok(Some(factorize_scalar_values_single(
values.into_no_null_iter(),
column.len(),
)))
}
dtype if is_string_like_dtype(dtype) => factorize_string_values_single(series).map(Some),
dtype if is_fast_key_dtype(dtype) => Ok(None),
dtype => Err(RangeFrameError::UnsupportedKeyDtype {
column: column_name.to_owned(),
dtype: dtype.to_string(),
}),
}
}
fn factorize_single_column_pair(
left: &Column,
right: &Column,
column_name: &str,
) -> Result<Option<(Vec<u32>, Vec<u32>)>> {
ensure_no_nulls(left, column_name)?;
ensure_no_nulls(right, column_name)?;
let left_series = left.as_materialized_series();
let right_series = right.as_materialized_series();
match (left.dtype(), right.dtype()) {
(DataType::Boolean, DataType::Boolean) => {
let left_values = left_series.bool().expect("validated boolean dtype");
let right_values = right_series.bool().expect("validated boolean dtype");
Ok(Some(factorize_scalar_values_pair(
left_values.into_no_null_iter(),
right_values.into_no_null_iter(),
left.len(),
right.len(),
)))
}
(left_dtype, right_dtype)
if is_signed_integer_dtype(left_dtype) && is_signed_integer_dtype(right_dtype) =>
{
let left_cast = cast_series(left_series, DataType::Int64)?;
let right_cast = cast_series(right_series, DataType::Int64)?;
let left_values = left_cast.i64().expect("cast to Int64 succeeded");
let right_values = right_cast.i64().expect("cast to Int64 succeeded");
Ok(Some(factorize_scalar_values_pair(
left_values.into_no_null_iter(),
right_values.into_no_null_iter(),
left.len(),
right.len(),
)))
}
(left_dtype, right_dtype)
if is_unsigned_integer_dtype(left_dtype) && is_unsigned_integer_dtype(right_dtype) =>
{
let left_cast = cast_series(left_series, DataType::UInt64)?;
let right_cast = cast_series(right_series, DataType::UInt64)?;
let left_values = left_cast.u64().expect("cast to UInt64 succeeded");
let right_values = right_cast.u64().expect("cast to UInt64 succeeded");
Ok(Some(factorize_scalar_values_pair(
left_values.into_no_null_iter(),
right_values.into_no_null_iter(),
left.len(),
right.len(),
)))
}
(left_dtype, right_dtype)
if is_string_like_dtype(left_dtype) && is_string_like_dtype(right_dtype) =>
{
factorize_string_values_pair(left_series, right_series).map(Some)
}
_ => {
if is_fast_key_dtype(left.dtype()) && is_fast_key_dtype(right.dtype()) {
return Ok(None);
}
Err(RangeFrameError::UnsupportedKeyDtype {
column: column_name.to_owned(),
dtype: format!("{} vs {}", left.dtype(), right.dtype()),
})
}
}
}
fn factorize_scalar_values_single<K, I>(values: I, len: usize) -> Vec<u32>
where
K: Eq + Hash,
I: IntoIterator<Item = K>,
{
let mut ids = HashMap::<K, u32>::with_capacity(len);
let mut next_id = 0_u32;
let mut out = Vec::with_capacity(len);
for value in values {
let id = *ids.entry(value).or_insert_with(|| {
let id = next_id;
next_id = next_id.saturating_add(1);
id
});
out.push(id);
}
out
}
fn factorize_scalar_values_pair<K, I, J>(
left_values: I,
right_values: J,
left_len: usize,
right_len: usize,
) -> (Vec<u32>, Vec<u32>)
where
K: Eq + Hash,
I: IntoIterator<Item = K>,
J: IntoIterator<Item = K>,
{
let mut ids = HashMap::<K, u32>::with_capacity(left_len.saturating_add(right_len));
let mut next_id = 0_u32;
let mut left_ids = Vec::with_capacity(left_len);
for value in left_values {
let id = *ids.entry(value).or_insert_with(|| {
let id = next_id;
next_id = next_id.saturating_add(1);
id
});
left_ids.push(id);
}
let mut right_ids = Vec::with_capacity(right_len);
for value in right_values {
let id = *ids.entry(value).or_insert_with(|| {
let id = next_id;
next_id = next_id.saturating_add(1);
id
});
right_ids.push(id);
}
(left_ids, right_ids)
}
fn factorize_string_values_single(series: &Series) -> Result<Vec<u32>> {
let values = borrowed_string_values(series)?;
let encoded = Categorical32Chunked::from_str_iter(
series.name().clone(),
DataType::from_categories(Categories::global()),
values.map(Some),
)?;
Ok(encoded.physical().into_no_null_iter().collect())
}
fn factorize_string_values_pair<'a>(
left: &'a Series,
right: &'a Series,
) -> Result<(Vec<u32>, Vec<u32>)> {
let left_values = borrowed_string_values(left)?;
let right_values = borrowed_string_values(right)?;
let encoded = Categorical32Chunked::from_str_iter(
left.name().clone(),
DataType::from_categories(Categories::global()),
left_values.chain(right_values).map(Some),
)?;
let mut ids = encoded.physical().into_no_null_iter();
let left_ids = ids.by_ref().take(left.len()).collect();
let right_ids = ids.collect();
Ok((left_ids, right_ids))
}
pub(crate) fn borrowed_string_values<'a>(
series: &'a Series,
) -> Result<Box<dyn Iterator<Item = &'a str> + 'a>> {
match series.dtype() {
DataType::String => Ok(Box::new(
series
.str()
.expect("validated string dtype")
.into_no_null_iter(),
)),
DataType::Categorical(_, _) | DataType::Enum(_, _) => {
match series.dtype().cat_physical()? {
CategoricalPhysical::U8 => Ok(Box::new(
series
.cat8()
.expect("validated categorical dtype")
.iter_str()
.map(expect_non_null_string),
)),
CategoricalPhysical::U16 => Ok(Box::new(
series
.cat16()
.expect("validated categorical dtype")
.iter_str()
.map(expect_non_null_string),
)),
CategoricalPhysical::U32 => Ok(Box::new(
series
.cat32()
.expect("validated categorical dtype")
.iter_str()
.map(expect_non_null_string),
)),
}
}
_ => unreachable!("borrowed_string_values only supports string-like dtypes"),
}
}
fn expect_non_null_string(value: Option<&str>) -> &str {
value.expect("validated string-like values contain no nulls")
}
fn factorize_hashed_single(df: &mut DataFrame) -> Result<Vec<u32>> {
let hasher = PlSeedableRandomStateQuality::seed_from_u64(0);
let hashes = df.hash_rows(Some(hasher))?;
let mut groups = HashMap::<u64, Vec<SingleKeyEntry>>::with_capacity(df.height());
let mut next_id = 0_u32;
let mut out = Vec::with_capacity(df.height());
for (row_idx, hash) in hashes.into_no_null_iter().enumerate() {
let entries = groups.entry(hash).or_default();
let mut existing_id = None;
for entry in entries.iter().copied() {
if rows_equal_single(row_idx, entry.row_idx, df)? {
existing_id = Some(entry.group_id);
break;
}
}
let group_id = existing_id.unwrap_or_else(|| {
let group_id = next_id;
next_id = next_id.saturating_add(1);
entries.push(SingleKeyEntry { group_id, row_idx });
group_id
});
out.push(group_id);
}
Ok(out)
}
fn factorize_hashed_pair(
left: &mut DataFrame,
right: &mut DataFrame,
) -> Result<(Vec<u32>, Vec<u32>)> {
let hasher = PlSeedableRandomStateQuality::seed_from_u64(0);
let left_hashes = left.hash_rows(Some(hasher.clone()))?;
let right_hashes = right.hash_rows(Some(hasher))?;
let mut groups = HashMap::<u64, Vec<PairKeyEntry>>::with_capacity(
left.height().saturating_add(right.height()),
);
let mut next_id = 0_u32;
let mut left_ids = Vec::with_capacity(left.height());
for (row_idx, hash) in left_hashes.into_no_null_iter().enumerate() {
let id = resolve_or_insert_hashed_row(
hash,
Side::Left,
row_idx,
left,
right,
&mut groups,
&mut next_id,
)?;
left_ids.push(id);
}
let mut right_ids = Vec::with_capacity(right.height());
for (row_idx, hash) in right_hashes.into_no_null_iter().enumerate() {
let id = resolve_or_insert_hashed_row(
hash,
Side::Right,
row_idx,
left,
right,
&mut groups,
&mut next_id,
)?;
right_ids.push(id);
}
Ok((left_ids, right_ids))
}
fn resolve_or_insert_hashed_row(
hash: u64,
side: Side,
row_idx: usize,
left: &DataFrame,
right: &DataFrame,
groups: &mut HashMap<u64, Vec<PairKeyEntry>>,
next_id: &mut u32,
) -> Result<u32> {
let entries = groups.entry(hash).or_default();
for entry in entries.iter().copied() {
if rows_equal_pair(side, row_idx, entry.side, entry.row_idx, left, right)? {
return Ok(entry.group_id);
}
}
let group_id = *next_id;
*next_id = next_id.saturating_add(1);
entries.push(PairKeyEntry {
group_id,
side,
row_idx,
});
Ok(group_id)
}
fn rows_equal_single(left_row_idx: usize, right_row_idx: usize, df: &DataFrame) -> Result<bool> {
for column in df.columns() {
if column.get(left_row_idx)? != column.get(right_row_idx)? {
return Ok(false);
}
}
Ok(true)
}
fn rows_equal_pair(
left_side: Side,
left_row_idx: usize,
right_side: Side,
right_row_idx: usize,
left: &DataFrame,
right: &DataFrame,
) -> Result<bool> {
let left_df = match left_side {
Side::Left => left,
Side::Right => right,
};
let right_df = match right_side {
Side::Left => left,
Side::Right => right,
};
for (left_column, right_column) in left_df.columns().iter().zip(right_df.columns().iter()) {
if left_column.get(left_row_idx)? != right_column.get(right_row_idx)? {
return Ok(false);
}
}
Ok(true)
}
fn factorize_single_fallback(df: &DataFrame, match_by: &[String]) -> Result<Vec<u32>> {
let mut ids = HashMap::<RowKey, u32>::new();
let mut next_id = 0_u32;
factorize_one_fallback(df, match_by, &mut ids, &mut next_id)
}
fn factorize_pair_fallback_by(
left: &DataFrame,
right: &DataFrame,
left_match_by: &[String],
right_match_by: &[String],
) -> Result<(Vec<u32>, Vec<u32>)> {
let mut ids = HashMap::<RowKey, u32>::new();
let mut next_id = 0_u32;
let left_ids = factorize_one_fallback(left, left_match_by, &mut ids, &mut next_id)?;
let right_ids = factorize_one_fallback(right, right_match_by, &mut ids, &mut next_id)?;
Ok((left_ids, right_ids))
}
fn factorize_one_fallback(
df: &DataFrame,
match_by: &[String],
ids: &mut HashMap<RowKey, u32>,
next_id: &mut u32,
) -> Result<Vec<u32>> {
let mut out = Vec::with_capacity(df.height());
for row_idx in 0..df.height() {
let key = row_key(df, match_by, row_idx)?;
let id = match ids.get(&key) {
Some(id) => *id,
None => {
let id = *next_id;
ids.insert(key, id);
*next_id = next_id.saturating_add(1);
id
}
};
out.push(id);
}
Ok(out)
}
fn row_key(df: &DataFrame, match_by: &[String], row_idx: usize) -> Result<RowKey> {
let mut values = Vec::with_capacity(match_by.len());
for column_name in match_by {
let column = df.column(column_name)?;
let value = column.get(row_idx)?;
if value.is_null() {
return Err(RangeFrameError::NullKeyValue {
column: column_name.clone(),
});
}
values.push(to_key_value(column_name, value)?);
}
Ok(RowKey(values))
}
fn to_key_value(column_name: &str, value: AnyValue<'_>) -> Result<KeyValue> {
match value {
AnyValue::Boolean(v) => Ok(KeyValue::Bool(v)),
AnyValue::Int8(v) => Ok(KeyValue::Int(v as i64)),
AnyValue::Int16(v) => Ok(KeyValue::Int(v as i64)),
AnyValue::Int32(v) => Ok(KeyValue::Int(v as i64)),
AnyValue::Int64(v) => Ok(KeyValue::Int(v)),
AnyValue::UInt8(v) => Ok(KeyValue::UInt(v as u64)),
AnyValue::UInt16(v) => Ok(KeyValue::UInt(v as u64)),
AnyValue::UInt32(v) => Ok(KeyValue::UInt(v as u64)),
AnyValue::UInt64(v) => Ok(KeyValue::UInt(v)),
AnyValue::String(v) => Ok(KeyValue::Str(v.to_owned())),
AnyValue::StringOwned(v) => Ok(KeyValue::Str(v.as_str().to_owned())),
other => Err(RangeFrameError::UnsupportedKeyDtype {
column: column_name.to_owned(),
dtype: format!("{other:?}"),
}),
}
}
fn is_signed_integer_dtype(dtype: &DataType) -> bool {
matches!(
dtype,
DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64
)
}
fn is_unsigned_integer_dtype(dtype: &DataType) -> bool {
matches!(
dtype,
DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64
)
}
fn is_string_like_dtype(dtype: &DataType) -> bool {
matches!(
dtype,
DataType::String | DataType::Categorical(_, _) | DataType::Enum(_, _)
)
}
fn is_fast_key_dtype(dtype: &DataType) -> bool {
matches!(
dtype,
DataType::Boolean
| DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64
| DataType::String
| DataType::Categorical(_, _)
| DataType::Enum(_, _)
)
}
#[cfg(test)]
mod tests {
use polars_core::prelude::{Column, DataFrame, DataType, NamedFrom, Series};
use polars_dtype::categorical::Categories;
use super::{factorize, factorize_pair, factorize_pair_by};
fn make_df(columns: Vec<Series>) -> DataFrame {
let columns = columns.into_iter().map(Column::from).collect::<Vec<_>>();
DataFrame::new_infer_height(columns).unwrap()
}
#[test]
fn factorize_single_matches_supported_multi_column_keys() {
let df = make_df(vec![
Series::new("Chrom".into(), &["chr1", "chr2", "chr1", "chr3"]),
Series::new("Sample".into(), &[1_i32, 2, 1, 3]),
]);
let ids = factorize(&df, &["Chrom".to_owned(), "Sample".to_owned()]).unwrap();
assert_eq!(ids, vec![0, 1, 0, 2]);
}
#[test]
fn factorize_pair_matches_supported_multi_column_keys() {
let left = make_df(vec![
Series::new("Chrom".into(), &["chr1", "chr2", "chr1"]),
Series::new("Sample".into(), &[1_i32, 2, 1]),
]);
let right = make_df(vec![
Series::new("Chrom".into(), &["chr2", "chr1", "chr3"]),
Series::new("Sample".into(), &[2_i64, 1, 3]),
]);
let (left_ids, right_ids) =
factorize_pair(&left, &right, &["Chrom".to_owned(), "Sample".to_owned()]).unwrap();
assert_eq!(left_ids, vec![0, 1, 0]);
assert_eq!(right_ids, vec![1, 0, 2]);
}
#[test]
fn factorize_pair_by_matches_differently_named_columns() {
let left = make_df(vec![
Series::new("Chrom".into(), &["chr1", "chr2", "chr1"]),
Series::new("Sample".into(), &[1_i32, 2, 1]),
]);
let right = make_df(vec![
Series::new("chromosome".into(), &["chr2", "chr1", "chr3"]),
Series::new("sample_id".into(), &[2_i64, 1, 3]),
]);
let (left_ids, right_ids) = factorize_pair_by(
&left,
&right,
&["Chrom".to_owned(), "Sample".to_owned()],
&["chromosome".to_owned(), "sample_id".to_owned()],
)
.unwrap();
assert_eq!(left_ids, vec![0, 1, 0]);
assert_eq!(right_ids, vec![1, 0, 2]);
}
#[test]
fn factorize_pair_normalizes_categorical_and_string_keys() {
let left_chrom = Series::new("Chrom".into(), &["chr1", "chr2", "chr1"])
.cast(&DataType::from_categories(Categories::global()))
.unwrap();
let right_chrom = Series::new("Chrom".into(), &["chr2", "chr1", "chr3"]);
let left = make_df(vec![left_chrom]);
let right = make_df(vec![right_chrom]);
let (left_ids, right_ids) = factorize_pair(&left, &right, &["Chrom".to_owned()]).unwrap();
assert_eq!(left_ids, vec![0, 1, 0]);
assert_eq!(right_ids, vec![1, 0, 2]);
}
#[test]
fn factorize_pair_single_integer_key_uses_shared_ids_across_widths() {
let left = make_df(vec![Series::new("Chrom".into(), &[1_i32, 2, 1, 3])]);
let right = make_df(vec![Series::new("Chrom".into(), &[2_i64, 4, 1])]);
let (left_ids, right_ids) = factorize_pair(&left, &right, &["Chrom".to_owned()]).unwrap();
assert_eq!(left_ids, vec![0, 1, 0, 2]);
assert_eq!(right_ids, vec![1, 3, 0]);
}
#[test]
fn factorize_single_column_normalizes_categorical_and_string_keys() {
let chrom = Series::new("Chrom".into(), &["chr1", "chr2", "chr1", "chr3"])
.cast(&DataType::from_categories(Categories::global()))
.unwrap();
let df = make_df(vec![chrom]);
let ids = factorize(&df, &["Chrom".to_owned()]).unwrap();
assert_eq!(ids, vec![0, 1, 0, 2]);
}
}