Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Fixed error in dict equality #333

Merged
merged 1 commit into from
Aug 24, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions src/array/dictionary/iterator.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use crate::array::Array;
use crate::bitmap::utils::{zip_validity, ZipValidity};
use crate::scalar::Scalar;
use crate::trusted_len::TrustedLen;

use super::{DictionaryArray, DictionaryKey};
Expand All @@ -23,7 +24,7 @@ impl<'a, K: DictionaryKey> DictionaryValuesIter<'a, K> {
}

impl<'a, K: DictionaryKey> Iterator for DictionaryValuesIter<'a, K> {
type Item = Box<dyn Array>;
type Item = Box<dyn Scalar>;

#[inline]
fn next(&mut self) -> Option<Self::Item> {
Expand Down Expand Up @@ -56,10 +57,10 @@ impl<'a, K: DictionaryKey> DoubleEndedIterator for DictionaryValuesIter<'a, K> {
}

type ValuesIter<'a, K> = DictionaryValuesIter<'a, K>;
type ZipIter<'a, K> = ZipValidity<'a, Box<dyn Array>, ValuesIter<'a, K>>;
type ZipIter<'a, K> = ZipValidity<'a, Box<dyn Scalar>, ValuesIter<'a, K>>;

impl<'a, K: DictionaryKey> IntoIterator for &'a DictionaryArray<K> {
type Item = Option<Box<dyn Array>>;
type Item = Option<Box<dyn Scalar>>;
type IntoIter = ZipIter<'a, K>;

fn into_iter(self) -> Self::IntoIter {
Expand Down
7 changes: 4 additions & 3 deletions src/array/dictionary/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use std::sync::Arc;
use crate::{
bitmap::Bitmap,
datatypes::DataType,
scalar::{new_scalar, Scalar},
types::{NativeType, NaturalDataType},
};

Expand Down Expand Up @@ -90,11 +91,11 @@ impl<K: DictionaryKey> DictionaryArray<K> {
&self.values
}

/// Returns the values of the [`DictionaryArray`].
/// Returns the value of the [`DictionaryArray`] at position `i`.
#[inline]
pub fn value(&self, index: usize) -> Box<dyn Array> {
pub fn value(&self, index: usize) -> Box<dyn Scalar> {
let index = self.keys.value(index).to_usize().unwrap();
self.values.clone().slice(index, 1)
new_scalar(self.values.as_ref(), index)
}
}

Expand Down
11 changes: 10 additions & 1 deletion src/array/equal/dictionary.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,14 @@
use crate::array::{Array, DictionaryArray, DictionaryKey};

pub(super) fn equal<K: DictionaryKey>(lhs: &DictionaryArray<K>, rhs: &DictionaryArray<K>) -> bool {
lhs.data_type() == rhs.data_type() && lhs.len() == rhs.len() && lhs.iter().eq(rhs.iter())
if !(lhs.data_type() == rhs.data_type() && lhs.len() == rhs.len()) {
return false;
};

// if x is not valid and y is but its child is not, the slots are equal.
lhs.iter().zip(rhs.iter()).all(|(x, y)| match (&x, &y) {
(None, Some(y)) => !y.is_valid(),
(Some(x), None) => !x.is_valid(),
_ => x == y,
})
}
65 changes: 53 additions & 12 deletions tests/it/array/equal/dictionary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,55 +4,96 @@ use arrow2::array::*;

use super::test_equal;

fn create_dictionary_array(values: &[&str], keys: &[Option<i16>]) -> DictionaryArray<i16> {
fn create_dictionary_array(values: &[Option<&str>], keys: &[Option<i16>]) -> DictionaryArray<i16> {
let keys = Int16Array::from(keys);
let values = Utf8Array::<i32>::from_slice(values);
let values = Utf8Array::<i32>::from(values);

DictionaryArray::from_data(keys, Arc::new(values))
}

#[test]
fn dictionary_equal() {
// (a, b, c), (0, 1, 0, 2) => (a, b, a, c)
let a = create_dictionary_array(&["a", "b", "c"], &[Some(0), Some(1), Some(0), Some(2)]);
let a = create_dictionary_array(
&[Some("a"), Some("b"), Some("c")],
&[Some(0), Some(1), Some(0), Some(2)],
);
// different representation (values and keys are swapped), same result
let b = create_dictionary_array(&["a", "c", "b"], &[Some(0), Some(2), Some(0), Some(1)]);
let b = create_dictionary_array(
&[Some("a"), Some("c"), Some("b")],
&[Some(0), Some(2), Some(0), Some(1)],
);
test_equal(&a, &b, true);

// different len
let b = create_dictionary_array(&["a", "c", "b"], &[Some(0), Some(2), Some(1)]);
let b = create_dictionary_array(
&[Some("a"), Some("c"), Some("b")],
&[Some(0), Some(2), Some(1)],
);
test_equal(&a, &b, false);

// different key
let b = create_dictionary_array(&["a", "c", "b"], &[Some(0), Some(2), Some(0), Some(0)]);
let b = create_dictionary_array(
&[Some("a"), Some("c"), Some("b")],
&[Some(0), Some(2), Some(0), Some(0)],
);
test_equal(&a, &b, false);

// different values, same keys
let b = create_dictionary_array(&["a", "b", "d"], &[Some(0), Some(1), Some(0), Some(2)]);
let b = create_dictionary_array(
&[Some("a"), Some("b"), Some("d")],
&[Some(0), Some(1), Some(0), Some(2)],
);
test_equal(&a, &b, false);
}

#[test]
fn dictionary_equal_null() {
// (a, b, c), (1, 2, 1, 3) => (a, b, a, c)
let a = create_dictionary_array(&["a", "b", "c"], &[Some(0), None, Some(0), Some(2)]);
let a = create_dictionary_array(
&[Some("a"), Some("b"), Some("c")],
&[Some(0), None, Some(0), Some(2)],
);

// equal to self
test_equal(&a, &a, true);

// different representation (values and keys are swapped), same result
let b = create_dictionary_array(&["a", "c", "b"], &[Some(0), None, Some(0), Some(1)]);
let b = create_dictionary_array(
&[Some("a"), Some("c"), Some("b")],
&[Some(0), None, Some(0), Some(1)],
);
test_equal(&a, &b, true);

// different null position
let b = create_dictionary_array(&["a", "c", "b"], &[Some(0), Some(2), Some(0), None]);
let b = create_dictionary_array(
&[Some("a"), Some("c"), Some("b")],
&[Some(0), Some(2), Some(0), None],
);
test_equal(&a, &b, false);

// different key
let b = create_dictionary_array(&["a", "c", "b"], &[Some(0), None, Some(0), Some(0)]);
let b = create_dictionary_array(
&[Some("a"), Some("c"), Some("b")],
&[Some(0), None, Some(0), Some(0)],
);
test_equal(&a, &b, false);

// different values, same keys
let b = create_dictionary_array(&["a", "b", "d"], &[Some(0), None, Some(0), Some(2)]);
let b = create_dictionary_array(
&[Some("a"), Some("b"), Some("d")],
&[Some(0), None, Some(0), Some(2)],
);
test_equal(&a, &b, false);

// different nulls in keys and values
let a = create_dictionary_array(
&[Some("a"), Some("b"), None],
&[Some(0), None, Some(0), Some(2)],
);
let b = create_dictionary_array(
&[Some("a"), Some("b"), Some("c")],
&[Some(0), None, Some(0), None],
);
test_equal(&a, &b, true);
}