o41-training-data 0.1.6

use crate::{
    columns,
    config::{Field, Source},
    storage::{Client, Input, Object},
};
use anyhow::{Result, ensure};
#[derive(Default)]
struct Location {
    record: usize,
    byte: u64,
    group: usize,
}
use arrow::record_batch::RecordBatch;
use bytes::{Buf, Bytes};
use rmpv::Value;
use std::{
    io::{BufRead, BufReader, Cursor, Read},
    ops::Range,
    sync::{Arc, Mutex},
};

/// Parquet tail read with the footer; it holds the metadata of most shards.
const TAIL: u64 = 64 * 1024;
/// Window for reads outside prefetched bytes, such as Parquet page headers.
const HEADER_WINDOW: u64 = 64 * 1024;
/// Bytes a sample reads ahead: a record, or a dictionary and first data page.
const SAMPLE_BYTES: u64 = 2 * 1024 * 1024;

fn validate(value: &Value) -> Result<()> {
    match value {
        Value::Map(m) => {
            for (k, v) in m {
                ensure!(k.as_str().is_some(), "map key must be UTF-8 string");
                validate(v)?;
            }
        }
        Value::Array(a) => {
            for v in a {
                validate(v)?;
            }
        }
        Value::String(s) => ensure!(s.as_str().is_some(), "invalid UTF-8"),
        _ => (),
    }
    Ok(())
}
fn msgpack(
    input: &mut impl Read,
    fields: &[Field],
    limit: usize,
    emit: &mut impl FnMut(Value) -> Result<bool>,
    location: &mut Location,
) -> Result<()> {
    loop {
        let mut header = [0u8; 8];
        if input.read(&mut header[..1])? == 0 {
            break;
        }
        input.read_exact(&mut header[1..])?;
        let size = u64::from_be_bytes(header);
        ensure!(
            size > 0 && size <= (256 * 1024 * 1024).min(limit as u64),
            "invalid frame length"
        );
        let mut compressed = vec![0; size as usize];
        input.read_exact(&mut compressed)?;
        let decompress = crate::profile::Timer::new(1);
        let frame_size = zstd::zstd_safe::find_frame_compressed_size(&compressed)
            .map_err(|e| anyhow::anyhow!("invalid zstd frame: {e}"))?;
        ensure!(
            frame_size == compressed.len(),
            "extra data after zstd frame"
        );
        let declared = zstd::zstd_safe::get_frame_content_size(&compressed)
            .map_err(|e| anyhow::anyhow!("invalid zstd content size: {e:?}"))?;
        ensure!(
            declared.is_none_or(|n| n <= limit as u64),
            "payload exceeds maxRecordBytes"
        );
        let mut payload = Vec::new();
        let mut decoder = zstd::stream::read::Decoder::new(compressed.as_slice())?;
        decoder.window_log_max(limit.max(1024).next_power_of_two().trailing_zeros())?;
        if let Err(error) = decoder.take(limit as u64 + 1).read_to_end(&mut payload) {
            location.byte += payload.len() as u64;
            return Err(error.into());
        }
        ensure!(payload.len() <= limit, "payload exceeds maxRecordBytes");
        drop(decompress);
        let parse = crate::profile::Timer::new(2);
        let mut cursor = Cursor::new(&payload);
        let row = (|| -> Result<Value> {
            let count = rmp::decode::read_map_len(&mut cursor)?;
            let mut values = Vec::new();
            for _ in 0..count {
                let key = rmpv::decode::read_value(&mut cursor)?;
                let name = key
                    .as_str()
                    .ok_or_else(|| anyhow::anyhow!("map key must be UTF-8 string"))?;
                if fields.iter().any(|f| f.from == name) {
                    values.push((key, rmpv::decode::read_value(&mut cursor)?));
                } else {
                    use serde::Deserialize;
                    serde::de::IgnoredAny::deserialize(&mut rmp_serde::Deserializer::new(
                        &mut cursor,
                    ))?;
                }
            }
            ensure!(
                cursor.position() == payload.len() as u64,
                "extra MessagePack data"
            );
            Ok(Value::Map(values))
        })();
        location.byte += cursor.position();
        let row = row?;
        ensure!(row.is_map(), "record must be a map");
        validate(&row)?;
        drop(parse);
        if !emit(row)? {
            break;
        }
        location.record += 1;
    }
    Ok(())
}
fn json_value(value: &jiter::JsonValue<'_>, strings: Option<&[Value]>) -> Result<Value> {
    use jiter::JsonValue as J;
    Ok(match value {
        J::Null => Value::Nil,
        J::Bool(b) => Value::Boolean(*b),
        J::Int(i) => Value::from(*i),
        J::BigInt(i) => {
            let digits = i.to_string();
            if let Ok(n) = digits.parse::<i64>() {
                Value::from(n)
            } else {
                Value::from(digits.parse::<u64>()?)
            }
        }
        J::Float(f) => Value::F64(*f),
        J::Str(s) => match strings {
            Some(values) => values[s.parse::<usize>()?].clone(),
            None => Value::from(s.as_ref()),
        },
        J::Array(a) => Value::Array(
            a.iter()
                .map(|v| json_value(v, strings))
                .collect::<Result<_>>()?,
        ),
        J::Object(m) => {
            let mut map = std::collections::BTreeMap::new();
            for (k, v) in m.iter() {
                let key = match strings {
                    Some(values) => values[k.parse::<usize>()?].as_str().ok_or_else(|| {
                        anyhow::anyhow!("Arrow field names require valid Unicode")
                    })?,
                    None => k.as_ref(),
                };
                map.insert(key, json_value(v, strings)?);
            }
            Value::Map(map.into_iter().map(|(k, v)| (Value::from(k), v)).collect())
        }
        #[allow(unreachable_patterns)]
        _ => Value::Nil,
    })
}
// serde_json's byte visitor preserves lone surrogates as WTF-8. Replace every
// string token with an index while jiter parses Python's non-finite numbers.
// Replacing all strings prevents placeholder collisions with user values.
struct StringBytes;
impl<'de> serde::de::Visitor<'de> for StringBytes {
    type Value = Vec<u8>;
    fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("a JSON string")
    }
    fn visit_bytes<E: serde::de::Error>(self, bytes: &[u8]) -> std::result::Result<Vec<u8>, E> {
        Ok(bytes.to_vec())
    }
}
fn json_tokens(line: &str) -> Result<(String, Vec<Value>)> {
    use serde::Deserializer;
    let mut out = String::with_capacity(line.len());
    let mut strings = Vec::new();
    let mut i = 0;
    let bytes = line.as_bytes();
    let mut start = 0;
    while i < bytes.len() {
        if bytes[i] != b'"' {
            i += 1;
            continue;
        }
        out.push_str(&line[start..i]);
        let token_start = i;
        i += 1;
        while i < bytes.len() {
            if bytes[i] == b'\\' {
                i = (i + 2).min(bytes.len());
            } else if bytes[i] == b'"' {
                i += 1;
                break;
            } else {
                i += 1;
            }
        }
        let mut decoder = serde_json::Deserializer::from_str(&line[token_start..i]);
        let decoded = decoder.deserialize_bytes(StringBytes)?;
        decoder.end()?;
        out.push_str(&format!("\"{}\"", strings.len()));
        strings.push(match String::from_utf8(decoded) {
            Ok(s) => Value::from(s),
            Err(e) => Value::Binary(e.into_bytes()),
        });
        start = i;
    }
    out.push_str(&line[start..]);
    Ok((out, strings))
}
fn selected_json(line: &str, fields: &[Field], strings: Option<&[Value]>) -> Result<Value> {
    let mut parser = jiter::Jiter::new(line.as_bytes()).with_allow_inf_nan();
    let mut key = parser.next_object()?;
    let mut values = Vec::new();
    while let Some(name) = key {
        // A key that is not valid Unicode cannot name a projected field.
        let name = match strings {
            Some(values) => values[name.parse::<usize>()?].as_str(),
            None => Some(name),
        };
        match name.filter(|name| fields.iter().any(|f| f.from == *name)) {
            Some(name) => {
                let name = name.to_owned();
                let value = {
                    let _parse = crate::profile::Timer::new(2);
                    parser.next_value()?
                };
                let _value = crate::profile::Timer::new(3);
                values.push((Value::from(name), json_value(&value, strings)?));
            }
            None => {
                let _parse = crate::profile::Timer::new(2);
                parser.next_skip()?;
            }
        }
        key = parser.next_key()?;
    }
    parser.finish()?;
    Ok(Value::Map(values))
}
fn jsonl(
    input: Input,
    fields: &[Field],
    limit: usize,
    emit: &mut impl FnMut(Value) -> Result<bool>,
    location: &mut Location,
) -> Result<()> {
    let mut decoder = flate2::bufread::GzDecoder::new(BufReader::new(input));
    {
        let mut lines = BufReader::new(&mut decoder);
        let mut bytes = Vec::new();
        loop {
            bytes.clear();
            let decompress = crate::profile::Timer::new(1);
            let read = (&mut lines)
                .take(limit as u64 + 1)
                .read_until(b'\n', &mut bytes);
            location.byte += bytes.len() as u64;
            let n = read?;
            let line = std::str::from_utf8(&bytes)?;
            ensure!(n <= limit, "JSON line exceeds maxRecordBytes");
            drop(decompress);
            if n == 0 {
                break;
            }
            {
                match selected_json(line, fields, None) {
                    Ok(row) => {
                        if !emit(row)? {
                            return Ok(());
                        }
                        location.record += 1;
                    }
                    Err(_) => {
                        // jiter rejects lone surrogates even in skipped strings.
                        // Normalize string tokens, then still skip unused arrays.
                        let Ok((normalized, strings)) = json_tokens(line) else {
                            continue;
                        };
                        match selected_json(&normalized, fields, Some(&strings)) {
                            Ok(row) => {
                                if !emit(row)? {
                                    return Ok(());
                                }
                                location.record += 1;
                            }
                            Err(error) => {
                                let mut syntax =
                                    jiter::Jiter::new(normalized.as_bytes()).with_allow_inf_nan();
                                if normalized.trim_start().starts_with('{')
                                    && syntax.next_skip().is_ok()
                                    && syntax.finish().is_ok()
                                {
                                    return Err(error);
                                }
                            }
                        }
                    }
                }
                continue;
            }
        }
    }
    ensure!(
        decoder.into_inner().fill_buf()?.is_empty(),
        "extra data after gzip member"
    );
    Ok(())
}

