lance_testing/
datagen.rs

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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

//! Data generation utilities for unit tests

use std::collections::HashSet;
use std::sync::Arc;
use std::{iter::repeat_with, ops::Range};

use arrow_array::types::ArrowPrimitiveType;
use arrow_array::{
    Float32Array, Int32Array, PrimitiveArray, RecordBatch, RecordBatchIterator, RecordBatchReader,
};
use arrow_schema::{DataType, Field, Schema as ArrowSchema};
use lance_arrow::{fixed_size_list_type, ArrowFloatType, FixedSizeListArrayExt};
use num_traits::{real::Real, FromPrimitive};
use rand::distributions::uniform::SampleUniform;
use rand::{
    distributions::Uniform, prelude::Distribution, rngs::StdRng, seq::SliceRandom, Rng, SeedableRng,
};

pub trait ArrayGenerator {
    fn generate(&mut self, length: usize) -> Arc<dyn arrow_array::Array>;
    fn data_type(&self) -> &DataType;
    fn name(&self) -> Option<&str>;
}

pub struct IncrementingInt32 {
    name: Option<String>,
    current: i32,
    step: i32,
}

impl Default for IncrementingInt32 {
    fn default() -> Self {
        Self {
            name: None,
            current: 0,
            step: 1,
        }
    }
}

impl IncrementingInt32 {
    pub fn new() -> Self {
        Default::default()
    }

    pub fn start(mut self, start: i32) -> Self {
        self.current = start;
        self
    }

    pub fn step(mut self, step: i32) -> Self {
        self.step = step;
        self
    }

    pub fn named(mut self, name: impl Into<String>) -> Self {
        self.name = Some(name.into());
        self
    }
}

impl ArrayGenerator for IncrementingInt32 {
    fn generate(&mut self, length: usize) -> Arc<dyn arrow_array::Array> {
        let mut values = Vec::with_capacity(length);
        for _ in 0..length {
            values.push(self.current);
            self.current += self.step;
        }
        Arc::new(Int32Array::from(values))
    }

    fn name(&self) -> Option<&str> {
        self.name.as_deref()
    }

    fn data_type(&self) -> &DataType {
        &DataType::Int32
    }
}

pub struct RandomVector {
    name: Option<String>,
    vec_width: i32,
    data_type: DataType,
}

impl Default for RandomVector {
    fn default() -> Self {
        Self {
            name: None,
            vec_width: 4,
            data_type: fixed_size_list_type(4, DataType::Float32),
        }
    }
}

impl RandomVector {
    pub fn new() -> Self {
        Default::default()
    }

    pub fn vec_width(mut self, vec_width: i32) -> Self {
        self.vec_width = vec_width;
        self.data_type = fixed_size_list_type(self.vec_width, DataType::Float32);
        self
    }

    pub fn named(mut self, name: String) -> Self {
        self.name = Some(name);
        self
    }
}

impl ArrayGenerator for RandomVector {
    fn generate(&mut self, length: usize) -> Arc<dyn arrow_array::Array> {
        let values = generate_random_array(length * (self.vec_width as usize));
        Arc::new(
            <arrow_array::FixedSizeListArray as FixedSizeListArrayExt>::try_new_from_values(
                values,
                self.vec_width,
            )
            .expect("Create fixed size list"),
        )
    }

    fn name(&self) -> Option<&str> {
        self.name.as_deref()
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }
}

#[derive(Default)]
pub struct BatchGenerator {
    generators: Vec<Box<dyn ArrayGenerator>>,
}

impl BatchGenerator {
    pub fn new() -> Self {
        Default::default()
    }

    pub fn col(mut self, gen: Box<dyn ArrayGenerator>) -> Self {
        self.generators.push(gen);
        self
    }

    fn gen_batch(&mut self, num_rows: u32) -> RecordBatch {
        let mut fields = Vec::with_capacity(self.generators.len());
        let mut arrays = Vec::with_capacity(self.generators.len());
        for (field_index, gen) in self.generators.iter_mut().enumerate() {
            let arr = gen.generate(num_rows as usize);
            let default_name = format!("field_{}", field_index);
            let name = gen.name().unwrap_or(&default_name);
            fields.push(Field::new(name, arr.data_type().clone(), true));
            arrays.push(arr);
        }
        let schema = Arc::new(ArrowSchema::new(fields));
        RecordBatch::try_new(schema, arrays).unwrap()
    }

