Spaces:

Brianferrell787
/

range

Running on Zero

RANGE: clean package identity, example taxonomy API, Gradio 6.27.0

c4aee46 verified about 1 month ago

12.7 kB

// 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(())
}

	// 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(())
	}