Repository: tobymao/sqlglot

Issue:
Substring type inference missing binary case
**Fully reproducible code snippet**


```python
>>> optimize(sqlglot.parse_one("select col from tbl",read='spark'),schema={"tbl": {"col": "BINARY"}}).expressions[0].type
DataType(this=Type.BINARY, nested=False)
>>> optimize(sqlglot.parse_one("select substring(col, 2, 3) as x from tbl",read='spark'),schema={"tbl": {"col": "BINARY"}}).expressions[0].type
DataType(this=Type.VARCHAR)
```

In the second expression, the expected output type is `Type.BINARY`. In spark shell, this works fine (spark 3.5.1):

```scala
scala> val df = spark.sql("select cast('thing' as binary)")
df: org.apache.spark.sql.DataFrame = [CAST(thing AS BINARY): binary]

scala> val df = spark.sql("select substring(cast('thing' as binary), 1, 3)")
df: org.apache.spark.sql.DataFrame = [substring(CAST(thing AS BINARY), 1, 3): binary]

```

From my understanding, the return data type for `Substring` is hardcoded to VARCHAR: https://github.com/tobymao/sqlglot/blob/fcaae87d50653600048917945652cdf463bcc3cf/sqlglot/dialects/dialect.py#L609-L617


Discussion:

Patch:
diff --git a/sqlglot/dialects/spark2.py b/sqlglot/dialects/spark2.py
index 017a6f6d..00daca06 100644
--- a/sqlglot/dialects/spark2.py
+++ b/sqlglot/dialects/spark2.py
@@ -111,6 +111,11 @@ def temporary_storage_provider(expression: exp.Expression) -> exp.Expression:
 
 
 class Spark2(Hive):
+    ANNOTATORS = {
+        **Hive.ANNOTATORS,
+        exp.Substring: lambda self, e: self._annotate_by_args(e, "this"),
+    }
+
     class Parser(Hive.Parser):
         TRIM_PATTERN_FIRST = True
 


Test patch:
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index d6e11a9c..fe5a4d7c 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -1345,3 +1345,26 @@ FROM READ_CSV('tests/fixtures/optimizer/tpc-h/nation.csv.gz', 'delimiter', '|')
         self.assertEqual(4, normalization_distance(gen_expr(2), max_=100))
         self.assertEqual(18, normalization_distance(gen_expr(3), max_=100))
         self.assertEqual(110, normalization_distance(gen_expr(10), max_=100))
+
+    def test_custom_annotators(self):
+        # In Spark hierarchy, SUBSTRING result type is dependent on input expr type
+        for dialect in ("spark2", "spark", "databricks"):
+            for expr_type_pair in (
+                ("col", "STRING"),
+                ("col", "BINARY"),
+                ("'str_literal'", "STRING"),
+                ("CAST('str_literal' AS BINARY)", "BINARY"),
+            ):
+                with self.subTest(
+                    f"Testing {dialect}'s SUBSTRING() result type for {expr_type_pair}"
+                ):
+                    expr, type = expr_type_pair
+                    ast = parse_one(f"SELECT substring({expr}, 2, 3) AS x FROM tbl", read=dialect)
+
+                    subst_type = (
+                        optimizer.optimize(ast, schema={"tbl": {"col": type}}, dialect=dialect)
+                        .expressions[0]
+                        .type
+                    )
+
+                    self.assertEqual(subst_type.sql(dialect), exp.DataType.build(type).sql(dialect))