Function partitioned_file_groups

datafusion/core/src/test/mod.rs:116–202

Returns file groups [`Vec `] for scanning `partitions` of `filename`

(
    path: &str,
    filename: &str,
    partitions: usize,
    file_format: &Arc<dyn FileFormat>,
    file_compression_type: FileCompressionType,
    work_dir: &Path,
)

/// Returns file groups [`Vec<FileGroup>`] for scanning `partitions` of `filename`
pub fn partitioned_file_groups(
    path: &str,
    filename: &str,
    partitions: usize,
    file_format: &Arc<dyn FileFormat>,
    file_compression_type: FileCompressionType,
    work_dir: &Path,
) -> Result<Vec<FileGroup>> {
    let path = format!("{path}/{filename}");

    let mut writers = vec![];
    let mut files = vec![];
    for i in 0..partitions {
        let filename = format!(
            "partition-{}{}",
            i,
            file_format
                .get_ext_with_compression(&file_compression_type)
                .unwrap()
        );
        let filename = work_dir.join(filename);

        let file = File::create(&filename).unwrap();

        let encoder: Box<dyn Write + Send> = match file_compression_type.to_owned() {
            FileCompressionType::UNCOMPRESSED => Box::new(file),
            #[cfg(feature = "compression")]
            FileCompressionType::GZIP => {
                Box::new(GzEncoder::new(file, GzCompression::default()))
            }
            #[cfg(feature = "compression")]
            FileCompressionType::XZ => Box::new(XzEncoder::new(file, 9)),
            #[cfg(feature = "compression")]
            FileCompressionType::ZSTD => {
                let encoder = ZstdEncoder::new(file, 0)
                    .map_err(|e| DataFusionError::External(Box::new(e)))?
                    .auto_finish();
                Box::new(encoder)
            }
            #[cfg(feature = "compression")]
            FileCompressionType::BZIP2 => {
                Box::new(BzEncoder::new(file, BzCompression::default()))
            }
            #[cfg(not(feature = "compression"))]
            FileCompressionType::GZIP
            | FileCompressionType::BZIP2
            | FileCompressionType::XZ
            | FileCompressionType::ZSTD => {
                panic!("Compression is not supported in this build")
            }
        };

        let writer = BufWriter::new(encoder);
        writers.push(writer);
        files.push(filename);
    }

    let f = File::open(path)?;