edgefirst_client/coco/

arrow.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright © 2025 Au-Zone Technologies. All Rights Reserved.

//! COCO to EdgeFirst Arrow IPC or Parquet conversion.
//!
//! Provides high-performance conversion between COCO JSON and EdgeFirst Arrow
//! format, supporting async operations and progress tracking.

use super::{
    convert::{
        box2d_to_coco_bbox, coco_bbox_to_box2d, coco_segmentation_to_mask_data,
        coco_segmentation_to_polygon, polygon_to_coco_polygon,
    },
    reader::{CocoReader, read_coco_directory},
    types::{CocoCategory, CocoDataset, CocoImage, CocoIndex, CocoInfo, CocoSegmentation},
    writer::{CocoDatasetBuilder, CocoWriteOptions, CocoWriter},
};
use crate::{Annotation, Box2d, Error, Polygon, Progress, Sample};
use polars::prelude::*;
use std::{
    collections::{BTreeMap, HashMap},
    path::{Path, PathBuf},
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
};
use tokio::sync::{Semaphore, mpsc::Sender};

/// Schema version written into Arrow IPC file metadata.
pub const SCHEMA_VERSION: &str = "2026.04";

/// Polygon rings for a single row: each ring is a vec of `(x, y)` coordinate pairs.
type PolygonRings = Vec<Vec<(f32, f32)>>;

/// Options for COCO to Arrow conversion.
///
/// Construct with `..Default::default()` (as every call site in this
/// workspace does) rather than a fully-specified struct literal: new fields
/// may be added here in a minor release, following this crate's existing
/// convention for `*Options` structs.
#[derive(Debug, Clone)]
pub struct CocoToArrowOptions {
    /// Include segmentation masks in output.
    pub include_masks: bool,
    /// Group name for all samples (e.g., "train", "val").
    pub group: Option<String>,
    /// Maximum number of parallel workers.
    pub max_workers: usize,
    /// Stage the images referenced by the COCO file into the EdgeFirst
    /// on-disk layout next to the output annotation file.
    pub images_dir: Option<PathBuf>,
    /// Symlink instead of copy (large datasets).
    pub link_images: bool,
}

impl Default for CocoToArrowOptions {
    fn default() -> Self {
        Self {
            include_masks: true,
            group: None,
            max_workers: max_workers(),
            images_dir: None,
            link_images: false,
        }
    }
}

/// Options for Arrow to COCO conversion.
#[derive(Debug, Clone)]
pub struct ArrowToCocoOptions {
    /// Filter by group names (empty = all).
    pub groups: Vec<String>,
    /// Include segmentation masks in output.
    pub include_masks: bool,
    /// COCO info section.
    pub info: Option<CocoInfo>,
    /// Pretty-print the output JSON.
    pub pretty: bool,
}

impl Default for ArrowToCocoOptions {
    fn default() -> Self {
        Self {
            groups: vec![],
            include_masks: true,
            info: None,
            pretty: false,
        }
    }
}

/// Determine maximum number of parallel workers.
fn max_workers() -> usize {
    std::env::var("MAX_COCO_WORKERS")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or_else(|| {
            let cpus = std::thread::available_parallelism()
                .map(|n| n.get())
                .unwrap_or(4);
            (cpus / 2).clamp(2, 8)
        })
}

/// Convert COCO annotations to the EdgeFirst Dataset Format.
///
/// This is a high-performance async conversion that uses parallel workers
/// for parsing and transforming annotations.
///
/// # Arguments
/// * `coco_path` - Path to COCO annotation JSON/ZIP or a standard extracted
///   directory; directory inputs combine inferred train/val groups
/// * `output_path` - Output Arrow IPC or Parquet path, selected by extension
/// * `options` - Conversion options
/// * `progress` - Optional progress channel
///
/// # Returns
/// Number of samples converted
pub async fn coco_to_arrow<P: AsRef<Path>>(
    coco_path: P,
    output_path: P,
    options: &CocoToArrowOptions,
    progress: Option<Sender<Progress>>,
) -> Result<usize, Error> {
    let coco_path = coco_path.as_ref();
    let output_path = output_path.as_ref();

    // A directory is the standard multi-split COCO layout. Process each
    // annotation file independently so train/val image or annotation IDs may
    // overlap without one split silently replacing rows from another.
    let mut sources: Vec<(CocoDataset, Option<String>, Option<String>)> = if coco_path.is_dir() {
        read_coco_directory(coco_path, &Default::default())?
            .into_iter()
            .map(|(dataset, inferred_group)| {
                let group = options.group.clone().or(Some(inferred_group.clone()));
                (dataset, group, Some(inferred_group))
            })
            .collect()
    } else {
        let reader = CocoReader::new();
        let dataset = if coco_path.extension().is_some_and(|e| e == "zip") {
            reader.read_annotations_zip(coco_path)?
        } else {
            reader.read_json(coco_path)?
        };
        vec![(dataset, options.group.clone(), None)]
    };

    let total_images: usize = sources
        .iter()
        .map(|(dataset, _, _)| dataset.images.len())
        .sum();
    let categories = collect_compatible_categories(&sources)?;
    let mut image_file_names = Vec::with_capacity(total_images);

    // Send initial progress
    if let Some(ref p) = progress {
        let _ = p
            .send(Progress {
                current: 0,
                total: total_images,
                status: None,
            })
            .await;
    }

    // Process images in parallel
    let sem = Arc::new(Semaphore::new(options.max_workers));
    let current = Arc::new(AtomicUsize::new(0));
    let mut tasks = Vec::with_capacity(total_images);
    for (dataset, group, inferred_group) in sources.drain(..) {
        let split_dir = if coco_path.is_dir() {
            inferred_group
                .as_deref()
                .and_then(|group| find_split_image_dir(coco_path, group))
        } else {
            None
        };
        image_file_names.extend(dataset.images.iter().map(|image| {
            split_dir
                .as_ref()
                .map(|dir| dir.join(&image.file_name).to_string_lossy().into_owned())
                .unwrap_or_else(|| image.file_name.clone())
        }));

        let index = Arc::new(CocoIndex::from_dataset(&dataset));
        for image in dataset.images {
            let sem = sem.clone();
            let index = index.clone();
            let current = current.clone();
            let progress = progress.clone();
            let group = group.clone();
            let include_masks = options.include_masks;

            let task = tokio::spawn(async move {
                let _permit = sem.acquire().await.map_err(Error::SemaphoreError)?;

                let samples =
                    convert_image_annotations(&image, &index, include_masks, group.as_deref());

                let c = current.fetch_add(1, Ordering::SeqCst) + 1;
                if let Some(ref p) = progress {
                    let _ = p
                        .send(Progress {
                            current: c,
                            total: total_images,
                            status: None,
                        })
                        .await;
                }

                Ok::<Vec<Sample>, Error>(samples)
            });

            tasks.push(task);
        }
    }

    // Collect all samples
    let mut all_samples = Vec::with_capacity(total_images);
    for task in tasks {
        let samples = task.await??;
        all_samples.extend(samples);
    }

    // Convert to DataFrame
    let mut df = crate::samples_dataframe(&all_samples)?;

    // Build schema-level metadata
    let mut metadata: BTreeMap<PlSmallStr, PlSmallStr> = BTreeMap::new();
    metadata.insert(
        PlSmallStr::from("schema_version"),
        PlSmallStr::from(SCHEMA_VERSION),
    );

    // Build category_metadata JSON from all categories.
    // Includes id, frequency, and any LVIS fields (synset, synonyms, def).
    // All categories are stored so that categories without annotations
    // (e.g., those only referenced in neg_category_ids) can be
    // reconstructed during Arrow→COCO export.
    if !categories.is_empty() {
        let cat_meta: HashMap<String, serde_json::Value> = categories
            .iter()
            .map(|c| {
                let mut entry = serde_json::Map::new();
                entry.insert("id".to_string(), serde_json::json!(c.id));
                if let Some(ref f) = c.frequency {
                    entry.insert(
                        "frequency".to_string(),
                        serde_json::Value::String(f.clone()),
                    );
                }
                if let Some(ref s) = c.synset {
                    entry.insert("synset".to_string(), serde_json::Value::String(s.clone()));
                }
                if let Some(ref syns) = c.synonyms {
                    entry.insert("synonyms".to_string(), serde_json::json!(syns));
                }
                if let Some(ref d) = c.def {
                    entry.insert(
                        "definition".to_string(),
                        serde_json::Value::String(d.clone()),
                    );
                }
                if let Some(ref sc) = c.supercategory {
                    entry.insert(
                        "supercategory".to_string(),
                        serde_json::Value::String(sc.clone()),
                    );
                }
                // Note: image_count and instance_count are intentionally not
                // stored — they are recomputable statistics that can be derived
                // from the annotations at any time.
                (c.name.clone(), serde_json::Value::Object(entry))
            })
            .collect();

        let json = serde_json::to_string(&cat_meta).unwrap_or_default();
        metadata.insert(
            PlSmallStr::from("category_metadata"),
            PlSmallStr::from(json.as_str()),
        );
    }

    // Write labels metadata: sorted list of category names by category_id.
    if !categories.is_empty() {
        let mut cats: Vec<_> = categories.iter().collect();
        cats.sort_by_key(|c| c.id);
        let labels: Vec<String> = cats.iter().map(|c| c.name.clone()).collect();
        let labels_json = serde_json::to_string(&labels).unwrap_or_default();
        metadata.insert(PlSmallStr::from("labels"), PlSmallStr::from(labels_json));
    }

    write_dataset(&mut df, output_path, metadata)?;

    if let Some(images_dir) = &options.images_dir {
        let report = stage_images(
            output_path,
            images_dir,
            options.link_images,
            &image_file_names,
        );
        log::info!("Staged {} images", report.staged);
        if !report.missing.is_empty() || !report.failed.is_empty() || report.collisions > 0 {
            const MAX_LISTED: usize = 5;
            let mut msg = String::new();
            if !report.missing.is_empty() {
                let listed: Vec<&str> = report
                    .missing
                    .iter()
                    .take(MAX_LISTED)
                    .map(String::as_str)
                    .collect();
                msg.push_str(&format!(
                    "{} images referenced by {} were not found in {}: {}",
                    report.missing.len(),
                    coco_path.display(),
                    images_dir.display(),
                    listed.join(", "),
                ));
                if report.missing.len() > MAX_LISTED {
                    msg.push_str(&format!(" (+{} more)", report.missing.len() - MAX_LISTED));
                }
            }
            if !report.failed.is_empty() {
                let listed: Vec<&str> = report
                    .failed
                    .iter()
                    .take(MAX_LISTED)
                    .map(String::as_str)
                    .collect();
                if !msg.is_empty() {
                    msg.push_str("; ");
                }
                msg.push_str(&format!(
                    "{} image(s) failed to stage (copy/link I/O error, see warnings above): {}",
                    report.failed.len(),
                    listed.join(", "),
                ));
                if report.failed.len() > MAX_LISTED {
                    msg.push_str(&format!(" (+{} more)", report.failed.len() - MAX_LISTED));
                }
            }
            if report.collisions > 0 {
                msg.push_str(&format!(
                    "; {} destination-basename collision(s) skipped (see warnings above)",
                    report.collisions
                ));
            }
            log::warn!("{}", msg);
        }
    }

    Ok(all_samples.len())
}

