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