fix joins for missing key columns

src/DataFrame/Errors.hs

@@ -11,6 +11,7 @@ import qualified Data.Vector.Unboxed as VU
 
 import Control.Exception
 import Data.Array
+import qualified Data.List as L
 import Data.Typeable (Typeable)
 import DataFrame.Display.Terminal.Colours
 import Type.Reflection (TypeRep)
@@ -29,7 +30,7 @@ data DataFrameException where
         TypeErrorContext a b ->
         DataFrameException
     AggregatedAndNonAggregatedException :: T.Text -> T.Text -> DataFrameException
-    ColumnNotFoundException :: T.Text -> T.Text -> [T.Text] -> DataFrameException
+    ColumnsNotFoundException :: [T.Text] -> T.Text -> [T.Text] -> DataFrameException
     EmptyDataSetException :: T.Text -> DataFrameException
     InternalException :: T.Text -> DataFrameException
     NonColumnReferenceException :: T.Text -> DataFrameException
@@ -51,7 +52,7 @@ instance Show DataFrameException where
                 (errorColumnName context)
                 (callingFunctionName context)
                 errorString
-    show (ColumnNotFoundException columnName callPoint availableColumns) = columnNotFound columnName callPoint availableColumns
+    show (ColumnsNotFoundException columnNames callPoint availableColumns) = columnsNotFound columnNames callPoint availableColumns
     show (EmptyDataSetException callPoint) = emptyDataSetError callPoint
     show (WrongQuantileNumberException q) = wrongQuantileNumberError q
     show (WrongQuantileIndexException qs q) = wrongQuantileIndexError qs q
@@ -65,15 +66,46 @@ instance Show DataFrameException where
             ++ T.unpack expr2
 
 columnNotFound :: T.Text -> T.Text -> [T.Text] -> String
-columnNotFound name callPoint columns =
+columnNotFound missingColumn = columnsNotFound [missingColumn]
+
+columnsNotFound :: [T.Text] -> T.Text -> [T.Text] -> String
+columnsNotFound missingColumns callPoint availableColumns =
     red "\n\n[ERROR] "
-        ++ "Column not found: "
-        ++ T.unpack name
+        ++ missingColumnsLabel missingColumns
+        ++ ": "
+        ++ T.unpack (T.intercalate ", " missingColumns)
         ++ " for operation "
         ++ T.unpack callPoint
-        ++ "\n\tDid you mean "
-        ++ T.unpack (guessColumnName name columns)
-        ++ "?\n\n"
+        ++ formatSuggestions missingColumns availableColumns
+        ++ "\n\n"
+  where
+    missingColumnsLabel [_] = "Column not found"
+    missingColumnsLabel _ = "Columns not found"
+
+    formatSuggestions [missingColumn] columns =
+        case guessColumnName missingColumn columns of
+            "" -> ""
+            guessed ->
+                "\n\tDid you mean "
+                    ++ T.unpack guessed
+                    ++ "?"
+    formatSuggestions names columns =
+        case traverse (`suggestColumnName` columns) names of
+            Just guessedColumns
+                | not (null guessedColumns) ->
+                    "\n\tDid you mean "
+                        ++ formatColumnSuggestions guessedColumns
+                        ++ "?"
+            _ -> ""
+
+    suggestColumnName missingColumn columns = case guessColumnName missingColumn columns of
+        "" -> Nothing
+        guessed -> Just guessed
+
+    formatColumnSuggestions guessedColumns =
+        "["
+            ++ L.intercalate ", " (map (show . T.unpack) guessedColumns)
+            ++ "]"
 
 typeMismatchError :: String -> String -> String
 typeMismatchError givenType expectedType =

src/DataFrame/IO/Parquet.hs

@@ -19,7 +19,7 @@ import qualified Data.Text as T
 import Data.Text.Encoding
 import Data.Time
 import Data.Time.Clock.POSIX (posixSecondsToUTCTime)
-import DataFrame.Errors (DataFrameException (ColumnNotFoundException))
+import DataFrame.Errors (DataFrameException (ColumnsNotFoundException))
 import DataFrame.Internal.Binary (littleEndianWord32)
 import qualified DataFrame.Internal.Column as DI
 import DataFrame.Internal.DataFrame (DataFrame)