/// Collect the shared category vocabulary for a multi-split conversion.
///
/// EdgeFirst metadata is keyed by category name while rows retain the numeric
/// COCO category ID in `label_index`. A disagreement between splits therefore
/// cannot be represented faithfully and is rejected instead of silently
/// selecting whichever split happened to be read last.
fn collect_compatible_categories(
    sources: &[(CocoDataset, Option<String>, Option<String>)],
) -> Result<Vec<CocoCategory>, Error> {
    let mut by_name: BTreeMap<String, CocoCategory> = BTreeMap::new();
    let mut names_by_id: BTreeMap<u32, String> = BTreeMap::new();

    for (dataset, _, _) in sources {
        for category in &dataset.categories {
            if let Some(existing_name) = names_by_id.get(&category.id)
                && existing_name != &category.name
            {
                return Err(Error::CocoError(format!(
                    "COCO splits disagree on category {}: '{}' versus '{}'",
                    category.id, existing_name, category.name
                )));
            }
            if let Some(existing) = by_name.get(&category.name) {
                if existing.id != category.id {
                    return Err(Error::CocoError(format!(
                        "COCO splits disagree on category '{}': id {} versus {}",
                        category.name, existing.id, category.id
                    )));
                }

                let conflicting_field = [
                    (
                        "supercategory",
                        existing.supercategory.as_ref() != category.supercategory.as_ref(),
                    ),
                    (
                        "synset",
                        existing.synset.as_ref() != category.synset.as_ref(),
                    ),
                    (
                        "frequency",
                        existing.frequency.as_ref() != category.frequency.as_ref(),
                    ),
                    (
                        "synonyms",
                        existing.synonyms.as_ref() != category.synonyms.as_ref(),
                    ),
                    ("definition", existing.def.as_ref() != category.def.as_ref()),
                ]
                .into_iter()
                .find_map(|(field, differs)| differs.then_some(field));

                if let Some(field) = conflicting_field {
                    return Err(Error::CocoError(format!(
                        "COCO splits disagree on category '{}' metadata field '{}'",
                        category.name, field
                    )));
                }
            }
            names_by_id.insert(category.id, category.name.clone());
            by_name
                .entry(category.name.clone())
                .or_insert_with(|| category.clone());
        }
    }

    Ok(by_name.into_values().collect())
}

/// Find the immediate COCO image directory associated with an inferred group.
///
/// Standard layouts use `train2017`, `val2017`, and `test2017`. Matching after
/// trimming the year also supports other year suffixes without hard-coding one
/// COCO release.
fn find_split_image_dir(coco_root: &Path, group: &str) -> Option<PathBuf> {
    let mut candidates: Vec<PathBuf> = std::fs::read_dir(coco_root)
        .ok()?
        .filter_map(Result::ok)
        // `Path::is_dir` follows symlinks, which is common for large local
        // COCO image trees stored on another volume.
        .filter(|entry| entry.path().is_dir())
        .filter_map(|entry| {
            let name = entry.file_name();
            let name_str = name.to_str()?;
            (name_str.trim_end_matches(char::is_numeric) == group).then(|| PathBuf::from(name))
        })
        .collect();
    candidates.sort();
    candidates.into_iter().next()
}

/// Outcome of staging images referenced by a COCO dataset into the output
/// annotation file's sibling image directory.
struct StageReport {
    /// Number of images copied/linked (or already present from a prior run).
    staged: usize,
    /// COCO `file_name` values that could not be found under `images_dir`.
    missing: Vec<String>,
    /// COCO `file_name` values that were found under `images_dir` but failed
    /// to copy/link into the staging directory (see warnings logged at the
    /// point of failure for the underlying I/O error).
    failed: Vec<String>,
    /// Number of `file_name`s that mapped to a destination basename already
    /// claimed *within this run* by a different source path, and were
    /// therefore skipped rather than silently overwriting/aliasing another
    /// image (e.g. `train/000001.jpg` and `val/000001.jpg`).
    collisions: usize,
}

/// Stage images referenced by `file_names` into the EdgeFirst on-disk
/// layout next to `output_path`.
///
/// For an output annotation path `<dir>/<stem>.<ext>`, images are staged
/// into the sibling directory `<dir>/<stem>/`, which is where
/// `resolve_files_with_container` expects to find them.
///
/// Source resolution per `file_name` tries `images_dir.join(file_name)`
/// first (COCO `file_name` may carry subpaths), then falls back to
/// `images_dir.join(basename)`. The destination always uses the basename.
/// Destination files that already exist from a prior run are left
/// untouched, making re-runs idempotent. Two `file_name`s that resolve to
/// the same destination basename but different source paths within the
/// same run (e.g. `train/000001.jpg` and `val/000001.jpg`) are a
/// collision: only the first is staged, the rest are counted in
/// `StageReport::collisions` and logged, never silently treated as staged.
/// Missing sources and copy/link I/O failures are both reported
/// non-fatally via `StageReport`, in its `missing` and `failed` fields
/// respectively.
fn stage_images(
    output_path: &Path,
    images_dir: &Path,
    link: bool,
    file_names: &[String],
) -> StageReport {
    let parent = output_path.parent().unwrap_or_else(|| Path::new("."));
    let stem = output_path
        .file_stem()
        .and_then(std::ffi::OsStr::to_str)
        .unwrap_or("dataset");
    let dest_dir = parent.join(stem);

    if let Err(e) = std::fs::create_dir_all(&dest_dir) {
        log::warn!(
            "Failed to create image staging directory {}: {}",
            dest_dir.display(),
            e
        );
        return StageReport {
            staged: 0,
            missing: file_names.to_vec(),
            failed: Vec::new(),
            collisions: 0,
        };
    }

    let mut staged = 0;
    let mut missing = Vec::new();
    let mut failed = Vec::new();
    let mut collisions = 0;

    // Basenames claimed by a source path within this run, so a second
    // `file_name` mapping to the same basename from a *different* source
    // can be detected as a collision rather than counted as staged just
    // because the destination now exists.
    let mut claimed_this_run: HashMap<std::ffi::OsString, (PathBuf, String)> = HashMap::new();

    for file_name in file_names {
        let basename = match Path::new(file_name).file_name() {
            Some(b) => b,
            None => {
                missing.push(file_name.clone());
                continue;
            }
        };

        let candidate = images_dir.join(file_name);
        let src = if candidate.exists() {
            candidate
        } else {
            images_dir.join(basename)
        };

        if let Some((claimed_src, claimed_file_name)) = claimed_this_run.get(basename) {
            if claimed_src == &src {
                // Same source referenced twice (e.g. by two annotations) — harmless.
                staged += 1;
            } else {
                log::warn!(
                    "Destination basename '{}' collision: '{}' was already staged from '{}' this run; skipping '{}' ({})",
                    basename.to_string_lossy(),
                    claimed_file_name,
                    claimed_src.display(),
                    file_name,
                    src.display(),
                );
                collisions += 1;
            }
            continue;
        }

        let dest = dest_dir.join(basename);

        // Skip files already staged from a prior run.
        if dest.symlink_metadata().is_ok() {
            staged += 1;
            claimed_this_run.insert(basename.to_os_string(), (src, file_name.clone()));
            continue;
        }

        if !src.exists() {
            missing.push(file_name.clone());
            continue;
        }

        let result = if link {
            stage_via_link(&src, &dest)
        } else {
            std::fs::copy(&src, &dest).map(|_| ())
        };

        match result {
            Ok(()) => {
                staged += 1;
                claimed_this_run.insert(basename.to_os_string(), (src, file_name.clone()));
            }
            Err(e) => {
                log::warn!("Failed to stage image '{}': {}", file_name, e);
                failed.push(file_name.clone());
            }
        }
    }

    StageReport {
        staged,
        missing,
        failed,
        collisions,
    }
}

/// Symlink `dest` to the canonicalized `src` (Unix/macOS). Staging is a
/// Linux/macOS workflow; other platforms fall back to copying.
#[cfg(unix)]
fn stage_via_link(src: &Path, dest: &Path) -> std::io::Result<()> {
    let canonical = src.canonicalize()?;
    std::os::unix::fs::symlink(canonical, dest)
}

#[cfg(not(unix))]
fn stage_via_link(src: &Path, dest: &Path) -> std::io::Result<()> {
    std::fs::copy(src, dest).map(|_| ())
}

/// Write a DataFrame plus file-level metadata to `output_path`.
///
/// Format is selected from the output path extension: `.parquet` writes
/// Apache Parquet with the metadata stored as footer key-value pairs;
/// anything else (including `.arrow`/`.ipc`) writes Arrow IPC with the
/// metadata stored as custom schema metadata.
pub fn write_dataset(
    df: &mut DataFrame,
    output_path: &Path,
    metadata: BTreeMap<PlSmallStr, PlSmallStr>,
) -> Result<(), Error> {
    if let Some(parent) = output_path.parent()
        && !parent.as_os_str().is_empty()
    {
        std::fs::create_dir_all(parent)?;
    }
    let ext = output_path
        .extension()
        .and_then(std::ffi::OsStr::to_str)
        .unwrap_or("");
    let mut file = std::fs::File::create(output_path)?;
    match ext {
        "parquet" => {
            let kv = metadata
                .into_iter()
                .map(|(k, v)| (k.to_string(), v.to_string()))
                .collect();
            ParquetWriter::new(&mut file)
                .with_key_value_metadata(Some(KeyValueMetadata::from_static(kv)))
                .finish(df)?;
        }
        _ => {
            let mut writer = IpcWriter::new(&mut file);
            writer.set_custom_schema_metadata(Arc::new(metadata));
            writer.finish(df)?;
        }
    }
    Ok(())
}

/// Convert a single image's annotations to EdgeFirst samples.
fn convert_image_annotations(
    image: &CocoImage,
    index: &CocoIndex,
    include_masks: bool,
    group: Option<&str>,
) -> Vec<Sample> {
    let annotations = index.annotations_for_image(image.id);
    let sample_name = sample_name_from_filename(&image.file_name);

    // Translate LVIS image-level fields to label_index lists
    let neg_label_indices = image.neg_category_ids.as_ref().map(|ids| {
        ids.iter()
            .filter_map(|&id| index.label_index(id).map(|idx| idx as u32))
            .collect::<Vec<u32>>()
    });
    let not_exhaustive_label_indices = image.not_exhaustive_category_ids.as_ref().map(|ids| {
        ids.iter()
            .filter_map(|&id| index.label_index(id).map(|idx| idx as u32))
            .collect::<Vec<u32>>()
    });

    let mut samples: Vec<Sample> = annotations
        .iter()
        .filter_map(|ann| {
            let label = index.label_name(ann.category_id)?;
            let label_index = index.label_index(ann.category_id);

            // Convert bbox
            let box2d = coco_bbox_to_box2d(&ann.bbox, image.width, image.height);

            // Convert segmentation based on type:
            // - Polygon → annotation.polygon (normalized coords)
            // - RLE/CompressedRle → annotation.mask (PNG-encoded MaskData)
            let (polygon, mask) = if include_masks {
                if let Some(seg) = &ann.segmentation {
                    match seg {
                        CocoSegmentation::Polygon(_) => {
                            let poly =
                                coco_segmentation_to_polygon(seg, image.width, image.height).ok();
                            (poly, None)
                        }
                        CocoSegmentation::Rle(_) | CocoSegmentation::CompressedRle(_) => {
                            let mask_data = coco_segmentation_to_mask_data(seg).ok().flatten();
                            (None, mask_data)
                        }
                    }
                } else {
                    (None, None)
                }
            } else {
                (None, None)
            };

            let mut annotation = Annotation::new();
            annotation.set_name(Some(sample_name.clone()));
            annotation.set_object_id(Some(ann.id.to_string()));
            annotation.set_label(Some(label.to_string()));
            annotation.set_label_index(label_index);
            annotation.set_box2d(Some(box2d));
            annotation.set_polygon(polygon);
            annotation.set_mask(mask);
            annotation.set_group(group.map(String::from));
            annotation.set_iscrowd(Some(ann.iscrowd != 0));
            annotation.set_category_frequency(index.frequency(ann.category_id).map(String::from));

            // Map COCO score to appropriate geometry score field
            if let Some(score) = ann.score {
                let score_f32 = score as f32;
                if annotation.mask().is_some() {
                    annotation.set_mask_score(Some(score_f32));
                } else if annotation.polygon().is_some() {
                    annotation.set_polygon_score(Some(score_f32));
                } else {
                    annotation.set_box2d_score(Some(score_f32));
                }
            }

            let mut sample = Sample {
                image_name: Some(sample_name.clone()),
                width: Some(image.width),
                height: Some(image.height),
                group: group.map(String::from),
                annotations: vec![annotation],
                ..Default::default()
            };
            sample.neg_label_indices = neg_label_indices.clone();
            sample.not_exhaustive_label_indices = not_exhaustive_label_indices.clone();

            Some(sample)
        })
        .collect();

    // Emit a placeholder row for any image with no annotations so the image is
    // never dropped from the dataset. This preserves the image's group (dataset
    // split) for every image, and carries any LVIS neg/exhaustive category data
    // for images that have verified-negative labels but no positive annotations.
    if samples.is_empty() {
        let mut sample = Sample {
            image_name: Some(sample_name.clone()),
            width: Some(image.width),
            height: Some(image.height),
            group: group.map(String::from),
            ..Default::default()
        };
        sample.neg_label_indices = neg_label_indices;
        sample.not_exhaustive_label_indices = not_exhaustive_label_indices;
        samples.push(sample);
    }

    samples
}

