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)
@@ -179,8 +179,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

@@ -161,16 +161,16 @@ getColumn name df
 
 {- | 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 ->
@@ -617,8 +617,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

@@ -169,8 +169,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)
@@ -203,7 +203,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

@@ -394,7 +394,7 @@ cloneColumn original new df
     | null df = throw (EmptyDataSetException "cloneColumn")
     | otherwise = fromMaybe
         ( throw $
-            ColumnNotFoundException original "cloneColumn" (M.keys $ columnIndices df)
+            ColumnsNotFoundException [original] "cloneColumn" (M.keys $ columnIndices df)
         )
         $ do
             column <- getColumn original df
@@ -485,7 +485,7 @@ renameSafe ::
 renameSafe orig new df
     | null df = throw (EmptyDataSetException "rename")
     | otherwise = 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)
@@ -859,7 +859,8 @@ columnAsVector expr df
         (Col name) -> 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)
         _ -> case interpret df expr of
             Left e -> throw e
             Right (TColumn col) -> toVector col
@@ -876,7 +877,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
@@ -893,7 +894,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
@@ -910,7 +914,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
@@ -923,7 +930,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