@@ -22,16 +22,36 @@ std::tuple<uint32_t, uint32_t, size_t> InMemStaticGraphReformatStore::load_impl(
2222 // get file size
2323 size_t check_file_size = file_reader.GetFileSize ();
2424
25- file_reader.Read (file_offset, sizeof (size_t ), (char *)&expected_file_size);
26- file_offset += sizeof (size_t );
27- file_reader.Read (file_offset, sizeof (uint32_t ), (char *)&_max_observed_degree);
28- file_offset += sizeof (uint32_t );
29- file_reader.Read (file_offset, sizeof (uint32_t ), (char *)&start);
30- file_offset += sizeof (uint32_t );
31- file_reader.Read (file_offset, sizeof (size_t ), (char *)&file_frozen_pts);
32- file_offset += sizeof (size_t );
33- file_reader.Read (file_offset, sizeof (size_t ), (char *)&num_points);
34- file_offset += sizeof (size_t );
25+ // Create fixed-size aligned buffer for 100,000 sectors
26+ const size_t SECTOR_SIZE = 4096 ;
27+ const size_t BUFFER_SECTORS = 100000 ;
28+ const size_t BUFFER_SIZE = BUFFER_SECTORS * SECTOR_SIZE; // ~409.6 MB
29+
30+ char * aligned_buffer = nullptr ;
31+ alloc_aligned ((void **)&aligned_buffer, BUFFER_SIZE, SECTOR_SIZE);
32+
33+ try {
34+ // Read metadata header (aligned to sector boundary)
35+ size_t metadata_size = sizeof (size_t ) + sizeof (uint32_t ) + sizeof (uint32_t ) + sizeof (size_t ) + sizeof (size_t );
36+ size_t aligned_metadata_size = ROUND_UP (metadata_size, SECTOR_SIZE);
37+
38+ if (!file_reader.Read (0 , aligned_metadata_size, aligned_buffer)) {
39+ throw diskann::ANNException (" Failed to read graph metadata" 1 , __FUNCSIG__, __FILE__, __LINE__);
40+ }
41+
42+ // Extract metadata from aligned buffer
43+ char * metadata_ptr = aligned_buffer;
44+ expected_file_size = *reinterpret_cast <size_t *>(metadata_ptr);
45+ metadata_ptr += sizeof (size_t );
46+ _max_observed_degree = *reinterpret_cast <uint32_t *>(metadata_ptr);
47+ metadata_ptr += sizeof (uint32_t );
48+ start = *reinterpret_cast <uint32_t *>(metadata_ptr);
49+ metadata_ptr += sizeof (uint32_t );
50+ file_frozen_pts = *reinterpret_cast <size_t *>(metadata_ptr);
51+ metadata_ptr += sizeof (size_t );
52+ num_points = *reinterpret_cast <size_t *>(metadata_ptr);
53+
54+ file_offset = metadata_size;
3555
3656 // max observed degree is the max degree of the graph
3757 _max_range_of_graph = _max_observed_degree;
@@ -53,20 +73,105 @@ std::tuple<uint32_t, uint32_t, size_t> InMemStaticGraphReformatStore::load_impl(
5373
5474 diskann::cout << " Loading vamana graph " " ..."
5575
76+ // Load node index with chunked aligned reads
5677 _node_index.resize (num_points + 1 );
57- file_reader.Read (file_offset, _node_index.size () * sizeof (size_t ), (char *)_node_index.data ());
58- file_offset += _node_index.size () * sizeof (size_t );
78+ size_t node_index_size = _node_index.size () * sizeof (size_t );
79+
80+ // Calculate aligned offset for node index data
81+ size_t aligned_node_offset = ROUND_DOWN (file_offset, SECTOR_SIZE);
82+ size_t node_offset_adjustment = file_offset - aligned_node_offset;
83+ size_t total_node_bytes_to_read = node_index_size + node_offset_adjustment;
84+
85+ // Read node index in chunks
86+ size_t node_bytes_read = 0 ;
87+ size_t node_elements_copied = 0 ;
88+
89+ while (node_bytes_read < total_node_bytes_to_read) {
90+ size_t chunk_size = std::min (BUFFER_SIZE, ROUND_UP (total_node_bytes_to_read - node_bytes_read, SECTOR_SIZE));
91+ size_t current_offset = aligned_node_offset + node_bytes_read;
92+
93+ if (!file_reader.Read (current_offset, chunk_size, aligned_buffer)) {
94+ throw diskann::ANNException (" Failed to read node index chunk" 1 , __FUNCSIG__, __FILE__, __LINE__);
95+ }
96+
97+ const char * chunk_data_start = aligned_buffer;
98+ if (node_bytes_read == 0 ) {
99+ // First chunk - account for offset
100+ chunk_data_start += node_offset_adjustment;
101+ }
102+
103+ size_t available_data = chunk_size;
104+ if (node_bytes_read == 0 ) {
105+ available_data -= node_offset_adjustment;
106+ }
107+
108+ size_t elements_in_chunk = std::min (available_data / sizeof (size_t ), _node_index.size () - node_elements_copied);
109+ std::memcpy (_node_index.data () + node_elements_copied, chunk_data_start, elements_in_chunk * sizeof (size_t ));
110+
111+ node_elements_copied += elements_in_chunk;
112+ node_bytes_read += (elements_in_chunk * sizeof (size_t ));
113+ if (node_bytes_read == 0 && node_offset_adjustment > 0 ) {
114+ node_bytes_read = node_offset_adjustment;
115+ }
116+ }
117+
118+ file_offset += node_index_size;
59119
60120 _graph_size = _node_index[num_points] * sizeof (std::uint32_t );
61121
62122 size_t total_neighbors = _node_index[num_points];
63123 // add one more slot than actually need to avoid read invaild address
64124 // while the last point is no neighbor
65125 _graph.resize (total_neighbors + 1 );
66- file_reader.Read (file_offset, _graph_size, (char *)_graph.data ());
67-
126+
127+ // Load graph data with chunked aligned reads
128+ size_t aligned_graph_offset = ROUND_DOWN (file_offset, SECTOR_SIZE);
129+ size_t graph_offset_adjustment = file_offset - aligned_graph_offset;
130+ size_t total_graph_bytes_to_read = _graph_size + graph_offset_adjustment;
131+
132+ // Read graph data in chunks
133+ size_t graph_bytes_read = 0 ;
134+ size_t graph_elements_copied = 0 ;
135+
136+ while (graph_bytes_read < total_graph_bytes_to_read) {
137+ size_t chunk_size = std::min (BUFFER_SIZE, ROUND_UP (total_graph_bytes_to_read - graph_bytes_read, SECTOR_SIZE));
138+ size_t current_offset = aligned_graph_offset + graph_bytes_read;
139+
140+ if (!file_reader.Read (current_offset, chunk_size, aligned_buffer)) {
141+ throw diskann::ANNException (" Failed to read graph data chunk" 1 , __FUNCSIG__, __FILE__, __LINE__);
142+ }
143+
144+ const char * chunk_data_start = aligned_buffer;
145+ if (graph_bytes_read == 0 ) {
146+ // First chunk - account for offset
147+ chunk_data_start += graph_offset_adjustment;
148+ }
149+
150+ size_t available_data = chunk_size;
151+ if (graph_bytes_read == 0 ) {
152+ available_data -= graph_offset_adjustment;
153+ }
154+
155+ size_t elements_in_chunk = std::min (available_data / sizeof (std::uint32_t ), _graph.size () - graph_elements_copied);
156+ std::memcpy (_graph.data () + graph_elements_copied, chunk_data_start, elements_in_chunk * sizeof (std::uint32_t ));
157+
158+ graph_elements_copied += elements_in_chunk;
159+ graph_bytes_read += (elements_in_chunk * sizeof (std::uint32_t ));
160+ if (graph_bytes_read == 0 && graph_offset_adjustment > 0 ) {
161+ graph_bytes_read = graph_offset_adjustment;
162+ }
163+ }
164+
165+ aligned_free (aligned_buffer);
166+
68167 diskann::cout << " done. Index has " " nodes and " " out-edges, _start is set to "
69168 << std::endl;
169+ } catch (...) {
170+ aligned_free (aligned_buffer);
171+ throw ;
172+ }
173+
174+
70175 return std::make_tuple (static_cast <std::uint32_t >(num_points), start, file_frozen_pts);
71176}
72177
0 commit comments