/// Extract sample name from image filename.
fn sample_name_from_filename(filename: &str) -> String {
    Path::new(filename)
        .file_stem()
        .and_then(|s| s.to_str())
        .map(String::from)
        .unwrap_or_else(|| filename.to_string())
}

/// Convert EdgeFirst Arrow format to COCO annotations.
///
/// Reads an Arrow file and produces COCO JSON output. LVIS extension fields
/// are preserved when present in the Arrow file: `neg_category_ids`,
/// `not_exhaustive_category_ids`, category `frequency`, annotation `iscrowd`,
/// `supercategory`, and category metadata (`synset`, `synonyms`, `def`).
///
/// # Arguments
/// * `arrow_path` - Path to EdgeFirst Arrow file
/// * `output_path` - Output COCO JSON file path
/// * `options` - Conversion options
/// * `progress` - Optional progress channel
///
/// # Returns
/// Number of annotations converted
pub async fn arrow_to_coco<P: AsRef<Path>>(
    arrow_path: P,
    output_path: P,
    options: &ArrowToCocoOptions,
    progress: Option<Sender<Progress>>,
) -> Result<usize, Error> {
    let arrow_path = arrow_path.as_ref();
    let output_path = output_path.as_ref();

    // Read the DataFrame and its file-level metadata (schema_version,
    // category_metadata, labels). Accepts .arrow/.ipc/.parquet.
    let (df, metadata) = crate::format::read_dataset_dataframe(arrow_path)?;
    let schema_version = metadata.get("schema_version").cloned();
    let category_metadata_json = metadata.get("category_metadata").cloned();
    let labels_metadata_json = metadata.get("labels").cloned();

    // Determine format version: absent → 2025.10, present → use value
    let is_legacy = schema_version.is_none();

    // Get group column for filtering
    let groups_to_filter: std::collections::HashSet<_> = options.groups.iter().cloned().collect();

    let total_rows = df.height();

    if let Some(ref p) = progress {
        let _ = p
            .send(Progress {
                current: 0,
                total: total_rows,
                status: None,
            })
            .await;
    }

    // Extract columns - all at once for O(n) instead of O(n²) per-row access
    let names: Vec<String> = df
        .column("name")?
        .str()?
        .iter()
        .map(|s| s.unwrap_or_default().to_string())
        .collect();

    let labels: Vec<String> = df
        .column("label")
        .ok()
        .and_then(|c| c.cast(&DataType::String).ok())
        .map(|c| {
            c.str()
                .ok()
                .map(|s| {
                    s.iter()
                        .map(|v| v.unwrap_or_default().to_string())
                        .collect()
                })
                .unwrap_or_else(|| vec![String::new(); total_rows])
        })
        .unwrap_or_else(|| vec![String::new(); total_rows]);

    let label_indices: Vec<Option<u64>> = df
        .column("label_index")
        .ok()
        .map(|c| {
            c.u64()
                .ok()
                .map(|s| s.iter().collect())
                .unwrap_or_else(|| vec![None; total_rows])
        })
        .unwrap_or_else(|| vec![None; total_rows]);

    // Get group column for filtering
    let groups: Vec<String> = df
        .column("group")
        .ok()
        .and_then(|c| c.cast(&DataType::String).ok())
        .map(|c| {
            c.str()
                .ok()
                .map(|s| {
                    s.iter()
                        .map(|v| v.unwrap_or_default().to_string())
                        .collect()
                })
                .unwrap_or_default()
        })
        .unwrap_or_else(|| vec!["".to_string(); total_rows]);

    // Extract all box2d values upfront (O(n) instead of O(n²))
    let box2ds = df
        .column("box2d")
        .ok()
        .map(extract_all_box2ds)
        .transpose()?
        .unwrap_or_else(|| vec![[0.0; 4]; total_rows]);

    // Extract segmentation data based on schema version
    //
    // 2025.10 (legacy): mask column is List(Float32) with NaN-separated polygon coords
    // 2026.04+:         polygon column is List(List(Float32)), mask column is Binary (PNG)
    let legacy_masks: Option<Vec<Vec<f32>>> = if is_legacy && options.include_masks {
        df.column("mask").ok().map(extract_all_masks).transpose()?
    } else {
        None
    };

    let polygons_2026: Option<Vec<Option<PolygonRings>>> = if !is_legacy && options.include_masks {
        df.column("polygon")
            .ok()
            .map(|c| extract_all_polygons(c, total_rows))
    } else {
        None
    };

    let mask_binary_2026: Option<Vec<Option<Vec<u8>>>> = if !is_legacy && options.include_masks {
        df.column("mask")
            .ok()
            .map(|c| extract_all_binary_masks(c, total_rows))
    } else {
        None
    };

    // Extract all sizes upfront if present
    let sizes = df
        .column("size")
        .ok()
        .and_then(|c| extract_all_sizes(c).ok());

    // Extract iscrowd column (optional, Boolean in 2026.04, UInt32 in older schemas)
    let iscrowds: Vec<u8> = df
        .column("iscrowd")
        .ok()
        .map(|c| {
            // Try Boolean first (2026.04 schema), then fall back to UInt32 (older schemas)
            if let Ok(bool_ca) = c.bool() {
                bool_ca
                    .iter()
                    .map(|v| if v.unwrap_or(false) { 1 } else { 0 })
                    .collect()
            } else {
                c.u32()
                    .ok()
                    .map(|s| s.iter().map(|v| v.unwrap_or(0) as u8).collect())
                    .unwrap_or_else(|| vec![0; total_rows])
            }
        })
        .unwrap_or_else(|| vec![0; total_rows]);

    // Extract category_frequency column (optional, Categorical/String)
    let category_frequencies: Vec<Option<String>> = df
        .column("category_frequency")
        .ok()
        .and_then(|c| c.cast(&DataType::String).ok())
        .map(|c| {
            c.str()
                .ok()
                .map(|s| s.iter().map(|v| v.map(String::from)).collect())
                .unwrap_or_else(|| vec![None; total_rows])
        })
        .unwrap_or_else(|| vec![None; total_rows]);

    // Extract neg_label_indices column (optional, List<UInt32>)
    let neg_label_indices: Vec<Option<Vec<u32>>> = df
        .column("neg_label_indices")
        .ok()
        .map(|c| extract_list_u32_column(c, total_rows))
        .unwrap_or_else(|| vec![None; total_rows]);

    // Extract not_exhaustive_label_indices column (optional, List<UInt32>)
    let not_exhaustive_label_indices: Vec<Option<Vec<u32>>> = df
        .column("not_exhaustive_label_indices")
        .ok()
        .map(|c| extract_list_u32_column(c, total_rows))
        .unwrap_or_else(|| vec![None; total_rows]);

    // Extract score columns (2026.04 schema)
    let box2d_scores: Vec<Option<f32>> = extract_f32_column(&df, "box2d_score", total_rows);
    let box3d_scores: Vec<Option<f32>> = extract_f32_column(&df, "box3d_score", total_rows);
    let polygon_scores: Vec<Option<f32>> = extract_f32_column(&df, "polygon_score", total_rows);
    let mask_scores: Vec<Option<f32>> = extract_f32_column(&df, "mask_score", total_rows);

    // Extract object_id column (optional, String) and parse to u64 where
    // possible. This preserves the source COCO/LVIS annotation `id` across
    // the Arrow→COCO round-trip so downstream tools (e.g., prompted-
    // segmentation workflows that key on ann.id) see the original IDs.
    //
    // Non-numeric object_ids — produced by datasets whose instances carry
    // string UUIDs rather than COCO numeric IDs — parse to None and fall
    // through to auto-generated IDs in the builder. This is intentional:
    // COCO requires numeric annotation IDs, and a UUID has no meaningful
    // numeric projection.
    let object_id_u64s: Vec<Option<u64>> = df
        .column("object_id")
        .ok()
        .and_then(|c| c.cast(&DataType::String).ok())
        .map(|c| {
            c.str()
                .ok()
                .map(|s| {
                    s.iter()
                        .map(|v| v.and_then(|s| s.parse::<u64>().ok()))
                        .collect()
                })
                .unwrap_or_else(|| vec![None; total_rows])
        })
        .unwrap_or_else(|| vec![None; total_rows]);

    // Build COCO dataset
    let mut builder = CocoDatasetBuilder::new();

    if let Some(info) = &options.info {
        builder = builder.info(info.clone());
    }

    // Group-filter predicate: returns true if this row should be skipped
    let skip_row = |i: usize| -> bool {
        !groups_to_filter.is_empty() && !groups_to_filter.contains(&groups[i])
    };

    // Track unique images and categories
    let mut image_dimensions: HashMap<String, (u32, u32)> = HashMap::new();
    let mut image_ids: HashMap<String, u64> = HashMap::new();
    let mut category_ids: HashMap<String, u32> = HashMap::new();

    // First pass: collect unique images and categories
    for i in 0..total_rows {
        if skip_row(i) {
            continue;
        }

        let name = &names[i];
        let label = &labels[i];

        // Get or estimate image dimensions
        if !image_ids.contains_key(name) {
            let (width, height) = sizes
                .as_ref()
                .and_then(|s| s.get(i).copied())
                .unwrap_or((0, 0));

            let id = builder.add_image(name, width, height);
            image_ids.insert(name.clone(), id);
            image_dimensions.insert(name.clone(), (width, height));
        }

        if !label.is_empty() && !category_ids.contains_key(label) {
            let id = if let Some(Some(idx)) = label_indices.get(i) {
                builder.add_category_with_id(*idx as u32, label, None)
            } else {
                builder.add_category(label, None)
            };
            category_ids.insert(label.clone(), id);
        }
    }

    // Second pass: create annotations
    let mut last_progress_update = 0;
    for i in 0..total_rows {
        if skip_row(i) {
            continue;
        }

        let name = &names[i];
        let label = &labels[i];

        // Skip sentinel rows (empty label = image with neg/exhaustive data but no annotations)
        if label.is_empty() {
            continue;
        }

        let image_id = *image_ids.get(name).unwrap_or(&0);
        let category_id = *category_ids.get(label).unwrap_or(&0);
        let (width, height) = *image_dimensions.get(name).unwrap_or(&(1, 1));

        // Convert box2d from Arrow center-normalized [cx, cy, w, h] to COCO format
        // Arrow stores center-point, Box2d expects top-left
        let bbox = box2ds.get(i).map(|box2d| {
            let cx = box2d[0];
            let cy = box2d[1];
            let w = box2d[2];
            let h = box2d[3];
            // Convert from center-point to top-left format
            let left = cx - w / 2.0;
            let top = cy - h / 2.0;
            let ef_box2d = Box2d::new(left, top, w, h);
            box2d_to_coco_bbox(&ef_box2d, width, height)
        });

        // Build segmentation based on schema version
        let segmentation = if options.include_masks {
            if is_legacy {
                // 2025.10: mask column contains NaN-separated flat polygon coords
                legacy_masks.as_ref().and_then(|m| {
                    m.get(i).and_then(|coords| {
                        if coords.is_empty() {
                            None
                        } else {
                            let rings = crate::unflatten_polygon_coordinates(coords);
                            let polygon = Polygon::new(rings);
                            let coco_poly = polygon_to_coco_polygon(&polygon, width, height);
                            if coco_poly.is_empty() {
                                None
                            } else {
                                Some(CocoSegmentation::Polygon(coco_poly))
                            }
                        }
                    })
                })
            } else {
                // 2026.04+: try mask (Binary/PNG → RLE) first, then polygon column
                let mask_seg = mask_binary_2026.as_ref().and_then(|masks| {
                    masks.get(i).and_then(|opt_bytes| {
                        opt_bytes
                            .as_ref()
                            .and_then(|png_bytes| png_to_rle_segmentation(png_bytes, i))
                    })
                });

                if mask_seg.is_some() {
                    mask_seg
                } else {
                    // Fall back to polygon column
                    polygons_2026.as_ref().and_then(|polys| {
                        polys.get(i).and_then(|opt_rings| {
                            opt_rings.as_ref().and_then(|rings| {
                                if rings.is_empty() {
                                    return None;
                                }
                                let polygon = Polygon::new(rings.clone());
                                let coco_poly = polygon_to_coco_polygon(&polygon, width, height);
                                if coco_poly.is_empty() {
                                    None
                                } else {
                                    Some(CocoSegmentation::Polygon(coco_poly))
                                }
                            })
                        })
                    })
                }
            }
        } else {
            None
        };

        // Determine the score: use first non-null from available score columns
        let score: Option<f64> = mask_scores[i]
            .or(polygon_scores[i])
            .or(box3d_scores[i])
            .or(box2d_scores[i])
            .map(|s| s as f64);

        if let Some(bbox) = bbox {
            let iscrowd = iscrowds[i];
            let ann_id = builder.add_annotation_with_id(
                object_id_u64s[i],
                image_id,
                category_id,
                bbox,
                segmentation,
                iscrowd,
            );

            // Set score on the annotation if present
            if let Some(score_val) = score {
                builder.set_annotation_score(ann_id, score_val);
            }
        }

        // Update progress every 1000 rows to reduce overhead
        if let Some(ref p) = progress
            && (i - last_progress_update >= 1000 || i == total_rows - 1)
        {
            let _ = p
                .send(Progress {
                    current: i + 1,
                    total: total_rows,
                    status: None,
                })
                .await;
            last_progress_update = i;
        }
    }

    // Send final progress event (may not have fired if last rows were filtered)
    if let Some(ref p) = progress
        && last_progress_update < total_rows.saturating_sub(1)
    {
        let _ = p
            .send(Progress {
                current: total_rows,
                total: total_rows,
                status: None,
            })
            .await;
    }

    // Third pass: set LVIS image-level fields (neg/not-exhaustive category IDs)
    // Since label_index == category_id, we can use the values directly.
    {
        let mut processed_images: std::collections::HashSet<u64> = std::collections::HashSet::new();
        for i in 0..total_rows {
            if skip_row(i) {
                continue;
            }
            let name = &names[i];
            if let Some(&image_id) = image_ids.get(name) {
                if !processed_images.insert(image_id) {
                    continue;
                }
                let neg = neg_label_indices[i].clone();
                let not_exhaustive = not_exhaustive_label_indices[i].clone();
                if neg.is_some() || not_exhaustive.is_some() {
                    builder.set_image_neg_categories(image_id, neg, not_exhaustive);
                }
            }
        }
    }

    // Set category frequency from the category_frequency column.
    // Build a map of category_name -> frequency from the first occurrence.
    {
        let mut freq_map: HashMap<String, String> = HashMap::new();
        for i in 0..total_rows {
            if skip_row(i) {
                continue;
            }
            let label = &labels[i];
            if !label.is_empty()
                && !freq_map.contains_key(label)
                && let Some(ref freq) = category_frequencies[i]
            {
                freq_map.insert(label.clone(), freq.clone());
            }
        }
        for (name, freq) in &freq_map {
            builder.set_category_metadata(name, None, Some(freq.clone()), None, None);
        }
    }

    // Set category metadata from file-level metadata JSON
    // (id, frequency, synset, synonyms, def, supercategory).
    // Also creates categories that exist in metadata but have no annotations
    // (e.g., categories only referenced in neg_category_ids).
    // set_category_metadata only updates fields that are Some, so frequency
    // set from the column above is preserved for categories that had annotations.
    if let Some(ref json_str) = category_metadata_json
        && let Ok(meta) = serde_json::from_str::<HashMap<String, serde_json::Value>>(json_str)
    {
        for (cat_name, value) in &meta {
            let supercategory = value.get("supercategory").and_then(|v| v.as_str());

            // If this category doesn't exist yet, create it with the stored id
            if !category_ids.contains_key(cat_name.as_str()) {
                let cat_id = value.get("id").and_then(|v| v.as_u64()).map(|id| id as u32);
                let id = if let Some(cat_id) = cat_id {
                    builder.add_category_with_id(cat_id, cat_name, supercategory)
                } else {
                    builder.add_category(cat_name, supercategory)
                };
                category_ids.insert(cat_name.clone(), id);
            } else {
                // Category already exists — set supercategory if present in metadata
                if let Some(sc) = supercategory {
                    builder.set_category_supercategory(cat_name, sc);
                }
            }

            let synset = value
                .get("synset")
                .and_then(|v| v.as_str())
                .map(String::from);
            let frequency = value
                .get("frequency")
                .and_then(|v| v.as_str())
                .map(String::from);
            let synonyms = value.get("synonyms").and_then(|v| {
                v.as_array().map(|arr| {
                    arr.iter()
                        .filter_map(|s| s.as_str().map(String::from))
                        .collect()
                })
            });
            let def = value
                .get("definition")
                .and_then(|v| v.as_str())
                .map(String::from);

            builder.set_category_metadata(cat_name, synset, frequency, synonyms, def);
        }
    }

    // Populate category names from labels metadata if categories weren't set
    // from category_metadata (e.g., older files that only have labels list).
    if category_metadata_json.is_none()
        && let Some(ref labels_json) = labels_metadata_json
        && let Ok(label_names) = serde_json::from_str::<Vec<String>>(labels_json)
    {
        for label_name in &label_names {
            if !category_ids.contains_key(label_name) {
                let id = builder.add_category(label_name, None);
                category_ids.insert(label_name.clone(), id);
            }
        }
    }

    let dataset = builder.build();
    let annotation_count = dataset.annotations.len();

    // Write output
    let writer = CocoWriter::with_options(CocoWriteOptions {
        pretty: options.pretty,
        ..Default::default()
    });
    writer.write_json(&dataset, output_path)?;

    Ok(annotation_count)
}

