{- | Apache Arrow IPC metadata types.
Arrow IPC uses FlatBuffers for metadata. These types mirror the Arrow
Schema.fbs definitions for schema, field, type, and record batch metadata.
-}
module Arrow.Types (
Schema (..),
Field (..),
DictionaryEncoding (..),
ArrowType (..),
Endianness (..),
Precision (..),
DateUnit (..),
TimeUnit (..),
IntervalUnit (..),
UnionMode (..),
Message (..),
RecordBatchDef (..),
BodyCompressionCodec (..),
Feature (..),
FieldNode (..),
Buffer (..),
-- * Smart constructors
--
{- | These thin wrappers default the optional record fields
('arrowMetadata', 'fieldMetadata', 'fieldDictionary',
'fieldChildren') so test fixtures and one-off encoders
don't have to spell out every slot. Adding a new optional
field to 'Schema' or 'Field' won't break callers that
went through these constructors.
-}
defaultSchema,
defaultField,
defaultLeafField,
-- * Decimal schema parameters
validateDecimalType,
-- * Schema fingerprint
schemaFingerprint,
schemaEquivalent,
) where
import Control.DeepSeq (NFData)
import Crypto.Hash qualified as Hash
import Data.Bits (shiftL, (.|.))
import Data.ByteArray qualified as BA
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as BL
import Data.Int (Int32, Int64)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding qualified as TE
import Data.Vector (Vector)
import Data.Vector qualified as V
import Data.Vector.Storable qualified as VS
import Data.Word (Word64, Word8)
import Foreign.Storable (Storable (..))
import GHC.Generics (Generic)
import Wireform.Builder qualified as BB
data Endianness = Little | Big
deriving stock (Show, Eq, Enum, Bounded, Generic)
deriving anyclass (NFData)
data Precision = Half | Single | DoublePrecision
deriving stock (Show, Eq, Enum, Bounded, Generic)
deriving anyclass (NFData)
data DateUnit = DateDay | DateMillisecond
deriving stock (Show, Eq, Enum, Bounded, Generic)
deriving anyclass (NFData)
data TimeUnit = Second | Millisecond | Microsecond | Nanosecond
deriving stock (Show, Eq, Enum, Bounded, Generic)
deriving anyclass (NFData)
data IntervalUnit = YearMonth | DayTime | MonthDayNano
deriving stock (Show, Eq, Enum, Bounded, Generic)
deriving anyclass (NFData)
data UnionMode = Sparse | Dense
deriving stock (Show, Eq, Enum, Bounded, Generic)
deriving anyclass (NFData)
data ArrowType
= ANull
| AInt !Int !Bool
| AFloatingPoint !Precision
| ABinary
| AUtf8
| ABool
| -- | Arrow format 1.5 signed 32-bit decimal (precision, scale).
ADecimal32 !Int !Int
| -- | Arrow format 1.5 signed 64-bit decimal (precision, scale).
ADecimal64 !Int !Int
| {- | 128-bit 'Decimal' (precision, scale). Kept as @ADecimal@
for source compatibility with the original schema; the
256-bit variant is 'ADecimal256'.
-}
ADecimal !Int !Int
| {- | 256-bit 'Decimal' (precision, scale). Arrow's on-wire
@Decimal@ type carries a @bitWidth@ field; we keep the two
widths as distinct 'ArrowType' constructors so the reader
can dispatch to the right 'ColumnArray' variant.
-}
ADecimal256 !Int !Int
| ADate !DateUnit
| ATime !TimeUnit !Int
| ATimestamp !TimeUnit !(Maybe Text)
| AInterval !IntervalUnit
| AList
| AStruct
| AUnion !UnionMode !(Vector Int32)
| AFixedSizeBinary !Int
| AFixedSizeList !Int
| AMap !Bool
| ADuration !TimeUnit
| ALargeBinary
| ALargeUtf8
| ALargeList
| -- The remaining constructors are post-V5 schema additions
-- (Arrow format version >= 1.4). The metadata writer in
-- "Arrow.FlatBufferIPC" understands them; a corresponding column
-- materializer / writer in "Arrow.Column" / "Arrow.Write" is not
-- yet wired up — the schema slot is enough for round-tripping
-- type information from external readers.
ARunEndEncoded
| ABinaryView
| AUtf8View
| AListView
| ALargeListView
deriving stock (Show, Eq, Generic)
deriving anyclass (NFData)
{- | Check decimal precision and the signed 32-bit scale carried by
Schema.fbs. General Arrow scales may be negative or exceed precision.
Other types are unaffected; this does not validate an entire schema.
-}
validateDecimalType :: ArrowType -> Either String ()
validateDecimalType ty = case ty of
ADecimal32 precision scale -> decimal 9 precision scale
ADecimal64 precision scale -> decimal 18 precision scale
ADecimal precision scale -> decimal 38 precision scale
ADecimal256 precision scale -> decimal 76 precision scale
_ -> Right ()
where
decimal maximumPrecision precision scale
| precision < 1 || precision > maximumPrecision =
Left ("Arrow.Types: decimal precision must be between 1 and " ++ show maximumPrecision)
| scale < fromIntegral (minBound :: Int32) || scale > fromIntegral (maxBound :: Int32) =
Left "Arrow.Types: decimal scale exceeds signed 32-bit schema field"
| otherwise = Right ()
data Field = Field
{ fieldName :: !Text
, fieldNullable :: !Bool
, fieldType :: !ArrowType
, fieldChildren :: !(Vector Field)
, fieldDictionary :: !(Maybe DictionaryEncoding)
{- ^ When non-'Nothing', the column is dictionary-encoded: the
record batch carries indices of type 'deIndexType', while
@fieldType@ / @fieldChildren@ describe the dictionary /values/,
which live in separate 'DictionaryBatch' messages keyed by 'deId'
(as in the Arrow spec).
-}
, fieldMetadata :: !(Vector (Text, Text))
{- ^ Arrow per-field @custom_metadata@ (Schema.fbs field 6).
Free-form key/value pairs that travel with the field
through encode + decode. Empty when the field carries
no annotation (the common case).
-}
}
deriving stock (Show, Eq, Generic)
deriving anyclass (NFData)
{- | Per @format/Schema.fbs@:
@
table DictionaryEncoding {
id: long;
indexType: Int;
isOrdered: bool;
dictionaryKind: DictionaryKind;
}
@
-}
data DictionaryEncoding = DictionaryEncoding
{ deId :: !Int64
, deIndexType :: !ArrowType -- always an AInt; default Int32 signed
, deIsOrdered :: !Bool
}
deriving stock (Show, Eq, Generic)
deriving anyclass (NFData)
{- | Arrow schema feature flags (Schema.fbs's @Feature@ enum).
'FeatureDictionaryReplacement' (1) advertises multiple full dictionaries
with the same id; 'FeatureCompressedBody' (2) advertises IPC body
compression. 'FeatureUnused' (0) is the FlatBuffers placeholder.
-}
data Feature
= FeatureUnused
| FeatureDictionaryReplacement
| FeatureCompressedBody
deriving stock (Show, Eq, Enum, Bounded, Generic)
deriving anyclass (NFData)
data Schema = Schema
{ arrowFields :: !(Vector Field)
, arrowEndianness :: !Endianness
, arrowMetadata :: !(Vector (Text, Text))
{- ^ Arrow schema-level @custom_metadata@ (Schema.fbs slot
2 of @table Schema@). Most files set zero or one
(@pandas@-style) annotation. Empty by default.
-}
, arrowFeatures :: !(Vector Feature)
{- ^ Arrow schema-level @features@ (Schema.fbs slot 3).
Empty by default. Flags survive IPC stream and file round-trips.
Producers emitting replacement dictionaries or compressed bodies can
advertise them with 'FeatureDictionaryReplacement' or
'FeatureCompressedBody'. Unknown feature values are rejected on read.
-}
}
deriving stock (Show, Eq, Generic)
deriving anyclass (NFData)
{- | One array's length and null count (Message.fbs @struct FieldNode@,
16 bytes). 'Storable' with the wire layout, so a record batch's nodes
are one 'VS.Vector'.
-}
data FieldNode = FieldNode
{ fnLength :: !Int64
, fnNullCount :: !Int64
}
deriving stock (Show, Eq, Generic)
deriving anyclass (NFData)
instance Storable FieldNode where
sizeOf _ = 16
alignment _ = 8
peek p = FieldNode <$> peekByteOff p 0 <*> peekByteOff p 8
poke p (FieldNode l n) = pokeByteOff p 0 l >> pokeByteOff p 8 n
{-# INLINE peek #-}
{-# INLINE poke #-}
{- | One body buffer's offset and length (Message.fbs @struct Buffer@,
16 bytes). 'Storable' with the wire layout.
-}
data Buffer = Buffer
{ bufOffset :: !Int64
, bufLength :: !Int64
}
deriving stock (Show, Eq, Generic)
deriving anyclass (NFData)
instance Storable Buffer where
sizeOf _ = 16
alignment _ = 8
peek p = Buffer <$> peekByteOff p 0 <*> peekByteOff p 8
poke p (Buffer o l) = pokeByteOff p 0 o >> pokeByteOff p 8 l
{-# INLINE peek #-}
{-# INLINE poke #-}
data RecordBatchDef = RecordBatchDef
{ rbLength :: !Int64
, rbNodes :: !(VS.Vector FieldNode)
-- ^ Field nodes in schema pre-order (one memcpy from the message).
, rbBuffers :: !(VS.Vector Buffer)
-- ^ Body buffers in spec layout order.
, rbVariadicBufferCounts :: !(Vector Int64)
{- ^ Per Arrow @format/Message.fbs@: when the schema contains
@Utf8View@ or @BinaryView@ fields each such field has a
variable number of additional data buffers for out-of-line
string payloads. The vector lists, in pre-order schema
traversal order, the number of variadic data buffers per
view column. Empty for schemas without view types.
-}
, rbBodyCompression :: !(Maybe BodyCompressionCodec)
{- ^ Per Arrow @format/Message.fbs@'s 'BodyCompression' table.
When 'Just', each buffer in the body is wrapped in an
@<i64 uncompressedLength><compressed bytes>@ envelope using
the named codec; readers see the original layout once
decompressed. 'Nothing' = uncompressed buffers (the wire
default).
-}
}
deriving stock (Show, Eq, Generic)
deriving anyclass (NFData)
{- | Body-compression codecs supported by Arrow IPC's
'BodyCompression' table (per @format/Message.fbs@).
-}
data BodyCompressionCodec
= {- | LZ4 frame format (@CompressionType = LZ4_FRAME = 0@).
Most widely supported; the default for pyarrow when
compression is enabled.
-}
LZ4Frame
| -- | Zstandard (@CompressionType = ZSTD = 1@).
BodyZstd
deriving stock (Show, Eq, Generic)
deriving anyclass (NFData)
{- | One Arrow IPC message header: the schema, a dictionary batch
(dictionary id, is-delta flag, the batch holding the dictionary
values), or a record batch. Bodies travel separately.
-}
data Message
= SchemaMessage !Schema
| DictionaryBatch !Int64 !Bool !RecordBatchDef
| RecordBatch !RecordBatchDef
deriving stock (Show, Eq, Generic)
deriving anyclass (NFData)
{- | Build a 'Schema' with little-endian + no custom_metadata.
Equivalent to @Schema fields Little V.empty@.
-}
defaultSchema :: Vector Field -> Schema
defaultSchema fs =
Schema
{ arrowFields = fs
, arrowEndianness = Little
, arrowMetadata = V.empty
, arrowFeatures = V.empty
}
{- | Build a 'Field' with no children, no dictionary, and no
custom_metadata.
-}
defaultLeafField :: Text -> Bool -> ArrowType -> Field
defaultLeafField name nullable ty =
Field
{ fieldName = name
, fieldNullable = nullable
, fieldType = ty
, fieldChildren = V.empty
, fieldDictionary = Nothing
, fieldMetadata = V.empty
}
{- | Build a 'Field' with explicit children but no dictionary
encoding and no custom_metadata.
-}
defaultField :: Text -> Bool -> ArrowType -> Vector Field -> Field
defaultField name nullable ty children =
Field
{ fieldName = name
, fieldNullable = nullable
, fieldType = ty
, fieldChildren = children
, fieldDictionary = Nothing
, fieldMetadata = V.empty
}
{- | A 64-bit deterministic fingerprint of an Arrow 'Schema',
computed as the high 64 bits of SHA-256 of a canonical
byte-encoding of the schema's structural skeleton (Avro-
style \"CRC-64-AVRO\"-shaped construction, but with SHA-256
as the underlying hash for collision resistance).
Two schemas that compare 'schemaEquivalent' produce the same
fingerprint; the construction is byte-deterministic across
GHC versions, processes, and architectures (no
'Data.Hashable' randomisation), so the fingerprint is safe
to persist as a cache key or schema-version identifier.
Equivalence ignores 'arrowMetadata' / 'arrowFeatures' and
per-field 'fieldMetadata' — those are annotations, not
structure. Endianness, field order, types, nullability, and
the full nested-children tree all participate.
-}
schemaFingerprint :: Schema -> Word64
schemaFingerprint sch =
let !canon = BL.toStrict (BB.toLazyByteString (buildSchema sch))
!digest = Hash.hash canon :: Hash.Digest Hash.SHA256
!bytes = BA.convert digest :: ByteString
in word64BE (BS.take 8 bytes)
where
word64BE bs =
let r i = fromIntegral (BS.index bs i) :: Word64
in (r 0 `shiftL` 56)
.|. (r 1 `shiftL` 48)
.|. (r 2 `shiftL` 40)
.|. (r 3 `shiftL` 32)
.|. (r 4 `shiftL` 24)
.|. (r 5 `shiftL` 16)
.|. (r 6 `shiftL` 8)
.|. r 7
buildSchema s =
BB.word8 (endianTag (arrowEndianness s))
<> buildVecField (arrowFields s)
buildVecField fs =
BB.int32BE (fromIntegral (V.length fs))
<> foldMap buildField (V.toList fs)
buildField f =
buildText (fieldName f)
<> BB.word8 (if fieldNullable f then 1 else 0)
<> buildText (T.pack (show (fieldType f)))
<> buildVecField (fieldChildren f)
buildText t =
let bs = TE.encodeUtf8 t
in BB.int32BE (fromIntegral (BS.length bs))
<> BB.byteString bs
endianTag :: Endianness -> Word8
endianTag Little = 0
endianTag Big = 1
{- | Structural equivalence, mirroring 'schemaFingerprint'.
Annotation fields ('arrowMetadata', 'arrowFeatures',
per-field 'fieldMetadata') are ignored.
-}
schemaEquivalent :: Schema -> Schema -> Bool
schemaEquivalent a b =
arrowEndianness a == arrowEndianness b
&& V.length (arrowFields a) == V.length (arrowFields b)
&& and (V.zipWith fieldEquivalent (arrowFields a) (arrowFields b))
fieldEquivalent :: Field -> Field -> Bool
fieldEquivalent a b =
fieldName a == fieldName b
&& fieldNullable a == fieldNullable b
&& fieldType a == fieldType b
&& V.length (fieldChildren a) == V.length (fieldChildren b)
&& and (V.zipWith fieldEquivalent (fieldChildren a) (fieldChildren b))