/// A column chunk read in slices: its bytes, the slice of it read last, and
/// where its next page header starts.
struct Sliced {
    chunk: Range<u64>,
    slice: Mutex<(u64, Bytes)>,
    header: Mutex<u64>,
}

/// A Parquet object with some byte ranges already fetched. Reads inside them
/// are served from memory; any other read goes to storage.
#[derive(Clone)]
struct ParquetObject {
    object: Arc<Mutex<Object>>,
    size: u64,
    fetched: Arc<Vec<(u64, Bytes)>>,
    /// The column chunks of a row group too large to fetch whole, the bytes
    /// a slice of one is fetched in, and maxRowGroupBytes.
    slices: Arc<Vec<Sliced>>,
    slice: u64,
    budget: u64,
}
impl ParquetObject {
    /// The bytes held in memory from `start` on. In a sliced column chunk, a
    /// read of `length` bytes past the slice fetches the next one from `start`.
    fn held(&self, start: u64, length: u64) -> parquet::errors::Result<Option<Bytes>> {
        if let Some((offset, bytes)) = self
            .fetched
            .iter()
            .find(|(offset, bytes)| *offset <= start && start < offset + bytes.len() as u64)
        {
            return Ok(Some(bytes.slice((start - offset) as usize..)));
        }
        let Some(Sliced { chunk, slice, .. }) =
            self.slices.iter().find(|s| s.chunk.contains(&start))
        else {
            return Ok(None);
        };
        let mut slice = slice.lock().unwrap();
        let (offset, bytes) = &*slice;
        if !(*offset <= start && (start + length).min(chunk.end) <= offset + bytes.len() as u64) {
            let end = (start + self.slice.max(length)).min(chunk.end);
            let bytes = self
                .object
                .lock()
                .unwrap()
                .range(start, end)
                .map_err(|e| parquet::errors::ParquetError::General(e.to_string()))?;
            *slice = (start, bytes);
        }
        Ok(Some(slice.1.slice((start - slice.0) as usize..)))
    }
}
/// A Thrift compact struct read from the front of a buffer, enough to measure
/// a Parquet page header, whose parser the parquet crate keeps private.
struct Compact<'a>(&'a [u8], usize);
impl Compact<'_> {
    fn byte(&mut self) -> Option<u8> {
        let byte = *self.0.get(self.1)?;
        self.1 += 1;
        Some(byte)
    }
    fn varint(&mut self) -> Option<u64> {
        let mut value = 0;
        for shift in (0..64).step_by(7) {
            let byte = self.byte()?;
            if shift == 63 && byte > 1 {
                return None;
            }
            value |= u64::from(byte & 0x7f) << shift;
            if byte < 0x80 {
                return Some(value);
            }
        }
        None
    }
    fn int(&mut self) -> Option<i64> {
        let value = self.varint()?;
        Some((value >> 1) as i64 ^ -((value & 1) as i64))
    }
    /// Skips a value of a compact type; a bool outside a field is one byte.
    fn skip(&mut self, kind: u8, depth: u32) -> Option<()> {
        if depth > 16 {
            return None;
        }
        match kind {
            1..=3 => drop(self.byte()?),
            4..=6 => drop(self.int()?),
            7 => self.1 = self.1.checked_add(8).filter(|&n| n <= self.0.len())?,
            8 => {
                let length = usize::try_from(self.varint()?).ok()?;
                self.1 = self.1.checked_add(length).filter(|&n| n <= self.0.len())?;
            }
            9 | 10 => {
                let header = self.byte()?;
                let count = match header >> 4 {
                    15 => self.varint()?,
                    n => u64::from(n),
                };
                for _ in 0..count {
                    self.skip(header & 0x0f, depth + 1)?;
                }
            }
            11 => {
                let count = self.varint()?;
                if count > 0 {
                    let kinds = self.byte()?;
                    for _ in 0..count {
                        self.skip(kinds >> 4, depth + 1)?;
                        self.skip(kinds & 0x0f, depth + 1)?;
                    }
                }
            }
            12 => self.fields(depth + 1, |_, _, _| None)?,
            _ => return None,
        }
        Some(())
    }
    /// Reads a struct's fields up to its stop byte, giving each field's id and
    /// type to `field`, which reads the value or returns `None` to skip it.
    fn fields(
        &mut self,
        depth: u32,
        mut field: impl FnMut(&mut Self, i64, u8) -> Option<()>,
    ) -> Option<()> {
        let mut id = 0;
        loop {
            let header = self.byte()?;
            let kind = header & 0x0f;
            if kind == 0 {
                return Some(());
            }
            id = match header >> 4 {
                0 => self.int()?,
                delta => id.checked_add(i64::from(delta))?,
            };
            let at = self.1;
            if field(self, id, kind).is_none() {
                self.1 = at;
                // A bool field's value is in its type.
                if kind > 2 {
                    self.skip(kind, depth)?;
                }
            }
        }
    }
}
/// A Parquet page header's length and the larger of its uncompressed and
/// compressed page sizes, or `None` if it does not read as one.
fn page_header(bytes: &[u8]) -> Option<(u64, u64)> {
    let mut header = Compact(bytes, 0);
    let mut sizes = [None, None];
    header.fields(0, |header, id, kind| {
        let size = sizes.get_mut(usize::try_from(id.checked_sub(2)?).ok()?)?;
        (kind == 5).then_some(())?;
        *size = Some(i32::try_from(header.int()?).ok()?);
        Some(())
    })?;
    let [Some(uncompressed), Some(compressed)] = sizes else {
        return None;
    };
    Some((
        (header.1 as u64).checked_add(u64::try_from(compressed).ok()?)?,
        u64::try_from(uncompressed.max(compressed)).ok()?,
    ))
}
impl parquet::file::reader::Length for ParquetObject {
    fn len(&self) -> u64 {
        self.size
    }
}
impl parquet::file::reader::ChunkReader for ParquetObject {
    type T = std::io::Chain<bytes::buf::Reader<Bytes>, Input>;
    fn get_read(&self, start: u64) -> parquet::errors::Result<Self::T> {
        let head = self
            .held(start, HEADER_WINDOW.min(self.slice))?
            .unwrap_or_default();
        // A sliced row group bounds memory by its slices only while each page,
        // held whole while decoded, fits in one. The decoder reads a header
        // ahead and may come back to its page, so only the next header in the
        // chunk is read here.
        let page = self.slice;
        if let Some(column) = self.slices.iter().find(|s| s.chunk.contains(&start))
            && let mut next = column.header.lock().unwrap()
            && *next == start
        {
            let header = page_header(&head);
            if let Some((length, size)) = header
                && size <= page
            {
                *next = start + length;
            } else {
                let what = header.map_or("an unreadable Parquet page header".into(), |(_, n)| {
                    format!("a Parquet page of {n} bytes")
                });
                return Err(parquet::errors::ParquetError::General(format!(
                    "{what} in a row group over {} bytes, where a page may hold {page}; \
                     write smaller row groups, such as pyarrow's row_group_size=4096",
                    self.budget
                )));
            }
        }
        let end = start + head.len() as u64;
        Ok(head
            .reader()
            .chain(Input::at(self.object.clone(), end, HEADER_WINDOW)))
    }
    fn get_bytes(&self, start: u64, length: usize) -> parquet::errors::Result<Bytes> {
        if let Some(bytes) = self.held(start, length as u64)?
            && length <= bytes.len()
        {
            return Ok(bytes.slice(..length));
        }
        self.object
            .lock()
            .unwrap()
            .range(start, start + length as u64)
            .map_err(|e| parquet::errors::ParquetError::General(e.to_string()))
    }
}
fn trim(value: &mut Value, widths: &[usize]) {
    if let Some((width, inner)) = widths.split_first()
        && let Value::Array(values) = value
    {
        values.truncate(*width);
        values.shrink_to_fit();
        for value in values {
            trim(value, inner);
        }
    }
}
pub struct Shard {
    pub offset: usize,
    pub columns: RecordBatch,
}
impl Shard {
    pub fn num_rows(&self) -> usize {
        self.offset + self.columns.num_rows()
    }
}
pub fn read(
    source: &Source,
    encoding: &str,
    client: &Client,
    fields: &[Field],
    start: usize,
    sample: bool,
    config: &crate::config::Config,
) -> Result<Shard> {
    let mut location = Location::default();
    let result = (|| -> Result<Shard> {
        let object = Object::open(source, client)?;
        if encoding == "parquet" {
            return parquet_read(
                object,
                fields,
                start,
                sample,
                config,
                source.rows,
                &mut location,
            );
        }
        let mut chunks = Vec::new();
        let mut rows = Vec::new();
        let mut count = 0;
        let mut emit = |mut row: Value| -> Result<bool> {
            count += 1;
            ensure!(
                count <= source.rows.unwrap_or(config.max_shard_rows),
                "shard exceeds declared rows or maxShardRows"
            );
            if count <= start {
                return Ok(true);
            }
            // Retain only the declared prefix on every axis while accumulating
            // one builder chunk, so retained rows stay within the memory bound.
            let Value::Map(ref mut map) = row else {
                anyhow::bail!("record must be a map");
            };
            for (key, value) in map {
                let mut widths: Vec<usize> = Vec::new();
                for field in fields
                    .iter()
                    .filter(|f| Some(f.from.as_str()) == key.as_str())
                {
                    for (axis, width) in field.dimensions().into_iter().enumerate() {
                        match widths.get_mut(axis) {
                            Some(current) => *current = (*current).max(width),
                            None => widths.push(width),
                        }
                    }
                }
                trim(value, &widths);
            }
            rows.push(row);
            if rows.len() == 128 || sample {
                chunks.push(columns::project_values(
                    &rows.iter().collect::<Vec<_>>(),
                    fields,
                )?);
                rows.clear();
            }
            Ok(!sample)
        };
        // A sample reads one record; a full read streams the object ahead.
        let mut input = if sample {
            Input::at(Arc::new(Mutex::new(object)), 0, SAMPLE_BYTES)
        } else {
            Input::stream(object)?
        };
        if encoding == "msgpack.zstf" {
            msgpack(
                &mut input,
                fields,
                config.max_record_bytes,
                &mut emit,
                &mut location,
            )?;
        } else {
            jsonl(
                input,
                fields,
                config.max_record_bytes,
                &mut emit,
                &mut location,
            )?;
        }
        if !sample
            && encoding == "msgpack.zstf"
            && let Some(n) = source.rows
        {
            ensure!(count == n, "MessagePack row count mismatch");
        }
        if !rows.is_empty() {
            chunks.push(columns::project_values(
                &rows.iter().collect::<Vec<_>>(),
                fields,
            )?);
        }
        let schema = columns::project_values(&[], fields)?.schema();
        Ok(Shard {
            offset: start,
            columns: arrow::compute::concat_batches(&schema, chunks.iter())?,
        })
    })();
    crate::profile::shard(
        &source.uri,
        sample,
        location.byte,
        result
            .as_ref()
            .map_or(0, |s| s.columns.get_array_memory_size() as u64),
    );
    result.map_err(|error| {
        let position = if encoding == "parquet" {
            format!("row_group={} row={}", location.group, location.record)
        } else {
            format!("record={} byte_offset={}", location.record, location.byte)
        };
        anyhow::anyhow!(
            "shard={} {position}: {error:#}",
            serde_json::to_string(&source.uri).unwrap()
        )
    })
}

