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