| 123456789101112131415161718192021222324252627282930313233343536373839 |
- """Utils for working with the parquet thrift models"""
- from ttypes import FieldRepetitionType
- class SchemaHelper(object):
- def __init__(self, schema_elements):
- self.schema_elements = schema_elements
- self.schema_elements_by_name = dict(
- [(se.name, se) for se in schema_elements])
- assert len(self.schema_elements) == len(self.schema_elements_by_name)
- def schema_element(self, name):
- """Get the schema element with the given name."""
- return self.schema_elements_by_name[name]
- def is_required(self, name):
- """Returns true iff the schema element with the given name is
- required"""
- return self.schema_element(name).repetition_type == FieldRepetitionType.REQUIRED
- def max_repetition_level(self, path):
- """get the max repetition level for the given schema path."""
- max_level = 0
- for part in path:
- se = self.schema_element(part)
- if se.repetition_type == FieldRepetitionType.REQUIRED:
- max_level += 1
- return max_level
- def max_definition_level(self, path):
- """get the max definition level for the given schema path."""
- max_level = 0
- for part in path:
- se = self.schema_element(part)
- if se.repetition_type != FieldRepetitionType.REQUIRED:
- max_level += 1
- return max_level
|