fix: Single column not found error when using typed dataframe.

Author: mchav

dataframe.cabal

@@ -27,7 +27,8 @@ extra-source-files: cbits/arrow_abi.h
                     tests/data/*.csv
                     tests/data/*.tsv
                     tests/data/*.parquet
-                    -- unstable_csv test data moved to dataframe-fastcsv
+                    tests/data/unstable_csv/*.csv
+                    tests/data/unstable_csv/*.tsv
                     -- tests/data/*.md
                     -- tests/data/*.bin
                     -- tests/data/*.json

src/DataFrame/Typed.hs

@@ -187,6 +187,7 @@ module DataFrame.Typed (
 
     -- * Schema type families (for advanced use)
     Lookup,
+    SafeLookup,
     HasName,
     SubsetSchema,
     ExcludeSchema,

src/DataFrame/Typed/Access.hs

@@ -21,7 +21,7 @@ import GHC.TypeLits (KnownSymbol, symbolVal)
 import DataFrame.Internal.Column (Columnable)
 import DataFrame.Internal.Expression (Expr (Col))
 import qualified DataFrame.Operations.Core as D
-import DataFrame.Typed.Schema (AssertPresent, Lookup)
+import DataFrame.Typed.Schema (AssertPresent, SafeLookup)
 import DataFrame.Typed.Types (TypedDataFrame (..))
 
 {- | Retrieve a column as a boxed 'Vector', with the type determined by
@@ -30,7 +30,7 @@ the schema. The column must exist (enforced at compile time).
 columnAsVector ::
     forall name cols a.
     ( KnownSymbol name
-    , a ~ Lookup name cols
+    , a ~ SafeLookup name cols
     , Columnable a
     , AssertPresent name cols
     ) =>
@@ -44,7 +44,7 @@ columnAsVector (TDF df) =
 columnAsList ::
     forall name cols a.
     ( KnownSymbol name
-    , a ~ Lookup name cols
+    , a ~ SafeLookup name cols
     , Columnable a
     , AssertPresent name cols
     ) =>

src/DataFrame/Typed/Expr.hs

@@ -155,7 +155,7 @@ import DataFrame.Internal.Nullable (
  )
 import DataFrame.Internal.Types (Promote, PromoteDiv)
 
-import DataFrame.Typed.Schema (AssertPresent, Lookup)
+import DataFrame.Typed.Schema (AssertPresent, SafeLookup)
 import DataFrame.Typed.Types (TExpr (..), TSortOrder (..))
 import Prelude hiding (maximum, minimum, sum)
 
@@ -172,7 +172,7 @@ salary = col \@\"salary\"
 col ::
     forall (name :: Symbol) cols a.
     ( KnownSymbol name
-    , a ~ Lookup name cols
+    , a ~ SafeLookup name cols
     , Columnable a
     , AssertPresent name cols
     ) =>

src/DataFrame/Typed/Operations.hs

@@ -328,7 +328,7 @@ replaceColumn ::
     forall name a cols.
     ( KnownSymbol name
     , Columnable a
-    , a ~ Lookup name cols
+    , a ~ SafeLookup name cols
     , AssertPresent name cols
     ) =>
     TExpr cols a -> TypedDataFrame cols -> TypedDataFrame cols

src/DataFrame/Typed/Schema.hs

@@ -16,6 +16,7 @@
 module DataFrame.Typed.Schema (
     -- * Type families for schema manipulation
     Lookup,
+    SafeLookup,
     HasName,
     RemoveColumn,
     Impute,
@@ -76,6 +77,14 @@ type family Lookup (name :: Symbol) (cols :: [Type]) :: Type where
         TypeError
             ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")
 
+-- | Like 'Lookup', but returns a harmless fallback ('Int') instead of
+-- 'TypeError' when the column is not found.  Use together with
+-- 'AssertPresent' so the error fires exactly once.
+type family SafeLookup (name :: Symbol) (cols :: [Type]) :: Type where
+    SafeLookup name (Column name a ': _) = a
+    SafeLookup name (Column _ _ ': rest) = SafeLookup name rest
+    SafeLookup name '[] = Int
+
 -- | Unwrap a Maybe from a type after we impute values.
 type family Impute (name :: Symbol) (cols :: [Type]) :: [Type] where
     Impute name (Column name (Maybe a) ': rest) = Column name a ': rest