netget/client/snowflake/
arrow.rs

1//! A bounded reader for the Arrow IPC **stream** a Snowflake server sends as `rowsetBase64`.
2//!
3//! Snowflake answers a query in one of two result formats: JSON (`rowset`, every cell a
4//! string) or Arrow (`rowsetBase64`, an Arrow IPC stream). A client can ask for JSON, but a
5//! server is free to answer Arrow, and fakesnow answers every SELECT in Arrow whatever it was
6//! asked. So this reads the subset of Arrow a Snowflake result uses — no more:
7//!
8//! * the stream framing (`0xFFFFFFFF`, metadata length, flatbuffer `Message`, body), with the
9//!   pre-0.15 framing (no continuation marker) accepted too;
10//! * a `Schema` of flat columns, where a column may be a `Struct` of flat children (Snowflake
11//!   sends timestamps as `{epoch, fraction[, timezone]}`) — and no deeper;
12//! * `RecordBatch`es of Int (8–64 bit), FloatingPoint (32/64), Bool, Utf8, Binary, Decimal128,
13//!   Date, Time, Timestamp, Null and that one level of Struct.
14//!
15//! Everything else is refused by name: dictionaries, compression, big-endian data, lists,
16//! maps, unions, views, 64-bit offsets, Decimal256. The flatbuffers are read by hand with every
17//! offset checked against the buffer, and nothing here recurses — a nested field is read in a
18//! second loop, and a field nested deeper than [`MAX_FIELD_NESTING`] is refused.
19
20use anyhow::{bail, ensure, Context, Result};
21
22/// Messages (schema, batches) in one stream.
23pub const MAX_MESSAGES: usize = 4096;
24/// Columns in one schema, counting a struct's children.
25pub const MAX_FIELDS: usize = 2048;
26/// How deep a field may sit inside another. Snowflake's timestamp structs are one level deep.
27pub const MAX_FIELD_NESTING: usize = 1;
28
29/// A column's physical type, as the schema declares it.
30#[derive(Debug, Clone, PartialEq)]
31pub enum ArrowType {
32    Null,
33    Int {
34        bits: u8,
35        signed: bool,
36    },
37    Float {
38        bits: u8,
39    },
40    Bool,
41    Utf8,
42    Binary,
43    Decimal {
44        precision: i32,
45        scale: i32,
46    },
47    /// `day`: Date32 in days, else Date64 in milliseconds.
48    Date {
49        day: bool,
50    },
51    /// The unit as a power of ten below a second (0, 3, 6, 9) and the width.
52    Time {
53        scale: u32,
54        bits: u8,
55    },
56    Timestamp {
57        scale: u32,
58        tz: Option<String>,
59    },
60    Struct,
61}
62
63/// One column of the schema, with its children if it is a struct.
64#[derive(Debug, Clone)]
65pub struct Field {
66    pub name: String,
67    pub data_type: ArrowType,
68    /// Snowflake's per-field metadata: `logicalType`, `scale`, `precision`, ...
69    pub metadata: Vec<(String, String)>,
70    pub children: Vec<Field>,
71}
72
73impl Field {
74    pub fn meta(&self, key: &str) -> Option<&str> {
75        self.metadata
76            .iter()
77            .find(|(k, _)| k == key)
78            .map(|(_, v)| v.as_str())
79    }
80}
81
82/// One cell, as the physical type holds it.
83#[derive(Debug, Clone, PartialEq)]
84pub enum Cell {
85    Null,
86    Int(i128),
87    Float(f64),
88    Bool(bool),
89    Text(String),
90    /// Binary data is never handed on as bytes; only its length.
91    Bytes(usize),
92    /// A struct's children, in schema order.
93    Struct(Vec<Cell>),
94}
95
96/// The decoded stream: the schema and up to `max_rows` rows.
97#[derive(Debug, Clone)]
98pub struct Decoded {
99    pub fields: Vec<Field>,
100    pub rows: Vec<Vec<Cell>>,
101    /// Rows the stream held, including those past `max_rows`.
102    pub total_rows: u64,
103}
104
105// ---------------------------------------------------------------------------------------------
106// Flatbuffers, by hand.
107
108#[derive(Clone, Copy)]
109struct Fb<'a> {
110    buf: &'a [u8],
111}
112
113impl<'a> Fb<'a> {
114    fn bytes(&self, pos: usize, n: usize) -> Result<&'a [u8]> {
115        let end = pos.checked_add(n).context("flatbuffer offset overflows")?;
116        self.buf
117            .get(pos..end)
118            .with_context(|| format!("flatbuffer read of {n} bytes at {pos} runs past its end"))
119    }
120    fn u8(&self, pos: usize) -> Result<u8> {
121        Ok(self.bytes(pos, 1)?[0])
122    }
123    fn u16(&self, pos: usize) -> Result<u16> {
124        let b = self.bytes(pos, 2)?;
125        Ok(u16::from_le_bytes([b[0], b[1]]))
126    }
127    fn i32(&self, pos: usize) -> Result<i32> {
128        let b = self.bytes(pos, 4)?;
129        Ok(i32::from_le_bytes([b[0], b[1], b[2], b[3]]))
130    }
131    fn u32(&self, pos: usize) -> Result<u32> {
132        Ok(self.i32(pos)? as u32)
133    }
134    fn i64(&self, pos: usize) -> Result<i64> {
135        let b = self.bytes(pos, 8)?;
136        let mut a = [0u8; 8];
137        a.copy_from_slice(b);
138        Ok(i64::from_le_bytes(a))
139    }
140    /// The position an unsigned offset stored at `pos` points to.
141    fn follow(&self, pos: usize) -> Result<usize> {
142        let target = pos
143            .checked_add(self.u32(pos)? as usize)
144            .context("flatbuffer offset overflows")?;
145        ensure!(
146            target < self.buf.len(),
147            "flatbuffer offset points past its end"
148        );
149        Ok(target)
150    }
151    fn table(&self, pos: usize) -> Result<Table<'a>> {
152        let soffset = self.i32(pos)? as i64;
153        let vt = pos as i64 - soffset;
154        ensure!(
155            vt >= 0 && (vt as usize) < self.buf.len(),
156            "flatbuffer vtable out of range"
157        );
158        let vt = vt as usize;
159        let vt_len = self.u16(vt)? as usize;
160        ensure!(
161            vt_len >= 4 && vt_len.is_multiple_of(2),
162            "malformed flatbuffer vtable"
163        );
164        self.bytes(vt, vt_len)?;
165        Ok(Table {
166            fb: *self,
167            pos,
168            vt,
169            vt_len,
170        })
171    }
172}
173
174#[derive(Clone, Copy)]
175struct Table<'a> {
176    fb: Fb<'a>,
177    pos: usize,
178    vt: usize,
179    vt_len: usize,
180}
181
182impl<'a> Table<'a> {
183    fn field(&self, slot: usize) -> Result<Option<usize>> {
184        let at = 4 + 2 * slot;
185        if at + 2 > self.vt_len {
186            return Ok(None);
187        }
188        match self.fb.u16(self.vt + at)? {
189            0 => Ok(None),
190            off => Ok(Some(self.pos + off as usize)),
191        }
192    }
193    fn u8_or(&self, slot: usize, default: u8) -> Result<u8> {
194        self.field(slot)?.map_or(Ok(default), |p| self.fb.u8(p))
195    }
196    fn i16_or(&self, slot: usize, default: i16) -> Result<i16> {
197        self.field(slot)?
198            .map_or(Ok(default), |p| self.fb.u16(p).map(|v| v as i16))
199    }
200    fn i32_or(&self, slot: usize, default: i32) -> Result<i32> {
201        self.field(slot)?.map_or(Ok(default), |p| self.fb.i32(p))
202    }
203    fn i64_or(&self, slot: usize, default: i64) -> Result<i64> {
204        self.field(slot)?.map_or(Ok(default), |p| self.fb.i64(p))
205    }
206    fn table(&self, slot: usize) -> Result<Option<Table<'a>>> {
207        match self.field(slot)? {
208            None => Ok(None),
209            Some(p) => Ok(Some(self.fb.table(self.fb.follow(p)?)?)),
210        }
211    }
212    /// A vector's element start and count, each element `elem` bytes wide.
213    fn vector(&self, slot: usize, elem: usize) -> Result<Option<(usize, usize)>> {
214        let Some(p) = self.field(slot)? else {
215            return Ok(None);
216        };
217        let at = self.fb.follow(p)?;
218        let n = self.fb.u32(at)? as usize;
219        let start = at + 4;
220        let len = n.checked_mul(elem).context("flatbuffer vector overflows")?;
221        self.fb.bytes(start, len)?;
222        Ok(Some((start, n)))
223    }
224    fn string(&self, slot: usize) -> Result<Option<String>> {
225        let Some((start, n)) = self.vector(slot, 1)? else {
226            return Ok(None);
227        };
228        let b = self.fb.bytes(start, n)?;
229        Ok(Some(
230            std::str::from_utf8(b)
231                .context("a flatbuffer string is not UTF-8")?
232                .to_string(),
233        ))
234    }
235    /// The tables of a vector of tables.
236    fn tables(&self, slot: usize, max: usize) -> Result<Vec<Table<'a>>> {
237        let Some((start, n)) = self.vector(slot, 4)? else {
238            return Ok(Vec::new());
239        };
240        ensure!(n <= max, "{n} entries where at most {max} are read");
241        (0..n)
242            .map(|i| self.fb.table(self.fb.follow(start + 4 * i)?))
243            .collect()
244    }
245}
246
247// ---------------------------------------------------------------------------------------------
248// Schema.
249
250const MESSAGE_SCHEMA: u8 = 1;
251const MESSAGE_DICTIONARY_BATCH: u8 = 2;
252const MESSAGE_RECORD_BATCH: u8 = 3;
253
254fn type_name(code: u8) -> &'static str {
255    match code {
256        1 => "Null",
257        2 => "Int",
258        3 => "FloatingPoint",
259        4 => "Binary",
260        5 => "Utf8",
261        6 => "Bool",
262        7 => "Decimal",
263        8 => "Date",
264        9 => "Time",
265        10 => "Timestamp",
266        11 => "Interval",
267        12 => "List",
268        13 => "Struct",
269        14 => "Union",
270        15 => "FixedSizeBinary",
271        16 => "FixedSizeList",
272        17 => "Map",
273        18 => "Duration",
274        19 => "LargeBinary",
275        20 => "LargeUtf8",
276        21 => "LargeList",
277        22 => "RunEndEncoded",
278        23 => "BinaryView",
279        24 => "Utf8View",
280        _ => "an unknown type",
281    }
282}
283
284fn unit_scale(unit: i16) -> Result<u32> {
285    Ok(match unit {
286        0 => 0,
287        1 => 3,
288        2 => 6,
289        3 => 9,
290        other => bail!("time unit {other} is not one Arrow defines"),
291    })
292}
293
294/// One field of the schema, without its children.
295fn read_field<'a>(t: &Table<'a>) -> Result<(Field, Vec<Table<'a>>)> {
296    let name = t.string(0)?.unwrap_or_default();
297    let type_code = t.u8_or(2, 0)?;
298    ensure!(
299        t.field(4)?.is_none(),
300        "column {name:?} is dictionary-encoded, which this client does not read"
301    );
302    let ty = t.table(3)?;
303    let data_type = match type_code {
304        1 => ArrowType::Null,
305        2 => {
306            let ty = ty.context("an Int type with no parameters")?;
307            let bits = ty.i32_or(0, 0)?;
308            ensure!(
309                matches!(bits, 8 | 16 | 32 | 64),
310                "column {name:?}: Int of {bits} bits"
311            );
312            ArrowType::Int {
313                bits: bits as u8,
314                signed: ty.u8_or(1, 0)? != 0,
315            }
316        }
317        3 => {
318            let precision = ty.map_or(Ok(0), |ty| ty.i16_or(0, 0))?;
319            ArrowType::Float {
320                bits: match precision {
321                    1 => 32,
322                    2 => 64,
323                    _ => bail!("column {name:?}: half-precision floats are not read"),
324                },
325            }
326        }
327        4 => ArrowType::Binary,
328        5 => ArrowType::Utf8,
329        6 => ArrowType::Bool,
330        7 => {
331            let ty = ty.context("a Decimal type with no parameters")?;
332            let bits = ty.i32_or(2, 128)?;
333            ensure!(
334                bits == 128,
335                "column {name:?}: Decimal{bits} is not read (Decimal128 is)"
336            );
337            ArrowType::Decimal {
338                precision: ty.i32_or(0, 0)?,
339                scale: ty.i32_or(1, 0)?,
340            }
341        }
342        8 => ArrowType::Date {
343            day: ty.map_or(Ok(1), |ty| ty.i16_or(0, 1))? == 0,
344        },
345        9 => {
346            let (unit, bits) = match ty {
347                Some(ty) => (ty.i16_or(0, 1)?, ty.i32_or(1, 32)?),
348                None => (1, 32),
349            };
350            ensure!(
351                matches!(bits, 32 | 64),
352                "column {name:?}: Time of {bits} bits"
353            );
354            ArrowType::Time {
355                scale: unit_scale(unit)?,
356                bits: bits as u8,
357            }
358        }
359        10 => {
360            let ty = ty.context("a Timestamp type with no parameters")?;
361            ArrowType::Timestamp {
362                scale: unit_scale(ty.i16_or(0, 0)?)?,
363                tz: ty.string(1)?,
364            }
365        }
366        13 => ArrowType::Struct,
367        other => bail!(
368            "column {name:?} is an Arrow {}, which this client does not read",
369            type_name(other)
370        ),
371    };
372    let mut metadata = Vec::new();
373    for kv in t.tables(6, 64)? {
374        metadata.push((
375            kv.string(0)?.unwrap_or_default(),
376            kv.string(1)?.unwrap_or_default(),
377        ));
378    }
379    let children = t.tables(5, MAX_FIELDS)?;
380    if data_type != ArrowType::Struct {
381        ensure!(
382            children.is_empty(),
383            "column {name:?} is not a struct but has children"
384        );
385    }
386    Ok((
387        Field {
388            name,
389            data_type,
390            metadata,
391            children: Vec::new(),
392        },
393        children,
394    ))
395}
396
397fn read_schema(schema: &Table<'_>) -> Result<Vec<Field>> {
398    ensure!(
399        schema.i16_or(0, 0)? == 0,
400        "the result is big-endian Arrow, which this client does not read"
401    );
402    let mut fields = Vec::new();
403    let mut count = 0usize;
404    for t in schema.tables(1, MAX_FIELDS)? {
405        let (mut field, children) = read_field(&t)?;
406        count += 1;
407        // One level of nesting, read in this loop rather than by recursion.
408        for child in children {
409            let (child, grandchildren) = read_field(&child)?;
410            ensure!(
411                grandchildren.is_empty() && child.data_type != ArrowType::Struct,
412                "column {:?} nests deeper than {MAX_FIELD_NESTING} level",
413                field.name
414            );
415            count += 1;
416            field.children.push(child);
417        }
418        ensure!(
419            count <= MAX_FIELDS,
420            "the schema has more than {MAX_FIELDS} fields"
421        );
422        fields.push(field);
423    }
424    Ok(fields)
425}
426
427// ---------------------------------------------------------------------------------------------
428// Record batches.
429
430struct Node {
431    length: usize,
432    null_count: usize,
433}
434
435struct BatchReader<'a> {
436    body: &'a [u8],
437    nodes: Vec<Node>,
438    buffers: Vec<(usize, usize)>,
439    next_node: usize,
440    next_buffer: usize,
441}
442
443impl<'a> BatchReader<'a> {
444    fn node(&mut self) -> Result<&Node> {
445        let n = self
446            .nodes
447            .get(self.next_node)
448            .context("the batch has fewer field nodes than the schema has fields")?;
449        self.next_node += 1;
450        Ok(n)
451    }
452    fn buffer(&mut self) -> Result<&'a [u8]> {
453        let (offset, length) = *self
454            .buffers
455            .get(self.next_buffer)
456            .context("the batch has fewer buffers than its columns need")?;
457        self.next_buffer += 1;
458        let end = offset
459            .checked_add(length)
460            .context("a buffer's extent overflows")?;
461        self.body
462            .get(offset..end)
463            .context("a buffer runs past the message body")
464    }
465}
466
467fn valid(bitmap: &[u8], null_count: usize, i: usize) -> Result<bool> {
468    if bitmap.is_empty() {
469        ensure!(
470            null_count == 0,
471            "a column with nulls has no validity bitmap"
472        );
473        return Ok(true);
474    }
475    let byte = bitmap
476        .get(i / 8)
477        .context("the validity bitmap is shorter than the column")?;
478    Ok(byte & (1 << (i % 8)) != 0)
479}
480
481fn fixed(values: &[u8], width: usize, i: usize) -> Result<&[u8]> {
482    values
483        .get(i * width..(i + 1) * width)
484        .context("a value buffer is shorter than the column")
485}
486
487fn read_int(b: &[u8], signed: bool) -> i128 {
488    let mut a = [0u8; 16];
489    a[..b.len()].copy_from_slice(b);
490    if signed && b.last().is_some_and(|x| x & 0x80 != 0) {
491        for x in &mut a[b.len()..] {
492            *x = 0xFF;
493        }
494    }
495    i128::from_le_bytes(a)
496}
497
498/// Read `take` rows of one flat (non-struct) column.
499fn read_flat(
500    r: &mut BatchReader<'_>,
501    ty: &ArrowType,
502    len: usize,
503    take: usize,
504) -> Result<Vec<Cell>> {
505    let node = r.node()?;
506    let (nlen, nulls) = (node.length, node.null_count);
507    ensure!(
508        nlen == len,
509        "a column holds {nlen} values in a batch of {len} rows"
510    );
511    ensure!(nulls <= nlen, "a column claims more nulls than values");
512    if *ty == ArrowType::Null {
513        return Ok(vec![Cell::Null; take]);
514    }
515    let bitmap = r.buffer()?;
516    let mut out = Vec::with_capacity(take);
517    match ty {
518        ArrowType::Utf8 | ArrowType::Binary => {
519            let offsets = r.buffer()?;
520            let data = r.buffer()?;
521            ensure!(
522                offsets.len() >= (len + 1) * 4,
523                "an offsets buffer is shorter than the column"
524            );
525            let at = |i: usize| -> Result<usize> {
526                let v = i32::from_le_bytes(fixed(offsets, 4, i)?.try_into()?);
527                ensure!(v >= 0, "a negative string offset");
528                Ok(v as usize)
529            };
530            for i in 0..take {
531                if !valid(bitmap, nulls, i)? {
532                    out.push(Cell::Null);
533                    continue;
534                }
535                let (s, e) = (at(i)?, at(i + 1)?);
536                ensure!(s <= e, "string offsets run backwards");
537                let bytes = data
538                    .get(s..e)
539                    .context("a string runs past its data buffer")?;
540                out.push(if *ty == ArrowType::Utf8 {
541                    Cell::Text(
542                        std::str::from_utf8(bytes)
543                            .context("a Utf8 value is not UTF-8")?
544                            .to_string(),
545                    )
546                } else {
547                    Cell::Bytes(bytes.len())
548                });
549            }
550        }
551        ArrowType::Bool => {
552            let values = r.buffer()?;
553            for i in 0..take {
554                if !valid(bitmap, nulls, i)? {
555                    out.push(Cell::Null);
556                    continue;
557                }
558                out.push(Cell::Bool(valid(values, 0, i)?));
559            }
560        }
561        _ => {
562            let values = r.buffer()?;
563            let (width, signed) = match ty {
564                ArrowType::Int { bits, signed } => (*bits as usize / 8, *signed),
565                ArrowType::Float { bits } => (*bits as usize / 8, true),
566                ArrowType::Decimal { .. } => (16, true),
567                ArrowType::Date { day } => (if *day { 4 } else { 8 }, true),
568                ArrowType::Time { bits, .. } => (*bits as usize / 8, true),
569                ArrowType::Timestamp { .. } => (8, true),
570                _ => bail!("unexpected column type"),
571            };
572            for i in 0..take {
573                if !valid(bitmap, nulls, i)? {
574                    out.push(Cell::Null);
575                    continue;
576                }
577                let b = fixed(values, width, i)?;
578                out.push(match ty {
579                    ArrowType::Float { bits: 32 } => {
580                        Cell::Float(f32::from_le_bytes(b.try_into()?) as f64)
581                    }
582                    ArrowType::Float { .. } => Cell::Float(f64::from_le_bytes(b.try_into()?)),
583                    _ => Cell::Int(read_int(b, signed)),
584                });
585            }
586        }
587    }
588    Ok(out)
589}
590
591fn read_batch(
592    batch: &Table<'_>,
593    body: &[u8],
594    fields: &[Field],
595    take: usize,
596) -> Result<(usize, Vec<Vec<Cell>>)> {
597    ensure!(
598        batch.field(3)?.is_none(),
599        "the result is compressed Arrow, which this client does not read"
600    );
601    let length = batch.i64_or(0, 0)?;
602    ensure!(length >= 0, "a batch with a negative length");
603    let length = length as usize;
604    let fb = batch.fb;
605    let (nstart, ncount) = batch.vector(1, 16)?.unwrap_or((0, 0));
606    let (bstart, bcount) = batch.vector(2, 16)?.unwrap_or((0, 0));
607    ensure!(
608        ncount <= MAX_FIELDS && bcount <= 3 * MAX_FIELDS,
609        "a batch with {ncount} nodes and {bcount} buffers"
610    );
611    let mut nodes = Vec::with_capacity(ncount);
612    for i in 0..ncount {
613        let (l, n) = (fb.i64(nstart + 16 * i)?, fb.i64(nstart + 16 * i + 8)?);
614        ensure!(l >= 0 && n >= 0, "a field node with a negative count");
615        nodes.push(Node {
616            length: l as usize,
617            null_count: n as usize,
618        });
619    }
620    let mut buffers = Vec::with_capacity(bcount);
621    for i in 0..bcount {
622        let (o, l) = (fb.i64(bstart + 16 * i)?, fb.i64(bstart + 16 * i + 8)?);
623        ensure!(o >= 0 && l >= 0, "a buffer with a negative extent");
624        buffers.push((o as usize, l as usize));
625    }
626    let mut r = BatchReader {
627        body,
628        nodes,
629        buffers,
630        next_node: 0,
631        next_buffer: 0,
632    };
633    let take = take.min(length);
634    let mut columns: Vec<Vec<Cell>> = Vec::with_capacity(fields.len());
635    for field in fields {
636        if field.data_type == ArrowType::Struct {
637            let node = r.node()?;
638            let (nlen, nulls) = (node.length, node.null_count);
639            ensure!(
640                nlen == length,
641                "a struct column holds {nlen} values in a batch of {length} rows"
642            );
643            let bitmap = r.buffer()?;
644            let mut children = Vec::with_capacity(field.children.len());
645            for child in &field.children {
646                children.push(read_flat(&mut r, &child.data_type, length, take)?);
647            }
648            let mut col = Vec::with_capacity(take);
649            for i in 0..take {
650                col.push(if valid(bitmap, nulls, i)? {
651                    Cell::Struct(children.iter().map(|c| c[i].clone()).collect())
652                } else {
653                    Cell::Null
654                });
655            }
656            columns.push(col);
657        } else {
658            columns.push(read_flat(&mut r, &field.data_type, length, take)?);
659        }
660    }
661    let rows = (0..take)
662        .map(|i| columns.iter().map(|c| c[i].clone()).collect())
663        .collect();
664    Ok((length, rows))
665}
666
667/// Decode an Arrow IPC stream, keeping at most `max_rows` rows.
668pub fn decode_stream(stream: &[u8], max_rows: usize) -> Result<Decoded> {
669    let mut pos = 0usize;
670    let mut fields: Option<Vec<Field>> = None;
671    let mut rows = Vec::new();
672    let mut total_rows = 0u64;
673    for message in 0.. {
674        ensure!(
675            message < MAX_MESSAGES,
676            "the stream has more than {MAX_MESSAGES} messages"
677        );
678        if stream.len() < pos + 4 {
679            break;
680        }
681        let word = |at: usize| -> Result<u32> {
682            let b = stream
683                .get(at..at + 4)
684                .context("the stream ends inside a message header")?;
685            Ok(u32::from_le_bytes([b[0], b[1], b[2], b[3]]))
686        };
687        let mut meta_len = word(pos)?;
688        pos += 4;
689        if meta_len == 0xFFFF_FFFF {
690            meta_len = word(pos)?;
691            pos += 4;
692        }
693        if meta_len == 0 {
694            break;
695        }
696        let meta_len = meta_len as usize;
697        ensure!(
698            meta_len <= stream.len() - pos.min(stream.len()),
699            "a message's metadata runs past the stream"
700        );
701        let meta = Fb {
702            buf: &stream[pos..pos + meta_len],
703        };
704        pos += meta_len;
705        let msg = meta.table(meta.follow(0)?)?;
706        let header_type = msg.u8_or(1, 0)?;
707        let body_len = msg.i64_or(3, 0)?;
708        ensure!(
709            body_len >= 0 && body_len as usize <= stream.len() - pos,
710            "a message body runs past the stream"
711        );
712        let body = &stream[pos..pos + body_len as usize];
713        pos += body_len as usize;
714        let header = msg.table(2)?.context("a message with no header")?;
715        match header_type {
716            MESSAGE_SCHEMA => {
717                ensure!(fields.is_none(), "the stream carries a second schema");
718                fields = Some(read_schema(&header)?);
719            }
720            MESSAGE_RECORD_BATCH => {
721                let schema = fields
722                    .as_ref()
723                    .context("a record batch before the schema")?;
724                let take = max_rows.saturating_sub(rows.len());
725                let (length, batch_rows) = read_batch(&header, body, schema, take)?;
726                total_rows += length as u64;
727                rows.extend(batch_rows);
728            }
729            MESSAGE_DICTIONARY_BATCH => {
730                bail!("the result carries an Arrow dictionary, which this client does not read")
731            }
732            other => bail!("an Arrow message of type {other}, which a result never carries"),
733        }
734    }
735    Ok(Decoded {
736        fields: fields.context("the Arrow stream has no schema")?,
737        rows,
738        total_rows,
739    })
740}
netget/client/snowflake/arrow.rs

netget/client/snowflake/
arrow.rs
