pipeline B 文本归一化优化-Daft-AtomGit

Msrc/daft-functions-utf8/src/lib.rs+3-0

	@@ -13,6 +13,7 @@ mod lower;
13	mod lpad;	13	mod lpad;
14	mod lstrip;	14	mod lstrip;
15	mod normalize;	15	mod normalize;
		16	+mod normalize_text;
16	pub(crate) mod pad;	17	pub(crate) mod pad;
17	mod regexp_count;	18	mod regexp_count;
18	mod regexp_extract;	19	mod regexp_extract;
	@@ -45,6 +46,7 @@ pub use lower::*;
45	pub use lpad::*;	46	pub use lpad::*;
46	pub use lstrip::*;	47	pub use lstrip::*;
47	pub use normalize::*;	48	pub use normalize::*;
		49	+pub use normalize_text::*;
48	pub use regexp_count::*;	50	pub use regexp_count::*;
49	pub use regexp_extract::*;	51	pub use regexp_extract::*;
50	pub use regexp_extract_all::*;	52	pub use regexp_extract_all::*;
	@@ -78,6 +80,7 @@ impl daft_dsl::functions::FunctionModule for Utf8Functions {
78	parent.add_fn(Lower);	80	parent.add_fn(Lower);
79	parent.add_fn(LPad);	81	parent.add_fn(LPad);
80	parent.add_fn(LStrip);	82	parent.add_fn(LStrip);
		83	+ parent.add_fn(NormalizeText);
81	parent.add_fn(Normalize);	84	parent.add_fn(Normalize);
82	parent.add_fn(RegexpCount);	85	parent.add_fn(RegexpCount);
83	parent.add_fn(RegexpExtract);	86	parent.add_fn(RegexpExtract);

Asrc/daft-functions-utf8/src/normalize_text.rs+265-0

@@ -0,0 +1,265 @@
	1	+use common_error::DaftResult;
	2	+use daft_core::{
	3	+ prelude::{AsArrow, DataType, Field, Schema, Utf8Array},
	4	+ series::{IntoSeries, Series},
	5	+};
	6	+use daft_dsl::{
	7	+ ExprRef,
	8	+ functions::{FunctionArgs, ScalarUDF, scalar::ScalarFn},
	9	+};
	10	+use serde::{Deserialize, Serialize};
	11	+
	12	+use crate::utils::{unary_utf8_evaluate, unary_utf8_to_field};
	13	+
	14	+#[derive(Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
	15	+pub struct NormalizeText;
	16	+
	17	+#[typetag::serde]
	18	+impl ScalarUDF for NormalizeText {
	19	+ fn name(&self) -> &'static str {
	20	+ "normalize_text"
	21	+ }
	22	+
	23	+ fn call(&self, inputs: FunctionArgs<Series>) -> DaftResult<Series> {
	24	+ unary_utf8_evaluate(inputs, \|series\| {
	25	+ series.with_utf8_array(\|arr\| normalize_text_impl(arr).map(IntoSeries::into_series))
	26	+ })
	27	+ }
	28	+
	29	+ fn get_return_field(
	30	+ &self,
	31	+ inputs: FunctionArgs<ExprRef>,
	32	+ schema: &Schema,
	33	+ ) -> DaftResult<Field> {
	34	+ unary_utf8_to_field(inputs, schema, self.name(), DataType::Utf8)
	35	+ }
	36	+
	37	+ fn docstring(&self) -> &'static str {
	38	+ "Internal fused strip, lowercase, and whitespace normalization"
	39	+ }
	40	+}
	41	+
	42	+#[must_use]
	43	+pub fn normalize_text(input: ExprRef) -> ExprRef {
	44	+ ScalarFn::builtin(NormalizeText, vec![input]).into()
	45	+}
	46	+
	47	+#[inline]
	48	+fn is_regex_whitespace_ascii(byte: u8) -> bool {
	49	+ matches!(byte, b' ' \| b'\t' \| b'\n' \| b'\r' \| 0x0c)
	50	+}
	51	+
	52	+#[inline]
	53	+fn trim_edge_literal_s(input: &str) -> &str {
	54	+ let bytes = input.as_bytes();
	55	+ let mut start = 0;
	56	+ if bytes.len() >= 2 && bytes[0] == b'\\' && bytes[1] == b's' {
	57	+ start = 2;
	58	+ while start < bytes.len() && bytes[start] == b's' {
	59	+ start += 1;
	60	+ }
	61	+ }
	62	+
	63	+ let mut end = bytes.len();
	64	+ if end > start && bytes[end - 1] == b's' {
	65	+ let mut run_start = end - 1;
	66	+ while run_start > start && bytes[run_start - 1] == b's' {
	67	+ run_start -= 1;
	68	+ }
	69	+ if run_start > start && bytes[run_start - 1] == b'\\' {
	70	+ end = run_start - 1;
	71	+ }
	72	+ }
	73	+ &input[start..end]
	74	+}
	75	+
	76	+#[inline]
	77	+fn append_ascii_normalized(input: &str, output: &mut Vec<u8>) {
	78	+ let input = trim_edge_literal_s(input);
	79	+ let mut pending_space = false;
	80	+ for &byte in input.as_bytes() {
	81	+ if is_regex_whitespace_ascii(byte) {
	82	+ pending_space = true;
	83	+ continue;
	84	+ }
	85	+ if pending_space {
	86	+ output.push(b' ');
	87	+ }
	88	+ output.push(byte.to_ascii_lowercase());
	89	+ pending_space = false;
	90	+ }
	91	+ if pending_space {
	92	+ output.push(b' ');
	93	+ }
	94	+}
	95	+
	96	+#[inline]
	97	+fn append_unicode_normalized(input: &str, output: &mut Vec<u8>) {
	98	+ let input = trim_edge_literal_s(input);
	99	+ let mut pending_space = false;
	100	+
	101	+ for character in input.chars() {
	102	+ if character.is_whitespace() {
	103	+ pending_space = true;
	104	+ continue;
	105	+ }
	106	+
	107	+ if pending_space {
	108	+ output.push(b' ');
	109	+ pending_space = false;
	110	+ }
	111	+
	112	+ for lowercase in character.to_lowercase() {
	113	+ let mut encoded = [0; 4];
	114	+ output.extend_from_slice(lowercase.encode_utf8(&mut encoded).as_bytes());
	115	+ }
	116	+ }
	117	+
	118	+ if pending_space {
	119	+ output.push(b' ');
	120	+ }
	121	+}
	122	+
	123	+fn normalize_text_impl(arr: &Utf8Array) -> DaftResult<Utf8Array> {
	124	+ let arrow = arr.as_arrow2();
	125	+ let mut offsets = daft_arrow::offset::Offsets::<i64>::with_capacity(arr.len());
	126	+ let mut values = Vec::with_capacity(arrow.values().len());
	127	+ let mut validity = daft_arrow::bitmap::MutableBitmap::with_capacity(arr.len());
	128	+
	129	+ for value in arr.into_iter() {
	130	+ let value_start = values.len();
	131	+ match value {
	132	+ None => validity.push(false),
	133	+ Some(value) => {
	134	+ validity.push(true);
	135	+ if value.is_ascii() {
	136	+ append_ascii_normalized(value, &mut values);
	137	+ } else {
	138	+ append_unicode_normalized(value, &mut values);
	139	+ }
	140	+ }
	141	+ }
	142	+ offsets.try_push_usize(values.len() - value_start)?;
	143	+ }
	144	+
	145	+ let validity: Option<daft_arrow::bitmap::Bitmap> = validity.into();
	146	+ let output = daft_arrow::array::Utf8Array::<i64>::new(
	147	+ arrow.data_type().clone(),
	148	+ offsets.into(),
	149	+ values.into(),
	150	+ validity,
	151	+ );
	152	+ Utf8Array::new(arr.field().clone().into(), Box::new(output))
	153	+}
	154	+
	155	+#[cfg(test)]
	156	+mod tests {
	157	+ use std::sync::LazyLock;
	158	+
	159	+ use super::*;
	160	+
	161	+ static SCHEME_A_EDGE_LITERAL_S_RE: LazyLock<regex::Regex> =
	162	+ LazyLock::new(\|\| regex::Regex::new(r"^\\s+\|\\s+$").unwrap());
	163	+ static SCHEME_A_WHITESPACE_RE: LazyLock<regex::Regex> =
	164	+ LazyLock::new(\|\| regex::Regex::new(r"\s+").unwrap());
	165	+
	166	+ fn scheme_a_regex_reference(input: &str) -> String {
	167	+ let stripped = SCHEME_A_EDGE_LITERAL_S_RE.replace_all(input, "");
	168	+ let lowered = stripped.to_lowercase();
	169	+ SCHEME_A_WHITESPACE_RE
	170	+ .replace_all(&lowered, " ")
	171	+ .into_owned()
	172	+ }
	173	+
	174	+ fn assert_matches_scheme_a(cases: &[String]) {
	175	+ let input = Utf8Array::from_iter("text", cases.iter().map(\|value\| Some(value.as_str())));
	176	+ let actual = normalize_text_impl(&input).unwrap();
	177	+
	178	+ for (idx, original) in cases.iter().enumerate() {
	179	+ let expected = scheme_a_regex_reference(original);
	180	+ assert_eq!(
	181	+ actual.get(idx),
	182	+ Some(expected.as_str()),
	183	+ "input {original:?}"
	184	+ );
	185	+ }
	186	+ }
	187	+
	188	+ #[test]
	189	+ fn normalizes_ascii_and_unicode_like_the_unfused_expression() {
	190	+ let input = Utf8Array::from_iter(
	191	+ "text",
	192	+ vec![
	193	+ Some(" HELLO\t\tWORLD "),
	194	+ Some("already normalized"),
	195	+ Some("\n\r\t"),
	196	+ Some(" ÉCOLE\u{2003}TEST "),
	197	+ Some("İSTANBUL"),
	198	+ Some(""),
	199	+ None,
	200	+ ]
	201	+ .into_iter(),
	202	+ );
	203	+ let actual = normalize_text_impl(&input).unwrap();
	204	+ let expected = [
	205	+ Some(" hello world "),
	206	+ Some("already normalized"),
	207	+ Some(" "),
	208	+ Some(" école test "),
	209	+ Some("i\u{307}stanbul"),
	210	+ Some(""),
	211	+ None,
	212	+ ];
	213	+ for (idx, expected) in expected.into_iter().enumerate() {
	214	+ assert_eq!(actual.get(idx), expected, "row {idx}");
	215	+ }
	216	+ }
	217	+
	218	+ #[test]
	219	+ fn ascii_whitespace_set_matches_regex_semantics() {
	220	+ let input = Utf8Array::from_iter("text", vec![Some(" A\t\n\r\x0b\x0c B ")].into_iter());
	221	+ let actual = normalize_text_impl(&input).unwrap();
	222	+ assert_eq!(actual.get(0), Some(" a \x0b b "));
	223	+ }
	224	+
	225	+ #[test]
	226	+ fn removes_only_the_edge_literal_backslash_s_pattern() {
	227	+ let input = Utf8Array::from_iter(
	228	+ "text",
	229	+ vec![Some(r"\sssHELLO\ss"), Some(r"keep\ssinside")].into_iter(),
	230	+ );
	231	+ let actual = normalize_text_impl(&input).unwrap();
	232	+ assert_eq!(actual.get(0), Some("hello"));
	233	+ assert_eq!(actual.get(1), Some(r"keep\ssinside"));
	234	+ }
	235	+
	236	+ #[test]
	237	+ fn matches_scheme_a_for_unicode_and_boundary_cases() {
	238	+ const UNICODE_WHITE_SPACE: [char; 25] = [
	239	+ '\u{0009}', '\u{000a}', '\u{000b}', '\u{000c}', '\u{000d}', '\u{0020}', '\u{0085}',
	240	+ '\u{00a0}', '\u{1680}', '\u{2000}', '\u{2001}', '\u{2002}', '\u{2003}', '\u{2004}',
	241	+ '\u{2005}', '\u{2006}', '\u{2007}', '\u{2008}', '\u{2009}', '\u{200a}', '\u{2028}',
	242	+ '\u{2029}', '\u{202f}', '\u{205f}', '\u{3000}',
	243	+ ];
	244	+
	245	+ let mut all_white_space = String::from("É");
	246	+ for whitespace in UNICODE_WHITE_SPACE {
	247	+ all_white_space.push(whitespace);
	248	+ all_white_space.push('X');
	249	+ }
	250	+
	251	+ assert_matches_scheme_a(&[
	252	+ all_white_space,
	253	+ "É\u{180e}\u{200b}\u{200c}\u{2060}\u{feff}COLE".into(),
	254	+ " ÉCOLE\u{2003}\u{2009}TEST ".into(),
	255	+ "İSTANBUL".into(),
	256	+ " 中文\u{3000}測試 ".into(),
	257	+ " 😀\t🚀 ".into(),
	258	+ " E\u{0301}\tA\u{0308} ".into(),
	259	+ r"\sssÉCOLE\ss".into(),
	260	+ r"keep\ssinside É".into(),
	261	+ "É\t \u{2003}\n\u{00a0}COLE".into(),
	262	+ "".into(),
	263	+ ]);
	264	+ }
	265	+}

