support lambda expression in presto

nooberfsh · Nov 30, 2021 · Dec 8, 2021 · Jan 19, 2022 · Dec 22, 2021 · Dec 28, 2021
commit 2c700dbaf16bc73b1585ddbb3cbeec5b0716334e
diff --git a/dialects/presto/src/Database/Sql/Presto/Parser.hs b/dialects/presto/src/Database/Sql/Presto/Parser.hs
@@ -438,6 +438,11 @@ columnAliasP = do
     (name, r) <- Tok.columnNameP
     makeColumnAlias r name
 
+lambdaParamP :: Parser (LambdaParam Range)
+lambdaParamP = do
+    (name, r) <- Tok.lambdaParamP
+    makeLambdaParam r name
+
 joinP :: Parser (Tablish RawNames Range -> Tablish RawNames Range)
 joinP = crossJoinP <|> regularJoinP <|> naturalJoinP
   where
@@ -696,7 +701,9 @@ selectionP idx = try selectStarP <|> do
             (name, r) <- Tok.columnNameP
             makeColumnAlias r name
 
-        , makeExprAlias expr idx'
+        , case expr of
+            LambdaExpr {} -> fail "Lambda expression should always be used inside a function"
+            _ -> makeExprAlias expr idx'
         ]
 
 countingSepBy1 :: (Integer -> Parser b) -> Parser c -> Parser [b]
@@ -720,6 +727,9 @@ makeTableAlias r alias = TableAlias r alias . TableAliasId <$> getNextCounter
 makeColumnAlias :: Range -> Text -> Parser (ColumnAlias Range)
 makeColumnAlias r alias = ColumnAlias r alias . ColumnAliasId <$> getNextCounter
 
+makeLambdaParam :: Range -> Text -> Parser (LambdaParam Range)
+makeLambdaParam r name = LambdaParam r name . LambdaParamId <$> getNextCounter
+
 makeDummyAlias :: Range -> Integer -> Parser (ColumnAlias Range)
 makeDummyAlias r idx = makeColumnAlias r $ TL.pack $ "_col" ++ show idx
 
@@ -743,6 +753,8 @@ makeExprAlias (FieldAccessExpr info _ _) idx = makeDummyAlias info idx
 makeExprAlias (ArrayAccessExpr info _ _) idx = makeDummyAlias info idx
 makeExprAlias (TypeCastExpr _ _ expr _) idx = makeExprAlias expr idx
 makeExprAlias (VariableSubstitutionExpr _) _ = fail "Unsupported variable substitution in Presto: unused expr-type in this dialect"
+makeExprAlias LambdaParamExpr {} _ = error "Unreachable, unresolved expr can not be lambda param"
+makeExprAlias LambdaExpr {} _ = error "Unreachable, selection parser should reject lambda expression"
 
 
 unOpP :: Text -> Parser (Expr RawNames Range -> Expr RawNames Range)
