Sf/add filtermap cases (finos#93)

Snowflake-Labs · Dec 13, 2023 · 94455b6 · 94455b6
1 parent 3477e76
commit 94455b6
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 7 deletions.
diff --git a/src/Morphir/Snowpark/Backend.elm b/src/Morphir/Snowpark/Backend.elm
@@ -45,6 +45,8 @@ import Morphir.Snowpark.MappingContext as MappingContext exposing (
 import Morphir.Snowpark.MappingContext exposing (isRecordWithComplexTypes)
 import Morphir.Snowpark.ReferenceUtils exposing (scalaPathToModule)
 import Morphir.Snowpark.Utils exposing (collectMaybeList)
+import Morphir.Snowpark.MappingContext exposing (isRecordWithSimpleTypes)
+import Morphir.Snowpark.MappingContext exposing (isAnonymousRecordWithSimpleTypes)
 
 type alias Options =
     {}
@@ -279,8 +281,21 @@ mapRecordCreation : Type () -> Dict.Dict (Name.Name) (Value ta (Type ())) -> Val
 mapRecordCreation tpe fields ctx =
     if isTypeRefToRecordWithComplexTypes tpe ctx.typesContextInfo then
         mapRecordCreationToCaseClassCreation tpe fields ctx
-    else
-        Scala.Literal (Scala.StringLit ("Record creation not converted"))
+    else 
+        if (isTypeRefToRecordWithSimpleTypes tpe ctx.typesContextInfo || 
+            isAnonymousRecordWithSimpleTypes tpe ctx.typesContextInfo) && 
+            isFunctionClassificationReturningDataFrameExpressions ctx.currentFunctionClassification  then
+            MappingContext.getFieldInfoIfRecordType tpe ctx.typesContextInfo 
+                |> Maybe.map (\fieldInfo -> collectMaybeList 
+                                                            ((\(fieldName, _) -> 
+                                                                (Dict.get fieldName fields) 
+                                                                    |> Maybe.map (\argExpr -> mapValue argExpr ctx)))
+                                                            fieldInfo  )
+                |> Maybe.withDefault Nothing
+                |> Maybe.map (applySnowparkFunc "array_construct")
+                |> Maybe.withDefault (Scala.Literal (Scala.StringLit ("Record creation not converted1")))
+        else
+            Scala.Literal (Scala.StringLit ("Record creation not converted2"))
 
 
 mapRecordCreationToCaseClassCreation : Type () -> Dict.Dict (Name.Name) (Value ta (Type ())) -> ValueMappingContext -> Scala.Value

diff --git a/src/Morphir/Snowpark/MapFunctionsMapping.elm b/src/Morphir/Snowpark/MapFunctionsMapping.elm
@@ -12,7 +12,9 @@ import Morphir.IR.Value as Value
 import Morphir.Snowpark.AggregateMapping as AggregateMapping
 import Morphir.Snowpark.Constants as Constants exposing (applySnowparkFunc)
 import Morphir.Snowpark.Operatorsmaps exposing (mapOperator)
-import Morphir.Snowpark.ReferenceUtils exposing (isTypeReferenceToSimpleTypesRecord, scalaPathToModule, getCustomTypeParameterFieldAccess)
+import Morphir.Snowpark.ReferenceUtils exposing (isTypeReferenceToSimpleTypesRecord
+            , scalaPathToModule
+            , getCustomTypeParameterFieldAccess)
 import Morphir.Snowpark.TypeRefMapping exposing (generateRecordTypeWrapperExpression,  generateCastIfPossible)
 import Morphir.Snowpark.MappingContext exposing (
             ValueMappingContext
@@ -23,9 +25,11 @@ import Morphir.Snowpark.MappingContext exposing (
             , isLocalFunctionName
             , getFieldInfoIfRecordType
             , addReplacementForIdentifier
+            , getLocalVariableIfDataFrameReference
             , isFunctionReturningDataFrameExpressions)
 import Morphir.Snowpark.Utils exposing (collectMaybeList)
 import Morphir.Snowpark.MappingContext exposing (isRecordWithSimpleTypes)
+import Morphir.Snowpark.ReferenceUtils exposing (getListTypeParameter)
 
 mapFunctionsMapping : ValueIR.Value ta (TypeIR.Type ()) -> Constants.MapValueType ta -> ValueMappingContext -> Scala.Value
 mapFunctionsMapping value mapValue ctx =
@@ -365,7 +369,27 @@ generateForListFilterMap predicate sourceRelation ctx mapValue =
               in
               Maybe.withDefault filterCall (generateProjectionForArrayColumnIfRequired  tpe ctx filterCall "result")
            _ ->
-              Scala.Literal (Scala.StringLit "Unsupported filterMap scenario")
+              let 
+                  tpe = Value.valueAttribute predicate
+                  recordReference =   getListTypeParameter (valueAttribute sourceRelation)
+                                        |> Maybe.map (\t -> (getLocalVariableIfDataFrameReference t ctx))
+                                        |> Maybe.withDefault Nothing
+                                        |> Maybe.map Scala.Variable
+                                        |> Maybe.withDefault (Scala.Literal Scala.NullLit)
+                  predicateApplication = 
+                    Scala.Apply (mapValue predicate ctx) [Scala.ArgValue Nothing recordReference]
+                  selectColumnAlias = 
+                       Scala.Apply (Scala.Select predicateApplication "as ") [ Scala.ArgValue Nothing resultId ]
+                  selectCall = 
+                       Scala.Apply (Scala.Select (mapValue sourceRelation ctx) "select") [Scala.ArgValue Nothing <| selectColumnAlias]
+                  resultId = 
+                       Scala.Literal <| Scala.StringLit "result"
+                  isNotNullCall = 
+                       Scala.Select (Constants.applySnowparkFunc "col" [ resultId ]) "is_not_null"
+                  filterCall = 
+                       Scala.Apply (Scala.Select selectCall "filter") [Scala.ArgValue Nothing isNotNullCall]
+              in
+              Maybe.withDefault filterCall (generateProjectionForArrayColumnIfRequired  tpe ctx filterCall "result")
      else 
         Scala.Literal (Scala.StringLit "Unsupported filterMap scenario")
 

diff --git a/src/Morphir/Snowpark/MappingContext.elm b/src/Morphir/Snowpark/MappingContext.elm
@@ -359,6 +359,8 @@ getFieldInfoIfRecordType tpe ctx =
                Just fieldInfo
              _ ->
                Nothing
+      Record _ fields ->
+         Just <| List.map (\field -> (field.name, field.tpe) ) fields 
       _ -> 
          Nothing
 
@@ -400,10 +402,11 @@ processFunctionDefinition definition ctx  =
       , (\_ -> 
             if List.all (\tpe -> (isDataFrameFriendlyType tpe ctx) ||
                                  (isTypeRefToRecordWithSimpleTypes tpe ctx) ||
-                                 (typeRefIsMaybeOf tpe (\t -> isTypeRefToRecordWithSimpleTypes t ctx)))
+                                 (isTypeRefToRecordWithSimpleTypes tpe ctx) ||
+                                 (typeRefIsMaybeOf tpe (\t -> (isTypeRefToRecordWithSimpleTypes t ctx) )))
                         inputTypes && 
                ((isDataFrameFriendlyType definition.outputType ctx) ||
-                  (typeRefIsMaybeOf definition.outputType (\t -> isTypeRefToRecordWithSimpleTypes t ctx)))
+                  (typeRefIsMaybeOf definition.outputType (\t -> (isTypeRefToRecordWithSimpleTypes t ctx) || (isAnonymousRecordWithSimpleTypes t ctx))))
                  then
                Just FromDfValuesToDfValues
             else Nothing)

diff --git a/src/Morphir/Snowpark/ReferenceUtils.elm b/src/Morphir/Snowpark/ReferenceUtils.elm
@@ -4,7 +4,8 @@ module Morphir.Snowpark.ReferenceUtils exposing (
     , isValueReferenceToSimpleTypesRecord
     , mapLiteral
     , scalaReferenceToUnionTypeCase
-    , getCustomTypeParameterFieldAccess)
+    , getCustomTypeParameterFieldAccess
+    , getListTypeParameter)
 
 import Morphir.IR.Name as Name
 import Morphir.IR.Type as IrType
@@ -70,3 +71,11 @@ scalaReferenceToUnionTypeCase typeName constructorName =
 getCustomTypeParameterFieldAccess : Int -> String
 getCustomTypeParameterFieldAccess paramIndex =
     "field" ++ (String.fromInt paramIndex)
+
+getListTypeParameter : IrType.Type () -> Maybe (IrType.Type ())
+getListTypeParameter tpe =
+    case tpe of
+        IrType.Reference _ ([ [ "morphir" ], [ "s", "d", "k" ] ], [ [ "list" ] ], [ "list" ]) [innertype] ->
+            Just innertype
+        _ ->
+            Nothing