Set up the Arrow writer for streaming writes (only on main process)
(self)
| 313 | return base_path |
| 314 | |
| 315 | def _setup_arrow_writer(self): |
| 316 | """Set up the Arrow writer for streaming writes (only on main process)""" |
| 317 | if self.process_index == 0: |
| 318 | logger.info(f"Setting up arrow writer...") |
| 319 | # Define schema for the Arrow file |
| 320 | fields = [ |
| 321 | pa.field('keys', pa.list_(pa.float16(), self.dimension)), |
| 322 | pa.field('vals', pa.int32()) |
| 323 | ] |
| 324 | schema = pa.schema(fields) |
| 325 | |
| 326 | # Create the Arrow file and writer using streaming format |
| 327 | self.arrow_file = pa.OSFile(self.arrow_file_path, 'wb') |
| 328 | self.arrow_writer = pa.ipc.new_stream(self.arrow_file, schema) |
| 329 | |
| 330 | def _save_step_data(self, keys, vals): |
| 331 | # Detach from computation graph and ensure contiguous memory |