@@ -856,6 +868,7 @@ primaryExprP = foldl (flip ($)) <$> baseP <*> many (arrayAccessP <|> structAcces
     baseP = choice
         [ extractPrimaryExprP
         , normalizePrimaryExprP
+        , try lambdaP
         , try substringPrimaryExprP -- try is because `substring` is both a special-form function and a regular function
         , try positionPrimaryExprP -- try is because `position` could be a column name / UDF function name
         , bareFuncPrimaryExprP
@@ -1125,6 +1138,22 @@ functionCallPrimaryExprP = do
         r' <- Tok.closeP
         return $ Filter (r <> r') expr
 
+lambdaP :: Parser (Expr RawNames Range) 
+lambdaP = do
+    (params, start) <- choice
+        [ do
+            s <- Tok.openP
+            params <- lambdaParamP `sepBy` Tok.commaP
+            _ <- Tok.closeP
+            return (params, s)
+        , do
+            p <- lambdaParamP  
+            return ([p], getInfo p)
+        ]
+    _ <- Tok.symbolP "->"
+    body <- exprP
+    return $ LambdaExpr (start <> getInfo body) params body
+
 overP :: Parser (OverSubExpr RawNames Range)
 overP = do
     start <- Tok.overP

diff --git a/dialects/presto/src/Database/Sql/Presto/Parser/Token.hs b/dialects/presto/src/Database/Sql/Presto/Parser/Token.hs
@@ -308,6 +308,16 @@ columnNameP = P.tokenPrim showTok posFromTok testTok
 
         _ -> Nothing
 
+lambdaParamP :: Parser (Text, Range)
+lambdaParamP = P.tokenPrim showTok posFromTok testTok
+  where
+    testTok (tok, s, e) = case tok of
+        TokWord True name -> Just (name, Range s e)
+        TokWord False name
+            | wordCanBeColumnName (wordInfo name) -> Just (name, Range s e)
+
+        _ -> Nothing
+
 structFieldNameP :: Parser (Text, Range)
 structFieldNameP = P.tokenPrim showTok posFromTok testTok
   where

diff --git a/dialects/presto/src/Database/Sql/Presto/Scanner.hs b/dialects/presto/src/Database/Sql/Presto/Scanner.hs
@@ -57,6 +57,7 @@ operators = sortBy (flip compare)
     , "||"
     , "(", ")", "[", "]", ",", ";"
     , "?"
+    , "->"
     ]
 
 isOperator :: Char -> Bool

diff --git a/dialects/vertica/src/Database/Sql/Vertica/Parser.hs b/dialects/vertica/src/Database/Sql/Vertica/Parser.hs
@@ -1022,8 +1022,6 @@ makeExprAlias (InSubqueryExpr info _ _) = makeDummyAlias info
 makeExprAlias (BetweenExpr info _ _ _) = makeDummyAlias info
 makeExprAlias (OverlapsExpr info _ _) = makeDummyAlias info
 makeExprAlias (AtTimeZoneExpr info _ _) = makeColumnAlias info "timezone" -- because reasons
-makeExprAlias LambdaParamExpr {} = error "Lambda param expression should always be used inside a lambda body"
-makeExprAlias LambdaExpr {} = error "Lambda expression should always be used inside a function"
 
 -- function expressions get the name of the function
 makeExprAlias (FunctionExpr info (QFunctionName _ _ name) _ _ _ _ _) = makeColumnAlias info name
@@ -1034,7 +1032,8 @@ makeExprAlias (FieldAccessExpr _ _ _) = fail "Unsupported struct access in Verti
 makeExprAlias (ArrayAccessExpr _ _ _) = fail "Unsupported array access in Vertica: unused datatype in this dialect"
 makeExprAlias (TypeCastExpr _ _ expr _) = makeExprAlias expr
 makeExprAlias (VariableSubstitutionExpr _) = fail "Unsupported variable substitution in Vertica: unused datatype in this dialect"
-
+makeExprAlias LambdaParamExpr {} = error "Unreachable, vertica does not support lambda"
+makeExprAlias LambdaExpr {} = error "Unreachable, vertica does not support lambda"
 
 aliasP :: Expr RawNames Range -> Parser (ColumnAlias Range)
 aliasP expr = choice

diff --git a/test/Database/Sql/Presto/Parser/Test.hs b/test/Database/Sql/Presto/Parser/Test.hs
@@ -410,6 +410,11 @@ testParser = test
             [ "SELECT RANK() OVER (x) FROM potato"
             , "WINDOW x AS (ORDER BY b);"
             ]
+
+        -- test lambda
+        , "SELECT numbers, transform(numbers, (n) -> n * n);"
+        , "SELECT transform(prices, n -> TRY_CAST(n AS VARCHAR) || '$');"
+        , "SELECT * FROM foo WHERE any_match(numbers, n ->  COALESCE(n, 0) > 100);"
        ]
 
     , "Exclude some broken examples" ~: map  (TestCase . parsesUnsuccessfully)
@@ -440,6 +445,8 @@ testParser = test
       , "WITH x AS (SELECT 1 AS a FROM dual)  SELECT * FROM X ORDER BY a  UNION  SELECT * FROM X            ;"
       , "SELECT 1 EXCEPT    ALL SELECT 1;"
       , "SELECT 1 INTERSECT ALL SELECT 1;"
+      , "SELECT n -> n;"
+      , "SELECT (n) -> n;"
       ]
 
     , ticket "T655130"
-Original file line number
+Diff line change
@@ Expand Up / @@ -57,6 +57,7 @@ operators = sortBy (flip compare) @@
         , "||"
         , "(", ")", "[", "]", ",", ";"
         , "?"
+        , "->"
         ]
     isOperator :: Char -> Bool
@@ Expand Down @@