geographdb-core 0.5.4

//! HuggingFace dataset loader for Rust-native corpus ingestion.
//!
//! Downloads shards from the HuggingFace Hub (Parquet, JSON, or JSONL) and
//! streams text rows directly into the graph builder pipeline — no Python
//! export step.

use anyhow::{Context, Result};
use arrow::array::{Array, StringArray};
use arrow::record_batch::RecordBatch;
use hf_hub::api::sync::ApiRepo;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use std::path::{Path, PathBuf};

/// Description of a HuggingFace dataset to ingest.
#[derive(Debug, Clone)]
pub struct HfDatasetSpec {
    /// Domain label used for node properties (e.g. "code", "math", "tool", "text").
    pub domain: String,
    /// HuggingFace repo id, e.g. `bigcode/the-stack-smol`.
    pub repo_id: String,
    /// Optional subdirectory / subset prefix inside the repo, e.g. `data/python`.
    pub subset: Option<String>,
    /// Column names to extract as text. If empty, the loader attempts domain defaults.
    pub text_columns: Vec<String>,
}

impl HfDatasetSpec {
    /// Parse `--dataset domain=repo_id:subset|col1,col2` style argument.
    ///
    /// The optional `|columns` suffix overrides the domain defaults.
    ///
    /// Examples:
    /// - `code=bigcode/the-stack-smol:data/python`
    /// - `math=meta-math/MetaMathQA`
    /// - `code=nickrosh/Evol-Instruct-Code-80k-v1|instruction,output`
    pub fn from_arg(arg: &str) -> Result<Self> {
        let (domain, rest) = arg.split_once('=').with_context(|| {
            format!("dataset arg must be domain=repo[:subset][|columns], got: {arg}")
        })?;

        let (repo_and_subset, column_override) = rest
            .split_once('|')
            .map_or((rest, None), |(r, c)| (r, Some(c)));
        let (repo_id, subset) = match repo_and_subset.split_once(':') {
            Some((repo, sub)) => (repo, Some(sub.to_string())),
            None => (repo_and_subset, None),
        };

        let text_columns = column_override
            .map(|s| {
                s.split(',')
                    .map(|c| c.trim().to_string())
                    .filter(|c| !c.is_empty())
                    .collect()
            })
            .unwrap_or_else(|| default_columns(domain));

        Ok(Self {
            domain: domain.to_string(),
            repo_id: repo_id.to_string(),
            subset,
            text_columns,
        })
    }
}

fn default_columns(domain: &str) -> Vec<String> {
    match domain {
        "code" => vec!["content".to_string()],
        "math" => vec!["query".to_string(), "response".to_string()],
        "tool" => vec!["text".to_string()],
        "text" => vec!["text".to_string()],
        _ => Vec::new(),
    }
}

/// Loader for a single HuggingFace dataset.
pub struct HfDatasetLoader {
    spec: HfDatasetSpec,
    repo: ApiRepo,
}

impl HfDatasetLoader {
    /// Create a loader from a parsed spec.
    pub fn new(spec: HfDatasetSpec) -> Result<Self> {
        let api = hf_hub::api::sync::Api::new().context("failed to initialize HuggingFace API")?;
        let repo = api.dataset(spec.repo_id.clone());
        Ok(Self { spec, repo })
    }

    /// Return the spec.
    pub fn spec(&self) -> &HfDatasetSpec {
        &self.spec
    }

    fn is_shard_filename(name: &str) -> bool {
        name.ends_with(".parquet") || name.ends_with(".json") || name.ends_with(".jsonl")
    }

    /// List shard filenames (Parquet, JSON, or JSONL) matching the optional subset prefix.
    pub fn list_shards(&self) -> Result<Vec<String>> {
        let info = self.repo.info().context("failed to fetch dataset info")?;
        let prefix = self.spec.subset.as_deref().unwrap_or("");

        let mut shards: Vec<String> = info
            .siblings
            .into_iter()
            .map(|s| s.rfilename)
            .filter(|f| Self::is_shard_filename(f))
            .filter(|f| prefix.is_empty() || f.starts_with(prefix))
            .collect();

        // Stable, deterministic order so repeated runs behave consistently.
        shards.sort();

        if shards.is_empty() {
            anyhow::bail!(
                "no parquet/json/jsonl shards found for {}{}",
                self.spec.repo_id,
                self.spec
                    .subset
                    .as_ref()
                    .map(|s| format!("::{s}"))
                    .unwrap_or_default()
            );
        }

        Ok(shards)
    }

    /// Download all matching shards into the HuggingFace cache and return local paths.
    pub fn download_shards(&self) -> Result<Vec<PathBuf>> {
        let shards = self.list_shards()?;
        let mut paths = Vec::with_capacity(shards.len());
        for filename in shards {
            let path = self
                .repo
                .get(&filename)
                .with_context(|| format!("failed to download shard {filename}"))?;
            paths.push(path);
        }
        Ok(paths)
    }

