rust/lib.rs

// The Python-list and Arrow entry points share one exact scanning core and ship
// in a single private extension module.
include!("core.rs");
mod curation;

use arrow_array::{
    Array, ArrayRef, BooleanArray, LargeStringArray, RecordBatch, StringArray,
    TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
    TimestampSecondArray, UInt32Array, UInt64Array, UInt8Array,
};
use arrow_schema::{DataType, Field, Schema, TimeUnit};
use pyo3_arrow::input::AnyRecordBatch;
use pyo3_arrow::PyRecordBatch;
use std::sync::{Arc, Mutex, OnceLock};

static ARROW_POOLS: OnceLock<Mutex<HashMap<usize, Arc<rayon::ThreadPool>>>> = OnceLock::new();

fn worker_pool(workers: usize) -> PyResult<Arc<rayon::ThreadPool>> {
    if workers == 0 {
        return Err(PyValueError::new_err("workers must be >= 1"));
    }
    let pools = ARROW_POOLS.get_or_init(|| Mutex::new(HashMap::new()));
    if let Some(pool) = pools
        .lock()
        .map_err(|_| PyValueError::new_err("Arrow worker-pool cache poisoned"))?
        .get(&workers)
        .cloned()
    {
        return Ok(pool);
    }

    let pool = Arc::new(
        ThreadPoolBuilder::new()
            .num_threads(workers)
            .thread_name(|index| format!("news-relevance-arrow-{index}"))
            .build()
            .map_err(|error| PyValueError::new_err(error.to_string()))?,
    );
    let mut guard = pools
        .lock()
        .map_err(|_| PyValueError::new_err("Arrow worker-pool cache poisoned"))?;
    Ok(guard.entry(workers).or_insert_with(|| pool.clone()).clone())
}

enum TextColumn<'a> {
    Utf8(&'a StringArray),
    LargeUtf8(&'a LargeStringArray),
}

impl<'a> TextColumn<'a> {
    #[inline(always)]
    fn value(&self, index: usize) -> &'a str {
        match self {
            Self::Utf8(values) => values.value(index),
            Self::LargeUtf8(values) => values.value(index),
        }
    }
}

