Write json to custom path

nautechsystems · Dec 22, 2024 · acf326a · acf326a
1 parent 9a2ed76
commit acf326a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/nautilus_core/persistence/src/backend/catalog.rs b/nautilus_core/persistence/src/backend/catalog.rs
@@ -88,15 +88,17 @@ impl ParquetDataCatalog {
     }
 
     #[must_use]
-    pub fn write_to_json<T>(&self, data: Vec<T>) -> PathBuf
+    pub fn write_to_json<T>(&self, data: Vec<T>, path: Option<PathBuf>) -> PathBuf
     where
         T: GetTsInit + Serialize,
     {
         let type_name = std::any::type_name::<T>().to_snake_case();
         Self::check_ascending_timestamps(&data, &type_name);
 
-        let path = self.make_path(&type_name, None);
-        let json_path = path.with_extension("json");
+        let json_path = path.unwrap_or_else(|| {
+            let path = self.make_path(&type_name, None);
+            path.with_extension("json")
+        });
 
         info!(
             "Writing {} records of {} data to {:?}",

diff --git a/nautilus_core/persistence/tests/test_catalog.rs b/nautilus_core/persistence/tests/test_catalog.rs
@@ -345,7 +345,7 @@ fn test_catalog_serialization_json_round_trip() {
     let quote_ticks: Vec<QuoteTick> = to_variant(quote_ticks);
 
     // Write to JSON using catalog
-    let json_path = catalog.write_to_json(quote_ticks.clone());
+    let json_path = catalog.write_to_json(quote_ticks.clone(), None);
 
     // Read back from JSON
     let json_str = std::fs::read_to_string(json_path).unwrap();
@@ -393,7 +393,7 @@ fn test_datafusion_parquet_round_trip() {
     let mut temp_file = std::fs::File::create(&temp_file_path).unwrap();
     {
         let writer_props = WriterProperties::builder()
-            .set_compression(Compression::ZSTD(ZstdLevel::default()))
+            .set_compression(Compression::SNAPPY)
             .set_max_row_group_size(1000)
             .build();