Msrc/daft-logical-plan/Cargo.toml+1-1

	@@ -14,6 +14,7 @@ daft-core = {path = "../daft-core", default-features = false}
14	daft-dsl = {path = "../daft-dsl", default-features = false}	14	daft-dsl = {path = "../daft-dsl", default-features = false}
15	daft-functions = {path = "../daft-functions", default-features = false}	15	daft-functions = {path = "../daft-functions", default-features = false}
16	daft-functions-list = {path = "../daft-functions-list", default-features = false}	16	daft-functions-list = {path = "../daft-functions-list", default-features = false}
		17	+daft-functions-utf8 = {path = "../daft-functions-utf8", default-features = false}
17	daft-functions-uri = {path = "../daft-functions-uri", default-features = false}	18	daft-functions-uri = {path = "../daft-functions-uri", default-features = false}
18	daft-recordbatch = {path = "../daft-recordbatch", default-features = false}	19	daft-recordbatch = {path = "../daft-recordbatch", default-features = false}
19	daft-schema = {path = "../daft-schema", default-features = false}	20	daft-schema = {path = "../daft-schema", default-features = false}
	@@ -33,7 +34,6 @@ uuid.workspace = true
33	[dev-dependencies]	34	[dev-dependencies]
34	daft-dsl = {path = "../daft-dsl", features = ["test-utils"]}	35	daft-dsl = {path = "../daft-dsl", features = ["test-utils"]}
35	daft-functions-binary = {path = "../daft-functions-binary", default-features = false}	36	daft-functions-binary = {path = "../daft-functions-binary", default-features = false}
36	-daft-functions-utf8 = {path = "../daft-functions-utf8", default-features = false}
37	indoc = "2"	37	indoc = "2"
38	pretty_assertions = {workspace = true}	38	pretty_assertions = {workspace = true}
39	rand = "0.8"	39	rand = "0.8"

