tbl-core 0.1.1

utilities for reading and modifying tabular files
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
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

use arrow::datatypes::Schema;
use arrow::record_batch::RecordBatch;
use futures::stream::StreamExt;
use parquet::arrow::arrow_writer::ArrowWriter;
use parquet::arrow::async_reader::ParquetRecordBatchStreamBuilder;
use parquet::file::properties::WriterProperties;
use std::io::BufWriter;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::fs::File;
use tokio::io::AsyncWriteExt;

/// create temporary path target similar to the final target path
pub fn create_tmp_target(path: &std::path::Path) -> PathBuf {
    let mut new_path = path.to_path_buf();
    let suffix = "_tmp";
    if let Some(stem) = path.file_stem() {
        let mut new_stem = stem.to_string_lossy().into_owned();
        new_stem.push_str(suffix);
        if let Some(extension) = path.extension() {
            new_stem.push('.');
            new_stem.push_str(&extension.to_string_lossy());
        }
        new_path.set_file_name(new_stem);
    }

    new_path
}

/// drop columns from parquet column
pub async fn drop_parquet_columns(
    input_path: PathBuf,
    output_path: PathBuf,
    columns_to_drop: Vec<String>,
    batch_size: usize,
) -> Result<(), crate::TblError> {
    let input_file = File::open(input_path).await?;
    let tmp_output_path = create_tmp_target(output_path.as_path());
    let mut output_file = File::create(&tmp_output_path).await?;
    let builder = ParquetRecordBatchStreamBuilder::new(input_file)
        .await?
        .with_batch_size(batch_size);
    let mut reader_stream = builder.build()?;
    let original_schema = reader_stream.schema().clone();

    // Create new schema without dropped columns
    let new_schema = Arc::new(Schema::new(
        original_schema
            .fields()
            .iter()
            .filter_map(|field| {
                if !columns_to_drop.contains(field.name()) {
                    Some(field.clone())
                } else {
                    None
                }
            })
            .collect::<Vec<_>>(),
    ));

    let writer_props = WriterProperties::builder().build();
    let mut buffer = Vec::new();
    let mut arrow_writer = ArrowWriter::try_new(
        BufWriter::new(&mut buffer),
        new_schema.clone(),
        Some(writer_props),
    )?;

    while let Some(batch) = reader_stream.next().await {
        let batch = batch?;
        let new_columns = batch
            .columns()
            .iter()
            .enumerate()
            .filter_map(|(i, col)| {
                if !columns_to_drop.contains(original_schema.field(i).name()) {
                    Some(col.clone())
                } else {
                    None
                }
            })
            .collect::<Vec<_>>();

        let new_batch = RecordBatch::try_new(new_schema.clone(), new_columns)?;
        arrow_writer.write(&new_batch)?;
    }

    arrow_writer.close()?;
    output_file.write_all(&buffer).await?;
    output_file.flush().await?;

    std::fs::rename(tmp_output_path, output_path)?;

    Ok(())
}

/// drop columns from multiple parquet files
pub async fn drop_parquets_columns(
    input_output_paths: Vec<(PathBuf, PathBuf)>,
    columns_to_drop: Vec<String>,
    batch_size: usize,
    max_concurrent: usize,
) -> Result<(), crate::TblError> {
    let semaphore = Arc::new(tokio::sync::Semaphore::new(max_concurrent));

    let results = futures::stream::iter(input_output_paths)
        .map(|(input, output)| {
            let columns_to_drop = columns_to_drop.clone();
            let sem = Arc::clone(&semaphore);
            async move {
                let _permit = sem.acquire().await?;
                drop_parquet_columns(input, output, columns_to_drop, batch_size).await
            }
        })
        .buffer_unordered(max_concurrent)
        .collect::<Vec<_>>()
        .await;

    // Check if any operations failed
    for result in results {
        result?;
    }

    Ok(())
}