parquet/tests/custom_writer.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use std::fs::File;
use std::{
    fs,
    io::{prelude::*, SeekFrom},
    sync::Arc,
};

use parquet::file::writer::TryClone;
use parquet::{
    basic::Repetition, basic::Type, file::properties::WriterProperties,
    file::writer::SerializedFileWriter, schema::types,
};
use std::env;

// Test creating some sort of custom writer to ensure the
// appropriate traits are exposed
struct CustomWriter {
    file: File,
}

impl Write for CustomWriter {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.file.write(buf)
    }
    fn flush(&mut self) -> std::io::Result<()> {
        self.file.flush()
    }
}

impl Seek for CustomWriter {
    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
        self.file.seek(pos)
    }
}

impl TryClone for CustomWriter {
    fn try_clone(&self) -> std::io::Result<Self> {
        use std::io::{Error, ErrorKind};
        Err(Error::new(ErrorKind::Other, "Clone not supported"))
    }
}

#[test]
fn test_custom_writer() {
    let schema = Arc::new(
        types::Type::group_type_builder("schema")
            .with_fields(&mut vec![Arc::new(
                types::Type::primitive_type_builder("col1", Type::INT32)
                    .with_repetition(Repetition::REQUIRED)
                    .build()
                    .unwrap(),
            )])
            .build()
            .unwrap(),
    );
    let props = Arc::new(WriterProperties::builder().build());

    let file = get_temp_file("test_custom_file_writer");
    let test_file = file.try_clone().unwrap();

    let writer = CustomWriter { file };

    // test is that this file can be created
    let file_writer = SerializedFileWriter::new(writer, schema, props).unwrap();
    std::mem::drop(file_writer);

    // ensure the file now exists and has non zero size
    let metadata = test_file.metadata().unwrap();
    assert!(metadata.len() > 0);
}

/// Returns file handle for a temp file in 'target' directory with a provided content
fn get_temp_file(file_name: &str) -> fs::File {
    // build tmp path to a file in "target/debug/testdata"
    let mut path_buf = env::current_dir().unwrap();
    path_buf.push("target");
    path_buf.push("debug");
    path_buf.push("testdata");
    fs::create_dir_all(&path_buf).unwrap();
    path_buf.push(file_name);

    File::create(path_buf).unwrap()
}