use std::slice;
use crate::ingress::{MAX_ARRAY_DIMS, MAX_NDARRAY_LEAF_ELEMS};
use crate::{Result, error};
use super::chunk::ValidityDescriptor;
use super::wire::{
F32_NULL, F64_NULL, I8_NULL, I16_NULL, I32_NULL, I64_NULL, QWP_TYPE_BOOLEAN, QWP_TYPE_BYTE,
QWP_TYPE_CHAR, QWP_TYPE_DATE, QWP_TYPE_DECIMAL64, QWP_TYPE_DECIMAL128, QWP_TYPE_DECIMAL256,
QWP_TYPE_DOUBLE, QWP_TYPE_DOUBLE_ARRAY, QWP_TYPE_FLOAT, QWP_TYPE_GEOHASH, QWP_TYPE_INT,
QWP_TYPE_IPV4, QWP_TYPE_LONG, QWP_TYPE_LONG256, QWP_TYPE_SHORT, QWP_TYPE_TIMESTAMP,
QWP_TYPE_TIMESTAMP_NANOS, QWP_TYPE_UUID, write_qwp_varint,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum NumpyDtype {
I64Direct,
F64Direct,
DateI64Direct,
TimestampMicrosDirect,
TimestampNanosDirect,
LongDirect,
UuidDirect,
Long256Direct,
Ipv4Direct,
CharDirect,
I8Direct,
I16Direct,
I32Direct,
I8WidenToI32,
I16WidenToI32,
I32WidenToI64,
U8WidenToI32,
U16WidenToI32,
U32WidenToI64,
U64WidenToI64,
F32Direct,
F16Widen,
Bool,
DatetimeSecToMicros,
DatetimeMinuteToMicros,
DatetimeHourToMicros,
DatetimeDayToMicros,
DatetimeWeekToMicros,
DatetimeMonthToMicros,
DatetimeYearToMicros,
Decimal64 {
scale: u8,
},
Decimal128 {
scale: u8,
},
Decimal256 {
scale: u8,
},
GeohashI8 {
bits: u8,
},
GeohashI16 {
bits: u8,
},
GeohashI32 {
bits: u8,
},
GeohashI64 {
bits: u8,
},
F64Ndarray {
ndim: u8,
shape: [u32; MAX_ARRAY_DIMS],
},
}
impl NumpyDtype {
pub fn wire_type(&self) -> u8 {
use NumpyDtype as D;
match self {
D::I8Direct => QWP_TYPE_BYTE,
D::I16Direct => QWP_TYPE_SHORT,
D::I32Direct
| D::I8WidenToI32
| D::I16WidenToI32
| D::U8WidenToI32
| D::U16WidenToI32 => QWP_TYPE_INT,
D::I64Direct
| D::LongDirect
| D::I32WidenToI64
| D::U32WidenToI64
| D::U64WidenToI64 => QWP_TYPE_LONG,
D::F64Direct => QWP_TYPE_DOUBLE,
D::F32Direct | D::F16Widen => QWP_TYPE_FLOAT,
D::Bool => QWP_TYPE_BOOLEAN,
D::DateI64Direct => QWP_TYPE_DATE,
D::TimestampMicrosDirect
| D::DatetimeSecToMicros
| D::DatetimeMinuteToMicros
| D::DatetimeHourToMicros
| D::DatetimeDayToMicros
| D::DatetimeWeekToMicros
| D::DatetimeMonthToMicros
| D::DatetimeYearToMicros => QWP_TYPE_TIMESTAMP,
D::TimestampNanosDirect => QWP_TYPE_TIMESTAMP_NANOS,
D::UuidDirect => QWP_TYPE_UUID,
D::Long256Direct => QWP_TYPE_LONG256,
D::Ipv4Direct => QWP_TYPE_IPV4,
D::CharDirect => QWP_TYPE_CHAR,
D::Decimal64 { .. } => QWP_TYPE_DECIMAL64,
D::Decimal128 { .. } => QWP_TYPE_DECIMAL128,
D::Decimal256 { .. } => QWP_TYPE_DECIMAL256,
D::GeohashI8 { .. }
| D::GeohashI16 { .. }
| D::GeohashI32 { .. }
| D::GeohashI64 { .. } => QWP_TYPE_GEOHASH,
D::F64Ndarray { .. } => QWP_TYPE_DOUBLE_ARRAY,
}
}
pub fn bytes_per_row(&self) -> usize {
use NumpyDtype as D;
match self {
D::Bool | D::I8Direct => 1,
D::I16Direct | D::CharDirect => 2,
D::I32Direct
| D::I8WidenToI32
| D::I16WidenToI32
| D::U8WidenToI32
| D::U16WidenToI32
| D::F32Direct
| D::F16Widen
| D::Ipv4Direct => 4,
D::I64Direct
| D::F64Direct
| D::LongDirect
| D::DateI64Direct
| D::TimestampMicrosDirect
| D::TimestampNanosDirect
| D::DatetimeSecToMicros
| D::DatetimeMinuteToMicros
| D::DatetimeHourToMicros
| D::DatetimeDayToMicros
| D::DatetimeWeekToMicros
| D::DatetimeMonthToMicros
| D::DatetimeYearToMicros
| D::I32WidenToI64
| D::U32WidenToI64
| D::U64WidenToI64
| D::Decimal64 { .. } => 8,
D::UuidDirect | D::Decimal128 { .. } => 16,
D::Long256Direct | D::Decimal256 { .. } => 32,
D::GeohashI8 { .. } => 1,
D::GeohashI16 { .. } => 2,
D::GeohashI32 { .. } => 4,
D::GeohashI64 { .. } => 8,
D::F64Ndarray { ndim, shape } => {
let nd = (*ndim as usize).min(shape.len());
let mut leaf: usize = 1;
for &d in &shape[..nd] {
leaf = leaf.saturating_mul(d as usize);
}
(1usize)
.saturating_add(4usize.saturating_mul(nd))
.saturating_add(8usize.saturating_mul(leaf))
}
}
}
pub fn validate(&self) -> Result<()> {
if let NumpyDtype::F64Ndarray { ndim, shape } = self {
let nd = *ndim as usize;
if nd == 0 {
return Err(error::fmt!(InvalidApiCall, "F64Ndarray ndim must be >= 1"));
}
if nd > MAX_ARRAY_DIMS {
return Err(error::fmt!(
InvalidApiCall,
"F64Ndarray ndim must be <= {} (MAX_ARRAY_DIMS), got {}",
MAX_ARRAY_DIMS,
nd
));
}
let mut leaf_count: usize = 1;
for (i, &dim) in shape[..nd].iter().enumerate() {
if dim == 0 {
return Err(error::fmt!(
InvalidApiCall,
"F64Ndarray shape[{}] must be >= 1, got 0",
i
));
}
leaf_count = leaf_count.checked_mul(dim as usize).ok_or_else(|| {
error::fmt!(InvalidApiCall, "F64Ndarray shape product overflows usize")
})?;
if leaf_count > MAX_NDARRAY_LEAF_ELEMS {
return Err(error::fmt!(
InvalidApiCall,
"F64Ndarray shape product exceeds MAX_NDARRAY_LEAF_ELEMS ({}) at dim {}",
MAX_NDARRAY_LEAF_ELEMS,
i
));
}
}
}
let geohash_bits = match self {
NumpyDtype::GeohashI8 { bits } => Some((*bits, 8u8)),
NumpyDtype::GeohashI16 { bits } => Some((*bits, 16u8)),
NumpyDtype::GeohashI32 { bits } => Some((*bits, 32u8)),
NumpyDtype::GeohashI64 { bits } => Some((*bits, 60u8)),
_ => None,
};
if let Some((bits, max_bits)) = geohash_bits
&& (bits == 0 || bits > max_bits)
{
return Err(error::fmt!(
InvalidApiCall,
"geohash bits must be in 1..={}, got {}",
max_bits,
bits
));
}
let decimal_scale = match self {
NumpyDtype::Decimal64 { scale } => Some((*scale, 18u8)),
NumpyDtype::Decimal128 { scale } => Some((*scale, 38u8)),
NumpyDtype::Decimal256 { scale } => Some((*scale, 76u8)),
_ => None,
};
if let Some((scale, max_scale)) = decimal_scale
&& scale > max_scale
{
return Err(error::fmt!(
InvalidApiCall,
"decimal scale must be <= {}, got {}",
max_scale,
scale
));
}
Ok(())
}
pub fn source_elem_size(&self) -> Result<usize> {
use NumpyDtype as D;
let n = match self {
D::I8Direct | D::I8WidenToI32 | D::U8WidenToI32 | D::Bool | D::GeohashI8 { .. } => 1,
D::I16Direct
| D::I16WidenToI32
| D::U16WidenToI32
| D::F16Widen
| D::CharDirect
| D::GeohashI16 { .. } => 2,
D::I32Direct
| D::I32WidenToI64
| D::U32WidenToI64
| D::F32Direct
| D::Ipv4Direct
| D::GeohashI32 { .. } => 4,
D::I64Direct
| D::F64Direct
| D::LongDirect
| D::DateI64Direct
| D::TimestampMicrosDirect
| D::TimestampNanosDirect
| D::U64WidenToI64
| D::DatetimeSecToMicros
| D::DatetimeMinuteToMicros
| D::DatetimeHourToMicros
| D::DatetimeDayToMicros
| D::DatetimeWeekToMicros
| D::DatetimeMonthToMicros
| D::DatetimeYearToMicros
| D::GeohashI64 { .. }
| D::Decimal64 { .. } => 8,
D::UuidDirect | D::Decimal128 { .. } => 16,
D::Long256Direct | D::Decimal256 { .. } => 32,
D::F64Ndarray { ndim, shape } => {
let nd = *ndim as usize;
if nd == 0 || nd > MAX_ARRAY_DIMS {
return Err(error::fmt!(
InvalidApiCall,
"F64Ndarray ndim must be in 1..={}, got {}",
MAX_ARRAY_DIMS,
nd
));
}
let leaf: usize = shape[..nd]
.iter()
.copied()
.map(|d| d as usize)
.try_fold(1usize, |acc, d| acc.checked_mul(d))
.ok_or_else(|| {
error::fmt!(InvalidApiCall, "F64Ndarray shape product overflows usize")
})?;
return leaf.checked_mul(8).ok_or_else(|| {
error::fmt!(InvalidApiCall, "F64Ndarray row size overflows usize")
});
}
};
Ok(n)
}
}
pub(crate) unsafe fn emit_into_wire(
out: &mut Vec<u8>,
dtype: NumpyDtype,
data: *const u8,
row_count: usize,
validity: Option<&ValidityDescriptor>,
) -> Result<()> {
use NumpyDtype as D;
match dtype {
D::I64Direct | D::LongDirect => unsafe {
emit_sentinel_le::<i64, 8>(
out,
data,
row_count,
validity,
I64_NULL.to_le_bytes(),
i64::to_le_bytes,
)
},
D::F64Direct => unsafe {
emit_sentinel_le::<f64, 8>(
out,
data,
row_count,
validity,
F64_NULL.to_le_bytes(),
f64::to_le_bytes,
)
},
D::CharDirect => unsafe {
emit_sentinel_le::<u16, 2>(out, data, row_count, validity, [0u8; 2], u16::to_le_bytes)
},
D::DateI64Direct => unsafe {
emit_bitmap_le::<i64, 8>(out, data, row_count, validity, i64::to_le_bytes)
},
D::TimestampMicrosDirect | D::TimestampNanosDirect => unsafe {
emit_bitmap_le::<i64, 8>(out, data, row_count, validity, i64::to_le_bytes)
},
D::Ipv4Direct => unsafe {
emit_bitmap_le::<u32, 4>(out, data, row_count, validity, u32::to_le_bytes)
},
D::UuidDirect => unsafe { emit_bitmap_fsb::<16>(out, data, row_count, validity) },
D::Long256Direct => unsafe { emit_bitmap_fsb::<32>(out, data, row_count, validity) },
D::I8Direct => unsafe {
emit_sentinel_le::<i8, 1>(out, data, row_count, validity, [I8_NULL as u8], |v| {
[v as u8]
})
},
D::I16Direct => unsafe {
emit_sentinel_le::<i16, 2>(
out,
data,
row_count,
validity,
I16_NULL.to_le_bytes(),
i16::to_le_bytes,
)
},
D::I32Direct => unsafe {
emit_sentinel_le::<i32, 4>(
out,
data,
row_count,
validity,
I32_NULL.to_le_bytes(),
i32::to_le_bytes,
)
},
D::I8WidenToI32 => unsafe {
emit_widen_i32_sentinel::<i8>(out, data, row_count, validity, I32_NULL, |v| v as i32)
},
D::I16WidenToI32 => unsafe {
emit_widen_i32_sentinel::<i16>(out, data, row_count, validity, I32_NULL, |v| v as i32)
},
D::I32WidenToI64 => unsafe {
emit_widen_i64_sentinel::<i32>(out, data, row_count, validity, I64_NULL, |v| v as i64)
},
D::U8WidenToI32 => unsafe {
emit_widen_i32_sentinel::<u8>(out, data, row_count, validity, I32_NULL, |v| v as i32)
},
D::U16WidenToI32 => unsafe {
emit_widen_i32_sentinel::<u16>(out, data, row_count, validity, I32_NULL, |v| v as i32)
},
D::U32WidenToI64 => unsafe {
emit_widen_i64_sentinel::<u32>(out, data, row_count, validity, I64_NULL, |v| v as i64)
},
D::U64WidenToI64 => unsafe { emit_u64_widen_i64_checked(out, data, row_count, validity)? },
D::F32Direct => unsafe {
emit_sentinel_le::<f32, 4>(
out,
data,
row_count,
validity,
F32_NULL.to_le_bytes(),
f32::to_le_bytes,
)
},
D::F16Widen => unsafe { emit_f16_to_f32(out, data, row_count, validity) },
D::Bool => unsafe { emit_bool(out, data, row_count, validity) },
D::DatetimeSecToMicros => unsafe {
emit_i64_to_micros(out, data, row_count, validity, "s", |v| {
v.checked_mul(1_000_000)
})?
},
D::DatetimeMinuteToMicros => unsafe {
emit_i64_to_micros(out, data, row_count, validity, "m", |v| {
v.checked_mul(60_000_000)
})?
},
D::DatetimeHourToMicros => unsafe {
emit_i64_to_micros(out, data, row_count, validity, "h", |v| {
v.checked_mul(3_600_000_000)
})?
},
D::DatetimeDayToMicros => unsafe {
emit_i64_to_micros(out, data, row_count, validity, "D", |v| {
v.checked_mul(86_400_000_000)
})?
},
D::DatetimeWeekToMicros => unsafe {
emit_i64_to_micros(out, data, row_count, validity, "W", |v| {
v.checked_mul(604_800_000_000)
})?
},
D::DatetimeMonthToMicros => unsafe {
let mut last: Option<(i64, i64)> = None;
emit_i64_to_micros(out, data, row_count, validity, "M", |v| {
if let Some((k, r)) = last
&& k == v
{
return Some(r);
}
let r = month_offset_to_micros(v)?;
last = Some((v, r));
Some(r)
})?
},
D::DatetimeYearToMicros => unsafe {
let mut last: Option<(i64, i64)> = None;
emit_i64_to_micros(out, data, row_count, validity, "Y", |v| {
if let Some((k, r)) = last
&& k == v
{
return Some(r);
}
let r = year_offset_to_micros(v)?;
last = Some((v, r));
Some(r)
})?
},
D::Decimal64 { scale } => unsafe {
emit_decimal::<8>(out, scale, data, row_count, validity)
},
D::Decimal128 { scale } => unsafe {
emit_decimal::<16>(out, scale, data, row_count, validity)
},
D::Decimal256 { scale } => unsafe {
emit_decimal::<32>(out, scale, data, row_count, validity)
},
D::GeohashI8 { bits } => unsafe {
emit_geohash::<1>(out, bits, data, row_count, validity)?
},
D::GeohashI16 { bits } => unsafe {
emit_geohash::<2>(out, bits, data, row_count, validity)?
},
D::GeohashI32 { bits } => unsafe {
emit_geohash::<4>(out, bits, data, row_count, validity)?
},
D::GeohashI64 { bits } => unsafe {
emit_geohash::<8>(out, bits, data, row_count, validity)?
},
D::F64Ndarray { ndim, shape } => unsafe {
emit_f64_ndarray(out, ndim, shape, data, row_count, validity)?
},
}
Ok(())
}
#[inline]
unsafe fn emit_sentinel_le<T, const N: usize>(
out: &mut Vec<u8>,
data: *const u8,
row_count: usize,
validity: Option<&ValidityDescriptor>,
sentinel: [u8; N],
to_le: impl Fn(T) -> [u8; N],
) where
T: Copy,
{
out.push(0);
out.reserve(N * row_count);
let typed = data as *const T;
let data_start = out.len();
if cfg!(target_endian = "little") {
if row_count > 0 {
let bytes = unsafe { slice::from_raw_parts(data, row_count * N) };
out.extend_from_slice(bytes);
}
} else {
for i in 0..row_count {
out.extend_from_slice(&to_le(unsafe { typed.add(i).read_unaligned() }));
}
}
let Some(v) = validity.filter(|v| v.has_nulls()) else {
return;
};
let mut i = 0usize;
while i < row_count {
let byte_idx = i / 8;
let bit_off = i % 8;
if bit_off == 0 && i + 8 <= row_count && unsafe { *v.bits.add(byte_idx) } == 0xFF {
i += 8;
continue;
}
if !unsafe { v.is_valid(i) } {
let off = data_start + i * N;
out[off..off + N].copy_from_slice(&sentinel);
}
i += 1;
}
}
#[inline]
unsafe fn emit_bitmap_le<T, const N: usize>(
out: &mut Vec<u8>,
data: *const u8,
row_count: usize,
validity: Option<&ValidityDescriptor>,
to_le: impl Fn(T) -> [u8; N],
) where
T: Copy,
{
let typed = data as *const T;
match validity.filter(|v| v.has_nulls()) {
None => {
out.push(0);
out.reserve(N * row_count);
if cfg!(target_endian = "little") {
if row_count > 0 {
let bytes = unsafe { slice::from_raw_parts(data, row_count * N) };
out.extend_from_slice(bytes);
}
} else {
for i in 0..row_count {
let value = unsafe { typed.add(i).read_unaligned() };
out.extend_from_slice(&to_le(value));
}
}
}
Some(v) => {
out.push(1);
unsafe { write_qwp_bitmap_from_validity(out, v) };
out.reserve(N * v.non_null_count);
for i in 0..row_count {
if unsafe { v.is_valid(i) } {
let value = unsafe { typed.add(i).read_unaligned() };
out.extend_from_slice(&to_le(value));
}
}
}
}
}
#[inline]
unsafe fn emit_bitmap_fsb<const N: usize>(
out: &mut Vec<u8>,
data: *const u8,
row_count: usize,
validity: Option<&ValidityDescriptor>,
) {
match validity.filter(|v| v.has_nulls()) {
None => {
out.push(0);
out.reserve(N * row_count);
if row_count > 0 {
let bytes = unsafe { slice::from_raw_parts(data, N * row_count) };
out.extend_from_slice(bytes);
}
}
Some(v) => {
out.push(1);
unsafe { write_qwp_bitmap_from_validity(out, v) };
out.reserve(N * v.non_null_count);
for i in 0..row_count {
if unsafe { v.is_valid(i) } {
let row_start = unsafe { data.add(i * N) };
let row = unsafe { slice::from_raw_parts(row_start, N) };
out.extend_from_slice(row);
}
}
}
}
}
#[inline]
unsafe fn emit_widen_i32_sentinel<T>(
out: &mut Vec<u8>,
data: *const u8,
row_count: usize,
validity: Option<&ValidityDescriptor>,
sentinel: i32,
widen: impl Fn(T) -> i32,
) where
T: Copy,
{
out.push(0);
out.reserve(4 * row_count);
let typed = data as *const T;
let sentinel_bytes = sentinel.to_le_bytes();
match validity {
None => {
for i in 0..row_count {
let v = unsafe { typed.add(i).read_unaligned() };
out.extend_from_slice(&widen(v).to_le_bytes());
}
}
Some(v) => {
for i in 0..row_count {
if unsafe { v.is_valid(i) } {
let raw = unsafe { typed.add(i).read_unaligned() };
out.extend_from_slice(&widen(raw).to_le_bytes());
} else {
out.extend_from_slice(&sentinel_bytes);
}
}
}
}
}
#[inline]
unsafe fn emit_widen_i64_sentinel<T>(
out: &mut Vec<u8>,
data: *const u8,
row_count: usize,
validity: Option<&ValidityDescriptor>,
sentinel: i64,
widen: impl Fn(T) -> i64,
) where
T: Copy,
{
out.push(0);
out.reserve(8 * row_count);
let typed = data as *const T;
let sentinel_bytes = sentinel.to_le_bytes();
match validity {
None => {
for i in 0..row_count {
let v = unsafe { typed.add(i).read_unaligned() };
out.extend_from_slice(&widen(v).to_le_bytes());
}
}
Some(v) => {
for i in 0..row_count {
if unsafe { v.is_valid(i) } {
let raw = unsafe { typed.add(i).read_unaligned() };
out.extend_from_slice(&widen(raw).to_le_bytes());
} else {
out.extend_from_slice(&sentinel_bytes);
}
}
}
}
}
#[inline]
fn u64_to_i64_checked(v: u64, row: usize) -> Result<i64> {
if v > i64::MAX as u64 {
return Err(error::fmt!(
InvalidApiCall,
"u64 value {} at row {} does not fit QuestDB LONG (max i64::MAX)",
v,
row
));
}
Ok(v as i64)
}
unsafe fn emit_u64_widen_i64_checked(
out: &mut Vec<u8>,
data: *const u8,
row_count: usize,
validity: Option<&ValidityDescriptor>,
) -> Result<()> {
let typed = data as *const u64;
if validity.is_none() && row_count > 0 {
let mut acc: u64 = 0;
for i in 0..row_count {
acc |= unsafe { typed.add(i).read_unaligned() };
}
if acc < (1u64 << 63) {
unsafe {
emit_widen_i64_sentinel::<u64>(out, data, row_count, validity, I64_NULL, |v| {
v as i64
})
};
return Ok(());
}
}
out.push(0);
out.reserve(8 * row_count);
let sentinel_bytes = I64_NULL.to_le_bytes();
match validity {
None => {
for i in 0..row_count {
let v = unsafe { typed.add(i).read_unaligned() };
out.extend_from_slice(&u64_to_i64_checked(v, i)?.to_le_bytes());
}
}
Some(v) => {
for i in 0..row_count {
if unsafe { v.is_valid(i) } {
let raw = unsafe { typed.add(i).read_unaligned() };
out.extend_from_slice(&u64_to_i64_checked(raw, i)?.to_le_bytes());
} else {
out.extend_from_slice(&sentinel_bytes);
}
}
}
}
Ok(())
}
unsafe fn emit_f16_to_f32(
out: &mut Vec<u8>,
data: *const u8,
row_count: usize,
validity: Option<&ValidityDescriptor>,
) {
out.push(0);
out.reserve(4 * row_count);
let typed = data as *const u16;
let sentinel = F32_NULL.to_le_bytes();
match validity {
None => {
for i in 0..row_count {
let bits = unsafe { typed.add(i).read_unaligned() };
out.extend_from_slice(&f16_bits_to_f32(bits).to_le_bytes());
}
}
Some(v) => {
for i in 0..row_count {
if unsafe { v.is_valid(i) } {
let bits = unsafe { typed.add(i).read_unaligned() };
out.extend_from_slice(&f16_bits_to_f32(bits).to_le_bytes());
} else {
out.extend_from_slice(&sentinel);
}
}
}
}
}
#[inline]
fn f16_bits_to_f32(bits: u16) -> f32 {
let sign = ((bits >> 15) as u32) << 31;
let exp = ((bits >> 10) & 0x1F) as u32;
let mant = (bits & 0x3FF) as u32;
let f32_bits = match exp {
0 => {
if mant == 0 {
sign
} else {
let mut m = mant;
let mut e: i32 = -14;
while (m & 0x400) == 0 {
m <<= 1;
e -= 1;
}
m &= 0x3FF;
let exp_f32 = ((e + 127) as u32) << 23;
sign | exp_f32 | (m << 13)
}
}
31 => {
sign | (0xFFu32 << 23) | (mant << 13)
}
_ => {
let exp_f32 = (exp + (127 - 15)) << 23;
sign | exp_f32 | (mant << 13)
}
};
f32::from_bits(f32_bits)
}
unsafe fn emit_bool(
out: &mut Vec<u8>,
data: *const u8,
row_count: usize,
validity: Option<&ValidityDescriptor>,
) {
out.push(0);
let bitmap_bytes = row_count.div_ceil(8);
out.reserve(bitmap_bytes);
if validity.is_none() {
let full_chunks = row_count / 8;
let tail = row_count % 8;
for chunk_idx in 0..full_chunks {
let base = chunk_idx * 8;
let src = unsafe { data.add(base) };
let b0 = unsafe { *src };
let b1 = unsafe { *src.add(1) };
let b2 = unsafe { *src.add(2) };
let b3 = unsafe { *src.add(3) };
let b4 = unsafe { *src.add(4) };
let b5 = unsafe { *src.add(5) };
let b6 = unsafe { *src.add(6) };
let b7 = unsafe { *src.add(7) };
let packed = u8::from(b0 != 0)
| (u8::from(b1 != 0) << 1)
| (u8::from(b2 != 0) << 2)
| (u8::from(b3 != 0) << 3)
| (u8::from(b4 != 0) << 4)
| (u8::from(b5 != 0) << 5)
| (u8::from(b6 != 0) << 6)
| (u8::from(b7 != 0) << 7);
out.push(packed);
}
if tail != 0 {
let base = full_chunks * 8;
let mut packed = 0u8;
for i in 0..tail {
let b = unsafe { *data.add(base + i) };
if b != 0 {
packed |= 1u8 << i;
}
}
out.push(packed);
}
return;
}
let v = validity.unwrap();
let mut packed = 0u8;
let mut bit_idx = 0u8;
for i in 0..row_count {
let raw = unsafe { *data.add(i) };
if unsafe { v.is_valid(i) } && raw != 0 {
packed |= 1u8 << bit_idx;
}
bit_idx += 1;
if bit_idx == 8 {
out.push(packed);
packed = 0;
bit_idx = 0;
}
}
if bit_idx != 0 {
out.push(packed);
}
}
#[inline]
unsafe fn emit_i64_to_micros<F>(
out: &mut Vec<u8>,
data: *const u8,
row_count: usize,
validity: Option<&ValidityDescriptor>,
unit_label: &str,
mut convert: F,
) -> Result<()>
where
F: FnMut(i64) -> Option<i64>,
{
let typed = data as *const i64;
let make_err = |i: usize, value: i64| {
error::fmt!(
InvalidApiCall,
"datetime64[{}] value at row {} ({}) overflows i64 when converted to microseconds",
unit_label,
i,
value
)
};
match validity.filter(|v| v.has_nulls()) {
None => {
out.push(0);
out.reserve(8 * row_count);
for i in 0..row_count {
let value = unsafe { typed.add(i).read_unaligned() };
let micros = if value == I64_NULL {
I64_NULL
} else {
convert(value).ok_or_else(|| make_err(i, value))?
};
out.extend_from_slice(µs.to_le_bytes());
}
}
Some(v) => {
out.push(1);
unsafe { write_qwp_bitmap_from_validity(out, v) };
out.reserve(8 * v.non_null_count);
for i in 0..row_count {
if !unsafe { v.is_valid(i) } {
continue;
}
let value = unsafe { typed.add(i).read_unaligned() };
let micros = if value == I64_NULL {
I64_NULL
} else {
convert(value).ok_or_else(|| make_err(i, value))?
};
out.extend_from_slice(µs.to_le_bytes());
}
}
}
Ok(())
}
fn year_offset_to_micros(year_offset: i64) -> Option<i64> {
if !(-292_277..=292_277).contains(&year_offset) {
return None;
}
let year = 1970 + year_offset;
let days = days_from_civil(year, 1, 1);
days.checked_mul(86_400_000_000)
}
fn month_offset_to_micros(month_offset: i64) -> Option<i64> {
let year_offset = month_offset.div_euclid(12);
let month_in_year = month_offset.rem_euclid(12) as u32 + 1; if !(-292_277..=292_277).contains(&year_offset) {
return None;
}
let year = 1970 + year_offset;
let days = days_from_civil(year, month_in_year, 1);
days.checked_mul(86_400_000_000)
}
fn days_from_civil(y: i64, m: u32, d: u32) -> i64 {
let y = if m <= 2 { y - 1 } else { y };
let era = if y >= 0 { y } else { y - 399 } / 400;
let yoe = (y - era * 400) as u64; let m_adj = if m > 2 { m - 3 } else { m + 9 } as u64;
let doy = (153 * m_adj + 2) / 5 + d as u64 - 1; let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy; era * 146_097 + doe as i64 - 719_468
}
#[inline]
unsafe fn emit_decimal<const N: usize>(
out: &mut Vec<u8>,
scale: u8,
data: *const u8,
row_count: usize,
validity: Option<&ValidityDescriptor>,
) {
match validity.filter(|v| v.has_nulls()) {
None => {
out.push(0);
out.reserve(1 + N * row_count);
out.push(scale);
if row_count > 0 {
let bytes = unsafe { slice::from_raw_parts(data, N * row_count) };
out.extend_from_slice(bytes);
}
}
Some(v) => {
out.push(1);
unsafe { write_qwp_bitmap_from_validity(out, v) };
out.reserve(1 + N * v.non_null_count);
out.push(scale);
for i in 0..row_count {
if unsafe { v.is_valid(i) } {
let row_start = unsafe { data.add(i * N) };
let row = unsafe { slice::from_raw_parts(row_start, N) };
out.extend_from_slice(row);
}
}
}
}
}
#[inline]
unsafe fn emit_geohash<const SRC: usize>(
out: &mut Vec<u8>,
bits: u8,
data: *const u8,
row_count: usize,
validity: Option<&ValidityDescriptor>,
) -> Result<()> {
let elem = (bits as usize).div_ceil(8);
if elem > SRC {
return Err(error::fmt!(
InvalidApiCall,
"numpy geohash bits ({bits}) exceeds source dtype width ({SRC} bytes)"
));
}
match validity.filter(|v| v.has_nulls()) {
None => {
out.push(0);
out.reserve(1 + elem * row_count);
write_qwp_varint(out, bits as u64);
if elem == SRC && row_count > 0 {
let bytes = unsafe { slice::from_raw_parts(data, SRC * row_count) };
out.extend_from_slice(bytes);
} else {
for i in 0..row_count {
let row_start = unsafe { data.add(i * SRC) };
let row = unsafe { slice::from_raw_parts(row_start, elem) };
out.extend_from_slice(row);
}
}
}
Some(v) => {
out.push(1);
unsafe { write_qwp_bitmap_from_validity(out, v) };
out.reserve(1 + elem * v.non_null_count);
write_qwp_varint(out, bits as u64);
for i in 0..row_count {
if unsafe { v.is_valid(i) } {
let row_start = unsafe { data.add(i * SRC) };
let row = unsafe { slice::from_raw_parts(row_start, elem) };
out.extend_from_slice(row);
}
}
}
}
Ok(())
}
#[inline]
unsafe fn emit_f64_ndarray(
out: &mut Vec<u8>,
ndim: u8,
shape: [u32; MAX_ARRAY_DIMS],
data: *const u8,
row_count: usize,
validity: Option<&ValidityDescriptor>,
) -> Result<()> {
let nd = ndim as usize;
if nd == 0 || nd > MAX_ARRAY_DIMS {
return Err(error::fmt!(
InvalidApiCall,
"F64Ndarray ndim {} must be in 1..={}",
nd,
MAX_ARRAY_DIMS
));
}
let leaf_count: usize = shape[..nd]
.iter()
.copied()
.map(|d| d as usize)
.try_fold(1usize, usize::checked_mul)
.ok_or_else(|| error::fmt!(InvalidApiCall, "F64Ndarray shape overflows usize"))?;
if leaf_count > MAX_NDARRAY_LEAF_ELEMS {
return Err(error::fmt!(
InvalidApiCall,
"F64Ndarray shape product {} exceeds MAX_NDARRAY_LEAF_ELEMS ({})",
leaf_count,
MAX_NDARRAY_LEAF_ELEMS
));
}
let row_payload = 1usize
.checked_add(4usize.saturating_mul(nd))
.and_then(|v| v.checked_add(8usize.saturating_mul(leaf_count)))
.ok_or_else(|| error::fmt!(InvalidApiCall, "F64Ndarray row payload overflows usize"))?;
let row_bytes = leaf_count
.checked_mul(8)
.ok_or_else(|| error::fmt!(InvalidApiCall, "F64Ndarray row size overflows usize"))?;
let validity = validity.filter(|v| v.has_nulls());
let non_null_rows = match validity {
None => {
out.push(0);
row_count
}
Some(v) => {
out.push(1);
unsafe { write_qwp_bitmap_from_validity(out, v) };
v.non_null_count
}
};
let reserve_bytes = non_null_rows.checked_mul(row_payload).ok_or_else(|| {
error::fmt!(
InvalidApiCall,
"F64Ndarray reservation overflows usize ({} rows * {} bytes/row)",
non_null_rows,
row_payload
)
})?;
out.try_reserve(reserve_bytes).map_err(|_| {
error::fmt!(
InvalidApiCall,
"F64Ndarray reservation of {} bytes failed",
reserve_bytes
)
})?;
let header_len = 1 + 4 * nd;
let mut header: [u8; 1 + 4 * MAX_ARRAY_DIMS] = [0u8; 1 + 4 * MAX_ARRAY_DIMS];
header[0] = ndim;
for (i, &d) in shape[..nd].iter().enumerate() {
let off = 1 + 4 * i;
header[off..off + 4].copy_from_slice(&d.to_le_bytes());
}
let header = &header[..header_len];
for row in 0..row_count {
if let Some(v) = validity
&& !unsafe { v.is_valid(row) }
{
continue;
}
out.extend_from_slice(header);
let src = unsafe { data.add(row * row_bytes) };
if cfg!(target_endian = "little") {
if row_bytes > 0 {
out.extend_from_slice(unsafe { slice::from_raw_parts(src, row_bytes) });
}
} else {
for i in 0..leaf_count {
let bits = unsafe { (src.add(i * 8) as *const u64).read_unaligned() };
out.extend_from_slice(&bits.to_le_bytes());
}
}
}
Ok(())
}
unsafe fn write_qwp_bitmap_from_validity(out: &mut Vec<u8>, v: &ValidityDescriptor) {
let src = unsafe { slice::from_raw_parts(v.bits, v.byte_len()) };
super::wire::write_qwp_bitmap_invert(out, src, v.bit_len);
}
#[cfg(test)]
mod tests {
use super::super::Validity;
use super::super::chunk::Chunk;
use super::super::encoder::{EncodeScratch, encode_chunk_into};
use super::*;
use crate::ingress::TimestampUnit;
use crate::ingress::buffer::SymbolGlobalDict;
fn encode(chunk: &Chunk<'_>) -> Vec<u8> {
let mut out = Vec::new();
let mut dict = SymbolGlobalDict::new();
let mut scratch = EncodeScratch::new();
encode_chunk_into(&mut out, chunk, &mut dict, &mut scratch, false).unwrap();
out
}
fn encode_err(chunk: &Chunk<'_>) -> crate::Error {
let mut out = Vec::new();
let mut dict = SymbolGlobalDict::new();
let mut scratch = EncodeScratch::new();
encode_chunk_into(&mut out, chunk, &mut dict, &mut scratch, false).unwrap_err()
}
#[test]
fn chunk_row_count_above_max_rejected_before_read() {
let buf = [0u8; 8];
let mut chunk = Chunk::new("t");
unsafe {
chunk
.push_numpy_deferred(
"v",
NumpyDtype::I8Direct,
buf.as_ptr(),
super::super::MAX_CHUNK_ROWS + 1,
None,
)
.unwrap();
}
let err = encode_err(&chunk);
assert_eq!(err.code(), crate::ErrorCode::InvalidApiCall);
assert!(err.msg().contains("MAX_CHUNK_ROWS"), "{}", err.msg());
}
#[test]
fn source_elem_size_matches_read_stride() {
use NumpyDtype as D;
assert_eq!(D::I8Direct.source_elem_size().unwrap(), 1);
assert_eq!(D::Bool.source_elem_size().unwrap(), 1);
assert_eq!(D::I8WidenToI32.source_elem_size().unwrap(), 1); assert_eq!(D::U8WidenToI32.source_elem_size().unwrap(), 1);
assert_eq!(D::F16Widen.source_elem_size().unwrap(), 2); assert_eq!(D::CharDirect.source_elem_size().unwrap(), 2);
assert_eq!(D::I32WidenToI64.source_elem_size().unwrap(), 4); assert_eq!(D::Ipv4Direct.source_elem_size().unwrap(), 4);
assert_eq!(D::F32Direct.source_elem_size().unwrap(), 4);
assert_eq!(D::I64Direct.source_elem_size().unwrap(), 8);
assert_eq!(D::U64WidenToI64.source_elem_size().unwrap(), 8);
assert_eq!(D::DatetimeSecToMicros.source_elem_size().unwrap(), 8);
assert_eq!(D::UuidDirect.source_elem_size().unwrap(), 16);
assert_eq!(D::Long256Direct.source_elem_size().unwrap(), 32);
assert_eq!(D::Decimal64 { scale: 0 }.source_elem_size().unwrap(), 8);
assert_eq!(D::Decimal128 { scale: 0 }.source_elem_size().unwrap(), 16);
assert_eq!(D::Decimal256 { scale: 0 }.source_elem_size().unwrap(), 32);
assert_eq!(D::GeohashI8 { bits: 1 }.source_elem_size().unwrap(), 1);
assert_eq!(D::GeohashI64 { bits: 1 }.source_elem_size().unwrap(), 8);
let mut shape = [0u32; MAX_ARRAY_DIMS];
shape[0] = 2;
shape[1] = 3;
assert_eq!(
D::F64Ndarray { ndim: 2, shape }.source_elem_size().unwrap(),
2 * 3 * 8
);
}
#[test]
fn geohash_dtype_rejects_invalid_bits() {
assert!(NumpyDtype::GeohashI8 { bits: 0 }.validate().is_err());
assert!(NumpyDtype::GeohashI8 { bits: 9 }.validate().is_err());
assert!(NumpyDtype::GeohashI64 { bits: 61 }.validate().is_err());
assert!(NumpyDtype::GeohashI8 { bits: 8 }.validate().is_ok());
assert!(NumpyDtype::GeohashI64 { bits: 60 }.validate().is_ok());
}
#[test]
fn decimal_dtype_rejects_scale_above_width_max() {
assert!(NumpyDtype::Decimal64 { scale: 18 }.validate().is_ok());
assert!(NumpyDtype::Decimal128 { scale: 38 }.validate().is_ok());
assert!(NumpyDtype::Decimal256 { scale: 76 }.validate().is_ok());
for dtype in [
NumpyDtype::Decimal64 { scale: 19 },
NumpyDtype::Decimal128 { scale: 39 },
NumpyDtype::Decimal256 { scale: 77 },
] {
let err = dtype.validate().unwrap_err();
assert_eq!(err.code(), crate::ErrorCode::InvalidApiCall);
assert!(err.msg().contains("decimal scale"), "{}", err.msg());
}
}
#[test]
fn f16_bits_to_f32_matches_half_crate_for_all_bit_patterns() {
for bits in 0u32..=0xFFFFu32 {
let bits = bits as u16;
let local = f16_bits_to_f32(bits);
let reference = half::f16::from_bits(bits).to_f32();
if reference.is_nan() {
assert!(local.is_nan(), "bits={:#06x}", bits);
} else {
assert_eq!(local.to_bits(), reference.to_bits(), "bits={:#06x}", bits);
}
}
}
#[test]
fn i8_direct_matches_column_i8() {
let src = [1i8, -2, 3];
let ts = [10i64, 20, 30];
let mut a = Chunk::new("t");
unsafe {
a.push_numpy_deferred(
"v",
NumpyDtype::I8Direct,
src.as_ptr() as *const u8,
src.len(),
None,
)
.unwrap();
}
a.at_nanos(&ts).unwrap();
let bytes_a = encode(&a);
let mut b = Chunk::new("t");
b.column_i8("v", &src, None).unwrap();
b.at_nanos(&ts).unwrap();
let bytes_b = encode(&b);
assert_eq!(
bytes_a, bytes_b,
"I8Direct must produce byte-identical wire to column_i8"
);
}
#[test]
fn i16_direct_matches_column_i16() {
let src = [1i16, -2, 3];
let ts = [10i64, 20, 30];
let mut a = Chunk::new("t");
unsafe {
a.push_numpy_deferred(
"v",
NumpyDtype::I16Direct,
src.as_ptr() as *const u8,
src.len(),
None,
)
.unwrap();
}
a.at_nanos(&ts).unwrap();
let bytes_a = encode(&a);
let mut b = Chunk::new("t");
b.column_i16("v", &src, None).unwrap();
b.at_nanos(&ts).unwrap();
let bytes_b = encode(&b);
assert_eq!(
bytes_a, bytes_b,
"I16Direct must produce byte-identical wire to column_i16"
);
}
#[test]
fn i32_direct_matches_column_i32() {
let src = [1i32, -2, 3];
let ts = [10i64, 20, 30];
let mut a = Chunk::new("t");
unsafe {
a.push_numpy_deferred(
"v",
NumpyDtype::I32Direct,
src.as_ptr() as *const u8,
src.len(),
None,
)
.unwrap();
}
a.at_nanos(&ts).unwrap();
let bytes_a = encode(&a);
let mut b = Chunk::new("t");
b.column_i32("v", &src, None).unwrap();
b.at_nanos(&ts).unwrap();
let bytes_b = encode(&b);
assert_eq!(
bytes_a, bytes_b,
"I32Direct must produce byte-identical wire to column_i32"
);
}
#[test]
fn u8_widen_matches_column_i32() {
let src = [0u8, 1, 200, 255];
let widened: [i32; 4] = [0, 1, 200, 255];
let ts = [10i64, 20, 30, 40];
let mut a = Chunk::new("t");
unsafe {
a.push_numpy_deferred("v", NumpyDtype::U8WidenToI32, src.as_ptr(), src.len(), None)
.unwrap();
}
a.at_nanos(&ts).unwrap();
let bytes_a = encode(&a);
let mut b = Chunk::new("t");
b.column_i32("v", &widened, None).unwrap();
b.at_nanos(&ts).unwrap();
let bytes_b = encode(&b);
assert_eq!(
bytes_a, bytes_b,
"U8WidenToI32 must produce byte-identical wire to column_i32 over the widened data"
);
}
#[test]
fn u16_widen_matches_column_i32() {
let src = [0u16, 1, 30000, 65535];
let widened: [i32; 4] = [0, 1, 30000, 65535];
let ts = [10i64, 20, 30, 40];
let mut a = Chunk::new("t");
unsafe {
a.push_numpy_deferred(
"v",
NumpyDtype::U16WidenToI32,
src.as_ptr() as *const u8,
src.len(),
None,
)
.unwrap();
}
a.at_nanos(&ts).unwrap();
let bytes_a = encode(&a);
let mut b = Chunk::new("t");
b.column_i32("v", &widened, None).unwrap();
b.at_nanos(&ts).unwrap();
let bytes_b = encode(&b);
assert_eq!(
bytes_a, bytes_b,
"U16WidenToI32 must produce byte-identical wire to column_i32 over the widened data"
);
}
#[test]
fn i8_widen_matches_column_i32() {
let src = [-128i8, -1, 0, 1, 127];
let widened: [i32; 5] = [-128, -1, 0, 1, 127];
let ts = [10i64, 20, 30, 40, 50];
let mut a = Chunk::new("t");
unsafe {
a.push_numpy_deferred(
"v",
NumpyDtype::I8WidenToI32,
src.as_ptr() as *const u8,
src.len(),
None,
)
.unwrap();
}
a.at_nanos(&ts).unwrap();
let bytes_a = encode(&a);
let mut b = Chunk::new("t");
b.column_i32("v", &widened, None).unwrap();
b.at_nanos(&ts).unwrap();
let bytes_b = encode(&b);
assert_eq!(
bytes_a, bytes_b,
"I8WidenToI32 must produce byte-identical wire to column_i32 over the widened data"
);
}
#[test]
fn i16_widen_matches_column_i32() {
let src = [i16::MIN, -1, 0, 1, i16::MAX];
let widened: [i32; 5] = [i16::MIN as i32, -1, 0, 1, i16::MAX as i32];
let ts = [10i64, 20, 30, 40, 50];
let mut a = Chunk::new("t");
unsafe {
a.push_numpy_deferred(
"v",
NumpyDtype::I16WidenToI32,
src.as_ptr() as *const u8,
src.len(),
None,
)
.unwrap();
}
a.at_nanos(&ts).unwrap();
let bytes_a = encode(&a);
let mut b = Chunk::new("t");
b.column_i32("v", &widened, None).unwrap();
b.at_nanos(&ts).unwrap();
let bytes_b = encode(&b);
assert_eq!(
bytes_a, bytes_b,
"I16WidenToI32 must produce byte-identical wire to column_i32 over the widened data"
);
}
#[test]
fn i32_widen_matches_column_i64() {
let src = [i32::MIN, -1, 0, 1, i32::MAX];
let widened: [i64; 5] = [i32::MIN as i64, -1, 0, 1, i32::MAX as i64];
let ts = [10i64, 20, 30, 40, 50];
let mut a = Chunk::new("t");
unsafe {
a.push_numpy_deferred(
"v",
NumpyDtype::I32WidenToI64,
src.as_ptr() as *const u8,
src.len(),
None,
)
.unwrap();
}
a.at_nanos(&ts).unwrap();
let bytes_a = encode(&a);
let mut b = Chunk::new("t");
b.column_i64("v", &widened, None).unwrap();
b.at_nanos(&ts).unwrap();
let bytes_b = encode(&b);
assert_eq!(
bytes_a, bytes_b,
"I32WidenToI64 must produce byte-identical wire to column_i64 over the widened data"
);
}
#[test]
fn u64_widen_within_i64_range_matches_column_i64() {
let src = [0u64, 42, i64::MAX as u64];
let widened: [i64; 3] = [0, 42, i64::MAX];
let ts = [10i64, 20, 30];
let mut a = Chunk::new("t");
unsafe {
a.push_numpy_deferred(
"v",
NumpyDtype::U64WidenToI64,
src.as_ptr() as *const u8,
src.len(),
None,
)
.unwrap();
}
a.at_nanos(&ts).unwrap();
let bytes_a = encode(&a);
let mut b = Chunk::new("t");
b.column_i64("v", &widened, None).unwrap();
b.at_nanos(&ts).unwrap();
let bytes_b = encode(&b);
assert_eq!(
bytes_a, bytes_b,
"U64WidenToI64 must produce signed LONG wire for values within i64::MAX"
);
}
#[test]
fn u64_widen_above_i64_max_rejects() {
let src = [i64::MAX as u64 + 1];
let ts = [10i64];
let mut chunk = Chunk::new("t");
unsafe {
chunk
.push_numpy_deferred(
"v",
NumpyDtype::U64WidenToI64,
src.as_ptr() as *const u8,
src.len(),
None,
)
.unwrap();
}
chunk.at_nanos(&ts).unwrap();
let err = encode_err(&chunk);
assert_eq!(err.code(), crate::ErrorCode::InvalidApiCall);
assert!(
err.msg().contains("does not fit QuestDB LONG"),
"{}",
err.msg()
);
}
#[test]
fn nullable_u64_widen_above_i64_max_rejects() {
let src = [0u64, i64::MAX as u64 + 1];
let ts = [10i64, 20];
let validity_bits = [0b0000_0010u8];
let validity = Validity::from_bitmap(&validity_bits, src.len()).unwrap();
let mut chunk = Chunk::new("t");
unsafe {
chunk
.push_numpy_deferred(
"v",
NumpyDtype::U64WidenToI64,
src.as_ptr() as *const u8,
src.len(),
Some(&validity),
)
.unwrap();
}
chunk.at_nanos(&ts).unwrap();
let err = encode_err(&chunk);
assert_eq!(err.code(), crate::ErrorCode::InvalidApiCall);
assert!(
err.msg().contains("does not fit QuestDB LONG"),
"{}",
err.msg()
);
}
#[test]
fn f32_direct_matches_column_f32() {
let src = [1.5f32, -2.25, 3.125, f32::NAN];
let ts = [10i64, 20, 30, 40];
let mut a = Chunk::new("t");
unsafe {
a.push_numpy_deferred(
"v",
NumpyDtype::F32Direct,
src.as_ptr() as *const u8,
src.len(),
None,
)
.unwrap();
}
a.at_nanos(&ts).unwrap();
let bytes_a = encode(&a);
let mut b = Chunk::new("t");
b.column_f32("v", &src, None).unwrap();
b.at_nanos(&ts).unwrap();
let bytes_b = encode(&b);
assert_eq!(
bytes_a, bytes_b,
"F32Direct must produce byte-identical wire to column_f32"
);
}
#[test]
fn bool_with_null_matches_column_bool() {
let raw = [1u8, 0, 1, 1];
let ts = [1i64, 2, 3, 4];
let v_bits = [0b0000_1011u8];
let v = Validity::from_bitmap(&v_bits, 4).unwrap();
let mut a = Chunk::new("t");
unsafe {
a.push_numpy_deferred("b", NumpyDtype::Bool, raw.as_ptr(), raw.len(), Some(&v))
.unwrap();
}
a.at_nanos(&ts).unwrap();
let bytes_a = encode(&a);
let mut packed = vec![0u8; raw.len().div_ceil(8)];
for (i, &b) in raw.iter().enumerate() {
if b != 0 {
packed[i / 8] |= 1u8 << (i % 8);
}
}
let mut b = Chunk::new("t");
b.column_bool("b", &packed, raw.len(), Some(&v)).unwrap();
b.at_nanos(&ts).unwrap();
let bytes_b = encode(&b);
assert_eq!(
bytes_a, bytes_b,
"Bool numpy emit must match column_bool over the equivalent packed bitmap"
);
}
#[test]
fn timestamp_nanos_direct_matches_column_ts_nanos() {
let src = [1_000i64, 2_000, 3_000];
let ts = [1i64, 2, 3];
let mut a = Chunk::new("t");
unsafe {
a.push_numpy_deferred(
"ts",
NumpyDtype::TimestampNanosDirect,
src.as_ptr() as *const u8,
src.len(),
None,
)
.unwrap();
}
a.at_nanos(&ts).unwrap();
let bytes_a = encode(&a);
let mut b = Chunk::new("t");
b.column_ts("ts", &src, TimestampUnit::Nanos, None).unwrap();
b.at_nanos(&ts).unwrap();
let bytes_b = encode(&b);
assert_eq!(
bytes_a, bytes_b,
"TimestampNanosDirect must produce byte-identical wire to column_ts(Nanos)"
);
}
fn encode_datetime_col(dtype: NumpyDtype, src_le_bytes: &[u8], row_count: usize) -> Vec<u8> {
let ts: Vec<i64> = (0..row_count as i64).collect();
let mut chunk = Chunk::new("t");
unsafe {
chunk
.push_numpy_deferred("v", dtype, src_le_bytes.as_ptr(), row_count, None)
.unwrap();
}
chunk.at_nanos(&ts).unwrap();
encode(&chunk)
}
fn encode_micros_col(values: &[i64]) -> Vec<u8> {
let ts: Vec<i64> = (0..values.len() as i64).collect();
let mut chunk = Chunk::new("t");
chunk
.column_ts("v", values, TimestampUnit::Micros, None)
.unwrap();
chunk.at_nanos(&ts).unwrap();
encode(&chunk)
}
#[test]
fn datetime_day_matches_column_ts_micros() {
let src = [0i64, 1, 18262]; let expected = [0i64, 86_400_000_000, 18262 * 86_400_000_000];
let raw: Vec<u8> = src.iter().flat_map(|v| v.to_le_bytes()).collect();
assert_eq!(
encode_datetime_col(NumpyDtype::DatetimeDayToMicros, &raw, src.len()),
encode_micros_col(&expected),
);
}
#[test]
fn datetime_nat_maps_to_null_not_error() {
let src = [0i64, i64::MIN, 1];
let expected = [0i64, i64::MIN, 86_400_000_000];
let raw: Vec<u8> = src.iter().flat_map(|v| v.to_le_bytes()).collect();
assert_eq!(
encode_datetime_col(NumpyDtype::DatetimeDayToMicros, &raw, src.len()),
encode_micros_col(&expected),
);
}
#[test]
fn datetime_hour_matches_column_ts_micros() {
let src = [0i64, 1, 24];
let expected = [0i64, 3_600_000_000, 24 * 3_600_000_000];
let raw: Vec<u8> = src.iter().flat_map(|v| v.to_le_bytes()).collect();
assert_eq!(
encode_datetime_col(NumpyDtype::DatetimeHourToMicros, &raw, src.len()),
encode_micros_col(&expected),
);
}
#[test]
fn datetime_minute_matches_column_ts_micros() {
let src = [0i64, 1, 60];
let expected = [0i64, 60_000_000, 60 * 60_000_000];
let raw: Vec<u8> = src.iter().flat_map(|v| v.to_le_bytes()).collect();
assert_eq!(
encode_datetime_col(NumpyDtype::DatetimeMinuteToMicros, &raw, src.len()),
encode_micros_col(&expected),
);
}
#[test]
fn datetime_year_matches_calendar() {
let src = [0i64, 50, -1];
let expected = [0i64, 18262 * 86_400_000_000, -365 * 86_400_000_000];
let raw: Vec<u8> = src.iter().flat_map(|v| v.to_le_bytes()).collect();
assert_eq!(
encode_datetime_col(NumpyDtype::DatetimeYearToMicros, &raw, src.len()),
encode_micros_col(&expected),
);
}
#[test]
fn datetime_month_matches_calendar() {
let src = [0i64, 1, 13, -1];
let expected = [
0i64,
31 * 86_400_000_000,
(365 + 31) * 86_400_000_000,
-31 * 86_400_000_000,
];
let raw: Vec<u8> = src.iter().flat_map(|v| v.to_le_bytes()).collect();
assert_eq!(
encode_datetime_col(NumpyDtype::DatetimeMonthToMicros, &raw, src.len()),
encode_micros_col(&expected),
);
}
#[test]
fn datetime_year_out_of_range_rejected() {
let bad = [10_000_000i64]; let ts = [1i64];
let mut chunk = Chunk::new("t");
unsafe {
chunk
.push_numpy_deferred(
"ts",
NumpyDtype::DatetimeYearToMicros,
bad.as_ptr() as *const u8,
bad.len(),
None,
)
.unwrap();
}
chunk.at_nanos(&ts).unwrap();
let err = {
let mut out = Vec::new();
let mut dict = SymbolGlobalDict::new();
let mut scratch = EncodeScratch::new();
encode_chunk_into(&mut out, &chunk, &mut dict, &mut scratch, false).unwrap_err()
};
assert_eq!(err.code(), crate::ErrorCode::InvalidApiCall);
assert!(err.msg().contains("overflows"));
}
#[test]
fn datetime_sec_overflow_rejected() {
let bad = [i64::MAX];
let ts = [1i64];
let mut chunk = Chunk::new("t");
unsafe {
chunk
.push_numpy_deferred(
"ts",
NumpyDtype::DatetimeSecToMicros,
bad.as_ptr() as *const u8,
bad.len(),
None,
)
.unwrap();
}
chunk.at_nanos(&ts).unwrap();
let err = {
let mut out = Vec::new();
let mut dict = SymbolGlobalDict::new();
let mut scratch = EncodeScratch::new();
encode_chunk_into(&mut out, &chunk, &mut dict, &mut scratch, false).unwrap_err()
};
assert_eq!(err.code(), crate::ErrorCode::InvalidApiCall);
assert!(err.msg().contains("overflows"));
}
#[test]
fn f64_ndarray_1d_no_validity_layout() {
let rows: [f64; 6] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
let ts = [10i64, 20];
let mut shape = [0u32; MAX_ARRAY_DIMS];
shape[0] = 3;
let mut chunk = Chunk::new("t");
unsafe {
chunk
.push_numpy_deferred(
"v",
NumpyDtype::F64Ndarray { ndim: 1, shape },
rows.as_ptr() as *const u8,
2,
None,
)
.unwrap();
}
chunk.at_nanos(&ts).unwrap();
let bytes = encode(&chunk);
let mut body: Vec<u8> = Vec::new();
body.push(0u8); for row_chunk in rows.chunks_exact(3) {
body.push(1u8); body.extend_from_slice(&3u32.to_le_bytes()); for &v in row_chunk {
body.extend_from_slice(&v.to_le_bytes());
}
}
assert!(
bytes.windows(body.len()).any(|w| w == body.as_slice()),
"expected ndarray column body subsequence in encoded frame"
);
}
#[test]
fn f16_bits_to_f32_known_values() {
assert_eq!(f16_bits_to_f32(0x0000), 0.0f32);
assert_eq!(f16_bits_to_f32(0x8000).to_bits(), (-0.0f32).to_bits());
assert_eq!(f16_bits_to_f32(0x3C00), 1.0f32);
assert_eq!(f16_bits_to_f32(0xC000), -2.0f32);
assert!(f16_bits_to_f32(0x7C00).is_infinite() && f16_bits_to_f32(0x7C00) > 0.0);
let v = f16_bits_to_f32(0x0001);
assert_eq!(v, 2.0f32.powi(-24));
}
}