teeny_vision/mnist/
dataset.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17//! Reader for HuggingFace MNIST parquet files.
18//!
19//! Schema (from `ylecun/mnist`):
20//!   - `image`: struct<bytes: binary, path: utf8>  — PNG-encoded 28×28 grayscale
21//!   - `label`: int64                              — class 0-9
22
23use std::{
24    fs::File,
25    path::{Path, PathBuf},
26};
27
28use anyhow::{Result, anyhow};
29use arrow_array::{Array, BinaryArray, Int64Array, LargeBinaryArray, StructArray};
30use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
31
32/// A loaded batch of MNIST samples.
33pub struct MnistBatch {
34    /// Pixel data in NCHW layout `[batch, 1, 28, 28]`, normalised to `[0.0, 1.0]`.
35    pub images: Vec<f32>,
36    /// Integer class labels `[0, 9]`, one per sample.
37    pub labels: Vec<u8>,
38    /// Number of samples in this batch.
39    pub batch_size: usize,
40}
41
42impl MnistBatch {
43    /// Total number of f32 elements: `batch_size × 1 × 28 × 28`.
44    pub fn image_numel(&self) -> usize {
45        self.batch_size * 28 * 28
46    }
47}
48
49/// Reads MNIST samples from a HuggingFace parquet file.
50pub struct MnistDataset {
51    path: PathBuf,
52    num_rows: usize,
53}
54
55impl MnistDataset {
56    /// Open a parquet file and read its metadata.
57    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
58        let path = path.as_ref().to_path_buf();
59        let file = File::open(&path)?;
60        let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
61        let num_rows = builder.metadata().file_metadata().num_rows() as usize;
62        Ok(Self { path, num_rows })
63    }
64
65    /// Total number of rows in the dataset.
66    pub fn len(&self) -> usize {
67        self.num_rows
68    }
69
70    /// Whether the dataset has zero rows.
71    pub fn is_empty(&self) -> bool {
72        self.num_rows == 0
73    }
74
75    /// Read `count` consecutive rows starting at `offset` (0-indexed).
76    ///
77    /// The returned batch is truncated if `offset + count > len()`.
78    pub fn read_batch(&self, offset: usize, count: usize) -> Result<MnistBatch> {
79        let actual = count.min(self.num_rows.saturating_sub(offset));
80        if actual == 0 {
81            return Err(anyhow!(
82                "offset {offset} is at or past the end of the dataset (len={})",
83                self.num_rows
84            ));
85        }
86
87        let file = File::open(&self.path)?;
88        let reader = ParquetRecordBatchReaderBuilder::try_new(file)?
89            .with_offset(offset)
90            .with_limit(actual)
91            .with_batch_size(actual)
92            .build()?;
93
94        let batch = reader
95            .into_iter()
96            .next()
97            .ok_or_else(|| anyhow!("parquet reader returned no record batches"))??;
98
99        decode_batch(&batch, actual)
100    }
101}
102
103fn decode_batch(batch: &arrow_array::RecordBatch, n: usize) -> Result<MnistBatch> {
104    // ── image column — struct<bytes: binary|large_binary, path: utf8> ──────────
105    let image_struct = batch
106        .column_by_name("image")
107        .ok_or_else(|| anyhow!("column 'image' not found in parquet file"))?
108        .as_any()
109        .downcast_ref::<StructArray>()
110        .ok_or_else(|| anyhow!("column 'image' is not a StructArray"))?;
111
112    // The `bytes` child may be Binary or LargeBinary depending on how HF serialised it.
113    let bytes_field = image_struct
114        .column_by_name("bytes")
115        .ok_or_else(|| anyhow!("image struct has no 'bytes' field"))?;
116
117    // ── label column — int64 ───────────────────────────────────────────────────
118    let label_col = batch
119        .column_by_name("label")
120        .ok_or_else(|| anyhow!("column 'label' not found in parquet file"))?
121        .as_any()
122        .downcast_ref::<Int64Array>()
123        .ok_or_else(|| anyhow!("column 'label' is not Int64Array"))?;
124
125    let mut images = Vec::with_capacity(n * 28 * 28);
126    let mut labels = Vec::with_capacity(n);
127
128    for i in 0..n {
129        let png_bytes: &[u8] = if let Some(col) = bytes_field.as_any().downcast_ref::<BinaryArray>()
130        {
131            col.value(i)
132        } else if let Some(col) = bytes_field.as_any().downcast_ref::<LargeBinaryArray>() {
133            col.value(i)
134        } else {
135            return Err(anyhow!(
136                "'bytes' field has unexpected Arrow type: {:?}",
137                bytes_field.data_type()
138            ));
139        };
140
141        let img = image::load_from_memory(png_bytes)
142            .map_err(|e| anyhow!("failed to decode PNG for sample {i}: {e}"))?
143            .into_luma8();
144
145        if img.width() != 28 || img.height() != 28 {
146            return Err(anyhow!(
147                "unexpected image dimensions for sample {i}: {}×{}",
148                img.width(),
149                img.height()
150            ));
151        }
152
153        for px in img.pixels() {
154            images.push(px.0[0] as f32 / 255.0);
155        }
156
157        labels.push(label_col.value(i) as u8);
158    }
159
160    Ok(MnistBatch {
161        images,
162        labels,
163        batch_size: n,
164    })
165}
teeny_vision/mnist/dataset.rs

teeny_vision/mnist/
dataset.rs
