Skip to content

Commit

Permalink
Merge branch 'develop' into ngates/arrow-random-access
Browse files Browse the repository at this point in the history
  • Loading branch information
gatesn committed Mar 26, 2024
2 parents ab5e8ec + 704d3f3 commit 1dac391
Show file tree
Hide file tree
Showing 8 changed files with 39 additions and 64 deletions.
27 changes: 18 additions & 9 deletions pyvortex/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use pyo3::prelude::*;

use dtype::PyDType;
use vortex_schema::DType;
use vortex_schema::Signedness::{Signed, Unsigned};

use crate::array::*;
use crate::compress::PyCompressConfig;
Expand Down Expand Up @@ -56,6 +57,7 @@ fn _lib(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_class::<PyCompressConfig>()?;

m.add_function(wrap_pyfunction!(dtype_int, m)?)?;
m.add_function(wrap_pyfunction!(dtype_uint, m)?)?;
m.add_function(wrap_pyfunction!(dtype_float, m)?)?;
m.add_function(wrap_pyfunction!(dtype_bool, m)?)?;
m.add_function(wrap_pyfunction!(dtype_utf8, m)?)?;
Expand All @@ -70,23 +72,30 @@ fn dtype_bool(py: Python<'_>, nullable: bool) -> PyResult<Py<PyDType>> {
}

#[pyfunction(name = "int")]
#[pyo3(signature = (width = None, signed = true, nullable = false))]
fn dtype_int(
py: Python<'_>,
width: Option<u16>,
signed: bool,
nullable: bool,
) -> PyResult<Py<PyDType>> {
#[pyo3(signature = (width = None, nullable = false))]
fn dtype_int(py: Python<'_>, width: Option<u16>, nullable: bool) -> PyResult<Py<PyDType>> {
PyDType::wrap(
py,
DType::Int(width.unwrap_or(64).into(), Signed, nullable.into()),
)
}

#[pyfunction(name = "uint")]
#[pyo3(signature = (width = None, nullable = false))]
fn dtype_uint(py: Python<'_>, width: Option<u16>, nullable: bool) -> PyResult<Py<PyDType>> {
PyDType::wrap(
py,
DType::Int(width.unwrap_or(0).into(), signed.into(), nullable.into()),
DType::Int(width.unwrap_or(64).into(), Unsigned, nullable.into()),
)
}

#[pyfunction(name = "float")]
#[pyo3(signature = (width = None, nullable = false))]
fn dtype_float(py: Python<'_>, width: Option<i8>, nullable: bool) -> PyResult<Py<PyDType>> {
PyDType::wrap(py, DType::Float(width.unwrap_or(0).into(), nullable.into()))
PyDType::wrap(
py,
DType::Float(width.unwrap_or(64).into(), nullable.into()),
)
}

#[pyfunction(name = "utf8")]
Expand Down
7 changes: 4 additions & 3 deletions pyvortex/test/test_dtype.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@


def test_int():
assert str(vortex.int()) == "signed_int(_)"
assert str(vortex.int(32)) == "signed_int(32)"
assert str(vortex.int(32, signed=False)) == "unsigned_int(32)"
assert str(vortex.int()) == "int(64)"
assert str(vortex.int(32)) == "int(32)"
assert str(vortex.int(32, nullable=True)) == "int(32)?"
assert str(vortex.uint(32)) == "uint(32)"
assert str(vortex.float(16)) == "float(16)"
assert str(vortex.bool(nullable=True)) == "bool?"
7 changes: 2 additions & 5 deletions vortex-array/src/formatter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -120,15 +120,12 @@ mod test {
#[test]
fn display_primitive() {
let arr: ArrayRef = (0..100).collect::<Vec<i32>>().into_array();
assert_eq!(
format!("{}", arr),
"vortex.primitive(signed_int(32), len=100)"
);
assert_eq!(format!("{}", arr), "vortex.primitive(int(32), len=100)");
}

#[test]
fn tree_display_primitive() {
let arr: ArrayRef = (0..100).collect::<Vec<i32>>().into_array();
assert_eq!(display_tree(&arr), "root: vortex.primitive(signed_int(32), len=100) nbytes=400 B (100.00%)\n values: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]...\n validity: None\n")
assert_eq!(display_tree(&arr), "root: vortex.primitive(int(32), len=100) nbytes=400 B (100.00%)\n values: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]...\n validity: None\n")
}
}
11 changes: 4 additions & 7 deletions vortex-array/src/ptype.rs
Original file line number Diff line number Diff line change
Expand Up @@ -182,19 +182,16 @@ impl TryFrom<&DType> for PType {
use vortex_schema::Signedness::*;
match value {
Int(w, s, _) => match (w, s) {
(IntWidth::Unknown, Unknown | Signed) => Ok(PType::I64),
(IntWidth::_8, Unknown | Signed) => Ok(PType::I8),
(IntWidth::_16, Unknown | Signed) => Ok(PType::I16),
(IntWidth::_32, Unknown | Signed) => Ok(PType::I32),
(IntWidth::_64, Unknown | Signed) => Ok(PType::I64),
(IntWidth::Unknown, Unsigned) => Ok(PType::U64),
(IntWidth::_8, Signed) => Ok(PType::I8),
(IntWidth::_16, Signed) => Ok(PType::I16),
(IntWidth::_32, Signed) => Ok(PType::I32),
(IntWidth::_64, Signed) => Ok(PType::I64),
(IntWidth::_8, Unsigned) => Ok(PType::U8),
(IntWidth::_16, Unsigned) => Ok(PType::U16),
(IntWidth::_32, Unsigned) => Ok(PType::U32),
(IntWidth::_64, Unsigned) => Ok(PType::U64),
},
Float(f, _) => match f {
FloatWidth::Unknown => Ok(PType::F64),
FloatWidth::_16 => Ok(PType::F16),
FloatWidth::_32 => Ok(PType::F32),
FloatWidth::_64 => Ok(PType::F64),
Expand Down
21 changes: 4 additions & 17 deletions vortex-array/src/scalar/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -110,30 +110,17 @@ impl Scalar {
DType::Null => NullScalar::new().into(),
DType::Bool(_) => BoolScalar::none().into(),
DType::Int(w, s, _) => match (w, s) {
(IntWidth::Unknown, Signedness::Unknown | Signedness::Signed) => {
PrimitiveScalar::none::<i64>().into()
}
(IntWidth::_8, Signedness::Unknown | Signedness::Signed) => {
PrimitiveScalar::none::<i8>().into()
}
(IntWidth::_16, Signedness::Unknown | Signedness::Signed) => {
PrimitiveScalar::none::<i16>().into()
}
(IntWidth::_32, Signedness::Unknown | Signedness::Signed) => {
PrimitiveScalar::none::<i32>().into()
}
(IntWidth::_64, Signedness::Unknown | Signedness::Signed) => {
PrimitiveScalar::none::<i64>().into()
}
(IntWidth::Unknown, Signedness::Unsigned) => PrimitiveScalar::none::<u64>().into(),
(IntWidth::_8, Signedness::Signed) => PrimitiveScalar::none::<i8>().into(),
(IntWidth::_16, Signedness::Signed) => PrimitiveScalar::none::<i16>().into(),
(IntWidth::_32, Signedness::Signed) => PrimitiveScalar::none::<i32>().into(),
(IntWidth::_64, Signedness::Signed) => PrimitiveScalar::none::<i64>().into(),
(IntWidth::_8, Signedness::Unsigned) => PrimitiveScalar::none::<u8>().into(),
(IntWidth::_16, Signedness::Unsigned) => PrimitiveScalar::none::<u16>().into(),
(IntWidth::_32, Signedness::Unsigned) => PrimitiveScalar::none::<u32>().into(),
(IntWidth::_64, Signedness::Unsigned) => PrimitiveScalar::none::<u64>().into(),
},
DType::Decimal(_, _, _) => unimplemented!("DecimalScalar"),
DType::Float(w, _) => match w {
FloatWidth::Unknown => PrimitiveScalar::none::<f64>().into(),
FloatWidth::_16 => PrimitiveScalar::none::<f16>().into(),
FloatWidth::_32 => PrimitiveScalar::none::<f32>().into(),
FloatWidth::_64 => PrimitiveScalar::none::<f64>().into(),
Expand Down
9 changes: 3 additions & 6 deletions vortex-schema/flatbuffers/schema.fbs
Original file line number Diff line number Diff line change
Expand Up @@ -6,19 +6,16 @@ enum Nullability: byte {
enum Signedness: byte {
Signed,
Unsigned,
Unknown,
}

enum IntWidth: byte {
Unknown,
_8,
_16,
_32,
_64,
}

enum FloatWidth: byte {
Unknown,
_16,
_32,
_64,
Expand All @@ -31,8 +28,8 @@ table Bool {
}

table Int {
width: IntWidth = Unknown;
signedness: Signedness = Unknown;
width: IntWidth;
signedness: Signedness;
nullability: Nullability;
}

Expand All @@ -46,7 +43,7 @@ table Decimal {
}

table Float {
width: FloatWidth = Unknown;
width: FloatWidth;
nullability: Nullability;
}

Expand Down
15 changes: 4 additions & 11 deletions vortex-schema/src/dtype.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@ impl Display for Nullability {

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]
pub enum Signedness {
Unknown,
Unsigned,
Signed,
}
Expand All @@ -54,7 +53,6 @@ impl From<bool> for Signedness {
impl Display for Signedness {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Signedness::Unknown => write!(f, "unknown"),
Signedness::Unsigned => write!(f, "unsigned"),
Signedness::Signed => write!(f, "signed"),
}
Expand All @@ -63,7 +61,6 @@ impl Display for Signedness {

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]
pub enum IntWidth {
Unknown,
_8,
_16,
_32,
Expand All @@ -77,15 +74,14 @@ impl From<u16> for IntWidth {
16 => IntWidth::_16,
32 => IntWidth::_32,
64 => IntWidth::_64,
_ => IntWidth::Unknown,
_ => panic!("Invalid int width: {}", item),
}
}
}

impl Display for IntWidth {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
IntWidth::Unknown => write!(f, "_"),
IntWidth::_8 => write!(f, "8"),
IntWidth::_16 => write!(f, "16"),
IntWidth::_32 => write!(f, "32"),
Expand All @@ -96,7 +92,6 @@ impl Display for IntWidth {

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]
pub enum FloatWidth {
Unknown,
_16,
_32,
_64,
Expand All @@ -108,15 +103,14 @@ impl From<i8> for FloatWidth {
16 => FloatWidth::_16,
32 => FloatWidth::_32,
64 => FloatWidth::_64,
_ => FloatWidth::Unknown,
_ => panic!("Invalid float width: {}", item),
}
}
}

impl Display for FloatWidth {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
FloatWidth::Unknown => write!(f, "_"),
FloatWidth::_16 => write!(f, "16"),
FloatWidth::_32 => write!(f, "32"),
FloatWidth::_64 => write!(f, "64"),
Expand Down Expand Up @@ -209,9 +203,8 @@ impl Display for DType {
Null => write!(f, "null"),
Bool(n) => write!(f, "bool{}", n),
Int(w, s, n) => match s {
Unknown => write!(f, "int({}){}", w, n),
Unsigned => write!(f, "unsigned_int({}){}", w, n),
Signed => write!(f, "signed_int({}){}", w, n),
Unsigned => write!(f, "uint({}){}", w, n),
Signed => write!(f, "int({}){}", w, n),
},
Decimal(p, s, n) => write!(f, "decimal({}, {}){}", p, s, n),
Float(w, n) => write!(f, "float({}){}", w, n),
Expand Down
6 changes: 0 additions & 6 deletions vortex-schema/src/serde.rs
Original file line number Diff line number Diff line change
Expand Up @@ -301,7 +301,6 @@ impl TryFrom<FbNullability> for Nullability {
impl From<&IntWidth> for FbIntWidth {
fn from(value: &IntWidth) -> Self {
match value {
IntWidth::Unknown => FbIntWidth::Unknown,
IntWidth::_8 => FbIntWidth::_8,
IntWidth::_16 => FbIntWidth::_16,
IntWidth::_32 => FbIntWidth::_32,
Expand All @@ -315,7 +314,6 @@ impl TryFrom<FbIntWidth> for IntWidth {

fn try_from(value: FbIntWidth) -> SchemaResult<Self> {
match value {
FbIntWidth::Unknown => Ok(IntWidth::Unknown),
FbIntWidth::_8 => Ok(IntWidth::_8),
FbIntWidth::_16 => Ok(IntWidth::_16),
FbIntWidth::_32 => Ok(IntWidth::_32),
Expand All @@ -330,7 +328,6 @@ impl TryFrom<FbIntWidth> for IntWidth {
impl From<&Signedness> for FbSignedness {
fn from(value: &Signedness) -> Self {
match value {
Signedness::Unknown => FbSignedness::Unknown,
Signedness::Unsigned => FbSignedness::Unsigned,
Signedness::Signed => FbSignedness::Signed,
}
Expand All @@ -342,7 +339,6 @@ impl TryFrom<FbSignedness> for Signedness {

fn try_from(value: FbSignedness) -> SchemaResult<Self> {
match value {
FbSignedness::Unknown => Ok(Signedness::Unknown),
FbSignedness::Unsigned => Ok(Signedness::Unsigned),
FbSignedness::Signed => Ok(Signedness::Signed),
_ => Err(SchemaError::InvalidArgument(
Expand All @@ -355,7 +351,6 @@ impl TryFrom<FbSignedness> for Signedness {
impl From<&FloatWidth> for FbFloatWidth {
fn from(value: &FloatWidth) -> Self {
match value {
FloatWidth::Unknown => FbFloatWidth::Unknown,
FloatWidth::_16 => FbFloatWidth::_16,
FloatWidth::_32 => FbFloatWidth::_32,
FloatWidth::_64 => FbFloatWidth::_64,
Expand All @@ -368,7 +363,6 @@ impl TryFrom<FbFloatWidth> for FloatWidth {

fn try_from(value: FbFloatWidth) -> SchemaResult<Self> {
match value {
FbFloatWidth::Unknown => Ok(FloatWidth::Unknown),
FbFloatWidth::_16 => Ok(FloatWidth::_16),
FbFloatWidth::_32 => Ok(FloatWidth::_32),
FbFloatWidth::_64 => Ok(FloatWidth::_64),
Expand Down

0 comments on commit 1dac391

Please sign in to comment.