enum TimestampColumn<'a> {
    Second(&'a TimestampSecondArray),
    Millisecond(&'a TimestampMillisecondArray),
    Microsecond(&'a TimestampMicrosecondArray),
    Nanosecond(&'a TimestampNanosecondArray),
}

impl<'a> TimestampColumn<'a> {
    #[inline(always)]
    fn value_microseconds(&self, index: usize) -> Result<i64, String> {
        match self {
            Self::Second(values) => values
                .value(index)
                .checked_mul(1_000_000)
                .ok_or_else(|| "timestamp second value overflows microseconds".to_string()),
            Self::Millisecond(values) => values
                .value(index)
                .checked_mul(1_000)
                .ok_or_else(|| "timestamp millisecond value overflows microseconds".to_string()),
            Self::Microsecond(values) => Ok(values.value(index)),
            Self::Nanosecond(values) => Ok(values.value(index) / 1_000),
        }
    }
}

fn text_column(array: &dyn Array) -> Result<TextColumn<'_>, String> {
    if array.null_count() != 0 {
        return Err("Arrow text column must not contain nulls".to_string());
    }
    match array.data_type() {
        DataType::Utf8 => Ok(TextColumn::Utf8(
            array
                .as_any()
                .downcast_ref::<StringArray>()
                .ok_or_else(|| "failed to downcast Utf8 text column".to_string())?,
        )),
        DataType::LargeUtf8 => Ok(TextColumn::LargeUtf8(
            array
                .as_any()
                .downcast_ref::<LargeStringArray>()
                .ok_or_else(|| "failed to downcast LargeUtf8 text column".to_string())?,
        )),
        other => Err(format!(
            "Arrow text column must be Utf8 or LargeUtf8, got {other}"
        )),
    }
}

fn timestamp_column(array: &dyn Array) -> Result<TimestampColumn<'_>, String> {
    if array.null_count() != 0 {
        return Err("Arrow timestamp column must not contain nulls".to_string());
    }
    let (unit, timezone) = match array.data_type() {
        DataType::Timestamp(unit, timezone) => (unit, timezone),
        other => {
            return Err(format!(
                "Arrow timestamp column must be a timestamp type, got {other}"
            ))
        }
    };
    if timezone.is_none() {
        return Err(
            "Arrow-native fast path requires a timezone-aware timestamp column; \

             use America/New_York (or another explicit timezone) to preserve \

             exact effective-date semantics"
                .to_string(),
        );
    }
    match unit {
        TimeUnit::Second => Ok(TimestampColumn::Second(
            array
                .as_any()
                .downcast_ref::<TimestampSecondArray>()
                .ok_or_else(|| "failed to downcast second timestamp column".to_string())?,
        )),
        TimeUnit::Millisecond => Ok(TimestampColumn::Millisecond(
            array
                .as_any()
                .downcast_ref::<TimestampMillisecondArray>()
                .ok_or_else(|| "failed to downcast millisecond timestamp column".to_string())?,
        )),
        TimeUnit::Microsecond => Ok(TimestampColumn::Microsecond(
            array
                .as_any()
                .downcast_ref::<TimestampMicrosecondArray>()
                .ok_or_else(|| "failed to downcast microsecond timestamp column".to_string())?,
        )),
        TimeUnit::Nanosecond => Ok(TimestampColumn::Nanosecond(
            array
                .as_any()
                .downcast_ref::<TimestampNanosecondArray>()
                .ok_or_else(|| "failed to downcast nanosecond timestamp column".to_string())?,
        )),
    }
}

fn facts_to_record_batch(rows: Vec<FactRow>) -> Result<RecordBatch, String> {
    let mut row_index = Vec::with_capacity(rows.len());
    let mut article_status = Vec::with_capacity(rows.len());
    let mut sp500 = Vec::with_capacity(rows.len());
    let mut nasdaq100 = Vec::with_capacity(rows.len());
    let mut russell2000 = Vec::with_capacity(rows.len());
    let mut market_wide = Vec::with_capacity(rows.len());
    let mut observation_count = Vec::with_capacity(rows.len());
    let mut native_candidate = Vec::with_capacity(rows.len());
    let mut exact_candidate = Vec::with_capacity(rows.len());

    for (index, article, sp, ndx, rut, market, observations, native, exact) in rows {
        row_index.push(index as u64);
        article_status.push(article);
        sp500.push(sp);
        nasdaq100.push(ndx);
        russell2000.push(rut);
        market_wide.push(market);
        observation_count.push(observations as u32);
        native_candidate.push(native);
        exact_candidate.push(exact);
    }

    let fields = vec![
        Field::new("row_index", DataType::UInt64, false),
        Field::new("article_status", DataType::UInt8, false),
        Field::new("sp500", DataType::UInt8, false),
        Field::new("nasdaq100", DataType::UInt8, false),
        Field::new("russell2000", DataType::UInt8, false),
        Field::new("market_wide", DataType::UInt8, false),
        Field::new("observation_count", DataType::UInt32, false),
        Field::new("native_candidate", DataType::Boolean, false),
        Field::new("exact_candidate", DataType::Boolean, false),
    ];
    let columns: Vec<ArrayRef> = vec![
        Arc::new(UInt64Array::from(row_index)),
        Arc::new(UInt8Array::from(article_status)),
        Arc::new(UInt8Array::from(sp500)),
        Arc::new(UInt8Array::from(nasdaq100)),
        Arc::new(UInt8Array::from(russell2000)),
        Arc::new(UInt8Array::from(market_wide)),
        Arc::new(UInt32Array::from(observation_count)),
        Arc::new(BooleanArray::from(native_candidate)),
        Arc::new(BooleanArray::from(exact_candidate)),
    ];
    RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).map_err(|error| error.to_string())
}

#[pyclass(frozen, module = "range_curate._native")]
struct NativeArrowCompactFactScanner {
    inner: NativeCompactFactScanner,
}