    pub fn batch(&mut self, num_rows: i32) -> impl RecordBatchReader {
        let batch = self.gen_batch(num_rows as u32);
        let schema = batch.schema();
        RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema)
    }

    pub fn batches(&mut self, num_batches: u32, rows_per_batch: u32) -> impl RecordBatchReader {
        let batches = (0..num_batches)
            .map(|_| self.gen_batch(rows_per_batch))
            .collect::<Vec<_>>();
        let schema = batches[0].schema();
        RecordBatchIterator::new(batches.into_iter().map(Ok), schema)
    }
}

/// Returns a batch of data that has a column that can be used to create an ANN index
///
/// The indexable column will be named "indexable"
/// The batch will not be empty
/// There will only be one batch
///
/// There are no other assumptions it is safe to make about the returned reader
pub fn some_indexable_batch() -> impl RecordBatchReader {
    let x = Box::new(RandomVector::new().named("indexable".to_string()));
    BatchGenerator::new().col(x).batch(512)
}

/// Returns a non-empty batch of data
///
/// The batch will not be empty
/// There will only be one batch
///
/// There are no other assumptions it is safe to make about the returned reader
pub fn some_batch() -> impl RecordBatchReader {
    some_indexable_batch()
}

/// Create a random float32 array.
pub fn generate_random_array_with_seed<T: ArrowFloatType>(n: usize, seed: [u8; 32]) -> T::ArrayType
where
    T::Native: Real + FromPrimitive,
{
    let mut rng = StdRng::from_seed(seed);

    T::ArrayType::from(
        repeat_with(|| T::Native::from_f32(rng.gen::<f32>()).unwrap())
            .take(n)
            .collect::<Vec<_>>(),
    )
}

/// Create a random float32 array where each element is uniformly
/// distributed between [0..1]
pub fn generate_random_array(n: usize) -> Float32Array {
    let mut rng = rand::thread_rng();
    Float32Array::from_iter_values(repeat_with(|| rng.gen::<f32>()).take(n))
}

/// Create a random primitive array where each element is uniformly distributed a
/// given range.
pub fn generate_random_array_with_range<T: ArrowPrimitiveType>(
    n: usize,
    range: Range<T::Native>,
) -> PrimitiveArray<T>
where
    T::Native: SampleUniform,
{
    let mut rng = StdRng::from_seed([13; 32]);
    let distribution = Uniform::new(range.start, range.end);
    PrimitiveArray::<T>::from_iter_values(repeat_with(|| distribution.sample(&mut rng)).take(n))
}

/// Create a random float32 array where each element is uniformly
/// distributed across the given range
pub fn generate_scaled_random_array(n: usize, min: f32, max: f32) -> Float32Array {
    let mut rng = rand::thread_rng();
    let distribution = Uniform::new(min, max);
    Float32Array::from_iter_values(repeat_with(|| distribution.sample(&mut rng)).take(n))
}

pub fn sample_indices(range: Range<usize>, num_picks: u32) -> Vec<usize> {
    let mut rng = rand::thread_rng();
    let dist = Uniform::new(range.start, range.end);
    let ratio = num_picks as f32 / range.len() as f32;
    if ratio < 0.1_f32 && num_picks > 1000 {
        // We want to pick a large number of values from a big range.  Better to
        // use a set and potential retries
        let mut picked = HashSet::<usize>::with_capacity(num_picks as usize);
        let mut ordered_picked = Vec::with_capacity(num_picks as usize);
        while picked.len() < num_picks as usize {
            let val = dist.sample(&mut rng);
            if picked.insert(val) {
                ordered_picked.push(val);
            }
        }
        ordered_picked
    } else {
        // We want to pick most of the range, or a small number of values.  Go ahead
        // and just materialize the range and shuffle
        let mut values = Vec::from_iter(range);
        values.partial_shuffle(&mut rng, num_picks as usize);
        values.truncate(num_picks as usize);
        values
    }
}

pub fn sample_without_replacement<T: Copy>(choices: &[T], num_picks: u32) -> Vec<T> {
    let mut rng = rand::thread_rng();
    let mut shuffled = Vec::from(choices);
    shuffled.partial_shuffle(&mut rng, num_picks as usize);
    shuffled.truncate(num_picks as usize);
    shuffled
}