use std::{collections::HashMap, fmt, sync::Arc};
use crate::file::metadata::HeapSize;
use crate::format::SchemaElement;
use crate::basic::{
ColumnOrder, ConvertedType, LogicalType, Repetition, SortOrder, TimeUnit, Type as PhysicalType,
};
use crate::errors::{ParquetError, Result};
pub type TypePtr = Arc<Type>;
pub type SchemaDescPtr = Arc<SchemaDescriptor>;
pub type ColumnDescPtr = Arc<ColumnDescriptor>;
#[derive(Clone, Debug, PartialEq)]
pub enum Type {
PrimitiveType {
basic_info: BasicTypeInfo,
physical_type: PhysicalType,
type_length: i32,
scale: i32,
precision: i32,
},
GroupType {
basic_info: BasicTypeInfo,
fields: Vec<TypePtr>,
},
}
impl HeapSize for Type {
fn heap_size(&self) -> usize {
match self {
Type::PrimitiveType { basic_info, .. } => basic_info.heap_size(),
Type::GroupType { basic_info, fields } => basic_info.heap_size() + fields.heap_size(),
}
}
}
impl Type {
pub fn primitive_type_builder(name: &str, physical_type: PhysicalType) -> PrimitiveTypeBuilder {
PrimitiveTypeBuilder::new(name, physical_type)
}
pub fn group_type_builder(name: &str) -> GroupTypeBuilder {
GroupTypeBuilder::new(name)
}
pub fn get_basic_info(&self) -> &BasicTypeInfo {
match *self {
Type::PrimitiveType { ref basic_info, .. } => basic_info,
Type::GroupType { ref basic_info, .. } => basic_info,
}
}
pub fn name(&self) -> &str {
self.get_basic_info().name()
}
pub fn get_fields(&self) -> &[TypePtr] {
match *self {
Type::GroupType { ref fields, .. } => &fields[..],
_ => panic!("Cannot call get_fields() on a non-group type"),
}
}
pub fn get_physical_type(&self) -> PhysicalType {
match *self {
Type::PrimitiveType {
basic_info: _,
physical_type,
..
} => physical_type,
_ => panic!("Cannot call get_physical_type() on a non-primitive type"),
}
}
pub fn get_precision(&self) -> i32 {
match *self {
Type::PrimitiveType { precision, .. } => precision,
_ => panic!("Cannot call get_precision() on non-primitive type"),
}
}
pub fn get_scale(&self) -> i32 {
match *self {
Type::PrimitiveType { scale, .. } => scale,
_ => panic!("Cannot call get_scale() on non-primitive type"),
}
}
pub fn check_contains(&self, sub_type: &Type) -> bool {
let basic_match = self.get_basic_info().name() == sub_type.get_basic_info().name()
&& (self.is_schema() && sub_type.is_schema()
|| !self.is_schema()
&& !sub_type.is_schema()
&& self.get_basic_info().repetition()
== sub_type.get_basic_info().repetition());
match *self {
Type::PrimitiveType { .. } if basic_match && sub_type.is_primitive() => {
self.get_physical_type() == sub_type.get_physical_type()
}
Type::GroupType { .. } if basic_match && sub_type.is_group() => {
let mut field_map = HashMap::new();
for field in self.get_fields() {
field_map.insert(field.name(), field);
}
for field in sub_type.get_fields() {
if !field_map
.get(field.name())
.map(|tpe| tpe.check_contains(field))
.unwrap_or(false)
{
return false;
}
}
true
}
_ => false,
}
}
pub fn is_primitive(&self) -> bool {
matches!(*self, Type::PrimitiveType { .. })
}
pub fn is_group(&self) -> bool {
matches!(*self, Type::GroupType { .. })
}
pub fn is_schema(&self) -> bool {
match *self {
Type::GroupType { ref basic_info, .. } => !basic_info.has_repetition(),
_ => false,
}
}
pub fn is_optional(&self) -> bool {
self.get_basic_info().has_repetition()
&& self.get_basic_info().repetition() != Repetition::REQUIRED
}
pub(crate) fn is_list(&self) -> bool {
if self.is_group() {
let basic_info = self.get_basic_info();
if let Some(logical_type) = basic_info.logical_type() {
return logical_type == LogicalType::List;
}
return basic_info.converted_type() == ConvertedType::LIST;
}
false
}
pub(crate) fn has_single_repeated_child(&self) -> bool {
if self.is_group() {
let children = self.get_fields();
return children.len() == 1
&& children[0].get_basic_info().has_repetition()
&& children[0].get_basic_info().repetition() == Repetition::REPEATED;
}
false
}
}
pub struct PrimitiveTypeBuilder<'a> {
name: &'a str,
repetition: Repetition,
physical_type: PhysicalType,
converted_type: ConvertedType,
logical_type: Option<LogicalType>,
length: i32,
precision: i32,
scale: i32,
id: Option<i32>,
}
impl<'a> PrimitiveTypeBuilder<'a> {
pub fn new(name: &'a str, physical_type: PhysicalType) -> Self {
Self {
name,
repetition: Repetition::OPTIONAL,
physical_type,
converted_type: ConvertedType::NONE,
logical_type: None,
length: -1,
precision: -1,
scale: -1,
id: None,
}
}
pub fn with_repetition(self, repetition: Repetition) -> Self {
Self { repetition, ..self }
}
pub fn with_converted_type(self, converted_type: ConvertedType) -> Self {
Self {
converted_type,
..self
}
}
pub fn with_logical_type(self, logical_type: Option<LogicalType>) -> Self {
Self {
logical_type,
..self
}
}
pub fn with_length(self, length: i32) -> Self {
Self { length, ..self }
}
pub fn with_precision(self, precision: i32) -> Self {
Self { precision, ..self }
}
pub fn with_scale(self, scale: i32) -> Self {
Self { scale, ..self }
}
pub fn with_id(self, id: Option<i32>) -> Self {
Self { id, ..self }
}
pub fn build(self) -> Result<Type> {
let mut basic_info = BasicTypeInfo {
name: String::from(self.name),
repetition: Some(self.repetition),
converted_type: self.converted_type,
logical_type: self.logical_type.clone(),
id: self.id,
};
if self.physical_type == PhysicalType::FIXED_LEN_BYTE_ARRAY && self.length < 0 {
return Err(general_err!(
"Invalid FIXED_LEN_BYTE_ARRAY length: {} for field '{}'",
self.length,
self.name
));
}
if let Some(logical_type) = &self.logical_type {
if self.converted_type != ConvertedType::NONE {
if ConvertedType::from(self.logical_type.clone()) != self.converted_type {
return Err(general_err!(
"Logical type {:?} is incompatible with converted type {} for field '{}'",
logical_type,
self.converted_type,
self.name
));
}
} else {
basic_info.converted_type = self.logical_type.clone().into();
}
match (logical_type, self.physical_type) {
(LogicalType::Map, _) | (LogicalType::List, _) => {
return Err(general_err!(
"{:?} cannot be applied to a primitive type for field '{}'",
logical_type,
self.name
));
}
(LogicalType::Enum, PhysicalType::BYTE_ARRAY) => {}
(LogicalType::Decimal { scale, precision }, _) => {
if *scale != self.scale {
return Err(general_err!(
"DECIMAL logical type scale {} must match self.scale {} for field '{}'",
scale,
self.scale,
self.name
));
}
if *precision != self.precision {
return Err(general_err!(
"DECIMAL logical type precision {} must match self.precision {} for field '{}'",
precision,
self.precision,
self.name
));
}
self.check_decimal_precision_scale()?;
}
(LogicalType::Date, PhysicalType::INT32) => {}
(
LogicalType::Time {
unit: TimeUnit::MILLIS(_),
..
},
PhysicalType::INT32,
) => {}
(LogicalType::Time { unit, .. }, PhysicalType::INT64) => {
if *unit == TimeUnit::MILLIS(Default::default()) {
return Err(general_err!(
"Cannot use millisecond unit on INT64 type for field '{}'",
self.name
));
}
}
(LogicalType::Timestamp { .. }, PhysicalType::INT64) => {}
(LogicalType::Integer { bit_width, .. }, PhysicalType::INT32)
if *bit_width <= 32 => {}
(LogicalType::Integer { bit_width, .. }, PhysicalType::INT64)
if *bit_width == 64 => {}
(LogicalType::Unknown, PhysicalType::INT32) => {}
(LogicalType::String, PhysicalType::BYTE_ARRAY) => {}
(LogicalType::Json, PhysicalType::BYTE_ARRAY) => {}
(LogicalType::Bson, PhysicalType::BYTE_ARRAY) => {}
(LogicalType::Uuid, PhysicalType::FIXED_LEN_BYTE_ARRAY) if self.length == 16 => {}
(LogicalType::Uuid, PhysicalType::FIXED_LEN_BYTE_ARRAY) => {
return Err(general_err!(
"UUID cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(16) field",
self.name
))
}
(LogicalType::Float16, PhysicalType::FIXED_LEN_BYTE_ARRAY)
if self.length == 2 => {}
(LogicalType::Float16, PhysicalType::FIXED_LEN_BYTE_ARRAY) => {
return Err(general_err!(
"FLOAT16 cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(2) field",
self.name
))
}
(a, b) => {
return Err(general_err!(
"Cannot annotate {:?} from {} for field '{}'",
a,
b,
self.name
))
}
}
}
match self.converted_type {
ConvertedType::NONE => {}
ConvertedType::UTF8 | ConvertedType::BSON | ConvertedType::JSON => {
if self.physical_type != PhysicalType::BYTE_ARRAY {
return Err(general_err!(
"{} cannot annotate field '{}' because it is not a BYTE_ARRAY field",
self.converted_type,
self.name
));
}
}
ConvertedType::DECIMAL => {
self.check_decimal_precision_scale()?;
}
ConvertedType::DATE
| ConvertedType::TIME_MILLIS
| ConvertedType::UINT_8
| ConvertedType::UINT_16
| ConvertedType::UINT_32
| ConvertedType::INT_8
| ConvertedType::INT_16
| ConvertedType::INT_32 => {
if self.physical_type != PhysicalType::INT32 {
return Err(general_err!(
"{} cannot annotate field '{}' because it is not a INT32 field",
self.converted_type,
self.name
));
}
}
ConvertedType::TIME_MICROS
| ConvertedType::TIMESTAMP_MILLIS
| ConvertedType::TIMESTAMP_MICROS
| ConvertedType::UINT_64
| ConvertedType::INT_64 => {
if self.physical_type != PhysicalType::INT64 {
return Err(general_err!(
"{} cannot annotate field '{}' because it is not a INT64 field",
self.converted_type,
self.name
));
}
}
ConvertedType::INTERVAL => {
if self.physical_type != PhysicalType::FIXED_LEN_BYTE_ARRAY || self.length != 12 {
return Err(general_err!(
"INTERVAL cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(12) field",
self.name
));
}
}
ConvertedType::ENUM => {
if self.physical_type != PhysicalType::BYTE_ARRAY {
return Err(general_err!(
"ENUM cannot annotate field '{}' because it is not a BYTE_ARRAY field",
self.name
));
}
}
_ => {
return Err(general_err!(
"{} cannot be applied to primitive field '{}'",
self.converted_type,
self.name
));
}
}
Ok(Type::PrimitiveType {
basic_info,
physical_type: self.physical_type,
type_length: self.length,
scale: self.scale,
precision: self.precision,
})
}
#[inline]
fn check_decimal_precision_scale(&self) -> Result<()> {
match self.physical_type {
PhysicalType::INT32
| PhysicalType::INT64
| PhysicalType::BYTE_ARRAY
| PhysicalType::FIXED_LEN_BYTE_ARRAY => (),
_ => {
return Err(general_err!(
"DECIMAL can only annotate INT32, INT64, BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY"
));
}
}
if self.precision < 1 {
return Err(general_err!(
"Invalid DECIMAL precision: {}",
self.precision
));
}
if self.scale < 0 {
return Err(general_err!("Invalid DECIMAL scale: {}", self.scale));
}
if self.scale > self.precision {
return Err(general_err!(
"Invalid DECIMAL: scale ({}) cannot be greater than precision \
({})",
self.scale,
self.precision
));
}
match self.physical_type {
PhysicalType::INT32 => {
if self.precision > 9 {
return Err(general_err!(
"Cannot represent INT32 as DECIMAL with precision {}",
self.precision
));
}
}
PhysicalType::INT64 => {
if self.precision > 18 {
return Err(general_err!(
"Cannot represent INT64 as DECIMAL with precision {}",
self.precision
));
}
}
PhysicalType::FIXED_LEN_BYTE_ARRAY => {
let max_precision = (2f64.powi(8 * self.length - 1) - 1f64).log10().floor() as i32;
if self.precision > max_precision {
return Err(general_err!(
"Cannot represent FIXED_LEN_BYTE_ARRAY as DECIMAL with length {} and \
precision {}. The max precision can only be {}",
self.length,
self.precision,
max_precision
));
}
}
_ => (), }
Ok(())
}
}
pub struct GroupTypeBuilder<'a> {
name: &'a str,
repetition: Option<Repetition>,
converted_type: ConvertedType,
logical_type: Option<LogicalType>,
fields: Vec<TypePtr>,
id: Option<i32>,
}
impl<'a> GroupTypeBuilder<'a> {
pub fn new(name: &'a str) -> Self {
Self {
name,
repetition: None,
converted_type: ConvertedType::NONE,
logical_type: None,
fields: Vec::new(),
id: None,
}
}
pub fn with_repetition(mut self, repetition: Repetition) -> Self {
self.repetition = Some(repetition);
self
}
pub fn with_converted_type(self, converted_type: ConvertedType) -> Self {
Self {
converted_type,
..self
}
}
pub fn with_logical_type(self, logical_type: Option<LogicalType>) -> Self {
Self {
logical_type,
..self
}
}
pub fn with_fields(self, fields: Vec<TypePtr>) -> Self {
Self { fields, ..self }
}
pub fn with_id(self, id: Option<i32>) -> Self {
Self { id, ..self }
}
pub fn build(self) -> Result<Type> {
let mut basic_info = BasicTypeInfo {
name: String::from(self.name),
repetition: self.repetition,
converted_type: self.converted_type,
logical_type: self.logical_type.clone(),
id: self.id,
};
if self.logical_type.is_some() && self.converted_type == ConvertedType::NONE {
basic_info.converted_type = self.logical_type.into();
}
Ok(Type::GroupType {
basic_info,
fields: self.fields,
})
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct BasicTypeInfo {
name: String,
repetition: Option<Repetition>,
converted_type: ConvertedType,
logical_type: Option<LogicalType>,
id: Option<i32>,
}
impl HeapSize for BasicTypeInfo {
fn heap_size(&self) -> usize {
self.name.heap_size()
}
}
impl BasicTypeInfo {
pub fn name(&self) -> &str {
&self.name
}
pub fn has_repetition(&self) -> bool {
self.repetition.is_some()
}
pub fn repetition(&self) -> Repetition {
assert!(self.repetition.is_some());
self.repetition.unwrap()
}
pub fn converted_type(&self) -> ConvertedType {
self.converted_type
}
pub fn logical_type(&self) -> Option<LogicalType> {
self.logical_type.clone()
}
pub fn has_id(&self) -> bool {
self.id.is_some()
}
pub fn id(&self) -> i32 {
assert!(self.id.is_some());
self.id.unwrap()
}
}
#[derive(Clone, PartialEq, Debug, Eq, Hash)]
pub struct ColumnPath {
parts: Vec<String>,
}
impl HeapSize for ColumnPath {
fn heap_size(&self) -> usize {
self.parts.heap_size()
}
}
impl ColumnPath {
pub fn new(parts: Vec<String>) -> Self {
ColumnPath { parts }
}
pub fn string(&self) -> String {
self.parts.join(".")
}
pub fn append(&mut self, mut tail: Vec<String>) {
self.parts.append(&mut tail);
}
pub fn parts(&self) -> &[String] {
&self.parts
}
}
impl fmt::Display for ColumnPath {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{:?}", self.string())
}
}
impl From<Vec<String>> for ColumnPath {
fn from(parts: Vec<String>) -> Self {
ColumnPath { parts }
}
}
impl From<&str> for ColumnPath {
fn from(single_path: &str) -> Self {
let s = String::from(single_path);
ColumnPath::from(s)
}
}
impl From<String> for ColumnPath {
fn from(single_path: String) -> Self {
let v = vec![single_path];
ColumnPath { parts: v }
}
}
impl AsRef<[String]> for ColumnPath {
fn as_ref(&self) -> &[String] {
&self.parts
}
}
#[derive(Debug, PartialEq)]
pub struct ColumnDescriptor {
primitive_type: TypePtr,
max_def_level: i16,
max_rep_level: i16,
path: ColumnPath,
}
impl HeapSize for ColumnDescriptor {
fn heap_size(&self) -> usize {
self.primitive_type.heap_size() + self.path.heap_size()
}
}
impl ColumnDescriptor {
pub fn new(
primitive_type: TypePtr,
max_def_level: i16,
max_rep_level: i16,
path: ColumnPath,
) -> Self {
Self {
primitive_type,
max_def_level,
max_rep_level,
path,
}
}
#[inline]
pub fn max_def_level(&self) -> i16 {
self.max_def_level
}
#[inline]
pub fn max_rep_level(&self) -> i16 {
self.max_rep_level
}
pub fn path(&self) -> &ColumnPath {
&self.path
}
pub fn self_type(&self) -> &Type {
self.primitive_type.as_ref()
}
pub fn self_type_ptr(&self) -> TypePtr {
self.primitive_type.clone()
}
pub fn name(&self) -> &str {
self.primitive_type.name()
}
pub fn converted_type(&self) -> ConvertedType {
self.primitive_type.get_basic_info().converted_type()
}
pub fn logical_type(&self) -> Option<LogicalType> {
self.primitive_type.get_basic_info().logical_type()
}
pub fn physical_type(&self) -> PhysicalType {
match self.primitive_type.as_ref() {
Type::PrimitiveType { physical_type, .. } => *physical_type,
_ => panic!("Expected primitive type!"),
}
}
pub fn type_length(&self) -> i32 {
match self.primitive_type.as_ref() {
Type::PrimitiveType { type_length, .. } => *type_length,
_ => panic!("Expected primitive type!"),
}
}
pub fn type_precision(&self) -> i32 {
match self.primitive_type.as_ref() {
Type::PrimitiveType { precision, .. } => *precision,
_ => panic!("Expected primitive type!"),
}
}
pub fn type_scale(&self) -> i32 {
match self.primitive_type.as_ref() {
Type::PrimitiveType { scale, .. } => *scale,
_ => panic!("Expected primitive type!"),
}
}
pub fn sort_order(&self) -> SortOrder {
ColumnOrder::get_sort_order(
self.logical_type(),
self.converted_type(),
self.physical_type(),
)
}
}
#[derive(PartialEq)]
pub struct SchemaDescriptor {
schema: TypePtr,
leaves: Vec<ColumnDescPtr>,
leaf_to_base: Vec<usize>,
}
impl fmt::Debug for SchemaDescriptor {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SchemaDescriptor")
.field("schema", &self.schema)
.finish()
}
}
impl HeapSize for SchemaDescriptor {
fn heap_size(&self) -> usize {
self.schema.heap_size() + self.leaves.heap_size() + self.leaf_to_base.heap_size()
}
}
impl SchemaDescriptor {
pub fn new(tp: TypePtr) -> Self {
assert!(tp.is_group(), "SchemaDescriptor should take a GroupType");
let mut leaves = vec![];
let mut leaf_to_base = Vec::new();
for (root_idx, f) in tp.get_fields().iter().enumerate() {
let mut path = vec![];
build_tree(f, root_idx, 0, 0, &mut leaves, &mut leaf_to_base, &mut path);
}
Self {
schema: tp,
leaves,
leaf_to_base,
}
}
pub fn column(&self, i: usize) -> ColumnDescPtr {
assert!(
i < self.leaves.len(),
"Index out of bound: {} not in [0, {})",
i,
self.leaves.len()
);
self.leaves[i].clone()
}
pub fn columns(&self) -> &[ColumnDescPtr] {
&self.leaves
}
pub fn num_columns(&self) -> usize {
self.leaves.len()
}
pub fn get_column_root(&self, i: usize) -> &Type {
let result = self.column_root_of(i);
result.as_ref()
}
pub fn get_column_root_ptr(&self, i: usize) -> TypePtr {
let result = self.column_root_of(i);
result.clone()
}
pub fn get_column_root_idx(&self, leaf: usize) -> usize {
assert!(
leaf < self.leaves.len(),
"Index out of bound: {} not in [0, {})",
leaf,
self.leaves.len()
);
*self
.leaf_to_base
.get(leaf)
.unwrap_or_else(|| panic!("Expected a value for index {leaf} but found None"))
}
fn column_root_of(&self, i: usize) -> &TypePtr {
&self.schema.get_fields()[self.get_column_root_idx(i)]
}
pub fn root_schema(&self) -> &Type {
self.schema.as_ref()
}
pub fn root_schema_ptr(&self) -> TypePtr {
self.schema.clone()
}
pub fn name(&self) -> &str {
self.schema.name()
}
}
fn build_tree<'a>(
tp: &'a TypePtr,
root_idx: usize,
mut max_rep_level: i16,
mut max_def_level: i16,
leaves: &mut Vec<ColumnDescPtr>,
leaf_to_base: &mut Vec<usize>,
path_so_far: &mut Vec<&'a str>,
) {
assert!(tp.get_basic_info().has_repetition());
path_so_far.push(tp.name());
match tp.get_basic_info().repetition() {
Repetition::OPTIONAL => {
max_def_level += 1;
}
Repetition::REPEATED => {
max_def_level += 1;
max_rep_level += 1;
}
_ => {}
}
match tp.as_ref() {
Type::PrimitiveType { .. } => {
let mut path: Vec<String> = vec![];
path.extend(path_so_far.iter().copied().map(String::from));
leaves.push(Arc::new(ColumnDescriptor::new(
tp.clone(),
max_def_level,
max_rep_level,
ColumnPath::new(path),
)));
leaf_to_base.push(root_idx);
}
Type::GroupType { ref fields, .. } => {
for f in fields {
build_tree(
f,
root_idx,
max_rep_level,
max_def_level,
leaves,
leaf_to_base,
path_so_far,
);
path_so_far.pop();
}
}
}
}
pub fn from_thrift(elements: &[SchemaElement]) -> Result<TypePtr> {
let mut index = 0;
let mut schema_nodes = Vec::new();
while index < elements.len() {
let t = from_thrift_helper(elements, index)?;
index = t.0;
schema_nodes.push(t.1);
}
if schema_nodes.len() != 1 {
return Err(general_err!(
"Expected exactly one root node, but found {}",
schema_nodes.len()
));
}
Ok(schema_nodes.remove(0))
}
fn from_thrift_helper(elements: &[SchemaElement], index: usize) -> Result<(usize, TypePtr)> {
let is_root_node = index == 0;
if index >= elements.len() {
return Err(general_err!(
"Index out of bound, index = {}, len = {}",
index,
elements.len()
));
}
let element = &elements[index];
let converted_type = ConvertedType::try_from(element.converted_type)?;
let logical_type = element
.logical_type
.as_ref()
.map(|value| LogicalType::from(value.clone()));
let field_id = elements[index].field_id;
match elements[index].num_children {
None | Some(0) => {
if elements[index].repetition_type.is_none() {
return Err(general_err!(
"Repetition level must be defined for a primitive type"
));
}
let repetition = Repetition::try_from(elements[index].repetition_type.unwrap())?;
if let Some(type_) = elements[index].type_ {
let physical_type = PhysicalType::try_from(type_)?;
let length = elements[index].type_length.unwrap_or(-1);
let scale = elements[index].scale.unwrap_or(-1);
let precision = elements[index].precision.unwrap_or(-1);
let name = &elements[index].name;
let builder = Type::primitive_type_builder(name, physical_type)
.with_repetition(repetition)
.with_converted_type(converted_type)
.with_logical_type(logical_type)
.with_length(length)
.with_precision(precision)
.with_scale(scale)
.with_id(field_id);
Ok((index + 1, Arc::new(builder.build()?)))
} else {
let mut builder = Type::group_type_builder(&elements[index].name)
.with_converted_type(converted_type)
.with_logical_type(logical_type)
.with_id(field_id);
if !is_root_node {
builder = builder.with_repetition(repetition);
}
Ok((index + 1, Arc::new(builder.build().unwrap())))
}
}
Some(n) => {
let repetition = elements[index]
.repetition_type
.map(Repetition::try_from)
.transpose()?;
let mut fields = vec![];
let mut next_index = index + 1;
for _ in 0..n {
let child_result = from_thrift_helper(elements, next_index)?;
next_index = child_result.0;
fields.push(child_result.1);
}
let mut builder = Type::group_type_builder(&elements[index].name)
.with_converted_type(converted_type)
.with_logical_type(logical_type)
.with_fields(fields)
.with_id(field_id);
if let Some(rep) = repetition {
if !is_root_node {
builder = builder.with_repetition(rep);
}
}
Ok((next_index, Arc::new(builder.build().unwrap())))
}
}
}
pub fn to_thrift(schema: &Type) -> Result<Vec<SchemaElement>> {
if !schema.is_group() {
return Err(general_err!("Root schema must be Group type"));
}
let mut elements: Vec<SchemaElement> = Vec::new();
to_thrift_helper(schema, &mut elements);
Ok(elements)
}
fn to_thrift_helper(schema: &Type, elements: &mut Vec<SchemaElement>) {
match *schema {
Type::PrimitiveType {
ref basic_info,
physical_type,
type_length,
scale,
precision,
} => {
let element = SchemaElement {
type_: Some(physical_type.into()),
type_length: if type_length >= 0 {
Some(type_length)
} else {
None
},
repetition_type: Some(basic_info.repetition().into()),
name: basic_info.name().to_owned(),
num_children: None,
converted_type: basic_info.converted_type().into(),
scale: if scale >= 0 { Some(scale) } else { None },
precision: if precision >= 0 {
Some(precision)
} else {
None
},
field_id: if basic_info.has_id() {
Some(basic_info.id())
} else {
None
},
logical_type: basic_info.logical_type().map(|value| value.into()),
};
elements.push(element);
}
Type::GroupType {
ref basic_info,
ref fields,
} => {
let repetition = if basic_info.has_repetition() {
Some(basic_info.repetition().into())
} else {
None
};
let element = SchemaElement {
type_: None,
type_length: None,
repetition_type: repetition,
name: basic_info.name().to_owned(),
num_children: Some(fields.len() as i32),
converted_type: basic_info.converted_type().into(),
scale: None,
precision: None,
field_id: if basic_info.has_id() {
Some(basic_info.id())
} else {
None
},
logical_type: basic_info.logical_type().map(|value| value.into()),
};
elements.push(element);
for field in fields {
to_thrift_helper(field, elements);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::schema::parser::parse_message_type;
#[test]
fn test_primitive_type() {
let mut result = Type::primitive_type_builder("foo", PhysicalType::INT32)
.with_logical_type(Some(LogicalType::Integer {
bit_width: 32,
is_signed: true,
}))
.with_id(Some(0))
.build();
assert!(result.is_ok());
if let Ok(tp) = result {
assert!(tp.is_primitive());
assert!(!tp.is_group());
let basic_info = tp.get_basic_info();
assert_eq!(basic_info.repetition(), Repetition::OPTIONAL);
assert_eq!(
basic_info.logical_type(),
Some(LogicalType::Integer {
bit_width: 32,
is_signed: true
})
);
assert_eq!(basic_info.converted_type(), ConvertedType::INT_32);
assert_eq!(basic_info.id(), 0);
match tp {
Type::PrimitiveType { physical_type, .. } => {
assert_eq!(physical_type, PhysicalType::INT32);
}
_ => panic!(),
}
}
result = Type::primitive_type_builder("foo", PhysicalType::INT64)
.with_repetition(Repetition::REPEATED)
.with_logical_type(Some(LogicalType::Integer {
is_signed: true,
bit_width: 8,
}))
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: Cannot annotate Integer { bit_width: 8, is_signed: true } from INT64 for field 'foo'"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::INT64)
.with_repetition(Repetition::REPEATED)
.with_converted_type(ConvertedType::BSON)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: BSON cannot annotate field 'foo' because it is not a BYTE_ARRAY field"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::INT96)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::DECIMAL)
.with_precision(-1)
.with_scale(-1)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: DECIMAL can only annotate INT32, INT64, BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
.with_repetition(Repetition::REQUIRED)
.with_logical_type(Some(LogicalType::Decimal {
scale: 32,
precision: 12,
}))
.with_precision(-1)
.with_scale(-1)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: DECIMAL logical type scale 32 must match self.scale -1 for field 'foo'"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::DECIMAL)
.with_precision(-1)
.with_scale(-1)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: Invalid DECIMAL precision: -1"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::DECIMAL)
.with_precision(0)
.with_scale(-1)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: Invalid DECIMAL precision: 0"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::DECIMAL)
.with_precision(1)
.with_scale(-1)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(format!("{e}"), "Parquet error: Invalid DECIMAL scale: -1");
}
result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::DECIMAL)
.with_precision(1)
.with_scale(2)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: Invalid DECIMAL: scale (2) cannot be greater than precision (1)"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::DECIMAL)
.with_precision(1)
.with_scale(1)
.build();
assert!(result.is_ok());
result = Type::primitive_type_builder("foo", PhysicalType::INT32)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::DECIMAL)
.with_precision(18)
.with_scale(2)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: Cannot represent INT32 as DECIMAL with precision 18"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::INT64)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::DECIMAL)
.with_precision(32)
.with_scale(2)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: Cannot represent INT64 as DECIMAL with precision 32"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::DECIMAL)
.with_length(5)
.with_precision(12)
.with_scale(2)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: Cannot represent FIXED_LEN_BYTE_ARRAY as DECIMAL with length 5 and precision 12. The max precision can only be 11"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::INT64)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::UINT_8)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: UINT_8 cannot annotate field 'foo' because it is not a INT32 field"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::INT32)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::TIME_MICROS)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: TIME_MICROS cannot annotate field 'foo' because it is not a INT64 field"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::INTERVAL)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: INTERVAL cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(12) field"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::INTERVAL)
.with_length(1)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: INTERVAL cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(12) field"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::INT32)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::ENUM)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: ENUM cannot annotate field 'foo' because it is not a BYTE_ARRAY field"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::INT32)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::MAP)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: MAP cannot be applied to primitive field 'foo'"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::DECIMAL)
.with_length(-1)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: Invalid FIXED_LEN_BYTE_ARRAY length: -1 for field 'foo'"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
.with_repetition(Repetition::REQUIRED)
.with_logical_type(Some(LogicalType::Float16))
.with_length(2)
.build();
assert!(result.is_ok());
result = Type::primitive_type_builder("foo", PhysicalType::FLOAT)
.with_repetition(Repetition::REQUIRED)
.with_logical_type(Some(LogicalType::Float16))
.with_length(2)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: Cannot annotate Float16 from FLOAT for field 'foo'"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
.with_repetition(Repetition::REQUIRED)
.with_logical_type(Some(LogicalType::Float16))
.with_length(4)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: FLOAT16 cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(2) field"
);
}
result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
.with_repetition(Repetition::REQUIRED)
.with_logical_type(Some(LogicalType::Uuid))
.with_length(15)
.build();
assert!(result.is_err());
if let Err(e) = result {
assert_eq!(
format!("{e}"),
"Parquet error: UUID cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(16) field"
);
}
}
#[test]
fn test_group_type() {
let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32)
.with_converted_type(ConvertedType::INT_32)
.with_id(Some(0))
.build();
assert!(f1.is_ok());
let f2 = Type::primitive_type_builder("f2", PhysicalType::BYTE_ARRAY)
.with_converted_type(ConvertedType::UTF8)
.with_id(Some(1))
.build();
assert!(f2.is_ok());
let fields = vec![Arc::new(f1.unwrap()), Arc::new(f2.unwrap())];
let result = Type::group_type_builder("foo")
.with_repetition(Repetition::REPEATED)
.with_logical_type(Some(LogicalType::List))
.with_fields(fields)
.with_id(Some(1))
.build();
assert!(result.is_ok());
let tp = result.unwrap();
let basic_info = tp.get_basic_info();
assert!(tp.is_group());
assert!(!tp.is_primitive());
assert_eq!(basic_info.repetition(), Repetition::REPEATED);
assert_eq!(basic_info.logical_type(), Some(LogicalType::List));
assert_eq!(basic_info.converted_type(), ConvertedType::LIST);
assert_eq!(basic_info.id(), 1);
assert_eq!(tp.get_fields().len(), 2);
assert_eq!(tp.get_fields()[0].name(), "f1");
assert_eq!(tp.get_fields()[1].name(), "f2");
}
#[test]
fn test_column_descriptor() {
let result = test_column_descriptor_helper();
assert!(
result.is_ok(),
"Expected result to be OK but got err:\n {}",
result.unwrap_err()
);
}
fn test_column_descriptor_helper() -> Result<()> {
let tp = Type::primitive_type_builder("name", PhysicalType::BYTE_ARRAY)
.with_converted_type(ConvertedType::UTF8)
.build()?;
let descr = ColumnDescriptor::new(Arc::new(tp), 4, 1, ColumnPath::from("name"));
assert_eq!(descr.path(), &ColumnPath::from("name"));
assert_eq!(descr.converted_type(), ConvertedType::UTF8);
assert_eq!(descr.physical_type(), PhysicalType::BYTE_ARRAY);
assert_eq!(descr.max_def_level(), 4);
assert_eq!(descr.max_rep_level(), 1);
assert_eq!(descr.name(), "name");
assert_eq!(descr.type_length(), -1);
assert_eq!(descr.type_precision(), -1);
assert_eq!(descr.type_scale(), -1);
Ok(())
}
#[test]
fn test_schema_descriptor() {
let result = test_schema_descriptor_helper();
assert!(
result.is_ok(),
"Expected result to be OK but got err:\n {}",
result.unwrap_err()
);
}
fn test_schema_descriptor_helper() -> Result<()> {
let mut fields = vec![];
let inta = Type::primitive_type_builder("a", PhysicalType::INT32)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::INT_32)
.build()?;
fields.push(Arc::new(inta));
let intb = Type::primitive_type_builder("b", PhysicalType::INT64)
.with_converted_type(ConvertedType::INT_64)
.build()?;
fields.push(Arc::new(intb));
let intc = Type::primitive_type_builder("c", PhysicalType::BYTE_ARRAY)
.with_repetition(Repetition::REPEATED)
.with_converted_type(ConvertedType::UTF8)
.build()?;
fields.push(Arc::new(intc));
let item1 = Type::primitive_type_builder("item1", PhysicalType::INT64)
.with_repetition(Repetition::REQUIRED)
.with_converted_type(ConvertedType::INT_64)
.build()?;
let item2 = Type::primitive_type_builder("item2", PhysicalType::BOOLEAN).build()?;
let item3 = Type::primitive_type_builder("item3", PhysicalType::INT32)
.with_repetition(Repetition::REPEATED)
.with_converted_type(ConvertedType::INT_32)
.build()?;
let list = Type::group_type_builder("records")
.with_repetition(Repetition::REPEATED)
.with_converted_type(ConvertedType::LIST)
.with_fields(vec![Arc::new(item1), Arc::new(item2), Arc::new(item3)])
.build()?;
let bag = Type::group_type_builder("bag")
.with_repetition(Repetition::OPTIONAL)
.with_fields(vec![Arc::new(list)])
.build()?;
fields.push(Arc::new(bag));
let schema = Type::group_type_builder("schema")
.with_repetition(Repetition::REPEATED)
.with_fields(fields)
.build()?;
let descr = SchemaDescriptor::new(Arc::new(schema));
let nleaves = 6;
assert_eq!(descr.num_columns(), nleaves);
let ex_max_def_levels = [0, 1, 1, 2, 3, 3];
let ex_max_rep_levels = [0, 0, 1, 1, 1, 2];
for i in 0..nleaves {
let col = descr.column(i);
assert_eq!(col.max_def_level(), ex_max_def_levels[i], "{i}");
assert_eq!(col.max_rep_level(), ex_max_rep_levels[i], "{i}");
}
assert_eq!(descr.column(0).path().string(), "a");
assert_eq!(descr.column(1).path().string(), "b");
assert_eq!(descr.column(2).path().string(), "c");
assert_eq!(descr.column(3).path().string(), "bag.records.item1");
assert_eq!(descr.column(4).path().string(), "bag.records.item2");
assert_eq!(descr.column(5).path().string(), "bag.records.item3");
assert_eq!(descr.get_column_root(0).name(), "a");
assert_eq!(descr.get_column_root(3).name(), "bag");
assert_eq!(descr.get_column_root(4).name(), "bag");
assert_eq!(descr.get_column_root(5).name(), "bag");
Ok(())
}
#[test]
fn test_schema_build_tree_def_rep_levels() {
let message_type = "
message spark_schema {
REQUIRED INT32 a;
OPTIONAL group b {
OPTIONAL INT32 _1;
OPTIONAL INT32 _2;
}
OPTIONAL group c (LIST) {
REPEATED group list {
OPTIONAL INT32 element;
}
}
}
";
let schema = parse_message_type(message_type).expect("should parse schema");
let descr = SchemaDescriptor::new(Arc::new(schema));
assert_eq!(descr.column(0).max_def_level(), 0);
assert_eq!(descr.column(0).max_rep_level(), 0);
assert_eq!(descr.column(1).max_def_level(), 2);
assert_eq!(descr.column(1).max_rep_level(), 0);
assert_eq!(descr.column(2).max_def_level(), 2);
assert_eq!(descr.column(2).max_rep_level(), 0);
assert_eq!(descr.column(3).max_def_level(), 3);
assert_eq!(descr.column(3).max_rep_level(), 1);
}
#[test]
#[should_panic(expected = "Cannot call get_physical_type() on a non-primitive type")]
fn test_get_physical_type_panic() {
let list = Type::group_type_builder("records")
.with_repetition(Repetition::REPEATED)
.build()
.unwrap();
list.get_physical_type();
}
#[test]
fn test_get_physical_type_primitive() {
let f = Type::primitive_type_builder("f", PhysicalType::INT64)
.build()
.unwrap();
assert_eq!(f.get_physical_type(), PhysicalType::INT64);
let f = Type::primitive_type_builder("f", PhysicalType::BYTE_ARRAY)
.build()
.unwrap();
assert_eq!(f.get_physical_type(), PhysicalType::BYTE_ARRAY);
}
#[test]
fn test_check_contains_primitive_primitive() {
let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
.build()
.unwrap();
let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
.build()
.unwrap();
assert!(f1.check_contains(&f2));
let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
.with_converted_type(ConvertedType::UINT_8)
.build()
.unwrap();
let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
.with_converted_type(ConvertedType::UINT_16)
.build()
.unwrap();
assert!(f1.check_contains(&f2));
let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32)
.build()
.unwrap();
let f2 = Type::primitive_type_builder("f2", PhysicalType::INT32)
.build()
.unwrap();
assert!(!f1.check_contains(&f2));
let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
.build()
.unwrap();
let f2 = Type::primitive_type_builder("f", PhysicalType::INT64)
.build()
.unwrap();
assert!(!f1.check_contains(&f2));
let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
.with_repetition(Repetition::REQUIRED)
.build()
.unwrap();
let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
.with_repetition(Repetition::OPTIONAL)
.build()
.unwrap();
assert!(!f1.check_contains(&f2));
}
fn test_new_group_type(name: &str, repetition: Repetition, types: Vec<Type>) -> Type {
Type::group_type_builder(name)
.with_repetition(repetition)
.with_fields(types.into_iter().map(Arc::new).collect())
.build()
.unwrap()
}
#[test]
fn test_check_contains_group_group() {
let f1 = Type::group_type_builder("f").build().unwrap();
let f2 = Type::group_type_builder("f").build().unwrap();
assert!(f1.check_contains(&f2));
assert!(!f1.is_optional());
let f1 = test_new_group_type(
"f",
Repetition::REPEATED,
vec![
Type::primitive_type_builder("f1", PhysicalType::INT32)
.build()
.unwrap(),
Type::primitive_type_builder("f2", PhysicalType::INT64)
.build()
.unwrap(),
],
);
let f2 = test_new_group_type(
"f",
Repetition::REPEATED,
vec![
Type::primitive_type_builder("f1", PhysicalType::INT32)
.build()
.unwrap(),
Type::primitive_type_builder("f2", PhysicalType::INT64)
.build()
.unwrap(),
],
);
assert!(f1.check_contains(&f2));
let f1 = test_new_group_type(
"f",
Repetition::REPEATED,
vec![
Type::primitive_type_builder("f1", PhysicalType::INT32)
.build()
.unwrap(),
Type::primitive_type_builder("f2", PhysicalType::INT64)
.build()
.unwrap(),
],
);
let f2 = test_new_group_type(
"f",
Repetition::REPEATED,
vec![Type::primitive_type_builder("f2", PhysicalType::INT64)
.build()
.unwrap()],
);
assert!(f1.check_contains(&f2));
let f1 = Type::group_type_builder("f1").build().unwrap();
let f2 = Type::group_type_builder("f2").build().unwrap();
assert!(!f1.check_contains(&f2));
let f1 = Type::group_type_builder("f")
.with_repetition(Repetition::OPTIONAL)
.build()
.unwrap();
let f2 = Type::group_type_builder("f")
.with_repetition(Repetition::REPEATED)
.build()
.unwrap();
assert!(!f1.check_contains(&f2));
let f1 = test_new_group_type(
"f",
Repetition::REPEATED,
vec![
Type::primitive_type_builder("f1", PhysicalType::INT32)
.build()
.unwrap(),
Type::primitive_type_builder("f2", PhysicalType::INT64)
.build()
.unwrap(),
],
);
let f2 = test_new_group_type(
"f",
Repetition::REPEATED,
vec![
Type::primitive_type_builder("f1", PhysicalType::INT32)
.build()
.unwrap(),
Type::primitive_type_builder("f2", PhysicalType::BOOLEAN)
.build()
.unwrap(),
],
);
assert!(!f1.check_contains(&f2));
let f1 = test_new_group_type(
"f",
Repetition::REPEATED,
vec![
Type::primitive_type_builder("f1", PhysicalType::INT32)
.build()
.unwrap(),
Type::primitive_type_builder("f2", PhysicalType::INT64)
.build()
.unwrap(),
],
);
let f2 = test_new_group_type(
"f",
Repetition::REPEATED,
vec![Type::primitive_type_builder("f3", PhysicalType::INT32)
.build()
.unwrap()],
);
assert!(!f1.check_contains(&f2));
}
#[test]
fn test_check_contains_group_primitive() {
let f1 = Type::group_type_builder("f").build().unwrap();
let f2 = Type::primitive_type_builder("f", PhysicalType::INT64)
.build()
.unwrap();
assert!(!f1.check_contains(&f2));
assert!(!f2.check_contains(&f1));
let f1 = test_new_group_type(
"f",
Repetition::REPEATED,
vec![Type::primitive_type_builder("f1", PhysicalType::INT32)
.build()
.unwrap()],
);
let f2 = Type::primitive_type_builder("f1", PhysicalType::INT32)
.build()
.unwrap();
assert!(!f1.check_contains(&f2));
assert!(!f2.check_contains(&f1));
let f1 = test_new_group_type(
"a",
Repetition::REPEATED,
vec![
test_new_group_type(
"b",
Repetition::REPEATED,
vec![Type::primitive_type_builder("c", PhysicalType::INT32)
.build()
.unwrap()],
),
Type::primitive_type_builder("d", PhysicalType::INT64)
.build()
.unwrap(),
Type::primitive_type_builder("e", PhysicalType::BOOLEAN)
.build()
.unwrap(),
],
);
let f2 = test_new_group_type(
"a",
Repetition::REPEATED,
vec![test_new_group_type(
"b",
Repetition::REPEATED,
vec![Type::primitive_type_builder("c", PhysicalType::INT32)
.build()
.unwrap()],
)],
);
assert!(f1.check_contains(&f2)); assert!(!f2.check_contains(&f1)); }
#[test]
fn test_schema_type_thrift_conversion_err() {
let schema = Type::primitive_type_builder("col", PhysicalType::INT32)
.build()
.unwrap();
let thrift_schema = to_thrift(&schema);
assert!(thrift_schema.is_err());
if let Err(e) = thrift_schema {
assert_eq!(
format!("{e}"),
"Parquet error: Root schema must be Group type"
);
}
}
#[test]
fn test_schema_type_thrift_conversion() {
let message_type = "
message conversions {
REQUIRED INT64 id;
OPTIONAL FIXED_LEN_BYTE_ARRAY (2) f16 (FLOAT16);
OPTIONAL group int_array_Array (LIST) {
REPEATED group list {
OPTIONAL group element (LIST) {
REPEATED group list {
OPTIONAL INT32 element;
}
}
}
}
OPTIONAL group int_map (MAP) {
REPEATED group map (MAP_KEY_VALUE) {
REQUIRED BYTE_ARRAY key (UTF8);
OPTIONAL INT32 value;
}
}
OPTIONAL group int_Map_Array (LIST) {
REPEATED group list {
OPTIONAL group g (MAP) {
REPEATED group map (MAP_KEY_VALUE) {
REQUIRED BYTE_ARRAY key (UTF8);
OPTIONAL group value {
OPTIONAL group H {
OPTIONAL group i (LIST) {
REPEATED group list {
OPTIONAL DOUBLE element;
}
}
}
}
}
}
}
}
OPTIONAL group nested_struct {
OPTIONAL INT32 A;
OPTIONAL group b (LIST) {
REPEATED group list {
REQUIRED FIXED_LEN_BYTE_ARRAY (16) element;
}
}
}
}
";
let expected_schema = parse_message_type(message_type).unwrap();
let thrift_schema = to_thrift(&expected_schema).unwrap();
let result_schema = from_thrift(&thrift_schema).unwrap();
assert_eq!(result_schema, Arc::new(expected_schema));
}
#[test]
fn test_schema_type_thrift_conversion_decimal() {
let message_type = "
message decimals {
OPTIONAL INT32 field0;
OPTIONAL INT64 field1 (DECIMAL (18, 2));
OPTIONAL FIXED_LEN_BYTE_ARRAY (16) field2 (DECIMAL (38, 18));
OPTIONAL BYTE_ARRAY field3 (DECIMAL (9));
}
";
let expected_schema = parse_message_type(message_type).unwrap();
let thrift_schema = to_thrift(&expected_schema).unwrap();
let result_schema = from_thrift(&thrift_schema).unwrap();
assert_eq!(result_schema, Arc::new(expected_schema));
}
#[test]
fn test_schema_from_thrift_with_num_children_set() {
let message_type = "
message schema {
OPTIONAL BYTE_ARRAY id (UTF8);
OPTIONAL BYTE_ARRAY name (UTF8);
OPTIONAL BYTE_ARRAY message (UTF8);
OPTIONAL INT32 type (UINT_8);
OPTIONAL INT64 author_time (TIMESTAMP_MILLIS);
OPTIONAL INT64 __index_level_0__;
}
";
let expected_schema = parse_message_type(message_type).unwrap();
let mut thrift_schema = to_thrift(&expected_schema).unwrap();
for elem in &mut thrift_schema[..] {
if elem.num_children.is_none() {
elem.num_children = Some(0);
}
}
let result_schema = from_thrift(&thrift_schema).unwrap();
assert_eq!(result_schema, Arc::new(expected_schema));
}
#[test]
fn test_schema_from_thrift_root_has_repetition() {
let message_type = "
message schema {
OPTIONAL BYTE_ARRAY a (UTF8);
OPTIONAL INT32 b (UINT_8);
}
";
let expected_schema = parse_message_type(message_type).unwrap();
let mut thrift_schema = to_thrift(&expected_schema).unwrap();
thrift_schema[0].repetition_type = Some(Repetition::REQUIRED.into());
let result_schema = from_thrift(&thrift_schema).unwrap();
assert_eq!(result_schema, Arc::new(expected_schema));
}
#[test]
fn test_schema_from_thrift_group_has_no_child() {
let message_type = "message schema {}";
let expected_schema = parse_message_type(message_type).unwrap();
let mut thrift_schema = to_thrift(&expected_schema).unwrap();
thrift_schema[0].repetition_type = Some(Repetition::REQUIRED.into());
let result_schema = from_thrift(&thrift_schema).unwrap();
assert_eq!(result_schema, Arc::new(expected_schema));
}
}