fix unique/n_unique for Categorical

polars/polars-arrow/Cargo.toml

@@ -10,8 +10,8 @@ description = "Arrow interfaces for Polars DataFrame library"
 
 [dependencies]
 # arrow = { package = "arrow2", git = "https://github.com/jorgecarleitao/arrow2", rev = "e9a6c3ef7e1a328c298bd45e36ac2abf8ae44ebb", default-features = false }
-# arrow = { package = "arrow2", git = "https://github.com/ritchie46/arrow2", default-features = false, features = ["compute"], branch = "fn_to" }
-arrow = { package = "arrow2", version = "0.8", default-features = false }
+arrow = { package = "arrow2", git = "https://github.com/ritchie46/arrow2", default-features = false, features = ["compute"], branch = "offset_pub" }
+# arrow = { package = "arrow2", version = "0.8", default-features = false }
 num = "^0.4"
 thiserror = "^1.0"
 

polars/polars-core/Cargo.toml

@@ -158,7 +158,9 @@ unsafe_unwrap = "^0.1.0"
 
 [dependencies.arrow]
 package = "arrow2"
-version = "0.8"
+git = "https://github.com/ritchie46/arrow2"
+branch = "offset_pub"
+# version = "0.8"
 default-features = false
 features = [
   "compute_aggregate",

polars/polars-core/src/chunked_array/builder/from.rs

@@ -1,4 +1,4 @@
-use crate::prelude::{BooleanChunked, ChunkedArray, PolarsNumericType};
+use crate::prelude::*;
 use arrow::array::{BooleanArray, PrimitiveArray, Utf8Array};
 use std::sync::Arc;
 
@@ -11,6 +11,12 @@ impl<T: PolarsNumericType> From<(&str, PrimitiveArray<T::Native>)> for ChunkedAr
     }
 }
 
+impl<T: PolarsNumericType> From<&[T::Native]> for ChunkedArray<T> {
+    fn from(slice: &[T::Native]) -> Self {
+        ChunkedArray::new_from_slice("", slice)
+    }
+}
+
 impl From<(&str, BooleanArray)> for BooleanChunked {
     fn from(tpl: (&str, BooleanArray)) -> Self {
         let name = tpl.0;

polars/polars-core/src/chunked_array/categorical/builder.rs

@@ -174,12 +174,15 @@ impl CategoricalChunkedBuilder {
     }
 
     pub fn finish(self) -> ChunkedArray<CategoricalType> {
+        // both for the local and the global map, we own a map that has all unique keys
+        let bit_settings = 1u8 << 4;
+
         ChunkedArray {
             field: Arc::new(self.field),
             chunks: vec![self.array_builder.into_arc()],
             phantom: PhantomData,
             categorical_map: Some(Arc::new(self.reverse_mapping.finish())),
-            ..Default::default()
+            bit_settings,
         }
     }
 }

polars/polars-core/src/chunked_array/categorical/mod.rs

@@ -69,6 +69,13 @@ impl CategoricalChunked {
         self
     }
 
+    pub(crate) fn can_fast_unique(&self) -> bool {
+        self.bit_settings & 1 << 4 != 0 && self.chunks.len() == 1 && {
+            let arr = self.downcast_iter().next().unwrap();
+            arr.values().offset() == 0
+        }
+    }
+
     /// Create an `[Iterator]` that iterates over the `&str` values of the `[CategoricalChunked]`.
     pub fn iter_str(&self) -> CatIter<'_> {
         let iter = self.deref().into_iter();
@@ -178,4 +185,18 @@ mod test {
         assert_eq!(appended.str_value(4), "\"x\"");
         assert_eq!(appended.str_value(5), "\"y\"");
     }
+
+    #[test]
+    fn test_fast_unique() {
+        let mut s = Series::new("1", vec!["a", "b", "c"])
+            .cast(&DataType::Categorical)
+            .unwrap();
+
+        assert_eq!(s.n_unique().unwrap(), 3);
+        // make sure that it does not take the fast path after take/ slice
+        let out = s.take(&([1, 2].as_ref()).into()).unwrap();
+        assert_eq!(out.n_unique().unwrap(), 2);
+        let out = s.slice(1, 2);
+        assert_eq!(out.n_unique().unwrap(), 2);
+    }
 }

polars/polars-core/src/chunked_array/mod.rs

@@ -155,6 +155,8 @@ pub struct ChunkedArray<T> {
     /// third bit dtype list: fast_explode
     ///     - unset: unknown or not all arrays have at least one value
     ///     - set: all list arrays are filled (this allows for cheap explode)
+    /// fourth bit: original local categorical
+    ///             meaning that n_unique is the same as the cat map length
     pub(crate) bit_settings: u8,
 }
 

polars/polars-core/src/chunked_array/ops/unique/mod.rs

@@ -264,11 +264,17 @@ impl ChunkUnique<Utf8Type> for Utf8Chunked {
 impl ChunkUnique<CategoricalType> for CategoricalChunked {
     fn unique(&self) -> Result<Self> {
         let cat_map = self.categorical_map.as_ref().unwrap();
-        let mut ca = match &**cat_map {
-            RevMapping::Local(a) => UInt32Chunked::new_from_iter(self.name(), 0..(a.len() as u32)),
-            RevMapping::Global(map, _, _) => {
-                UInt32Chunked::new_from_iter(self.name(), map.keys().copied())
+        let mut ca = if self.can_fast_unique() {
+            match &**cat_map {
+                RevMapping::Local(a) => {
+                    UInt32Chunked::new_from_iter(self.name(), 0..(a.len() as u32))
+                }
+                RevMapping::Global(map, _, _) => {
+                    UInt32Chunked::new_from_iter(self.name(), map.keys().copied())
+                }
             }
+        } else {
+            self.deref().unique()?
         };
         ca.categorical_map = self.categorical_map.clone();
         Ok(ca.into())
@@ -289,7 +295,11 @@ impl ChunkUnique<CategoricalType> for CategoricalChunked {
         impl_value_counts!(self)
     }
     fn n_unique(&self) -> Result<usize> {
-        Ok(self.categorical_map.as_ref().unwrap().len())
+        if self.can_fast_unique() {
+            Ok(self.categorical_map.as_ref().unwrap().len())
+        } else {
+            self.deref().n_unique()
+        }
     }
     #[cfg(feature = "mode")]
     fn mode(&self) -> Result<Self> {

polars/polars-io/Cargo.toml

@@ -32,8 +32,8 @@ private = []
 ahash = "0.7"
 anyhow = "1.0"
 # arrow = { package = "arrow2", git = "https://github.com/jorgecarleitao/arrow2", rev = "e9a6c3ef7e1a328c298bd45e36ac2abf8ae44ebb", default-features = false }
-# arrow = { package = "arrow2", git = "https://github.com/ritchie46/arrow2", default-features = false, features = ["compute"], branch = "fn_to" }
-arrow = { package = "arrow2", version = "0.8", default-features = false }
+arrow = { package = "arrow2", git = "https://github.com/ritchie46/arrow2", default-features = false, features = ["compute"], branch = "offset_pub" }
+# arrow = { package = "arrow2", version = "0.8", default-features = false }
 csv-core = { version = "0.1.10", optional = true }
 dirs = "4.0"
 flate2 = { version = "1", optional = true, default-features = false }