crates/paimon/src/arrow/format/mod.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.

mod avro;
pub(crate) mod blob;
#[cfg(feature = "mosaic")]
mod mosaic;
mod orc;
mod parquet;
mod row;
#[cfg(feature = "vortex")]
mod vortex;

use crate::io::{FileRead, OutputFile};
use crate::spec::{DataField, Predicate};
use crate::table::{ArrowRecordBatchStream, RowRange};
use crate::Error;
use arrow_array::RecordBatch;
use arrow_schema::SchemaRef;
use async_trait::async_trait;

/// Predicates with the file-level field context needed for pushdown.
/// Only used by formats that support predicate pushdown (e.g. Parquet).
pub(crate) struct FilePredicates {
    /// Predicates with indices already remapped to file-level fields.
    pub predicates: Vec<Predicate>,
    /// File-level fields (full file schema), used for stats access and row filtering.
    pub file_fields: Vec<DataField>,
}

/// Format-agnostic file reader that produces Arrow RecordBatch streams.
///
/// Each implementation (Parquet, ORC, ...) handles:
/// - Column projection
/// - Predicate pushdown where supported (row-group/stripe pruning and, for
///   some formats, row-level filtering)
/// - Row range selection
#[async_trait]
pub(crate) trait FormatFileReader: Send + Sync {
    /// Read a single data file, returning a stream of RecordBatches
    /// containing only the projected columns (using names from the file's schema).
    ///
    /// `row_selection` is a pre-merged list of 0-based inclusive row ranges
    /// (DV + row_ranges already combined by the caller).
    async fn read_batch_stream(
        &self,
        reader: Box<dyn FileRead>,
        file_size: u64,
        read_fields: &[DataField],
        predicates: Option<&FilePredicates>,
        batch_size: Option<usize>,
        row_selection: Option<Vec<RowRange>>,
    ) -> crate::Result<ArrowRecordBatchStream>;
}

/// Format-agnostic file writer that streams Arrow RecordBatches directly to storage.
///
/// Each implementation (Parquet, ORC, ...) handles format-specific encoding.
/// Usage: create via [`create_format_writer`], call [`write`](FormatFileWriter::write)
/// for each batch, then [`close`](FormatFileWriter::close) to finalize the file.
#[async_trait]
pub(crate) trait FormatFileWriter: Send {
    /// Write a RecordBatch to the underlying storage.
    async fn write(&mut self, batch: &RecordBatch) -> crate::Result<()>;

    /// Number of bytes written so far (approximate, before close).
    fn num_bytes(&self) -> usize;

    /// Number of bytes buffered in the current row group (not yet flushed).
    fn in_progress_size(&self) -> usize;

    /// Flush the current row group to storage without closing the file.
    async fn flush(&mut self) -> crate::Result<()>;

    /// Flush and close the writer, finalizing the file on storage.
    /// Returns the total number of bytes written.
    async fn close(self: Box<Self>) -> crate::Result<u64>;
}

/// Create a format reader based on the file extension.
pub(crate) fn create_format_reader(
    path: &str,
    blob_as_descriptor: bool,
) -> crate::Result<Box<dyn FormatFileReader>> {
    let lower = path.to_ascii_lowercase();
    if lower.ends_with(".parquet") {
        Ok(Box::new(parquet::ParquetFormatReader))
    } else if lower.ends_with(".blob") {
        Ok(Box::new(blob::BlobFormatReader::new(
            path.to_string(),
            blob_as_descriptor,
        )))
    } else if lower.ends_with(".orc") {
        Ok(Box::new(orc::OrcFormatReader))
    } else if lower.ends_with(".avro") {
        Ok(Box::new(avro::AvroFormatReader))
    } else if lower.ends_with(".row") {
        Ok(Box::new(row::RowFormatReader))
    } else {
        #[cfg(feature = "mosaic")]
        if lower.ends_with(".mosaic") {
            return Ok(Box::new(mosaic::MosaicFormatReader));
        }
        #[cfg(feature = "vortex")]
        if lower.ends_with(".vortex") {
            return Ok(Box::new(vortex::VortexFormatReader));
        }
        Err(Error::Unsupported {
            message: format!(
                "unsupported file format: expected {}, got: {path}",
                supported_read_formats().join(", ")
            ),
        })
    }
}

fn supported_read_formats() -> Vec<&'static str> {
    vec![
        ".parquet",
        ".blob",
        ".orc",
        ".avro",
        ".row",
        #[cfg(feature = "mosaic")]
        ".mosaic",
        #[cfg(feature = "vortex")]
        ".vortex",
    ]
}

/// Create a format writer that streams directly to storage.
pub(crate) async fn create_format_writer(
    output: &OutputFile,
    schema: SchemaRef,
    compression: &str,
    zstd_level: i32,
    file_io: Option<crate::io::FileIO>,
    write_fields: Option<&[DataField]>,
) -> crate::Result<Box<dyn FormatFileWriter>> {
    let path = output.location();
    let lower = path.to_ascii_lowercase();
    if lower.ends_with(".parquet") {
        Ok(Box::new(
            parquet::ParquetFormatWriter::new(output, schema, compression, zstd_level).await?,
        ))
    } else if lower.ends_with(".blob") {
        Ok(Box::new(
            blob::BlobFormatWriter::new(output, file_io).await?,
        ))
    } else if lower.ends_with(".row") {
        let row_type = match write_fields {
            Some(fields) => fields.to_vec(),
            None => row::row_type_from_arrow_schema(&schema)?,
        };
        Ok(Box::new(
            row::RowFormatWriter::new(output, schema, row_type, zstd_level).await?,
        ))
    } else {
        #[cfg(feature = "vortex")]
        if lower.ends_with(".vortex") {
            return Ok(Box::new(
                vortex::VortexFormatWriter::new(output, schema).await?,
            ));
        }
        Err(Error::Unsupported {
            message: format!("unsupported write format: expected .parquet, .row, got: {path}"),
        })
    }
}