Msrc/daft-logical-plan/src/optimization/optimizer.rs+10-6

	@@ -9,12 +9,12 @@ use super::{
9	rules::{	9	rules::{
10	DetectMonotonicId, DropIntoBatches, DropRepartition, EliminateCrossJoin, EliminateOffsets,	10	DetectMonotonicId, DropIntoBatches, DropRepartition, EliminateCrossJoin, EliminateOffsets,
11	EliminateSubqueryAliasRule, EnrichWithStats, ExtractWindowFunction, FilterNullJoinKey,	11	EliminateSubqueryAliasRule, EnrichWithStats, ExtractWindowFunction, FilterNullJoinKey,
12	- LiftProjectFromAgg, MaterializeScans, OptimizerRule, PushDownAggregation,	12	+ FuseTextNormalization, LiftProjectFromAgg, MaterializeScans, OptimizerRule,
13	- PushDownAntiSemiJoin, PushDownFilter, PushDownJoinPredicate, PushDownLimit,	13	+ PushDownAggregation, PushDownAntiSemiJoin, PushDownFilter, PushDownJoinPredicate,
14	- PushDownProjection, PushDownShard, ReorderJoins, RewriteCountDistinct, RewriteOffset,	14	+ PushDownLimit, PushDownProjection, PushDownShard, ReorderJoins, RewriteCountDistinct,
15	- ShardScans, SimplifyExpressionsRule, SimplifyNullFilteredJoin, SplitExplodeFromProject,	15	+ RewriteOffset, ShardScans, SimplifyExpressionsRule, SimplifyNullFilteredJoin,
16	- SplitGranularProjection, SplitUDFs, SplitUDFsFromFilters, UnnestPredicateSubquery,	16	+ SplitExplodeFromProject, SplitGranularProjection, SplitUDFs, SplitUDFsFromFilters,
17	- UnnestScalarSubquery,	17	+ UnnestPredicateSubquery, UnnestScalarSubquery,
18	},	18	},
19	};	19	};
20	use crate::{LogicalPlan, optimization::rules::SplitVLLM};	20	use crate::{LogicalPlan, optimization::rules::SplitVLLM};
	@@ -141,6 +141,10 @@ impl OptimizerBuilder {
141	vec![Box::new(SimplifyExpressionsRule::new())],	141	vec![Box::new(SimplifyExpressionsRule::new())],
142	RuleExecutionStrategy::FixedPoint(None),	142	RuleExecutionStrategy::FixedPoint(None),
143	),	143	),
		144	+ RuleBatch::new(
		145	+ vec![Box::new(FuseTextNormalization::new())],
		146	+ RuleExecutionStrategy::FixedPoint(None),
		147	+ ),