    /// Read text strings from a single shard, dispatching on file extension.
    pub fn read_shard_texts(&self, path: &Path) -> Result<Vec<String>> {
        if self.spec.text_columns.is_empty() {
            anyhow::bail!("no text columns configured for domain {}", self.spec.domain);
        }

        let ext = path
            .extension()
            .and_then(|e| e.to_str())
            .unwrap_or("")
            .to_ascii_lowercase();

        match ext.as_str() {
            "parquet" => self.read_parquet_texts(path),
            "json" | "jsonl" => self.read_json_texts(path, ext == "jsonl"),
            other => anyhow::bail!(
                "unsupported shard extension '{other}' for {}",
                path.display()
            ),
        }
    }

    fn read_parquet_texts(&self, path: &Path) -> Result<Vec<String>> {
        let file = std::fs::File::open(path)
            .with_context(|| format!("failed to open shard {}", path.display()))?;
        let reader = ParquetRecordBatchReaderBuilder::try_new(file)
            .context("failed to create parquet reader")?
            .build()
            .context("failed to build parquet batch reader")?;

        let mut texts = Vec::new();
        for batch in reader {
            let batch = batch.context("failed to read parquet batch")?;
            self.extract_texts_from_batch(&batch, &mut texts)?;
        }
        Ok(texts)
    }

    fn read_json_texts(&self, path: &Path, jsonl: bool) -> Result<Vec<String>> {
        let bytes = std::fs::read(path)
            .with_context(|| format!("failed to read json shard {}", path.display()))?;

        let mut texts = Vec::new();
        if jsonl {
            for (line_no, line) in String::from_utf8_lossy(&bytes).lines().enumerate() {
                if line.trim().is_empty() {
                    continue;
                }
                let value: serde_json::Value = serde_json::from_str(line)
                    .with_context(|| format!("invalid JSON on line {}", line_no + 1))?;
                self.extract_texts_from_json_value(&value, &mut texts)?;
            }
        } else {
            let value: serde_json::Value = serde_json::from_slice(&bytes)
                .with_context(|| format!("invalid json shard {}", path.display()))?;
            if let Some(arr) = value.as_array() {
                for v in arr {
                    self.extract_texts_from_json_value(v, &mut texts)?;
                }
            } else {
                self.extract_texts_from_json_value(&value, &mut texts)?;
            }
        }
        Ok(texts)
    }

    fn extract_texts_from_json_value(
        &self,
        value: &serde_json::Value,
        out: &mut Vec<String>,
    ) -> Result<()> {
        let obj = value
            .as_object()
            .with_context(|| format!("expected json object, got {}", value))?;

        for col_name in &self.spec.text_columns {
            match obj.get(col_name) {
                Some(serde_json::Value::String(s)) => out.push(s.clone()),
                Some(v) => anyhow::bail!(
                    "column '{col_name}' is not a string in json shard (got {})",
                    v
                ),
                None => anyhow::bail!("missing column '{col_name}' in json shard"),
            }
        }
        Ok(())
    }

