dataframe-arrow-bridge 1.0.0.0 → 1.1.0.0
raw patch · 3 files changed
+26/−63 lines, 3 filesdep ~dataframe-coredep ~dataframe-expr-serializerdep ~dataframe-operations
Dependency ranges changed: dataframe-core, dataframe-expr-serializer, dataframe-operations, dataframe-parquet
Files
- dataframe-arrow-bridge.cabal +5/−5
- src/DataFrame/IO/Arrow.hs +16/−38
- src/DataFrame/IR.hs +5/−20
dataframe-arrow-bridge.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.4 name: dataframe-arrow-bridge-version: 1.0.0.0+version: 1.1.0.0 synopsis: Arrow C Data Interface bridge and plan IR for the dataframe ecosystem. description:@@ -48,13 +48,13 @@ aeson >= 0.11 && < 3, bytestring >= 0.11 && < 0.14, containers >= 0.6.7 && < 0.10,- dataframe-core >= 2.5 && < 2.6,+ dataframe-core >= 2.6 && < 2.7, dataframe-csv >= 2.3 && < 2.4,- dataframe-expr-serializer >= 1.2.1 && < 1.3,+ dataframe-expr-serializer >= 1.3 && < 1.4, dataframe-json >= 1.2.0.1 && < 1.3, dataframe-lazy >= 2.4.1 && < 2.5,- dataframe-operations >= 2.5 && < 2.6,- dataframe-parquet >= 1.5 && < 1.6,+ dataframe-operations >= 2.6 && < 2.7,+ dataframe-parquet >= 1.6 && < 1.7, dataframe-parsing >= 2.2 && < 2.3, text >= 2.1 && < 3, vector >= 0.13 && < 0.15
src/DataFrame/IO/Arrow.hs view
@@ -4,9 +4,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- | Convert a 'DataFrame' to Arrow C Data Interface structs for zero-copy- transfer to Python (or any other Arrow consumer).--}+-- | Convert a 'DataFrame' to Arrow C Data Interface structs, zero-copy. module DataFrame.IO.Arrow ( dataframeToArrow, columnToArrow,@@ -22,7 +20,6 @@ import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU import qualified DataFrame.Internal.Column as DI-import qualified DataFrame.Internal.Column.Bitmap as DI import Control.Monad (foldM_, forM, join, when, zipWithM_) import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))@@ -53,11 +50,7 @@ import Type.Reflection (typeRep) import DataFrame.Internal.Column (Column (..))-import DataFrame.Internal.DataFrame (DataFrame (..), fromNamedColumns)---- ------------------------------------------------------------------------------ Opaque phantom types for the Arrow structs--- ---------------------------------------------------------------------------+import DataFrame.Internal.DataFrame (DataFrame (..), fromColumns) data ArrowSchema data ArrowArray@@ -112,10 +105,6 @@ _arrayRelease = 64 _arrayPrivateData = 72 --- ------------------------------------------------------------------------------ Helpers--- ---------------------------------------------------------------------------- -- Write a Storable value at a byte offset from a base pointer. at :: (Storable a) => Ptr b -> Int -> a -> IO () at p off = poke (castPtr (p `plusPtr` off))@@ -124,10 +113,6 @@ readAt :: (Storable a) => Ptr b -> Int -> IO a readAt p off = peek (castPtr (p `plusPtr` off)) --- ------------------------------------------------------------------------------ Release callbacks (self-import trick for compile-time-constant FunPtr)--- ---------------------------------------------------------------------------- foreign export ccall "df_release_schema" releaseSchemaImpl :: Ptr ArrowSchema -> IO () @@ -153,9 +138,7 @@ let sp = castPtrToStablePtr rawPriv :: StablePtr (IO ()) join (deRefStablePtr sp) freeStablePtr sp- -- Arrow spec: release callback must set release to NULL to signal completion.- -- p here is Arrow C++'s internal copy of the struct (not our mallocBytes- -- allocation); our original allocation is freed inside the cleanup closure.+ -- Arrow spec: the release callback must set release to NULL when done. p `at` _schemaRelease $ (nullFunPtr :: FunPtr (Ptr ArrowSchema -> IO ())) releaseArrayImpl :: Ptr ArrowArray -> IO ()@@ -180,7 +163,7 @@ p `at` _schemaChildren $ (nullPtr :: Ptr ()) p `at` _schemaDictionary $ (nullPtr :: Ptr ()) p `at` _schemaRelease $ pReleaseSchema- -- Capture p so our original mallocBytes allocation is freed when release runs.+ -- Capture p so our mallocBytes allocation is freed on release. cleanup <- newStablePtr (free fmtStr >> free nameStr >> free p) p `at` _schemaPrivateData $ castStablePtrToPtr cleanup return p@@ -200,14 +183,12 @@ p `at` _arrayChildren $ (nullPtr :: Ptr ()) p `at` _arrayDictionary $ (nullPtr :: Ptr ()) p `at` _arrayRelease $ pReleaseArray- -- Capture p so our original mallocBytes allocation is freed when release runs.+ -- Capture p so our mallocBytes allocation is freed on release. cleanup <- newStablePtr (free bufArr >> extraCleanup >> free p) p `at` _arrayPrivateData $ castStablePtrToPtr cleanup return p -{- | Allocate an Arrow-format validity bitmap from a 'DI.Bitmap'.-Returns (ptr, nullCount). Caller must 'free' the pointer.--}+-- | Allocate an Arrow validity bitmap from a 'DI.Bitmap'; caller frees it. bitmapToPtr :: Int -> DI.Bitmap -> IO (Ptr Word8, Int) bitmapToPtr n bm = do let numBytes = max 1 ((n + 7) `div` 8)@@ -226,7 +207,7 @@ readArrowBitmap bitmapPtr n = VU.generateM ((n + 7) `div` 8) (peekElemOff bitmapPtr) columnToArrow :: T.Text -> Column -> IO (Ptr ArrowSchema, Ptr ArrowArray)-columnToArrow colName (UnboxedColumn _ (vec :: VU.Vector a))+columnToArrow colName (UnboxedColumn Nothing (vec :: VU.Vector a)) | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = do let n = VU.length vec dataPtr <- mallocArray (max 1 n) :: IO (Ptr Int64)@@ -234,7 +215,7 @@ sPtr <- makeLeafSchema "l" colName aPtr <- makeLeafArray n 0 [nullPtr, castPtr dataPtr] (free dataPtr) return (sPtr, aPtr)-columnToArrow colName (UnboxedColumn _ (vec :: VU.Vector a))+columnToArrow colName (UnboxedColumn Nothing (vec :: VU.Vector a)) | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = do let n = VU.length vec dataPtr <- mallocArray (max 1 n) :: IO (Ptr Double)@@ -320,7 +301,7 @@ columnToArrow colName (BoxedColumn (Just bm) (vec :: V.Vector a)) | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) = do let n = V.length vec- -- For null positions, use empty BS (null placeholder in vec is never evaluated)+ -- Null positions get an empty ByteString. bss = map (\i -> if DI.bitmapTestBit bm i then TE.encodeUtf8 (vec V.! i) else BS.empty)@@ -350,6 +331,10 @@ [castPtr bitmapPtr, castPtr offPtr, castPtr charsPtr] (free bitmapPtr >> free offPtr >> free charsPtr) return (sPtr, aPtr)+columnToArrow colName c@(PackedText _ _) =+ columnToArrow colName (DI.materializePacked c)+columnToArrow colName c@(MergedColumn _ _) =+ columnToArrow colName (DI.materializeMerged c) columnToArrow colName _ = error $ "DataFrame.IO.Arrow.columnToArrow: unsupported column type for '"@@ -385,10 +370,7 @@ topSchema `at` _schemaChildren $ childSArr topSchema `at` _schemaDictionary $ (nullPtr :: Ptr ()) topSchema `at` _schemaRelease $ pReleaseSchema- -- Do NOT loop over children here: Arrow C++ zeroes children[i]->release- -- during import, so reading it would yield a null function pointer.- -- Children are released independently by Arrow C++; their own cleanup- -- closures free their buffers and struct memory.+ -- Do not release children here: Arrow C++ zeroes their release on import. cleanupS <- newStablePtr $ do free childSArr free fmtStr@@ -396,7 +378,6 @@ free topSchema -- free our original mallocBytes allocation topSchema `at` _schemaPrivateData $ castStablePtrToPtr cleanupS - -- ── Top-level struct array ────────────────────────────────────────────── topArray <- mallocBytes arrowArraySize childAArr <- mallocArray ncols :: IO (Ptr (Ptr ArrowArray)) zipWithM_ (pokeElemOff childAArr) [0 ..] childAPtrs@@ -420,10 +401,7 @@ return (topSchema, topArray) -{- | Import an Arrow RecordBatch from raw C Data Interface pointers.- Copies all data into GC-managed Haskell vectors, then calls the- producer's release callbacks.--}+-- | Import an Arrow RecordBatch from raw C Data Interface pointers. arrowToDataframe :: Ptr () -> Ptr () -> IO DataFrame arrowToDataframe rawSchema rawArray = do let schemaPtr = castPtr rawSchema :: Ptr ArrowSchema@@ -441,7 +419,7 @@ relS <- readAt schemaPtr _schemaRelease :: IO (FunPtr (Ptr ArrowSchema -> IO ())) when (relS /= nullFunPtr) $ callRelSchema relS schemaPtr- return $ fromNamedColumns cols+ return $ fromColumns cols readArrowColumn :: Ptr ArrowSchema -> Ptr ArrowArray -> IO (T.Text, Column) readArrowColumn schemaPtr arrayPtr = do
src/DataFrame/IR.hs view
@@ -7,9 +7,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- | Intermediate Representation for DataFrame query plans.- JSON-decodable plan tree + interpreter.--}+-- | JSON-decodable query plan tree and its interpreter. module DataFrame.IR ( PlanNode (..), AggSpec (..),@@ -65,10 +63,6 @@ import DataFrame.Operations.Transformations (derive) import DataFrame.Schema (Schema, makeSchema, schemaType) --- ------------------------------------------------------------------------------ IR types--- ---------------------------------------------------------------------------- data AggSpec = AggSpec { aggName :: T.Text , aggFn :: T.Text@@ -96,7 +90,7 @@ | Drop Int PlanNode | DropLast Int PlanNode | Range Int Int PlanNode- | -- | joinType ("inner"|"left"|"right"|"outer"), shared key columns, left, right+ | -- | joinType ("inner"|"left"|"right"|"outer"), key columns, left, right Join T.Text [T.Text] PlanNode PlanNode | Describe PlanNode | -- | first column, second column, child plan@@ -106,18 +100,11 @@ | ReadJson FilePath | -- | path, separator (single character), child plan; runs as a terminal op WriteCsv FilePath PlanNode- | {- | path, schema (column name → type-tag map). Reads via the lazy- engine with predicate / projection pushdown; subsequent ops- currently still run eagerly on the materialized result.- -}+ | -- | path, schema map; reads via the lazy engine with pushdown ScanCsv FilePath [(T.Text, T.Text)] | ScanParquet FilePath [(T.Text, T.Text)] deriving (Show) --- ------------------------------------------------------------------------------ JSON decoding--- ---------------------------------------------------------------------------- instance FromJSON AggSpec where parseJSON = withObject "AggSpec" $ \o -> AggSpec@@ -194,7 +181,7 @@ executePlan reader (Derive name exprJson node) = do df <- executePlan reader node case decodeExprAny exprJson of- Right (SomeExpr _trep expr) -> return $ derive name expr df+ Right (SomeExpr expr) -> return $ derive name expr df Left err -> ioError $ userError $ "DataFrame.IR.Derive: " <> err executePlan reader (Exclude cols node) = exclude cols <$> executePlan reader node@@ -250,9 +237,7 @@ schema <- buildSchema schemaPairs Lazy.runDataFrame (Lazy.scanParquet schema (T.pack path)) -{- | Build a SortOrder from a column's runtime type.-Uses type dispatch to recover Ord for known column types.--}+-- | Build a SortOrder from a column's runtime type. mkSortOrder :: Bool -> T.Text -> Column -> SortOrder mkSortOrder isAsc name col = dispatchType (columnTypeRep col) where