1+ /*
2+ * Copyright 2025 PixelsDB.
3+ *
4+ * Licensed under the Apache License, Version 2.0 (the "License");
5+ * you may not use this file except in compliance with the License.
6+ * You may obtain a copy of the License at
7+ *
8+ * http://www.apache.org/licenses/LICENSE-2.0
9+ *
10+ * Unless required by applicable law or agreed to in writing, software
11+ * distributed under the License is distributed on an "AS IS" BASIS,
12+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+ * See the License for the specific language governing permissions and
14+ * limitations under the License.
15+ *
16+ */
17+
18+ package io .pixelsdb .pixels .daemon .index ;
19+
20+ import com .google .protobuf .ByteString ;
21+ import io .pixelsdb .pixels .common .index .IndexOption ;
22+ import io .pixelsdb .pixels .common .index .service .LocalIndexService ;
23+ import io .pixelsdb .pixels .common .utils .ConfigFactory ;
24+ import io .pixelsdb .pixels .index .IndexProto ;
25+ import org .junit .jupiter .api .*;
26+
27+ import java .util .Collections ;
28+ import java .util .List ;
29+
30+ import static org .junit .jupiter .api .Assertions .*;
31+ import static org .junit .jupiter .api .Assertions .assertNotNull ;
32+ import static org .junit .jupiter .api .Assertions .assertNull ;
33+
34+ public class TestIndexUpsert
35+ {
36+ private static LocalIndexService indexService ;
37+ private static final long TABLE_ID = 3622L ;
38+ private static final long PRIMARY_INDEX_ID = 3560L ;
39+ private static IndexOption indexOption ;
40+ private static IndexProto .PrimaryIndexEntry missingEntry ;
41+
42+ @ BeforeAll
43+ static void setup () throws Exception
44+ {
45+ ConfigFactory .Instance ().addProperty ("retina.upsert-mode.enabled" , "true" );
46+ indexService = LocalIndexService .Instance ();
47+ indexOption = IndexOption .builder ().vNodeId (0 ).build ();
48+ // 1. Enable Upsert Mode (Ensure your config utility supports this)
49+ // PixelsConfig.set("retina.index.upsert-mode.enabled", "true");
50+
51+ indexService .openIndex (TABLE_ID , PRIMARY_INDEX_ID , true , indexOption );
52+
53+ // Define an entry that is NOT in the database yet
54+ missingEntry = IndexProto .PrimaryIndexEntry .newBuilder ()
55+ .setRowId (9999L )
56+ .setIndexKey (IndexProto .IndexKey .newBuilder ()
57+ .setTableId (TABLE_ID )
58+ .setIndexId (PRIMARY_INDEX_ID )
59+ .setKey (ByteString .copyFromUtf8 ("ghost_key" ))
60+ .setTimestamp (System .currentTimeMillis ()))
61+ .setRowLocation (IndexProto .RowLocation .newBuilder ()
62+ .setFileId (10 )
63+ .setRgId (1 )
64+ .setRgRowOffset (500 ))
65+ .build ();
66+ }
67+
68+ @ Test
69+ @ Order (1 )
70+ @ DisplayName ("Test UPDATE on missing key (Should Insert)" )
71+ void testUpdateMissingKey () throws Exception
72+ {
73+ // In standard mode, this would throw IndexException.
74+ // In Upsert mode, it returns null and inserts the data.
75+ IndexProto .RowLocation prevLocation = indexService .updatePrimaryIndexEntry (missingEntry , indexOption );
76+
77+ // Assertions
78+ assertNull (prevLocation , "Upsert should return null when no previous entry exists" );
79+
80+ // Verify the data was actually inserted
81+ IndexProto .RowLocation currentLoc = indexService .lookupUniqueIndex (missingEntry .getIndexKey (), indexOption );
82+ assertNotNull (currentLoc , "Entry should have been inserted by the update call" );
83+ assertEquals (10 , currentLoc .getFileId ());
84+ }
85+
86+ @ Test
87+ @ Order (2 )
88+ @ DisplayName ("Test DELETE on missing key (Should Ignore)" )
89+ void testDeleteMissingKey () throws Exception
90+ {
91+ // Create a key that definitely doesn't exist
92+ IndexProto .IndexKey nonExistentKey = missingEntry .getIndexKey ().toBuilder ()
93+ .setKey (ByteString .copyFromUtf8 ("never_existed" ))
94+ .build ();
95+
96+ // In Upsert mode, this should NOT throw exception
97+ IndexProto .RowLocation deletedLocation = indexService .deletePrimaryIndexEntry (nonExistentKey , indexOption );
98+
99+ assertNull (deletedLocation , "Delete on missing key should return null in upsert mode" );
100+ }
101+
102+ @ Test
103+ @ Order (3 )
104+ @ DisplayName ("Test Batch UPDATE with missing keys" )
105+ void testBatchUpdateUpsert () throws Exception
106+ {
107+ List <IndexProto .PrimaryIndexEntry > entries = Collections .singletonList (
108+ missingEntry .toBuilder ()
109+ .setRowId (8888L )
110+ .setIndexKey (missingEntry .getIndexKey ().toBuilder ().setKey (ByteString .copyFromUtf8 ("batch_ghost" )))
111+ .build ()
112+ );
113+
114+ // Should execute without throwing exception
115+ List <IndexProto .RowLocation > prevLocations = indexService .updatePrimaryIndexEntries (TABLE_ID , PRIMARY_INDEX_ID , entries , indexOption );
116+
117+ // Even if some or all were missing, the returned list can be empty or only contain found locations
118+ assertNotNull (prevLocations );
119+ }
120+
121+ @ AfterAll
122+ static void tearDown () throws Exception
123+ {
124+ indexService .removeIndex (TABLE_ID , PRIMARY_INDEX_ID , true , indexOption );
125+ // PixelsConfig.set("retina.index.upsert-mode.enabled", "false");
126+ }
127+ }
0 commit comments