    /// Stream texts from all shards. The current implementation buffers per shard;
    /// true cross-shard streaming will be added once the API stabilizes.
    pub fn stream_texts(&self) -> Result<Box<dyn Iterator<Item = Result<String>> + '_>> {
        let paths = self.download_shards()?;
        let mut all = Vec::new();
        for path in paths {
            match self.read_shard_texts(&path) {
                Ok(texts) => all.extend(texts.into_iter().map(Ok)),
                Err(e) => all.push(Err(e)),
            }
        }
        Ok(Box::new(all.into_iter()))
    }

    fn extract_texts_from_batch(&self, batch: &RecordBatch, out: &mut Vec<String>) -> Result<()> {
        if self.spec.text_columns.is_empty() {
            anyhow::bail!("no text columns configured for domain {}", self.spec.domain);
        }

        for col_name in &self.spec.text_columns {
            let col = batch
                .column_by_name(col_name)
                .with_context(|| format!("missing column '{col_name}' in shard"))?;

            if let Some(arr) = col.as_any().downcast_ref::<StringArray>() {
                for i in 0..arr.len() {
                    if !arr.is_null(i) {
                        out.push(arr.value(i).to_string());
                    }
                }
            } else {
                anyhow::bail!(
                    "column '{col_name}' is not a UTF-8 string array (type: {:?})",
                    col.data_type()
                );
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow::array::ArrayRef;
    use parquet::arrow::ArrowWriter;
    use std::sync::Arc;

    #[test]
    fn spec_from_arg_parses_all_fields() {
        let spec = HfDatasetSpec::from_arg("code=bigcode/the-stack-smol:data/python").unwrap();
        assert_eq!(spec.domain, "code");
        assert_eq!(spec.repo_id, "bigcode/the-stack-smol");
        assert_eq!(spec.subset, Some("data/python".to_string()));
        assert_eq!(spec.text_columns, vec!["content".to_string()]);
    }

    #[test]
    fn spec_from_arg_without_subset() {
        let spec = HfDatasetSpec::from_arg("math=meta-math/MetaMathQA").unwrap();
        assert_eq!(spec.domain, "math");
        assert_eq!(spec.repo_id, "meta-math/MetaMathQA");
        assert_eq!(spec.subset, None);
        assert_eq!(
            spec.text_columns,
            vec!["query".to_string(), "response".to_string()]
        );
    }

    #[test]
    fn spec_from_arg_rejects_missing_equals() {
        let err = HfDatasetSpec::from_arg("bigcode/the-stack-smol").unwrap_err();
        assert!(err.to_string().contains("domain=repo"));
    }

    fn write_test_parquet(path: &Path, rows: &[(&str, &str)]) -> Result<()> {
        let content: Vec<&str> = rows.iter().map(|(_, text)| *text).collect();
        let ids: Vec<&str> = rows.iter().map(|(id, _)| *id).collect();

        let id_array: ArrayRef = Arc::new(StringArray::from(ids));
        let content_array: ArrayRef = Arc::new(StringArray::from(content));
        let batch = RecordBatch::try_from_iter(vec![("id", id_array), ("content", content_array)])?;

        let mut writer = ArrowWriter::try_new(std::fs::File::create(path)?, batch.schema(), None)?;
        writer.write(&batch)?;
        writer.close()?;
        Ok(())
    }

    #[test]
    fn read_shard_texts_extracts_configured_columns() {
        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("test.parquet");
        write_test_parquet(&path, &[("1", "fn main() {}"), ("2", "let x = 42;")]).unwrap();

        let loader = HfDatasetLoader::new(HfDatasetSpec {
            domain: "code".to_string(),
            repo_id: "test".to_string(),
            subset: None,
            text_columns: vec!["content".to_string()],
        })
        .unwrap();

        let texts = loader.read_shard_texts(&path).unwrap();
        assert_eq!(
            texts,
            vec!["fn main() {}".to_string(), "let x = 42;".to_string()]
        );
    }

    #[test]
    fn read_shard_texts_errors_on_missing_column() {
        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("test.parquet");
        write_test_parquet(&path, &[("1", "hello")]).unwrap();

        let loader = HfDatasetLoader::new(HfDatasetSpec {
            domain: "math".to_string(),
            repo_id: "test".to_string(),
            subset: None,
            text_columns: vec!["query".to_string()],
        })
        .unwrap();

        let err = loader.read_shard_texts(&path).unwrap_err();
        assert!(err.to_string().contains("missing column 'query'"));
    }

    #[test]
    fn spec_from_arg_with_column_override() {
        let spec =
            HfDatasetSpec::from_arg("code=nickrosh/Evol-Instruct-Code-80k-v1|instruction,output")
                .unwrap();
        assert_eq!(spec.domain, "code");
        assert_eq!(spec.repo_id, "nickrosh/Evol-Instruct-Code-80k-v1");
        assert_eq!(spec.subset, None);
        assert_eq!(
            spec.text_columns,
            vec!["instruction".to_string(), "output".to_string()]
        );
    }

    #[test]
    fn read_json_shard_extracts_text_columns() {
        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("test.json");
        std::fs::write(
            &path,
            r#"[{"instruction":"hello","output":"world"},{"instruction":"foo","output":"bar"}]"#,
        )
        .unwrap();

        let loader = HfDatasetLoader::new(HfDatasetSpec {
            domain: "code".to_string(),
            repo_id: "test".to_string(),
            subset: None,
            text_columns: vec!["instruction".to_string(), "output".to_string()],
        })
        .unwrap();

        let texts = loader.read_shard_texts(&path).unwrap();
        assert_eq!(
            texts,
            vec!["hello", "world", "foo", "bar"]
                .into_iter()
                .map(|s| s.to_string())
                .collect::<Vec<_>>()
        );
    }

    #[test]
    fn read_jsonl_shard_extracts_text_columns() {
        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("test.jsonl");
        std::fs::write(&path, "{\"text\":\"first\"}\n{\"text\":\"second\"}\n").unwrap();

        let loader = HfDatasetLoader::new(HfDatasetSpec {
            domain: "text".to_string(),
            repo_id: "test".to_string(),
            subset: None,
            text_columns: vec!["text".to_string()],
        })
        .unwrap();

        let texts = loader.read_shard_texts(&path).unwrap();
        assert_eq!(texts, vec!["first".to_string(), "second".to_string()]);
    }
}