#[pymethods]
impl NativeArrowCompactFactScanner {
    #[new]
    #[pyo3(signature = (

        patterns,

        pattern_consumers,

        consumers,

        exact_specs,

        alias_specs,

        lexical_specs,

        context_rules,

        conjunction_specs

    ))]
    fn new(
        patterns: Vec<String>,
        pattern_consumers: Vec<Vec<usize>>,
        consumers: Vec<(u8, usize, usize, usize)>,
        exact_specs: Vec<(String, String, u8, usize)>,
        alias_specs: Vec<(usize, i64, Option<i64>, i8, u8, bool, Option<usize>)>,
        lexical_specs: Vec<(usize, Vec<String>, bool, bool, bool, bool)>,
        context_rules: Vec<(usize, Vec<Vec<String>>, Vec<String>)>,
        conjunction_specs: Vec<(usize, usize, usize, usize)>,
    ) -> PyResult<Self> {
        Ok(Self {
            inner: NativeCompactFactScanner::new(
                patterns,
                pattern_consumers,
                consumers,
                exact_specs,
                alias_specs,
                lexical_specs,
                context_rules,
                conjunction_specs,
            )?,
        })
    }

    #[pyo3(signature = (data, text_column_name="text", timestamp_column_name="timestamp", workers=4))]
    fn screen_arrow(
        &self,
        py: Python<'_>,
        data: AnyRecordBatch,
        text_column_name: &str,
        timestamp_column_name: &str,
        workers: usize,
    ) -> PyResult<PyRecordBatch> {
        if workers == 0 {
            return Err(PyValueError::new_err("workers must be >= 1"));
        }
        let schema = data.schema()?;
        let text_index = schema
            .index_of(text_column_name)
            .map_err(|error| PyValueError::new_err(error.to_string()))?;
        let timestamp_index = schema
            .index_of(timestamp_column_name)
            .map_err(|error| PyValueError::new_err(error.to_string()))?;
        let mut reader = data.into_reader()?;

        let pool = worker_pool(workers)?;
        let inner = &self.inner;

        let result: Result<Vec<FactRow>, String> = py.detach(move || {
            let mut output = Vec::new();
            let mut row_base = 0usize;
            for batch_result in &mut reader {
                let batch = batch_result.map_err(|error| error.to_string())?;
                let texts = text_column(batch.column(text_index).as_ref())?;
                let timestamps = timestamp_column(batch.column(timestamp_index).as_ref())?;
                let batch_rows: Result<Vec<FactRow>, String> = pool.install(|| {
                    (0..batch.num_rows())
                        .into_par_iter()
                        .map_init(ScanScratch::default, |scratch, index| {
                            let epoch = timestamps.value_microseconds(index)?;
                            Ok(scan_one(
                                &inner.automaton,
                                &inner.pattern_consumers,
                                &inner.consumers,
                                &inner.exact_specs,
                                &inner.aliases,
                                &inner.lexical,
                                &inner.contexts,
                                &inner.conjunctions,
                                texts.value(index),
                                epoch,
                                row_base + index,
                                scratch,
                            ))
                        })
                        .filter_map(|value: Result<Option<FactRow>, String>| match value {
                            Ok(Some(row)) => Some(Ok(row)),
                            Ok(None) => None,
                            Err(error) => Some(Err(error)),
                        })
                        .collect()
                });
                output.extend(batch_rows?);
                row_base += batch.num_rows();
            }
            output.sort_unstable_by_key(|row| row.0);
            Ok(output)
        });

        let record_batch = facts_to_record_batch(result.map_err(PyValueError::new_err)?)
            .map_err(PyValueError::new_err)?;
        Ok(PyRecordBatch::new(record_batch))
    }
}

#[pymodule]
fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> {
    module.add_class::<curation::NativeCurationScanner>()?;
    module.add_class::<NativeCompactFactScanner>()?;
    module.add_class::<NativeArrowCompactFactScanner>()?;
    Ok(())
}