bask 0.1.1

Build Tasks
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

/*
 * SPDX-FileCopyrightText: 2026 Wavelens GmbH <info@wavelens.io>
 *
 * SPDX-License-Identifier: MIT OR Apache-2.0
 */
#![cfg(feature = "formats")]

use std::sync::Arc;

use arrow::array::{ArrayRef, Int64Array};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use arrow::record_batch::RecordBatch;
use bask::formats::{ArrowFormat, Format, ParquetFormat};

fn schema() -> SchemaRef {
    Arc::new(Schema::new(vec![Field::new("n", DataType::Int64, false)]))
}

fn batch(start: i64, len: i64) -> RecordBatch {
    let column = Int64Array::from((start..start + len).collect::<Vec<_>>());
    RecordBatch::try_new(schema(), vec![Arc::new(column) as ArrayRef]).unwrap()
}

fn roundtrip_rechunks(format: impl Format, ext: &str) {
    let path = std::env::temp_dir().join(format!("bask_roundtrip.{ext}"));

    let mut writer = format.open_writer(&path, schema()).unwrap();
    writer.write(&batch(0, 2500)).unwrap();
    writer.write(&batch(2500, 2500)).unwrap();
    writer.finish().unwrap();

    let mut reader = format.open_reader(&path, 1024).unwrap();
    let (mut rows, mut sum, mut chunks) = (0usize, 0i64, 0usize);
    while let Some(chunk) = reader.next_chunk().unwrap() {
        chunks += 1;
        assert!(
            chunk.num_rows() <= 1024,
            "chunk of {} exceeds chunk_rows",
            chunk.num_rows()
        );
        let column = chunk
            .column(0)
            .as_any()
            .downcast_ref::<Int64Array>()
            .unwrap();
        rows += chunk.num_rows();
        sum += column.iter().flatten().sum::<i64>();
    }
    std::fs::remove_file(&path).ok();

    assert_eq!(rows, 5000);
    assert_eq!(sum, (0..5000).sum::<i64>());
    assert!(
        chunks >= 5,
        "expected 5000 rows to re-chunk into >= 5 chunks, got {chunks}"
    );
}

#[test]
fn parquet_roundtrip_chunks() {
    roundtrip_rechunks(ParquetFormat, "parquet");
}

#[test]
fn arrow_roundtrip_chunks() {
    roundtrip_rechunks(ArrowFormat, "arrow");
}