diff --git a/polars_io_tools/io_sources/sql_utils.py b/polars_io_tools/io_sources/sql_utils.py index 74051a1..50dec87 100644 --- a/polars_io_tools/io_sources/sql_utils.py +++ b/polars_io_tools/io_sources/sql_utils.py @@ -296,6 +296,14 @@ def _handle_string_function(self, node: FunctionNode, input_exprs: list[sqlglot. self.result = sqlglot.exp.Upper(this=input_exprs[0]) elif node.function_type == StringFunctionType.LOWERCASE: self.result = sqlglot.exp.Lower(this=input_exprs[0]) + elif node.function_type == StringFunctionType.STRIP_CHARS: + self.result = sqlglot.exp.Trim(this=input_exprs[0]) + elif node.function_type == StringFunctionType.STRIP_CHARS_START: + self.result = sqlglot.exp.Trim(this=input_exprs[0], position="LEADING") + elif node.function_type == StringFunctionType.STRIP_CHARS_END: + self.result = sqlglot.exp.Trim(this=input_exprs[0], position="TRAILING") + elif node.function_type == StringFunctionType.CONCAT_HORIZONTAL and input_exprs: + self.result = sqlglot.exp.Concat(expressions=input_exprs) else: log.warning(f"Unsupported string function: {node.function_type}") self.result = None diff --git a/polars_io_tools/tests/io_sources/test_base.py b/polars_io_tools/tests/io_sources/test_base.py index 70152d8..7ba0a40 100644 --- a/polars_io_tools/tests/io_sources/test_base.py +++ b/polars_io_tools/tests/io_sources/test_base.py @@ -1088,6 +1088,30 @@ def test_string_predicates(tester): tester.assert_predicate_pushed_down(expr3) +def test_strip_chars_predicate_pushed_down(): + """str.strip_chars / str.strip_chars_start / str.strip_chars_end translate to TRIM / LTRIM / RTRIM.""" + df = pl.DataFrame({"padded": [" A ", " B ", " C ", " D ", "E"]}) + tracker = PredicateTracker(df) + + for expr in [ + pl.col("padded").str.strip_chars() == "A", + pl.col("padded").str.strip_chars_start() == "A ", + pl.col("padded").str.strip_chars_end() == " A", + ]: + tracker.assert_predicate_pushed_down(expr) + tracker.assert_results_match(expr) + + +def test_concat_horizontal_predicate_pushed_down(): + """pl.concat_str translates to SQL CONCAT(...).""" + df = pl.DataFrame({"first": ["A", "B", "C"], "last": ["X", "Y", "Z"]}) + tracker = PredicateTracker(df) + + expr = pl.concat_str([pl.col("first"), pl.lit("-"), pl.col("last")]) == "A-X" + tracker.assert_predicate_pushed_down(expr) + tracker.assert_results_match(expr) + + def test_fill_null_predicates(tester): for val in [0, 1, 9]: expr = pl.col("nullable").fill_null(val) > 5