@@ -160,8 +160,8 @@ _readParquetWithOpts extraConfig opts path = withFileBufferedOrSeekable extraCon
              in unless
                     (L.null missing)
                     ( throw
-                        ( ColumnNotFoundException
-                            (T.pack $ show missing)
+                        ( ColumnsNotFoundException
+                            missing
                             "readParquetWithOpts"
                             availableSelectedColumns
                         )

src/DataFrame/Internal/DataFrame.hs

@@ -154,16 +154,16 @@ getColumn name df = do
 
 {- | Retrieves a column by name from the dataframe, throwing an exception if not found.
 
-This is an unsafe version of 'getColumn' that throws 'ColumnNotFoundException'
+This is an unsafe version of 'getColumn' that throws 'ColumnsNotFoundException'
 if the column does not exist. Use this when you are certain the column exists.
 
 ==== __Throws__
 
-* 'ColumnNotFoundException' - if the column with the given name does not exist
+* 'ColumnsNotFoundException' - if the column with the given name does not exist
 -}
 unsafeGetColumn :: T.Text -> DataFrame -> Column
 unsafeGetColumn name df = case getColumn name df of
-    Nothing -> throw $ ColumnNotFoundException name "" (M.keys $ columnIndices df)
+    Nothing -> throw $ ColumnsNotFoundException [name] "" (M.keys $ columnIndices df)
     Just col -> col
 
 {- | Checks if the dataframe is empty (has no columns).

src/DataFrame/Internal/Interpreter.hs

@@ -482,7 +482,7 @@ eval _ (Lit v) = Right (Scalar v)
 eval (FlatCtx df) (Col name) =
     case getColumn name df of
         Nothing ->
-            Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df)
+            Left $ ColumnsNotFoundException [name] "" (M.keys $ columnIndices df)
         Just c
             | hasElemType @a c -> Right (Flat c)
             | otherwise ->
@@ -500,8 +500,8 @@ eval (GroupCtx gdf) (Col name) =
     case getColumn name (fullDataframe gdf) of
         Nothing ->
             Left $
-                ColumnNotFoundException
-                    name
+                ColumnsNotFoundException
+                    [name]
                     ""
                     (M.keys $ columnIndices $ fullDataframe gdf)
         Just c
@@ -524,14 +524,14 @@ eval (FlatCtx df) (CastWith name _tag onResult) =
     case getColumn name df of
         Nothing ->
             Left $
-                ColumnNotFoundException name "" (M.keys $ columnIndices df)
+                ColumnsNotFoundException [name] "" (M.keys $ columnIndices df)
         Just c -> Flat <$> promoteColumnWith onResult c
 eval (GroupCtx gdf) (CastWith name _tag onResult) =
     case getColumn name (fullDataframe gdf) of
         Nothing ->
             Left $
-                ColumnNotFoundException
-                    name
+                ColumnsNotFoundException
+                    [name]
                     ""
                     (M.keys $ columnIndices $ fullDataframe gdf)
         Just c -> do
@@ -579,8 +579,8 @@ eval (GroupCtx gdf) expr@(Agg (FoldAgg _ (Just seed) (f :: a -> b -> a)) (Col na
         case getColumn name (fullDataframe gdf) of
             Nothing ->
                 Left $
-                    ColumnNotFoundException
-                        name
+                    ColumnsNotFoundException
+                        [name]
                         ""
                         (M.keys $ columnIndices $ fullDataframe gdf)
             Just col ->
@@ -599,8 +599,8 @@ eval (GroupCtx gdf) expr@(Agg (FoldAgg _ Nothing (f :: a -> b -> a)) (Col name :
                 case getColumn name (fullDataframe gdf) of
                     Nothing ->
                         Left $
-                            ColumnNotFoundException
-                                name
+                            ColumnsNotFoundException
+                                [name]
                                 ""
                                 (M.keys $ columnIndices $ fullDataframe gdf)
                     Just col ->
@@ -618,8 +618,8 @@ eval
             case getColumn name (fullDataframe gdf) of
                 Nothing ->
                     Left $
-                        ColumnNotFoundException
-                            name
+                        ColumnsNotFoundException
+                            [name]
                             ""
                             (M.keys $ columnIndices $ fullDataframe gdf)
                 Just col ->

src/DataFrame/Internal/Row.hs

@@ -173,8 +173,8 @@ mkRowFromArgs names df i = V.map get (V.fromList names)
     get name = case getColumn name df of
         Nothing ->
             throw $
-                ColumnNotFoundException
-                    name
+                ColumnsNotFoundException
+                    [name]
                     "[INTERNAL] mkRowFromArgs"
                     (M.keys $ columnIndices df)
         Just (BoxedColumn column) -> toAny (column V.! i)
@@ -207,7 +207,7 @@ mkRowRep df names i = V.generate (L.length names) (\index -> get (names' V.! ind
             Just e -> toAny e
             Nothing -> throwError name
         Nothing ->
-            throw $ ColumnNotFoundException name "mkRowRep" (M.keys $ columnIndices df)
+            throw $ ColumnsNotFoundException [name] "mkRowRep" (M.keys $ columnIndices df)
 
 sortedIndexes' :: [Bool] -> V.Vector Row -> VU.Vector Int
 sortedIndexes' flipCompare rows = runST $ do

src/DataFrame/Operations/Aggregation.hs

@@ -49,8 +49,8 @@ groupBy ::
 groupBy names df
     | any (`notElem` columnNames df) names =
         throw $
-            ColumnNotFoundException
-                (T.pack $ show $ names L.\\ columnNames df)
+            ColumnsNotFoundException
+                (names L.\\ columnNames df)
                 "groupBy"
                 (columnNames df)
     | nRows df == 0 =

src/DataFrame/Operations/Core.hs

@@ -391,7 +391,7 @@ insertColumn name column d =
 cloneColumn :: T.Text -> T.Text -> DataFrame -> DataFrame
 cloneColumn original new df = fromMaybe
     ( throw $
-        ColumnNotFoundException original "cloneColumn" (M.keys $ columnIndices df)
+        ColumnsNotFoundException [original] "cloneColumn" (M.keys $ columnIndices df)
     )
     $ do
         column <- getColumn original df
@@ -480,7 +480,7 @@ renameMany = fold (uncurry rename)
 renameSafe ::
     T.Text -> T.Text -> DataFrame -> Either DataFrameException DataFrame
 renameSafe orig new df = fromMaybe
-    (Left $ ColumnNotFoundException orig "rename" (M.keys $ columnIndices df))
+    (Left $ ColumnsNotFoundException [orig] "rename" (M.keys $ columnIndices df))
     $ do
         columnIndex <- M.lookup orig (columnIndices df)
         let origRemoved = M.delete orig (columnIndices df)
@@ -856,7 +856,8 @@ columnAsVector ::
 columnAsVector (Col name) df = case getColumn name df of
     Just col -> toVector col
     Nothing ->
-        Left $ ColumnNotFoundException name "columnAsVector" (M.keys $ columnIndices df)
+        Left $
+            ColumnsNotFoundException [name] "columnAsVector" (M.keys $ columnIndices df)
 columnAsVector expr df = case interpret df expr of
     Left e -> throw e
     Right (TColumn col) -> toVector col
@@ -873,7 +874,7 @@ columnAsIntVector (Col name) df = case getColumn name df of
     Just col -> toIntVector col
     Nothing ->
         Left $
-            ColumnNotFoundException name "columnAsIntVector" (M.keys $ columnIndices df)
+            ColumnsNotFoundException [name] "columnAsIntVector" (M.keys $ columnIndices df)
 columnAsIntVector expr df = case interpret df expr of
     Left e -> throw e
     Right (TColumn col) -> toIntVector col
@@ -890,7 +891,10 @@ columnAsDoubleVector (Col name) df = case getColumn name df of
     Just col -> toDoubleVector col
     Nothing ->
         Left $
-            ColumnNotFoundException name "columnAsDoubleVector" (M.keys $ columnIndices df)
+            ColumnsNotFoundException
+                [name]
+                "columnAsDoubleVector"
+                (M.keys $ columnIndices df)
 columnAsDoubleVector expr df = case interpret df expr of
     Left e -> throw e
     Right (TColumn col) -> toDoubleVector col
@@ -907,7 +911,10 @@ columnAsFloatVector (Col name) df = case getColumn name df of
     Just col -> toFloatVector col
     Nothing ->
         Left $
-            ColumnNotFoundException name "columnAsFloatVector" (M.keys $ columnIndices df)
+            ColumnsNotFoundException
+                [name]
+                "columnAsFloatVector"
+                (M.keys $ columnIndices df)
 columnAsFloatVector expr df = case interpret df expr of
     Left e -> throw e
     Right (TColumn col) -> toFloatVector col
@@ -920,7 +927,10 @@ columnAsUnboxedVector (Col name) df = case getColumn name df of
     Just col -> toUnboxedVector col
     Nothing ->
         Left $
-            ColumnNotFoundException name "columnAsFloatVector" (M.keys $ columnIndices df)
+            ColumnsNotFoundException
+                [name]
+                "columnAsFloatVector"
+                (M.keys $ columnIndices df)
 columnAsUnboxedVector expr df = case interpret df expr of
     Left e -> throw e
     Right (TColumn col) -> toUnboxedVector col