field.py 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168
  1. from typing import List
  2. from redis import DataError
  3. class Field:
  4. NUMERIC = "NUMERIC"
  5. TEXT = "TEXT"
  6. WEIGHT = "WEIGHT"
  7. GEO = "GEO"
  8. TAG = "TAG"
  9. VECTOR = "VECTOR"
  10. SORTABLE = "SORTABLE"
  11. NOINDEX = "NOINDEX"
  12. AS = "AS"
  13. def __init__(
  14. self,
  15. name: str,
  16. args: List[str] = None,
  17. sortable: bool = False,
  18. no_index: bool = False,
  19. as_name: str = None,
  20. ):
  21. if args is None:
  22. args = []
  23. self.name = name
  24. self.args = args
  25. self.args_suffix = list()
  26. self.as_name = as_name
  27. if sortable:
  28. self.args_suffix.append(Field.SORTABLE)
  29. if no_index:
  30. self.args_suffix.append(Field.NOINDEX)
  31. if no_index and not sortable:
  32. raise ValueError("Non-Sortable non-Indexable fields are ignored")
  33. def append_arg(self, value):
  34. self.args.append(value)
  35. def redis_args(self):
  36. args = [self.name]
  37. if self.as_name:
  38. args += [self.AS, self.as_name]
  39. args += self.args
  40. args += self.args_suffix
  41. return args
  42. class TextField(Field):
  43. """
  44. TextField is used to define a text field in a schema definition
  45. """
  46. NOSTEM = "NOSTEM"
  47. PHONETIC = "PHONETIC"
  48. def __init__(
  49. self,
  50. name: str,
  51. weight: float = 1.0,
  52. no_stem: bool = False,
  53. phonetic_matcher: str = None,
  54. withsuffixtrie: bool = False,
  55. **kwargs,
  56. ):
  57. Field.__init__(self, name, args=[Field.TEXT, Field.WEIGHT, weight], **kwargs)
  58. if no_stem:
  59. Field.append_arg(self, self.NOSTEM)
  60. if phonetic_matcher and phonetic_matcher in [
  61. "dm:en",
  62. "dm:fr",
  63. "dm:pt",
  64. "dm:es",
  65. ]:
  66. Field.append_arg(self, self.PHONETIC)
  67. Field.append_arg(self, phonetic_matcher)
  68. if withsuffixtrie:
  69. Field.append_arg(self, "WITHSUFFIXTRIE")
  70. class NumericField(Field):
  71. """
  72. NumericField is used to define a numeric field in a schema definition
  73. """
  74. def __init__(self, name: str, **kwargs):
  75. Field.__init__(self, name, args=[Field.NUMERIC], **kwargs)
  76. class GeoField(Field):
  77. """
  78. GeoField is used to define a geo-indexing field in a schema definition
  79. """
  80. def __init__(self, name: str, **kwargs):
  81. Field.__init__(self, name, args=[Field.GEO], **kwargs)
  82. class TagField(Field):
  83. """
  84. TagField is a tag-indexing field with simpler compression and tokenization.
  85. See http://redisearch.io/Tags/
  86. """
  87. SEPARATOR = "SEPARATOR"
  88. CASESENSITIVE = "CASESENSITIVE"
  89. def __init__(
  90. self,
  91. name: str,
  92. separator: str = ",",
  93. case_sensitive: bool = False,
  94. withsuffixtrie: bool = False,
  95. **kwargs,
  96. ):
  97. args = [Field.TAG, self.SEPARATOR, separator]
  98. if case_sensitive:
  99. args.append(self.CASESENSITIVE)
  100. if withsuffixtrie:
  101. args.append("WITHSUFFIXTRIE")
  102. Field.__init__(self, name, args=args, **kwargs)
  103. class VectorField(Field):
  104. """
  105. Allows vector similarity queries against the value in this attribute.
  106. See https://oss.redis.com/redisearch/Vectors/#vector_fields.
  107. """
  108. def __init__(self, name: str, algorithm: str, attributes: dict, **kwargs):
  109. """
  110. Create Vector Field. Notice that Vector cannot have sortable or no_index tag,
  111. although it's also a Field.
  112. ``name`` is the name of the field.
  113. ``algorithm`` can be "FLAT" or "HNSW".
  114. ``attributes`` each algorithm can have specific attributes. Some of them
  115. are mandatory and some of them are optional. See
  116. https://oss.redis.com/redisearch/master/Vectors/#specific_creation_attributes_per_algorithm
  117. for more information.
  118. """
  119. sort = kwargs.get("sortable", False)
  120. noindex = kwargs.get("no_index", False)
  121. if sort or noindex:
  122. raise DataError("Cannot set 'sortable' or 'no_index' in Vector fields.")
  123. if algorithm.upper() not in ["FLAT", "HNSW"]:
  124. raise DataError(
  125. "Realtime vector indexing supporting 2 Indexing Methods:"
  126. "'FLAT' and 'HNSW'."
  127. )
  128. attr_li = []
  129. for key, value in attributes.items():
  130. attr_li.extend([key, value])
  131. Field.__init__(
  132. self, name, args=[Field.VECTOR, algorithm, len(attr_li), *attr_li], **kwargs
  133. )