Skip to content

Commit

Permalink
feat: add arrow_schema_utility
Browse files Browse the repository at this point in the history
This patch adds arrow_schema_utility::get_posql_compatible_schema() for
converting arrow schemas into posql compatible arrow schemas.
  • Loading branch information
henrymai committed Oct 23, 2024
1 parent b55765d commit aa8dac5
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 0 deletions.
37 changes: 37 additions & 0 deletions crates/proof-of-sql/src/base/database/arrow_schema_utility.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
//! This module provides utility functions for working with Arrow schemas in the context of Proof of SQL.
//! It includes functionality to convert Arrow schemas to PoSQL-compatible formats.

use alloc::sync::Arc;
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};

/// Converts an Arrow schema to a PoSQL-compatible schema.
///
/// This function takes an Arrow `SchemaRef` and returns a new `SchemaRef` where
/// floating-point data types (Float16, Float32, Float64) are converted to Decimal256(75, 30).
/// Other data types remain unchanged.
///
/// # Arguments
///
/// * `schema` - The input Arrow schema to convert.
///
/// # Returns
///
/// A new `SchemaRef` with PoSQL-compatible data types.
#[must_use]
pub fn get_posql_compatible_schema(schema: &SchemaRef) -> SchemaRef {
let new_fields: Vec<Field> = schema
.fields()
.iter()
.map(|field| {
let new_data_type = match field.data_type() {
DataType::Float16 | DataType::Float32 | DataType::Float64 => {
DataType::Decimal256(75, 30)
}
_ => field.data_type().clone(),
};
Field::new(field.name(), new_data_type, field.is_nullable())
})
.collect();

Arc::new(Schema::new(new_fields))
}
3 changes: 3 additions & 0 deletions crates/proof-of-sql/src/base/database/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,9 @@ mod record_batch_utility;
#[cfg(feature = "arrow")]
pub use record_batch_utility::ToArrow;

#[cfg(feature = "arrow")]
pub mod arrow_schema_utility;

#[cfg(all(test, feature = "arrow", feature = "test"))]
mod test_accessor_utility;
#[cfg(all(test, feature = "arrow", feature = "test"))]
Expand Down

0 comments on commit aa8dac5

Please sign in to comment.