diff --git a/changelog.d/20260818_000000_shane.obrien_take_drop.md b/changelog.d/20260818_000000_shane.obrien_take_drop.md new file mode 100644 index 00000000..42d14a1a --- /dev/null +++ b/changelog.d/20260818_000000_shane.obrien_take_drop.md @@ -0,0 +1,7 @@ +### Added + +- Added `take`/`drop` (and `take1`/`drop1`, `takeExpr`/`dropExpr`, `take1Expr`/`drop1Expr`) to `Rel8.Array`. + +### Fixed + +- `last` and `last1` used `array_lower` instead of `array_upper`, and so returned the first element of an array rather than the last. diff --git a/rel8-internal/src/Rel8/Internal/Expr/List.hs b/rel8-internal/src/Rel8/Internal/Expr/List.hs index 0b2fd319..1a683a93 100644 --- a/rel8-internal/src/Rel8/Internal/Expr/List.hs +++ b/rel8-internal/src/Rel8/Internal/Expr/List.hs @@ -9,6 +9,8 @@ module Rel8.Internal.Expr.List ( sindexExpr, slastExpr, lengthExpr, + takeExpr, + dropExpr, ) where -- base @@ -50,3 +52,11 @@ slastExpr info = mapPrimExpr (Prim.last info) lengthExpr :: Expr [a] -> Expr Int32 lengthExpr = mapPrimExpr (Prim.length) + + +takeExpr :: Expr Int32 -> Expr [a] -> Expr [a] +takeExpr n = mapPrimExpr (Prim.take (toPrimExpr n)) + + +dropExpr :: Expr Int32 -> Expr [a] -> Expr [a] +dropExpr n = mapPrimExpr (Prim.drop (toPrimExpr n)) diff --git a/rel8-internal/src/Rel8/Internal/Expr/NonEmpty.hs b/rel8-internal/src/Rel8/Internal/Expr/NonEmpty.hs index bb805f39..9059889f 100644 --- a/rel8-internal/src/Rel8/Internal/Expr/NonEmpty.hs +++ b/rel8-internal/src/Rel8/Internal/Expr/NonEmpty.hs @@ -9,6 +9,8 @@ module Rel8.Internal.Expr.NonEmpty ( sindex1Expr, slast1Expr, length1Expr, + take1Expr, + drop1Expr, ) where -- base @@ -51,3 +53,11 @@ slast1Expr info = mapPrimExpr (Prim.last info) length1Expr :: Expr (NonEmpty a) -> Expr Int32 length1Expr = mapPrimExpr (Prim.length) + + +take1Expr :: Expr Int32 -> Expr (NonEmpty a) -> Expr [a] +take1Expr n = mapPrimExpr (Prim.take (toPrimExpr n)) + + +drop1Expr :: Expr Int32 -> Expr (NonEmpty a) -> Expr [a] +drop1Expr n = mapPrimExpr (Prim.drop (toPrimExpr n)) diff --git a/rel8-internal/src/Rel8/Internal/Table/List.hs b/rel8-internal/src/Rel8/Internal/Table/List.hs index bc97654d..df16b7a3 100644 --- a/rel8-internal/src/Rel8/Internal/Table/List.hs +++ b/rel8-internal/src/Rel8/Internal/Table/List.hs @@ -19,6 +19,8 @@ module Rel8.Internal.Table.List , index , last , length + , take + , drop ) where @@ -26,19 +28,20 @@ where import Data.Functor.Identity (Identity (Identity)) import Data.Int (Int32) import Data.Kind ( Type ) -import Prelude hiding (head, last, length) +import Prelude hiding (drop, head, last, length, take) -- rel8 import Rel8.Internal.Expr ( Expr ) import Rel8.Internal.Expr.Array ( sappend, sempty, slistOf ) -import Rel8.Internal.Expr.List (lengthExpr, sheadExpr, sindexExpr, slastExpr) +import Rel8.Internal.Expr.List + ( dropExpr, lengthExpr, sheadExpr, sindexExpr, slastExpr, takeExpr ) import Rel8.Internal.Schema.Dict ( Dict( Dict ) ) import Rel8.Internal.Schema.HTable.List ( HListTable ) import Rel8.Internal.Schema.HTable.Vectorize ( hvectorize, hunvectorize , hnullify , happend, hempty - , hproject, hcolumn + , hproject, htraverseVectorP, hcolumn , First (..) ) import qualified Rel8.Internal.Schema.Kind as K @@ -191,3 +194,19 @@ length = getFirst . hunvectorize (\_ -> First . lengthExpr) . toColumns + + +-- | @'take' n as@ returns the first @n@ elements of @as@. +take :: Table Expr a => Expr Int32 -> ListTable Expr a -> ListTable Expr a +take n = + fromColumns . + htraverseVectorP (\_ -> takeExpr n) . + toColumns + + +-- | @'drop' n as@ returns the suffix of @as@ after the first @n@ elements. +drop :: Table Expr a => Expr Int32 -> ListTable Expr a -> ListTable Expr a +drop n = + fromColumns . + htraverseVectorP (\_ -> dropExpr n) . + toColumns diff --git a/rel8-internal/src/Rel8/Internal/Table/NonEmpty.hs b/rel8-internal/src/Rel8/Internal/Table/NonEmpty.hs index f6c8e5d7..83059dd1 100644 --- a/rel8-internal/src/Rel8/Internal/Table/NonEmpty.hs +++ b/rel8-internal/src/Rel8/Internal/Table/NonEmpty.hs @@ -19,6 +19,8 @@ module Rel8.Internal.Table.NonEmpty , index1 , last1 , length1 + , take1 + , drop1 ) where @@ -32,14 +34,15 @@ import Prelude hiding ( id ) -- rel8 import Rel8.Internal.Expr ( Expr ) import Rel8.Internal.Expr.Array ( sappend1, snonEmptyOf ) -import Rel8.Internal.Expr.NonEmpty (length1Expr, shead1Expr, sindex1Expr, slast1Expr) +import Rel8.Internal.Expr.NonEmpty + ( drop1Expr, length1Expr, shead1Expr, sindex1Expr, slast1Expr, take1Expr ) import Rel8.Internal.Schema.Dict ( Dict( Dict ) ) import Rel8.Internal.Schema.HTable.NonEmpty ( HNonEmptyTable ) import Rel8.Internal.Schema.HTable.Vectorize ( hvectorize, hunvectorize , hnullify , happend - , hproject, hcolumn + , hproject, htraverseVectorP, hcolumn , First (..) ) import qualified Rel8.Internal.Schema.Kind as K @@ -54,6 +57,7 @@ import Rel8.Internal.Table ) import Rel8.Internal.Table.Alternative ( AltTable, (<|>:) ) import Rel8.Internal.Table.Eq ( EqTable, eqTable ) +import Rel8.Internal.Table.List (ListTable) import Rel8.Internal.Table.Null (NullTable) import Rel8.Internal.Table.Ord ( OrdTable, ordTable ) import Rel8.Internal.Table.Projection @@ -187,3 +191,19 @@ length1 = getFirst . hunvectorize (\_ -> First . length1Expr) . toColumns + + +-- | @'take1' n as@ returns the first @n@ elements of @as@. +take1 :: Table Expr a => Expr Int32 -> NonEmptyTable Expr a -> ListTable Expr a +take1 n = + fromColumns . + htraverseVectorP (\_ -> take1Expr n) . + toColumns + + +-- | @'drop1' n as@ returns the suffix of @as@ after the first @n@ elements. +drop1 :: Table Expr a => Expr Int32 -> NonEmptyTable Expr a -> ListTable Expr a +drop1 n = + fromColumns . + htraverseVectorP (\_ -> drop1Expr n) . + toColumns diff --git a/rel8-internal/src/Rel8/Internal/Type/Array.hs b/rel8-internal/src/Rel8/Internal/Type/Array.hs index 9edc0c99..2ad97b2b 100644 --- a/rel8-internal/src/Rel8/Internal/Type/Array.hs +++ b/rel8-internal/src/Rel8/Internal/Type/Array.hs @@ -11,7 +11,7 @@ module Rel8.Internal.Type.Array , arrayTypeName , listTypeInformation , nonEmptyTypeInformation - , head, index, last, length + , head, index, last, length, take, drop ) where @@ -24,7 +24,7 @@ import Data.Bifunctor (first) import Data.Foldable (fold, toList) import Data.Functor.Contravariant ((>$<)) import Data.List.NonEmpty (NonEmpty, nonEmpty) -import Prelude hiding (head, last, length, null, repeat, zipWith) +import Prelude hiding (drop, head, last, length, null, repeat, take, zipWith) -- bytestring import Data.ByteString (ByteString) @@ -236,12 +236,32 @@ index :: TypeInformation a -> Opaleye.PrimExpr -> Opaleye.PrimExpr -> Opaleye.Pr index info i a = extractArrayElement info $ subscript (plus (lower a) i) a +take :: Opaleye.PrimExpr -> Opaleye.PrimExpr -> Opaleye.PrimExpr +take n a = slice (lowerBound a) (plus (lowerBound a) (minus n one)) a + + +drop :: Opaleye.PrimExpr -> Opaleye.PrimExpr -> Opaleye.PrimExpr +drop n a = slice (plus (lowerBound a) n) (upperBound a) a + + +slice :: Opaleye.PrimExpr -> Opaleye.PrimExpr -> Opaleye.PrimExpr -> Opaleye.PrimExpr +slice i j a = Opaleye.ArraySlice a i j + + lower :: Opaleye.PrimExpr -> Opaleye.PrimExpr lower a = Opaleye.FunExpr "array_lower" [a, one] +lowerBound :: Opaleye.PrimExpr -> Opaleye.PrimExpr +lowerBound a = Opaleye.FunExpr "coalesce" [lower a, one] + + upper :: Opaleye.PrimExpr -> Opaleye.PrimExpr -upper a = Opaleye.FunExpr "array_lower" [a, one] +upper a = Opaleye.FunExpr "array_upper" [a, one] + + +upperBound :: Opaleye.PrimExpr -> Opaleye.PrimExpr +upperBound a = Opaleye.FunExpr "coalesce" [upper a, zero] length :: Opaleye.PrimExpr -> Opaleye.PrimExpr @@ -258,3 +278,7 @@ zero = Opaleye.ConstExpr (Opaleye.IntegerLit 0) plus :: Opaleye.PrimExpr -> Opaleye.PrimExpr -> Opaleye.PrimExpr plus = Opaleye.BinExpr (Opaleye.:+) + + +minus :: Opaleye.PrimExpr -> Opaleye.PrimExpr -> Opaleye.PrimExpr +minus = Opaleye.BinExpr (Opaleye.:-) diff --git a/rel8/src/Rel8/Array.hs b/rel8/src/Rel8/Array.hs index af69cb40..de3cae54 100644 --- a/rel8/src/Rel8/Array.hs +++ b/rel8/src/Rel8/Array.hs @@ -10,6 +10,8 @@ module Rel8.Array , index, indexExpr , last, lastExpr , length, lengthExpr + , take, takeExpr + , drop, dropExpr , elem -- ** @NonEmptyTable@ @@ -18,6 +20,8 @@ module Rel8.Array , index1, index1Expr , last1, last1Expr , length1, length1Expr + , take1, take1Expr + , drop1, drop1Expr , elem1 -- ** Unsafe @@ -28,7 +32,7 @@ where -- base import Data.List.NonEmpty (NonEmpty) -import Prelude hiding (elem, head, last, length) +import Prelude hiding (drop, elem, head, last, length, take) -- rel8 import Rel8.Internal.Expr (Expr)