infino 0.5.6

A fast retrieval engine that stores data on object storage and runs SQL, full-text search, and vector search over it from a single system — search-on-Parquet.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Infino Authors

//! Classify ingested [`RecordBatch`] bytes into text/scalar vs vector columns.
//!
//! Used by the platform worker when attributing write meters for Stripe.
//! Vector columns are Arrow [`DataType::FixedSizeList`] (the engine's vector
//! column shape); every other column counts as text/scalar.

use arrow_array::{Array, RecordBatch};
use arrow_schema::DataType;

/// Split `batch`'s visible column footprint into `(text_bytes, vector_bytes)`.
///
/// Sizes use Arrow's slice-aware array-data memory size so sliced arrays are
/// not billed for unused backing capacity (falls back to full array memory
/// size if slice sizing fails). This attributes the ingested payload, not the
/// on-storage superfile footprint after encode.
pub fn classify_batch_bytes(batch: &RecordBatch) -> (u64, u64) {
    let mut text_bytes = 0u64;
    let mut vector_bytes = 0u64;
    for column in batch.columns() {
        let size = visible_array_bytes(column.as_ref());
        if matches!(column.data_type(), DataType::FixedSizeList(_, _)) {
            vector_bytes = vector_bytes.saturating_add(size);
        } else {
            text_bytes = text_bytes.saturating_add(size);
        }
    }
    (text_bytes, vector_bytes)
}

fn visible_array_bytes(column: &dyn Array) -> u64 {
    column
        .to_data()
        .get_slice_memory_size()
        .unwrap_or_else(|_| column.get_array_memory_size()) as u64
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use arrow_array::{FixedSizeListArray, Float32Array, LargeStringArray, RecordBatch};
    use arrow_schema::{DataType, Field, Schema};

    use super::classify_batch_bytes;

    #[test]
    fn splits_text_and_vector_columns() {
        let item = Arc::new(Field::new("item", DataType::Float32, true));
        let schema = Arc::new(Schema::new(vec![
            Field::new("body", DataType::LargeUtf8, false),
            Field::new("emb", DataType::FixedSizeList(Arc::clone(&item), 2), false),
        ]));
        let body = LargeStringArray::from(vec!["hello", "world"]);
        let flat = Float32Array::from(vec![1.0_f32, 2.0, 3.0, 4.0]);
        let emb = FixedSizeListArray::new(item, 2, Arc::new(flat), None);
        let batch =
            RecordBatch::try_new(schema, vec![Arc::new(body), Arc::new(emb)]).expect("batch");
        let (text, vector) = classify_batch_bytes(&batch);
        assert!(text > 0, "text columns must contribute bytes");
        assert!(vector > 0, "vector columns must contribute bytes");
    }

    #[test]
    fn sliced_column_does_not_bill_full_backing_capacity() {
        let full = LargeStringArray::from(vec!["aaaa", "bbbb", "cccc", "dddd"]);
        let schema = Arc::new(Schema::new(vec![Field::new(
            "body",
            DataType::LargeUtf8,
            false,
        )]));
        let full_batch =
            RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(full.clone())]).expect("full");
        let sliced_batch =
            RecordBatch::try_new(schema, vec![Arc::new(full.slice(1, 1))]).expect("sliced");
        let (full_text, _) = classify_batch_bytes(&full_batch);
        let (sliced_text, _) = classify_batch_bytes(&sliced_batch);
        assert!(
            sliced_text < full_text,
            "slice should bill less than full array ({sliced_text} vs {full_text})"
        );
    }
}