diff --git a/Cargo.toml b/Cargo.toml index 2c82857..41b6cce 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -98,7 +98,7 @@ extended_numeric_types = [] # Adds a cube object for stacking tables on an extra axis # Useful for time series, and group analytics -cube = [] +cube = ["hash"] # Adds a unified scalar type, that's useful for `Array` aggregations, and other use cases where you end up with one value. # However, it is one of several downcasting methods available in Rust, and when predominantly diff --git a/minarrow-py/Cargo.toml b/minarrow-py/Cargo.toml index e237a1d..e70d16a 100644 --- a/minarrow-py/Cargo.toml +++ b/minarrow-py/Cargo.toml @@ -25,7 +25,7 @@ pyo3 = { version = "0.29", features = ["abi3-py39"] } thiserror = "2" [features] -default = ["datetime", "large_string", "scalar_type", "value_type", "cube", "arrow_interop", "ndarray"] +default = ["datetime", "large_string", "matrix", "scalar_type", "value_type", "cube", "arrow_interop", "ndarray"] extension-module = ["pyo3/extension-module"] arrow_interop = ["dep:minarrow-pyo3", "minarrow-pyo3/datetime"] simd = ["minarrow/simd"] @@ -52,6 +52,8 @@ extended_numeric_types = ["minarrow/extended_numeric_types", "minarrow-pyo3?/ext large_string = ["minarrow/large_string"] scalar_type = ["minarrow/scalar_type"] value_type = ["minarrow/value_type"] +# Dense column-major f64 buffer for the routines that hand memory to LAPACK. +matrix = ["minarrow/matrix"] cube = ["minarrow/cube"] default_categorical_8 = ["minarrow/default_categorical_8", "minarrow-pyo3?/default_categorical_8"] # Mirrors the core crate, where `extended_categorical` implies `default_categorical_8`. diff --git a/minarrow-py/src/convert.rs b/minarrow-py/src/convert.rs index f772395..5862f2f 100644 --- a/minarrow-py/src/convert.rs +++ b/minarrow-py/src/convert.rs @@ -129,10 +129,18 @@ pub fn parse_dtype(name: &str) -> PyResult { "string" | "str" | "utf8" | "str32" => ArrowType::String, "large_string" | "largestring" | "str64" => ArrowType::LargeString, "bool" | "boolean" => ArrowType::Boolean, - #[cfg(any(not(feature = "default_categorical_8"), feature = "extended_categorical"))] - "categorical" | "category" | "cat" | "cat32" => { - ArrowType::Dictionary(CategoricalIndexType::UInt32) + "categorical" | "category" | "cat" => { + #[cfg(feature = "default_categorical_8")] + { + ArrowType::Dictionary(CategoricalIndexType::UInt8) + } + #[cfg(not(feature = "default_categorical_8"))] + { + ArrowType::Dictionary(CategoricalIndexType::UInt32) + } } + #[cfg(any(not(feature = "default_categorical_8"), feature = "extended_categorical"))] + "cat32" => ArrowType::Dictionary(CategoricalIndexType::UInt32), "cat8" => { #[cfg(feature = "default_categorical_8")] { diff --git a/minarrow-py/src/cube.rs b/minarrow-py/src/cube.rs new file mode 100644 index 0000000..32d20ed --- /dev/null +++ b/minarrow-py/src/cube.rs @@ -0,0 +1,190 @@ +// Copyright 2025 Peter Garfield Bower +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +//! A named set of tables over a common schema. +//! +//! Each table is a group, a period or a partition, and carries its own name. +//! Looking one up by name is O(1) rather than a scan. +//! +//! `group_by` is the leading case: one table per distinct key, named by that +//! key, so `cube["BTC"]` reads the group. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use minarrow::ffi::schema::Schema; +use minarrow::{Cube, Field, Table}; +use pyo3::exceptions::{PyIndexError, PyKeyError, PyTypeError}; +use pyo3::prelude::*; + +use crate::field::PySchema; +use crate::table::{PyTable, PyTableInner}; + +/// A named set of tables over a common schema. +#[pyclass(name = "Cube", module = "minarrow")] +pub struct PyCube(pub Arc); + +impl From for PyCube { + fn from(cube: Cube) -> Self { + PyCube(Arc::new(cube)) + } +} + +/// The table at `index`, as the Python object carrying it. +fn table_at(cube: &Cube, index: usize) -> Option { + cube.table(index) + .map(|t| PyTable(PyTableInner::Owned(t.clone()))) +} + +#[pymethods] +impl PyCube { + /// Build a cube from a list of tables, each keeping its own name. + /// + /// `index_by` names the columns forming the third dimension, e.g. the time + /// column when the tables are periods of a series. + #[new] + #[pyo3(signature = (tables, name=None, index_by=None))] + fn new( + tables: Vec>, + name: Option, + index_by: Option>, + ) -> PyResult { + let tables: Vec = tables + .iter() + .map(|t| t.borrow().0.as_view().to_table()) + .collect(); + let cube = Cube::from_tables(tables, name.unwrap_or_default(), index_by); + Ok(PyCube(Arc::new(cube))) + } + + #[getter] + fn name(&self) -> &str { + &self.0.name + } + + /// The number of tables the cube holds. + #[getter] + fn n_tables(&self) -> usize { + self.0.n_tables() + } + + /// The number of columns each table carries. + #[getter] + fn n_cols(&self) -> usize { + self.0.n_cols() + } + + /// The row count of each table, in order. + #[getter] + fn n_rows(&self) -> Vec { + self.0.n_rows() + } + + /// The name of each table, in order. For a grouped cube, the group keys. + #[getter] + fn names(&self) -> Vec { + self.0.table_names().into_iter().map(str::to_string).collect() + } + + #[getter] + fn columns(&self) -> Vec { + self.0.col_names().into_iter().map(str::to_string).collect() + } + + #[getter] + fn schema(&self) -> PySchema { + let fields: Vec = self.0.schema().iter().map(|f| (**f).clone()).collect(); + PySchema(Schema::new(fields, BTreeMap::new())) + } + + /// Every table, in order. + #[getter] + fn tables(&self) -> Vec { + (0..self.0.n_tables()) + .filter_map(|i| table_at(&self.0, i)) + .collect() + } + + /// The table at a position, or `None` where there is none. + fn table(&self, index: usize) -> Option { + table_at(&self.0, index) + } + + /// The number of tables. + fn __len__(&self) -> usize { + self.0.n_tables() + } + + /// Read a table by position with `cube[0]`, or by name with `cube["BTC"]`. + fn __getitem__(&self, key: &Bound<'_, PyAny>) -> PyResult { + if let Ok(name) = key.extract::() { + let name = name.as_str(); + let index = self.0.table_index(name).ok_or_else(|| { + PyKeyError::new_err(format!("no table named '{name}' in this cube")) + })?; + return table_at(&self.0, index).ok_or_else(|| { + PyKeyError::new_err(format!("no table named '{name}' in this cube")) + }); + } + if let Ok(index) = key.extract::() { + return table_at(&self.0, index).ok_or_else(|| { + PyIndexError::new_err(format!( + "table {index} is outside a cube of {} tables", + self.0.n_tables() + )) + }); + } + Err(PyTypeError::new_err( + "index a cube with a position or with a table name", + )) + } + + /// The table for a key, looked up on the key's own type. + /// + /// A cube split on an `Int32` column is reached with an `int`, and on a + /// `Float64` column with a `float`. This is the lookup a grouped cube is + /// built for; `cube[i]` reads by position and `cube["name"]` by name. + fn group(&self, key: &Bound<'_, PyAny>) -> PyResult { + let scalar = crate::convert::py_to_scalar(key)?; + let index = self + .0 + .resolve(&scalar) + .ok_or_else(|| PyKeyError::new_err(format!("no group keyed {scalar} in this cube")))?; + table_at(&self.0, index) + .ok_or_else(|| PyKeyError::new_err(format!("no group keyed {scalar} in this cube"))) + } + + /// Whether a table of this name is present. + fn __contains__(&self, name: &str) -> bool { + self.0.has_table(name) + } + + fn __iter__(slf: PyRef<'_, Self>) -> PyResult> { + let tables: Vec = slf.tables(); + let py = slf.py(); + let list = pyo3::types::PyList::new(py, tables)?; + Ok(list.as_any().try_iter()?.into_any().unbind()) + } + + fn __repr__(&self) -> String { + let rows: usize = self.0.n_rows().iter().sum(); + format!( + "Cube(name: {}, tables: {}, rows: {}, cols: {})", + self.0.name, + self.0.n_tables(), + rows, + self.0.n_cols() + ) + } +} diff --git a/minarrow-py/src/lib.rs b/minarrow-py/src/lib.rs index 41e0b95..02e88f6 100644 --- a/minarrow-py/src/lib.rs +++ b/minarrow-py/src/lib.rs @@ -25,8 +25,12 @@ mod chunked_array; mod chunked_ndarray; mod chunked_table; mod convert; +#[cfg(feature = "cube")] +mod cube; mod dtype; mod field; +#[cfg(feature = "matrix")] +mod matrix; #[cfg(feature = "ndarray")] mod ndarray; #[cfg(feature = "embed")] @@ -49,8 +53,12 @@ pub use chunked_array::PyChunkedArray; pub use chunked_ndarray::{PyChunkedNdArray, PyChunkedNdArrayInner}; pub use chunked_table::PyChunkedTable; pub use convert::{build_array, py_to_scalar, resolve_index, scalar_to_py}; +#[cfg(feature = "cube")] +pub use cube::PyCube; pub use dtype::{dtype_from_arrow, width_from_arrow, DType, TypeClass}; pub use field::{PyField, PySchema}; +#[cfg(feature = "matrix")] +pub use matrix::PyMatrix; #[cfg(feature = "ndarray")] pub use ndarray::{PyNdArray, PyNdArrayInner}; #[cfg(feature = "embed")] @@ -87,6 +95,10 @@ fn minarrow_py(m: &Bound<'_, PyModule>) -> PyResult<()> { } m.add_class::()?; m.add_class::()?; + #[cfg(feature = "cube")] + m.add_class::()?; + #[cfg(feature = "matrix")] + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/minarrow-py/src/matrix.rs b/minarrow-py/src/matrix.rs new file mode 100644 index 0000000..83c4730 --- /dev/null +++ b/minarrow-py/src/matrix.rs @@ -0,0 +1,234 @@ +// Copyright 2025 Peter Garfield Bower +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Dense `f64` matrix for the routines that hand memory to LAPACK. +//! +//! A `Table` holds each column in its own allocation, which suits scanning and +//! appending. Regression, decomposition and clustering instead want one +//! contiguous column-major buffer with a known stride, so they can pass +//! `(pointer, leading dimension)` straight to BLAS without repacking. `Matrix` +//! is that layout, and `from_table` is the boundary a caller crosses to reach +//! it. +//! +//! Columns are padded so each begins on a 64-byte boundary, which is why +//! `stride` may exceed `n_rows`. The padding is internal, so `n_rows`, +//! `n_cols`, indexing and iteration all read the logical shape. + +use minarrow::{Matrix, Table}; +use pyo3::exceptions::{PyIndexError, PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::PyTuple; + +use crate::table::{PyTable, PyTableInner}; + +/// A dense, column-major `f64` matrix. +#[pyclass(name = "Matrix", module = "minarrow")] +#[derive(Clone)] +pub struct PyMatrix(pub Matrix); + +impl From for PyMatrix { + fn from(matrix: Matrix) -> Self { + PyMatrix(matrix) + } +} + +#[pymethods] +impl PyMatrix { + /// Build a matrix from rows, each row a list of the same length. + /// + /// Rows are the form a literal is normally written in, matching the reading + /// `[[1, 2], [3, 4]]` gets elsewhere in Python. The buffer underneath is + /// column-major regardless, so `from_cols` skips the transposition when the + /// caller already holds columns. + #[new] + #[pyo3(signature = (rows, name=None))] + fn new(rows: Vec>, name: Option) -> PyResult { + if rows.is_empty() { + return Err(PyValueError::new_err("Matrix requires at least one row")); + } + let n_cols = rows[0].len(); + if n_cols == 0 { + return Err(PyValueError::new_err("Matrix requires at least one column")); + } + if let Some((i, row)) = rows.iter().enumerate().find(|(_, r)| r.len() != n_cols) { + return Err(PyValueError::new_err(format!( + "row {i} has {} values against {n_cols} in the first row", + row.len() + ))); + } + let cols: Vec> = (0..n_cols) + .map(|c| rows.iter().map(|r| r[c]).collect()) + .collect(); + let mut matrix = Matrix::from(cols.as_slice()); + if let Some(name) = name { + matrix.set_name(name); + } + Ok(PyMatrix(matrix)) + } + + /// Build a matrix from columns, which is the layout it already stores. + #[staticmethod] + #[pyo3(signature = (cols, name=None))] + fn from_cols(cols: Vec>, name: Option) -> PyResult { + if cols.is_empty() { + return Err(PyValueError::new_err("Matrix requires at least one column")); + } + let n_rows = cols[0].len(); + if let Some((i, col)) = cols.iter().enumerate().find(|(_, c)| c.len() != n_rows) { + return Err(PyValueError::new_err(format!( + "column {i} has {} values against {n_rows} in the first column", + col.len() + ))); + } + let mut matrix = Matrix::from(cols.as_slice()); + if let Some(name) = name { + matrix.set_name(name); + } + Ok(PyMatrix(matrix)) + } + + /// Build a matrix from a table's columns. + /// + /// Every column must be `f64` and free of nulls, since the layout has + /// nowhere to record a missing value and the routines reading it assume + /// each entry is a number. A table that does not meet this says which + /// column disagreed. + #[staticmethod] + fn from_table(table: PyRef<'_, PyTable>) -> PyResult { + Matrix::try_from(&table.0.as_view()) + .map(PyMatrix) + .map_err(|e| PyValueError::new_err(format!("{e}"))) + } + + #[getter] + fn n_rows(&self) -> usize { + self.0.n_rows + } + + #[getter] + fn n_cols(&self) -> usize { + self.0.n_cols + } + + #[getter] + fn shape<'py>(&self, py: Python<'py>) -> PyResult> { + PyTuple::new(py, [self.0.n_rows, self.0.n_cols]) + } + + /// The physical distance between the start of one column and the next. + /// + /// This is `n_rows` rounded up to a 64-byte boundary, and it is what BLAS + /// calls the leading dimension. + #[getter] + fn stride(&self) -> usize { + self.0.stride + } + + #[getter] + fn name(&self) -> Option<&str> { + self.0.name.as_deref() + } + + /// The number of rows, so a matrix reads like a sequence of them. + fn __len__(&self) -> usize { + self.0.n_rows + } + + /// Read a single entry with `m[row, col]`, or a whole row with `m[row]`. + fn __getitem__(&self, py: Python<'_>, key: &Bound<'_, PyAny>) -> PyResult> { + if let Ok((row, col)) = key.extract::<(usize, usize)>() { + if row >= self.0.n_rows || col >= self.0.n_cols { + return Err(PyIndexError::new_err(format!( + "({row}, {col}) is outside a {} by {} matrix", + self.0.n_rows, self.0.n_cols + ))); + } + return Ok(self.0.get(row, col).into_pyobject(py)?.into_any().unbind()); + } + if let Ok(row) = key.extract::() { + if row >= self.0.n_rows { + return Err(PyIndexError::new_err(format!( + "row {row} is outside a matrix of {} rows", + self.0.n_rows + ))); + } + return Ok(self.0.row(row).into_pyobject(py)?.into_any().unbind()); + } + Err(PyTypeError::new_err( + "index a matrix with a row, or with a (row, column) pair", + )) + } + + /// One column as a list of values. + fn col(&self, col: usize) -> PyResult> { + if col >= self.0.n_cols { + return Err(PyIndexError::new_err(format!( + "column {col} is outside a matrix of {} columns", + self.0.n_cols + ))); + } + Ok(self.0.col(col).to_vec()) + } + + /// One row as a list of values, gathered across the columns. + fn row(&self, row: usize) -> PyResult> { + if row >= self.0.n_rows { + return Err(PyIndexError::new_err(format!( + "row {row} is outside a matrix of {} rows", + self.0.n_rows + ))); + } + Ok(self.0.row(row)) + } + + /// The matrix with rows and columns exchanged, in a fresh buffer. + fn transpose(&self) -> Self { + PyMatrix(self.0.transpose()) + } + + #[getter] + fn T(&self) -> Self { + PyMatrix(self.0.transpose()) + } + + /// The rows at the given positions, in the order given. + fn extract_rows(&self, indices: Vec) -> PyResult { + if let Some(&i) = indices.iter().find(|&&i| i >= self.0.n_rows) { + return Err(PyIndexError::new_err(format!( + "row {i} is outside a matrix of {} rows", + self.0.n_rows + ))); + } + Ok(PyMatrix(self.0.extract_rows(&indices))) + } + + /// The same values as a table, each column named by its position. + fn to_table(&self) -> PyTable { + let table: Table = self.0.clone().to_table_gen(); + PyTable(PyTableInner::from(table)) + } + + fn __repr__(&self) -> String { + match &self.0.name { + Some(name) => format!( + "Matrix('{name}', {} rows x {} cols, f64)", + self.0.n_rows, self.0.n_cols + ), + None => format!( + "Matrix({} rows x {} cols, f64)", + self.0.n_rows, self.0.n_cols + ), + } + } +} diff --git a/src/structs/cube.rs b/src/structs/cube.rs index 1709390..b38bc50 100644 --- a/src/structs/cube.rs +++ b/src/structs/cube.rs @@ -50,7 +50,7 @@ use crate::aliases::CubeV; use crate::enums::{error::MinarrowError, shape_dim::ShapeDim}; use crate::ffi::arrow_dtype::ArrowType; use crate::traits::{concatenate::Concatenate, shape::Shape}; -use crate::{Field, Table}; +use crate::{Field, Scalar, Table}; // Global counter for unnamed cube instances static UNNAMED_COUNTER: AtomicUsize = AtomicUsize::new(1); @@ -86,7 +86,11 @@ pub struct Cube { pub third_dim_index: Option>, // O(1) lookup from group key to position in the tables vec. // Maintained alongside tables for fast group resolution across batches. - pub resolver: HashMap, + // + // The key is a `Scalar`, so a group taken from an `Int32` column resolves + // on an `Int32` and hands that type back. A table that carries a name + // rather than a key resolves on the name as a `String32`. + pub resolver: HashMap, } impl Cube { @@ -113,7 +117,7 @@ impl Cube { if let Some(cols) = cols { let table = Table::new(name.clone(), Some(cols)); - resolver.insert(table.name.clone(), 0); + resolver.insert(Scalar::String32(table.name.clone()), 0); tables.push(Arc::new(table)); } @@ -135,7 +139,7 @@ impl Cube { let arc_tables: Vec> = tables.into_iter().map(Arc::new).collect(); let mut resolver = HashMap::new(); for (i, t) in arc_tables.iter().enumerate() { - resolver.insert(t.name.clone(), i); + resolver.insert(Scalar::String32(t.name.clone()), i); } Self { tables: arc_tables, @@ -145,6 +149,65 @@ impl Cube { } } + /// Splits a table into one table per distinct value of a key column. + /// + /// The column is read in order, and each value not seen before opens a + /// group. Every row then joins the group its key names, so the tables come + /// back in the order their keys first appear. + /// + /// The key column is dropped from each group, since it would hold one value + /// repeated down every row. That value is the group's key in the resolver + /// and its name, so `cube.resolve(&Scalar::Int32(2))` reaches the group + /// without a scan and without the key losing its type. + /// + /// The column name is recorded as the cube's third dimension. + pub fn from_table( + table: &Table, + key_col: &str, + name: impl Into, + ) -> Result { + let key_idx = table + .col_name_index(key_col) + .ok_or_else(|| MinarrowError::ShapeError { + message: format!("Cube::from_table: key column '{key_col}' not found"), + })?; + let key_array = &table.cols[key_idx].array; + + // First appearance decides a group's position, and the rows of each + // group are gathered before any table is built so a row is visited once. + let mut keys: Vec = Vec::new(); + let mut rows: Vec> = Vec::new(); + let mut seen: HashMap = HashMap::new(); + for row in 0..table.n_rows { + let key = key_array.get_scalar(row).unwrap_or(Scalar::Null); + match seen.get(&key) { + Some(&group) => rows[group].push(row), + None => { + seen.insert(key.clone(), keys.len()); + keys.push(key); + rows.push(vec![row]); + } + } + } + + let mut tables = Vec::with_capacity(keys.len()); + let mut resolver = HashMap::with_capacity(keys.len()); + for (group, key) in keys.into_iter().enumerate() { + let mut group_table = table.gather_rows(&rows[group]); + group_table.remove_col(key_col); + group_table.set_name(key.to_string()); + resolver.insert(key, group); + tables.push(Arc::new(group_table)); + } + + Ok(Self { + tables, + name: name.into(), + third_dim_index: Some(vec![key_col.to_string()]), + resolver, + }) + } + /// Constructs a new, empty Cube with a globally unique name. pub fn new_empty() -> Self { let id = UNNAMED_COUNTER.fetch_add(1, Ordering::Relaxed); @@ -182,7 +245,7 @@ impl Cube { } let idx = self.tables.len(); - self.resolver.insert(table.name.clone(), idx); + self.resolver.insert(Scalar::String32(table.name.clone()), idx); self.tables.push(Arc::new(table)); } @@ -248,7 +311,7 @@ impl Cube { pub fn remove_table_at(&mut self, idx: usize) -> bool { if idx < self.tables.len() { let removed = self.tables.remove(idx); - self.resolver.remove(&removed.name); + self.resolver.remove(&Scalar::String32(removed.name.clone())); // Rebuild resolver indices for entries after the removed position for (_, pos) in self.resolver.iter_mut() { if *pos > idx { @@ -263,7 +326,7 @@ impl Cube { /// Removes a table by name. pub fn remove_table(&mut self, name: &str) -> bool { - if let Some(idx) = self.resolver.remove(name) { + if let Some(idx) = self.resolver.remove(&Scalar::String32(name.to_string())) { self.tables.remove(idx); // Rebuild resolver indices for entries after the removed position for (_, pos) in self.resolver.iter_mut() { @@ -386,15 +449,24 @@ impl Cube { } /// Resolve a group key to its table position. - pub fn resolve(&self, key: &str) -> Option { + /// + /// The key is the value the group was built on, so an `Int32` column + /// resolves on an `Int32`. A table carrying a name rather than a key + /// resolves on `Scalar::String32`, which [`Self::resolve_name`] wraps. + pub fn resolve(&self, key: &Scalar) -> Option { self.resolver.get(key).copied() } + /// Resolve a table position by name. + pub fn resolve_name(&self, name: &str) -> Option { + self.resolve(&Scalar::String32(name.to_string())) + } + /// Rebuild the resolver from the current tables vec. pub fn rebuild_resolver(&mut self) { self.resolver.clear(); for (i, t) in self.tables.iter().enumerate() { - self.resolver.insert(t.name.clone(), i); + self.resolver.insert(Scalar::String32(t.name.clone()), i); } } @@ -469,7 +541,7 @@ impl Cube { .collect(); let mut resolver = HashMap::new(); for (i, t) in tables.iter().enumerate() { - resolver.insert(t.name.clone(), i); + resolver.insert(Scalar::String32(t.name.clone()), i); } let name = format!("{}[{}, {})", self.name, offset, offset + len); Cube { @@ -638,7 +710,7 @@ impl Concatenate for Cube { result_tables.extend(other.tables); let mut resolver = HashMap::new(); for (i, t) in result_tables.iter().enumerate() { - resolver.insert(t.name.clone(), i); + resolver.insert(Scalar::String32(t.name.clone()), i); } Ok(Cube { @@ -699,7 +771,7 @@ impl crate::traits::selection::ColumnSelection for Cube { let mut resolver = HashMap::new(); for (i, t) in tables.iter().enumerate() { - resolver.insert(t.name.clone(), i); + resolver.insert(Scalar::String32(t.name.clone()), i); } Cube { tables, @@ -710,8 +782,8 @@ impl crate::traits::selection::ColumnSelection for Cube { } fn get(&self, field: &str) -> Option> { - let idx = self.resolver.get(field)?; - self.tables.get(*idx).cloned() + let idx = self.resolve_name(field)?; + self.tables.get(idx).cloned() } fn col_ix(&self, idx: usize) -> Option { @@ -1006,7 +1078,7 @@ mod tests { tables: vec![Arc::new(table)], name: "test".to_string(), third_dim_index: Some(vec!["timestamp".to_string()]), - resolver: HashMap::from([("single".to_string(), 0)]), + resolver: HashMap::from([(Scalar::String32("single".to_string()), 0)]), }; assert_eq!(cube.n_tables(), 1); assert_eq!(cube.n_cols(), 2); @@ -1072,4 +1144,32 @@ mod tests { names.sort_unstable(); assert_eq!(names, vec!["a", "b"]); } + + /// Splitting on a key gives one table per distinct value, in first-seen + /// order, with the key column gone and the key resolving on its own type. + #[test] + fn from_table_splits_on_a_key_column() { + let table = build_test_table("t", &[2, 1, 2, 3, 1], &[true, false, true, false, true]); + let cube = Cube::from_table(&table, "ints", "grouped").unwrap(); + + assert_eq!(cube.n_tables(), 3); + assert_eq!(cube.table_names(), vec!["2", "1", "3"], "first-seen order"); + assert_eq!(cube.n_rows(), vec![2, 2, 1]); + + // The key column is gone, and only it. + assert_eq!(cube.col_names(), vec!["bools"]); + + // The key resolves as the type it was read from, not as text. + let at = cube.resolve(&Scalar::Int32(2)).expect("Int32 key resolves"); + assert_eq!(cube.tables[at].n_rows, 2); + assert_eq!(cube.resolve(&Scalar::Int32(9)), None); + assert_eq!( + cube.resolve(&Scalar::String32("2".to_string())), + None, + "an Int32 key is not the same key as the text that prints it" + ); + + assert_eq!(cube.third_dim_index().unwrap(), &["ints"]); + } + }