144	// --- Filter out null join keys ---	148	// --- Filter out null join keys ---
145	// This rule should be run once, before any filter pushdown rules.	149	// This rule should be run once, before any filter pushdown rules.
146	RuleBatch::new(	150	RuleBatch::new(

Asrc/daft-logical-plan/src/optimization/rules/fuse_text_normalization.rs+120-0

@@ -0,0 +1,120 @@
	1	+use std::sync::Arc;
	2	+
	3	+use common_error::DaftResult;
	4	+use common_treenode::{Transformed, TreeNode};
	5	+use daft_core::lit::Literal;
	6	+use daft_dsl::{
	7	+ Expr, ExprRef,
	8	+ functions::{BuiltinScalarFn, ScalarUDF, scalar::ScalarFn},
	9	+};
	10	+use daft_functions_utf8::{Lower, RegexpReplace, normalize_text};
	11	+
	12	+use super::OptimizerRule;
	13	+use crate::LogicalPlan;
	14	+
	15	+#[derive(Default, Debug)]
	16	+pub struct FuseTextNormalization;
	17	+
	18	+impl FuseTextNormalization {
	19	+ pub fn new() -> Self {
	20	+ Self
	21	+ }
	22	+}
	23	+
	24	+fn as_builtin<F: ScalarUDF>(expr: &ExprRef) -> Option<&BuiltinScalarFn> {
	25	+ match expr.as_ref() {
	26	+ Expr::ScalarFn(ScalarFn::Builtin(function)) if function.is_function_type::<F>() => {
	27	+ Some(function)
	28	+ }
	29	+ _ => None,
	30	+ }
	31	+}
	32	+
	33	+fn arg(function: &BuiltinScalarFn, index: usize) -> Option<&ExprRef> {
	34	+ function.inputs.iter().nth(index).map(\|arg\| arg.inner())
	35	+}
	36	+
	37	+fn is_utf8_literal(expr: &ExprRef, expected: &str) -> bool {
	38	+ matches!(expr.as_ref(), Expr::Literal(Literal::Utf8(value)) if value == expected)
	39	+}
	40	+
	41	+fn fuse_text_normalization(expr: ExprRef) -> DaftResult<Transformed<ExprRef>> {
	42	+ let Some(outer) = as_builtin::<RegexpReplace>(&expr) else {
	43	+ return Ok(Transformed::no(expr));
	44	+ };
	45	+ let (Some(lower_expr), Some(outer_pattern), Some(outer_replacement)) =
	46	+ (arg(outer, 0), arg(outer, 1), arg(outer, 2))
	47	+ else {
	48	+ return Ok(Transformed::no(expr));
	49	+ };
	50	+ let Some(lower) = as_builtin::<Lower>(lower_expr) else {
	51	+ return Ok(Transformed::no(expr));
	52	+ };
	53	+ let Some(inner_expr) = arg(lower, 0) else {
	54	+ return Ok(Transformed::no(expr));
	55	+ };
	56	+ let Some(inner) = as_builtin::<RegexpReplace>(inner_expr) else {
	57	+ return Ok(Transformed::no(expr));
	58	+ };
	59	+ let (Some(input), Some(inner_pattern), Some(inner_replacement)) =
	60	+ (arg(inner, 0), arg(inner, 1), arg(inner, 2))
	61	+ else {
	62	+ return Ok(Transformed::no(expr));
	63	+ };
	64	+
	65	+ let exact = is_utf8_literal(outer_pattern, r"\s+")
	66	+ && is_utf8_literal(outer_replacement, " ")
	67	+ && is_utf8_literal(inner_pattern, r"^\\s+\|\\s+$")
	68	+ && is_utf8_literal(inner_replacement, "");
	69	+ if exact {
	70	+ Ok(Transformed::yes(normalize_text(input.clone())))
	71	+ } else {
	72	+ Ok(Transformed::no(expr))
	73	+ }
	74	+}
	75	+
	76	+impl OptimizerRule for FuseTextNormalization {
	77	+ fn try_optimize(&self, plan: Arc<LogicalPlan>) -> DaftResult<Transformed<Arc<LogicalPlan>>> {
	78	+ plan.transform(\|plan\| {
	79	+ plan.map_expressions(\|expr, _schema\| expr.transform_up(fuse_text_normalization))
	80	+ })
	81	+ }
	82	+}
	83	+
	84	+#[cfg(test)]
	85	+mod tests {
	86	+ use daft_dsl::{lit, resolved_col};
	87	+ use daft_functions_utf8::{NormalizeText, lower, replace};
	88	+
	89	+ use super::*;
	90	+
	91	+ fn pipeline_normalization(input: ExprRef) -> ExprRef {
	92	+ replace(
	93	+ lower(replace(input, lit(r"^\\s+\|\\s+$"), lit(""), true)),
	94	+ lit(r"\s+"),
	95	+ lit(" "),
	96	+ true,
	97	+ )
	98	+ }
	99	+
	100	+ #[test]
	101	+ fn fuses_exact_pipeline_expression() {
	102	+ let input = resolved_col("text");
	103	+ let transformed = fuse_text_normalization(pipeline_normalization(input.clone())).unwrap();
	104	+ assert!(transformed.transformed);
	105	+ let function = as_builtin::<NormalizeText>(&transformed.data).unwrap();
	106	+ assert_eq!(arg(function, 0), Some(&input));
	107	+ }
	108	+
	109	+ #[test]
	110	+ fn rejects_different_pattern_or_replacement() {
	111	+ let input = resolved_col("text");
	112	+ let expressions = [
	113	+ replace(lower(input.clone()), lit(r"\s*"), lit(" "), true),
	114	+ replace(lower(input), lit(r"\s+"), lit("_"), true),
	115	+ ];
	116	+ for expression in expressions {
	117	+ assert!(!fuse_text_normalization(expression).unwrap().transformed);
	118	+ }
	119	+ }
	120	+}

Msrc/daft-logical-plan/src/optimization/rules/mod.rs+2-0

	@@ -7,6 +7,7 @@ mod eliminate_subquery_alias;
7	mod enrich_with_stats;	7	mod enrich_with_stats;
8	mod extract_window_function;	8	mod extract_window_function;
9	mod filter_null_join_key;	9	mod filter_null_join_key;
		10	+mod fuse_text_normalization;
10	mod granular_projections;	11	mod granular_projections;
11	mod lift_project_from_agg;	12	mod lift_project_from_agg;
12	mod materialize_scans;	13	mod materialize_scans;
	@@ -37,6 +38,7 @@ pub use eliminate_subquery_alias::EliminateSubqueryAliasRule;
37	pub use enrich_with_stats::EnrichWithStats;	38	pub use enrich_with_stats::EnrichWithStats;
38	pub use extract_window_function::ExtractWindowFunction;	39	pub use extract_window_function::ExtractWindowFunction;
39	pub use filter_null_join_key::FilterNullJoinKey;	40	pub use filter_null_join_key::FilterNullJoinKey;
		41	+pub use fuse_text_normalization::FuseTextNormalization;
40	pub use granular_projections::SplitGranularProjection;	42	pub use granular_projections::SplitGranularProjection;
41	pub use lift_project_from_agg::LiftProjectFromAgg;	43	pub use lift_project_from_agg::LiftProjectFromAgg;
42	pub use materialize_scans::MaterializeScans;	44	pub use materialize_scans::MaterializeScans;