heisenberg_data_processing/embedded/
mod.rs

1use std::{path::PathBuf, sync::LazyLock};
2
3use polars::prelude::*;
4use serde::{Deserialize, Serialize};
5use tracing::info;
6
7#[cfg(any(test, feature = "test-data"))]
8use crate::raw::test_data::create_test_data;
9use crate::{
10    DataSource, Result,
11    processed::{generate_processed_data, save_processed_data_to_parquet},
12    raw::fetch::TempData,
13};
14
15// Sweaty workaround to let us use the same string as a static here and also in the include_bytes! macro
16#[macro_export]
17macro_rules! embedded_file_paths {
18    (admin) => {
19        "embedded_admin_search.parquet"
20    };
21    (place) => {
22        "embedded_place_search.parquet"
23    };
24    (metadata) => {
25        "embedded_data_metadata.json"
26    };
27}
28
29// Also provide constants for runtime use
30pub static ADMIN_DATA_PATH: &str = embedded_file_paths!(admin);
31pub static PLACE_DATA_PATH: &str = embedded_file_paths!(place);
32pub static METADATA_PATH: &str = embedded_file_paths!(metadata);
33
34static EMBEDDED_DIR_DEFAULT: &str = "src/data/embedded";
35
36pub static EMBEDDED_DIR: LazyLock<PathBuf> = LazyLock::new(|| {
37    std::env::var("EMBEDDED_DIR")
38        .map_or_else(|_| PathBuf::from(EMBEDDED_DIR_DEFAULT), PathBuf::from)
39});
40
41/// Generate embedded data from cities15000.zip and write as Rust source files
42pub fn generate_embedded_dataset(data_source: DataSource) -> Result<()> {
43    #[cfg(feature = "download-data")]
44    {
45        crate::raw::fetch::download_data(&data_source)
46            .and_then(|temp_data| embed_data(temp_data, data_source))?;
47        Ok(())
48    }
49    #[cfg(all(not(feature = "download-data"), any(test, feature = "test-data")))]
50    {
51        tracing::warn!("download_data feature not enabled, falling back to test data");
52        generate_test_data_rust_code()
53    }
54
55    #[cfg(all(not(feature = "download-data"), not(any(test, feature = "test-data"))))]
56    {
57        compile_error!(
58            "Either 'download_data' or 'test_data' feature must be enabled to generate embedded datasets"
59        )
60    }
61}
62
63/// Generate embedded data from test data and write as Rust source files
64#[cfg(any(test, feature = "test-data"))]
65pub fn generate_test_data_rust_code() -> Result<()> {
66    tracing::info!("Generating embedded dataset from test data");
67
68    let test_temp_files = create_test_data();
69    embed_data(test_temp_files, DataSource::TestData)
70}
71
72fn embed_data(temp_data: TempData, data_source: DataSource) -> Result<()> {
73    info!("Generating processed data for {}", data_source);
74
75    std::fs::create_dir_all(EMBEDDED_DIR.as_path())?;
76
77    let (admin_df, place_df) = generate_processed_data(temp_data)?;
78
79    let metadata = EmbeddedMetadata::from_dfs(&admin_df, &place_df, data_source)?;
80    let metadata_path = EMBEDDED_DIR.join(METADATA_PATH);
81    metadata.write_to_file(&metadata_path)?;
82
83    let admin_path = EMBEDDED_DIR.join(ADMIN_DATA_PATH);
84    let place_path = EMBEDDED_DIR.join(PLACE_DATA_PATH);
85
86    save_processed_data_to_parquet(admin_df.clone(), &admin_path)?;
87    save_processed_data_to_parquet(place_df.clone(), &place_path)?;
88    Ok(())
89}
90
91#[derive(Clone, Debug, Deserialize, Serialize)]
92pub struct DataFrameMetadata {
93    pub rows: usize,
94    pub size_bytes: usize,
95}
96
97#[derive(Clone, Debug, Deserialize, Serialize)]
98pub struct EmbeddedMetadata {
99    pub version: String,
100    pub source: DataSource,
101    pub generated_at: String,
102    pub admin_df: DataFrameMetadata,
103    pub place_df: DataFrameMetadata,
104}
105
106impl EmbeddedMetadata {
107    pub const VERSION: &str = env!("CARGO_PKG_VERSION");
108
109    #[must_use]
110    pub fn new(
111        source: DataSource,
112        admin_rows: usize,
113        place_rows: usize,
114        admin_size_bytes: usize,
115        place_size_bytes: usize,
116    ) -> Self {
117        let admin_df_metadata = DataFrameMetadata {
118            rows: admin_rows,
119            size_bytes: admin_size_bytes,
120        };
121        let place_df_metadata = DataFrameMetadata {
122            rows: place_rows,
123            size_bytes: place_size_bytes,
124        };
125        Self {
126            version: Self::VERSION.to_string(),
127            source,
128            generated_at: chrono::Utc::now().to_rfc3339(),
129            admin_df: admin_df_metadata,
130            place_df: place_df_metadata,
131        }
132    }
133
134    pub fn from_dfs(
135        admin_df: &DataFrame,
136        place_df: &DataFrame,
137        data_source: DataSource,
138    ) -> Result<Self> {
139        Ok(Self::new(
140            data_source,
141            admin_df.height(),
142            place_df.height(),
143            admin_df.estimated_size(),
144            place_df.estimated_size(),
145        ))
146    }
147
148    pub fn to_json(&self) -> Result<String> {
149        serde_json::to_string_pretty(self).map_err(Into::into)
150    }
151
152    pub fn write_to_file(&self, path: &std::path::Path) -> Result<()> {
153        let json = self.to_json()?;
154        std::fs::write(path, json).map_err(Into::into)
155    }
156
157    pub fn load_from_file(path: &std::path::Path) -> Result<Self> {
158        let content = std::fs::read_to_string(path)?;
159        serde_json::from_str(&content).map_err(Into::into)
160    }
161
162    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
163        serde_json::from_slice(bytes).map_err(Into::into)
164    }
165}
heisenberg_data_processing/embedded/mod.rs

heisenberg_data_processing/embedded/
mod.rs