/// Extract all box2d values from a column at once (O(n) instead of O(n²)).
fn extract_all_box2ds(col: &Column) -> Result<Vec<[f32; 4]>, Error> {
    let arr = col.array()?;
    let mut result = Vec::with_capacity(arr.len());

    for inner in arr.amortized_iter() {
        let values = if let Some(inner) = inner {
            let series = inner.as_ref();
            let vals: Vec<f32> = series
                .f32()
                .map_err(|e| Error::CocoError(format!("box2d cast error: {}", e)))?
                .iter()
                .map(|v| v.unwrap_or(0.0))
                .collect();

            if vals.len() == 4 {
                [vals[0], vals[1], vals[2], vals[3]]
            } else {
                [0.0, 0.0, 0.0, 0.0]
            }
        } else {
            [0.0, 0.0, 0.0, 0.0]
        };
        result.push(values);
    }

    Ok(result)
}

/// Extract all mask coordinates from a column at once (O(n) instead of O(n²)).
fn extract_all_masks(col: &Column) -> Result<Vec<Vec<f32>>, Error> {
    let list = col.list()?;
    let mut result = Vec::with_capacity(list.len());

    for i in 0..list.len() {
        let coords = match list.get_as_series(i) {
            Some(series) => series
                .f32()
                .map_err(|e| Error::CocoError(format!("mask cast error: {}", e)))?
                .iter()
                .map(|v| v.unwrap_or(f32::NAN))
                .collect(),
            None => vec![],
        };
        result.push(coords);
    }

    Ok(result)
}

/// Extract all image sizes from a column at once.
fn extract_all_sizes(col: &Column) -> Result<Vec<(u32, u32)>, Error> {
    let arr = col.array()?;
    let mut result = Vec::with_capacity(arr.len());

    for inner in arr.amortized_iter() {
        let size = if let Some(inner) = inner {
            let series = inner.as_ref();
            let values: Vec<u32> = series
                .u32()
                .map_err(|e| Error::CocoError(format!("size cast error: {}", e)))?
                .iter()
                .map(|v| v.unwrap_or(0))
                .collect();

            if values.len() >= 2 {
                (values[0], values[1])
            } else {
                (0, 0)
            }
        } else {
            (0, 0)
        };
        result.push(size);
    }

    Ok(result)
}

/// Extract a List<UInt32> column into a vector of optional Vec<u32>.
fn extract_list_u32_column(col: &Column, total_rows: usize) -> Vec<Option<Vec<u32>>> {
    col.list()
        .ok()
        .map(|list| {
            (0..list.len())
                .map(|i| {
                    list.get_as_series(i).and_then(|series| {
                        series
                            .u32()
                            .ok()
                            .map(|ca| ca.iter().flatten().collect::<Vec<u32>>())
                    })
                })
                .collect()
        })
        .unwrap_or_else(|| vec![None; total_rows])
}

/// Extract polygon rings from a `List(List(Float32))` column (2026.04 schema).
///
/// Each row is an optional list of rings; each ring is a list of flat `[x, y, x, y, ...]`
/// coordinate pairs.
fn extract_all_polygons(col: &Column, total_rows: usize) -> Vec<Option<PolygonRings>> {
    let outer_list = match col.list() {
        Ok(l) => l,
        Err(_) => return vec![None; total_rows],
    };

    let mut result = Vec::with_capacity(total_rows);
    for i in 0..outer_list.len() {
        let rings = outer_list.get_as_series(i).and_then(|ring_series| {
            let inner_list = ring_series.list().ok()?;
            let mut rings = Vec::new();
            for j in 0..inner_list.len() {
                if let Some(coords_series) = inner_list.get_as_series(j)
                    && let Ok(f32_ca) = coords_series.f32()
                {
                    let coords: Vec<f32> = f32_ca.iter().map(|v| v.unwrap_or(0.0)).collect();
                    // Convert flat [x, y, x, y, ...] to Vec<(f32, f32)>
                    let points: Vec<(f32, f32)> = coords
                        .chunks(2)
                        .filter(|c| c.len() == 2)
                        .map(|c| (c[0], c[1]))
                        .collect();
                    if !points.is_empty() {
                        rings.push(points);
                    }
                }
            }
            if rings.is_empty() { None } else { Some(rings) }
        });
        result.push(rings);
    }
    result
}

