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 Some(name) = entry.name().strip_suffix(".npy") else {
1146 continue;
1147 };
1148 if entry.is_dir() || name.is_empty() {
1149 continue;
1150 }
1151 let stored = entry.compression() == ::zip::CompressionMethod::Stored;
1152 members.push(Member {
1153 name: name.to_string(),
1154 index,
1155 stored_at: if stored { entry.data_start() } else { None },
1156 size: entry.size(),
1157 });
1158 }
1159 Ok(members)
1160}
1161
1162fn member_header(path: &Path, member: &Member) -> std::result::Result<Header, String> {
1165 let mut archive = archive_of(path).map_err(|e| e.to_string())?;
1166 let entry = archive.by_index(member.index).map_err(|e| e.to_string())?;
1167 let mut reader = entry.take(12);
1168 let mut head = Vec::with_capacity(12);
1169 reader.read_to_end(&mut head).map_err(|e| e.to_string())?;
1170 let (prefix, len) = header_len(&head)?;
1171 let mut entry = reader.into_inner();
1172 let mut rest = Vec::with_capacity(len);
1173 (&mut entry)
1174 .take((prefix + len - head.len()) as u64)
1175 .read_to_end(&mut rest)
1176 .map_err(|e| e.to_string())?;
1177 head.extend(rest);
1178 parse_header(&head)
1179}
1180
1181pub fn tables(path: &Path) -> Result<Vec<Table>> {
1183 let members = members(path)?;
1184 Ok(members
1185 .iter()
1186 .map(|member| {
1187 let columns: Vec<String> = member_header(path, member)
1188 .ok()
1189 .and_then(|header| plan(&header, &member.name).ok())
1190 .map(|plan| plan.columns.iter().map(|c| c.name.to_string()).collect())
1191 .unwrap_or_default();
1192 Table::plain(&member.name, "array", columns)
1193 })
1194 .collect())
1195}
1196
1197pub fn schema_preview(
1200 path: &Path,
1201 member: Option<&str>,
1202) -> Option<crate::home::discover::SchemaPreview> {
1203 let (header, name) = match member {
1204 Some(name) => {
1205 let found = members(path).ok()?.into_iter().find(|m| m.name == name)?;
1206 (member_header(path, &found).ok()?, found.name)
1207 }
1208 None if is_archive_name(path) => {
1209 let all = members(path).ok()?;
1210 let [one] = all.as_slice() else {
1211 return None;
1212 };
1213 (member_header(path, one).ok()?, one.name.clone())
1214 }
1215 None => {
1216 let mut file = std::fs::File::open(path).ok()?;
1217 let mut head = [0u8; 12];
1218 file.read_exact(&mut head).ok()?;
1219 let (prefix, len) = header_len(&head).ok()?;
1220 let mut bytes = head.to_vec();
1221 file.take((prefix + len - 12) as u64)
1222 .read_to_end(&mut bytes)
1223 .ok()?;
1224 (parse_header(&bytes).ok()?, stem(path))
1225 }
1226 };
1227 let plan = plan(&header, &name).ok()?;
1228 Some(
1229 plan.columns
1230 .iter()
1231 .map(|c| (c.name.to_string(), c.dtype()))
1232 .collect(),
1233 )
1234}
1235
1236fn opened(
1238 array: Array,
1239 archive: Option<(&Path, usize)>,
1240 compressed: bool,
1241) -> crate::formats::members::Opened {
1242 let rows = array.records.rows();
1243 crate::formats::members::Opened {
1244 window: Some((array.records.clone(), rows)),
1245 detail: Some(Arc::new(detail(&array.header, archive, compressed))),
1246 other_tables: Vec::new(),
1247 notes: array
1248 .notes
1249 .into_iter()
1250 .map(|summary| crate::formats::text_formats::note(summary, "the array".to_string()))
1251 .collect(),
1252 units: Vec::new(),
1253 indexing: None,
1254 numbering: None,
1255 }
1256}
1257
1258pub(crate) fn convert(
1261 path: &Path,
1262 name: &str,
1263 options: &crate::OpenOptions,
1264 writer: &crate::loading::unfinished::Writer,
1265 read: &std::sync::atomic::AtomicU64,
1266) -> Result<(
1267 crate::cloud::download::TempDownload,
1268 LazyFrame,
1269 crate::formats::members::Opened,
1270)> {
1271 use std::io::Write;
1272 use std::sync::atomic::Ordering;
1273 let all = members(path)?;
1274 let member = all
1275 .iter()
1276 .find(|m| m.name == name)
1277 .ok_or_else(|| FileError::new(path, format!("no array \"{name}\"")))?;
1278 let mut archive = archive_of(path)?;
1279 let entry = archive
1280 .by_index(member.index)
1281 .map_err(|e| FileError::new(path, format!("array \"{name}\": {e}")))?;
1282 let compressed = entry.compressed_size();
1283 let Some((mut file, claim)) = writer.create(|| {
1284 crate::cloud::download::TempDownload::create(options.temp_dir.as_deref(), Some("npy"))
1285 })?
1286 else {
1287 return Err(eyre!("Reading was stopped."));
1288 };
1289 let mut reader = entry.take(member.size);
1292 let mut chunk = vec![0u8; 1 << 16];
1293 let mut written = 0u64;
1294 loop {
1295 if writer.stopped() {
1296 return Err(eyre!("Reading was stopped."));
1297 }
1298 let n = match reader.read(&mut chunk) {
1299 Ok(0) => break,
1300 Ok(n) => n,
1301 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
1302 Err(e) => return Err(FileError::new(path, format!("array \"{name}\": {e}")).into()),
1303 };
1304 file.write_all(&chunk[..n])?;
1305 written += n as u64;
1306 if member.size > 0 {
1307 let share = (compressed as u128 * n as u128 / member.size as u128) as u64;
1309 read.fetch_add(share, Ordering::Relaxed);
1310 }
1311 }
1312 file.flush()?;
1313 let held = crate::cloud::download::TempDownload::held(file, Some(claim));
1314 let bytes = Bytes::map(held.path())?;
1315 let array = open_in(Arc::new(bytes), 0, written as usize, name)
1316 .map_err(|e| FileError::new(path, format!("array \"{name}\" is not read: {e}")))?;
1317 let tables = arrays(&all);
1318 let lf = array.records.lazy();
1319 let mut opened = opened(array, Some((path, all.len())), true);
1320 opened.other_tables = crate::formats::members::others(&tables, name);
1321 Ok((held, lf, opened))
1322}
1323
1324fn scan(input: crate::formats::readers::ScanIn<'_>) -> Result<crate::loading::scan::Scan> {
1328 let path = &input.path().to_path_buf();
1329 let wanted = input.options.table.as_deref();
1330 let mut head = [0u8; 8];
1331 let is_npy = std::fs::File::open(path)
1332 .and_then(|mut f| f.read_exact(&mut head))
1333 .is_ok()
1334 && looks_like(&head);
1335 if is_npy {
1336 if let Some(wanted) = wanted {
1337 return Err(FileError::new(
1338 path,
1339 format!(
1340 "the file is one array; --table \"{wanted}\" picks an array of an .npz archive."
1341 ),
1342 )
1343 .into());
1344 }
1345 let array = open_file(path)?;
1346 let lf = array.records.lazy();
1347 return Ok(opened(array, None, false).scan(input, lf));
1348 }
1349 if !is_archive_name(path) {
1350 return Err(FileError::new(
1351 path,
1352 "not a NumPy file: an .npy file starts with \\x93NUMPY, and an .npz is a zip of them.",
1353 )
1354 .into());
1355 }
1356 let members = members(path)?;
1357 let tables = arrays(&members);
1358 let picked = match crate::formats::members::pick(tables.clone(), wanted, path, "")? {
1359 crate::formats::sqlite::Pick::One(table) => table,
1360 crate::formats::sqlite::Pick::Several(tables) => {
1361 return Ok(crate::formats::members::several(&input, tables));
1362 }
1363 };
1364 let member = members
1365 .iter()
1366 .find(|m| m.name == picked.name)
1367 .expect("picked from the members");
1368 let Some(at) = member.stored_at else {
1369 return Ok(crate::loading::scan::Scan::Unpack {
1370 file: path.clone(),
1371 member: member.name.clone(),
1372 format: input.format,
1373 });
1374 };
1375 let bytes = Bytes::map(path).map_err(|e| in_file(path, e.into()))?;
1376 let at = usize::try_from(at).map_err(|_| FileError::new(path, "the file is too large"))?;
1377 let len = usize::try_from(member.size).unwrap_or(usize::MAX);
1378 let array = open_in(Arc::new(bytes), at, len, &member.name)
1379 .map_err(|e| FileError::new(path, format!("array \"{}\" is not read: {e}", member.name)))?;
1380 let lf = array.records.lazy();
1381 let mut opened = opened(array, Some((path, members.len())), false);
1382 opened.other_tables = crate::formats::members::others(&tables, &member.name);
1383 Ok(opened.scan(input, lf))
1384}
1385
1386fn arrays(members: &[Member]) -> Vec<Table> {
1388 members
1389 .iter()
1390 .map(|m| Table::plain(&m.name, "array", Vec::<String>::new()))
1391 .collect()
1392}
1393
1394#[cfg(test)]
1395mod tests {
1396 use super::*;
1397
1398 #[test]
1400 fn a_header_past_the_limit_names_its_key() {
1401 let mut head = b"\x93NUMPY\x02\x00".to_vec();
1402 let past = crate::limits::get().npy_header_bytes.bytes() as u32 + 1;
1403 head.extend_from_slice(&past.to_le_bytes());
1404 let refused = header_len(&head).unwrap_err();
1405 assert!(refused.contains("limits.npy_header_bytes"), "{refused}");
1406 head[8..12].copy_from_slice(&64u32.to_le_bytes());
1407 assert_eq!(header_len(&head), Ok((12, 64)));
1408 }
1409
1410 #[test]
1412 fn errors_name_the_file() {
1413 use crate::formats::readers::bad_input::{assert_shape, each_names_its_file, opening};
1414 let mut v9 = npy("'<f8'", false, "(1,)", &[0; 8]);
1415 v9[6] = 9;
1416 each_names_its_file(
1417 crate::FileFormat::Numpy,
1418 &[
1419 ("text.npy", b"hello there", "Not a NumPy file"),
1420 ("text.npz", b"hello there", "Not a NumPy archive"),
1421 ("v9.npy", &v9, "version 9"),
1422 ("type.npy", &npy("'<q9'", false, "(1,)", &[0; 8]), "<q9"),
1423 ],
1424 );
1425 let dir = tempfile::tempdir().unwrap();
1426 let options = crate::OpenOptions {
1427 table: Some("x".into()),
1428 ..Default::default()
1429 };
1430 let one = npy("'<f8'", false, "(1,)", &[0; 8]);
1431 let message = opening(
1432 dir.path(),
1433 "a.npy",
1434 &one,
1435 crate::FileFormat::Numpy,
1436 &options,
1437 )
1438 .expect("--table refused");
1439 eprintln!("{message}");
1440 assert_shape(&message, &dir.path().join("a.npy"));
1441 assert!(message.contains("--table \"x\""), "{message}");
1442 }
1443
1444 #[test]
1447 fn a_fortran_array_too_large_to_address_is_refused() {
1448 let bytes = npy("'<f4'", true, "(4000000000000000100, 4)", &[]);
1449 let len = bytes.len();
1450 assert!(open_in(Arc::new(Bytes::Owned(bytes)), 0, len, "a").is_err());
1451 }
1452
1453 fn npy(descr: &str, fortran: bool, shape: &str, data: &[u8]) -> Vec<u8> {
1454 let mut header = format!(
1455 "{{'descr': {descr}, 'fortran_order': {}, 'shape': {shape}, }}",
1456 if fortran { "True" } else { "False" }
1457 );
1458 while (10 + header.len() + 1) % 64 != 0 {
1459 header.push(' ');
1460 }
1461 header.push('\n');
1462 let mut out = MAGIC.to_vec();
1463 out.extend([1, 0]);
1464 out.extend((header.len() as u16).to_le_bytes());
1465 out.extend(header.as_bytes());
1466 out.extend(data);
1467 out
1468 }
1469
1470 fn read(bytes: Vec<u8>, name: &str) -> DataFrame {
1471 let len = bytes.len();
1472 let array = open_in(Arc::new(Bytes::Owned(bytes)), 0, len, name).unwrap();
1473 array.records.lazy().collect().unwrap()
1474 }
1475
1476 #[test]
1477 fn literals_parse() {
1478 let parsed = parse_literal("{'a': [('x', '<f8', (3,)), (('t', 'y'), '|b1')], 'b': (2L,), 'c': True, 'd': None, 'e': -4}")
1479 .unwrap();
1480 assert_eq!(
1481 parsed.get("b"),
1482 Some(&Literal::Tuple(vec![Literal::Int(2)]))
1483 );
1484 assert_eq!(parsed.get("c"), Some(&Literal::Bool(true)));
1485 assert_eq!(parsed.get("e"), Some(&Literal::Int(-4)));
1486 assert_eq!(parse_literal("(1)").unwrap(), Literal::Int(1));
1487 assert_eq!(parse_literal("()").unwrap(), Literal::Tuple(Vec::new()));
1488 assert!(parse_literal("{'a': 1").is_err());
1489 assert!(parse_literal(&"[".repeat(100)).is_err());
1490 assert!(parse_literal("__import__('os')").is_err());
1491 }
1492
1493 #[test]
1494 fn a_vector_is_one_column_named_for_the_array() {
1495 let data: Vec<u8> = [1.5f64, -2.0, 3.25]
1496 .iter()
1497 .flat_map(|v| v.to_le_bytes())
1498 .collect();
1499 let df = read(npy("'<f8'", false, "(3,)", &data), "prices");
1500 assert_eq!(df.get_column_names(), ["prices"]);
1501 assert_eq!(
1502 df.column("prices").unwrap().f64().unwrap().to_vec(),
1503 [Some(1.5), Some(-2.0), Some(3.25)]
1504 );
1505 }
1506
1507 #[test]
1508 fn a_grid_in_either_order() {
1509 let c: Vec<u8> = [1i32, 2, 3, 4, 5, 6]
1511 .iter()
1512 .flat_map(|v| v.to_le_bytes())
1513 .collect();
1514 let f: Vec<u8> = [1i32, 4, 2, 5, 3, 6]
1515 .iter()
1516 .flat_map(|v| v.to_le_bytes())
1517 .collect();
1518 for df in [
1519 read(npy("'<i4'", false, "(2, 3)", &c), "g"),
1520 read(npy("'<i4'", true, "(2, 3)", &f), "g"),
1521 ] {
1522 assert_eq!(df.get_column_names(), ["0", "1", "2"]);
1523 assert_eq!(
1524 df.column("2").unwrap().i32().unwrap().to_vec(),
1525 [Some(3), Some(6)]
1526 );
1527 }
1528 }
1529
1530 #[test]
1531 fn structured_fields_mixed_endian_subarrays_and_padding() {
1532 let mut data = Vec::new();
1534 for (ts, px, v) in [(1u64, 0.5f32, [7i16, 8]), (2, 1.5, [9, 10])] {
1535 data.extend(ts.to_le_bytes());
1536 data.extend(px.to_be_bytes());
1537 data.extend([0u8; 3]);
1538 data.extend(v.iter().flat_map(|x| x.to_le_bytes()));
1539 }
1540 let descr = "[('ts', '<u8'), ('px', '>f4'), ('', '|V3'), ('v', '<i2', (2,))]";
1541 let df = read(npy(descr, false, "(2,)", &data), "t");
1542 assert_eq!(df.get_column_names(), ["ts", "px", "v"]);
1543 assert_eq!(
1544 df.column("px").unwrap().f32().unwrap().to_vec(),
1545 [Some(0.5), Some(1.5)]
1546 );
1547 assert_eq!(
1548 df.column("v").unwrap().dtype(),
1549 &DataType::Array(Box::new(DataType::Int16), 2)
1550 );
1551 }
1552
1553 #[test]
1554 fn the_dict_form_honors_offsets_and_itemsize() {
1555 let mut data = Vec::new();
1557 for (a, b) in [(10i32, 20i32), (30, 40)] {
1558 data.extend(b.to_le_bytes());
1559 data.extend(a.to_le_bytes());
1560 data.extend([0u8; 4]);
1561 }
1562 let descr =
1563 "{'names': ['a', 'b'], 'formats': ['<i4', '<i4'], 'offsets': [4, 0], 'itemsize': 12}";
1564 let df = read(npy(descr, false, "(2,)", &data), "t");
1565 assert_eq!(
1566 df.column("a").unwrap().i32().unwrap().to_vec(),
1567 [Some(10), Some(30)]
1568 );
1569 assert_eq!(
1570 df.column("b").unwrap().i32().unwrap().to_vec(),
1571 [Some(20), Some(40)]
1572 );
1573 }
1574
1575 #[test]
1576 fn strings_times_and_bools() {
1577 let mut data = Vec::new();
1578 for c in "hé".chars().chain(std::iter::once('\0')) {
1580 data.extend((c as u32).to_le_bytes());
1581 }
1582 data.extend(b"ab\0\0");
1583 data.extend(86_400i64.to_le_bytes());
1584 data.extend(1500i64.to_le_bytes());
1585 data.push(1);
1586 let descr = "[('u', '<U3'), ('s', '|S4'), ('t', '<M8[s]'), ('d', '<m8[ms]'), ('b', '|b1')]";
1587 let df = read(npy(descr, false, "(1,)", &data), "t");
1588 assert_eq!(df.column("u").unwrap().str().unwrap().get(0), Some("hé"));
1589 assert_eq!(df.column("s").unwrap().str().unwrap().get(0), Some("ab"));
1590 assert_eq!(
1591 df.column("t").unwrap().dtype(),
1592 &DataType::Datetime(TimeUnit::Milliseconds, None)
1593 );
1594 assert_eq!(
1595 df.column("d").unwrap().dtype(),
1596 &DataType::Duration(TimeUnit::Milliseconds)
1597 );
1598 assert_eq!(df.column("b").unwrap().bool().unwrap().get(0), Some(true));
1599 }
1600
1601 #[test]
1602 fn what_is_refused() {
1603 let refused = |descr: &str, shape: &str| {
1604 let bytes = npy(descr, false, shape, &[0u8; 64]);
1605 let len = bytes.len();
1606 open_in(Arc::new(Bytes::Owned(bytes)), 0, len, "x")
1607 .err()
1608 .unwrap()
1609 };
1610 assert!(refused("'|O'", "(1,)").contains("does not unpickle"));
1611 assert!(refused("'<f8'", "(1, 1, 1)").contains("shape (1, 1, 1)"));
1612 assert!(refused("'<f16'", "(1,)").contains("not read"));
1613 }
1614
1615 #[test]
1616 fn a_short_file_keeps_the_rows_it_has() {
1617 let data: Vec<u8> = [1i16, 2, 3].iter().flat_map(|v| v.to_le_bytes()).collect();
1618 let bytes = npy("'<i2'", false, "(10,)", &data);
1619 let len = bytes.len();
1620 let array = open_in(Arc::new(Bytes::Owned(bytes)), 0, len, "x").unwrap();
1621 assert_eq!(array.records.rows(), 3);
1622 assert_eq!(array.notes.len(), 1);
1623 }
1624}