schema.py 1.4 KB

123456789101112131415161718192021222324252627282930313233343536373839
  1. """Utils for working with the parquet thrift models"""
  2. from ttypes import FieldRepetitionType
  3. class SchemaHelper(object):
  4. def __init__(self, schema_elements):
  5. self.schema_elements = schema_elements
  6. self.schema_elements_by_name = dict(
  7. [(se.name, se) for se in schema_elements])
  8. assert len(self.schema_elements) == len(self.schema_elements_by_name)
  9. def schema_element(self, name):
  10. """Get the schema element with the given name."""
  11. return self.schema_elements_by_name[name]
  12. def is_required(self, name):
  13. """Returns true iff the schema element with the given name is
  14. required"""
  15. return self.schema_element(name).repetition_type == FieldRepetitionType.REQUIRED
  16. def max_repetition_level(self, path):
  17. """get the max repetition level for the given schema path."""
  18. max_level = 0
  19. for part in path:
  20. se = self.schema_element(part)
  21. if se.repetition_type == FieldRepetitionType.REQUIRED:
  22. max_level += 1
  23. return max_level
  24. def max_definition_level(self, path):
  25. """get the max definition level for the given schema path."""
  26. max_level = 0
  27. for part in path:
  28. se = self.schema_element(part)
  29. if se.repetition_type != FieldRepetitionType.REQUIRED:
  30. max_level += 1
  31. return max_level