/// Extract binary mask data from a `Binary` column (2026.04 schema — PNG bytes).
fn extract_all_binary_masks(col: &Column, total_rows: usize) -> Vec<Option<Vec<u8>>> {
    let binary_ca = match col.binary() {
        Ok(b) => b,
        Err(_) => return vec![None; total_rows],
    };

    (0..binary_ca.len())
        .map(|i| binary_ca.get(i).map(|bytes| bytes.to_vec()))
        .collect()
}

/// Extract an optional Float32 column by name.
fn extract_f32_column(df: &DataFrame, name: &str, total_rows: usize) -> Vec<Option<f32>> {
    df.column(name)
        .ok()
        .and_then(|c| c.f32().ok())
        .map(|ca| ca.iter().collect())
        .unwrap_or_else(|| vec![None; total_rows])
}

/// Decode a PNG mask (Binary column bytes) into COCO RLE segmentation.
///
/// Validates the PNG, decodes pixels, binarizes if needed (8-bit or 16-bit),
/// and encodes as COCO RLE. Returns `None` for empty or invalid data (with
/// a warning log for invalid cases).
fn png_to_rle_segmentation(png_bytes: &[u8], row_index: usize) -> Option<CocoSegmentation> {
    if png_bytes.is_empty() {
        return None;
    }

    let mask_data = match crate::MaskData::from_png_checked(png_bytes.to_vec()) {
        Ok(m) => m,
        Err(e) => {
            log::warn!("Skipping invalid PNG mask at row {}: {}", row_index, e);
            return None;
        }
    };

    let mw = mask_data.width();
    let mh = mask_data.height();
    let bit_depth = mask_data.bit_depth();

    let decoded = match mask_data.decode() {
        Ok(d) => d,
        Err(e) => {
            log::warn!("Failed to decode PNG mask at row {}: {}", row_index, e);
            return None;
        }
    };

    let binary_mask = match bit_depth {
        1 => decoded,
        8 => {
            log::warn!(
                "Binarizing 8-bit mask for row {} — score data is lost",
                row_index
            );
            decoded
                .iter()
                .map(|&v| if v >= 128 { 1 } else { 0 })
                .collect()
        }
        16 => {
            log::warn!(
                "Binarizing 16-bit mask for row {} — score data is lost",
                row_index
            );
            decoded
                .chunks(2)
                .map(|pair| {
                    let val = if pair.len() == 2 {
                        u16::from_be_bytes([pair[0], pair[1]])
                    } else {
                        0
                    };
                    if val >= 32768 { 1u8 } else { 0u8 }
                })
                .collect()
        }
        _ => decoded,
    };

    match super::convert::encode_rle(&binary_mask, mw, mh) {
        Ok(rle) => Some(CocoSegmentation::Rle(rle)),
        Err(e) => {
            log::warn!("Failed to encode RLE for row {}: {}", row_index, e);
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::coco::{CocoAnnotation, CocoCategory, CocoDataset};
    use tempfile::TempDir;

    // =========================================================================
    // unflatten_polygon_coords tests
    // =========================================================================

    #[test]
    fn test_unflatten_polygon_coords_empty() {
        let coords: Vec<f32> = vec![];
        let result = crate::unflatten_polygon_coordinates(&coords);
        assert!(result.is_empty());
    }

    #[test]
    fn test_unflatten_polygon_coords_single_polygon() {
        // Simple rectangle: 4 points
        let coords = vec![0.1, 0.2, 0.3, 0.2, 0.3, 0.4, 0.1, 0.4];
        let result = crate::unflatten_polygon_coordinates(&coords);

        assert_eq!(result.len(), 1);
        assert_eq!(result[0].len(), 4);
        assert_eq!(result[0][0], (0.1, 0.2));
        assert_eq!(result[0][3], (0.1, 0.4));
    }

    #[test]
    fn test_unflatten_polygon_coords_multiple_polygons() {
        // Two triangles separated by NaN
        let coords = vec![
            0.1,
            0.1,
            0.2,
            0.1,
            0.15,
            0.2,      // First triangle
            f32::NAN, // Separator
            0.5,
            0.5,
            0.6,
            0.5,
            0.55,
            0.6, // Second triangle
        ];
        let result = crate::unflatten_polygon_coordinates(&coords);

        assert_eq!(result.len(), 2);
        assert_eq!(result[0].len(), 3);
        assert_eq!(result[1].len(), 3);
        assert_eq!(result[0][0], (0.1, 0.1));
        assert_eq!(result[1][0], (0.5, 0.5));
    }

    #[test]
    fn test_unflatten_polygon_coords_leading_nan() {
        // NaN at the start should be handled gracefully
        let coords = vec![f32::NAN, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6];
        let result = crate::unflatten_polygon_coordinates(&coords);

        assert_eq!(result.len(), 1);
        assert_eq!(result[0].len(), 3);
    }

    #[test]
    fn test_unflatten_polygon_coords_trailing_nan() {
        // NaN at the end
        let coords = vec![0.1, 0.2, 0.3, 0.4, f32::NAN];
        let result = crate::unflatten_polygon_coordinates(&coords);

        assert_eq!(result.len(), 1);
        assert_eq!(result[0].len(), 2);
    }

    #[test]
    fn test_unflatten_polygon_coords_consecutive_nans() {
        // Multiple NaNs in a row
        let coords = vec![0.1, 0.2, f32::NAN, f32::NAN, 0.3, 0.4];
        let result = crate::unflatten_polygon_coordinates(&coords);

        assert_eq!(result.len(), 2);
        assert_eq!(result[0].len(), 1);
        assert_eq!(result[1].len(), 1);
    }

    #[test]
    fn test_unflatten_polygon_coords_odd_values() {
        // Odd number of coordinates (trailing x without y)
        let coords = vec![0.1, 0.2, 0.3, 0.4, 0.5];
        let result = crate::unflatten_polygon_coordinates(&coords);

        assert_eq!(result.len(), 1);
        assert_eq!(result[0].len(), 2); // Only complete pairs
    }

    // =========================================================================
    // convert_image_annotations tests
    // =========================================================================

    #[test]
    fn test_convert_image_annotations_basic() {
        let image = CocoImage {
            id: 1,
            width: 640,
            height: 480,
            file_name: "test_image.jpg".to_string(),
            ..Default::default()
        };

        let dataset = CocoDataset {
            images: vec![image.clone()],
            categories: vec![CocoCategory {
                id: 1,
                name: "cat".to_string(),
                supercategory: Some("animal".to_string()),
                ..Default::default()
            }],
            annotations: vec![CocoAnnotation {
                id: 42,
                image_id: 1,
                category_id: 1,
                bbox: [100.0, 100.0, 200.0, 200.0],
                area: 40000.0,
                iscrowd: 0,
                segmentation: None,
                score: None,
            }],
            ..Default::default()
        };

        let index = CocoIndex::from_dataset(&dataset);
        let samples = convert_image_annotations(&image, &index, true, Some("train"));

        assert_eq!(samples.len(), 1);
        assert_eq!(samples[0].image_name, Some("test_image".to_string()));
        assert_eq!(samples[0].group, Some("train".to_string()));
        assert_eq!(samples[0].annotations.len(), 1);
        assert_eq!(samples[0].annotations[0].label(), Some(&"cat".to_string()));
        assert_eq!(
            samples[0].annotations[0].object_id(),
            Some(&"42".to_string()),
            "object_id must be populated from COCO annotation id to enable \
             prediction-to-prompt linking in prompted-segmentation workflows",
        );
    }

    #[test]
    fn test_convert_image_annotations_with_mask() {
        let image = CocoImage {
            id: 1,
            width: 100,
            height: 100,
            file_name: "masked.jpg".to_string(),
            ..Default::default()
        };

        let dataset = CocoDataset {
            images: vec![image.clone()],
            categories: vec![CocoCategory {
                id: 1,
                name: "object".to_string(),
                supercategory: None,
                ..Default::default()
            }],
            annotations: vec![CocoAnnotation {
                id: 1,
                image_id: 1,
                category_id: 1,
                bbox: [10.0, 10.0, 50.0, 50.0],
                area: 2500.0,
                iscrowd: 0,
                segmentation: Some(CocoSegmentation::Polygon(vec![vec![
                    10.0, 10.0, 60.0, 10.0, 60.0, 60.0, 10.0, 60.0,
                ]])),
                score: None,
            }],
            ..Default::default()
        };

        let index = CocoIndex::from_dataset(&dataset);

        // With masks enabled
        let samples_with_mask = convert_image_annotations(&image, &index, true, None);
        assert!(samples_with_mask[0].annotations[0].polygon().is_some());

        // With masks disabled
        let samples_no_mask = convert_image_annotations(&image, &index, false, None);
        assert!(samples_no_mask[0].annotations[0].polygon().is_none());
    }

    #[test]
    fn test_convert_image_annotations_object_id_from_lvis_large_id() {
        // LVIS v1.0 annotation IDs are u64 and routinely exceed 32-bit range
        // (the public release goes well past 2 billion). This test guards
        // against any future change that silently truncates on the path from
        // CocoAnnotation.id (u64) to Annotation.object_id (String).
        let image = CocoImage {
            id: 397133,
            width: 640,
            height: 480,
            file_name: "000000397133.jpg".to_string(),
            ..Default::default()
        };

        let large_id: u64 = 9_876_543_210;
        let dataset = CocoDataset {
            images: vec![image.clone()],
            categories: vec![CocoCategory {
                id: 16,
                name: "dog".to_string(),
                synset: Some("dog.n.01".to_string()),
                frequency: Some("f".to_string()),
                ..Default::default()
            }],
            annotations: vec![CocoAnnotation {
                id: large_id,
                image_id: 397133,
                category_id: 16,
                bbox: [192.81, 224.8, 74.73, 33.43],
                area: 1035.7,
                iscrowd: 0,
                segmentation: None,
                score: None,
            }],
            ..Default::default()
        };

        let index = CocoIndex::from_dataset(&dataset);
        let samples = convert_image_annotations(&image, &index, true, None);

        assert_eq!(samples.len(), 1);
        assert_eq!(samples[0].annotations.len(), 1);
        assert_eq!(
            samples[0].annotations[0].object_id(),
            Some(&large_id.to_string()),
        );
    }

    #[test]
    fn test_convert_image_annotations_no_annotations() {
        let image = CocoImage {
            id: 1,
            width: 640,
            height: 480,
            file_name: "empty.jpg".to_string(),
            ..Default::default()
        };

        let dataset = CocoDataset {
            images: vec![image.clone()],
            categories: vec![],
            annotations: vec![],
            ..Default::default()
        };

        let index = CocoIndex::from_dataset(&dataset);
        let samples = convert_image_annotations(&image, &index, true, None);

        // An image with no annotations must still emit one placeholder row so
        // the image is never dropped from the dataset.
        assert_eq!(samples.len(), 1);
        assert_eq!(samples[0].image_name, Some("empty".to_string()));
        assert!(samples[0].annotations.is_empty());
        assert_eq!(samples[0].group, None);
    }

    // =========================================================================
    // sample_name_from_filename tests
    // =========================================================================

    #[test]
    fn test_sample_name_from_filename() {
        assert_eq!(
            sample_name_from_filename("000000397133.jpg"),
            "000000397133"
        );
        assert_eq!(sample_name_from_filename("train2017/image.jpg"), "image");
        assert_eq!(sample_name_from_filename("test"), "test");
    }

    #[test]
    fn test_sample_name_from_filename_nested_path() {
        assert_eq!(
            sample_name_from_filename("a/b/c/deep_image.png"),
            "deep_image"
        );
    }

    #[test]
    fn test_sample_name_from_filename_no_extension() {
        assert_eq!(sample_name_from_filename("no_extension"), "no_extension");
    }

    // =========================================================================
    // Options tests
    // =========================================================================

    #[test]
    fn test_coco_to_arrow_options_default() {
        let options = CocoToArrowOptions::default();
        assert!(options.include_masks);
        assert!(options.group.is_none());
        assert!(options.max_workers >= 2);
    }

    #[test]
    fn test_arrow_to_coco_options_default() {
        let options = ArrowToCocoOptions::default();
        assert!(options.groups.is_empty());
        assert!(options.include_masks);
        assert!(options.info.is_none());
        assert!(!options.pretty);
    }

    #[test]
    fn test_max_workers() {
        let workers = max_workers();
        assert!(workers >= 2);
        assert!(workers <= 8);
    }

    #[tokio::test]
    async fn test_coco_to_arrow_minimal() {
        let temp_dir = TempDir::new().unwrap();

        // Create minimal COCO JSON
        let coco_json = r#"{
            "images": [
                {"id": 1, "width": 640, "height": 480, "file_name": "test.jpg"}
            ],
            "annotations": [
                {"id": 1, "image_id": 1, "category_id": 1, "bbox": [10, 20, 100, 80], "area": 8000, "iscrowd": 0}
            ],
            "categories": [
                {"id": 1, "name": "person", "supercategory": "human"}
            ]
        }"#;

        let coco_path = temp_dir.path().join("test.json");
        std::fs::write(&coco_path, coco_json).unwrap();

        let arrow_path = temp_dir.path().join("output.arrow");

        let options = CocoToArrowOptions::default();
        let count = coco_to_arrow(&coco_path, &arrow_path, &options, None)
            .await
            .unwrap();

        assert_eq!(count, 1);
        assert!(arrow_path.exists());

        // Verify Arrow contents
        let mut file = std::fs::File::open(&arrow_path).unwrap();
        let df = IpcReader::new(&mut file).finish().unwrap();
        assert_eq!(df.height(), 1);
    }

    #[tokio::test]
    async fn test_arrow_to_coco_roundtrip() {
        let temp_dir = TempDir::new().unwrap();

        // Create COCO JSON
        let original = CocoDataset {
            images: vec![CocoImage {
                id: 1,
                width: 640,
                height: 480,
                file_name: "test.jpg".to_string(),
                ..Default::default()
            }],
            annotations: vec![CocoAnnotation {
                id: 1,
                image_id: 1,
                category_id: 1,
                bbox: [100.0, 50.0, 200.0, 150.0],
                area: 30000.0,
                iscrowd: 0,
                segmentation: Some(CocoSegmentation::Polygon(vec![vec![
                    100.0, 50.0, 300.0, 50.0, 300.0, 200.0, 100.0, 200.0,
                ]])),
                score: None,
            }],
            categories: vec![CocoCategory {
                id: 1,
                name: "person".to_string(),
                supercategory: Some("human".to_string()),
                ..Default::default()
            }],
            ..Default::default()
        };

        // Write original COCO
        let coco_path = temp_dir.path().join("original.json");
        let writer = CocoWriter::new();
        writer.write_json(&original, &coco_path).unwrap();

        // Convert to Arrow
        let arrow_path = temp_dir.path().join("converted.arrow");
        let options = CocoToArrowOptions::default();
        coco_to_arrow(&coco_path, &arrow_path, &options, None)
            .await
            .unwrap();

        // Convert back to COCO
        let restored_path = temp_dir.path().join("restored.json");
        let options = ArrowToCocoOptions {
            pretty: true,
            ..Default::default()
        };
        arrow_to_coco(&arrow_path, &restored_path, &options, None)
            .await
            .unwrap();

        // Verify restored data
        let contents = std::fs::read_to_string(&restored_path).unwrap();
        assert!(
            contents.lines().count() > 1,
            "pretty output should span multiple lines"
        );
        let reader = CocoReader::new();
        let restored = reader.read_json(&restored_path).unwrap();

        assert_eq!(restored.images.len(), 1);
        assert_eq!(restored.annotations.len(), 1);
        assert_eq!(restored.categories.len(), 1);

        // Check category name preserved
        assert_eq!(restored.categories[0].name, "person");
    }

    #[tokio::test]
    async fn test_arrow_to_coco_roundtrip_preserves_annotation_id() {
        // Asserts the COCO/LVIS annotation `id` survives the full
        // JSON → Arrow → JSON round-trip. The IDs deliberately mix a
        // small value (1) with a 33-bit value (9_876_543_210) to catch
        // any future regression that silently truncates to u32 along
        // the path.
        let temp_dir = TempDir::new().unwrap();

        let large_id: u64 = 9_876_543_210;
        let original = CocoDataset {
            images: vec![CocoImage {
                id: 1,
                width: 640,
                height: 480,
                file_name: "test.jpg".to_string(),
                ..Default::default()
            }],
            annotations: vec![
                CocoAnnotation {
                    id: 1,
                    image_id: 1,
                    category_id: 1,
                    bbox: [10.0, 20.0, 100.0, 80.0],
                    area: 8000.0,
                    iscrowd: 0,
                    segmentation: None,
                    score: None,
                },
                CocoAnnotation {
                    id: large_id,
                    image_id: 1,
                    category_id: 1,
                    bbox: [200.0, 200.0, 100.0, 100.0],
                    area: 10000.0,
                    iscrowd: 0,
                    segmentation: None,
                    score: None,
                },
            ],
            categories: vec![CocoCategory {
                id: 1,
                name: "person".to_string(),
                supercategory: Some("human".to_string()),
                ..Default::default()
            }],
            ..Default::default()
        };

        let coco_path = temp_dir.path().join("original.json");
        let writer = CocoWriter::new();
        writer.write_json(&original, &coco_path).unwrap();

        let arrow_path = temp_dir.path().join("converted.arrow");
        coco_to_arrow(
            &coco_path,
            &arrow_path,
            &CocoToArrowOptions::default(),
            None,
        )
        .await
        .unwrap();

        let restored_path = temp_dir.path().join("restored.json");
        arrow_to_coco(
            &arrow_path,
            &restored_path,
            &ArrowToCocoOptions::default(),
            None,
        )
        .await
        .unwrap();

        let restored = CocoReader::new().read_json(&restored_path).unwrap();
        assert_eq!(restored.annotations.len(), 2);

        let restored_ids: std::collections::HashSet<u64> =
            restored.annotations.iter().map(|a| a.id).collect();
        assert!(
            restored_ids.contains(&1),
            "small annotation id (1) must round-trip; got {restored_ids:?}"
        );
        assert!(
            restored_ids.contains(&large_id),
            "33-bit LVIS-scale annotation id ({large_id}) must round-trip; got {restored_ids:?}"
        );
    }

    // =========================================================================
    // Arrow IPC file metadata tests
    // =========================================================================

    #[tokio::test]
    async fn test_coco_to_arrow_schema_version_metadata() {
        let temp_dir = TempDir::new().unwrap();

        // Create minimal COCO JSON (no LVIS fields)
        let coco_json = r#"{
            "images": [
                {"id": 1, "width": 640, "height": 480, "file_name": "test.jpg"}
            ],
            "annotations": [
                {"id": 1, "image_id": 1, "category_id": 1, "bbox": [10, 20, 100, 80], "area": 8000, "iscrowd": 0}
            ],
            "categories": [
                {"id": 1, "name": "person", "supercategory": "human"}
            ]
        }"#;

        let coco_path = temp_dir.path().join("test.json");
        std::fs::write(&coco_path, coco_json).unwrap();

        let arrow_path = temp_dir.path().join("output.arrow");
        let options = CocoToArrowOptions::default();
        coco_to_arrow(&coco_path, &arrow_path, &options, None)
            .await
            .unwrap();

        // Read back and verify schema_version metadata
        let mut file = std::fs::File::open(&arrow_path).unwrap();
        let mut reader = IpcReader::new(&mut file);
        let custom_meta = reader.custom_metadata().unwrap();
        assert!(custom_meta.is_some(), "custom metadata should be present");

        let meta = custom_meta.unwrap();
        assert_eq!(
            meta.get(&PlSmallStr::from("schema_version")),
            Some(&PlSmallStr::from(SCHEMA_VERSION)),
            "schema_version metadata should be '2026.04'"
        );

        // category_metadata is always present when there are categories
        assert!(
            meta.contains_key(&PlSmallStr::from("category_metadata")),
            "category_metadata should be present even without LVIS fields"
        );
    }

    #[tokio::test]
    async fn test_coco_to_arrow_category_metadata_lvis() {
        let temp_dir = TempDir::new().unwrap();

        // Create COCO JSON with LVIS category fields
        let coco_json = r#"{
            "images": [
                {"id": 1, "width": 640, "height": 480, "file_name": "test.jpg"}
            ],
            "annotations": [
                {"id": 1, "image_id": 1, "category_id": 1, "bbox": [10, 20, 100, 80], "area": 8000, "iscrowd": 0},
                {"id": 2, "image_id": 1, "category_id": 2, "bbox": [50, 60, 80, 40], "area": 3200, "iscrowd": 0}
            ],
            "categories": [
                {
                    "id": 1,
                    "name": "aerosol_can",
                    "synset": "aerosol.n.02",
                    "synonyms": ["aerosol_can", "spray_can"],
                    "def": "a dispenser that holds a substance under pressure"
                },
                {
                    "id": 2,
                    "name": "person",
                    "supercategory": "human"
                }
            ]
        }"#;

        let coco_path = temp_dir.path().join("lvis.json");
        std::fs::write(&coco_path, coco_json).unwrap();

        let arrow_path = temp_dir.path().join("lvis_output.arrow");
        let options = CocoToArrowOptions::default();
        coco_to_arrow(&coco_path, &arrow_path, &options, None)
            .await
            .unwrap();

        // Read back and verify metadata
        let mut file = std::fs::File::open(&arrow_path).unwrap();
        let mut reader = IpcReader::new(&mut file);
        let custom_meta = reader.custom_metadata().unwrap();
        assert!(custom_meta.is_some(), "custom metadata should be present");

        let meta = custom_meta.unwrap();

        // schema_version is always present
        assert_eq!(
            meta.get(&PlSmallStr::from("schema_version")),
            Some(&PlSmallStr::from(SCHEMA_VERSION)),
        );

        // category_metadata should be present (aerosol_can has LVIS fields)
        let cat_meta_str = meta
            .get(&PlSmallStr::from("category_metadata"))
            .expect("category_metadata should be present for LVIS data");

        let cat_meta: HashMap<String, serde_json::Value> =
            serde_json::from_str(cat_meta_str.as_str()).unwrap();

        // Both categories should be present (all categories are now stored)
        assert!(
            cat_meta.contains_key("aerosol_can"),
            "aerosol_can should be in category_metadata"
        );
        assert!(
            cat_meta.contains_key("person"),
            "person should also be in category_metadata"
        );

        // Verify aerosol_can entry contents
        let aerosol = cat_meta.get("aerosol_can").unwrap();
        assert_eq!(
            aerosol.get("synset").and_then(|v| v.as_str()),
            Some("aerosol.n.02")
        );
        assert_eq!(
            aerosol.get("definition").and_then(|v| v.as_str()),
            Some("a dispenser that holds a substance under pressure")
        );
        let synonyms = aerosol.get("synonyms").and_then(|v| v.as_array()).unwrap();
        assert_eq!(synonyms.len(), 2);
        assert_eq!(synonyms[0].as_str(), Some("aerosol_can"));
        assert_eq!(synonyms[1].as_str(), Some("spray_can"));
    }

    // =========================================================================
    // LVIS round-trip tests
    // =========================================================================

    #[tokio::test]
    async fn test_coco_arrow_roundtrip_lvis_supercategory() {
        let temp_dir = TempDir::new().unwrap();

        // Create COCO JSON with supercategory
        let coco_json = r#"{
            "images": [
                {"id": 1, "width": 640, "height": 480, "file_name": "test.jpg"}
            ],
            "annotations": [
                {"id": 1, "image_id": 1, "category_id": 1, "bbox": [10, 20, 100, 80], "area": 8000, "iscrowd": 0}
            ],
            "categories": [
                {"id": 1, "name": "person", "supercategory": "human"}
            ]
        }"#;

        let coco_path = temp_dir.path().join("original.json");
        std::fs::write(&coco_path, coco_json).unwrap();

        // Convert to Arrow
        let arrow_path = temp_dir.path().join("converted.arrow");
        let options = CocoToArrowOptions::default();
        coco_to_arrow(&coco_path, &arrow_path, &options, None)
            .await
            .unwrap();

        // Convert back to COCO
        let restored_path = temp_dir.path().join("restored.json");
        let options = ArrowToCocoOptions::default();
        arrow_to_coco(&arrow_path, &restored_path, &options, None)
            .await
            .unwrap();

        // Verify supercategory is preserved
        let reader = CocoReader::new();
        let restored = reader.read_json(&restored_path).unwrap();

        assert_eq!(restored.categories.len(), 1);
        assert_eq!(restored.categories[0].name, "person");
        assert_eq!(
            restored.categories[0].supercategory,
            Some("human".to_string()),
            "supercategory should survive COCO→Arrow→COCO round-trip"
        );
    }

    #[tokio::test]
    async fn test_coco_arrow_roundtrip_neg_categories_no_annotations() {
        let temp_dir = TempDir::new().unwrap();

        // Create COCO JSON: image has neg_category_ids but NO annotations
        let coco_json = r#"{
            "images": [
                {
                    "id": 1,
                    "width": 640,
                    "height": 480,
                    "file_name": "empty.jpg",
                    "neg_category_ids": [1, 2]
                }
            ],
            "annotations": [],
            "categories": [
                {"id": 1, "name": "cat", "supercategory": "animal"},
                {"id": 2, "name": "dog", "supercategory": "animal"}
            ]
        }"#;

        let coco_path = temp_dir.path().join("original.json");
        std::fs::write(&coco_path, coco_json).unwrap();

        // Convert to Arrow
        let arrow_path = temp_dir.path().join("converted.arrow");
        let options = CocoToArrowOptions::default();
        let sample_count = coco_to_arrow(&coco_path, &arrow_path, &options, None)
            .await
            .unwrap();

        // Should have 1 sentinel sample (image with neg data but no annotations)
        assert_eq!(
            sample_count, 1,
            "sentinel row should be emitted for image with neg data"
        );

        // Convert back to COCO
        let restored_path = temp_dir.path().join("restored.json");
        let options = ArrowToCocoOptions::default();
        arrow_to_coco(&arrow_path, &restored_path, &options, None)
            .await
            .unwrap();

        // Verify neg_category_ids survived the round-trip
        let reader = CocoReader::new();
        let restored = reader.read_json(&restored_path).unwrap();

        assert_eq!(restored.images.len(), 1);
        assert_eq!(restored.annotations.len(), 0, "no annotations expected");
        assert_eq!(restored.categories.len(), 2, "both categories should exist");

        let neg = restored.images[0].neg_category_ids.as_ref();
        assert!(
            neg.is_some(),
            "neg_category_ids should survive round-trip for zero-annotation image"
        );
        let neg_ids = neg.unwrap();
        assert_eq!(neg_ids.len(), 2, "should have 2 neg categories");
        assert!(neg_ids.contains(&1), "neg_category_ids should contain 1");
        assert!(neg_ids.contains(&2), "neg_category_ids should contain 2");

        // Verify supercategory survives for annotation-free categories
        for cat in &restored.categories {
            assert_eq!(
                cat.supercategory,
                Some("animal".to_string()),
                "supercategory should survive round-trip for annotation-free category '{}'",
                cat.name
            );
        }
    }

    #[test]
    fn test_convert_image_annotations_neg_only_no_annotations() {
        let image = CocoImage {
            id: 1,
            width: 640,
            height: 480,
            file_name: "neg_only.jpg".to_string(),
            neg_category_ids: Some(vec![1, 2]),
            ..Default::default()
        };

        let dataset = CocoDataset {
            images: vec![image.clone()],
            categories: vec![
                CocoCategory {
                    id: 1,
                    name: "cat".to_string(),
                    supercategory: Some("animal".to_string()),
                    ..Default::default()
                },
                CocoCategory {
                    id: 2,
                    name: "dog".to_string(),
                    supercategory: Some("animal".to_string()),
                    ..Default::default()
                },
            ],
            annotations: vec![],
            ..Default::default()
        };

        let index = CocoIndex::from_dataset(&dataset);
        let samples = convert_image_annotations(&image, &index, true, None);

        // Should emit 1 sentinel sample (no annotations but has neg data)
        assert_eq!(
            samples.len(),
            1,
            "sentinel row should be emitted for neg-only image"
        );
        assert_eq!(samples[0].image_name, Some("neg_only".to_string()));
        assert!(
            samples[0].annotations.is_empty(),
            "sentinel should have no annotations"
        );
        assert!(
            samples[0].neg_label_indices.is_some(),
            "sentinel should preserve neg_label_indices"
        );
        assert_eq!(samples[0].neg_label_indices.as_ref().unwrap().len(), 2);
    }

    #[test]
    fn test_convert_image_annotations_no_annotations_emits_placeholder() {
        // A plain image with NO annotations and NO LVIS neg/exhaustive fields
        // must still emit one placeholder sample so the image is never dropped
        // and its dataset split (group) is preserved.
        let image = CocoImage {
            id: 1,
            width: 640,
            height: 480,
            file_name: "empty.jpg".to_string(),
            ..Default::default()
        };

        let dataset = CocoDataset {
            images: vec![image.clone()],
            categories: vec![CocoCategory {
                id: 1,
                name: "person".to_string(),
                ..Default::default()
            }],
            annotations: vec![],
            ..Default::default()
        };

        let index = CocoIndex::from_dataset(&dataset);
        let samples = convert_image_annotations(&image, &index, true, Some("train"));

        assert_eq!(
            samples.len(),
            1,
            "placeholder row must be emitted for an unannotated image"
        );
        assert_eq!(samples[0].image_name, Some("empty".to_string()));
        assert!(
            samples[0].annotations.is_empty(),
            "placeholder should have no annotations"
        );
        assert_eq!(
            samples[0].group,
            Some("train".to_string()),
            "group must be preserved on the placeholder row"
        );
    }

    #[tokio::test]
    async fn test_coco_to_arrow_includes_unannotated_images() {
        // coco-to-arrow must emit one row per image even when some images have
        // no annotations, so dataset splits (group) cover every image.
        let temp_dir = TempDir::new().unwrap();

        let coco_json = r#"{
            "images": [
                {"id": 1, "width": 640, "height": 480, "file_name": "annotated.jpg"},
                {"id": 2, "width": 640, "height": 480, "file_name": "empty.jpg"}
            ],
            "annotations": [
                {"id": 1, "image_id": 1, "category_id": 1, "bbox": [10, 20, 100, 80], "area": 8000, "iscrowd": 0}
            ],
            "categories": [
                {"id": 1, "name": "person", "supercategory": "human"}
            ]
        }"#;

        let coco_path = temp_dir.path().join("test.json");
        std::fs::write(&coco_path, coco_json).unwrap();
        let arrow_path = temp_dir.path().join("out.arrow");

        let options = CocoToArrowOptions {
            group: Some("train".to_string()),
            ..Default::default()
        };
        let count = coco_to_arrow(&coco_path, &arrow_path, &options, None)
            .await
            .unwrap();

        // 1 annotation row + 1 placeholder row for the unannotated image.
        assert_eq!(count, 2, "every image must produce at least one row");

        let mut file = std::fs::File::open(&arrow_path).unwrap();
        let df = IpcReader::new(&mut file).finish().unwrap();
        assert_eq!(df.height(), 2);

        // The unannotated image must appear with its group set and a null label.
        let names = df.column("name").unwrap().str().unwrap();
        let empty_row = (0..df.height()).find(|&i| names.get(i) == Some("empty"));
        assert!(
            empty_row.is_some(),
            "unannotated image 'empty' must appear in the Arrow output"
        );
        let i = empty_row.unwrap();
        let group_col = df.column("group").unwrap().cast(&DataType::String).unwrap();
        assert_eq!(
            group_col.str().unwrap().get(i),
            Some("train"),
            "group must be set on the unannotated image's row"
        );
        let label_col = df.column("label").unwrap().cast(&DataType::String).unwrap();
        assert_eq!(
            label_col.str().unwrap().get(i),
            None,
            "unannotated image row must have a null label"
        );
    }

    // =========================================================================
    // Parquet output tests
    // =========================================================================

    const MINIMAL_COCO_JSON: &str = r#"{
        "images": [
            {"id": 1, "width": 640, "height": 480, "file_name": "test1.jpg"},
            {"id": 2, "width": 320, "height": 240, "file_name": "test2.jpg"}
        ],
        "annotations": [
            {"id": 1, "image_id": 1, "category_id": 1, "bbox": [10, 20, 100, 80], "area": 8000, "iscrowd": 0},
            {"id": 2, "image_id": 2, "category_id": 2, "bbox": [5, 5, 50, 40], "area": 2000, "iscrowd": 0}
        ],
        "categories": [
            {"id": 1, "name": "person", "supercategory": "human"},
            {"id": 2, "name": "cat", "supercategory": "animal"}
        ]
    }"#;

    #[tokio::test]
    async fn coco_to_parquet_roundtrip() {
        let dir = TempDir::new().unwrap();
        let coco = dir.path().join("instances.json");
        std::fs::write(&coco, MINIMAL_COCO_JSON).unwrap();
        let out = dir.path().join("out.parquet");

        let n = coco_to_arrow(&coco, &out, &CocoToArrowOptions::default(), None)
            .await
            .unwrap();
        assert_eq!(n, 2);

        let file = std::fs::File::open(&out).unwrap();
        let mut reader = ParquetReader::new(file);

        // KV metadata: read the parquet footer key-value pairs.
        let metadata = reader.get_metadata().unwrap().clone();
        let kv = metadata
            .key_value_metadata
            .as_ref()
            .expect("parquet footer must carry key-value metadata");
        let get = |key: &str| {
            kv.iter()
                .find(|e| e.key == key)
                .and_then(|e| e.value.clone())
        };
        assert_eq!(
            get("schema_version").as_deref(),
            Some(SCHEMA_VERSION),
            "schema_version metadata should be '2026.04'"
        );
        assert!(get("labels").is_some(), "labels metadata should be present");
        assert!(
            get("category_metadata").is_some(),
            "category_metadata metadata should be present"
        );

        let df = reader.finish().unwrap();
        assert_eq!(df.height(), 2);
        let cols = df.get_column_names();
        assert!(cols.iter().any(|c| c.as_str() == "name"));
        assert!(cols.iter().any(|c| c.as_str() == "label"));
        assert!(cols.iter().any(|c| c.as_str() == "box2d"));
    }

    #[tokio::test]
    async fn parquet_to_coco_roundtrip() {
        let dir = TempDir::new().unwrap();
        let coco = dir.path().join("instances.json");
        std::fs::write(&coco, MINIMAL_COCO_JSON).unwrap();
        let parquet = dir.path().join("out.parquet");

        let n = coco_to_arrow(&coco, &parquet, &CocoToArrowOptions::default(), None)
            .await
            .unwrap();
        assert_eq!(n, 2);

        let roundtrip_json = dir.path().join("roundtrip.json");
        let converted = arrow_to_coco(
            &parquet,
            &roundtrip_json,
            &ArrowToCocoOptions::default(),
            None,
        )
        .await
        .unwrap();
        assert_eq!(converted, 2, "arrow_to_coco must accept a .parquet input");

        let output: serde_json::Value =
            serde_json::from_str(&std::fs::read_to_string(&roundtrip_json).unwrap()).unwrap();

        let categories = output["categories"].as_array().unwrap();
        assert_eq!(categories.len(), 2, "categories must round-trip");
        let category_names: std::collections::HashSet<_> = categories
            .iter()
            .map(|c| c["name"].as_str().unwrap().to_string())
            .collect();
        assert!(category_names.contains("person"));
        assert!(category_names.contains("cat"));

        let annotations = output["annotations"].as_array().unwrap();
        assert_eq!(annotations.len(), 2, "annotation count must round-trip");
    }

    #[tokio::test]
    async fn coco_directory_combines_splits_for_arrow_and_parquet() {
        let dir = TempDir::new().unwrap();
        let coco_root = dir.path().join("coco");
        let annotations = coco_root.join("annotations");
        std::fs::create_dir_all(&annotations).unwrap();
        std::fs::create_dir_all(coco_root.join("train2017")).unwrap();
        #[cfg(unix)]
        {
            std::fs::create_dir_all(coco_root.join("val-images")).unwrap();
            std::os::unix::fs::symlink(coco_root.join("val-images"), coco_root.join("val2017"))
                .unwrap();
        }
        #[cfg(not(unix))]
        std::fs::create_dir_all(coco_root.join("val2017")).unwrap();

        // Deliberately reuse image and annotation IDs across splits. Processing
        // each source independently must retain both rather than deduplicating
        // one split by numeric ID.
        let split_json = |file_name: &str, annotated: bool| {
            let annotations = if annotated {
                r#"[{"id": 1, "image_id": 1, "category_id": 1, "bbox": [10, 20, 100, 80], "area": 8000, "iscrowd": 0}]"#
            } else {
                "[]"
            };
            format!(
                r#"{{
                    "images": [{{"id": 1, "width": 640, "height": 480, "file_name": "{file_name}"}}],
                    "annotations": {annotations},
                    "categories": [{{"id": 1, "name": "person", "supercategory": "human"}}]
                }}"#
            )
        };
        std::fs::write(
            annotations.join("instances_train2017.json"),
            split_json("train.jpg", true),
        )
        .unwrap();
        std::fs::write(
            annotations.join("instances_val2017.json"),
            split_json("val.jpg", false),
        )
        .unwrap();
        std::fs::write(coco_root.join("train2017/train.jpg"), b"train").unwrap();
        std::fs::write(coco_root.join("val2017/val.jpg"), b"val").unwrap();

        for extension in ["arrow", "parquet"] {
            let dataset_name = format!("combined_{extension}");
            let output_dir = dir.path().join(&dataset_name);
            let output = output_dir.join(format!("{dataset_name}.{extension}"));
            let options = CocoToArrowOptions {
                images_dir: Some(coco_root.clone()),
                link_images: true,
                ..Default::default()
            };
            let count = coco_to_arrow(&coco_root, &output, &options, None)
                .await
                .unwrap();
            assert_eq!(count, 2, "annotated and empty images must both be retained");

            let (df, metadata) = crate::format::read_dataset_dataframe(&output).unwrap();
            assert_eq!(df.height(), 2);
            assert_eq!(
                metadata.get("schema_version").map(String::as_str),
                Some(SCHEMA_VERSION)
            );

            let groups = df.column("group").unwrap().cast(&DataType::String).unwrap();
            let groups: std::collections::HashSet<_> =
                groups.str().unwrap().iter().flatten().collect();
            assert_eq!(groups, std::collections::HashSet::from(["train", "val"]));

            let names = df.column("name").unwrap().str().unwrap();
            assert!(names.iter().flatten().any(|name| name == "train"));
            assert!(names.iter().flatten().any(|name| name == "val"));

            let staged = output_dir.join(&dataset_name);
            assert_eq!(std::fs::read(staged.join("train.jpg")).unwrap(), b"train");
            assert_eq!(std::fs::read(staged.join("val.jpg")).unwrap(), b"val");
            assert!(
                crate::format::validate_dataset_structure(&output_dir)
                    .unwrap()
                    .is_empty(),
                "linked image staging should pass offline validation"
            );
        }
    }

    #[tokio::test]
    async fn coco_directory_rejects_incompatible_split_categories() {
        let dir = TempDir::new().unwrap();
        let annotations = dir.path().join("annotations");
        std::fs::create_dir_all(&annotations).unwrap();
        let make_json = |category: &str| {
            format!(
                r#"{{
                    "images": [{{"id": 1, "width": 1, "height": 1, "file_name": "image.jpg"}}],
                    "annotations": [],
                    "categories": [{{"id": 1, "name": "{category}"}}]
                }}"#
            )
        };
        std::fs::write(
            annotations.join("instances_train2017.json"),
            make_json("person"),
        )
        .unwrap();
        std::fs::write(
            annotations.join("instances_val2017.json"),
            make_json("vehicle"),
        )
        .unwrap();

        let output = dir.path().join("output.arrow");
        let error = coco_to_arrow(dir.path(), &output, &CocoToArrowOptions::default(), None)
            .await
            .unwrap_err();
        assert!(
            error
                .to_string()
                .contains("COCO splits disagree on category 1")
        );
        assert!(!output.exists());
    }

    #[test]
    fn compatible_categories_reject_conflicting_persisted_metadata() {
        let category = CocoCategory {
            id: 1,
            name: "person".to_string(),
            supercategory: Some("human".to_string()),
            synset: Some("person.n.01".to_string()),
            frequency: Some("f".to_string()),
            synonyms: Some(vec!["person".to_string(), "human".to_string()]),
            def: Some("a human being".to_string()),
            ..Default::default()
        };
        let conflicting = [
            (
                "supercategory",
                CocoCategory {
                    supercategory: Some("animal".to_string()),
                    ..category.clone()
                },
            ),
            (
                "synset",
                CocoCategory {
                    synset: Some("person.n.02".to_string()),
                    ..category.clone()
                },
            ),
            (
                "frequency",
                CocoCategory {
                    frequency: Some("c".to_string()),
                    ..category.clone()
                },
            ),
            (
                "synonyms",
                CocoCategory {
                    synonyms: Some(vec!["individual".to_string()]),
                    ..category.clone()
                },
            ),
            (
                "definition",
                CocoCategory {
                    def: Some("a different definition".to_string()),
                    ..category.clone()
                },
            ),
        ];

        for (field, conflicting_category) in conflicting {
            let sources = vec![
                (
                    CocoDataset {
                        categories: vec![category.clone()],
                        ..Default::default()
                    },
                    Some("train".to_string()),
                    None,
                ),
                (
                    CocoDataset {
                        categories: vec![conflicting_category],
                        ..Default::default()
                    },
                    Some("val".to_string()),
                    None,
                ),
            ];

            let error = collect_compatible_categories(&sources).unwrap_err();
            assert!(
                error
                    .to_string()
                    .contains(&format!("metadata field '{field}'")),
                "unexpected error for {field}: {error}"
            );
        }
    }

    // =========================================================================
    // Image staging tests
    // =========================================================================

    const STAGING_COCO_JSON: &str = r#"{
        "images": [
            {"id": 1, "width": 640, "height": 480, "file_name": "img1.jpg"},
            {"id": 2, "width": 320, "height": 240, "file_name": "img2.jpg"}
        ],
        "annotations": [],
        "categories": []
    }"#;

    #[tokio::test]
    async fn staging_copies_referenced_images_and_reports_missing() {
        let dir = TempDir::new().unwrap();
        let src = dir.path().join("src");
        std::fs::create_dir(&src).unwrap();
        std::fs::write(src.join("img1.jpg"), b"jpegdata").unwrap();

        let coco = dir.path().join("instances.json");
        std::fs::write(&coco, STAGING_COCO_JSON).unwrap();

        let out = dir.path().join("ds/ds.arrow");
        let options = CocoToArrowOptions {
            images_dir: Some(src),
            ..Default::default()
        };

        coco_to_arrow(&coco, &out, &options, None).await.unwrap();

        assert!(dir.path().join("ds/ds/img1.jpg").exists());
        assert!(
            !dir.path().join("ds/ds/img2.jpg").exists(),
            "missing source images must be skipped non-fatally"
        );
    }

    #[tokio::test]
    async fn staging_links_when_requested() {
        let dir = TempDir::new().unwrap();
        let src = dir.path().join("src");
        std::fs::create_dir(&src).unwrap();
        std::fs::write(src.join("img1.jpg"), b"jpegdata").unwrap();

        let coco = dir.path().join("instances.json");
        std::fs::write(&coco, STAGING_COCO_JSON).unwrap();

        let out = dir.path().join("ds/ds.arrow");
        let options = CocoToArrowOptions {
            images_dir: Some(src),
            link_images: true,
            ..Default::default()
        };

        coco_to_arrow(&coco, &out, &options, None).await.unwrap();

        let staged = dir.path().join("ds/ds/img1.jpg");
        #[cfg(unix)]
        assert!(
            std::fs::symlink_metadata(&staged)
                .unwrap()
                .file_type()
                .is_symlink(),
            "--link must produce a symlink, not a copy"
        );
        assert!(staged.exists());
    }

    #[tokio::test]
    async fn staging_is_idempotent_on_rerun() {
        let dir = TempDir::new().unwrap();
        let src = dir.path().join("src");
        std::fs::create_dir(&src).unwrap();
        std::fs::write(src.join("img1.jpg"), b"jpegdata").unwrap();

        let coco = dir.path().join("instances.json");
        std::fs::write(&coco, STAGING_COCO_JSON).unwrap();

        let out = dir.path().join("ds/ds.arrow");
        let options = CocoToArrowOptions {
            images_dir: Some(src),
            ..Default::default()
        };

        coco_to_arrow(&coco, &out, &options, None).await.unwrap();
        coco_to_arrow(&coco, &out, &options, None)
            .await
            .expect("re-running staging over an already-staged output must not error");

        let staged = dir.path().join("ds/ds/img1.jpg");
        assert!(staged.exists());
        assert_eq!(
            std::fs::read(&staged).unwrap(),
            b"jpegdata",
            "re-run must leave exactly one, unmodified copy"
        );
    }

    #[test]
    fn staging_reports_collision_on_duplicate_basename_from_different_sources() {
        let dir = TempDir::new().unwrap();
        let src = dir.path().join("src");
        std::fs::create_dir_all(src.join("train")).unwrap();
        std::fs::create_dir_all(src.join("val")).unwrap();
        std::fs::write(src.join("train/000001.jpg"), b"train-data").unwrap();
        std::fs::write(src.join("val/000001.jpg"), b"val-data").unwrap();

        let out = dir.path().join("ds/ds.arrow");
        let file_names = vec!["train/000001.jpg".to_string(), "val/000001.jpg".to_string()];

        let report = stage_images(&out, &src, false, &file_names);

        assert_eq!(
            report.staged, 1,
            "only the first source claiming the basename must be staged"
        );
        assert_eq!(
            report.collisions, 1,
            "the second file_name must be counted as a collision, not staged"
        );
        assert!(report.missing.is_empty());

        let dest = dir.path().join("ds/ds/000001.jpg");
        assert_eq!(
            std::fs::read(&dest).unwrap(),
            b"train-data",
            "the destination must retain the first-staged source's bytes, never the second's"
        );
    }
}