Merge pull request #310 from moj-analytical-services/issue_309

Parse StructType columns in case expressions
moj-analytical-services · Mar 2, 2022 · ff503de · ff503de
2 parents cedcca5 + 2af239f
commit ff503de
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "splink"
-version = "2.1.8"
+version = "2.1.9"
 description = "Implementation of Fellegi-Sunter's canonical model of record linkage in Apache Spark, including EM algorithm to estimate parameters"
 authors = ["Robin Linacre <robinlinacre@hotmail.com>", "Sam Lindsay", "Theodore Manassis"]
 license = "MIT"

diff --git a/splink/parse_case_statement.py b/splink/parse_case_statement.py
@@ -1,6 +1,6 @@
 import sqlglot
 from sqlglot.errors import ParseError
-from sqlglot.expressions import Case, Column, Alias
+from sqlglot.expressions import Case, Column, Alias, Bracket
 import re
 
 
@@ -11,6 +11,8 @@ def get_columns_used_from_sql(sql):
         subtree = tup[0]
         if type(subtree) is Column:
             column_names.add(subtree.sql())
+        elif type(subtree) is Bracket:
+            column_names.add(subtree.text("this"))
 
     return list(column_names)