From d7dfe35f10ad1e4be7f78dab30ef4dc3779560d5 Mon Sep 17 00:00:00 2001 From: prrao87 Date: Tue, 10 Nov 2026 14:21:16 +0000 Subject: [PATCH] perf(query): evaluate CONTAINS, STARTS WITH, ENDS WITH with Arrow kernels The string predicate UDF broadcast a literal search string to a full array and called str::contains per row, building a new substring searcher each time. Call Arrow's contains/starts_with/ends_with kernels instead: a literal stays a scalar and one searcher serves the batch. These are the kernels DataFusion's LIKE and contains() use. Bound parameters, null and non-string operands keep their Cypher semantics. Fixes #146 --- .../datafusion_planner/scalar_functions.rs | 41 +++++++++---------- python/python/tests/test_string_parameters.py | 5 +++ 2 files changed, 25 insertions(+), 21 deletions(-) diff --git a/crates/lance-csr/src/datafusion_planner/scalar_functions.rs b/crates/lance-csr/src/datafusion_planner/scalar_functions.rs index 8f8bdbc..526393e 100644 --- a/crates/lance-csr/src/datafusion_planner/scalar_functions.rs +++ b/crates/lance-csr/src/datafusion_planner/scalar_functions.rs @@ -7,9 +7,9 @@ use std::sync::Arc; use arrow::array::{ - new_null_array, Array, AsArray, BooleanArray, Int64Array, ListBuilder, StringArray, - StringBuilder, + new_null_array, Array, AsArray, Datum, Int64Array, ListBuilder, StringArray, StringBuilder, }; +use arrow::compute::kernels::comparison; use arrow::compute::{cast, cast_with_options, CastOptions}; use arrow_schema::{DataType, Field}; use datafusion::common::{exec_err, plan_err, ScalarValue}; @@ -268,6 +268,24 @@ impl ScalarUDFImpl for CypherScalar { Ok(result) }; } + if matches!( + self.name.as_str(), + "__string_contains" | "__string_starts_with" | "__string_ends_with" + ) { + // Arrow's kernels build one substring searcher per batch, and a literal + // pattern stays a scalar instead of being repeated for every row. + let text = cast(&args.args[0].to_array(args.number_rows)?, &DataType::Utf8)?; + let pattern: Box = match &args.args[1] { + ColumnarValue::Scalar(p) => Box::new(p.cast_to(&DataType::Utf8)?.to_scalar()?), + ColumnarValue::Array(p) => Box::new(cast(p, &DataType::Utf8)?), + }; + let result = match self.name.as_str() { + "__string_contains" => comparison::contains(&text, pattern.as_ref()), + "__string_starts_with" => comparison::starts_with(&text, pattern.as_ref()), + _ => comparison::ends_with(&text, pattern.as_ref()), + }?; + return Ok(ColumnarValue::Array(Arc::new(result))); + } let arrays = args .args .iter() @@ -282,25 +300,6 @@ impl ScalarUDFImpl for CypherScalar { }) .collect::, _>>()?; let strings = arrays[0].as_string::(); - if matches!( - self.name.as_str(), - "__string_contains" | "__string_starts_with" | "__string_ends_with" - ) { - let patterns = arrays[1].as_string::(); - let result: BooleanArray = strings - .iter() - .zip(patterns.iter()) - .map(|(text, pattern)| { - text.zip(pattern) - .map(|(text, pattern)| match self.name.as_str() { - "__string_contains" => text.contains(pattern), - "__string_starts_with" => text.starts_with(pattern), - _ => text.ends_with(pattern), - }) - }) - .collect(); - return Ok(ColumnarValue::Array(Arc::new(result))); - } if self.name == "size" { return Ok(ColumnarValue::Array(Arc::new( strings diff --git a/python/python/tests/test_string_parameters.py b/python/python/tests/test_string_parameters.py index 4dd88d5..c72f1e1 100644 --- a/python/python/tests/test_string_parameters.py +++ b/python/python/tests/test_string_parameters.py @@ -49,3 +49,8 @@ def test_bound_string_operands(csr: bool) -> None: assert engine.execute( "MATCH (p:Person {id: $id}) RETURN p.id AS id", {"id": 2} ).to_pylist() == [{"id": 2}] + # The search string can also be a column, one value per row. + assert engine.execute( + "MATCH (p:Person), (q:Person) WHERE q.id IN [2, 3] " + "AND p.name CONTAINS q.name RETURN p.id AS id, q.id AS q ORDER BY id, q" + ).to_pylist() == [{"id": 1, "q": 2}, {"id": 2, "q": 2}]