fn parquet_read(
    object: Object,
    fields: &[Field],
    start: usize,
    sample: bool,
    config: &crate::config::Config,
    declared_rows: Option<usize>,
    location: &mut Location,
) -> Result<Shard> {
    use parquet::arrow::{
        ProjectionMask,
        arrow_reader::{
            ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReaderBuilder, RowSelection,
            RowSelector,
        },
    };
    let size = object.size;
    let object = Arc::new(Mutex::new(object));
    let input = |fetched, slices: Vec<Range<u64>>| ParquetObject {
        object: object.clone(),
        size,
        fetched: Arc::new(fetched),
        // Three slices of a column can be alive at once: one still held by the
        // page being decoded, one by the reader of the next page's header, and
        // the one being fetched. A page fits in a slice, so all of them
        // together stay within the budget.
        slice: (config.max_row_group_bytes / (3 * slices.len().max(1))) as u64,
        budget: config.max_row_group_bytes as u64,
        slices: Arc::new(
            slices
                .into_iter()
                .map(|chunk| Sliced {
                    header: Mutex::new(chunk.start),
                    slice: Mutex::new((0, Bytes::new())),
                    chunk,
                })
                .collect(),
        ),
    };
    let tail = size.saturating_sub(TAIL);
    let footer = object.lock().unwrap().range(tail, size)?;
    let metadata = ArrowReaderMetadata::load(
        &input(vec![(tail, footer)], vec![]),
        ArrowReaderOptions::default(),
    )?;
    let total = metadata.metadata().file_metadata().num_rows() as usize;
    ensure!(
        declared_rows.is_none_or(|n| n == total),
        "Parquet row count mismatch"
    );
    ensure!(
        total <= config.max_shard_rows && start <= total,
        "Parquet exceeds maxShardRows or cursor bounds"
    );
    let indices: Vec<_> = metadata
        .schema()
        .fields()
        .iter()
        .enumerate()
        .filter(|(_, f)| fields.iter().any(|p| p.from == *f.name()))
        .map(|(i, _)| i)
        .collect();
    let mask = ProjectionMask::roots(metadata.metadata().file_metadata().schema_descr(), indices);
    // Plan the row groups to read with the bytes of their projected column
    // chunks, then fetch those ahead of decoding. A row group whose projected
    // columns exceed maxRowGroupBytes is not fetched ahead but read in slices
    // as it is decoded, which holds at most maxRowGroupBytes of it, a sample
    // too, so that it checks the pages it reads.
    let mut skipped = 0;
    let mut plan = Vec::new();
    for (i, group) in metadata.metadata().row_groups().iter().enumerate() {
        let rows = group.num_rows() as usize;
        if skipped + rows <= start {
            skipped += rows;
            continue;
        }
        location.group = i;
        location.record = start.saturating_sub(skipped);
        let columns: Vec<_> = group
            .columns()
            .iter()
            .enumerate()
            .filter(|(i, _)| mask.leaf_included(*i))
            .map(|(_, c)| c)
            .collect();
        let bytes: i64 = columns.iter().map(|c| c.uncompressed_size()).sum();
        ensure!(bytes >= 0, "invalid Parquet row group size");
        let sliced = bytes as usize > config.max_row_group_bytes;
        let mut ranges = Vec::new();
        for c in columns {
            let offset = c.dictionary_page_offset().unwrap_or(c.data_page_offset());
            let length = c.compressed_size();
            ensure!(offset >= 0 && length >= 0, "invalid Parquet column chunk");
            let length = if sample && !sliced {
                length.min(SAMPLE_BYTES as i64)
            } else {
                length
            };
            ranges.push(offset as u64..(offset + length) as u64);
        }
        plan.push((i, location.record, rows, ranges, sliced));
        if sample {
            break;
        }
        skipped += rows;
    }
    let mut fetched = object.lock().unwrap().prefetch(
        plan.iter()
            .map(|(.., ranges, sliced)| if *sliced { vec![] } else { ranges.clone() })
            .collect(),
    )?;
    let mut chunks = Vec::new();
    for (i, first, rows, ranges, sliced) in plan {
        location.group = i;
        location.record = first;
        let bodies = fetched
            .next()
            .ok_or_else(|| anyhow::anyhow!("Parquet prefetch stopped"))??;
        let selection = RowSelection::from(vec![
            RowSelector::skip(first),
            RowSelector::select(if sample { 1 } else { rows - first }),
        ]);
        let reader = ParquetRecordBatchReaderBuilder::new_with_metadata(
            if sliced {
                input(vec![], ranges)
            } else {
                input(ranges.iter().map(|r| r.start).zip(bodies).collect(), vec![])
            },
            metadata.clone(),
        )
        .with_projection(mask.clone())
        .with_row_groups(vec![i])
        .with_row_selection(selection)
        .with_batch_size(if sample { 1 } else { 128 })
        .build()?;
        for batch in reader {
            let batch = batch?;
            chunks.push(columns::project(&batch, fields)?);
            location.record += batch.num_rows();
            if sample {
                break;
            }
        }
    }
    let schema = columns::project_values(&[], fields)?.schema();
    Ok(Shard {
        offset: start,
        columns: arrow::compute::concat_batches(&schema, chunks.iter())?,
    })
}