diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,20 @@ +Copyright (c) 2026 Michael Chavinda + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be included +in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/dataframe-th.cabal b/dataframe-th.cabal new file mode 100644 --- /dev/null +++ b/dataframe-th.cabal @@ -0,0 +1,42 @@ +cabal-version: 2.4 +name: dataframe-th +version: 1.0.0.0 + +synopsis: Record-based Template Haskell splices for the dataframe ecosystem. +description: + The IO-agnostic part of the @dataframe@ TH machinery — splices that + read schemas from Haskell record types, not from files. File-based + splices live in @dataframe-csv-th@ and @dataframe-parquet-th@. + +bug-reports: https://github.com/mchav/dataframe/issues +license: MIT +license-file: LICENSE +author: Michael Chavinda +maintainer: mschavinda@gmail.com +copyright: (c) 2024-2025 Michael Chavinda +category: Data +tested-with: GHC ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2 + +common warnings + ghc-options: + -Wincomplete-patterns + -Wincomplete-uni-patterns + -Wunused-imports + -Wunused-local-binds + +library + import: warnings + exposed-modules: + DataFrame.Internal.Schema.TH + DataFrame.TH.Records + DataFrame.Typed.TH.Records + build-depends: base >= 4 && < 5, + containers >= 0.6.7 && < 0.9, + dataframe-core ^>= 1.0, + dataframe-operations ^>= 1.0, + dataframe-parsing ^>= 1.0, + template-haskell >= 2.0 && < 3, + text >= 2.0 && < 3, + vector ^>= 0.13 + hs-source-dirs: src + default-language: Haskell2010 diff --git a/src/DataFrame/Internal/Schema/TH.hs b/src/DataFrame/Internal/Schema/TH.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Internal/Schema/TH.hs @@ -0,0 +1,123 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskellQuotes #-} +{-# LANGUAGE TypeApplications #-} + +{- | +Template-Haskell @deriveSchema@ splice for "DataFrame.Internal.Schema". +Kept in a separate module so the runtime schema types in +"DataFrame.Internal.Schema" do not pull in @template-haskell@. +-} +module DataFrame.Internal.Schema.TH ( + deriveSchema, + camelToSnake, +) where + +import Data.Char (isUpper, toLower, toUpper) +import qualified Data.Text as T +import Language.Haskell.TH + +import DataFrame.Internal.Expression (Expr) +import DataFrame.Internal.Schema (Schema, makeSchema, schemaType) +import DataFrame.Operators (col) + +{- | Auto-generate a runtime 'Schema' (and per-column @'Expr'@ accessors) +from a record ADT. + +The splice reifies the record, applies @camelCase -> snake_case@ to each +record-selector name, and emits: + +* a top-level @\Schema :: 'Schema'@ binding suitable + for passing to 'DataFrame.IO.CSV.readCsvWithSchema'. +* one @\\ :: 'Expr' /ty/@ binding + per field, so you can refer to columns in expression DSL code by name + without writing @col \@/ty/ "snake_case_name"@ at every call site. + +The data type must have exactly one record constructor; sum types or +positional constructors fail the splice with a descriptive error. Field +types must satisfy @('Columnable' a, 'Read' a)@ — the same constraints +'schemaType' already requires. +-} +deriveSchema :: Name -> DecsQ +deriveSchema tyName = do + info <- reify tyName + fields <- extractRecordFields tyName info + let entries = + [ (camelToSnake fieldBase, fieldBase, fTy) + | (fName, _bang, fTy) <- fields + , let fieldBase = nameBase fName + ] + schemaName = mkName (lowerFirst (nameBase tyName) ++ "Schema") + prefix = lowerFirst (nameBase tyName) + tupleE (colName, _, fTy) = + TupE + [ Just (AppE (VarE 'T.pack) (LitE (StringL colName))) + , Just (AppTypeE (VarE 'schemaType) fTy) + ] + schemaBody = + AppE (VarE 'makeSchema) (ListE (map tupleE entries)) + schemaDecls = + [ SigD schemaName (ConT ''Schema) + , ValD (VarP schemaName) (NormalB schemaBody) [] + ] + accessorDecls = + concat + [ [ SigD accName (AppT (ConT ''Expr) fTy) + , ValD + (VarP accName) + ( NormalB + ( AppE + (VarE 'col) + ( AppE + (VarE 'T.pack) + (LitE (StringL colName)) + ) + ) + ) + [] + ] + | (colName, fieldBase, fTy) <- entries + , let accName = mkName (prefix ++ upperFirst fieldBase) + ] + pure (schemaDecls ++ accessorDecls) + +extractRecordFields :: Name -> Info -> Q [VarBangType] +extractRecordFields _ (TyConI dec) = case dec of + DataD _ _ _ _ [RecC _ fs] _ -> pure fs + NewtypeD _ _ _ _ (RecC _ fs) _ -> pure fs + DataD _ n _ _ _ _ -> + fail $ + "deriveSchema: " + ++ show n + ++ " must have exactly one record constructor" + NewtypeD _ n _ _ _ _ -> + fail $ + "deriveSchema: " ++ show n ++ " newtype must use record syntax" + other -> + fail $ + "deriveSchema: unsupported declaration: " ++ show other +extractRecordFields tyName _ = + fail $ + "deriveSchema: " + ++ show tyName + ++ " is not a data/newtype declaration" + +{- | @camelCase -> snake_case@. Lowercases the first character then prefixes +@\'_\'@ before every uppercase character (lowercased). +-} +camelToSnake :: String -> String +camelToSnake [] = [] +camelToSnake (c : cs) = toLower c : go cs + where + go [] = [] + go (x : xs) + | isUpper x = '_' : toLower x : go xs + | otherwise = x : go xs + +lowerFirst :: String -> String +lowerFirst [] = [] +lowerFirst (c : cs) = toLower c : cs + +upperFirst :: String -> String +upperFirst [] = [] +upperFirst (c : cs) = toUpper c : cs diff --git a/src/DataFrame/TH/Records.hs b/src/DataFrame/TH/Records.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/TH/Records.hs @@ -0,0 +1,112 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} + +{- | +Module : DataFrame.TH.Records +License : MIT + +Record-based 'DataFrame' splices — the IO-agnostic core of the +'DataFrame.TH' family. Splices that read CSV / Parquet files at compile +time live in @DataFrame.TH.CSV@ (in @dataframe-csv-th@) and +@DataFrame.TH.Parquet@ (in @dataframe-parquet-th@). +-} +module DataFrame.TH.Records ( + -- * Declare one binding per column + declareColumns, + declareColumnsWithPrefix, + declareColumnsWithPrefix', + + -- * Type-string parser (exposed for testing) + typeFromString, +) where + +import Data.Function (on) +import Data.Functor ((<&>)) +import qualified Data.List as L +import qualified Data.Map as M +import qualified Data.Text as T + +import Control.Monad (forM) +import Language.Haskell.TH +import qualified Language.Haskell.TH.Syntax as TH + +import DataFrame.Functions (sanitize) +import DataFrame.Internal.Column (columnTypeString) +import DataFrame.Internal.DataFrame ( + DataFrame (..), + unsafeGetColumn, + ) +import DataFrame.Internal.Expression (Expr) +import DataFrame.Operators (col) +import Prelude as P + +typeFromString :: [String] -> Q Type +typeFromString [] = fail "No type specified" +typeFromString [t0] = do + let t = trim t0 + case stripBrackets t of + Just inner -> typeFromString [inner] <&> AppT ListT + Nothing + | t == "Text" || t == "Data.Text.Text" || t == "T.Text" -> + pure (ConT ''T.Text) + | otherwise -> do + m <- lookupTypeName t + case m of + Just tyName -> pure (ConT tyName) + Nothing -> fail $ "Unsupported type: " ++ t0 +typeFromString [tycon, t1] = AppT <$> typeFromString [tycon] <*> typeFromString [t1] +typeFromString [tycon, t1, t2] = + (\outer a b -> AppT (AppT outer a) b) + <$> typeFromString [tycon] + <*> typeFromString [t1] + <*> typeFromString [t2] +typeFromString s = fail $ "Unsupported types: " ++ unwords s + +trim :: String -> String +trim = dropWhile (== ' ') . reverse . dropWhile (== ' ') . reverse + +stripBrackets :: String -> Maybe String +stripBrackets s = + case s of + ('[' : rest) + | P.not (null rest) && last rest == ']' -> + Just (init rest) + _ -> Nothing + +{- | Splice a binding for every column of @df@, named after the column. +Column names that are not valid Haskell identifiers are sanitized +(see 'DataFrame.Functions.sanitize'). +-} +declareColumns :: DataFrame -> DecsQ +declareColumns = declareColumnsWithPrefix' Nothing + +-- | Like 'declareColumns' but prefixes every binding name with @prefix_@. +declareColumnsWithPrefix :: T.Text -> DataFrame -> DecsQ +declareColumnsWithPrefix prefix = declareColumnsWithPrefix' (Just prefix) + +-- | Like 'declareColumnsWithPrefix' but takes an optional prefix. +declareColumnsWithPrefix' :: Maybe T.Text -> DataFrame -> DecsQ +declareColumnsWithPrefix' prefix df = + let + names = (map fst . L.sortBy (compare `on` snd) . M.toList . columnIndices) df + types = map (columnTypeString . (`unsafeGetColumn` df)) names + specs = + zipWith + ( \colName type_ -> + ( colName + , maybe "" (sanitize . (<> "_")) prefix <> sanitize colName + , type_ + ) + ) + names + types + in + fmap concat $ forM specs $ \(raw, nm, tyStr) -> do + ty <- typeFromString (words tyStr) + let n = mkName (T.unpack nm) + sig <- sigD n [t|Expr $(pure ty)|] + val <- valD (varP n) (normalB [|col $(TH.lift raw)|]) [] + pure [sig, val] diff --git a/src/DataFrame/Typed/TH/Records.hs b/src/DataFrame/Typed/TH/Records.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/TH/Records.hs @@ -0,0 +1,269 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskellQuotes #-} +{-# LANGUAGE TypeApplications #-} + +{- | +Module : DataFrame.Typed.TH.Records +License : MIT + +Record-based typed-API schema derivation. The IO-agnostic core of the +@DataFrame.Typed.TH@ family. CSV-file-based derivation lives in +"DataFrame.Typed.TH.CSV" (in @dataframe-csv-th@). +-} +module DataFrame.Typed.TH.Records ( + -- * Schema inference from an existing DataFrame value + deriveSchema, + + -- * ADT-based schema derivation + deriveSchemaFromType, + deriveSchemaFromTypeWith, + SchemaOptions (..), + defaultSchemaOptions, + camelToSnake, + + -- * Re-export for TH splices + TypedDataFrame, + Column, + + -- * Shared helpers (used by sibling TH modules) + getSchemaInfo, + findDuplicate, +) where + +import Control.Monad (when) +import qualified Data.List as L +import qualified Data.Map as M +import qualified Data.Text as T +import qualified Data.Vector as VB + +import Language.Haskell.TH + +import qualified DataFrame.Internal.Column as C +import qualified DataFrame.Internal.DataFrame as D +import DataFrame.Typed.Record ( + HasSchema, + Schema, + fromColumns, + requireColumn, + toColumns, + ) +import DataFrame.Typed.Types (Column, TypedDataFrame) +import DataFrame.Typed.Util (camelToSnake) + +deriveSchema :: String -> D.DataFrame -> DecsQ +deriveSchema typeName df = do + let cols = getSchemaInfo df + let names = map fst cols + case findDuplicate names of + Just dup -> fail $ "Duplicate column name in DataFrame: " ++ T.unpack dup + Nothing -> pure () + colTypes <- mapM mkColumnType cols + let schemaType = foldr (\t acc -> PromotedConsT `AppT` t `AppT` acc) PromotedNilT colTypes + let synName = mkName typeName + pure [TySynD synName [] schemaType] + +getSchemaInfo :: D.DataFrame -> [(T.Text, String)] +getSchemaInfo df = + let orderedNames = + map fst $ + L.sortBy (\(_, a) (_, b) -> compare a b) $ + M.toList (D.columnIndices df) + in map (\name -> (name, getColumnTypeStr name df)) orderedNames + +getColumnTypeStr :: T.Text -> D.DataFrame -> String +getColumnTypeStr name df = case D.getColumn name df of + Just col -> C.columnTypeString col + Nothing -> error $ "Column not found: " ++ T.unpack name + +mkColumnType :: (T.Text, String) -> Q Type +mkColumnType (name, tyStr) = do + ty <- parseTypeString tyStr + let nameLit = LitT (StrTyLit (T.unpack name)) + pure $ ConT ''Column `AppT` nameLit `AppT` ty + +parseTypeString :: String -> Q Type +parseTypeString "Int" = pure $ ConT ''Int +parseTypeString "Double" = pure $ ConT ''Double +parseTypeString "Float" = pure $ ConT ''Float +parseTypeString "Bool" = pure $ ConT ''Bool +parseTypeString "Char" = pure $ ConT ''Char +parseTypeString "String" = pure $ ConT ''String +parseTypeString "Text" = pure $ ConT ''T.Text +parseTypeString "Integer" = pure $ ConT ''Integer +parseTypeString s + | "Maybe " `L.isPrefixOf` s = do + inner <- parseTypeString (L.drop 6 s) + pure $ ConT ''Maybe `AppT` inner +parseTypeString s = fail $ "Unsupported column type in schema inference: " ++ s + +findDuplicate :: (Eq a) => [a] -> Maybe a +findDuplicate [] = Nothing +findDuplicate (x : xs) + | x `elem` xs = Just x + | otherwise = findDuplicate xs + +-- | Options controlling 'deriveSchemaFromTypeWith'. +data SchemaOptions = SchemaOptions + { nameTransform :: String -> String + -- ^ Map each record selector name to a column name. Default: 'camelToSnake'. + , schemaTypeName :: Maybe String + -- ^ Override the generated type synonym name. Default: @\Schema@. + , generateInstance :: Bool + {- ^ When @True@ (default), also generate a 'HasSchema' instance so the + record can be turned into a 'DataFrame' via 'fromRecords' / 'toRecords'. + -} + } + +defaultSchemaOptions :: SchemaOptions +defaultSchemaOptions = + SchemaOptions + { nameTransform = camelToSnake + , schemaTypeName = Nothing + , generateInstance = True + } + +deriveSchemaFromType :: Name -> DecsQ +deriveSchemaFromType = deriveSchemaFromTypeWith defaultSchemaOptions + +deriveSchemaFromTypeWith :: SchemaOptions -> Name -> DecsQ +deriveSchemaFromTypeWith opts tyName = do + info <- reify tyName + (conName, vbts) <- extractRecord tyName info + when (Prelude.null vbts) $ + fail $ + "deriveSchemaFromType: record " + ++ show tyName + ++ " has no fields" + let fields = + [ (nameTransform opts (nameBase fName), fName, ty) + | (fName, _bang, ty) <- vbts + ] + colNames = [c | (c, _, _) <- fields] + case findDuplicate colNames of + Just dup -> + fail $ + "deriveSchemaFromType: duplicate transformed column name " + ++ show dup + ++ " (consider customizing nameTransform via deriveSchemaFromTypeWith)" + Nothing -> pure () + let synName = case schemaTypeName opts of + Just s -> mkName s + Nothing -> mkName (nameBase tyName ++ "Schema") + let columnTypes = + [ ConT ''Column `AppT` LitT (StrTyLit colName) `AppT` ty + | (colName, _, ty) <- fields + ] + schemaType = + foldr + (\t acc -> PromotedConsT `AppT` t `AppT` acc) + PromotedNilT + columnTypes + synDec = TySynD synName [] schemaType + if generateInstance opts + then do + inst <- mkHasSchemaInstance tyName schemaType conName fields + pure [synDec, inst] + else pure [synDec] + +extractRecord :: Name -> Info -> Q (Name, [VarBangType]) +extractRecord _ (TyConI dec) = case dec of + DataD _ _ _ _ [RecC conName fs] _ -> pure (conName, fs) + NewtypeD _ _ _ _ (RecC conName fs) _ -> pure (conName, fs) + DataD _ name _ _ _ _ -> + fail $ + "deriveSchemaFromType: " + ++ show name + ++ " must have exactly one record constructor" + NewtypeD _ name _ _ _ _ -> + fail $ + "deriveSchemaFromType: " + ++ show name + ++ " newtype must use record syntax" + other -> + fail $ + "deriveSchemaFromType: unsupported declaration: " ++ show other +extractRecord tyName _ = + fail $ + "deriveSchemaFromType: " ++ show tyName ++ " is not a data/newtype declaration" + +mkHasSchemaInstance :: + Name -> Type -> Name -> [(String, Name, Type)] -> Q Dec +mkHasSchemaInstance tyName schemaType conName fields = do + toClause <- mkToColumnsClause fields + fromClause <- mkFromColumnsClause conName fields + let instType = ConT ''HasSchema `AppT` ConT tyName + schemaInst = + TySynInstD + (TySynEqn Nothing (ConT ''Schema `AppT` ConT tyName) schemaType) + pure $ + InstanceD + Nothing + [] + instType + [ schemaInst + , FunD 'toColumns [toClause] + , FunD 'fromColumns [fromClause] + ] + +mkToColumnsClause :: [(String, Name, Type)] -> Q Clause +mkToColumnsClause fields = do + rs <- newName "rs" + let mkPair (colName, fieldFn, _ty) = + let nameE = AppE (VarE 'T.pack) (LitE (StringL colName)) + colE = + AppE + (VarE 'C.fromList) + ( AppE + (AppE (VarE 'map) (VarE fieldFn)) + (VarE rs) + ) + in TupE [Just nameE, Just colE] + listExp = ListE (map mkPair fields) + pure $ Clause [VarP rs] (NormalB listExp) [] + +mkFromColumnsClause :: Name -> [(String, Name, Type)] -> Q Clause +mkFromColumnsClause conName fields = do + df <- newName "df" + iN <- newName "i" + nN <- newName "n" + vNames <- mapM (\k -> newName ("v" ++ show (k :: Int))) [0 .. length fields - 1] + let mkBind v (colName, _, ty) = + let nameE = AppE (VarE 'T.pack) (LitE (StringL colName)) + callE = + AppE + ( AppE + (AppTypeE (VarE 'requireColumn) ty) + nameE + ) + (VarE df) + in BindS (VarP v) callE + binds = zipWith mkBind vNames fields + firstV = case vNames of + (v0 : _) -> v0 + [] -> error "mkFromColumnsClause: empty fields (should have failed earlier)" + lengthBind = + LetS + [ ValD + (VarP nN) + (NormalB (AppE (VarE 'VB.length) (VarE firstV))) + [] + ] + indexE v = + AppE (AppE (VarE 'VB.unsafeIndex) (VarE v)) (VarE iN) + elemE = + foldl + (\acc v -> AppE acc (indexE v)) + (ConE conName) + vNames + nMinus1E = + InfixE + (Just (VarE nN)) + (VarE '(-)) + (Just (LitE (IntegerL 1))) + rangeE = ArithSeqE (FromToR (LitE (IntegerL 0)) nMinus1E) + compExp = CompE [BindS (VarP iN) rangeE, NoBindS elemE] + rightE = AppE (ConE 'Right) compExp + body = DoE Nothing (binds ++ [lengthBind, NoBindS rightE]) + pure $ Clause [VarP df] (NormalB body) []