use std::fs::File;
use arrow_array::{
Array, BinaryArray, BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array,
LargeBinaryArray, StringArray, TimestampMicrosecondArray,
};
use arrow_schema::DataType as ArrowDataType;
use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
use crate::catalog::TableMetadata;
use crate::executor::{ExecutorError, Result};
use crate::planner::types::ResolvedType;
use crate::storage::SqlValue;
use super::{BulkReader, CopyField, CopySchema};
pub struct ParquetReader {
schema: CopySchema,
target_types: Vec<ResolvedType>,
reader: ParquetRecordBatchReader,
buffer: Option<Vec<Vec<SqlValue>>>,
}
impl ParquetReader {
pub fn open(path: &str, table_meta: &TableMetadata, _header: bool) -> Result<Self> {
let file = File::open(path)
.map_err(|e| ExecutorError::BulkLoad(format!("failed to open parquet: {e}")))?;
let builder = ParquetRecordBatchReaderBuilder::try_new(file).map_err(|e| {
ExecutorError::BulkLoad(format!("failed to read parquet metadata: {e}"))
})?;
let arrow_schema = builder.schema();
let mut fields = Vec::with_capacity(arrow_schema.fields().len());
for f in arrow_schema.fields() {
let ty = map_arrow_type(f.data_type())?;
fields.push(CopyField {
name: Some(f.name().clone()),
data_type: Some(ty),
});
}
let reader = builder
.with_batch_size(1024)
.build()
.map_err(|e| ExecutorError::BulkLoad(format!("failed to build parquet reader: {e}")))?;
let target_types: Vec<ResolvedType> = table_meta
.columns
.iter()
.map(|c| c.data_type.clone())
.collect();
Ok(Self {
schema: CopySchema { fields },
target_types,
reader,
buffer: None,
})
}
}
impl BulkReader for ParquetReader {
fn schema(&self) -> &CopySchema {
&self.schema
}
fn next_batch(&mut self, max_rows: usize) -> Result<Option<Vec<Vec<SqlValue>>>> {
let max_rows = max_rows.max(1);
if let Some(mut buffered) = self.buffer.take() {
if buffered.len() > max_rows {
let rest = buffered.split_off(max_rows);
self.buffer = Some(rest);
}
return Ok(Some(buffered));
}
let maybe_batch = self.reader.next();
let batch = match maybe_batch {
Some(b) => b.map_err(|e| {
ExecutorError::BulkLoad(format!("failed to read parquet batch: {e}"))
})?,
None => return Ok(None),
};
let mut rows: Vec<Vec<SqlValue>> = Vec::with_capacity(batch.num_rows());
for row_idx in 0..batch.num_rows() {
let mut row = Vec::with_capacity(self.schema.fields.len());
for col_idx in 0..self.schema.fields.len() {
let value = arrow_value_to_sql(
batch.column(col_idx).as_ref(),
batch.schema().field(col_idx).data_type(),
self.target_types
.get(col_idx)
.ok_or_else(|| ExecutorError::BulkLoad("missing target type".into()))?,
row_idx,
)?;
row.push(value);
}
rows.push(row);
}
if rows.len() > max_rows {
let rest = rows.split_off(max_rows);
self.buffer = Some(rest);
}
Ok(Some(rows))
}
}
fn map_arrow_type(dt: &ArrowDataType) -> Result<ResolvedType> {
match dt {
ArrowDataType::Int32 => Ok(ResolvedType::Integer),
ArrowDataType::Int64 => Ok(ResolvedType::BigInt),
ArrowDataType::Float32 => Ok(ResolvedType::Float),
ArrowDataType::Float64 => Ok(ResolvedType::Double),
ArrowDataType::Boolean => Ok(ResolvedType::Boolean),
ArrowDataType::Utf8 => Ok(ResolvedType::Text),
ArrowDataType::Binary | ArrowDataType::LargeBinary => Ok(ResolvedType::Blob),
ArrowDataType::Timestamp(arrow_schema::TimeUnit::Microsecond, _) => {
Ok(ResolvedType::Timestamp)
}
other => Err(ExecutorError::BulkLoad(format!(
"unsupported parquet/arrow type: {other:?}"
))),
}
}
fn arrow_value_to_sql(
array: &dyn Array,
dt: &ArrowDataType,
expected: &ResolvedType,
row_idx: usize,
) -> Result<SqlValue> {
if array.is_null(row_idx) {
return Ok(SqlValue::Null);
}
match (dt, expected) {
(ArrowDataType::Int32, ResolvedType::Integer) => {
let arr = array.as_any().downcast_ref::<Int32Array>().unwrap();
Ok(SqlValue::Integer(arr.value(row_idx)))
}
(ArrowDataType::Int32, ResolvedType::BigInt) => {
let arr = array.as_any().downcast_ref::<Int32Array>().unwrap();
Ok(SqlValue::BigInt(arr.value(row_idx) as i64))
}
(ArrowDataType::Int32, ResolvedType::Float) => {
let arr = array.as_any().downcast_ref::<Int32Array>().unwrap();
Ok(SqlValue::Float(arr.value(row_idx) as f32))
}
(ArrowDataType::Int32, ResolvedType::Double) => {
let arr = array.as_any().downcast_ref::<Int32Array>().unwrap();
Ok(SqlValue::Double(arr.value(row_idx) as f64))
}
(ArrowDataType::Int64, ResolvedType::BigInt) => {
let arr = array.as_any().downcast_ref::<Int64Array>().unwrap();
Ok(SqlValue::BigInt(arr.value(row_idx)))
}
(ArrowDataType::Int64, ResolvedType::Double) => {
let arr = array.as_any().downcast_ref::<Int64Array>().unwrap();
Ok(SqlValue::Double(arr.value(row_idx) as f64))
}
(ArrowDataType::Float32, ResolvedType::Float) => {
let arr = array.as_any().downcast_ref::<Float32Array>().unwrap();
Ok(SqlValue::Float(arr.value(row_idx)))
}
(ArrowDataType::Float32, ResolvedType::Double) => {
let arr = array.as_any().downcast_ref::<Float32Array>().unwrap();
Ok(SqlValue::Double(arr.value(row_idx) as f64))
}
(ArrowDataType::Float64, ResolvedType::Double) => {
let arr = array.as_any().downcast_ref::<Float64Array>().unwrap();
Ok(SqlValue::Double(arr.value(row_idx)))
}
(ArrowDataType::Boolean, ResolvedType::Boolean) => {
let arr = array.as_any().downcast_ref::<BooleanArray>().unwrap();
Ok(SqlValue::Boolean(arr.value(row_idx)))
}
(ArrowDataType::Utf8, ResolvedType::Text) => {
let arr = array.as_any().downcast_ref::<StringArray>().unwrap();
Ok(SqlValue::Text(arr.value(row_idx).to_string()))
}
(ArrowDataType::Binary, ResolvedType::Blob) => {
let arr = array.as_any().downcast_ref::<BinaryArray>().unwrap();
Ok(SqlValue::Blob(arr.value(row_idx).to_vec()))
}
(ArrowDataType::LargeBinary, ResolvedType::Blob) => {
let arr = array.as_any().downcast_ref::<LargeBinaryArray>().unwrap();
Ok(SqlValue::Blob(arr.value(row_idx).to_vec()))
}
(
ArrowDataType::Timestamp(arrow_schema::TimeUnit::Microsecond, _),
ResolvedType::Timestamp,
) => {
let arr = array
.as_any()
.downcast_ref::<TimestampMicrosecondArray>()
.unwrap();
Ok(SqlValue::Timestamp(arr.value(row_idx)))
}
_ => Err(ExecutorError::BulkLoad(format!(
"parquet field type {:?} does not match expected {:?}",
dt, expected
))),
}
}