diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Revision history for beam-duckdb + +## 0.1.0.0 -- 2026-02-26 + +* First version. Released on an unsuspecting world. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,20 @@ +Copyright (c) 2026 Laurent P. René de Cotret + +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/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,12 @@ +# `beam-duckdb`: Beam backend for the DuckDB analytics database + +`beam-duckdb` is a beam backend for +the [DuckDB analytics database](https://duckdb.org/). + +`beam-duckdb` extends the set of capabilities provided by `beam-core` in a few key +ways. Most importantly, sources of data _other_ than database tables can be +loaded by DuckDB and used for queries like a regular table. Currently, `beam-duckdb` +supports loading Parquet files, Apache Iceberg tables, and CSV files. + +`beam-duckdb` is the most recent backend for beam; do not hesitate to raise an issue +if you'd like us to add support for something that isn't currently covered! diff --git a/beam-duckdb.cabal b/beam-duckdb.cabal new file mode 100644 --- /dev/null +++ b/beam-duckdb.cabal @@ -0,0 +1,69 @@ +cabal-version: 3.0 +name: beam-duckdb +version: 0.1.0.0 +synopsis: DuckDB backend for Beam +description: Beam driver for DuckDB, an analytics-focused open-source in-process database. +license: MIT +license-file: LICENSE +author: Laurent P. René de Cotret +maintainer: laurent.decotret@outlook.com +copyright: (c) Laurent P. René de Cotret +category: Database +build-type: Simple +extra-doc-files: README.md + CHANGELOG.md +extra-source-files: tests/data/*.py + tests/data/*.parquet + tests/data/*.csv + tests/data/lineitem_iceberg/data/*.parquet + tests/data/lineitem_iceberg/metadata/*.avro + tests/data/lineitem_iceberg/metadata/*.json + tests/data/lineitem_iceberg/metadata/version-hint.text + +common common + default-language: Haskell2010 + ghc-options: -Wall + -Wcompat + -Widentities + -Wincomplete-uni-patterns + -Wincomplete-record-updates + -Wredundant-constraints + -fhide-source-paths + -Wpartial-fields +library + import: common + exposed-modules: Database.Beam.DuckDB + other-modules: Database.Beam.DuckDB.Backend + Database.Beam.DuckDB.Syntax + Database.Beam.DuckDB.Syntax.Builder + Database.Beam.DuckDB.Syntax.Extensions.DataSource + build-depends: base >=4.11 && <5 + , beam-core ^>=0.10 + , beam-migrate ^>=0.5 + , bytestring >=0.10 && <0.13 + , duckdb-simple ^>=0.1 + , dlist >=0.8 && <1.1 + , free >=4.12 && <5.3 + , text >=1.0 && <2.2 + , time >=1.6 && <1.16 + , transformers >=0.3 && <0.7 + , uuid-types ^>=1.0 + hs-source-dirs: src + +test-suite beam-duckdb-test + import: common + other-modules: Database.Beam.DuckDB.Test.Extensions + Database.Beam.DuckDB.Test.Query + type: exitcode-stdio-1.0 + hs-source-dirs: tests + main-is: Main.hs + build-depends: base >=4.11 && <5 + , beam-core + , beam-duckdb + , duckdb-simple + , hedgehog + , tasty + , tasty-hedgehog + , tasty-hunit + , text + , time diff --git a/src/Database/Beam/DuckDB.hs b/src/Database/Beam/DuckDB.hs new file mode 100644 --- /dev/null +++ b/src/Database/Beam/DuckDB.hs @@ -0,0 +1,216 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | DuckDB is a powerful in-process database specialized in analytical processing workloads. +-- +-- The @beam-duckdb@ library is built atop of @duckdb-simple@, which is +-- used for connection management, transaction support, serialization, and +-- deserialization. +-- +-- @beam-duckdb@ supports most beam features as well as many DuckDB-specific +-- features, such as support for reading Parquet files and Apache Iceberg tables. +module Database.Beam.DuckDB + ( -- * Executing DuckDB queries + runBeamDuckDB, + runBeamDuckDBDebug, + + -- * Backend datatype + DuckDB, + DuckDBM, + + -- * DuckDB-specific functionality + + -- ** Data sources + DataSourceEntity, + dataSource, + modifyDataSourceFields, + allFromDataSource_, + + -- *** Parquet + parquet, + + -- *** Apache Iceberg tables + icebergTable, + icebergTableWith, + IcebergTableOptions (..), + defaultIcebergTableOptions, + + -- *** CSV + csv, + csvWith, + CSVOptions (..), + defaultCSVOptions, + ) +where + +import Control.Exception (SomeException (..)) +import Control.Monad (void) +import Control.Monad.Free.Church (runF) +import Control.Monad.IO.Class (MonadIO (..)) +import Control.Monad.Trans.Class (lift) +import Control.Monad.Trans.Reader (ReaderT (..), ask) +import Control.Monad.Trans.State.Strict (StateT (..), get, put) +import qualified Data.DList as DL +import Data.Data (cast) +import Data.Maybe (mapMaybe) +import Data.Text (Text) +import qualified Data.Text as Text +import qualified Data.Text.Lazy as Text.Lazy +import qualified Data.Text.Lazy.Builder as Builder +import Database.Beam.Backend + ( BeamRowReadError (..), + ColumnParseError (..), + FromBackendRow, + FromBackendRowF (..), + FromBackendRowM (..), + MonadBeam, + ) +import Database.Beam.Backend.SQL (FromBackendRow (..), MonadBeam (..)) +import Database.Beam.DuckDB.Backend (DuckDB) +import Database.Beam.DuckDB.Syntax + ( DuckDBCommandSyntax (..), + ) +import Database.Beam.DuckDB.Syntax.Builder + ( DuckDBSyntax (..), + SomeField (..), + withPlaceholder, + ) +import Database.Beam.DuckDB.Syntax.Extensions.DataSource +import Database.DuckDB.Simple (Connection, FromRow, Query (Query), ResultError (..), RowParser, ToField (toField), ToRow (toRow), bind, execute, nextRow, withStatement) +import Database.DuckDB.Simple.FromRow (FromRow (..), RowParser (..), field) +import Database.DuckDB.Simple.Ok (Ok (..)) + +-- | 'MonadBeam' instance inside which DuckDB queries are run. See the +-- for more information +newtype DuckDBM a + = DuckDBM + { -- | Run an IO action with access to a DuckDB connection and a debug logging + -- function, called or each query submitted on the connection. + runDuckDBM :: ReaderT (Text -> IO (), Connection) IO a + } + deriving (Monad, Functor, Applicative, MonadIO, MonadFail) + +-- | Execute queries on a DuckDB 'Connection'. +-- +-- To trace the exact SQL statements being generated, +-- use 'runBeamDuckDBDebug' instead. +runBeamDuckDB :: Connection -> DuckDBM a -> IO a +runBeamDuckDB = runBeamDuckDBDebug (\_ -> pure ()) + +runBeamDuckDBDebug :: (Text -> IO ()) -> Connection -> DuckDBM a -> IO a +runBeamDuckDBDebug debug conn action = + runReaderT (runDuckDBM action) (debug, conn) + +newtype BeamDuckDBParams = BeamDuckDBParams [SomeField] + +instance ToRow BeamDuckDBParams where + toRow (BeamDuckDBParams x) = map (\(SomeField f) -> toField f) x + +instance MonadBeam DuckDB DuckDBM where + runNoReturn (DuckDBCommandSyntax (DuckDBSyntax cmd vals)) = + DuckDBM $ do + (logger, conn) <- ask + let cmdString = Text.Lazy.toStrict (Builder.toLazyText (withPlaceholder cmd)) + liftIO (logger (cmdString <> ";\n-- With values: " <> Text.pack (show (DL.toList vals)))) + liftIO (void $ execute conn (Query cmdString) (BeamDuckDBParams (DL.toList vals))) + + runReturningMany (DuckDBCommandSyntax (DuckDBSyntax cmd vals)) action = + DuckDBM $ do + (logger, conn) <- ask + let cmdString = Text.Lazy.toStrict (Builder.toLazyText (withPlaceholder cmd)) + liftIO $ do + logger (cmdString <> ";\n-- With values: " <> Text.pack (show (DL.toList vals))) + withStatement conn (Query cmdString) $ \stmt -> + do + bind stmt (map toField (DL.toList vals)) + let nextRow' = + liftIO (nextRow stmt) >>= \case + Nothing -> pure Nothing + Just (BeamDuckDBRow row) -> pure row + runReaderT (runDuckDBM (action nextRow')) (logger, conn) + +newtype BeamDuckDBRow a = BeamDuckDBRow a + +instance (FromBackendRow DuckDB a) => FromRow (BeamDuckDBRow a) where + fromRow = BeamDuckDBRow <$> runF fromBackendRow' finish step + where + FromBackendRowM fromBackendRow' = fromBackendRow + + translateErrors :: Maybe Int -> SomeException -> Maybe SomeException + translateErrors col (SomeException e) = + case cast e of + Just + ( ConversionFailed + { errSQLType = typeString, + errHaskellType = hsString, + errMessage = msg + } + ) -> + Just + ( SomeException + ( BeamRowReadError + col + ( ColumnTypeMismatch + (Text.unpack hsString) + (Text.unpack typeString) + ("conversion failed: " ++ Text.unpack msg) + ) + ) + ) + Just (UnexpectedNull {}) -> + Just (SomeException (BeamRowReadError col ColumnUnexpectedNull)) + Just + ( Incompatible + { errSQLType = typeString, + errHaskellType = hsString, + errMessage = msg + } + ) -> + Just + ( SomeException + ( BeamRowReadError + col + ( ColumnTypeMismatch + (Text.unpack hsString) + (Text.unpack typeString) + ("incompatible: " ++ Text.unpack msg) + ) + ) + ) + Nothing -> Nothing + + finish = pure + + step :: forall a'. FromBackendRowF DuckDB (RowParser a') -> RowParser a' + step (ParseOneField next) = + RowParser $ ReaderT $ \ro -> StateT $ \st@(col, _) -> + case runStateT (runReaderT (runRowParser field) ro) st of + Ok (x, st') -> runStateT (runReaderT (runRowParser (next x)) ro) st' + Errors errs -> Errors (mapMaybe (translateErrors (Just col)) errs) + step (Alt (FromBackendRowM a) (FromBackendRowM b) next) = do + RowParser $ do + let RowParser a' = runF a finish step + RowParser b' = runF b finish step + + st <- lift get + ro <- ask + case runStateT (runReaderT a' ro) st of + Ok (ra, st') -> do + lift $ put st' + runRowParser (next ra) + Errors aErrs -> + case runStateT (runReaderT b' ro) st of + Ok (rb, st') -> do + lift $ put st' + runRowParser (next rb) + Errors bErrs -> + lift (lift (Errors (aErrs ++ bErrs))) + step (FailParseWith err) = RowParser (lift (lift (Errors [SomeException err]))) diff --git a/src/Database/Beam/DuckDB/Backend.hs b/src/Database/Beam/DuckDB/Backend.hs new file mode 100644 --- /dev/null +++ b/src/Database/Beam/DuckDB/Backend.hs @@ -0,0 +1,198 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +module Database.Beam.DuckDB.Backend (DuckDB) where + +import Data.ByteString (ByteString) +import Data.Data (Proxy (Proxy)) +import Data.Functor (($>)) +import Data.Int (Int16, Int32, Int64, Int8) +import Data.Text (Text) +import Data.Time (Day, LocalTime, TimeOfDay, UTCTime) +import Data.UUID.Types (UUID) +import Data.Word (Word16, Word32, Word64, Word8) +import Database.Beam (HasQBuilder, HasSqlEqualityCheck, HasSqlInTable (..), HasSqlQuantifiedEqualityCheck) +import Database.Beam.Backend ( + BeamBackend (..), + BeamSqlBackend, + BeamSqlBackendIsString, + BeamSqlBackendSyntax, + FromBackendRow, + SqlNull (..), + parseOneField, + ) +import Database.Beam.Backend.SQL (FromBackendRow (..)) +import Database.Beam.DuckDB.Syntax ( + DuckDBCommandSyntax, + DuckDBExpressionSyntax (..), + ) +import Database.Beam.DuckDB.Syntax.Builder ( + commas, + emit, + parens, + ) +import Database.Beam.Query.SQL92 (buildSql92Query') +import Database.Beam.Query.Types (HasQBuilder (..)) +import Database.DuckDB.Simple (Null) +import Database.DuckDB.Simple.FromField (FromField) + +data DuckDB + +type instance BeamSqlBackendSyntax DuckDB = DuckDBCommandSyntax + +instance BeamSqlBackend DuckDB + +instance HasQBuilder DuckDB where + buildSqlQuery = buildSql92Query' True + +instance HasSqlInTable DuckDB where + inRowValuesE Proxy e es = + DuckDBExpressionSyntax $ + mconcat + [ parens $ fromDuckDBExpression e + , emit " IN " + , parens $ emit "VALUES " <> commas (map fromDuckDBExpression es) + ] + +instance BeamSqlBackendIsString DuckDB Text + +instance BeamSqlBackendIsString DuckDB String + +instance BeamBackend DuckDB where + type BackendFromField DuckDB = FromField + +instance FromBackendRow DuckDB SqlNull where + fromBackendRow = parseOneField @DuckDB @Null $> SqlNull + +instance FromBackendRow DuckDB Bool + +instance FromBackendRow DuckDB Float + +instance FromBackendRow DuckDB Double + +instance FromBackendRow DuckDB Integer + +instance FromBackendRow DuckDB Int + +instance FromBackendRow DuckDB Int8 + +instance FromBackendRow DuckDB Int16 + +instance FromBackendRow DuckDB Int32 + +instance FromBackendRow DuckDB Int64 + +instance FromBackendRow DuckDB Word + +instance FromBackendRow DuckDB Word8 + +instance FromBackendRow DuckDB Word16 + +instance FromBackendRow DuckDB Word32 + +instance FromBackendRow DuckDB Word64 + +instance FromBackendRow DuckDB Text + +instance FromBackendRow DuckDB ByteString + +instance FromBackendRow DuckDB UUID + +instance FromBackendRow DuckDB Day + +instance FromBackendRow DuckDB TimeOfDay + +instance FromBackendRow DuckDB LocalTime + +instance FromBackendRow DuckDB UTCTime + +instance HasSqlEqualityCheck DuckDB Bool + +instance HasSqlEqualityCheck DuckDB Float + +instance HasSqlEqualityCheck DuckDB Double + +instance HasSqlEqualityCheck DuckDB Integer + +instance HasSqlEqualityCheck DuckDB Int + +instance HasSqlEqualityCheck DuckDB Int8 + +instance HasSqlEqualityCheck DuckDB Int16 + +instance HasSqlEqualityCheck DuckDB Int32 + +instance HasSqlEqualityCheck DuckDB Int64 + +instance HasSqlEqualityCheck DuckDB Word + +instance HasSqlEqualityCheck DuckDB Word8 + +instance HasSqlEqualityCheck DuckDB Word16 + +instance HasSqlEqualityCheck DuckDB Word32 + +instance HasSqlEqualityCheck DuckDB Word64 + +instance HasSqlEqualityCheck DuckDB Text + +instance HasSqlEqualityCheck DuckDB ByteString + +instance HasSqlEqualityCheck DuckDB UUID + +instance HasSqlEqualityCheck DuckDB Day + +instance HasSqlEqualityCheck DuckDB TimeOfDay + +instance HasSqlEqualityCheck DuckDB LocalTime + +instance HasSqlEqualityCheck DuckDB UTCTime + +instance HasSqlQuantifiedEqualityCheck DuckDB Bool + +instance HasSqlQuantifiedEqualityCheck DuckDB Float + +instance HasSqlQuantifiedEqualityCheck DuckDB Double + +instance HasSqlQuantifiedEqualityCheck DuckDB Integer + +instance HasSqlQuantifiedEqualityCheck DuckDB Int + +instance HasSqlQuantifiedEqualityCheck DuckDB Int8 + +instance HasSqlQuantifiedEqualityCheck DuckDB Int16 + +instance HasSqlQuantifiedEqualityCheck DuckDB Int32 + +instance HasSqlQuantifiedEqualityCheck DuckDB Int64 + +instance HasSqlQuantifiedEqualityCheck DuckDB Word + +instance HasSqlQuantifiedEqualityCheck DuckDB Word8 + +instance HasSqlQuantifiedEqualityCheck DuckDB Word16 + +instance HasSqlQuantifiedEqualityCheck DuckDB Word32 + +instance HasSqlQuantifiedEqualityCheck DuckDB Word64 + +instance HasSqlQuantifiedEqualityCheck DuckDB Text + +instance HasSqlQuantifiedEqualityCheck DuckDB ByteString + +instance HasSqlQuantifiedEqualityCheck DuckDB UUID + +instance HasSqlQuantifiedEqualityCheck DuckDB Day + +instance HasSqlQuantifiedEqualityCheck DuckDB TimeOfDay + +instance HasSqlQuantifiedEqualityCheck DuckDB LocalTime + +instance HasSqlQuantifiedEqualityCheck DuckDB UTCTime diff --git a/src/Database/Beam/DuckDB/Syntax.hs b/src/Database/Beam/DuckDB/Syntax.hs new file mode 100644 --- /dev/null +++ b/src/Database/Beam/DuckDB/Syntax.hs @@ -0,0 +1,589 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} +-- TODO: clean up unused top binds +{-# OPTIONS_GHC -Wno-unused-top-binds #-} + +module Database.Beam.DuckDB.Syntax ( + -- * Command + DuckDBCommandSyntax (..), + + -- * Concrete syntaxes + DuckDBSelectSyntax (..), + DuckDBFromSyntax (..), + DuckDBExpressionSyntax (..), + DuckDBInsertSyntax (..), + DuckDBUpdateSyntax (..), + DuckDBDeleteSyntax (..), + DuckDBTableNameSyntax (..), + DuckDBOnConflictSyntax (..), +) +where + +import Data.Coerce (coerce) +import Data.Int (Int16, Int32, Int64, Int8) +import Data.String (fromString) +import Data.Text (Text) +import qualified Data.Text as Text +import Data.Word (Word16, Word32, Word64, Word8) +import Database.Beam.Backend ( + HasSqlValueSyntax (..), + IsSql92AggregationExpressionSyntax (..), + IsSql92AggregationSetQuantifierSyntax (..), + IsSql92DataTypeSyntax (..), + IsSql92DeleteSyntax (..), + IsSql92ExpressionSyntax (..), + IsSql92ExtractFieldSyntax (..), + IsSql92FieldNameSyntax (..), + IsSql92FromOuterJoinSyntax (..), + IsSql92FromSyntax (..), + IsSql92GroupingSyntax (..), + IsSql92InsertSyntax (..), + IsSql92InsertValuesSyntax (..), + IsSql92OrderingSyntax (..), + IsSql92ProjectionSyntax (..), + IsSql92QuantifierSyntax (..), + IsSql92SelectSyntax (..), + IsSql92SelectTableSyntax (..), + IsSql92Syntax (..), + IsSql92TableNameSyntax (..), + IsSql92TableSourceSyntax (..), + IsSql92UpdateSyntax (..), + SqlNull (..), + ) +import Database.Beam.DuckDB.Syntax.Builder (DuckDBSyntax, commas, emit, emitIntegral, emitRealFloat, emitValue, parens, quotedIdentifier, spaces) +import Database.Beam.Migrate.Serialization (BeamSerializedDataType) +import Database.DuckDB.Simple (Null (Null), ToField) + +newtype DuckDBCommandSyntax = DuckDBCommandSyntax {fromDuckDBSyntax :: DuckDBSyntax} + +-- +instance IsSql92Syntax DuckDBCommandSyntax where + type Sql92SelectSyntax DuckDBCommandSyntax = DuckDBSelectSyntax + type Sql92InsertSyntax DuckDBCommandSyntax = DuckDBInsertSyntax + type Sql92UpdateSyntax DuckDBCommandSyntax = DuckDBUpdateSyntax + type Sql92DeleteSyntax DuckDBCommandSyntax = DuckDBDeleteSyntax + + selectCmd = DuckDBCommandSyntax . fromDuckDBSelect + insertCmd = DuckDBCommandSyntax . fromDuckDBInsert + updateCmd = DuckDBCommandSyntax . fromDuckDBUpdate + deleteCmd = DuckDBCommandSyntax . fromDuckDBDelete + +newtype DuckDBTableNameSyntax = DuckDBTableNameSyntax {fromDuckDBTableName :: DuckDBSyntax} + +-- | DuckDB @SELECT@ syntax +newtype DuckDBSelectSyntax = DuckDBSelectSyntax {fromDuckDBSelect :: DuckDBSyntax} + +newtype DuckDBSelectTableSyntax = DuckDBSelectTableSyntax {fromDuckDBSelectTable :: DuckDBSyntax} + +data DuckDBOrderingSyntax = DuckDBOrderingSyntax + { duckDBOrdering :: DuckDBSyntax + , -- DuckDB re-uses the Postgres parser, and therefore + -- has the same null ordering properties + duckDBNullOrdering :: Maybe DuckDBNullOrdering + } + +data DuckDBNullOrdering + = DuckDBNullOrderingNullsFirst + | DuckDBNullOrderingNullsLast + deriving (Show, Eq) + +newtype DuckDBExpressionSyntax = DuckDBExpressionSyntax {fromDuckDBExpression :: DuckDBSyntax} deriving (Eq) + +newtype DuckDBTableSourceSyntax = DuckDBTableSourceSyntax {fromDuckDBTableSource :: DuckDBSyntax} + +newtype DuckDBProjectionSyntax = DuckDBProjectionSyntax {fromDuckDBProjection :: DuckDBSyntax} + +newtype DuckDBFromSyntax = DuckDBFromSyntax {fromDuckDBFrom :: DuckDBSyntax} + +newtype DuckDBValueSyntax = DuckDBValueSyntax {fromDuckDBValue :: DuckDBSyntax} + +newtype DuckDBFieldNameSyntax = DuckDBFieldNameSyntax {fromDuckDBFieldName :: DuckDBSyntax} + +newtype DuckDBComparisonQuantifierSyntax = DuckDBComparisonQuantifierSyntax {fromDuckDBComparisonQuantifier :: DuckDBSyntax} + +data DuckDBDataTypeSyntax = DuckDBDataTypeSyntax + { duckDBDataType :: DuckDBSyntax + , duckDBDataTypeSerialized :: BeamSerializedDataType + } + +newtype DuckDBExtractFieldSyntax = DuckDBExtractFieldSyntax {fromDuckDBExtractField :: DuckDBSyntax} + +newtype DuckDBGroupingSyntax = DuckDBGroupingSyntax {fromDuckDBGrouping :: DuckDBSyntax} + +newtype DuckDBSelectSetQuantifierSyntax = DuckDBSelectSetQuantifierSyntax {fromDuckDBSelectSetQuantifier :: DuckDBSyntax} + +newtype DuckDBAggregationSetQuantifierSyntax = DuckDBAggregationSetQuantifierSyntax {fromDuckDBAggregationSetQuantifier :: DuckDBSyntax} + +instance IsSql92AggregationExpressionSyntax DuckDBExpressionSyntax where + type Sql92AggregationSetQuantifierSyntax DuckDBExpressionSyntax = DuckDBAggregationSetQuantifierSyntax + + countAllE = DuckDBExpressionSyntax (emit "COUNT(*)") + countE = aggFunc "COUNT" + avgE = aggFunc "AVG" + sumE = aggFunc "SUM" + minE = aggFunc "MIN" + maxE = aggFunc "MAX" + +aggFunc :: Text -> Maybe DuckDBAggregationSetQuantifierSyntax -> DuckDBExpressionSyntax -> DuckDBExpressionSyntax +aggFunc fn q e = + DuckDBExpressionSyntax $ + emit fn <> emit "(" <> maybe mempty (\qInner -> fromDuckDBAggregationSetQuantifier qInner <> emit " ") q <> fromDuckDBExpression e <> emit ")" + +instance IsSql92SelectTableSyntax DuckDBSelectTableSyntax where + type Sql92SelectTableSelectSyntax DuckDBSelectTableSyntax = DuckDBSelectSyntax + type Sql92SelectTableExpressionSyntax DuckDBSelectTableSyntax = DuckDBExpressionSyntax + type Sql92SelectTableProjectionSyntax DuckDBSelectTableSyntax = DuckDBProjectionSyntax + type Sql92SelectTableFromSyntax DuckDBSelectTableSyntax = DuckDBFromSyntax + type Sql92SelectTableGroupingSyntax DuckDBSelectTableSyntax = DuckDBGroupingSyntax + type Sql92SelectTableSetQuantifierSyntax DuckDBSelectTableSyntax = DuckDBAggregationSetQuantifierSyntax + + selectTableStmt setQuantifier proj from where_ grouping having = + DuckDBSelectTableSyntax $ + mconcat + [ emit "SELECT " + , maybe mempty ((<> emit " ") . fromDuckDBAggregationSetQuantifier) setQuantifier + , fromDuckDBProjection proj + , maybe mempty ((emit " FROM " <>) . fromDuckDBFrom) from + , maybe mempty ((emit " WHERE " <>) . fromDuckDBExpression) where_ + , maybe mempty ((emit " GROUP BY " <>) . fromDuckDBGrouping) grouping + , maybe mempty ((emit " HAVING " <>) . fromDuckDBExpression) having + ] + + unionTables unionAll = tableOp (if unionAll then "UNION ALL" else "UNION") + intersectTables intersectAcc = tableOp (if intersectAcc then "INTERSECT ALL" else "INTERSECT") + exceptTable exceptAll = tableOp (if exceptAll then "EXCEPT ALL" else "EXCEPT") + +tableOp :: Text -> DuckDBSelectTableSyntax -> DuckDBSelectTableSyntax -> DuckDBSelectTableSyntax +tableOp op a b = + DuckDBSelectTableSyntax $ + fromDuckDBSelectTable a <> spaces (emit op) <> fromDuckDBSelectTable b + +instance IsSql92QuantifierSyntax DuckDBComparisonQuantifierSyntax where + quantifyOverAll = DuckDBComparisonQuantifierSyntax (emit "ALL") + quantifyOverAny = DuckDBComparisonQuantifierSyntax (emit "ANY") + +instance IsSql92ExtractFieldSyntax DuckDBExtractFieldSyntax where + secondsField = DuckDBExtractFieldSyntax (emit "SECOND") + minutesField = DuckDBExtractFieldSyntax (emit "MINUTE") + hourField = DuckDBExtractFieldSyntax (emit "HOUR") + dayField = DuckDBExtractFieldSyntax (emit "DAY") + monthField = DuckDBExtractFieldSyntax (emit "MONTH") + yearField = DuckDBExtractFieldSyntax (emit "YEAR") + +instance IsSql92AggregationSetQuantifierSyntax DuckDBAggregationSetQuantifierSyntax where + setQuantifierDistinct = DuckDBAggregationSetQuantifierSyntax $ emit "DISTINCT" + setQuantifierAll = DuckDBAggregationSetQuantifierSyntax $ emit "ALL" + +instance IsSql92AggregationSetQuantifierSyntax DuckDBSelectSetQuantifierSyntax where + setQuantifierDistinct = DuckDBSelectSetQuantifierSyntax $ emit "DISTINCT" + setQuantifierAll = DuckDBSelectSetQuantifierSyntax $ emit "ALL" + +instance IsSql92GroupingSyntax DuckDBGroupingSyntax where + type Sql92GroupingExpressionSyntax DuckDBGroupingSyntax = DuckDBExpressionSyntax + + groupByExpressions es = + DuckDBGroupingSyntax $ + commas (map fromDuckDBExpression es) + +instance IsSql92TableNameSyntax DuckDBTableNameSyntax where + tableName Nothing t = DuckDBTableNameSyntax (quotedIdentifier t) + tableName (Just s) t = DuckDBTableNameSyntax (quotedIdentifier s <> emit "." <> quotedIdentifier t) + +instance IsSql92TableSourceSyntax DuckDBTableSourceSyntax where + type Sql92TableSourceSelectSyntax DuckDBTableSourceSyntax = DuckDBSelectSyntax + type Sql92TableSourceExpressionSyntax DuckDBTableSourceSyntax = DuckDBExpressionSyntax + type Sql92TableSourceTableNameSyntax DuckDBTableSourceSyntax = DuckDBTableNameSyntax + + tableNamed = DuckDBTableSourceSyntax . fromDuckDBTableName + tableFromSubSelect s = DuckDBTableSourceSyntax $ emit "(" <> fromDuckDBSelect s <> emit ")" + tableFromValues vss = + DuckDBTableSourceSyntax . parens $ + emit "VALUES " + <> commas + ( map + (parens . commas . map fromDuckDBExpression) + vss + ) + +instance IsSql92FromSyntax DuckDBFromSyntax where + type Sql92FromExpressionSyntax DuckDBFromSyntax = DuckDBExpressionSyntax + type Sql92FromTableSourceSyntax DuckDBFromSyntax = DuckDBTableSourceSyntax + + fromTable tableSrc Nothing = coerce tableSrc + fromTable tableSrc (Just (nm, colNms)) = + DuckDBFromSyntax $ + coerce tableSrc + <> emit " AS " + <> quotedIdentifier nm + <> maybe mempty (parens . commas . map quotedIdentifier) colNms + + innerJoin a b Nothing = DuckDBFromSyntax (fromDuckDBFrom a <> emit " CROSS JOIN " <> fromDuckDBFrom b) + innerJoin a b (Just e) = join "INNER JOIN" a b (Just e) + + leftJoin = join "LEFT JOIN" + rightJoin = join "RIGHT JOIN" + +join :: Text -> DuckDBFromSyntax -> DuckDBFromSyntax -> Maybe DuckDBExpressionSyntax -> DuckDBFromSyntax +join joinType a b Nothing = + DuckDBFromSyntax $ + fromDuckDBFrom a <> emit (" " <> joinType <> " ") <> fromDuckDBFrom b <> emit " ON TRUE" +join joinType a b (Just on) = + DuckDBFromSyntax $ + fromDuckDBFrom a + <> emit (" " <> joinType <> " ") + <> fromDuckDBFrom b + <> emit " ON " + <> fromDuckDBExpression on + +instance IsSql92FromOuterJoinSyntax DuckDBFromSyntax where + outerJoin = join "FULL OUTER JOIN" + +instance IsSql92FieldNameSyntax DuckDBFieldNameSyntax where + qualifiedField a b = + DuckDBFieldNameSyntax $ + quotedIdentifier a <> emit "." <> quotedIdentifier b + unqualifiedField = DuckDBFieldNameSyntax . quotedIdentifier + +instance IsSql92DataTypeSyntax DuckDBDataTypeSyntax where + domainType nm = + DuckDBDataTypeSyntax + (quotedIdentifier nm) + (domainType nm) + + charType prec charSet = + DuckDBDataTypeSyntax + (emit "CHAR" <> optPrec prec <> optCharSet charSet) + (charType prec charSet) + varCharType prec charSet = + DuckDBDataTypeSyntax + (emit "VARCHAR" <> optPrec prec <> optCharSet charSet) + (varCharType prec charSet) + nationalCharType prec = + DuckDBDataTypeSyntax + (emit "NATIONAL CHAR" <> optPrec prec) + (nationalCharType prec) + nationalVarCharType prec = + DuckDBDataTypeSyntax + (emit "NATIONAL CHARACTER VARYING" <> optPrec prec) + (nationalVarCharType prec) + + bitType prec = + DuckDBDataTypeSyntax + (emit "BIT" <> optPrec prec) + (bitType prec) + varBitType prec = + DuckDBDataTypeSyntax + (emit "BIT VARYING" <> optPrec prec) + (varBitType prec) + + numericType prec = + DuckDBDataTypeSyntax + (emit "NUMERIC" <> optNumericPrec prec) + (numericType prec) + decimalType prec = + DuckDBDataTypeSyntax + (emit "DECIMAL" <> optNumericPrec prec) + (decimalType prec) + + intType = DuckDBDataTypeSyntax (emit "INT") intType + smallIntType = DuckDBDataTypeSyntax (emit "SMALLINT") smallIntType + + floatType prec = + DuckDBDataTypeSyntax + (emit "FLOAT" <> optPrec prec) + (floatType prec) + doubleType = DuckDBDataTypeSyntax (emit "DOUBLE PRECISION") doubleType + realType = DuckDBDataTypeSyntax (emit "REAL") realType + dateType = DuckDBDataTypeSyntax (emit "DATE") dateType + timeType prec withTz = + DuckDBDataTypeSyntax + (emit "TIME" <> optPrec prec <> if withTz then emit " WITH TIME ZONE" else mempty) + (timeType prec withTz) + timestampType prec withTz = + DuckDBDataTypeSyntax + (emit "TIMESTAMP" <> optPrec prec <> if withTz then emit " WITH TIME ZONE" else mempty) + (timestampType prec withTz) + +optPrec :: Maybe Word -> DuckDBSyntax +optPrec Nothing = mempty +optPrec (Just x) = emit "(" <> emit (fromString (show x)) <> emit ")" + +optCharSet :: Maybe Text -> DuckDBSyntax +optCharSet Nothing = mempty +optCharSet (Just cs) = emit " CHARACTER SET " <> emit cs + +optNumericPrec :: Maybe (Word, Maybe Word) -> DuckDBSyntax +optNumericPrec Nothing = mempty +optNumericPrec (Just (prec, Nothing)) = optPrec (Just prec) +optNumericPrec (Just (prec, Just dec)) = emit "(" <> emit (fromString (show prec)) <> emit ", " <> emit (fromString (show dec)) <> emit ")" + +instance HasSqlValueSyntax DuckDBValueSyntax Int8 where + sqlValueSyntax i = DuckDBValueSyntax (emitIntegral i) + +instance HasSqlValueSyntax DuckDBValueSyntax Int16 where + sqlValueSyntax i = DuckDBValueSyntax (emitIntegral i) + +instance HasSqlValueSyntax DuckDBValueSyntax Int32 where + sqlValueSyntax i = DuckDBValueSyntax (emitIntegral i) + +instance HasSqlValueSyntax DuckDBValueSyntax Int64 where + sqlValueSyntax i = DuckDBValueSyntax (emitIntegral i) + +instance HasSqlValueSyntax DuckDBValueSyntax Word8 where + sqlValueSyntax i = DuckDBValueSyntax (emitIntegral i) + +instance HasSqlValueSyntax DuckDBValueSyntax Word16 where + sqlValueSyntax i = DuckDBValueSyntax (emitIntegral i) + +instance HasSqlValueSyntax DuckDBValueSyntax Word32 where + sqlValueSyntax i = DuckDBValueSyntax (emitIntegral i) + +instance HasSqlValueSyntax DuckDBValueSyntax Word64 where + sqlValueSyntax i = DuckDBValueSyntax (emitIntegral i) + +instance HasSqlValueSyntax DuckDBValueSyntax Float where + sqlValueSyntax f = DuckDBValueSyntax (emitRealFloat f) + +instance HasSqlValueSyntax DuckDBValueSyntax Double where + sqlValueSyntax f = DuckDBValueSyntax (emitRealFloat f) + +instance HasSqlValueSyntax DuckDBValueSyntax Bool where + sqlValueSyntax = DuckDBValueSyntax . emitValue + +instance HasSqlValueSyntax DuckDBValueSyntax SqlNull where + sqlValueSyntax _ = DuckDBValueSyntax (emitValue Null) + +instance HasSqlValueSyntax DuckDBValueSyntax String where + sqlValueSyntax = sqlValueSyntax . Text.pack + +instance HasSqlValueSyntax DuckDBValueSyntax Text where + sqlValueSyntax = DuckDBValueSyntax . emitValue + +instance {-# OVERLAPPABLE #-} (ToField a, Eq a) => HasSqlValueSyntax DuckDBValueSyntax a where + sqlValueSyntax = DuckDBValueSyntax . emitValue + +instance (HasSqlValueSyntax DuckDBValueSyntax x) => HasSqlValueSyntax DuckDBValueSyntax (Maybe x) where + sqlValueSyntax (Just x) = sqlValueSyntax x + sqlValueSyntax Nothing = sqlValueSyntax SqlNull + +instance IsSql92ExpressionSyntax DuckDBExpressionSyntax where + type Sql92ExpressionValueSyntax DuckDBExpressionSyntax = DuckDBValueSyntax + type Sql92ExpressionSelectSyntax DuckDBExpressionSyntax = DuckDBSelectSyntax + type Sql92ExpressionFieldNameSyntax DuckDBExpressionSyntax = DuckDBFieldNameSyntax + type Sql92ExpressionQuantifierSyntax DuckDBExpressionSyntax = DuckDBComparisonQuantifierSyntax + type Sql92ExpressionCastTargetSyntax DuckDBExpressionSyntax = DuckDBDataTypeSyntax + type Sql92ExpressionExtractFieldSyntax DuckDBExpressionSyntax = DuckDBExtractFieldSyntax + + addE = binOp "+" + subE = binOp "-" + mulE = binOp "*" + divE = binOp "/" + modE = binOp "%" + orE = binOp "OR" + andE = binOp "AND" + likeE = binOp "LIKE" + overlapsE = binOp "OVERLAPS" + + eqE = compOp "=" + neqE = compOp "<>" + ltE = compOp "<" + gtE = compOp ">" + leE = compOp "<=" + geE = compOp ">=" + + negateE = unOp "-" + notE = unOp "NOT" + + isNotNullE = postFix "IS NOT NULL" + isNullE = postFix "IS NULL" + + -- SQLite doesn't handle tri-state booleans properly + isTrueE = postFix "IS 1" + isNotTrueE = postFix "IS NOT 1" + isFalseE = postFix "IS 0" + isNotFalseE = postFix "IS NOT 0" + isUnknownE = postFix "IS NULL" + isNotUnknownE = postFix "IS NOT NULL" + + existsE select = DuckDBExpressionSyntax (emit "EXISTS " <> parens (fromDuckDBSelect select)) + uniqueE select = DuckDBExpressionSyntax (emit "UNIQUE " <> parens (fromDuckDBSelect select)) + + betweenE a b c = + DuckDBExpressionSyntax + ( parens (fromDuckDBExpression a) + <> emit " BETWEEN " + <> parens (fromDuckDBExpression b) + <> emit " AND " + <> parens (fromDuckDBExpression c) + ) + + valueE = DuckDBExpressionSyntax . fromDuckDBValue + + rowE vs = DuckDBExpressionSyntax (parens (commas (map fromDuckDBExpression vs))) + fieldE = DuckDBExpressionSyntax . fromDuckDBFieldName + + subqueryE = DuckDBExpressionSyntax . parens . fromDuckDBSelect + + positionE needle haystack = + DuckDBExpressionSyntax $ + emit "POSITION" <> parens (parens (fromDuckDBExpression needle) <> emit " IN " <> parens (fromDuckDBExpression haystack)) + nullIfE a b = + DuckDBExpressionSyntax $ + emit "NULLIF" <> parens (fromDuckDBExpression a <> emit ", " <> fromDuckDBExpression b) + absE x = DuckDBExpressionSyntax (emit "ABS" <> parens (fromDuckDBExpression x)) + bitLengthE x = DuckDBExpressionSyntax (emit "8 * LENGTH" <> parens (emit "CAST" <> parens (parens (fromDuckDBExpression x) <> emit " AS BLOB"))) + charLengthE x = DuckDBExpressionSyntax (emit "LENGTH" <> parens (fromDuckDBExpression x)) + octetLengthE x = DuckDBExpressionSyntax (emit "LENGTH" <> parens (emit "CAST" <> parens (parens (fromDuckDBExpression x) <> emit " AS BLOB"))) + lowerE x = DuckDBExpressionSyntax (emit "LOWER" <> parens (fromDuckDBExpression x)) + upperE x = DuckDBExpressionSyntax (emit "UPPER" <> parens (fromDuckDBExpression x)) + trimE x = DuckDBExpressionSyntax (emit "TRIM" <> parens (fromDuckDBExpression x)) + coalesceE es = DuckDBExpressionSyntax (emit "COALESCE" <> parens (commas (map fromDuckDBExpression es))) + extractE field from = DuckDBExpressionSyntax (emit "EXTRACT(" <> fromDuckDBExtractField field <> emit " FROM (" <> fromDuckDBExpression from <> emit "))") + castE e t = DuckDBExpressionSyntax (emit "CAST" <> parens (parens (fromDuckDBExpression e) <> emit " AS " <> duckDBDataType t)) + caseE cases else_ = + DuckDBExpressionSyntax $ + emit "CASE " + <> foldMap (\(cond, res) -> emit "WHEN " <> fromDuckDBExpression cond <> emit " THEN " <> fromDuckDBExpression res <> emit " ") cases + <> emit "ELSE " + <> fromDuckDBExpression else_ + <> emit " END" + + currentTimestampE = DuckDBExpressionSyntax (emit "CURRENT_TIMESTAMP") + + defaultE = DuckDBExpressionSyntax (emit "DEFAULT") + inE e es = DuckDBExpressionSyntax (parens (fromDuckDBExpression e) <> emit " IN " <> parens (commas (map fromDuckDBExpression es))) + inSelectE e sel = + DuckDBExpressionSyntax (parens (fromDuckDBExpression e) <> emit " IN " <> parens (fromDuckDBSelect sel)) + +binOp :: Text -> DuckDBExpressionSyntax -> DuckDBExpressionSyntax -> DuckDBExpressionSyntax +binOp op a b = + DuckDBExpressionSyntax $ + parens (fromDuckDBExpression a) <> emit " " <> emit op <> emit " " <> parens (fromDuckDBExpression b) + +compOp :: + Text -> + Maybe DuckDBComparisonQuantifierSyntax -> + DuckDBExpressionSyntax -> + DuckDBExpressionSyntax -> + DuckDBExpressionSyntax +compOp op quantifier a b = + DuckDBExpressionSyntax $ + parens (fromDuckDBExpression a) + <> emit op + <> maybe mempty (\q -> emit " " <> fromDuckDBComparisonQuantifier q <> emit " ") quantifier + <> parens (fromDuckDBExpression b) + +unOp, postFix :: Text -> DuckDBExpressionSyntax -> DuckDBExpressionSyntax +unOp op a = + DuckDBExpressionSyntax (emit op <> parens (fromDuckDBExpression a)) +postFix op a = + DuckDBExpressionSyntax (parens (fromDuckDBExpression a) <> emit " " <> emit op) + +instance IsSql92ProjectionSyntax DuckDBProjectionSyntax where + type Sql92ProjectionExpressionSyntax DuckDBProjectionSyntax = DuckDBExpressionSyntax + + projExprs exprs = + DuckDBProjectionSyntax $ + commas + ( map + ( \(expr, nm) -> + fromDuckDBExpression expr + <> maybe mempty (\nm' -> emit " AS " <> quotedIdentifier nm') nm + ) + exprs + ) + +instance IsSql92OrderingSyntax DuckDBOrderingSyntax where + type Sql92OrderingExpressionSyntax DuckDBOrderingSyntax = DuckDBExpressionSyntax + + ascOrdering e = DuckDBOrderingSyntax (fromDuckDBExpression e <> emit " ASC") Nothing + descOrdering e = DuckDBOrderingSyntax (fromDuckDBExpression e <> emit " DESC") Nothing + +instance IsSql92SelectSyntax DuckDBSelectSyntax where + type Sql92SelectSelectTableSyntax DuckDBSelectSyntax = DuckDBSelectTableSyntax + type Sql92SelectOrderingSyntax DuckDBSelectSyntax = DuckDBOrderingSyntax + + selectStmt tbl ordering limit offset = + DuckDBSelectSyntax $ + mconcat + [ fromDuckDBSelectTable tbl + , case ordering of + [] -> mempty + ordering' -> emit " ORDER BY " <> commas (map duckDBOrdering ordering') + , maybe mempty (emit . fromString . (" LIMIT " <>) . show) limit + , maybe mempty (emit . fromString . (" OFFSET " <>) . show) offset + ] + +-- | DuckDB @ON CONFLICT@ syntax +newtype DuckDBOnConflictSyntax = DuckDBOnConflictSyntax {fromDuckDBOnConflict :: DuckDBSyntax} + +-- | SQLite @INSERT@ syntax. +newtype DuckDBInsertSyntax = DuckDBInsertSyntax {fromDuckDBInsert :: DuckDBSyntax} + +newtype DuckDBInsertValuesSyntax = DuckDBInsertValuesSyntax {fromDuckDBInsertValues :: DuckDBSyntax} + +instance IsSql92InsertSyntax DuckDBInsertSyntax where + type Sql92InsertTableNameSyntax DuckDBInsertSyntax = DuckDBTableNameSyntax + type Sql92InsertValuesSyntax DuckDBInsertSyntax = DuckDBInsertValuesSyntax + + insertStmt tblName fields values = + DuckDBInsertSyntax $ + emit "INSERT INTO " + <> fromDuckDBTableName tblName + <> emit "(" + <> commas (map quotedIdentifier fields) + <> emit ") " + <> fromDuckDBInsertValues values + +instance IsSql92InsertValuesSyntax DuckDBInsertValuesSyntax where + type Sql92InsertValuesExpressionSyntax DuckDBInsertValuesSyntax = DuckDBExpressionSyntax + type Sql92InsertValuesSelectSyntax DuckDBInsertValuesSyntax = DuckDBSelectSyntax + + insertSqlExpressions es = + DuckDBInsertValuesSyntax $ + emit "VALUES " + <> commas + ( map + (\values -> emit "(" <> commas (coerce values) <> emit ")") + es + ) + insertFromSql (DuckDBSelectSyntax a) = DuckDBInsertValuesSyntax a + +-- | DuckDB @UPDATE@ syntax +newtype DuckDBUpdateSyntax = DuckDBUpdateSyntax {fromDuckDBUpdate :: DuckDBSyntax} + +instance IsSql92UpdateSyntax DuckDBUpdateSyntax where + type Sql92UpdateFieldNameSyntax DuckDBUpdateSyntax = DuckDBFieldNameSyntax + type Sql92UpdateExpressionSyntax DuckDBUpdateSyntax = DuckDBExpressionSyntax + type Sql92UpdateTableNameSyntax DuckDBUpdateSyntax = DuckDBTableNameSyntax + + updateStmt tbl fields where_ = + DuckDBUpdateSyntax $ + emit "UPDATE " + <> fromDuckDBTableName tbl + <> ( case fields of + [] -> mempty + fs -> + emit " SET " + <> commas (map (\(field, val) -> fromDuckDBFieldName field <> emit "=" <> fromDuckDBExpression val) fs) + ) + <> maybe mempty (\whereInner -> emit " WHERE " <> fromDuckDBExpression whereInner) where_ + +-- | DuckDB @DELETE@ syntax +newtype DuckDBDeleteSyntax = DuckDBDeleteSyntax {fromDuckDBDelete :: DuckDBSyntax} + +instance IsSql92DeleteSyntax DuckDBDeleteSyntax where + type Sql92DeleteExpressionSyntax DuckDBDeleteSyntax = DuckDBExpressionSyntax + type Sql92DeleteTableNameSyntax DuckDBDeleteSyntax = DuckDBTableNameSyntax + + deleteStmt tbl alias where_ = + DuckDBDeleteSyntax $ + emit "DELETE FROM " + <> fromDuckDBTableName tbl + <> maybe mempty (\alias_ -> emit " AS " <> quotedIdentifier alias_) alias + <> maybe mempty (\whereInner -> emit " WHERE " <> fromDuckDBExpression whereInner) where_ + + deleteSupportsAlias _ = True diff --git a/src/Database/Beam/DuckDB/Syntax/Builder.hs b/src/Database/Beam/DuckDB/Syntax/Builder.hs new file mode 100644 --- /dev/null +++ b/src/Database/Beam/DuckDB/Syntax/Builder.hs @@ -0,0 +1,126 @@ +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE OverloadedStrings #-} + +module Database.Beam.DuckDB.Syntax.Builder + ( DuckDBSyntax (..), + SomeField (..), + emitChar, + emit, + emitIntegral, + emitRealFloat, + emit', + emitValue, + spaces, + parens, + sepBy, + commas, + quotedIdentifier, + withPlaceholder, + ) +where + +import Data.DList (DList) +import qualified Data.DList as DL +import Data.Text (Text) +import qualified Data.Text as Text +import qualified Data.Text.Lazy as Text.Lazy +import Data.Text.Lazy.Builder (Builder) +import qualified Data.Text.Lazy.Builder as Builder +import qualified Data.Text.Lazy.Builder.Int as Builder +import qualified Data.Text.Lazy.Builder.RealFloat as Builder +import Database.Beam.Backend (Sql92DisplaySyntax (..)) +import Database.DuckDB.Simple (ToField (toField)) +import Database.DuckDB.Simple.ToField (renderFieldBinding) + +data SomeField = forall a. (ToField a, Eq a) => SomeField a + +instance Show SomeField where + show (SomeField f) = renderFieldBinding (toField f) + +instance ToField SomeField where + toField (SomeField f) = toField f + +data DuckDBSyntax + = DuckDBSyntax + ((SomeField -> Builder) -> Builder) + (DList SomeField) + +emitChar :: Char -> DuckDBSyntax +emitChar c = DuckDBSyntax (const (Builder.singleton c)) mempty + +emit :: Text -> DuckDBSyntax +emit t = DuckDBSyntax (const (Builder.fromText t)) mempty + +emitIntegral :: (Integral a) => a -> DuckDBSyntax +emitIntegral i = DuckDBSyntax (const (Builder.decimal i)) mempty + +emitRealFloat :: (RealFloat f) => f -> DuckDBSyntax +emitRealFloat f = DuckDBSyntax (const (Builder.realFloat f)) mempty + +emit' :: (Show a) => a -> DuckDBSyntax +emit' s = DuckDBSyntax (const (Builder.fromString (show s))) mempty + +-- | Emit a properly escaped value into the syntax +emitValue :: (ToField a, Eq a) => a -> DuckDBSyntax +emitValue v = + DuckDBSyntax + ($ SomeField v) + (DL.singleton (SomeField v)) + +spaces :: DuckDBSyntax -> DuckDBSyntax +spaces a = emit " " <> a <> emit " " + +parens :: DuckDBSyntax -> DuckDBSyntax +parens a = emit "(" <> a <> emit ")" + +sepBy :: DuckDBSyntax -> [DuckDBSyntax] -> DuckDBSyntax +sepBy _ [] = mempty +sepBy _ [x] = x +sepBy sep (x : xs) = x <> foldMap (sep <>) xs + +commas :: [DuckDBSyntax] -> DuckDBSyntax +commas = sepBy (emit ", ") + +quotedIdentifier :: Text -> DuckDBSyntax +quotedIdentifier txt = emit "\"" <> DuckDBSyntax (\_ -> Builder.fromText (duckDBEscape txt)) mempty <> emit "\"" + +instance Show DuckDBSyntax where + show (DuckDBSyntax s vs) = + mconcat + [ "DuckDBSyntax (", + show (Builder.toLazyText (withPlaceholder s)), + ") ", + show vs + ] + +instance Eq DuckDBSyntax where + DuckDBSyntax sx vx == DuckDBSyntax sy vy = + -- TODO: is there a cleaner way to do this? Is it even important to have this instance? + show vx == show vy + && withPlaceholder sx == withPlaceholder sy + +instance Semigroup DuckDBSyntax where + DuckDBSyntax sx vx <> DuckDBSyntax sy vy = + DuckDBSyntax (\v -> sx v <> sy v) (vx <> vy) + +instance Monoid DuckDBSyntax where + mempty = DuckDBSyntax (const mempty) mempty + +instance Sql92DisplaySyntax DuckDBSyntax where + displaySyntax = Text.unpack . duckDBRenderSyntaxScript + +withPlaceholder :: ((SomeField -> Builder) -> Builder) -> Builder +withPlaceholder build = build (const (Builder.singleton '?')) + +-- | A best effort attempt to implement the escaping rules of DuckDB. This is +-- never used to escape data sent to the database; only for emitting scripts or +-- displaying syntax to the user. +duckDBEscape :: Text -> Text +duckDBEscape = Text.concatMap (\c -> if c == '"' then "\"\"" else Text.singleton c) + +duckDBRenderSyntaxScript :: DuckDBSyntax -> Text +duckDBRenderSyntaxScript (DuckDBSyntax s _) = + Text.Lazy.toStrict . Builder.toLazyText . s $ render + where + render :: SomeField -> Builder + render (SomeField f) = Builder.fromString (renderFieldBinding (toField f)) diff --git a/src/Database/Beam/DuckDB/Syntax/Extensions/DataSource.hs b/src/Database/Beam/DuckDB/Syntax/Extensions/DataSource.hs new file mode 100644 --- /dev/null +++ b/src/Database/Beam/DuckDB/Syntax/Extensions/DataSource.hs @@ -0,0 +1,345 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +module Database.Beam.DuckDB.Syntax.Extensions.DataSource + ( -- ** Specifying data sources as part of the database + DataSourceEntity, + DataSource, + dataSource, + modifyDataSourceFields, + + -- *** Data source configurations + csv, + csvWith, + CSVOptions (..), + defaultCSVOptions, + icebergTable, + icebergTableWith, + IcebergTableOptions (..), + defaultIcebergTableOptions, + parquet, + + -- ** Querying data from a CSV file + allFromDataSource_, + ) +where + +import Control.Monad.Free (liftF) +import Data.Kind (Type) +import Data.List.NonEmpty (NonEmpty ((:|))) +import qualified Data.List.NonEmpty as NonEmpty +import Data.Maybe (catMaybes) +import Data.Monoid (Endo (..)) +import Data.Proxy (Proxy (..)) +import Data.Text (Text) +import qualified Data.Text as Text +import Database.Beam (Beamable, DatabaseEntity, EntityModification, FieldModification, QExpr, Table, TableField, TableSettings, defTblFieldSettings, withTableModification) +import Database.Beam.DuckDB.Backend (DuckDB) +import Database.Beam.DuckDB.Syntax (DuckDBFromSyntax (..)) +import Database.Beam.DuckDB.Syntax.Builder (DuckDBSyntax, emit, emit', emitChar, quotedIdentifier, sepBy) +import Database.Beam.Query.Internal (Q (..), QF (..), tableFieldsToExpressions) +import Database.Beam.Schema.Tables + ( Columnar' (..), + DatabaseEntity (..), + EntityModification (..), + FieldRenamer (..), + GDefaultTableFieldSettings, + IsDatabaseEntity (..), + RenamableField (..), + RenamableWithRule (..), + changeBeamRep, + ) +import GHC.Generics (Generic (Rep)) + +-- | A phantom type tag for entities which are external sources of data for DuckDB +-- databases. This data can come from files (e.g. Parquet), other table formats +-- (e.g. Apache Iceberg), or even other databases. +data DataSourceEntity (table :: (Type -> Type) -> Type) + +-- | Opaque type representing the different data sources for a 'DataSourceEntity'. +-- +-- The only way to create a 'DataSource' is to use one of the helper functions +-- such as 'parquet', 'csv', 'icebergTable', and their variants. +data DataSource + = -- The constructors are not exposed so that we can extend the list + -- of supported data sources without breakage. + UndefinedDataSource FilePath -- Only used by dbEntityAuto. Assume this is a text file. + | Parquet (NonEmpty FilePath) -- no Parquet options supported yet + | CSV (NonEmpty FilePath) CSVOptions + | ApacheIceberg FilePath IcebergTableOptions + +dataSource :: + DataSource -> + EntityModification (DatabaseEntity DuckDB db) DuckDB (DataSourceEntity table) +dataSource src = + EntityModification $ Endo $ \(DatabaseEntity desc) -> + DatabaseEntity + desc + { source = src + } + +-- | Declare a CSV file(s) or glob(s) as the source of data for a database. +-- +-- Note: Be careful to sanitize any user input that is then used as +-- a data source. +csv :: + -- | File path(s) or glob(s) + NonEmpty FilePath -> + DataSource +csv = flip csvWith defaultCSVOptions + +-- | Declare a CSV file (or files) as the source of data for a database. +-- +-- See 'csv' if you want to use default options. +-- +-- Note: Be careful to sanitize any user input that is then used as +-- a data source. +csvWith :: + -- | File path(s) or glob(s) + NonEmpty FilePath -> + -- | CSV options. + CSVOptions -> + DataSource +csvWith = CSV + +-- | Options affecting the handling of CSV files. +-- +-- See 'defaultCSVOptions' for default options. +data CSVOptions = CSVOptions + { -- | Character used to initiate comments. Lines starting with a comment character + -- (optionally preceded by space characters) are completely ignored; other + -- lines containing a comment character are parsed only up to that point. + -- Default: empty. + comment :: Maybe Char, + -- | Delimiter character used to separate columns within each line, e.g., , ; \t. + -- The delimiter character can be up to 4 bytes, e.g., 🦆. Default: @","@. + delim :: Maybe Text, + -- | First line of each file contains the column names. Default: false. + header :: Maybe Bool, + -- | Ignore any parsing errors encountered. Default: false. + ignoreErrors :: Maybe Bool + } + deriving (Eq, Show) + +defaultCSVOptions :: CSVOptions +defaultCSVOptions = + CSVOptions + { comment = Nothing, + delim = Nothing, + header = Nothing, + ignoreErrors = Nothing + } + +-- | Define an Apache Iceberg table with default table options. +-- +-- See 'icebergTableWith' if you want to change the default options. +-- +-- Note: Be careful to sanitize any user input that is then used as +-- a data source. +icebergTable :: + -- | File path + FilePath -> + DataSource +icebergTable = flip icebergTableWith defaultIcebergTableOptions + +-- | Define an Apache Iceberg table with options. +-- +-- See 'icebergTable' if you want to use default options. +-- +-- Note: Be careful to sanitize any user input that is then used as +-- a data source. +icebergTableWith :: + -- | File path + FilePath -> + -- | Iceberg table options. + IcebergTableOptions -> + DataSource +icebergTableWith = ApacheIceberg + +-- | Options affecting the handling of Apache Iceberg tables. +-- +-- See 'defaultIcebergTableOptions' for default options. +data IcebergTableOptions = IcebergTableOptions + { -- | Allow scanning Iceberg tables that are moved. Default: False + allowMovedPaths :: Maybe Bool, + -- | Provides an explicit version string, hint file, or guessing. Default: @"?"@ + version :: Maybe Text + } + deriving (Eq, Show) + +defaultIcebergTableOptions :: IcebergTableOptions +defaultIcebergTableOptions = IcebergTableOptions Nothing Nothing + +-- | Declare a Parquet file(s) or glob(s) as the source of data for a database. +-- +-- Note: Be careful to sanitize any user input that is then used as +-- a data source. +parquet :: + -- | File paths or glob + NonEmpty FilePath -> + DataSource +parquet = Parquet + +instance + (Beamable tbl) => + RenamableWithRule + ( FieldRenamer + (DatabaseEntityDescriptor DuckDB (DataSourceEntity tbl)) + ) + where + renamingFields renamer = + FieldRenamer $ \tbl -> + tbl + { tableSettings = + changeBeamRep + ( \(Columnar' tblField :: Columnar' (TableField tbl) a) -> + Columnar' + ( renameField + (Proxy @(TableField tbl)) + (Proxy @a) + renamer + tblField + ) :: + Columnar' (TableField tbl) a + ) + $ tableSettings tbl + } + +instance (Beamable table) => IsDatabaseEntity DuckDB (DataSourceEntity table) where + data DatabaseEntityDescriptor DuckDB (DataSourceEntity table) + = DataSourceEntityDescriptor + { source :: !DataSource, + name :: !Text, -- Only exists so that renaming this entity works. However, it serves no purpose + tableSettings :: !(TableSettings table) + } + + type + DatabaseEntityDefaultRequirements DuckDB (DataSourceEntity table) = + ( GDefaultTableFieldSettings (Rep (TableSettings table) ()), + Generic (TableSettings table), + Table table, + Beamable table + ) + type + DatabaseEntityRegularRequirements DuckDB (DataSourceEntity table) = + (Table table, Beamable table) + + -- \| This 'dbEntityName' is useless. Changing it with 'modifyEntityName' does effectively nothing. + dbEntityName f vw = fmap (\t' -> vw {name = t'}) (f (name vw)) + + -- \| Schema doesn't apply to data sources + dbEntitySchema f vw = fmap (const vw) (f Nothing) + + -- \| By default, create a data sources as if the provided name is the path + -- to a text file. You should point the data source to the right + -- location by using 'parquet', 'csv', 'icebergTable', or any other + -- relevant function. + dbEntityAuto nm = + DataSourceEntityDescriptor + { source = UndefinedDataSource (Text.unpack nm), + name = nm, + tableSettings = defTblFieldSettings + } + +-- | This is the equivalent of 'all_', but for pulling data +-- from a 'DataSourceEntity'. +allFromDataSource_ :: + (Beamable table) => + DatabaseEntity DuckDB db (DataSourceEntity table) -> + Q DuckDB db s (table (QExpr DuckDB s)) +allFromDataSource_ (DatabaseEntity desc) = + Q $ + liftF + ( QAll + ( \_ nm -> + DuckDBFromSyntax $ + emitDataSource (source desc) + <> emit " AS " + <> quotedIdentifier nm + ) + (tableFieldsToExpressions (tableSettings desc)) + (const Nothing) + snd + ) + +quotePath :: FilePath -> DuckDBSyntax +quotePath path = mconcat [emitChar '\'', emit (Text.pack path), emitChar '\''] + +-- Render a list of filepaths +-- +-- >>> emitPaths (NonEmpty.singleton "hello/world.txt") +-- DuckDBSyntax ("'hello/world.txt'") fromList [] +-- >>> emitPaths ("hello/world.txt" :| ["foo/bar.csv"]) +-- DuckDBSyntax ("['hello/world.txt', 'foo/bar.csv']") fromList [] +emitPaths :: NonEmpty FilePath -> DuckDBSyntax +emitPaths (path :| []) = quotePath path +emitPaths paths = + emitChar '[' <> sepBy (emit ", ") (NonEmpty.toList $ fmap quotePath paths) <> emitChar ']' + +emitDataSource :: DataSource -> DuckDBSyntax +emitDataSource (UndefinedDataSource path) = + emit "read_text(" + <> quotePath path + <> emit ")" +emitDataSource (Parquet paths) = + emit "read_parquet(" + <> emitPaths paths + <> emitChar ')' +emitDataSource (CSV paths options) = + emit "read_csv(" + <> emitPaths paths + <> emitCSVOptions options + <> emitChar ')' +emitDataSource (ApacheIceberg path options) = + emit "iceberg_scan(" + <> emitPaths (NonEmpty.singleton path) + <> emitIcebergOptions options + <> emitChar ')' + +emitCSVOptions :: CSVOptions -> DuckDBSyntax +emitCSVOptions options@(CSVOptions mComment mDelim mHeader mIgnoreErrors) = + if options == defaultCSVOptions + then mempty + else + emit ", " + <> sepBy + (emit ", ") + ( catMaybes + [ fmap (\x -> emit "comment='" <> emitChar x <> emitChar '\'') mComment, + fmap (\x -> emit "delim='" <> emit x <> emitChar '\'') mDelim, + fmap (\x -> emit "header=" <> emit' x) mHeader, + fmap (\x -> emit "ignore_errors=" <> emit' x) mIgnoreErrors + ] + ) + +emitIcebergOptions :: IcebergTableOptions -> DuckDBSyntax +emitIcebergOptions options@(IcebergTableOptions mAllowMovedPaths mVersion) = + if options == defaultIcebergTableOptions + then mempty + else + emit ", " + <> sepBy + (emit ", ") + ( catMaybes + [ fmap (\x -> emit "allow_moved_paths=" <> emit' x) mAllowMovedPaths, + fmap (\x -> emit "version='" <> emit x <> emitChar '\'') mVersion + ] + ) + +-- | Construct an 'EntityModification' to rename the fields of a 'DataSourceEntity' +modifyDataSourceFields :: + (Beamable tbl) => + tbl (FieldModification (TableField tbl)) -> + EntityModification (DatabaseEntity DuckDB db) be (DataSourceEntity tbl) +modifyDataSourceFields modFields = + EntityModification + ( Endo + ( \(DatabaseEntity tbl) -> + DatabaseEntity tbl {tableSettings = withTableModification modFields (tableSettings tbl)} + ) + ) diff --git a/tests/Database/Beam/DuckDB/Test/Extensions.hs b/tests/Database/Beam/DuckDB/Test/Extensions.hs new file mode 100644 --- /dev/null +++ b/tests/Database/Beam/DuckDB/Test/Extensions.hs @@ -0,0 +1,365 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} + +module Database.Beam.DuckDB.Test.Extensions (tests) where + +import Data.Int (Int32) +import Data.List.NonEmpty (NonEmpty ((:|))) +import Data.Text (Text) +import Data.Time (Day, fromGregorian) +import Database.Beam + ( Beamable, + Columnar, + Database, + DatabaseSettings, + Generic, + Identity, + Table (..), + aggregate_, + as_, + countAll_, + dbModification, + defaultDbSettings, + guard_, + min_, + runSelectReturningList, + runSelectReturningOne, + select, + tableModification, + val_, + withDbModification, + (==.), + ) +import Database.Beam.DuckDB + ( CSVOptions (..), + DataSourceEntity, + DuckDB, + allFromDataSource_, + allowMovedPaths, + csvWith, + dataSource, + defaultCSVOptions, + defaultIcebergTableOptions, + icebergTableWith, + modifyDataSourceFields, + parquet, + runBeamDuckDB, + ) +import Database.DuckDB.Simple (Connection, withConnection) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Extensions" + [ testGroup + "Parquet" + [ testCountFromParquet, + testQueryFromParquet, + testCountFromMultipleParquetFiles, + testQueryFromMultipleParquetFiles + ], + testGroup + "Apache Iceberg" + [ testCountFromIceberg, + testParseDateFromIceberg + ], + testGroup + "CSV" + [ testCountFromCSV, + testQueryFromCSV + ] + ] + +testCountFromParquet :: TestTree +testCountFromParquet = testCase "Counting records from a Parquet file" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningList $ + select $ do + exam <- allFromDataSource_ (_dbExams testDb) + pure (_examName exam) + results @?= ["alice", "bob", "carol", "dave"] + +testQueryFromParquet :: TestTree +testQueryFromParquet = testCase "Query from a Parquet file" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningOne $ + select $ do + exam <- allFromDataSource_ (_dbExams testDb) + guard_ (_examId exam ==. 1) + pure (_examName exam) + results @?= Just "alice" + +testCountFromMultipleParquetFiles :: TestTree +testCountFromMultipleParquetFiles = testCase "Counting records from multiple Parquet file" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningList $ + select $ do + exam <- allFromDataSource_ (_dbExamsMulti testDb) + pure (_examName exam) + results @?= ["alice", "bob", "carol", "dave", "erika", "francis", "genevieve", "hugo"] + +testQueryFromMultipleParquetFiles :: TestTree +testQueryFromMultipleParquetFiles = testCase "Query from multiple Parquet files" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningOne $ + select $ do + exam <- allFromDataSource_ (_dbExamsMulti testDb) + guard_ (_examId exam ==. 5) + pure (_examName exam) + results @?= Just "erika" + +testCountFromIceberg :: TestTree +testCountFromIceberg = testCase "Counting records from an Apache Iceberg table" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningOne $ + select $ + aggregate_ (\_ -> as_ @Int32 countAll_) (allFromDataSource_ (_dbLineItems testDb)) + results @?= Just 51793 -- From DuckDB's documentation + +testParseDateFromIceberg :: TestTree +testParseDateFromIceberg = testCase "Test parsing date columns from Apache Iceberg table" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningOne $ + select $ + aggregate_ (min_ . _lineitemShipdate) (allFromDataSource_ (_dbLineItems testDb)) + + results @?= Just (Just (fromGregorian 1992 01 04)) + +testCountFromCSV :: TestTree +testCountFromCSV = testCase "Counting records from a CSV file" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningOne $ + select $ + aggregate_ (\_ -> as_ @Int32 countAll_) (allFromDataSource_ (_dbFlights testDb)) + results @?= Just 3 + +testQueryFromCSV :: TestTree +testQueryFromCSV = testCase "Query from a CSV file" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningOne $ + select $ do + flight <- allFromDataSource_ (_dbFlights testDb) + guard_ (_flightDate flight ==. val_ (fromGregorian 1988 01 01)) + pure flight + results + @?= Just + ( Flight + { _flightDate = fromGregorian 1988 01 01, + _flightUniqueCarrier = "AA", + _flightOriginCity = "New York, NY", + _flightDestCity = "Los Angeles, CA" + } + ) + +data ExamT f = Exam + { _examId :: Columnar f Int32, + _examName :: Columnar f Text, + _examScore :: Columnar f Double, + _examDate :: Columnar f Day + } + deriving (Generic) + +type Exam = ExamT Identity + +type ExamId = PrimaryKey ExamT Identity + +deriving instance Show ExamId + +deriving instance Eq ExamId + +deriving instance Ord ExamId + +deriving instance Show Exam + +deriving instance Eq Exam + +deriving instance Ord Exam + +instance Beamable ExamT + +instance Table ExamT where + data PrimaryKey ExamT f = ExamId (Columnar f Int32) + deriving (Generic) + primaryKey = ExamId . _examId + +instance Beamable (PrimaryKey ExamT) + +data LineItemT f = Lineitem + { _lineitemOrderkey :: Columnar f Int32, + _lineitemPartkey :: Columnar f Int32, + _lineitemSuppkey :: Columnar f Int32, + _lineitemLinenumber :: Columnar f Int32, + _lineitemQuantity :: Columnar f Double, + _lineitemExtendedprice :: Columnar f Double, + _lineitemDiscount :: Columnar f Double, + _lineitemTax :: Columnar f Double, + _lineitemReturnflag :: Columnar f Text, + _lineitemLinestatus :: Columnar f Text, + _lineitemShipdate :: Columnar f Day, + _lineitemCommitdate :: Columnar f Day, + _lineitemReceiptdate :: Columnar f Day, + _lineitemShipinstruct :: Columnar f Text, + _lineitemShipmode :: Columnar f Text, + _lineitemComment :: Columnar f Text + } + deriving (Generic) + +type Lineitem = LineItemT Identity + +deriving instance Show Lineitem + +deriving instance Eq Lineitem + +instance Beamable LineItemT + +instance Table LineItemT where + data PrimaryKey LineItemT f + = LineItemKey (Columnar f Int32) (Columnar f Int32) + deriving (Generic) + primaryKey li = + LineItemKey (_lineitemOrderkey li) (_lineitemLinenumber li) + +instance Beamable (PrimaryKey LineItemT) + +deriving instance Show (PrimaryKey LineItemT Identity) + +deriving instance Eq (PrimaryKey LineItemT Identity) + +data FlightT f = Flight + { _flightDate :: Columnar f Day, + _flightUniqueCarrier :: Columnar f Text, + _flightOriginCity :: Columnar f Text, + _flightDestCity :: Columnar f Text + } + deriving (Generic) + +type Flight = FlightT Identity + +deriving instance Show Flight + +deriving instance Eq Flight + +instance Beamable FlightT + +instance Table FlightT where + data PrimaryKey FlightT f + = FlightKey + (Columnar f Day) + (Columnar f Text) + (Columnar f Text) + (Columnar f Text) + deriving (Generic) + primaryKey fl = + FlightKey + (_flightDate fl) + (_flightUniqueCarrier fl) + (_flightOriginCity fl) + (_flightDestCity fl) + +instance Beamable (PrimaryKey FlightT) + +deriving instance Show (PrimaryKey FlightT Identity) + +deriving instance Eq (PrimaryKey FlightT Identity) + +data TestDB f = TestDB + { _dbExams :: f (DataSourceEntity ExamT), + _dbExamsMulti :: f (DataSourceEntity ExamT), -- set up with multiple parquet files + _dbLineItems :: f (DataSourceEntity LineItemT), + _dbFlights :: f (DataSourceEntity FlightT) + } + deriving (Generic, Database DuckDB) + +testDb :: DatabaseSettings DuckDB TestDB +testDb = + defaultDbSettings + `withDbModification` (dbModification @_ @DuckDB) + { _dbExams = + dataSource + (parquet ("tests/data/test1.parquet" :| [])) + <> modifyDataSourceFields + tableModification + { _examId = "id", + _examName = "name", + _examScore = "score", + _examDate = "exam_date" + }, + _dbExamsMulti = + dataSource + ( parquet + ("tests/data/test1.parquet" :| ["tests/data/test2.parquet"]) + ) + <> modifyDataSourceFields + tableModification + { _examId = "id", + _examName = "name", + _examScore = "score", + _examDate = "exam_date" + }, + _dbLineItems = + dataSource + ( icebergTableWith + "tests/data/lineitem_iceberg" + (defaultIcebergTableOptions {allowMovedPaths = Just True}) + ) + <> modifyDataSourceFields + tableModification + { _lineitemOrderkey = "l_orderkey", + _lineitemPartkey = "l_partkey", + _lineitemSuppkey = "l_suppkey", + _lineitemLinenumber = "l_linenumber", + _lineitemQuantity = "l_quantity", + _lineitemExtendedprice = "l_extendedprice", + _lineitemDiscount = "l_discount", + _lineitemTax = "l_tax", + _lineitemReturnflag = "l_returnflag", + _lineitemLinestatus = "l_linestatus", + _lineitemShipdate = "l_shipdate", + _lineitemCommitdate = "l_commitdate", + _lineitemReceiptdate = "l_receiptdate", + _lineitemShipinstruct = "l_shipinstruct", + _lineitemShipmode = "l_shipmode", + _lineitemComment = "l_comment" + }, + _dbFlights = + dataSource + ( csvWith + ("tests/data/flights.csv" :| []) + ( defaultCSVOptions + { header = Just True, + comment = Just '#', + delim = Just "|", + ignoreErrors = Just False + } + ) + ) + <> modifyDataSourceFields + tableModification + { _flightDate = "FlightDate", + _flightUniqueCarrier = "UniqueCarrier", + _flightOriginCity = "OriginCityName", + _flightDestCity = "DestCityName" + } + } + +withTestDb :: + (Connection -> IO a) -> + IO a +withTestDb = + withConnection ":memory:" diff --git a/tests/Database/Beam/DuckDB/Test/Query.hs b/tests/Database/Beam/DuckDB/Test/Query.hs new file mode 100644 --- /dev/null +++ b/tests/Database/Beam/DuckDB/Test/Query.hs @@ -0,0 +1,447 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeFamilies #-} + +module Database.Beam.DuckDB.Test.Query (tests) where + +import Control.Monad (void) +import Data.Int (Int32) +import Data.List (nubBy, sort, sortOn) +import Data.Text (Text) +import Database.Beam + ( Beamable, + Columnar, + Database, + DatabaseSettings, + Generic, + Identity, + SqlEq ((==.)), + SqlOrd ((>.)), + SqlValable (val_), + Table (..), + TableEntity, + all_, + dbModification, + defaultDbSettings, + guard_, + insert, + insertValues, + leftJoin_, + modifyTableFields, + related_, + runInsert, + runSelectReturningList, + select, + tableModification, + withDbModification, + (>=.), + ) +import Database.Beam.DuckDB (DuckDB, runBeamDuckDB) +import Database.DuckDB.Simple (Connection, execute_, withConnection) +import Hedgehog (Gen, annotate, evalIO, forAll, property, (===)) +import qualified Hedgehog.Gen as Gen +import qualified Hedgehog.Range as Range +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.Hedgehog (testProperty) + +tests :: TestTree +tests = + testGroup + "Query" + [ testGroup + "Selection" + [ testSelectAll, + testSelectWithFilter, + testSelectEquality + ], + testGroup + "Projection" + [ testProjectSingleColumn, + testProjectMultipleColumns, + testProjectWithExpression + ], + testGroup + "Join" + [ testInnerJoin, + testMultiInnerJoin, + testMixingJoinsWithFilters, + testLeftJoin + ] + ] + +testSelectAll :: TestTree +testSelectAll = testProperty "selecting all users should return the users initially inserted" $ property $ do + users <- forAll genUsers + results <- evalIO $ + withTestDb users [] [] $ \conn -> + runBeamDuckDB conn $ + runSelectReturningList $ + select (all_ (_dbUsers testDb)) + sortOn _userId results === sortOn _userId users + +testSelectWithFilter :: TestTree +testSelectWithFilter = testProperty "selecting users satisfying a condition works as expected" $ property $ do + users <- forAll genUsers + ageThreshold <- forAll genAge + results <- evalIO $ + withTestDb users [] [] $ \conn -> + runBeamDuckDB conn $ + runSelectReturningList $ + select $ do + u <- all_ (_dbUsers testDb) + guard_ (_userAge u >. val_ ageThreshold) + pure u + let expected = filter (\u -> _userAge u > ageThreshold) users + sortOn _userId results === sortOn _userId expected + +testSelectEquality :: TestTree +testSelectEquality = testProperty "" $ property $ do + users <- forAll genUsers + target <- forAll (Gen.element users) + results <- evalIO $ + withTestDb users [] [] $ \conn -> + runBeamDuckDB conn $ + runSelectReturningList $ + select $ do + u <- all_ (_dbUsers testDb) + guard_ (_userId u ==. val_ (_userId target)) + pure u + results === [target] + +testProjectSingleColumn :: TestTree +testProjectSingleColumn = testProperty "Single column projection works as expected" $ property $ do + users <- forAll genUsers + results <- evalIO $ + withTestDb users [] [] $ \conn -> + runBeamDuckDB conn $ + runSelectReturningList $ + select $ do + u <- all_ (_dbUsers testDb) + pure (_userName u) + sort results === sort (map _userName users) + +testProjectMultipleColumns :: TestTree +testProjectMultipleColumns = testProperty "Multiple column projection works as expected" $ property $ do + users <- forAll genUsers + results <- evalIO $ + withTestDb users [] [] $ \conn -> + runBeamDuckDB conn $ + runSelectReturningList $ + select $ do + u <- all_ (_dbUsers testDb) + pure (_userName u, _userAge u) + let expected = map (\u -> (_userName u, _userAge u)) users + sort results === sort expected + +testProjectWithExpression :: TestTree +testProjectWithExpression = testProperty "Projection using an expression works as expected" $ property $ do + products <- forAll genProducts + results <- evalIO $ + withTestDb [] products [] $ \conn -> + runBeamDuckDB conn $ + runSelectReturningList $ + select $ do + p <- all_ (_dbProducts testDb) + pure (_productId p, _productPrice p * val_ 2) + let expected = map (\p -> (_productId p, _productPrice p * 2)) products + sortOn fst results === sortOn fst expected + +testInnerJoin :: TestTree +testInnerJoin = testProperty "Inner joins work as expected" $ property $ do + users <- forAll genUsers + products <- forAll genProducts + orders <- forAll (genOrders users products) + results <- evalIO $ + withTestDb users products orders $ \conn -> + runBeamDuckDB conn $ + runSelectReturningList $ + select $ do + o <- all_ (_dbOrders testDb) + u <- related_ (_dbUsers testDb) (_orderUserId o) + pure (_userName u, _orderQuantity o) + + let expected = do + o <- orders + let UserId uid = _orderUserId o + u <- filter (\u -> _userId u == uid) users + pure (_userName u, _orderQuantity o) + sort results === sort expected + +testMultiInnerJoin :: TestTree +testMultiInnerJoin = testProperty "Multi-inner joins work as expected" $ property $ do + users <- forAll genUsers + products <- forAll genProducts + orders <- forAll (genOrders users products) + results <- evalIO $ + withTestDb users products orders $ \conn -> + runBeamDuckDB conn $ + runSelectReturningList $ + select $ do + o <- all_ (_dbOrders testDb) + u <- related_ (_dbUsers testDb) (_orderUserId o) + p <- related_ (_dbProducts testDb) (_orderProductId o) + pure (_userName u, _productName p, _orderQuantity o) + let expected = do + o <- orders + let UserId uid = _orderUserId o + ProductId pid = _orderProductId o + u <- filter (\u -> _userId u == uid) users + p <- filter (\p -> _productId p == pid) products + pure (_userName u, _productName p, _orderQuantity o) + sort results === sort expected + +testMixingJoinsWithFilters :: TestTree +testMixingJoinsWithFilters = testProperty "Mixing joins and filters works as expected" $ property $ do + users <- forAll genUsers + products <- forAll genProducts + orders <- forAll (genOrders users products) + minQty <- forAll (Gen.int32 (Range.linear 1 50)) + results <- evalIO $ + withTestDb users products orders $ \conn -> + runBeamDuckDB conn $ + runSelectReturningList $ + select $ do + o <- all_ (_dbOrders testDb) + guard_ (_orderQuantity o >=. val_ minQty) + u <- related_ (_dbUsers testDb) (_orderUserId o) + pure (_userName u, _orderQuantity o) + let expected = do + o <- filter (\o -> _orderQuantity o >= minQty) orders + let UserId uid = _orderUserId o + u <- filter (\u -> _userId u == uid) users + pure (_userName u, _orderQuantity o) + sort results === sort expected + +testLeftJoin :: TestTree +testLeftJoin = testProperty "Left joins work as expected" $ property $ do + users <- forAll genUsers + products <- forAll genProducts + -- Generate orders for only a subset of users + let halfUsers = take (length users `div` 2) users + orders <- forAll (genOrders halfUsers products) + results <- evalIO $ + withTestDb users products orders $ \conn -> + runBeamDuckDB conn $ + runSelectReturningList $ + select $ do + u <- all_ (_dbUsers testDb) + o <- + leftJoin_ + (all_ (_dbOrders testDb)) + (\o -> _orderUserId o ==. primaryKey u) + pure (_userId u, _orderQuantity o) + + let resultUserIds = nubBy (\a b -> fst a == fst b) (sortOn fst results) + length resultUserIds === length users + + let usersWithoutOrders = filter (\u -> _userId u `notElem` map _userId halfUsers) users + nothingRows = filter (\(uid, _) -> uid `elem` map _userId usersWithoutOrders) results + annotate "Users without orders should have Nothing quantity" + mapM_ (\(_, mq) -> mq === Nothing) nothingRows + +data UserT f = User + { _userId :: Columnar f Int32, + _userName :: Columnar f Text, + _userAge :: Columnar f Int32 + } + deriving (Generic) + +type User = UserT Identity + +type UserId = PrimaryKey UserT Identity + +deriving instance Show UserId + +deriving instance Eq UserId + +deriving instance Ord UserId + +deriving instance Show User + +deriving instance Eq User + +deriving instance Ord User + +instance Beamable UserT + +instance Table UserT where + data PrimaryKey UserT f = UserId (Columnar f Int32) + deriving (Generic) + primaryKey = UserId . _userId + +instance Beamable (PrimaryKey UserT) + +data ProductT f = Product + { _productId :: Columnar f Int32, + _productName :: Columnar f Text, + _productPrice :: Columnar f Int32 -- cents, to avoid floating point + } + deriving (Generic) + +type Product = ProductT Identity + +type ProductId = PrimaryKey ProductT Identity + +deriving instance Show ProductId + +deriving instance Eq ProductId + +deriving instance Ord ProductId + +deriving instance Show Product + +deriving instance Eq Product + +deriving instance Ord Product + +instance Beamable ProductT + +instance Table ProductT where + data PrimaryKey ProductT f = ProductId (Columnar f Int32) + deriving (Generic) + primaryKey = ProductId . _productId + +instance Beamable (PrimaryKey ProductT) + +data OrderT f = Order + { _orderId :: Columnar f Int32, + _orderUserId :: PrimaryKey UserT f, + _orderProductId :: PrimaryKey ProductT f, + _orderQuantity :: Columnar f Int32 + } + deriving (Generic) + +type Order = OrderT Identity + +type OrderId = PrimaryKey OrderT Identity + +deriving instance Show OrderId + +deriving instance Eq OrderId + +deriving instance Ord OrderId + +deriving instance Show Order + +deriving instance Eq Order + +deriving instance Ord Order + +instance Beamable OrderT + +instance Table OrderT where + data PrimaryKey OrderT f = OrderId (Columnar f Int32) + deriving (Generic) + primaryKey = OrderId . _orderId + +instance Beamable (PrimaryKey OrderT) + +data TestDB f = TestDB + { _dbUsers :: f (TableEntity UserT), + _dbProducts :: f (TableEntity ProductT), + _dbOrders :: f (TableEntity OrderT) + } + deriving (Generic, Database be) + +testDb :: DatabaseSettings DuckDB TestDB +testDb = + defaultDbSettings + `withDbModification` dbModification + { _dbUsers = + modifyTableFields + tableModification + { _userId = "id", + _userName = "name", + _userAge = "age" + }, + _dbProducts = + modifyTableFields + tableModification + { _productId = "id", + _productName = "name", + _productPrice = "price" + }, + _dbOrders = + modifyTableFields + tableModification + { _orderId = "id", + _orderUserId = UserId "user_id", + _orderProductId = ProductId "product_id", + _orderQuantity = "quantity" + } + } + +genName :: Gen Text +genName = Gen.text (Range.linear 1 50) Gen.alphaNum + +genAge :: Gen Int32 +genAge = Gen.int32 (Range.linear 1 120) + +genPrice :: Gen Int32 +genPrice = Gen.int32 (Range.linear 100 100000) + +genQuantity :: Gen Int32 +genQuantity = Gen.int32 (Range.linear 1 100) + +genUsers :: Gen [User] +genUsers = do + n <- Gen.int (Range.linear 3 20) + traverse (\i -> User (fromIntegral i) <$> genName <*> genAge) [1 .. n] + +genProducts :: Gen [Product] +genProducts = do + n <- Gen.int (Range.linear 2 10) + traverse (\i -> Product (fromIntegral i) <$> genName <*> genPrice) [1 .. n] + +genOrders :: [User] -> [Product] -> Gen [Order] +genOrders users products = do + n <- Gen.int (Range.linear 1 (length users * length products)) + traverse + ( \i -> do + uid <- Gen.element (map _userId users) + pid <- Gen.element (map _productId products) + Order (fromIntegral i) (UserId uid) (ProductId pid) <$> genQuantity + ) + [1 .. n] + +createTables :: Connection -> IO () +createTables conn = do + void $ + execute_ + conn + "CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT NOT NULL, age INTEGER NOT NULL)" + void $ + execute_ + conn + "CREATE TABLE products (id INTEGER PRIMARY KEY, name TEXT NOT NULL, price INTEGER NOT NULL)" + void $ + execute_ + conn + "CREATE TABLE orders (\ + \ id INTEGER PRIMARY KEY, \ + \ user_id INTEGER NOT NULL REFERENCES users(id), \ + \ product_id INTEGER NOT NULL REFERENCES products(id), \ + \ quantity INTEGER NOT NULL)" + +seedData :: Connection -> [User] -> [Product] -> [Order] -> IO () +seedData conn users products orders = runBeamDuckDB conn $ do + runInsert $ insert (_dbUsers testDb) $ insertValues users + runInsert $ insert (_dbProducts testDb) $ insertValues products + runInsert $ insert (_dbOrders testDb) $ insertValues orders + +-- Run a test with a fresh in-memory DB populated with the given data +withTestDb :: + [User] -> + [Product] -> + [Order] -> + (Connection -> IO a) -> + IO a +withTestDb users products orders action = + withConnection ":memory:" $ \conn -> do + createTables conn + seedData conn users products orders + action conn diff --git a/tests/Main.hs b/tests/Main.hs new file mode 100644 --- /dev/null +++ b/tests/Main.hs @@ -0,0 +1,14 @@ +module Main (main) where + +import qualified Database.Beam.DuckDB.Test.Extensions (tests) +import qualified Database.Beam.DuckDB.Test.Query (tests) +import Test.Tasty (defaultMain, testGroup) + +main :: IO () +main = + defaultMain $ + testGroup + "beam-duckdb tests" + [ Database.Beam.DuckDB.Test.Query.tests + , Database.Beam.DuckDB.Test.Extensions.tests + ] diff --git a/tests/data/flights.csv b/tests/data/flights.csv new file mode 100644 --- /dev/null +++ b/tests/data/flights.csv @@ -0,0 +1,4 @@ +FlightDate|UniqueCarrier|OriginCityName|DestCityName +1988-01-01|AA|New York, NY|Los Angeles, CA +1988-01-02|AA|New York, NY|Houston, TX +1988-01-03|AA|New York, NY|Seattle, WA diff --git a/tests/data/gen_parquet.py b/tests/data/gen_parquet.py new file mode 100644 --- /dev/null +++ b/tests/data/gen_parquet.py @@ -0,0 +1,40 @@ +#!/usr/bin/env -S uv run +# /// script +# requires-python = ">=3.12" +# dependencies = ["pyarrow"] +# /// + +import pyarrow as pa +import pyarrow.parquet as pq +from datetime import date +from pathlib import Path + +table1 = pa.table({ + "id": pa.array([1, 2, 3, 4], type=pa.int32()), + "name": pa.array(["alice", "bob", "carol", "dave"], type=pa.utf8()), + "score": pa.array([97.5, 82.3, 91.0, 76.8], type=pa.float64()), + "passed": pa.array([True, True, True, False], type=pa.bool_()), + "exam_date": pa.array( + [date(2025, 1, 10), date(2025, 1, 11), date(2025, 1, 10), date(2025, 1, 12)], + type=pa.date32(), + ), +}) + +table2 = pa.table({ + "id": pa.array([5, 6, 7, 8], type=pa.int32()), + "name": pa.array(["erika", "francis", "genevieve", "hugo"], type=pa.utf8()), + "score": pa.array([96.5, 72.3, 81.0, 76.8], type=pa.float64()), + "passed": pa.array([True, True, True, False], type=pa.bool_()), + "exam_date": pa.array( + [date(2025, 1, 10), date(2025, 1, 11), date(2025, 1, 10), date(2025, 1, 12)], + type=pa.date32(), + ), +}) + +out1 = Path(__file__).parent / "test1.parquet" +pq.write_table(table1, out1) +print(f"Wrote {out1} ({out1.stat().st_size} bytes, {table1.num_rows} rows)") + +out2 = Path(__file__).parent / "test2.parquet" +pq.write_table(table2, out2) +print(f"Wrote {out2} ({out2.stat().st_size} bytes, {table2.num_rows} rows)") diff --git a/tests/data/lineitem_iceberg/data/00000-411-0792dcfe-4e25-4ca3-8ada-175286069a47-00001.parquet b/tests/data/lineitem_iceberg/data/00000-411-0792dcfe-4e25-4ca3-8ada-175286069a47-00001.parquet new file mode 100644 Binary files /dev/null and b/tests/data/lineitem_iceberg/data/00000-411-0792dcfe-4e25-4ca3-8ada-175286069a47-00001.parquet differ diff --git a/tests/data/lineitem_iceberg/data/00041-414-f3c73457-bbd6-4b92-9c15-17b241171b16-00001.parquet b/tests/data/lineitem_iceberg/data/00041-414-f3c73457-bbd6-4b92-9c15-17b241171b16-00001.parquet new file mode 100644 Binary files /dev/null and b/tests/data/lineitem_iceberg/data/00041-414-f3c73457-bbd6-4b92-9c15-17b241171b16-00001.parquet differ diff --git a/tests/data/lineitem_iceberg/metadata/10eaca8a-1e1c-421e-ad6d-b232e5ee23d3-m0.avro b/tests/data/lineitem_iceberg/metadata/10eaca8a-1e1c-421e-ad6d-b232e5ee23d3-m0.avro new file mode 100644 Binary files /dev/null and b/tests/data/lineitem_iceberg/metadata/10eaca8a-1e1c-421e-ad6d-b232e5ee23d3-m0.avro differ diff --git a/tests/data/lineitem_iceberg/metadata/10eaca8a-1e1c-421e-ad6d-b232e5ee23d3-m1.avro b/tests/data/lineitem_iceberg/metadata/10eaca8a-1e1c-421e-ad6d-b232e5ee23d3-m1.avro new file mode 100644 Binary files /dev/null and b/tests/data/lineitem_iceberg/metadata/10eaca8a-1e1c-421e-ad6d-b232e5ee23d3-m1.avro differ diff --git a/tests/data/lineitem_iceberg/metadata/cf3d0be5-cf70-453d-ad8f-48fdc412e608-m0.avro b/tests/data/lineitem_iceberg/metadata/cf3d0be5-cf70-453d-ad8f-48fdc412e608-m0.avro new file mode 100644 Binary files /dev/null and b/tests/data/lineitem_iceberg/metadata/cf3d0be5-cf70-453d-ad8f-48fdc412e608-m0.avro differ diff --git a/tests/data/lineitem_iceberg/metadata/snap-3776207205136740581-1-cf3d0be5-cf70-453d-ad8f-48fdc412e608.avro b/tests/data/lineitem_iceberg/metadata/snap-3776207205136740581-1-cf3d0be5-cf70-453d-ad8f-48fdc412e608.avro new file mode 100644 Binary files /dev/null and b/tests/data/lineitem_iceberg/metadata/snap-3776207205136740581-1-cf3d0be5-cf70-453d-ad8f-48fdc412e608.avro differ diff --git a/tests/data/lineitem_iceberg/metadata/snap-7635660646343998149-1-10eaca8a-1e1c-421e-ad6d-b232e5ee23d3.avro b/tests/data/lineitem_iceberg/metadata/snap-7635660646343998149-1-10eaca8a-1e1c-421e-ad6d-b232e5ee23d3.avro new file mode 100644 Binary files /dev/null and b/tests/data/lineitem_iceberg/metadata/snap-7635660646343998149-1-10eaca8a-1e1c-421e-ad6d-b232e5ee23d3.avro differ diff --git a/tests/data/lineitem_iceberg/metadata/v1.metadata.json b/tests/data/lineitem_iceberg/metadata/v1.metadata.json new file mode 100644 --- /dev/null +++ b/tests/data/lineitem_iceberg/metadata/v1.metadata.json @@ -0,0 +1,142 @@ +{ + "format-version" : 2, + "table-uuid" : "a319422b-6f8c-44d0-90ba-96242d9a1d7b", + "location" : "./lineitem_iceberg", + "last-sequence-number" : 1, + "last-updated-ms" : 1676473674504, + "last-column-id" : 16, + "current-schema-id" : 0, + "schemas" : [ { + "type" : "struct", + "schema-id" : 0, + "fields" : [ { + "id" : 1, + "name" : "l_orderkey", + "required" : false, + "type" : "int" + }, { + "id" : 2, + "name" : "l_partkey", + "required" : false, + "type" : "int" + }, { + "id" : 3, + "name" : "l_suppkey", + "required" : false, + "type" : "int" + }, { + "id" : 4, + "name" : "l_linenumber", + "required" : false, + "type" : "int" + }, { + "id" : 5, + "name" : "l_quantity", + "required" : false, + "type" : "int" + }, { + "id" : 6, + "name" : "l_extendedprice", + "required" : false, + "type" : "decimal(15, 2)" + }, { + "id" : 7, + "name" : "l_discount", + "required" : false, + "type" : "decimal(15, 2)" + }, { + "id" : 8, + "name" : "l_tax", + "required" : false, + "type" : "decimal(15, 2)" + }, { + "id" : 9, + "name" : "l_returnflag", + "required" : false, + "type" : "string" + }, { + "id" : 10, + "name" : "l_linestatus", + "required" : false, + "type" : "string" + }, { + "id" : 11, + "name" : "l_shipdate", + "required" : false, + "type" : "date" + }, { + "id" : 12, + "name" : "l_commitdate", + "required" : false, + "type" : "date" + }, { + "id" : 13, + "name" : "l_receiptdate", + "required" : false, + "type" : "date" + }, { + "id" : 14, + "name" : "l_shipinstruct", + "required" : false, + "type" : "string" + }, { + "id" : 15, + "name" : "l_shipmode", + "required" : false, + "type" : "string" + }, { + "id" : 16, + "name" : "l_comment", + "required" : false, + "type" : "string" + } ] + } ], + "default-spec-id" : 0, + "partition-specs" : [ { + "spec-id" : 0, + "fields" : [ ] + } ], + "last-partition-id" : 999, + "default-sort-order-id" : 0, + "sort-orders" : [ { + "order-id" : 0, + "fields" : [ ] + } ], + "properties" : { + "owner" : "root", + "write.update.mode" : "merge-on-read" + }, + "current-snapshot-id" : 3776207205136740581, + "refs" : { + "main" : { + "snapshot-id" : 3776207205136740581, + "type" : "branch" + } + }, + "snapshots" : [ { + "sequence-number" : 1, + "snapshot-id" : 3776207205136740581, + "timestamp-ms" : 1676473674504, + "summary" : { + "operation" : "append", + "spark.app.id" : "local-1676472783435", + "added-data-files" : "1", + "added-records" : "60175", + "added-files-size" : "1390176", + "changed-partition-count" : "1", + "total-records" : "60175", + "total-files-size" : "1390176", + "total-data-files" : "1", + "total-delete-files" : "0", + "total-position-deletes" : "0", + "total-equality-deletes" : "0" + }, + "manifest-list" : "lineitem_iceberg/metadata/snap-3776207205136740581-1-cf3d0be5-cf70-453d-ad8f-48fdc412e608.avro", + "schema-id" : 0 + } ], + "snapshot-log" : [ { + "timestamp-ms" : 1676473674504, + "snapshot-id" : 3776207205136740581 + } ], + "metadata-log" : [ ] +} diff --git a/tests/data/lineitem_iceberg/metadata/v2.metadata.json b/tests/data/lineitem_iceberg/metadata/v2.metadata.json new file mode 100644 --- /dev/null +++ b/tests/data/lineitem_iceberg/metadata/v2.metadata.json @@ -0,0 +1,172 @@ +{ + "format-version" : 2, + "table-uuid" : "a319422b-6f8c-44d0-90ba-96242d9a1d7b", + "location" : "./lineitem_iceberg", + "last-sequence-number" : 2, + "last-updated-ms" : 1676473694730, + "last-column-id" : 16, + "current-schema-id" : 0, + "schemas" : [ { + "type" : "struct", + "schema-id" : 0, + "fields" : [ { + "id" : 1, + "name" : "l_orderkey", + "required" : false, + "type" : "int" + }, { + "id" : 2, + "name" : "l_partkey", + "required" : false, + "type" : "int" + }, { + "id" : 3, + "name" : "l_suppkey", + "required" : false, + "type" : "int" + }, { + "id" : 4, + "name" : "l_linenumber", + "required" : false, + "type" : "int" + }, { + "id" : 5, + "name" : "l_quantity", + "required" : false, + "type" : "int" + }, { + "id" : 6, + "name" : "l_extendedprice", + "required" : false, + "type" : "decimal(15, 2)" + }, { + "id" : 7, + "name" : "l_discount", + "required" : false, + "type" : "decimal(15, 2)" + }, { + "id" : 8, + "name" : "l_tax", + "required" : false, + "type" : "decimal(15, 2)" + }, { + "id" : 9, + "name" : "l_returnflag", + "required" : false, + "type" : "string" + }, { + "id" : 10, + "name" : "l_linestatus", + "required" : false, + "type" : "string" + }, { + "id" : 11, + "name" : "l_shipdate", + "required" : false, + "type" : "date" + }, { + "id" : 12, + "name" : "l_commitdate", + "required" : false, + "type" : "date" + }, { + "id" : 13, + "name" : "l_receiptdate", + "required" : false, + "type" : "date" + }, { + "id" : 14, + "name" : "l_shipinstruct", + "required" : false, + "type" : "string" + }, { + "id" : 15, + "name" : "l_shipmode", + "required" : false, + "type" : "string" + }, { + "id" : 16, + "name" : "l_comment", + "required" : false, + "type" : "string" + } ] + } ], + "default-spec-id" : 0, + "partition-specs" : [ { + "spec-id" : 0, + "fields" : [ ] + } ], + "last-partition-id" : 999, + "default-sort-order-id" : 0, + "sort-orders" : [ { + "order-id" : 0, + "fields" : [ ] + } ], + "properties" : { + "owner" : "root", + "write.update.mode" : "merge-on-read" + }, + "current-snapshot-id" : 7635660646343998149, + "refs" : { + "main" : { + "snapshot-id" : 7635660646343998149, + "type" : "branch" + } + }, + "snapshots" : [ { + "sequence-number" : 1, + "snapshot-id" : 3776207205136740581, + "timestamp-ms" : 1676473674504, + "summary" : { + "operation" : "append", + "spark.app.id" : "local-1676472783435", + "added-data-files" : "1", + "added-records" : "60175", + "added-files-size" : "1390176", + "changed-partition-count" : "1", + "total-records" : "60175", + "total-files-size" : "1390176", + "total-data-files" : "1", + "total-delete-files" : "0", + "total-position-deletes" : "0", + "total-equality-deletes" : "0" + }, + "manifest-list" : "lineitem_iceberg/metadata/snap-3776207205136740581-1-cf3d0be5-cf70-453d-ad8f-48fdc412e608.avro", + "schema-id" : 0 + }, { + "sequence-number" : 2, + "snapshot-id" : 7635660646343998149, + "parent-snapshot-id" : 3776207205136740581, + "timestamp-ms" : 1676473694730, + "summary" : { + "operation" : "overwrite", + "spark.app.id" : "local-1676472783435", + "added-data-files" : "1", + "deleted-data-files" : "1", + "added-records" : "51793", + "deleted-records" : "60175", + "added-files-size" : "1208539", + "removed-files-size" : "1390176", + "changed-partition-count" : "1", + "total-records" : "51793", + "total-files-size" : "1208539", + "total-data-files" : "1", + "total-delete-files" : "0", + "total-position-deletes" : "0", + "total-equality-deletes" : "0" + }, + "manifest-list" : "lineitem_iceberg/metadata/snap-7635660646343998149-1-10eaca8a-1e1c-421e-ad6d-b232e5ee23d3.avro", + "schema-id" : 0 + } ], + "snapshot-log" : [ { + "timestamp-ms" : 1676473674504, + "snapshot-id" : 3776207205136740581 + }, { + "timestamp-ms" : 1676473694730, + "snapshot-id" : 7635660646343998149 + } ], + "metadata-log" : [ { + "timestamp-ms" : 1676473674504, + "metadata-file" : "lineitem_iceberg/metadata/v1.metadata.json" + } ] +} diff --git a/tests/data/lineitem_iceberg/metadata/version-hint.text b/tests/data/lineitem_iceberg/metadata/version-hint.text new file mode 100644 --- /dev/null +++ b/tests/data/lineitem_iceberg/metadata/version-hint.text @@ -0,0 +1,1 @@ +2 diff --git a/tests/data/test.parquet b/tests/data/test.parquet new file mode 100644 Binary files /dev/null and b/tests/data/test.parquet differ diff --git a/tests/data/test1.parquet b/tests/data/test1.parquet new file mode 100644 Binary files /dev/null and b/tests/data/test1.parquet differ diff --git a/tests/data/test2.parquet b/tests/data/test2.parquet new file mode 100644 Binary files /dev/null and b/tests/data/test2.parquet differ