diff --git a/src/uproot/const.py b/src/uproot/const.py index a4a1feeca..798c03fd5 100644 --- a/src/uproot/const.py +++ b/src/uproot/const.py @@ -12,6 +12,8 @@ # determines when a file is "big" kStartBigFile = 2000000000 +kMaxTBasketBytes = numpy.iinfo(numpy.int32).max + # used in unmarshaling kByteCountMask = numpy.int64(0x40000000) kByteCountVMask = numpy.int64(0x4000) diff --git a/src/uproot/writing/_cascadetree.py b/src/uproot/writing/_cascadetree.py index 6019af610..7df7a85e7 100644 --- a/src/uproot/writing/_cascadetree.py +++ b/src/uproot/writing/_cascadetree.py @@ -1395,6 +1395,13 @@ def write_np_basket(self, sink, branch_name, compression, array): fObjlen = len(uncompressed_data) fNbytes = fKeylen + len(compressed_data) + if max(fObjlen, fNbytes) > uproot.const.kMaxTBasketBytes: + raise ValueError( + f"Numpy array data of branch {branch_name} has an uncompressed size of {fObjlen} bytes " + f"and a compressed size of {fNbytes} bytes, which is too large to fit in a TBasket. " + "Neither of these sizes can exceed 2 GiB." + ) + parent_location = self._directory.key.location # FIXME: is this correct? location = self._freesegments.allocate(fNbytes, dry_run=False) @@ -1470,6 +1477,13 @@ def write_jagged_basket(self, sink, branch_name, compression, array, offsets): fObjlen = len(uncompressed_data) fNbytes = fKeylen + len(compressed_data) + if max(fObjlen, fNbytes) > uproot.const.kMaxTBasketBytes: + raise ValueError( + f"Jagged array data of branch {branch_name} has an uncompressed size of {fObjlen} bytes " + f"and a compressed size of {fNbytes} bytes, which is too large to fit in a TBasket. " + "Neither of these sizes can exceed 2 GiB." + ) + parent_location = self._directory.key.location # FIXME: is this correct? location = self._freesegments.allocate(fNbytes, dry_run=False) @@ -1559,6 +1573,13 @@ def write_string_basket(self, sink, branch_name, compression, array, offsets): fObjlen = len(uncompressed_data) fNbytes = fKeylen + len(compressed_data) + if max(fObjlen, fNbytes) > uproot.const.kMaxTBasketBytes: + raise ValueError( + f"String data of branch {branch_name} has an uncompressed size of {fObjlen} bytes " + f"and a compressed size of {fNbytes} bytes, which is too large to fit in a TBasket. " + "Neither of these sizes can exceed 2 GiB." + ) + parent_location = self._directory.key.location # FIXME: is this correct? location = self._freesegments.allocate(fNbytes, dry_run=False)