schema.py 1.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940
  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([(se.name, se) for se in schema_elements])
  7. assert len(self.schema_elements) == len(self.schema_elements_by_name)
  8. def schema_element(self, name):
  9. """Get the schema element with the given name."""
  10. return self.schema_elements_by_name[name]
  11. def is_required(self, name):
  12. """Returns true iff the schema element with the given name is required"""
  13. return self.schema_element(name).repetition_type == FieldRepetitionType.REQUIRED
  14. def max_repetition_level(self, path):
  15. """get the max repetition level for the given schema path."""
  16. max_level = 0
  17. for part in path:
  18. se = self.schema_element(part)
  19. if se.repetition_type == FieldRepetitionType.REQUIRED:
  20. max_level += 1
  21. return max_level
  22. def max_definition_level(self, path):
  23. """get the max definition level for the given schema path."""
  24. max_level = 0
  25. for part in path:
  26. se = self.schema_element(part)
  27. if se.repetition_type != FieldRepetitionType.REQUIRED:
  28. max_level += 1
  29. return max_level