Add take and drop operations to Rel8.Array

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.

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))

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))

rel8-internal/src/Rel8/Internal/Table/List.hs

@@ -19,26 +19,29 @@ module Rel8.Internal.Table.List
   , index
   , last
   , length
+  , take
+  , drop
   )
 where
 
 -- base
 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

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

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.:-)

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)