1use std::io::Read;
21use std::path::Path;
22use std::sync::Arc;
23
24use color_eyre::Result;
25use color_eyre::eyre::eyre;
26
27use crate::error_display::{FileError, in_file};
28use polars::prelude::*;
29
30use crate::formats::fixed_records::{Bytes, ColumnLayout, FixedRecords, Logical, Null, Physical};
31use crate::formats::model_files::MetaValue;
32use crate::formats::sqlite::Table;
33use crate::formats::text_formats::Detail;
34
35pub(crate) const READER: crate::formats::readers::Reader = crate::formats::readers::Reader {
37 convert: Some(|input| {
39 let ([file], Some(name)) = (input.files, input.options.table.as_deref()) else {
40 return Err(eyre!("Open one array of an archive at a time."));
41 };
42 let (held, lf, opened) = convert(file, name, input.options, input.writer, input.read)?;
43 Ok((
44 crate::formats::segments::Converted {
45 lf,
46 files: vec![held],
47 notes: opened.notes,
48 other_tables: opened.other_tables,
49 },
50 opened.detail,
51 ))
52 }),
53 scan,
54 python: Some(crate::export::python_script::Python {
55 call: "pl.from_numpy",
56 eager: true,
57 glob_flag: false,
58 arguments: Some(crate::export::python_script::numpy_arguments),
59 }),
60 signatures: &[
61 crate::formats::readers::Signature {
62 says: |head, _| looks_like(head),
63 kind: crate::formats::readers::Kind::Magic,
64 trusted: crate::formats::readers::EVERYWHERE,
65 },
66 crate::formats::readers::Signature {
68 says: |head, file| file.is_some_and(|file| is_archive(file, head)),
69 kind: crate::formats::readers::Kind::Magic,
70 trusted: crate::formats::readers::Trusted {
71 pipe: false,
72 open: crate::formats::readers::Unnamed::Any,
73 listing: false,
74 tables: true,
75 },
76 },
77 ],
78 tables: Some(tables),
79 table_schema: Some(schema_preview),
80 ..crate::formats::readers::BASE
81};
82
83pub const MAGIC: &[u8; 6] = b"\x93NUMPY";
85
86const MAX_DEPTH: usize = 32;
88
89pub const MAX_COLUMNS: usize = 4096;
91
92pub const MAX_GRID: usize = 1024;
95
96const MAX_MEMBERS: usize = 100_000;
98
99pub fn looks_like(head: &[u8]) -> bool {
101 head.starts_with(MAGIC)
102}
103
104pub fn is_archive(path: &Path, head: &[u8]) -> bool {
106 is_archive_name(path) && (head.starts_with(b"PK\x03\x04") || head.starts_with(b"PK\x05\x06"))
107}
108
109fn is_archive_name(path: &Path) -> bool {
110 path.extension()
111 .is_some_and(|e| e.eq_ignore_ascii_case("npz"))
112}
113
114#[derive(Debug, Clone, PartialEq)]
118pub enum Literal {
119 Str(String),
120 Int(i128),
121 Bool(bool),
122 None,
123 List(Vec<Literal>),
124 Tuple(Vec<Literal>),
125 Dict(Vec<(Literal, Literal)>),
126}
127
128impl Literal {
129 fn as_str(&self) -> Option<&str> {
130 match self {
131 Self::Str(s) => Some(s),
132 _ => None,
133 }
134 }
135
136 fn items(&self) -> Option<&[Literal]> {
137 match self {
138 Self::List(items) | Self::Tuple(items) => Some(items),
139 _ => None,
140 }
141 }
142
143 fn get(&self, key: &str) -> Option<&Literal> {
144 match self {
145 Self::Dict(pairs) => pairs
146 .iter()
147 .find(|(k, _)| k.as_str() == Some(key))
148 .map(|(_, v)| v),
149 _ => None,
150 }
151 }
152}
153
154pub fn parse_literal(text: &str) -> std::result::Result<Literal, String> {
157 let mut parser = Parser {
158 text: text.as_bytes(),
159 at: 0,
160 };
161 let value = parser.value(0)?;
162 parser.space();
163 if parser.at != parser.text.len() {
164 return Err(format!("unexpected text at byte {}", parser.at));
165 }
166 Ok(value)
167}
168
169struct Parser<'a> {
170 text: &'a [u8],
171 at: usize,
172}
173
174impl Parser<'_> {
175 fn space(&mut self) {
176 while self
177 .text
178 .get(self.at)
179 .is_some_and(|b| b.is_ascii_whitespace())
180 {
181 self.at += 1;
182 }
183 }
184
185 fn peek(&mut self) -> Option<u8> {
186 self.space();
187 self.text.get(self.at).copied()
188 }
189
190 fn value(&mut self, depth: usize) -> std::result::Result<Literal, String> {
191 if depth > MAX_DEPTH {
192 return Err("nested too deep".into());
193 }
194 match self.peek() {
195 Some(b'\'' | b'"') => self.string().map(Literal::Str),
196 Some(b'[') => self
197 .sequence(b']', depth)
198 .map(|(items, _)| Literal::List(items)),
199 Some(b'(') => {
200 let (items, comma) = self.sequence(b')', depth)?;
201 Ok(match (items.len(), comma) {
203 (1, false) => items.into_iter().next().expect("one item"),
204 _ => Literal::Tuple(items),
205 })
206 }
207 Some(b'{') => self.dict(depth),
208 Some(b'-' | b'+' | b'0'..=b'9') => self.int(),
209 Some(_) => self.word(),
210 None => Err("the header ends early".into()),
211 }
212 }
213
214 fn string(&mut self) -> std::result::Result<String, String> {
215 let quote = self.text[self.at];
216 self.at += 1;
217 let mut out = Vec::new();
218 loop {
219 let Some(&b) = self.text.get(self.at) else {
220 return Err("a string does not end".into());
221 };
222 self.at += 1;
223 match b {
224 b if b == quote => break,
225 b'\\' => {
226 let Some(&e) = self.text.get(self.at) else {
227 return Err("a string does not end".into());
228 };
229 self.at += 1;
230 match e {
231 b'n' => out.push(b'\n'),
232 b't' => out.push(b'\t'),
233 b'r' => out.push(b'\r'),
234 b'0' => out.push(0),
235 b'x' => {
236 let hex = self
237 .text
238 .get(self.at..self.at + 2)
239 .and_then(|h| std::str::from_utf8(h).ok())
240 .and_then(|h| u8::from_str_radix(h, 16).ok())
241 .ok_or("a bad \\x escape")?;
242 self.at += 2;
243 out.push(hex);
244 }
245 other => out.push(other),
246 }
247 }
248 b => out.push(b),
249 }
250 }
251 Ok(String::from_utf8_lossy(&out).into_owned())
252 }
253
254 fn sequence(
255 &mut self,
256 close: u8,
257 depth: usize,
258 ) -> std::result::Result<(Vec<Literal>, bool), String> {
259 self.at += 1;
260 let mut items = Vec::new();
261 let mut comma = false;
262 loop {
263 if self.peek() == Some(close) {
264 self.at += 1;
265 return Ok((items, comma));
266 }
267 items.push(self.value(depth + 1)?);
268 comma = false;
269 match self.peek() {
270 Some(b',') => {
271 self.at += 1;
272 comma = true;
273 }
274 Some(b) if b == close => {}
275 _ => {
276 return Err(format!(
277 "expected , or {} at byte {}",
278 close as char, self.at
279 ));
280 }
281 }
282 }
283 }
284
285 fn dict(&mut self, depth: usize) -> std::result::Result<Literal, String> {
286 self.at += 1;
287 let mut pairs = Vec::new();
288 loop {
289 if self.peek() == Some(b'}') {
290 self.at += 1;
291 return Ok(Literal::Dict(pairs));
292 }
293 let key = self.value(depth + 1)?;
294 if self.peek() != Some(b':') {
295 return Err(format!("expected : at byte {}", self.at));
296 }
297 self.at += 1;
298 let value = self.value(depth + 1)?;
299 pairs.push((key, value));
300 match self.peek() {
301 Some(b',') => self.at += 1,
302 Some(b'}') => {}
303 _ => return Err(format!("expected , or }} at byte {}", self.at)),
304 }
305 }
306 }
307
308 fn int(&mut self) -> std::result::Result<Literal, String> {
309 let start = self.at;
310 if matches!(self.text.get(self.at), Some(b'-' | b'+')) {
311 self.at += 1;
312 }
313 while self.text.get(self.at).is_some_and(u8::is_ascii_digit) {
314 self.at += 1;
315 }
316 let digits = std::str::from_utf8(&self.text[start..self.at]).unwrap_or_default();
317 if matches!(self.text.get(self.at), Some(b'L' | b'l')) {
319 self.at += 1;
320 }
321 digits
322 .parse::<i128>()
323 .map(Literal::Int)
324 .map_err(|_| format!("a bad number at byte {start}"))
325 }
326
327 fn word(&mut self) -> std::result::Result<Literal, String> {
328 let start = self.at;
329 while self.text.get(self.at).is_some_and(u8::is_ascii_alphabetic) {
330 self.at += 1;
331 }
332 match &self.text[start..self.at] {
333 b"True" => Ok(Literal::Bool(true)),
334 b"False" => Ok(Literal::Bool(false)),
335 b"None" => Ok(Literal::None),
336 _ => Err(format!("unexpected text at byte {start}")),
337 }
338 }
339}
340
341#[derive(Debug, Clone, PartialEq)]
345pub enum Dtype {
346 Simple(Simple),
347 Struct {
349 fields: Vec<FieldDef>,
350 itemsize: usize,
351 },
352}
353
354#[derive(Debug, Clone, PartialEq)]
356pub struct FieldDef {
357 pub name: String,
358 pub dtype: Dtype,
359 pub shape: Vec<usize>,
361 pub offset: usize,
362}
363
364#[derive(Debug, Clone, PartialEq)]
366pub struct Simple {
367 pub text: String,
369 pub kind: u8,
370 pub size: usize,
372 pub big: bool,
373 pub unit: Option<(String, i64)>,
375}
376
377impl Dtype {
378 pub fn itemsize(&self) -> usize {
380 match self {
381 Self::Simple(s) => s.size,
382 Self::Struct { itemsize, .. } => *itemsize,
383 }
384 }
385
386 pub fn describe(&self) -> String {
388 match self {
389 Self::Simple(s) => s.text.clone(),
390 Self::Struct { fields, itemsize } => {
391 format!("struct of {} fields, {itemsize} bytes", fields.len())
392 }
393 }
394 }
395}
396
397pub fn parse_dtype(descr: &Literal) -> std::result::Result<Dtype, String> {
399 dtype_of(descr, 0)
400}
401
402fn dtype_of(descr: &Literal, depth: usize) -> std::result::Result<Dtype, String> {
403 if depth > MAX_DEPTH {
404 return Err("a structured type nested too deep".into());
405 }
406 match descr {
407 Literal::Str(text) => simple(text).map(Dtype::Simple),
408 Literal::List(items) => {
409 let mut fields = Vec::new();
411 let mut offset = 0usize;
412 for item in items {
413 let parts = item.items().ok_or("a field is not a tuple")?;
414 let (name, field_descr, shape) = match parts {
415 [name, descr] => (name, descr, None),
416 [name, descr, shape] => (name, descr, Some(shape)),
417 _ => return Err("a field is not (name, type[, shape])".into()),
418 };
419 let name = field_name(name)?;
420 let dtype = dtype_of(field_descr, depth + 1)?;
421 let shape = shape.map(shape_of).transpose()?.unwrap_or_default();
422 let count = count_of(&shape)?;
423 let size = dtype
424 .itemsize()
425 .checked_mul(count)
426 .ok_or("a field is too large")?;
427 if !name.is_empty() {
428 fields.push(FieldDef {
429 name,
430 dtype,
431 shape,
432 offset,
433 });
434 }
435 offset = offset.checked_add(size).ok_or("a type is too large")?;
436 }
437 Ok(Dtype::Struct {
438 fields,
439 itemsize: offset,
440 })
441 }
442 Literal::Dict(_) => {
443 let list = |key: &str| {
445 descr
446 .get(key)
447 .and_then(Literal::items)
448 .ok_or_else(|| format!("a structured type without {key}"))
449 };
450 let names = list("names")?;
451 let formats = list("formats")?;
452 if names.len() != formats.len() {
453 return Err("names and formats differ in length".into());
454 }
455 let offsets: Vec<usize> = match descr.get("offsets") {
456 Some(offsets) => offsets
457 .items()
458 .ok_or("offsets is not a list")?
459 .iter()
460 .map(|o| match o {
461 Literal::Int(n) => usize::try_from(*n).map_err(|_| "a bad offset".into()),
462 _ => Err("a bad offset".to_string()),
463 })
464 .collect::<std::result::Result<_, _>>()?,
465 None => Vec::new(),
466 };
467 let mut fields = Vec::new();
468 let mut next = 0usize;
469 for (i, (name, format)) in names.iter().zip(formats).enumerate() {
470 let name = field_name(name)?;
471 let (dtype, shape) = match format.items() {
472 Some([inner, shape]) => (dtype_of(inner, depth + 1)?, shape_of(shape)?),
474 _ => (dtype_of(format, depth + 1)?, Vec::new()),
475 };
476 let size = dtype
477 .itemsize()
478 .checked_mul(count_of(&shape)?)
479 .ok_or("a field is too large")?;
480 let offset = offsets.get(i).copied().unwrap_or(next);
481 next = offset.checked_add(size).ok_or("a type is too large")?;
482 if !name.is_empty() {
483 fields.push(FieldDef {
484 name,
485 dtype,
486 shape,
487 offset,
488 });
489 }
490 }
491 let end = fields
492 .iter()
493 .map(|f| f.offset + f.dtype.itemsize() * count_of(&f.shape).unwrap_or(1))
494 .max()
495 .unwrap_or(0);
496 let itemsize = match descr.get("itemsize") {
497 Some(Literal::Int(n)) => usize::try_from(*n).map_err(|_| "a bad itemsize")?,
498 _ => end,
499 };
500 if itemsize < end {
501 return Err(format!(
502 "a field ends at byte {end}, past the element's {itemsize}"
503 ));
504 }
505 Ok(Dtype::Struct { fields, itemsize })
506 }
507 _ => Err("descr is not a type".into()),
508 }
509}
510
511fn field_name(name: &Literal) -> std::result::Result<String, String> {
513 match name {
514 Literal::Str(s) => Ok(s.clone()),
515 Literal::Tuple(pair) if pair.len() == 2 => pair[1]
516 .as_str()
517 .map(str::to_string)
518 .ok_or_else(|| "a field's name is not text".into()),
519 _ => Err("a field's name is not text".into()),
520 }
521}
522
523fn shape_of(shape: &Literal) -> std::result::Result<Vec<usize>, String> {
524 match shape {
525 Literal::Int(n) => Ok(vec![usize::try_from(*n).map_err(|_| "a bad dimension")?]),
526 Literal::Tuple(items) | Literal::List(items) => items
527 .iter()
528 .map(|d| match d {
529 Literal::Int(n) => usize::try_from(*n).map_err(|_| "a bad dimension".into()),
530 _ => Err("a dimension is not a number".to_string()),
531 })
532 .collect(),
533 _ => Err("a shape is not a tuple".into()),
534 }
535}
536
537fn count_of(shape: &[usize]) -> std::result::Result<usize, String> {
538 shape
539 .iter()
540 .try_fold(1usize, |n, d| n.checked_mul(*d))
541 .ok_or_else(|| "a shape is too large".into())
542}
543
544fn simple(text: &str) -> std::result::Result<Simple, String> {
546 let bytes = text.as_bytes();
547 let (order, rest) = match bytes.first() {
548 Some(b @ (b'<' | b'>' | b'|' | b'=' | b'!')) => (*b, &text[1..]),
549 _ => (b'=', text),
550 };
551 let big = match order {
552 b'>' | b'!' => true,
553 b'=' => cfg!(target_endian = "big"),
554 _ => false,
555 };
556 let mut chars = rest.chars();
557 let kind = chars.next().ok_or("an empty type")?;
558 if !kind.is_ascii() {
559 return Err(format!("the type {text} is not one NumPy writes"));
560 }
561 let rest = chars.as_str();
562 let (digits, unit) = match rest.split_once('[') {
563 Some((digits, unit)) => (
564 digits,
565 Some(
566 unit.strip_suffix(']')
567 .ok_or_else(|| format!("the type {text} has no ]"))?,
568 ),
569 ),
570 None => (rest, None),
571 };
572 let n: usize = if digits.is_empty() {
573 0
574 } else {
575 digits
576 .parse()
577 .map_err(|_| format!("the type {text} is not one NumPy writes"))?
578 };
579 let kind = kind as u8;
580 let size = match kind {
581 b'?' => 1,
582 b'U' => n.checked_mul(4).ok_or("a string type is too long")?,
583 b'O' => 8,
584 _ => n,
585 };
586 let unit = match unit {
587 Some(unit) => {
588 let split = unit
589 .find(|c: char| !c.is_ascii_digit())
590 .unwrap_or(unit.len());
591 let multiplier = match &unit[..split] {
592 "" => 1,
593 m => m.parse().map_err(|_| format!("the unit of {text}"))?,
594 };
595 Some((unit[split..].to_string(), multiplier))
596 }
597 None => None,
598 };
599 Ok(Simple {
600 text: text.to_string(),
601 kind,
602 size,
603 big,
604 unit,
605 })
606}
607
608fn read_as(s: &Simple) -> std::result::Result<(Physical, usize, Logical, Option<Null>), String> {
611 let plain = Logical::Plain;
612 let unsupported = || format!("the type {} is not read", s.text);
613 Ok(match (s.kind, s.size) {
614 (b'b' | b'?', 1) => (Physical::Bool, 1, plain, None),
615 (b'i', 1 | 2 | 4 | 8) => (Physical::Signed(s.size as u8), 1, plain, None),
616 (b'u', 1 | 2 | 4 | 8) => (Physical::Unsigned(s.size as u8), 1, plain, None),
617 (b'f', 2 | 4 | 8) => (Physical::Float(s.size as u8), 1, plain, None),
618 (b'c', 8 | 16) => (Physical::Float((s.size / 2) as u8), 2, plain, None),
620 (b'S' | b'a', 1..) => (Physical::Text, 1, plain, None),
621 (b'U', 4..) => (Physical::Utf32 { big_endian: s.big }, 1, plain, None),
622 (b'V', 1..) => (Physical::Raw, 1, plain, None),
623 (b'O', _) => {
624 return Err(
625 "it holds Python objects (dtype O), which only unpickling could read; datui does not unpickle"
626 .into(),
627 );
628 }
629 (b'M', 8) => {
630 let logical = match time_unit(s.unit.as_ref()) {
631 Some(Step::Days(1)) => Logical::Days { epoch_days: 0 },
632 Some(Step::Days(n)) => Logical::Timestamp {
633 unit: TimeUnit::Milliseconds,
634 multiplier: n.checked_mul(86_400_000).ok_or_else(unsupported)?,
635 epoch: 0,
636 },
637 Some(Step::Unit(unit, multiplier)) => Logical::Timestamp {
638 unit,
639 multiplier,
640 epoch: 0,
641 },
642 None => plain,
645 };
646 (Physical::Signed(8), 1, logical, Some(Null::Min))
647 }
648 (b'm', 8) => {
649 let logical = match time_unit(s.unit.as_ref()) {
650 Some(Step::Days(n)) => Logical::Duration {
651 unit: TimeUnit::Milliseconds,
652 multiplier: n.checked_mul(86_400_000).ok_or_else(unsupported)?,
653 },
654 Some(Step::Unit(unit, multiplier)) => Logical::Duration { unit, multiplier },
655 None => plain,
656 };
657 (Physical::Signed(8), 1, logical, Some(Null::Min))
658 }
659 _ => return Err(unsupported()),
660 })
661}
662
663enum Step {
664 Unit(TimeUnit, i64),
666 Days(i64),
668}
669
670fn time_unit(unit: Option<&(String, i64)>) -> Option<Step> {
672 let (name, n) = unit?;
673 let n = *n;
674 let ms = |per: i64| Some(Step::Unit(TimeUnit::Milliseconds, per.checked_mul(n)?));
675 match name.as_str() {
676 "ns" => Some(Step::Unit(TimeUnit::Nanoseconds, n)),
677 "us" | "\u{b5}s" => Some(Step::Unit(TimeUnit::Microseconds, n)),
678 "ms" => Some(Step::Unit(TimeUnit::Milliseconds, n)),
679 "s" => ms(1000),
680 "m" => ms(60_000),
681 "h" => ms(3_600_000),
682 "D" => Some(Step::Days(n)),
683 "W" => Some(Step::Days(n.checked_mul(7)?)),
684 _ => None,
685 }
686}
687
688#[derive(Debug, Clone, PartialEq)]
692pub struct Header {
693 pub version: (u8, u8),
694 pub dtype: Dtype,
695 pub fortran: bool,
696 pub shape: Vec<usize>,
697 pub data_offset: usize,
699}
700
701pub fn header_len(head: &[u8]) -> std::result::Result<(usize, usize), String> {
704 if !looks_like(head) {
705 return Err("not a NumPy .npy file: no \\x93NUMPY at the start".into());
706 }
707 let major = *head.get(6).ok_or("the header is cut short")?;
708 let (prefix, len) = match major {
709 1 => (
710 10,
711 u16::from_le_bytes(
712 head.get(8..10)
713 .ok_or("the header is cut short")?
714 .try_into()
715 .expect("two bytes"),
716 ) as usize,
717 ),
718 2 | 3 => (
719 12,
720 u32::from_le_bytes(
721 head.get(8..12)
722 .ok_or("the header is cut short")?
723 .try_into()
724 .expect("four bytes"),
725 ) as usize,
726 ),
727 v => {
728 return Err(format!(
729 "format version {v} is not one datui reads (1, 2 or 3)"
730 ));
731 }
732 };
733 let most = crate::limits::get().npy_header_bytes.bytes();
736 if len as u64 > most {
737 let size = crate::numfmt::bytes;
738 return Err(format!(
739 "its header is {}; datui reads headers up to {}, and limits.npy_header_bytes raises it",
740 size(len as u64),
741 size(most)
742 ));
743 }
744 Ok((prefix, len))
745}
746
747pub fn parse_header(bytes: &[u8]) -> std::result::Result<Header, String> {
749 let (prefix, len) = header_len(bytes)?;
750 let text = bytes
751 .get(prefix..prefix + len)
752 .ok_or("the header is cut short")?;
753 let text: String = match bytes[6] {
755 3 => String::from_utf8_lossy(text).into_owned(),
756 _ => text.iter().map(|&b| char::from(b)).collect(),
757 };
758 let dict = parse_literal(text.trim_end_matches(['\n', ' ', '\0']))
759 .map_err(|e| format!("its header is not a Python dict: {e}"))?;
760 if !matches!(dict, Literal::Dict(_)) {
761 return Err("its header is not a Python dict".into());
762 }
763 let descr = dict.get("descr").ok_or("its header has no descr")?;
764 let dtype = parse_dtype(descr)?;
765 let fortran = match dict.get("fortran_order") {
766 Some(Literal::Bool(b)) => *b,
767 None => false,
768 _ => return Err("fortran_order is not True or False".into()),
769 };
770 let shape = shape_of(dict.get("shape").ok_or("its header has no shape")?)?;
771 count_of(&shape)?;
772 Ok(Header {
773 version: (bytes[6], bytes[7]),
774 dtype,
775 fortran,
776 shape,
777 data_offset: prefix + len,
778 })
779}
780
781pub fn shape_text(shape: &[usize]) -> String {
783 match shape {
784 [one] => format!("({one},)"),
785 _ => format!(
786 "({})",
787 shape
788 .iter()
789 .map(usize::to_string)
790 .collect::<Vec<_>>()
791 .join(", ")
792 ),
793 }
794}
795
796#[derive(Debug)]
801pub struct Plan {
802 pub columns: Vec<ColumnLayout>,
803 pub rows: usize,
804 pub bytes: u128,
806}
807
808pub fn plan(header: &Header, name: &str) -> std::result::Result<Plan, String> {
810 let itemsize = header.dtype.itemsize();
811 if itemsize == 0 {
812 return Err("its elements take no bytes".into());
813 }
814 let elements = count_of(&header.shape)?;
815 let bytes = elements as u128 * itemsize as u128;
816 let refuse_shape = || {
817 format!(
818 "it has {} dimensions, shape {}; datui reads 1-D and 2-D arrays",
819 header.shape.len(),
820 shape_text(&header.shape)
821 )
822 };
823 match (&header.dtype, header.shape.as_slice()) {
824 (_, [] | [_]) => {
825 let mut columns = Vec::new();
826 let name = match header.dtype {
828 Dtype::Struct { .. } => "",
829 Dtype::Simple(_) => name,
830 };
831 leaves(&header.dtype, name, 0, itemsize, &mut columns, 0)?;
832 Ok(Plan {
833 columns,
834 rows: header.shape.first().copied().unwrap_or(1),
835 bytes,
836 })
837 }
838 (Dtype::Simple(_), [rows, cols]) => {
839 let (rows, cols) = (*rows, *cols);
840 let mut columns = Vec::new();
841 if cols > MAX_GRID {
842 if header.fortran {
843 return Err(format!(
844 "it is {} in Fortran order, with more than {MAX_GRID} columns",
845 shape_text(&header.shape)
846 ));
847 }
848 let mut leaf = Vec::new();
849 leaves(&header.dtype, "values", 0, itemsize, &mut leaf, 0)?;
850 let mut column = leaf.remove(0);
851 column.count = column.count.checked_mul(cols).ok_or("a row is too large")?;
852 column.stride = itemsize.checked_mul(cols).ok_or("a row is too large")?;
853 columns.push(column);
854 } else {
855 for c in 0..cols {
856 let mut leaf = Vec::new();
857 leaves(&header.dtype, &c.to_string(), 0, itemsize, &mut leaf, 0)?;
858 let mut column = leaf.remove(0);
859 if header.fortran {
860 column.start = (c * rows)
863 .checked_mul(itemsize)
864 .ok_or("the array is too large")?;
865 column.stride = itemsize;
866 } else {
867 column.start = c * itemsize;
868 column.stride = itemsize * cols;
869 }
870 columns.push(column);
871 }
872 }
873 Ok(Plan {
874 columns,
875 rows,
876 bytes,
877 })
878 }
879 (Dtype::Struct { .. }, [_, _]) => Err(format!(
880 "it is a structured array of shape {}; datui reads structured arrays of one dimension",
881 shape_text(&header.shape)
882 )),
883 _ => Err(refuse_shape()),
884 }
885}
886
887fn leaves(
890 dtype: &Dtype,
891 name: &str,
892 offset: usize,
893 stride: usize,
894 out: &mut Vec<ColumnLayout>,
895 depth: usize,
896) -> std::result::Result<(), String> {
897 if depth > MAX_DEPTH {
898 return Err("a structured type nested too deep".into());
899 }
900 match dtype {
901 Dtype::Simple(s) => {
902 let (physical, values, logical, null) =
903 read_as(s).map_err(|e| format!("column \"{name}\": {e}"))?;
904 let width = s.size / values;
905 out.push(ColumnLayout {
906 name: name.into(),
907 source: 0,
908 start: offset,
909 stride,
910 width,
911 count: values,
912 physical,
913 big_endian: s.big,
914 null,
915 logical,
916 });
917 }
918 Dtype::Struct { fields, .. } => {
919 for field in fields {
920 let full = match name {
921 "" => field.name.clone(),
922 _ => format!("{name}.{}", field.name),
923 };
924 let count = count_of(&field.shape)?;
925 let at = offset
926 .checked_add(field.offset)
927 .ok_or("a field is too far in")?;
928 match &field.dtype {
929 Dtype::Simple(_) => {
930 leaves(&field.dtype, &full, at, stride, out, depth + 1)?;
931 if count != 1
932 && let Some(column) = out.last_mut()
933 {
934 column.count = column
936 .count
937 .checked_mul(count)
938 .ok_or("a subarray is too large")?;
939 }
940 }
941 Dtype::Struct { .. } if count == 1 => {
942 leaves(&field.dtype, &full, at, stride, out, depth + 1)?;
943 }
944 Dtype::Struct { itemsize, .. } => {
945 for k in 0..count {
946 let at = k
947 .checked_mul(*itemsize)
948 .and_then(|o| at.checked_add(o))
949 .ok_or("a field is too far in")?;
950 leaves(
951 &field.dtype,
952 &format!("{full}[{k}]"),
953 at,
954 stride,
955 out,
956 depth + 1,
957 )?;
958 if out.len() > MAX_COLUMNS {
959 break;
960 }
961 }
962 }
963 }
964 if out.len() > MAX_COLUMNS {
965 return Err(format!("its fields make more than {MAX_COLUMNS} columns"));
966 }
967 }
968 }
969 }
970 Ok(())
971}
972
973pub struct Array {
977 pub records: Arc<FixedRecords>,
978 pub header: Header,
979 pub notes: Vec<String>,
980}
981
982pub fn open_in(
985 bytes: Arc<Bytes>,
986 at: usize,
987 len: usize,
988 name: &str,
989) -> std::result::Result<Array, String> {
990 let end = at.checked_add(len).filter(|&end| end <= bytes.len());
991 let Some(end) = end else {
992 return Err("the array runs past the end of the file".into());
993 };
994 let npy = &bytes.as_slice()[at..end];
995 let (prefix, header_bytes) = header_len(npy)?;
996 if npy.len() < prefix + header_bytes {
997 return Err("the header is cut short".into());
998 }
999 let header = parse_header(npy)?;
1000 let plan = plan(&header, name)?;
1001 let data = header.data_offset;
1002 let on_hand = (npy.len() - data) as u128;
1003 let mut notes = Vec::new();
1004 if on_hand < plan.bytes {
1005 notes.push(format!(
1006 "file cut short: {} of {} bytes for shape {} {} rows past the end left out",
1007 crate::numfmt::group_chrome(on_hand as usize),
1008 crate::numfmt::group_chrome(usize::try_from(plan.bytes).unwrap_or(usize::MAX)),
1009 shape_text(&header.shape),
1010 crate::glyphs::get().middot
1011 ));
1012 }
1013 if plan.rows > crate::formats::row_index::MAX_ROWS {
1014 notes.push(format!(
1015 "{} rows {} first {} shown",
1016 crate::numfmt::group_chrome(plan.rows),
1017 crate::glyphs::get().middot,
1018 crate::numfmt::group_chrome(crate::formats::row_index::MAX_ROWS)
1019 ));
1020 }
1021 let rows_on_hand = |c: &ColumnLayout| -> usize {
1024 let cell = c.width * c.count;
1025 let room = (end - at).saturating_sub(data + c.start);
1026 if room < cell || c.stride == 0 {
1027 0
1028 } else {
1029 (room - cell) / c.stride + 1
1030 }
1031 };
1032 let rows = plan
1033 .columns
1034 .iter()
1035 .map(rows_on_hand)
1036 .min()
1037 .unwrap_or(plan.rows)
1038 .min(plan.rows);
1039 let columns = plan
1040 .columns
1041 .into_iter()
1042 .map(|c| ColumnLayout {
1043 start: c.start + at + data,
1044 ..c
1045 })
1046 .collect();
1047 let records = FixedRecords::new(vec![bytes], columns, rows).map_err(|e| e.to_string())?;
1048 Ok(Array {
1049 records: Arc::new(records),
1050 header,
1051 notes,
1052 })
1053}
1054
1055pub fn open_file(path: &Path) -> Result<Array> {
1057 let bytes = Bytes::map(path).map_err(|e| in_file(path, e.into()))?;
1058 let len = bytes.len();
1059 open_in(Arc::new(bytes), 0, len, &stem(path)).map_err(|e| FileError::new(path, e).into())
1060}
1061
1062fn stem(path: &Path) -> String {
1063 path.file_stem()
1064 .map(|s| s.to_string_lossy().into_owned())
1065 .unwrap_or_else(|| "values".to_string())
1066}
1067
1068pub fn detail(header: &Header, archive: Option<(&Path, usize)>, compressed: bool) -> Detail {
1070 let mut lines = vec![
1071 format!("Shape: {}", shape_text(&header.shape)),
1072 format!("Type: {}", header.dtype.describe()),
1073 format!(
1074 "Order: {}",
1075 if header.fortran {
1076 "Fortran (column-major)"
1077 } else {
1078 "C (row-major)"
1079 }
1080 ),
1081 format!("Version: {}.{}", header.version.0, header.version.1),
1082 ];
1083 if let Some((archive, arrays)) = archive {
1084 lines.push(format!(
1085 "Archive: {}, {}{}",
1086 archive.display(),
1087 crate::formats::text_formats::count(arrays as u64, "array", "arrays"),
1088 if compressed {
1089 "; this one compressed, read from a copy"
1090 } else {
1091 ""
1092 }
1093 ));
1094 }
1095 let mut list = Vec::new();
1096 if let Dtype::Struct { fields, .. } = &header.dtype {
1097 let rows = fields.iter().map(|f| {
1098 let mut text = f.dtype.describe();
1099 if !f.shape.is_empty() {
1100 text.push_str(&format!(" {}", shape_text(&f.shape)));
1101 }
1102 text.push_str(&format!(" at byte {}", f.offset));
1103 (f.name.clone(), MetaValue::Text(text))
1104 });
1105 list = crate::formats::text_formats::capped_list(rows, fields.len());
1106 }
1107 Detail {
1108 tab: crate::formats::text_formats::tab(crate::FileFormat::Numpy),
1109 lines,
1110 list_title: "Fields",
1111 list,
1112 first: false,
1113 ..Default::default()
1114 }
1115}
1116
1117#[derive(Debug, Clone)]
1121pub struct Member {
1122 pub name: String,
1124 pub index: usize,
1126 pub stored_at: Option<u64>,
1128 pub size: u64,
1129}
1130
1131fn archive_of(path: &Path) -> Result<::zip::ZipArchive<std::fs::File>> {
1132 let file = std::fs::File::open(path).map_err(|e| in_file(path, e.into()))?;
1133 ::zip::ZipArchive::new(file)
1134 .map_err(|e| FileError::new(path, format!("not a NumPy archive: {e}")).into())
1135}
1136
1137pub fn members(path: &Path) -> Result<Vec<Member>> {
1139 let mut archive = archive_of(path)?;
1140 let mut members = Vec::new();
1141 for index in 0..archive.len().min(MAX_MEMBERS) {
1142 let Ok(entry) = archive.by_index_raw(index) else {
1143 continue;
1144 };
1145 let Ok(full) = entry.name() else {
1146 continue;
1147 };
1148 let Some(name) = full.strip_suffix(".npy") else {
1149 continue;
1150 };
1151 if entry.is_dir() || name.is_empty() {
1152 continue;
1153 }
1154 let stored = entry.compression() == ::zip::CompressionMethod::Stored;
1155 members.push(Member {
1156 name: name.to_string(),
1157 index,
1158 stored_at: if stored { entry.data_start() } else { None },
1159 size: entry.size(),
1160 });
1161 }
1162 Ok(members)
1163}
1164
1165fn member_header(path: &Path, member: &Member) -> std::result::Result<Header, String> {
1168 let mut archive = archive_of(path).map_err(|e| e.to_string())?;
1169 let entry = archive.by_index(member.index).map_err(|e| e.to_string())?;
1170 let mut reader = entry.take(12);
1171 let mut head = Vec::with_capacity(12);
1172 reader.read_to_end(&mut head).map_err(|e| e.to_string())?;
1173 let (prefix, len) = header_len(&head)?;
1174 let mut entry = reader.into_inner();
1175 let mut rest = Vec::with_capacity(len);
1176 (&mut entry)
1177 .take((prefix + len - head.len()) as u64)
1178 .read_to_end(&mut rest)
1179 .map_err(|e| e.to_string())?;
1180 head.extend(rest);
1181 parse_header(&head)
1182}
1183
1184pub fn tables(path: &Path) -> Result<Vec<Table>> {
1186 let members = members(path)?;
1187 Ok(members
1188 .iter()
1189 .map(|member| {
1190 let columns: Vec<String> = member_header(path, member)
1191 .ok()
1192 .and_then(|header| plan(&header, &member.name).ok())
1193 .map(|plan| plan.columns.iter().map(|c| c.name.to_string()).collect())
1194 .unwrap_or_default();
1195 Table::plain(&member.name, "array", columns)
1196 })
1197 .collect())
1198}
1199
1200pub fn schema_preview(
1203 path: &Path,
1204 member: Option<&str>,
1205) -> Option<crate::home::discover::SchemaPreview> {
1206 let (header, name) = match member {
1207 Some(name) => {
1208 let found = members(path).ok()?.into_iter().find(|m| m.name == name)?;
1209 (member_header(path, &found).ok()?, found.name)
1210 }
1211 None if is_archive_name(path) => {
1212 let all = members(path).ok()?;
1213 let [one] = all.as_slice() else {
1214 return None;
1215 };
1216 (member_header(path, one).ok()?, one.name.clone())
1217 }
1218 None => {
1219 let mut file = std::fs::File::open(path).ok()?;
1220 let mut head = [0u8; 12];
1221 file.read_exact(&mut head).ok()?;
1222 let (prefix, len) = header_len(&head).ok()?;
1223 let mut bytes = head.to_vec();
1224 file.take((prefix + len - 12) as u64)
1225 .read_to_end(&mut bytes)
1226 .ok()?;
1227 (parse_header(&bytes).ok()?, stem(path))
1228 }
1229 };
1230 let plan = plan(&header, &name).ok()?;
1231 Some(
1232 plan.columns
1233 .iter()
1234 .map(|c| (c.name.to_string(), c.dtype()))
1235 .collect(),
1236 )
1237}
1238
1239fn opened(
1241 array: Array,
1242 archive: Option<(&Path, usize)>,
1243 compressed: bool,
1244) -> crate::formats::members::Opened {
1245 let rows = array.records.rows();
1246 crate::formats::members::Opened {
1247 window: Some((array.records.clone(), rows)),
1248 detail: Some(Arc::new(detail(&array.header, archive, compressed))),
1249 other_tables: Vec::new(),
1250 notes: array
1251 .notes
1252 .into_iter()
1253 .map(|summary| crate::formats::text_formats::note(summary, "the array".to_string()))
1254 .collect(),
1255 units: Vec::new(),
1256 indexing: None,
1257 numbering: None,
1258 }
1259}
1260
1261pub(crate) fn convert(
1264 path: &Path,
1265 name: &str,
1266 options: &crate::OpenOptions,
1267 writer: &crate::loading::unfinished::Writer,
1268 read: &std::sync::atomic::AtomicU64,
1269) -> Result<(
1270 crate::cloud::download::TempDownload,
1271 LazyFrame,
1272 crate::formats::members::Opened,
1273)> {
1274 use std::io::Write;
1275 use std::sync::atomic::Ordering;
1276 let all = members(path)?;
1277 let member = all
1278 .iter()
1279 .find(|m| m.name == name)
1280 .ok_or_else(|| FileError::new(path, format!("no array \"{name}\"")))?;
1281 let mut archive = archive_of(path)?;
1282 let entry = archive
1283 .by_index(member.index)
1284 .map_err(|e| FileError::new(path, format!("array \"{name}\": {e}")))?;
1285 let compressed = entry.compressed_size();
1286 let Some((mut file, claim)) = writer.create(|| {
1287 crate::cloud::download::TempDownload::create(options.temp_dir.as_deref(), Some("npy"))
1288 })?
1289 else {
1290 return Err(eyre!("Reading was stopped."));
1291 };
1292 let mut reader = entry.take(member.size);
1295 let mut chunk = vec![0u8; 1 << 16];
1296 let mut written = 0u64;
1297 loop {
1298 if writer.stopped() {
1299 return Err(eyre!("Reading was stopped."));
1300 }
1301 let n = match reader.read(&mut chunk) {
1302 Ok(0) => break,
1303 Ok(n) => n,
1304 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
1305 Err(e) => return Err(FileError::new(path, format!("array \"{name}\": {e}")).into()),
1306 };
1307 file.write_all(&chunk[..n])?;
1308 written += n as u64;
1309 if member.size > 0 {
1310 let share = (compressed as u128 * n as u128 / member.size as u128) as u64;
1312 read.fetch_add(share, Ordering::Relaxed);
1313 }
1314 }
1315 file.flush()?;
1316 let held = crate::cloud::download::TempDownload::held(file, Some(claim));
1317 let bytes = Bytes::map(held.path())?;
1318 let array = open_in(Arc::new(bytes), 0, written as usize, name)
1319 .map_err(|e| FileError::new(path, format!("array \"{name}\" is not read: {e}")))?;
1320 let tables = arrays(&all);
1321 let lf = array.records.lazy();
1322 let mut opened = opened(array, Some((path, all.len())), true);
1323 opened.other_tables = crate::formats::members::others(&tables, name);
1324 Ok((held, lf, opened))
1325}
1326
1327fn scan(input: crate::formats::readers::ScanIn<'_>) -> Result<crate::loading::scan::Scan> {
1331 let path = &input.path().to_path_buf();
1332 let wanted = input.options.table.as_deref();
1333 let mut head = [0u8; 8];
1334 let is_npy = std::fs::File::open(path)
1335 .and_then(|mut f| f.read_exact(&mut head))
1336 .is_ok()
1337 && looks_like(&head);
1338 if is_npy {
1339 if let Some(wanted) = wanted {
1340 return Err(FileError::new(
1341 path,
1342 format!(
1343 "the file is one array; --table \"{wanted}\" picks an array of an .npz archive."
1344 ),
1345 )
1346 .into());
1347 }
1348 let array = open_file(path)?;
1349 let lf = array.records.lazy();
1350 return Ok(opened(array, None, false).scan(input, lf));
1351 }
1352 if !is_archive_name(path) {
1353 return Err(FileError::new(
1354 path,
1355 "not a NumPy file: an .npy file starts with \\x93NUMPY, and an .npz is a zip of them.",
1356 )
1357 .into());
1358 }
1359 let members = members(path)?;
1360 let tables = arrays(&members);
1361 let picked = match crate::formats::members::pick(tables.clone(), wanted, path, "")? {
1362 crate::formats::sqlite::Pick::One(table) => table,
1363 crate::formats::sqlite::Pick::Several(tables) => {
1364 return Ok(crate::formats::members::several(&input, tables));
1365 }
1366 };
1367 let member = members
1368 .iter()
1369 .find(|m| m.name == picked.name)
1370 .expect("picked from the members");
1371 let Some(at) = member.stored_at else {
1372 return Ok(crate::loading::scan::Scan::Unpack {
1373 file: path.clone(),
1374 member: member.name.clone(),
1375 format: input.format,
1376 });
1377 };
1378 let bytes = Bytes::map(path).map_err(|e| in_file(path, e.into()))?;
1379 let at = usize::try_from(at).map_err(|_| FileError::new(path, "the file is too large"))?;
1380 let len = usize::try_from(member.size).unwrap_or(usize::MAX);
1381 let array = open_in(Arc::new(bytes), at, len, &member.name)
1382 .map_err(|e| FileError::new(path, format!("array \"{}\" is not read: {e}", member.name)))?;
1383 let lf = array.records.lazy();
1384 let mut opened = opened(array, Some((path, members.len())), false);
1385 opened.other_tables = crate::formats::members::others(&tables, &member.name);
1386 Ok(opened.scan(input, lf))
1387}
1388
1389fn arrays(members: &[Member]) -> Vec<Table> {
1391 members
1392 .iter()
1393 .map(|m| Table::plain(&m.name, "array", Vec::<String>::new()))
1394 .collect()
1395}
1396
1397#[cfg(test)]
1398mod tests {
1399 use super::*;
1400
1401 #[test]
1403 fn a_header_past_the_limit_names_its_key() {
1404 let mut head = b"\x93NUMPY\x02\x00".to_vec();
1405 let past = crate::limits::get().npy_header_bytes.bytes() as u32 + 1;
1406 head.extend_from_slice(&past.to_le_bytes());
1407 let refused = header_len(&head).unwrap_err();
1408 assert!(refused.contains("limits.npy_header_bytes"), "{refused}");
1409 head[8..12].copy_from_slice(&64u32.to_le_bytes());
1410 assert_eq!(header_len(&head), Ok((12, 64)));
1411 }
1412
1413 #[test]
1415 fn errors_name_the_file() {
1416 use crate::formats::readers::bad_input::{assert_shape, each_names_its_file, opening};
1417 let mut v9 = npy("'<f8'", false, "(1,)", &[0; 8]);
1418 v9[6] = 9;
1419 each_names_its_file(
1420 crate::FileFormat::Numpy,
1421 &[
1422 ("text.npy", b"hello there", "Not a NumPy file"),
1423 ("text.npz", b"hello there", "Not a NumPy archive"),
1424 ("v9.npy", &v9, "version 9"),
1425 ("type.npy", &npy("'<q9'", false, "(1,)", &[0; 8]), "<q9"),
1426 ],
1427 );
1428 let dir = tempfile::tempdir().unwrap();
1429 let options = crate::OpenOptions {
1430 table: Some("x".into()),
1431 ..Default::default()
1432 };
1433 let one = npy("'<f8'", false, "(1,)", &[0; 8]);
1434 let message = opening(
1435 dir.path(),
1436 "a.npy",
1437 &one,
1438 crate::FileFormat::Numpy,
1439 &options,
1440 )
1441 .expect("--table refused");
1442 eprintln!("{message}");
1443 assert_shape(&message, &dir.path().join("a.npy"));
1444 assert!(message.contains("--table \"x\""), "{message}");
1445 }
1446
1447 #[test]
1450 fn a_fortran_array_too_large_to_address_is_refused() {
1451 let bytes = npy("'<f4'", true, "(4000000000000000100, 4)", &[]);
1452 let len = bytes.len();
1453 assert!(open_in(Arc::new(Bytes::Owned(bytes)), 0, len, "a").is_err());
1454 }
1455
1456 fn npy(descr: &str, fortran: bool, shape: &str, data: &[u8]) -> Vec<u8> {
1457 let mut header = format!(
1458 "{{'descr': {descr}, 'fortran_order': {}, 'shape': {shape}, }}",
1459 if fortran { "True" } else { "False" }
1460 );
1461 while (10 + header.len() + 1) % 64 != 0 {
1462 header.push(' ');
1463 }
1464 header.push('\n');
1465 let mut out = MAGIC.to_vec();
1466 out.extend([1, 0]);
1467 out.extend((header.len() as u16).to_le_bytes());
1468 out.extend(header.as_bytes());
1469 out.extend(data);
1470 out
1471 }
1472
1473 fn read(bytes: Vec<u8>, name: &str) -> DataFrame {
1474 let len = bytes.len();
1475 let array = open_in(Arc::new(Bytes::Owned(bytes)), 0, len, name).unwrap();
1476 array.records.lazy().collect().unwrap()
1477 }
1478
1479 #[test]
1480 fn literals_parse() {
1481 let parsed = parse_literal("{'a': [('x', '<f8', (3,)), (('t', 'y'), '|b1')], 'b': (2L,), 'c': True, 'd': None, 'e': -4}")
1482 .unwrap();
1483 assert_eq!(
1484 parsed.get("b"),
1485 Some(&Literal::Tuple(vec![Literal::Int(2)]))
1486 );
1487 assert_eq!(parsed.get("c"), Some(&Literal::Bool(true)));
1488 assert_eq!(parsed.get("e"), Some(&Literal::Int(-4)));
1489 assert_eq!(parse_literal("(1)").unwrap(), Literal::Int(1));
1490 assert_eq!(parse_literal("()").unwrap(), Literal::Tuple(Vec::new()));
1491 assert!(parse_literal("{'a': 1").is_err());
1492 assert!(parse_literal(&"[".repeat(100)).is_err());
1493 assert!(parse_literal("__import__('os')").is_err());
1494 }
1495
1496 #[test]
1497 fn a_vector_is_one_column_named_for_the_array() {
1498 let data: Vec<u8> = [1.5f64, -2.0, 3.25]
1499 .iter()
1500 .flat_map(|v| v.to_le_bytes())
1501 .collect();
1502 let df = read(npy("'<f8'", false, "(3,)", &data), "prices");
1503 assert_eq!(df.get_column_names(), ["prices"]);
1504 assert_eq!(
1505 df.column("prices").unwrap().f64().unwrap().to_vec(),
1506 [Some(1.5), Some(-2.0), Some(3.25)]
1507 );
1508 }
1509
1510 #[test]
1511 fn a_grid_in_either_order() {
1512 let c: Vec<u8> = [1i32, 2, 3, 4, 5, 6]
1514 .iter()
1515 .flat_map(|v| v.to_le_bytes())
1516 .collect();
1517 let f: Vec<u8> = [1i32, 4, 2, 5, 3, 6]
1518 .iter()
1519 .flat_map(|v| v.to_le_bytes())
1520 .collect();
1521 for df in [
1522 read(npy("'<i4'", false, "(2, 3)", &c), "g"),
1523 read(npy("'<i4'", true, "(2, 3)", &f), "g"),
1524 ] {
1525 assert_eq!(df.get_column_names(), ["0", "1", "2"]);
1526 assert_eq!(
1527 df.column("2").unwrap().i32().unwrap().to_vec(),
1528 [Some(3), Some(6)]
1529 );
1530 }
1531 }
1532
1533 #[test]
1534 fn structured_fields_mixed_endian_subarrays_and_padding() {
1535 let mut data = Vec::new();
1537 for (ts, px, v) in [(1u64, 0.5f32, [7i16, 8]), (2, 1.5, [9, 10])] {
1538 data.extend(ts.to_le_bytes());
1539 data.extend(px.to_be_bytes());
1540 data.extend([0u8; 3]);
1541 data.extend(v.iter().flat_map(|x| x.to_le_bytes()));
1542 }
1543 let descr = "[('ts', '<u8'), ('px', '>f4'), ('', '|V3'), ('v', '<i2', (2,))]";
1544 let df = read(npy(descr, false, "(2,)", &data), "t");
1545 assert_eq!(df.get_column_names(), ["ts", "px", "v"]);
1546 assert_eq!(
1547 df.column("px").unwrap().f32().unwrap().to_vec(),
1548 [Some(0.5), Some(1.5)]
1549 );
1550 assert_eq!(
1551 df.column("v").unwrap().dtype(),
1552 &DataType::Array(Box::new(DataType::Int16), 2)
1553 );
1554 }
1555
1556 #[test]
1557 fn the_dict_form_honors_offsets_and_itemsize() {
1558 let mut data = Vec::new();
1560 for (a, b) in [(10i32, 20i32), (30, 40)] {
1561 data.extend(b.to_le_bytes());
1562 data.extend(a.to_le_bytes());
1563 data.extend([0u8; 4]);
1564 }
1565 let descr =
1566 "{'names': ['a', 'b'], 'formats': ['<i4', '<i4'], 'offsets': [4, 0], 'itemsize': 12}";
1567 let df = read(npy(descr, false, "(2,)", &data), "t");
1568 assert_eq!(
1569 df.column("a").unwrap().i32().unwrap().to_vec(),
1570 [Some(10), Some(30)]
1571 );
1572 assert_eq!(
1573 df.column("b").unwrap().i32().unwrap().to_vec(),
1574 [Some(20), Some(40)]
1575 );
1576 }
1577
1578 #[test]
1579 fn strings_times_and_bools() {
1580 let mut data = Vec::new();
1581 for c in "hé".chars().chain(std::iter::once('\0')) {
1583 data.extend((c as u32).to_le_bytes());
1584 }
1585 data.extend(b"ab\0\0");
1586 data.extend(86_400i64.to_le_bytes());
1587 data.extend(1500i64.to_le_bytes());
1588 data.push(1);
1589 let descr = "[('u', '<U3'), ('s', '|S4'), ('t', '<M8[s]'), ('d', '<m8[ms]'), ('b', '|b1')]";
1590 let df = read(npy(descr, false, "(1,)", &data), "t");
1591 assert_eq!(df.column("u").unwrap().str().unwrap().get(0), Some("hé"));
1592 assert_eq!(df.column("s").unwrap().str().unwrap().get(0), Some("ab"));
1593 assert_eq!(
1594 df.column("t").unwrap().dtype(),
1595 &DataType::Datetime(TimeUnit::Milliseconds, None)
1596 );
1597 assert_eq!(
1598 df.column("d").unwrap().dtype(),
1599 &DataType::Duration(TimeUnit::Milliseconds)
1600 );
1601 assert_eq!(df.column("b").unwrap().bool().unwrap().get(0), Some(true));
1602 }
1603
1604 #[test]
1605 fn what_is_refused() {
1606 let refused = |descr: &str, shape: &str| {
1607 let bytes = npy(descr, false, shape, &[0u8; 64]);
1608 let len = bytes.len();
1609 open_in(Arc::new(Bytes::Owned(bytes)), 0, len, "x")
1610 .err()
1611 .unwrap()
1612 };
1613 assert!(refused("'|O'", "(1,)").contains("does not unpickle"));
1614 assert!(refused("'<f8'", "(1, 1, 1)").contains("shape (1, 1, 1)"));
1615 assert!(refused("'<f16'", "(1,)").contains("not read"));
1616 }
1617
1618 #[test]
1619 fn a_short_file_keeps_the_rows_it_has() {
1620 let data: Vec<u8> = [1i16, 2, 3].iter().flat_map(|v| v.to_le_bytes()).collect();
1621 let bytes = npy("'<i2'", false, "(10,)", &data);
1622 let len = bytes.len();
1623 let array = open_in(Arc::new(Bytes::Owned(bytes)), 0, len, "x").unwrap();
1624 assert_eq!(array.records.rows(), 3);
1625 assert_eq!(array.notes.len(), 1);
1626 }
1627}