diff --git a/dataframe-arrow-bridge.cabal b/dataframe-arrow-bridge.cabal --- a/dataframe-arrow-bridge.cabal +++ b/dataframe-arrow-bridge.cabal @@ -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 diff --git a/src/DataFrame/IO/Arrow.hs b/src/DataFrame/IO/Arrow.hs --- a/src/DataFrame/IO/Arrow.hs +++ b/src/DataFrame/IO/Arrow.hs @@ -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 diff --git a/src/DataFrame/IR.hs b/src/DataFrame/IR.hs --- a/src/DataFrame/IR.hs +++ b/src/DataFrame/IR.hs @@ -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