11from enum import Enum
2+ from typing import *
23
3- class _ConditionOp (Enum ):
4- eq = 1
5- notEq = 2
6- contains = 3
7- startsWith = 4
8- endsWith = 5
9- gt = 6
10- greaterOrEq = 7
11- lt = 8
12- lessOrEq = 9
13- between = 10
4+
5+ class _QueryConditionOp (Enum ):
6+ EQ = 1
7+ NOT_EQ = 2
8+ CONTAINS = 3
9+ STARTS_WITH = 4
10+ ENDS_WITH = 5
11+ GT = 6
12+ GTE = 7
13+ LT = 8
14+ LTE = 9
15+ BETWEEN = 10
16+ NEAREST_NEIGHBOR = 11
1417
1518
1619class QueryCondition :
17- def __init__ (self , property_id : int , op : _ConditionOp , value , value_b = None , case_sensitive : bool = True ):
20+ def __init__ (self , property_id : int , op : _QueryConditionOp , args : Dict [str , Any ]):
21+ if op not in self ._get_op_map ():
22+ raise Exception (f"Invalid query condition op with ID: { op } " )
23+
1824 self ._property_id = property_id
1925 self ._op = op
20- self ._value = value
21- self ._value_b = value_b
22- self ._case_sensitive = case_sensitive
23-
24- def apply (self , builder : 'QueryBuilder' ):
25- if self ._op == _ConditionOp .eq :
26- if isinstance (self ._value , str ):
27- builder .equals_string (self ._property_id , self ._value , self ._case_sensitive )
28- elif isinstance (self ._value , int ):
29- builder .equals_int (self ._property_id , self ._value )
30- else :
31- raise Exception ("Unsupported type for 'eq': " + str (type (self ._value )))
32-
33- elif self ._op == _ConditionOp .notEq :
34- if isinstance (self ._value , str ):
35- builder .not_equals_string (self ._property_id , self ._value , self ._case_sensitive )
36- elif isinstance (self ._value , int ):
37- builder .not_equals_int (self ._property_id , self ._value )
38- else :
39- raise Exception ("Unsupported type for 'notEq': " + str (type (self ._value )))
40-
41- elif self ._op == _ConditionOp .contains :
42- if isinstance (self ._value , str ):
43- builder .contains_string (self ._property_id , self ._value , self ._case_sensitive )
44- else :
45- raise Exception ("Unsupported type for 'contains': " + str (type (self ._value )))
46-
47- elif self ._op == _ConditionOp .startsWith :
48- if isinstance (self ._value , str ):
49- builder .starts_with_string (self ._property_id , self ._value , self ._case_sensitive )
50- else :
51- raise Exception ("Unsupported type for 'startsWith': " + str (type (self ._value )))
52-
53- elif self ._op == _ConditionOp .endsWith :
54- if isinstance (self ._value , str ):
55- builder .ends_with_string (self ._property_id , self ._value , self ._case_sensitive )
56- else :
57- raise Exception ("Unsupported type for 'endsWith': " + str (type (self ._value )))
58-
59- elif self ._op == _ConditionOp .gt :
60- if isinstance (self ._value , str ):
61- builder .greater_than_string (self ._property_id , self ._value , self ._case_sensitive )
62- elif isinstance (self ._value , int ):
63- builder .greater_than_int (self ._property_id , self ._value )
64- else :
65- raise Exception ("Unsupported type for 'gt': " + str (type (self ._value )))
66-
67- elif self ._op == _ConditionOp .greaterOrEq :
68- if isinstance (self ._value , str ):
69- builder .greater_or_equal_string (self ._property_id , self ._value , self ._case_sensitive )
70- elif isinstance (self ._value , int ):
71- builder .greater_or_equal_int (self ._property_id , self ._value )
72- else :
73- raise Exception ("Unsupported type for 'greaterOrEq': " + str (type (self ._value )))
74-
75- elif self ._op == _ConditionOp .lt :
76- if isinstance (self ._value , str ):
77- builder .less_than_string (self ._property_id , self ._value , self ._case_sensitive )
78- elif isinstance (self ._value , int ):
79- builder .less_than_int (self ._property_id , self ._value )
80- else :
81- raise Exception ("Unsupported type for 'lt': " + str (type (self ._value )))
82-
83- elif self ._op == _ConditionOp .lessOrEq :
84- if isinstance (self ._value , str ):
85- builder .less_or_equal_string (self ._property_id , self ._value , self ._case_sensitive )
86- elif isinstance (self ._value , int ):
87- builder .less_or_equal_int (self ._property_id , self ._value )
88- else :
89- raise Exception ("Unsupported type for 'lessOrEq': " + str (type (self ._value )))
90-
91- elif self ._op == _ConditionOp .between :
92- if isinstance (self ._value , int ):
93- builder .between_2ints (self ._property_id , self ._value , self ._value_b )
94- else :
95- raise Exception ("Unsupported type for 'between': " + str (type (self ._value )))
26+ self ._args = args
27+
28+ def _get_op_map (self ):
29+ return {
30+ _QueryConditionOp .EQ : self ._apply_eq ,
31+ _QueryConditionOp .NOT_EQ : self ._apply_not_eq ,
32+ _QueryConditionOp .CONTAINS : self ._apply_contains ,
33+ _QueryConditionOp .STARTS_WITH : self ._apply_starts_with ,
34+ _QueryConditionOp .ENDS_WITH : self ._apply_ends_with ,
35+ _QueryConditionOp .GT : self ._apply_gt ,
36+ _QueryConditionOp .GTE : self ._apply_gte ,
37+ _QueryConditionOp .LT : self ._apply_lt ,
38+ _QueryConditionOp .LTE : self ._apply_lte ,
39+ _QueryConditionOp .BETWEEN : self ._apply_between ,
40+ _QueryConditionOp .NEAREST_NEIGHBOR : self ._apply_nearest_neighbor
41+ # ... new query condition here ... :)
42+ }
43+
44+ def _apply_eq (self , qb : 'QueryBuilder' ):
45+ value = self ._args ['value' ]
46+ case_sensitive = self ._args ['case_sensitive' ]
47+ if isinstance (value , str ):
48+ qb .equals_string (self ._property_id , value , case_sensitive )
49+ elif isinstance (value , int ):
50+ qb .equals_int (self ._property_id , value )
51+ else :
52+ raise Exception (f"Unsupported type for 'EQ': { type (value )} " )
53+
54+ def _apply_not_eq (self , qb : 'QueryBuilder' ):
55+ value = self ._args ['value' ]
56+ case_sensitive = self ._args ['case_sensitive' ]
57+ if isinstance (value , str ):
58+ qb .not_equals_string (self ._property_id , value , case_sensitive )
59+ elif isinstance (value , int ):
60+ qb .not_equals_int (self ._property_id , value )
61+ else :
62+ raise Exception (f"Unsupported type for 'NOT_EQ': { type (value )} " )
63+
64+ def _apply_contains (self , qb : 'QueryBuilder' ):
65+ value = self ._args ['value' ]
66+ case_sensitive = self ._args ['case_sensitive' ]
67+ if isinstance (value , str ):
68+ qb .contains_string (self ._property_id , value , case_sensitive )
69+ else :
70+ raise Exception (f"Unsupported type for 'CONTAINS': { type (self_value )} " )
71+
72+ def _apply_starts_with (self , qb : 'QueryBuilder' ):
73+ value = self ._args ['value' ]
74+ case_sensitive = self ._args ['case_sensitive' ]
75+ if isinstance (value , str ):
76+ qb .starts_with_string (self ._property_id , value , case_sensitive )
77+ else :
78+ raise Exception (f"Unsupported type for 'STARTS_WITH': { type (value )} " )
79+
80+ def _apply_ends_with (self , qb : 'QueryBuilder' ):
81+ value = self ._args ['value' ]
82+ case_sensitive = self ._args ['case_sensitive' ]
83+ if isinstance (value , str ):
84+ qb .ends_with_string (self ._property_id , value , case_sensitive )
85+ else :
86+ raise Exception (f"Unsupported type for 'ENDS_WITH': { type (value )} " )
87+
88+ def _apply_gt (self , qb : 'QueryBuilder' ):
89+ value = self ._args ['value' ]
90+ case_sensitive = self ._args ['case_sensitive' ]
91+ if isinstance (value , str ):
92+ qb .greater_than_string (self ._property_id , value , case_sensitive )
93+ elif isinstance (value , int ):
94+ qb .greater_than_int (self ._property_id , value )
95+ else :
96+ raise Exception (f"Unsupported type for 'GT': { type (value )} " )
97+
98+ def _apply_gte (self , qb : 'QueryBuilder' ):
99+ value = self ._args ['value' ]
100+ case_sensitive = self ._args ['case_sensitive' ]
101+ if isinstance (value , str ):
102+ qb .greater_or_equal_string (self ._property_id , value , case_sensitive )
103+ elif isinstance (value , int ):
104+ qb .greater_or_equal_int (self ._property_id , value )
105+ else :
106+ raise Exception (f"Unsupported type for 'GTE': { type (value )} " )
107+
108+ def _apply_lt (self , qb : 'QueryCondition' ):
109+ value = self ._args ['value' ]
110+ case_sensitive = self ._args ['case_sensitive' ]
111+ if isinstance (value , str ):
112+ qb .less_than_string (self ._property_id , value , case_sensitive )
113+ elif isinstance (value , int ):
114+ qb .less_than_int (self ._property_id , value )
115+ else :
116+ raise Exception ("Unsupported type for 'LT': " + str (type (value )))
117+
118+ def _apply_lte (self , qb : 'QueryBuilder' ):
119+ value = self ._args ['value' ]
120+ case_sensitive = self ._args ['case_sensitive' ]
121+ if isinstance (value , str ):
122+ qb .less_or_equal_string (self ._property_id , value , case_sensitive )
123+ elif isinstance (value , int ):
124+ qb .less_or_equal_int (self ._property_id , value )
125+ else :
126+ raise Exception (f"Unsupported type for 'LTE': { type (value )} " )
127+
128+ def _apply_between (self , qb : 'QueryBuilder' ):
129+ a = self ._args ['a' ]
130+ b = self ._args ['b' ]
131+ if isinstance (a , int ):
132+ qb .between_2ints (self ._property_id , a , b )
133+ else :
134+ raise Exception (f"Unsupported type for 'BETWEEN': { type (a )} " )
135+
136+ def _apply_nearest_neighbor (self , qb : 'QueryCondition' ):
137+ query_vector = self ._args ['query_vector' ]
138+ element_count = self ._args ['element_count' ]
139+
140+ if len (query_vector ) == 0 :
141+ raise Exception ("query_vector can't be empty" )
142+
143+ is_float_vector = False
144+ is_float_vector |= isinstance (query_vector , np .ndarray ) and query_vector .dtype == np .float32
145+ is_float_vector |= isinstance (query_vector , list ) and type (query_vector [0 ]) == float
146+ if is_float_vector :
147+ qb .nearest_neighbors_f32 (self ._property_id , query_vector , element_count )
148+ else :
149+ raise Exception (f"Unsupported type for 'NEAREST_NEIGHBOR': { type (query_vector )} " )
150+
151+ def apply (self , qb : 'QueryBuilder' ):
152+ self ._get_op_map ()[self ._op ](qb )
0 commit comments