Vendor dependencies

This commit is contained in:
2026-08-01 16:11:49 +03:00
parent 7f139a0241
commit 6b5e7f0f8b
29706 changed files with 9575646 additions and 0 deletions
File diff suppressed because one or more lines are too long
+939
View File
@@ -0,0 +1,939 @@
# This file is automatically @generated by Cargo.
# It is not intended for manual editing.
version = 4
[[package]]
name = "ahash"
version = "0.8.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75"
dependencies = [
"cfg-if",
"const-random",
"getrandom 0.3.4",
"once_cell",
"version_check",
"zerocopy",
]
[[package]]
name = "aho-corasick"
version = "1.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301"
dependencies = [
"memchr",
]
[[package]]
name = "alloca"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4"
dependencies = [
"cc",
]
[[package]]
name = "android_system_properties"
version = "0.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311"
dependencies = [
"libc",
]
[[package]]
name = "anes"
version = "0.1.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299"
[[package]]
name = "anstyle"
version = "1.0.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000"
[[package]]
name = "arrow-array"
version = "57.3.1"
dependencies = [
"ahash",
"arrow-buffer",
"arrow-data",
"arrow-schema",
"chrono",
"chrono-tz",
"criterion",
"half",
"hashbrown",
"num-complex",
"num-integer",
"num-traits",
"rand",
]
[[package]]
name = "arrow-buffer"
version = "57.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0d18b89b4c4f4811d0858175e79541fe98e33e18db3b011708bc287b1240593f"
dependencies = [
"bytes",
"half",
"num-bigint",
"num-traits",
]
[[package]]
name = "arrow-data"
version = "57.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c1683705c63dcf0d18972759eda48489028cbbff67af7d6bef2c6b7b74ab778a"
dependencies = [
"arrow-buffer",
"arrow-schema",
"half",
"num-integer",
"num-traits",
]
[[package]]
name = "arrow-schema"
version = "57.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e4cf0d4a6609679e03002167a61074a21d7b1ad9ea65e462b2c0a97f8a3b2bc6"
dependencies = [
"bitflags",
]
[[package]]
name = "autocfg"
version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
[[package]]
name = "bitflags"
version = "2.11.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3"
[[package]]
name = "bumpalo"
version = "3.20.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb"
[[package]]
name = "bytes"
version = "1.11.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33"
[[package]]
name = "cast"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5"
[[package]]
name = "cc"
version = "1.2.62"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a1dce859f0832a7d088c4f1119888ab94ef4b5d6795d1ce05afb7fe159d79f98"
dependencies = [
"find-msvc-tools",
"shlex",
]
[[package]]
name = "cfg-if"
version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
[[package]]
name = "chrono"
version = "0.4.44"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0"
dependencies = [
"iana-time-zone",
"num-traits",
"windows-link",
]
[[package]]
name = "chrono-tz"
version = "0.10.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3"
dependencies = [
"chrono",
"phf",
]
[[package]]
name = "ciborium"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e"
dependencies = [
"ciborium-io",
"ciborium-ll",
"serde",
]
[[package]]
name = "ciborium-io"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757"
[[package]]
name = "ciborium-ll"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9"
dependencies = [
"ciborium-io",
"half",
]
[[package]]
name = "clap"
version = "4.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51"
dependencies = [
"clap_builder",
]
[[package]]
name = "clap_builder"
version = "4.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f"
dependencies = [
"anstyle",
"clap_lex",
]
[[package]]
name = "clap_lex"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9"
[[package]]
name = "const-random"
version = "0.1.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359"
dependencies = [
"const-random-macro",
]
[[package]]
name = "const-random-macro"
version = "0.1.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e"
dependencies = [
"getrandom 0.2.17",
"once_cell",
"tiny-keccak",
]
[[package]]
name = "core-foundation-sys"
version = "0.8.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
[[package]]
name = "criterion"
version = "0.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3"
dependencies = [
"alloca",
"anes",
"cast",
"ciborium",
"clap",
"criterion-plot",
"itertools",
"num-traits",
"oorandom",
"page_size",
"regex",
"serde",
"serde_json",
"tinytemplate",
"walkdir",
]
[[package]]
name = "criterion-plot"
version = "0.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea"
dependencies = [
"cast",
"itertools",
]
[[package]]
name = "crunchy"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5"
[[package]]
name = "either"
version = "1.15.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
[[package]]
name = "find-msvc-tools"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582"
[[package]]
name = "futures-core"
version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d"
[[package]]
name = "futures-task"
version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393"
[[package]]
name = "futures-util"
version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6"
dependencies = [
"futures-core",
"futures-task",
"pin-project-lite",
"slab",
]
[[package]]
name = "getrandom"
version = "0.2.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0"
dependencies = [
"cfg-if",
"libc",
"wasi",
]
[[package]]
name = "getrandom"
version = "0.3.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd"
dependencies = [
"cfg-if",
"libc",
"r-efi",
"wasip2",
]
[[package]]
name = "half"
version = "2.7.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b"
dependencies = [
"cfg-if",
"crunchy",
"num-traits",
"zerocopy",
]
[[package]]
name = "hashbrown"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
[[package]]
name = "iana-time-zone"
version = "0.1.65"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470"
dependencies = [
"android_system_properties",
"core-foundation-sys",
"iana-time-zone-haiku",
"js-sys",
"log",
"wasm-bindgen",
"windows-core",
]
[[package]]
name = "iana-time-zone-haiku"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f"
dependencies = [
"cc",
]
[[package]]
name = "itertools"
version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186"
dependencies = [
"either",
]
[[package]]
name = "itoa"
version = "1.0.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
[[package]]
name = "js-sys"
version = "0.3.98"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "67df7112613f8bfd9150013a0314e196f4800d3201ae742489d999db2f979f08"
dependencies = [
"cfg-if",
"futures-util",
"once_cell",
"wasm-bindgen",
]
[[package]]
name = "libc"
version = "0.2.186"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "libm"
version = "0.2.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981"
[[package]]
name = "log"
version = "0.4.29"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
[[package]]
name = "memchr"
version = "2.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
[[package]]
name = "num-bigint"
version = "0.4.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9"
dependencies = [
"num-integer",
"num-traits",
]
[[package]]
name = "num-complex"
version = "0.4.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
dependencies = [
"num-traits",
]
[[package]]
name = "num-integer"
version = "0.1.46"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f"
dependencies = [
"num-traits",
]
[[package]]
name = "num-traits"
version = "0.2.19"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
dependencies = [
"autocfg",
"libm",
]
[[package]]
name = "once_cell"
version = "1.21.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
[[package]]
name = "oorandom"
version = "11.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e"
[[package]]
name = "page_size"
version = "0.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da"
dependencies = [
"libc",
"winapi",
]
[[package]]
name = "phf"
version = "0.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "913273894cec178f401a31ec4b656318d95473527be05c0752cc41cdc32be8b7"
dependencies = [
"phf_shared",
]
[[package]]
name = "phf_shared"
version = "0.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "06005508882fb681fd97892ecff4b7fd0fee13ef1aa569f8695dae7ab9099981"
dependencies = [
"siphasher",
]
[[package]]
name = "pin-project-lite"
version = "0.2.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd"
[[package]]
name = "ppv-lite86"
version = "0.2.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9"
dependencies = [
"zerocopy",
]
[[package]]
name = "proc-macro2"
version = "1.0.106"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
dependencies = [
"unicode-ident",
]
[[package]]
name = "quote"
version = "1.0.45"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
dependencies = [
"proc-macro2",
]
[[package]]
name = "r-efi"
version = "5.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
[[package]]
name = "rand"
version = "0.9.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea"
dependencies = [
"rand_chacha",
"rand_core",
]
[[package]]
name = "rand_chacha"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
dependencies = [
"ppv-lite86",
"rand_core",
]
[[package]]
name = "rand_core"
version = "0.9.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c"
dependencies = [
"getrandom 0.3.4",
]
[[package]]
name = "regex"
version = "1.12.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276"
dependencies = [
"aho-corasick",
"memchr",
"regex-automata",
"regex-syntax",
]
[[package]]
name = "regex-automata"
version = "0.4.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f"
dependencies = [
"aho-corasick",
"memchr",
"regex-syntax",
]
[[package]]
name = "regex-syntax"
version = "0.8.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a"
[[package]]
name = "rustversion"
version = "1.0.22"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
[[package]]
name = "same-file"
version = "1.0.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502"
dependencies = [
"winapi-util",
]
[[package]]
name = "serde"
version = "1.0.228"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e"
dependencies = [
"serde_core",
"serde_derive",
]
[[package]]
name = "serde_core"
version = "1.0.228"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad"
dependencies = [
"serde_derive",
]
[[package]]
name = "serde_derive"
version = "1.0.228"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "serde_json"
version = "1.0.149"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86"
dependencies = [
"itoa",
"memchr",
"serde",
"serde_core",
"zmij",
]
[[package]]
name = "shlex"
version = "1.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64"
[[package]]
name = "siphasher"
version = "1.0.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649"
[[package]]
name = "slab"
version = "0.4.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5"
[[package]]
name = "syn"
version = "2.0.117"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99"
dependencies = [
"proc-macro2",
"quote",
"unicode-ident",
]
[[package]]
name = "tiny-keccak"
version = "2.0.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237"
dependencies = [
"crunchy",
]
[[package]]
name = "tinytemplate"
version = "1.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc"
dependencies = [
"serde",
"serde_json",
]
[[package]]
name = "unicode-ident"
version = "1.0.24"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
[[package]]
name = "version_check"
version = "0.9.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
[[package]]
name = "walkdir"
version = "2.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b"
dependencies = [
"same-file",
"winapi-util",
]
[[package]]
name = "wasi"
version = "0.11.1+wasi-snapshot-preview1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
[[package]]
name = "wasip2"
version = "1.0.3+wasi-0.2.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6"
dependencies = [
"wit-bindgen",
]
[[package]]
name = "wasm-bindgen"
version = "0.2.121"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "49ace1d07c165b0864824eee619580c4689389afa9dc9ed3a4c75040d82e6790"
dependencies = [
"cfg-if",
"once_cell",
"rustversion",
"wasm-bindgen-macro",
"wasm-bindgen-shared",
]
[[package]]
name = "wasm-bindgen-macro"
version = "0.2.121"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8e68e6f4afd367a562002c05637acb8578ff2dea1943df76afb9e83d177c8578"
dependencies = [
"quote",
"wasm-bindgen-macro-support",
]
[[package]]
name = "wasm-bindgen-macro-support"
version = "0.2.121"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d95a9ec35c64b2a7cb35d3fead40c4238d0940c86d107136999567a4703259f2"
dependencies = [
"bumpalo",
"proc-macro2",
"quote",
"syn",
"wasm-bindgen-shared",
]
[[package]]
name = "wasm-bindgen-shared"
version = "0.2.121"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c4e0100b01e9f0d03189a92b96772a1fb998639d981193d7dbab487302513441"
dependencies = [
"unicode-ident",
]
[[package]]
name = "winapi"
version = "0.3.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
dependencies = [
"winapi-i686-pc-windows-gnu",
"winapi-x86_64-pc-windows-gnu",
]
[[package]]
name = "winapi-i686-pc-windows-gnu"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
[[package]]
name = "winapi-util"
version = "0.1.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
dependencies = [
"windows-sys",
]
[[package]]
name = "winapi-x86_64-pc-windows-gnu"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
[[package]]
name = "windows-core"
version = "0.62.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb"
dependencies = [
"windows-implement",
"windows-interface",
"windows-link",
"windows-result",
"windows-strings",
]
[[package]]
name = "windows-implement"
version = "0.60.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "windows-interface"
version = "0.59.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "windows-link"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
[[package]]
name = "windows-result"
version = "0.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5"
dependencies = [
"windows-link",
]
[[package]]
name = "windows-strings"
version = "0.5.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091"
dependencies = [
"windows-link",
]
[[package]]
name = "windows-sys"
version = "0.61.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc"
dependencies = [
"windows-link",
]
[[package]]
name = "wit-bindgen"
version = "0.57.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e"
[[package]]
name = "zerocopy"
version = "0.8.48"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eed437bf9d6692032087e337407a86f04cd8d6a16a37199ed57949d415bd68e9"
dependencies = [
"zerocopy-derive",
]
[[package]]
name = "zerocopy-derive"
version = "0.8.48"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "zmij"
version = "1.0.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa"
+147
View File
@@ -0,0 +1,147 @@
# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO
#
# When uploading crates to the registry Cargo will automatically
# "normalize" Cargo.toml files for maximal compatibility
# with all versions of Cargo and also rewrite `path` dependencies
# to registry (e.g., crates.io) dependencies.
#
# If you are reading this file be aware that the original Cargo.toml
# will likely look very different (and much more reasonable).
# See Cargo.toml.orig for the original contents.
[package]
edition = "2024"
rust-version = "1.85"
name = "arrow-array"
version = "57.3.1"
authors = ["Apache Arrow <dev@arrow.apache.org>"]
build = false
include = [
"benches/*.rs",
"src/**/*.rs",
"Cargo.toml",
"LICENSE.txt",
"NOTICE.txt",
]
autolib = false
autobins = false
autoexamples = false
autotests = false
autobenches = false
description = "Array abstractions for Apache Arrow"
homepage = "https://github.com/apache/arrow-rs"
readme = false
keywords = ["arrow"]
license = "Apache-2.0"
repository = "https://github.com/apache/arrow-rs"
resolver = "2"
[package.metadata.docs.rs]
all-features = true
[features]
ffi = [
"arrow-schema/ffi",
"arrow-data/ffi",
]
force_validate = []
[lib]
name = "arrow_array"
path = "src/lib.rs"
bench = false
[[bench]]
name = "decimal_overflow"
path = "benches/decimal_overflow.rs"
harness = false
[[bench]]
name = "fixed_size_list_array"
path = "benches/fixed_size_list_array.rs"
harness = false
[[bench]]
name = "occupancy"
path = "benches/occupancy.rs"
harness = false
[[bench]]
name = "record_batch"
path = "benches/record_batch.rs"
harness = false
[[bench]]
name = "union_array"
path = "benches/union_array.rs"
harness = false
[[bench]]
name = "view_types"
path = "benches/view_types.rs"
harness = false
[dependencies.arrow-buffer]
version = "57.3.1"
[dependencies.arrow-data]
version = "57.3.1"
[dependencies.arrow-schema]
version = "57.3.1"
[dependencies.chrono]
version = "0.4.40"
features = ["clock"]
default-features = false
[dependencies.chrono-tz]
version = "0.10"
optional = true
[dependencies.half]
version = "2.1"
features = ["num-traits"]
default-features = false
[dependencies.hashbrown]
version = "0.16.0"
default-features = false
[dependencies.num-complex]
version = "0.4.6"
features = ["std"]
default-features = false
[dependencies.num-integer]
version = "0.1.46"
features = ["std"]
default-features = false
[dependencies.num-traits]
version = "0.2.19"
features = ["std"]
default-features = false
[dev-dependencies.criterion]
version = "0.8.0"
default-features = false
[dev-dependencies.rand]
version = "0.9"
features = [
"std",
"std_rng",
"thread_rng",
]
default-features = false
[target.'cfg(not(target_arch = "wasm32"))'.dependencies.ahash]
version = "0.8"
features = ["runtime-rng"]
default-features = false
[target.'cfg(target_arch = "wasm32")'.dependencies.ahash]
version = "0.8"
features = ["compile-time-rng"]
default-features = false
+86
View File
@@ -0,0 +1,86 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
[package]
name = "arrow-array"
version = { workspace = true }
description = "Array abstractions for Apache Arrow"
homepage = { workspace = true }
repository = { workspace = true }
authors = { workspace = true }
license = { workspace = true }
keywords = { workspace = true }
include = { workspace = true }
edition = { workspace = true }
rust-version = { workspace = true }
[lib]
name = "arrow_array"
bench = false
[target.'cfg(target_arch = "wasm32")'.dependencies]
ahash = { version = "0.8", default-features = false, features = ["compile-time-rng"] }
[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
ahash = { version = "0.8", default-features = false, features = ["runtime-rng"] }
[dependencies]
arrow-buffer = { workspace = true }
arrow-schema = { workspace = true }
arrow-data = { workspace = true }
chrono = { workspace = true }
chrono-tz = { version = "0.10", optional = true }
num-complex = { version = "0.4.6", default-features = false, features = ["std"] }
num-integer = { version = "0.1.46", default-features = false, features = ["std"] }
num-traits = { version = "0.2.19", default-features = false, features = ["std"] }
half = { version = "2.1", default-features = false, features = ["num-traits"] }
hashbrown = { version = "0.16.0", default-features = false }
[package.metadata.docs.rs]
all-features = true
[features]
ffi = ["arrow-schema/ffi", "arrow-data/ffi"]
force_validate = []
[dev-dependencies]
rand = { version = "0.9", default-features = false, features = ["std", "std_rng", "thread_rng"] }
criterion = { workspace = true, default-features = false }
[[bench]]
name = "occupancy"
harness = false
[[bench]]
name = "view_types"
harness = false
[[bench]]
name = "fixed_size_list_array"
harness = false
[[bench]]
name = "decimal_overflow"
harness = false
[[bench]]
name = "union_array"
harness = false
[[bench]]
name = "record_batch"
harness = false
+202
View File
@@ -0,0 +1,202 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
+84
View File
@@ -0,0 +1,84 @@
Apache Arrow
Copyright 2016-2019 The Apache Software Foundation
This product includes software developed at
The Apache Software Foundation (http://www.apache.org/).
This product includes software from the SFrame project (BSD, 3-clause).
* Copyright (C) 2015 Dato, Inc.
* Copyright (c) 2009 Carnegie Mellon University.
This product includes software from the Feather project (Apache 2.0)
https://github.com/wesm/feather
This product includes software from the DyND project (BSD 2-clause)
https://github.com/libdynd
This product includes software from the LLVM project
* distributed under the University of Illinois Open Source
This product includes software from the google-lint project
* Copyright (c) 2009 Google Inc. All rights reserved.
This product includes software from the mman-win32 project
* Copyright https://code.google.com/p/mman-win32/
* Licensed under the MIT License;
This product includes software from the LevelDB project
* Copyright (c) 2011 The LevelDB Authors. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* Moved from Kudu http://github.com/cloudera/kudu
This product includes software from the CMake project
* Copyright 2001-2009 Kitware, Inc.
* Copyright 2012-2014 Continuum Analytics, Inc.
* All rights reserved.
This product includes software from https://github.com/matthew-brett/multibuild (BSD 2-clause)
* Copyright (c) 2013-2016, Matt Terry and Matthew Brett; all rights reserved.
This product includes software from the Ibis project (Apache 2.0)
* Copyright (c) 2015 Cloudera, Inc.
* https://github.com/cloudera/ibis
This product includes software from Dremio (Apache 2.0)
* Copyright (C) 2017-2018 Dremio Corporation
* https://github.com/dremio/dremio-oss
This product includes software from Google Guava (Apache 2.0)
* Copyright (C) 2007 The Guava Authors
* https://github.com/google/guava
This product include software from CMake (BSD 3-Clause)
* CMake - Cross Platform Makefile Generator
* Copyright 2000-2019 Kitware, Inc. and Contributors
The web site includes files generated by Jekyll.
--------------------------------------------------------------------------------
This product includes code from Apache Kudu, which includes the following in
its NOTICE file:
Apache Kudu
Copyright 2016 The Apache Software Foundation
This product includes software developed at
The Apache Software Foundation (http://www.apache.org/).
Portions of this software were developed at
Cloudera, Inc (http://www.cloudera.com/).
--------------------------------------------------------------------------------
This product includes code from Apache ORC, which includes the following in
its NOTICE file:
Apache ORC
Copyright 2013-2019 The Apache Software Foundation
This product includes software developed by The Apache Software
Foundation (http://www.apache.org/).
This product includes software developed by Hewlett-Packard:
(c) Copyright [2014-2015] Hewlett-Packard Development Company, L.P
@@ -0,0 +1,54 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use arrow_array::builder::{Decimal128Builder, Decimal256Builder};
use arrow_buffer::i256;
use criterion::*;
use std::hint;
fn criterion_benchmark(c: &mut Criterion) {
let len = 8192;
let mut builder_128 = Decimal128Builder::with_capacity(len);
let mut builder_256 = Decimal256Builder::with_capacity(len);
for i in 0..len {
if i % 10 == 0 {
builder_128.append_value(i128::MAX);
builder_256.append_value(i256::from_i128(i128::MAX));
} else {
builder_128.append_value(i as i128);
builder_256.append_value(i256::from_i128(i as i128));
}
}
let array_128 = builder_128.finish();
let array_256 = builder_256.finish();
c.bench_function("validate_decimal_precision_128", |b| {
b.iter(|| hint::black_box(array_128.validate_decimal_precision(8)));
});
c.bench_function("null_if_overflow_precision_128", |b| {
b.iter(|| hint::black_box(array_128.null_if_overflow_precision(8)));
});
c.bench_function("validate_decimal_precision_256", |b| {
b.iter(|| hint::black_box(array_256.validate_decimal_precision(8)));
});
c.bench_function("null_if_overflow_precision_256", |b| {
b.iter(|| hint::black_box(array_256.null_if_overflow_precision(8)));
});
}
criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);
@@ -0,0 +1,51 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use arrow_array::{Array, FixedSizeListArray, Int32Array};
use arrow_schema::Field;
use criterion::*;
use rand::{Rng, rng};
use std::{hint, sync::Arc};
fn gen_fsl(len: usize, value_len: usize) -> FixedSizeListArray {
let mut rng = rng();
let values = Arc::new(Int32Array::from(
(0..len).map(|_| rng.random::<i32>()).collect::<Vec<_>>(),
));
let field = Arc::new(Field::new_list_field(values.data_type().clone(), true));
FixedSizeListArray::new(field, value_len as i32, values, None)
}
fn criterion_benchmark(c: &mut Criterion) {
let len = 4096;
for value_len in [1, 32, 1024] {
let fsl = gen_fsl(len, value_len);
c.bench_function(
&format!("fixed_size_list_array(len: {len}, value_len: {value_len})"),
|b| {
b.iter(|| {
for i in 0..len / value_len {
hint::black_box(fsl.value(i));
}
});
},
);
}
}
criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);
+57
View File
@@ -0,0 +1,57 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use arrow_array::types::Int32Type;
use arrow_array::{DictionaryArray, Int32Array};
use arrow_buffer::NullBuffer;
use criterion::*;
use rand::{Rng, rng};
use std::{hint, sync::Arc};
fn gen_dict(
len: usize,
values_len: usize,
occupancy: f64,
null_percent: f64,
) -> DictionaryArray<Int32Type> {
let mut rng = rng();
let values = Int32Array::from(vec![0; values_len]);
let max_key = (values_len as f64 * occupancy) as i32;
let keys = (0..len).map(|_| rng.random_range(0..max_key)).collect();
let nulls = (0..len).map(|_| !rng.random_bool(null_percent)).collect();
let keys = Int32Array::new(keys, Some(NullBuffer::new(nulls)));
DictionaryArray::new(keys, Arc::new(values))
}
fn criterion_benchmark(c: &mut Criterion) {
for values in [10, 100, 512] {
for occupancy in [1., 0.5, 0.1] {
for null_percent in [0.0, 0.1, 0.5, 0.9] {
let dict = gen_dict(1024, values, occupancy, null_percent);
c.bench_function(&format!("occupancy(values: {values}, occupancy: {occupancy}, null_percent: {null_percent})"), |b| {
b.iter(|| {
hint::black_box(&dict).occupancy()
});
});
}
}
}
}
criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);
@@ -0,0 +1,93 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use arrow_array::{ArrayRef, Int64Array, RecordBatch, RecordBatchOptions};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use num_integer::Integer;
use std::hint::black_box;
use std::sync::Arc;
fn make_record_batch(column_count: usize, row_count: usize) -> RecordBatch {
let fields = (0..column_count)
.map(|i| Field::new(format!("col_{}", i), DataType::Int64, i.is_even()))
.collect::<Vec<_>>();
let columns = fields
.iter()
.map(|_| {
let array_ref: ArrayRef = Arc::new(Int64Array::from_value(0, row_count));
array_ref
})
.collect::<Vec<_>>();
let schema = Schema::new(fields);
let mut options = RecordBatchOptions::new();
options.row_count = Some(row_count);
RecordBatch::try_new_with_options(SchemaRef::new(schema), columns, &options).unwrap()
}
fn project_benchmark(
c: &mut Criterion,
column_count: usize,
row_count: usize,
projection_size: usize,
) {
let input = make_input(column_count, row_count, projection_size);
c.bench_with_input(
BenchmarkId::new(
"project",
format!(
"{:?}x{:?} -> {:?}x{:?}",
input.0.num_columns(),
input.0.num_rows(),
input.1.len(),
input.0.num_rows()
),
),
&input,
|b, (rb, projection)| {
b.iter(|| black_box(rb.project(projection).unwrap()));
},
);
}
fn make_input(
column_count: usize,
row_count: usize,
projection_size: usize,
) -> (RecordBatch, Vec<usize>) {
let rb = make_record_batch(column_count, row_count);
let projection = (0..projection_size).collect::<Vec<_>>();
(rb, projection)
}
fn criterion_benchmark(c: &mut Criterion) {
[10, 100, 1000].iter().for_each(|&column_count| {
[1, column_count / 2, column_count - 1]
.iter()
.for_each(|&projection_size| {
project_benchmark(c, column_count, 8192, projection_size);
})
});
}
criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);
+80
View File
@@ -0,0 +1,80 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use std::{hint, iter::repeat_with, sync::Arc};
use arrow_array::{Array, ArrayRef, Int32Array, UnionArray};
use arrow_buffer::{NullBuffer, ScalarBuffer};
use arrow_schema::{DataType, Field, UnionFields};
use criterion::*;
use rand::{Rng, rng};
fn array_with_nulls() -> ArrayRef {
let mut rng = rng();
let values = ScalarBuffer::from_iter(repeat_with(|| rng.random()).take(4096));
// nulls with at least one null and one valid
let nulls: NullBuffer = [true, false]
.into_iter()
.chain(repeat_with(|| rng.random()))
.take(4096)
.collect();
Arc::new(Int32Array::new(values.clone(), Some(nulls)))
}
fn array_without_nulls() -> ArrayRef {
let mut rng = rng();
let values = ScalarBuffer::from_iter(repeat_with(|| rng.random()).take(4096));
Arc::new(Int32Array::new(values.clone(), None))
}
fn criterion_benchmark(c: &mut Criterion) {
for with_nulls in 1..12 {
for without_nulls in [0, 1, 10] {
c.bench_function(
&format!("union logical nulls 4096 {with_nulls} children with nulls, {without_nulls} without nulls"),
|b| {
let type_ids = 0..with_nulls+without_nulls;
let fields = UnionFields::try_new(
type_ids.clone(),
type_ids.clone().map(|i| Field::new(format!("f{i}"), DataType::Int32, true)),
).unwrap();
let array = UnionArray::try_new(
fields,
type_ids.cycle().take(4096).collect(),
None,
std::iter::repeat_n(array_with_nulls(), with_nulls as usize)
.chain(std::iter::repeat_n(array_without_nulls(), without_nulls as usize))
.collect(),
)
.unwrap();
b.iter(|| hint::black_box(array.logical_nulls()))
},
);
}
}
}
criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);
+112
View File
@@ -0,0 +1,112 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use arrow_array::StringViewArray;
use criterion::{Criterion, criterion_group, criterion_main};
use std::hint::black_box;
fn gen_view_array(size: usize) -> StringViewArray {
StringViewArray::from_iter((0..size).map(|v| match v % 3 {
0 => Some("small"),
1 => Some("larger than 12 bytes array"),
2 => None,
_ => unreachable!("unreachable"),
}))
}
fn gen_view_array_without_nulls(size: usize) -> StringViewArray {
StringViewArray::from_iter((0..size).map(|v| {
let s = match v % 3 {
0 => "small".to_string(), // < 12 bytes
1 => "larger than 12 bytes array".to_string(), // >12 bytes
2 => "x".repeat(300), // 300 bytes (>256)
_ => unreachable!(),
};
Some(s)
}))
}
fn criterion_benchmark(c: &mut Criterion) {
let array = gen_view_array(100_000);
c.bench_function("view types slice", |b| {
b.iter(|| {
black_box(array.slice(0, 100_000 / 2));
});
});
c.bench_function("gc view types all[100000]", |b| {
b.iter(|| {
black_box(array.gc());
});
});
let sliced = array.slice(0, 100_000 / 2);
c.bench_function("gc view types slice half[100000]", |b| {
b.iter(|| {
black_box(sliced.gc());
});
});
let array = gen_view_array_without_nulls(100_000);
c.bench_function("gc view types all without nulls[100000]", |b| {
b.iter(|| {
black_box(array.gc());
});
});
let sliced = array.slice(0, 100_000 / 2);
c.bench_function("gc view types slice half without nulls[100000]", |b| {
b.iter(|| {
black_box(sliced.gc());
});
});
let array = gen_view_array(8000);
c.bench_function("gc view types all[8000]", |b| {
b.iter(|| {
black_box(array.gc());
});
});
let sliced = array.slice(0, 8000 / 2);
c.bench_function("gc view types slice half[8000]", |b| {
b.iter(|| {
black_box(sliced.gc());
});
});
let array = gen_view_array_without_nulls(8000);
c.bench_function("gc view types all without nulls[8000]", |b| {
b.iter(|| {
black_box(array.gc());
});
});
let sliced = array.slice(0, 8000 / 2);
c.bench_function("gc view types slice half without nulls[8000]", |b| {
b.iter(|| {
black_box(sliced.gc());
});
});
}
criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);
+903
View File
@@ -0,0 +1,903 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use arrow_buffer::{ArrowNativeType, IntervalDayTime, IntervalMonthDayNano, i256};
use arrow_schema::ArrowError;
use half::f16;
use num_complex::ComplexFloat;
use std::cmp::Ordering;
/// Trait for [`ArrowNativeType`] that adds checked and unchecked arithmetic operations,
/// and totally ordered comparison operations
///
/// The APIs with `_wrapping` suffix do not perform overflow-checking. For integer
/// types they will wrap around the boundary of the type. For floating point types they
/// will overflow to INF or -INF preserving the expected sign value
///
/// Note `div_wrapping` and `mod_wrapping` will panic for integer types if `rhs` is zero
/// although this may be subject to change <https://github.com/apache/arrow-rs/issues/2647>
///
/// The APIs with `_checked` suffix perform overflow-checking. For integer types
/// these will return `Err` instead of wrapping. For floating point types they will
/// overflow to INF or -INF preserving the expected sign value
///
/// Comparison of integer types is as per normal integer comparison rules, floating
/// point values are compared as per IEEE 754's totalOrder predicate see [`f32::total_cmp`]
///
pub trait ArrowNativeTypeOp: ArrowNativeType {
/// The additive identity
const ZERO: Self;
/// The multiplicative identity
const ONE: Self;
/// The minimum value and identity for the `max` aggregation.
/// Note that the aggregation uses the total order predicate for floating point values,
/// which means that this value is a negative NaN.
const MIN_TOTAL_ORDER: Self;
/// The maximum value and identity for the `min` aggregation.
/// Note that the aggregation uses the total order predicate for floating point values,
/// which means that this value is a positive NaN.
const MAX_TOTAL_ORDER: Self;
/// Checked addition operation
fn add_checked(self, rhs: Self) -> Result<Self, ArrowError>;
/// Wrapping addition operation
fn add_wrapping(self, rhs: Self) -> Self;
/// Checked subtraction operation
fn sub_checked(self, rhs: Self) -> Result<Self, ArrowError>;
/// Wrapping subtraction operation
fn sub_wrapping(self, rhs: Self) -> Self;
/// Checked multiplication operation
fn mul_checked(self, rhs: Self) -> Result<Self, ArrowError>;
/// Wrapping multiplication operation
fn mul_wrapping(self, rhs: Self) -> Self;
/// Checked division operation
fn div_checked(self, rhs: Self) -> Result<Self, ArrowError>;
/// Wrapping division operation
fn div_wrapping(self, rhs: Self) -> Self;
/// Checked remainder operation
fn mod_checked(self, rhs: Self) -> Result<Self, ArrowError>;
/// Wrapping remainder operation
fn mod_wrapping(self, rhs: Self) -> Self;
/// Checked negation operation
fn neg_checked(self) -> Result<Self, ArrowError>;
/// Wrapping negation operation
fn neg_wrapping(self) -> Self;
/// Checked exponentiation operation
fn pow_checked(self, exp: u32) -> Result<Self, ArrowError>;
/// Wrapping exponentiation operation
fn pow_wrapping(self, exp: u32) -> Self;
/// Returns true if zero else false
fn is_zero(self) -> bool;
/// Compare operation
fn compare(self, rhs: Self) -> Ordering;
/// Equality operation
fn is_eq(self, rhs: Self) -> bool;
/// Not equal operation
#[inline]
fn is_ne(self, rhs: Self) -> bool {
!self.is_eq(rhs)
}
/// Less than operation
#[inline]
fn is_lt(self, rhs: Self) -> bool {
self.compare(rhs).is_lt()
}
/// Less than equals operation
#[inline]
fn is_le(self, rhs: Self) -> bool {
self.compare(rhs).is_le()
}
/// Greater than operation
#[inline]
fn is_gt(self, rhs: Self) -> bool {
self.compare(rhs).is_gt()
}
/// Greater than equals operation
#[inline]
fn is_ge(self, rhs: Self) -> bool {
self.compare(rhs).is_ge()
}
}
macro_rules! native_type_op {
($t:tt) => {
native_type_op!($t, 0, 1);
};
($t:tt, $zero:expr, $one: expr) => {
native_type_op!($t, $zero, $one, $t::MIN, $t::MAX);
};
($t:tt, $zero:expr, $one: expr, $min: expr, $max: expr) => {
impl ArrowNativeTypeOp for $t {
const ZERO: Self = $zero;
const ONE: Self = $one;
const MIN_TOTAL_ORDER: Self = $min;
const MAX_TOTAL_ORDER: Self = $max;
#[inline]
fn add_checked(self, rhs: Self) -> Result<Self, ArrowError> {
self.checked_add(rhs).ok_or_else(|| {
ArrowError::ArithmeticOverflow(format!(
"Overflow happened on: {:?} + {:?}",
self, rhs
))
})
}
#[inline]
fn add_wrapping(self, rhs: Self) -> Self {
self.wrapping_add(rhs)
}
#[inline]
fn sub_checked(self, rhs: Self) -> Result<Self, ArrowError> {
self.checked_sub(rhs).ok_or_else(|| {
ArrowError::ArithmeticOverflow(format!(
"Overflow happened on: {:?} - {:?}",
self, rhs
))
})
}
#[inline]
fn sub_wrapping(self, rhs: Self) -> Self {
self.wrapping_sub(rhs)
}
#[inline]
fn mul_checked(self, rhs: Self) -> Result<Self, ArrowError> {
self.checked_mul(rhs).ok_or_else(|| {
ArrowError::ArithmeticOverflow(format!(
"Overflow happened on: {:?} * {:?}",
self, rhs
))
})
}
#[inline]
fn mul_wrapping(self, rhs: Self) -> Self {
self.wrapping_mul(rhs)
}
#[inline]
fn div_checked(self, rhs: Self) -> Result<Self, ArrowError> {
if rhs.is_zero() {
Err(ArrowError::DivideByZero)
} else {
self.checked_div(rhs).ok_or_else(|| {
ArrowError::ArithmeticOverflow(format!(
"Overflow happened on: {:?} / {:?}",
self, rhs
))
})
}
}
#[inline]
fn div_wrapping(self, rhs: Self) -> Self {
self.wrapping_div(rhs)
}
#[inline]
fn mod_checked(self, rhs: Self) -> Result<Self, ArrowError> {
if rhs.is_zero() {
Err(ArrowError::DivideByZero)
} else {
self.checked_rem(rhs).ok_or_else(|| {
ArrowError::ArithmeticOverflow(format!(
"Overflow happened on: {:?} % {:?}",
self, rhs
))
})
}
}
#[inline]
fn mod_wrapping(self, rhs: Self) -> Self {
self.wrapping_rem(rhs)
}
#[inline]
fn neg_checked(self) -> Result<Self, ArrowError> {
self.checked_neg().ok_or_else(|| {
ArrowError::ArithmeticOverflow(format!("Overflow happened on: - {:?}", self))
})
}
#[inline]
fn pow_checked(self, exp: u32) -> Result<Self, ArrowError> {
self.checked_pow(exp).ok_or_else(|| {
ArrowError::ArithmeticOverflow(format!(
"Overflow happened on: {:?} ^ {exp:?}",
self
))
})
}
#[inline]
fn pow_wrapping(self, exp: u32) -> Self {
self.wrapping_pow(exp)
}
#[inline]
fn neg_wrapping(self) -> Self {
self.wrapping_neg()
}
#[inline]
fn is_zero(self) -> bool {
self == Self::ZERO
}
#[inline]
fn compare(self, rhs: Self) -> Ordering {
self.cmp(&rhs)
}
#[inline]
fn is_eq(self, rhs: Self) -> bool {
self == rhs
}
}
};
}
native_type_op!(i8);
native_type_op!(i16);
native_type_op!(i32);
native_type_op!(i64);
native_type_op!(i128);
native_type_op!(u8);
native_type_op!(u16);
native_type_op!(u32);
native_type_op!(u64);
native_type_op!(i256, i256::ZERO, i256::ONE);
native_type_op!(IntervalDayTime, IntervalDayTime::ZERO, IntervalDayTime::ONE);
native_type_op!(
IntervalMonthDayNano,
IntervalMonthDayNano::ZERO,
IntervalMonthDayNano::ONE
);
macro_rules! native_type_float_op {
($t:tt, $zero:expr, $one:expr, $min:expr, $max:expr) => {
impl ArrowNativeTypeOp for $t {
const ZERO: Self = $zero;
const ONE: Self = $one;
const MIN_TOTAL_ORDER: Self = $min;
const MAX_TOTAL_ORDER: Self = $max;
#[inline]
fn add_checked(self, rhs: Self) -> Result<Self, ArrowError> {
Ok(self + rhs)
}
#[inline]
fn add_wrapping(self, rhs: Self) -> Self {
self + rhs
}
#[inline]
fn sub_checked(self, rhs: Self) -> Result<Self, ArrowError> {
Ok(self - rhs)
}
#[inline]
fn sub_wrapping(self, rhs: Self) -> Self {
self - rhs
}
#[inline]
fn mul_checked(self, rhs: Self) -> Result<Self, ArrowError> {
Ok(self * rhs)
}
#[inline]
fn mul_wrapping(self, rhs: Self) -> Self {
self * rhs
}
#[inline]
fn div_checked(self, rhs: Self) -> Result<Self, ArrowError> {
if rhs.is_zero() {
Err(ArrowError::DivideByZero)
} else {
Ok(self / rhs)
}
}
#[inline]
fn div_wrapping(self, rhs: Self) -> Self {
self / rhs
}
#[inline]
fn mod_checked(self, rhs: Self) -> Result<Self, ArrowError> {
if rhs.is_zero() {
Err(ArrowError::DivideByZero)
} else {
Ok(self % rhs)
}
}
#[inline]
fn mod_wrapping(self, rhs: Self) -> Self {
self % rhs
}
#[inline]
fn neg_checked(self) -> Result<Self, ArrowError> {
Ok(-self)
}
#[inline]
fn neg_wrapping(self) -> Self {
-self
}
#[inline]
fn pow_checked(self, exp: u32) -> Result<Self, ArrowError> {
Ok(self.powi(exp as i32))
}
#[inline]
fn pow_wrapping(self, exp: u32) -> Self {
self.powi(exp as i32)
}
#[inline]
fn is_zero(self) -> bool {
self == $zero
}
#[inline]
fn compare(self, rhs: Self) -> Ordering {
<$t>::total_cmp(&self, &rhs)
}
#[inline]
fn is_eq(self, rhs: Self) -> bool {
// Equivalent to `self.total_cmp(&rhs).is_eq()`
// but LLVM isn't able to realise this is bitwise equality
// https://rust.godbolt.org/z/347nWGxoW
self.to_bits() == rhs.to_bits()
}
}
};
}
// the smallest/largest bit patterns for floating point numbers are NaN, but differ from the canonical NAN constants.
// See test_float_total_order_min_max for details.
native_type_float_op!(
f16,
f16::ZERO,
f16::ONE,
f16::from_bits(-1 as _),
f16::from_bits(i16::MAX as _)
);
// from_bits is not yet stable as const fn, see https://github.com/rust-lang/rust/issues/72447
native_type_float_op!(
f32,
0.,
1.,
unsafe {
// Need to allow in clippy because
// current MSRV (Minimum Supported Rust Version) is `1.85.0` but this item is stable since `1.87.0`
#[allow(unnecessary_transmutes)]
std::mem::transmute(-1_i32)
},
unsafe {
// Need to allow in clippy because
// current MSRV (Minimum Supported Rust Version) is `1.85.0` but this item is stable since `1.87.0`
#[allow(unnecessary_transmutes)]
std::mem::transmute(i32::MAX)
}
);
native_type_float_op!(
f64,
0.,
1.,
unsafe {
// Need to allow in clippy because
// current MSRV (Minimum Supported Rust Version) is `1.85.0` but this item is stable since `1.87.0`
#[allow(unnecessary_transmutes)]
std::mem::transmute(-1_i64)
},
unsafe {
// Need to allow in clippy because
// current MSRV (Minimum Supported Rust Version) is `1.85.0` but this item is stable since `1.87.0`
#[allow(unnecessary_transmutes)]
std::mem::transmute(i64::MAX)
}
);
#[cfg(test)]
mod tests {
use super::*;
macro_rules! assert_approx_eq {
( $x: expr, $y: expr ) => {{ assert_approx_eq!($x, $y, 1.0e-4) }};
( $x: expr, $y: expr, $tol: expr ) => {{
let x_val = $x;
let y_val = $y;
let diff = f64::from((x_val - y_val).abs());
assert!(
diff <= $tol,
"{} != {} (with tolerance = {})",
x_val,
y_val,
$tol
);
}};
}
#[test]
fn test_native_type_is_zero() {
assert!(0_i8.is_zero());
assert!(0_i16.is_zero());
assert!(0_i32.is_zero());
assert!(0_i64.is_zero());
assert!(0_i128.is_zero());
assert!(i256::ZERO.is_zero());
assert!(0_u8.is_zero());
assert!(0_u16.is_zero());
assert!(0_u32.is_zero());
assert!(0_u64.is_zero());
assert!(f16::ZERO.is_zero());
assert!(0.0_f32.is_zero());
assert!(0.0_f64.is_zero());
}
#[test]
fn test_native_type_comparison() {
// is_eq
assert!(8_i8.is_eq(8_i8));
assert!(8_i16.is_eq(8_i16));
assert!(8_i32.is_eq(8_i32));
assert!(8_i64.is_eq(8_i64));
assert!(8_i128.is_eq(8_i128));
assert!(i256::from_parts(8, 0).is_eq(i256::from_parts(8, 0)));
assert!(8_u8.is_eq(8_u8));
assert!(8_u16.is_eq(8_u16));
assert!(8_u32.is_eq(8_u32));
assert!(8_u64.is_eq(8_u64));
assert!(f16::from_f32(8.0).is_eq(f16::from_f32(8.0)));
assert!(8.0_f32.is_eq(8.0_f32));
assert!(8.0_f64.is_eq(8.0_f64));
// is_ne
assert!(8_i8.is_ne(1_i8));
assert!(8_i16.is_ne(1_i16));
assert!(8_i32.is_ne(1_i32));
assert!(8_i64.is_ne(1_i64));
assert!(8_i128.is_ne(1_i128));
assert!(i256::from_parts(8, 0).is_ne(i256::from_parts(1, 0)));
assert!(8_u8.is_ne(1_u8));
assert!(8_u16.is_ne(1_u16));
assert!(8_u32.is_ne(1_u32));
assert!(8_u64.is_ne(1_u64));
assert!(f16::from_f32(8.0).is_ne(f16::from_f32(1.0)));
assert!(8.0_f32.is_ne(1.0_f32));
assert!(8.0_f64.is_ne(1.0_f64));
// is_lt
assert!(8_i8.is_lt(10_i8));
assert!(8_i16.is_lt(10_i16));
assert!(8_i32.is_lt(10_i32));
assert!(8_i64.is_lt(10_i64));
assert!(8_i128.is_lt(10_i128));
assert!(i256::from_parts(8, 0).is_lt(i256::from_parts(10, 0)));
assert!(8_u8.is_lt(10_u8));
assert!(8_u16.is_lt(10_u16));
assert!(8_u32.is_lt(10_u32));
assert!(8_u64.is_lt(10_u64));
assert!(f16::from_f32(8.0).is_lt(f16::from_f32(10.0)));
assert!(8.0_f32.is_lt(10.0_f32));
assert!(8.0_f64.is_lt(10.0_f64));
// is_gt
assert!(8_i8.is_gt(1_i8));
assert!(8_i16.is_gt(1_i16));
assert!(8_i32.is_gt(1_i32));
assert!(8_i64.is_gt(1_i64));
assert!(8_i128.is_gt(1_i128));
assert!(i256::from_parts(8, 0).is_gt(i256::from_parts(1, 0)));
assert!(8_u8.is_gt(1_u8));
assert!(8_u16.is_gt(1_u16));
assert!(8_u32.is_gt(1_u32));
assert!(8_u64.is_gt(1_u64));
assert!(f16::from_f32(8.0).is_gt(f16::from_f32(1.0)));
assert!(8.0_f32.is_gt(1.0_f32));
assert!(8.0_f64.is_gt(1.0_f64));
}
#[test]
fn test_native_type_add() {
// add_wrapping
assert_eq!(8_i8.add_wrapping(2_i8), 10_i8);
assert_eq!(8_i16.add_wrapping(2_i16), 10_i16);
assert_eq!(8_i32.add_wrapping(2_i32), 10_i32);
assert_eq!(8_i64.add_wrapping(2_i64), 10_i64);
assert_eq!(8_i128.add_wrapping(2_i128), 10_i128);
assert_eq!(
i256::from_parts(8, 0).add_wrapping(i256::from_parts(2, 0)),
i256::from_parts(10, 0)
);
assert_eq!(8_u8.add_wrapping(2_u8), 10_u8);
assert_eq!(8_u16.add_wrapping(2_u16), 10_u16);
assert_eq!(8_u32.add_wrapping(2_u32), 10_u32);
assert_eq!(8_u64.add_wrapping(2_u64), 10_u64);
assert_eq!(
f16::from_f32(8.0).add_wrapping(f16::from_f32(2.0)),
f16::from_f32(10.0)
);
assert_eq!(8.0_f32.add_wrapping(2.0_f32), 10_f32);
assert_eq!(8.0_f64.add_wrapping(2.0_f64), 10_f64);
// add_checked
assert_eq!(8_i8.add_checked(2_i8).unwrap(), 10_i8);
assert_eq!(8_i16.add_checked(2_i16).unwrap(), 10_i16);
assert_eq!(8_i32.add_checked(2_i32).unwrap(), 10_i32);
assert_eq!(8_i64.add_checked(2_i64).unwrap(), 10_i64);
assert_eq!(8_i128.add_checked(2_i128).unwrap(), 10_i128);
assert_eq!(
i256::from_parts(8, 0)
.add_checked(i256::from_parts(2, 0))
.unwrap(),
i256::from_parts(10, 0)
);
assert_eq!(8_u8.add_checked(2_u8).unwrap(), 10_u8);
assert_eq!(8_u16.add_checked(2_u16).unwrap(), 10_u16);
assert_eq!(8_u32.add_checked(2_u32).unwrap(), 10_u32);
assert_eq!(8_u64.add_checked(2_u64).unwrap(), 10_u64);
assert_eq!(
f16::from_f32(8.0).add_checked(f16::from_f32(2.0)).unwrap(),
f16::from_f32(10.0)
);
assert_eq!(8.0_f32.add_checked(2.0_f32).unwrap(), 10_f32);
assert_eq!(8.0_f64.add_checked(2.0_f64).unwrap(), 10_f64);
}
#[test]
fn test_native_type_sub() {
// sub_wrapping
assert_eq!(8_i8.sub_wrapping(2_i8), 6_i8);
assert_eq!(8_i16.sub_wrapping(2_i16), 6_i16);
assert_eq!(8_i32.sub_wrapping(2_i32), 6_i32);
assert_eq!(8_i64.sub_wrapping(2_i64), 6_i64);
assert_eq!(8_i128.sub_wrapping(2_i128), 6_i128);
assert_eq!(
i256::from_parts(8, 0).sub_wrapping(i256::from_parts(2, 0)),
i256::from_parts(6, 0)
);
assert_eq!(8_u8.sub_wrapping(2_u8), 6_u8);
assert_eq!(8_u16.sub_wrapping(2_u16), 6_u16);
assert_eq!(8_u32.sub_wrapping(2_u32), 6_u32);
assert_eq!(8_u64.sub_wrapping(2_u64), 6_u64);
assert_eq!(
f16::from_f32(8.0).sub_wrapping(f16::from_f32(2.0)),
f16::from_f32(6.0)
);
assert_eq!(8.0_f32.sub_wrapping(2.0_f32), 6_f32);
assert_eq!(8.0_f64.sub_wrapping(2.0_f64), 6_f64);
// sub_checked
assert_eq!(8_i8.sub_checked(2_i8).unwrap(), 6_i8);
assert_eq!(8_i16.sub_checked(2_i16).unwrap(), 6_i16);
assert_eq!(8_i32.sub_checked(2_i32).unwrap(), 6_i32);
assert_eq!(8_i64.sub_checked(2_i64).unwrap(), 6_i64);
assert_eq!(8_i128.sub_checked(2_i128).unwrap(), 6_i128);
assert_eq!(
i256::from_parts(8, 0)
.sub_checked(i256::from_parts(2, 0))
.unwrap(),
i256::from_parts(6, 0)
);
assert_eq!(8_u8.sub_checked(2_u8).unwrap(), 6_u8);
assert_eq!(8_u16.sub_checked(2_u16).unwrap(), 6_u16);
assert_eq!(8_u32.sub_checked(2_u32).unwrap(), 6_u32);
assert_eq!(8_u64.sub_checked(2_u64).unwrap(), 6_u64);
assert_eq!(
f16::from_f32(8.0).sub_checked(f16::from_f32(2.0)).unwrap(),
f16::from_f32(6.0)
);
assert_eq!(8.0_f32.sub_checked(2.0_f32).unwrap(), 6_f32);
assert_eq!(8.0_f64.sub_checked(2.0_f64).unwrap(), 6_f64);
}
#[test]
fn test_native_type_mul() {
// mul_wrapping
assert_eq!(8_i8.mul_wrapping(2_i8), 16_i8);
assert_eq!(8_i16.mul_wrapping(2_i16), 16_i16);
assert_eq!(8_i32.mul_wrapping(2_i32), 16_i32);
assert_eq!(8_i64.mul_wrapping(2_i64), 16_i64);
assert_eq!(8_i128.mul_wrapping(2_i128), 16_i128);
assert_eq!(
i256::from_parts(8, 0).mul_wrapping(i256::from_parts(2, 0)),
i256::from_parts(16, 0)
);
assert_eq!(8_u8.mul_wrapping(2_u8), 16_u8);
assert_eq!(8_u16.mul_wrapping(2_u16), 16_u16);
assert_eq!(8_u32.mul_wrapping(2_u32), 16_u32);
assert_eq!(8_u64.mul_wrapping(2_u64), 16_u64);
assert_eq!(
f16::from_f32(8.0).mul_wrapping(f16::from_f32(2.0)),
f16::from_f32(16.0)
);
assert_eq!(8.0_f32.mul_wrapping(2.0_f32), 16_f32);
assert_eq!(8.0_f64.mul_wrapping(2.0_f64), 16_f64);
// mul_checked
assert_eq!(8_i8.mul_checked(2_i8).unwrap(), 16_i8);
assert_eq!(8_i16.mul_checked(2_i16).unwrap(), 16_i16);
assert_eq!(8_i32.mul_checked(2_i32).unwrap(), 16_i32);
assert_eq!(8_i64.mul_checked(2_i64).unwrap(), 16_i64);
assert_eq!(8_i128.mul_checked(2_i128).unwrap(), 16_i128);
assert_eq!(
i256::from_parts(8, 0)
.mul_checked(i256::from_parts(2, 0))
.unwrap(),
i256::from_parts(16, 0)
);
assert_eq!(8_u8.mul_checked(2_u8).unwrap(), 16_u8);
assert_eq!(8_u16.mul_checked(2_u16).unwrap(), 16_u16);
assert_eq!(8_u32.mul_checked(2_u32).unwrap(), 16_u32);
assert_eq!(8_u64.mul_checked(2_u64).unwrap(), 16_u64);
assert_eq!(
f16::from_f32(8.0).mul_checked(f16::from_f32(2.0)).unwrap(),
f16::from_f32(16.0)
);
assert_eq!(8.0_f32.mul_checked(2.0_f32).unwrap(), 16_f32);
assert_eq!(8.0_f64.mul_checked(2.0_f64).unwrap(), 16_f64);
}
#[test]
fn test_native_type_div() {
// div_wrapping
assert_eq!(8_i8.div_wrapping(2_i8), 4_i8);
assert_eq!(8_i16.div_wrapping(2_i16), 4_i16);
assert_eq!(8_i32.div_wrapping(2_i32), 4_i32);
assert_eq!(8_i64.div_wrapping(2_i64), 4_i64);
assert_eq!(8_i128.div_wrapping(2_i128), 4_i128);
assert_eq!(
i256::from_parts(8, 0).div_wrapping(i256::from_parts(2, 0)),
i256::from_parts(4, 0)
);
assert_eq!(8_u8.div_wrapping(2_u8), 4_u8);
assert_eq!(8_u16.div_wrapping(2_u16), 4_u16);
assert_eq!(8_u32.div_wrapping(2_u32), 4_u32);
assert_eq!(8_u64.div_wrapping(2_u64), 4_u64);
assert_eq!(
f16::from_f32(8.0).div_wrapping(f16::from_f32(2.0)),
f16::from_f32(4.0)
);
assert_eq!(8.0_f32.div_wrapping(2.0_f32), 4_f32);
assert_eq!(8.0_f64.div_wrapping(2.0_f64), 4_f64);
// div_checked
assert_eq!(8_i8.div_checked(2_i8).unwrap(), 4_i8);
assert_eq!(8_i16.div_checked(2_i16).unwrap(), 4_i16);
assert_eq!(8_i32.div_checked(2_i32).unwrap(), 4_i32);
assert_eq!(8_i64.div_checked(2_i64).unwrap(), 4_i64);
assert_eq!(8_i128.div_checked(2_i128).unwrap(), 4_i128);
assert_eq!(
i256::from_parts(8, 0)
.div_checked(i256::from_parts(2, 0))
.unwrap(),
i256::from_parts(4, 0)
);
assert_eq!(8_u8.div_checked(2_u8).unwrap(), 4_u8);
assert_eq!(8_u16.div_checked(2_u16).unwrap(), 4_u16);
assert_eq!(8_u32.div_checked(2_u32).unwrap(), 4_u32);
assert_eq!(8_u64.div_checked(2_u64).unwrap(), 4_u64);
assert_eq!(
f16::from_f32(8.0).div_checked(f16::from_f32(2.0)).unwrap(),
f16::from_f32(4.0)
);
assert_eq!(8.0_f32.div_checked(2.0_f32).unwrap(), 4_f32);
assert_eq!(8.0_f64.div_checked(2.0_f64).unwrap(), 4_f64);
}
#[test]
fn test_native_type_mod() {
// mod_wrapping
assert_eq!(9_i8.mod_wrapping(2_i8), 1_i8);
assert_eq!(9_i16.mod_wrapping(2_i16), 1_i16);
assert_eq!(9_i32.mod_wrapping(2_i32), 1_i32);
assert_eq!(9_i64.mod_wrapping(2_i64), 1_i64);
assert_eq!(9_i128.mod_wrapping(2_i128), 1_i128);
assert_eq!(
i256::from_parts(9, 0).mod_wrapping(i256::from_parts(2, 0)),
i256::from_parts(1, 0)
);
assert_eq!(9_u8.mod_wrapping(2_u8), 1_u8);
assert_eq!(9_u16.mod_wrapping(2_u16), 1_u16);
assert_eq!(9_u32.mod_wrapping(2_u32), 1_u32);
assert_eq!(9_u64.mod_wrapping(2_u64), 1_u64);
assert_eq!(
f16::from_f32(9.0).mod_wrapping(f16::from_f32(2.0)),
f16::from_f32(1.0)
);
assert_eq!(9.0_f32.mod_wrapping(2.0_f32), 1_f32);
assert_eq!(9.0_f64.mod_wrapping(2.0_f64), 1_f64);
// mod_checked
assert_eq!(9_i8.mod_checked(2_i8).unwrap(), 1_i8);
assert_eq!(9_i16.mod_checked(2_i16).unwrap(), 1_i16);
assert_eq!(9_i32.mod_checked(2_i32).unwrap(), 1_i32);
assert_eq!(9_i64.mod_checked(2_i64).unwrap(), 1_i64);
assert_eq!(9_i128.mod_checked(2_i128).unwrap(), 1_i128);
assert_eq!(
i256::from_parts(9, 0)
.mod_checked(i256::from_parts(2, 0))
.unwrap(),
i256::from_parts(1, 0)
);
assert_eq!(9_u8.mod_checked(2_u8).unwrap(), 1_u8);
assert_eq!(9_u16.mod_checked(2_u16).unwrap(), 1_u16);
assert_eq!(9_u32.mod_checked(2_u32).unwrap(), 1_u32);
assert_eq!(9_u64.mod_checked(2_u64).unwrap(), 1_u64);
assert_eq!(
f16::from_f32(9.0).mod_checked(f16::from_f32(2.0)).unwrap(),
f16::from_f32(1.0)
);
assert_eq!(9.0_f32.mod_checked(2.0_f32).unwrap(), 1_f32);
assert_eq!(9.0_f64.mod_checked(2.0_f64).unwrap(), 1_f64);
}
#[test]
fn test_native_type_neg() {
// neg_wrapping
assert_eq!(8_i8.neg_wrapping(), -8_i8);
assert_eq!(8_i16.neg_wrapping(), -8_i16);
assert_eq!(8_i32.neg_wrapping(), -8_i32);
assert_eq!(8_i64.neg_wrapping(), -8_i64);
assert_eq!(8_i128.neg_wrapping(), -8_i128);
assert_eq!(i256::from_parts(8, 0).neg_wrapping(), i256::from_i128(-8));
assert_eq!(8_u8.neg_wrapping(), u8::MAX - 7_u8);
assert_eq!(8_u16.neg_wrapping(), u16::MAX - 7_u16);
assert_eq!(8_u32.neg_wrapping(), u32::MAX - 7_u32);
assert_eq!(8_u64.neg_wrapping(), u64::MAX - 7_u64);
assert_eq!(f16::from_f32(8.0).neg_wrapping(), f16::from_f32(-8.0));
assert_eq!(8.0_f32.neg_wrapping(), -8_f32);
assert_eq!(8.0_f64.neg_wrapping(), -8_f64);
// neg_checked
assert_eq!(8_i8.neg_checked().unwrap(), -8_i8);
assert_eq!(8_i16.neg_checked().unwrap(), -8_i16);
assert_eq!(8_i32.neg_checked().unwrap(), -8_i32);
assert_eq!(8_i64.neg_checked().unwrap(), -8_i64);
assert_eq!(8_i128.neg_checked().unwrap(), -8_i128);
assert_eq!(
i256::from_parts(8, 0).neg_checked().unwrap(),
i256::from_i128(-8)
);
assert!(8_u8.neg_checked().is_err());
assert!(8_u16.neg_checked().is_err());
assert!(8_u32.neg_checked().is_err());
assert!(8_u64.neg_checked().is_err());
assert_eq!(
f16::from_f32(8.0).neg_checked().unwrap(),
f16::from_f32(-8.0)
);
assert_eq!(8.0_f32.neg_checked().unwrap(), -8_f32);
assert_eq!(8.0_f64.neg_checked().unwrap(), -8_f64);
}
#[test]
fn test_native_type_pow() {
// pow_wrapping
assert_eq!(8_i8.pow_wrapping(2_u32), 64_i8);
assert_eq!(8_i16.pow_wrapping(2_u32), 64_i16);
assert_eq!(8_i32.pow_wrapping(2_u32), 64_i32);
assert_eq!(8_i64.pow_wrapping(2_u32), 64_i64);
assert_eq!(8_i128.pow_wrapping(2_u32), 64_i128);
assert_eq!(
i256::from_parts(8, 0).pow_wrapping(2_u32),
i256::from_parts(64, 0)
);
assert_eq!(8_u8.pow_wrapping(2_u32), 64_u8);
assert_eq!(8_u16.pow_wrapping(2_u32), 64_u16);
assert_eq!(8_u32.pow_wrapping(2_u32), 64_u32);
assert_eq!(8_u64.pow_wrapping(2_u32), 64_u64);
assert_approx_eq!(f16::from_f32(8.0).pow_wrapping(2_u32), f16::from_f32(64.0));
assert_approx_eq!(8.0_f32.pow_wrapping(2_u32), 64_f32);
assert_approx_eq!(8.0_f64.pow_wrapping(2_u32), 64_f64);
// pow_checked
assert_eq!(8_i8.pow_checked(2_u32).unwrap(), 64_i8);
assert_eq!(8_i16.pow_checked(2_u32).unwrap(), 64_i16);
assert_eq!(8_i32.pow_checked(2_u32).unwrap(), 64_i32);
assert_eq!(8_i64.pow_checked(2_u32).unwrap(), 64_i64);
assert_eq!(8_i128.pow_checked(2_u32).unwrap(), 64_i128);
assert_eq!(
i256::from_parts(8, 0).pow_checked(2_u32).unwrap(),
i256::from_parts(64, 0)
);
assert_eq!(8_u8.pow_checked(2_u32).unwrap(), 64_u8);
assert_eq!(8_u16.pow_checked(2_u32).unwrap(), 64_u16);
assert_eq!(8_u32.pow_checked(2_u32).unwrap(), 64_u32);
assert_eq!(8_u64.pow_checked(2_u32).unwrap(), 64_u64);
assert_approx_eq!(
f16::from_f32(8.0).pow_checked(2_u32).unwrap(),
f16::from_f32(64.0)
);
assert_approx_eq!(8.0_f32.pow_checked(2_u32).unwrap(), 64_f32);
assert_approx_eq!(8.0_f64.pow_checked(2_u32).unwrap(), 64_f64);
}
#[test]
fn test_float_total_order_min_max() {
assert!(<f64 as ArrowNativeTypeOp>::MIN_TOTAL_ORDER.is_lt(f64::NEG_INFINITY));
assert!(<f64 as ArrowNativeTypeOp>::MAX_TOTAL_ORDER.is_gt(f64::INFINITY));
assert!(<f64 as ArrowNativeTypeOp>::MIN_TOTAL_ORDER.is_nan());
assert!(<f64 as ArrowNativeTypeOp>::MIN_TOTAL_ORDER.is_sign_negative());
assert!(<f64 as ArrowNativeTypeOp>::MIN_TOTAL_ORDER.is_lt(-f64::NAN));
assert!(<f64 as ArrowNativeTypeOp>::MAX_TOTAL_ORDER.is_nan());
assert!(<f64 as ArrowNativeTypeOp>::MAX_TOTAL_ORDER.is_sign_positive());
assert!(<f64 as ArrowNativeTypeOp>::MAX_TOTAL_ORDER.is_gt(f64::NAN));
assert!(<f32 as ArrowNativeTypeOp>::MIN_TOTAL_ORDER.is_lt(f32::NEG_INFINITY));
assert!(<f32 as ArrowNativeTypeOp>::MAX_TOTAL_ORDER.is_gt(f32::INFINITY));
assert!(<f32 as ArrowNativeTypeOp>::MIN_TOTAL_ORDER.is_nan());
assert!(<f32 as ArrowNativeTypeOp>::MIN_TOTAL_ORDER.is_sign_negative());
assert!(<f32 as ArrowNativeTypeOp>::MIN_TOTAL_ORDER.is_lt(-f32::NAN));
assert!(<f32 as ArrowNativeTypeOp>::MAX_TOTAL_ORDER.is_nan());
assert!(<f32 as ArrowNativeTypeOp>::MAX_TOTAL_ORDER.is_sign_positive());
assert!(<f32 as ArrowNativeTypeOp>::MAX_TOTAL_ORDER.is_gt(f32::NAN));
assert!(<f16 as ArrowNativeTypeOp>::MIN_TOTAL_ORDER.is_lt(f16::NEG_INFINITY));
assert!(<f16 as ArrowNativeTypeOp>::MAX_TOTAL_ORDER.is_gt(f16::INFINITY));
assert!(<f16 as ArrowNativeTypeOp>::MIN_TOTAL_ORDER.is_nan());
assert!(<f16 as ArrowNativeTypeOp>::MIN_TOTAL_ORDER.is_sign_negative());
assert!(<f16 as ArrowNativeTypeOp>::MIN_TOTAL_ORDER.is_lt(-f16::NAN));
assert!(<f16 as ArrowNativeTypeOp>::MAX_TOTAL_ORDER.is_nan());
assert!(<f16 as ArrowNativeTypeOp>::MAX_TOTAL_ORDER.is_sign_positive());
assert!(<f16 as ArrowNativeTypeOp>::MAX_TOTAL_ORDER.is_gt(f16::NAN));
}
}
@@ -0,0 +1,631 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::types::{ByteArrayType, GenericBinaryType};
use crate::{Array, GenericByteArray, GenericListArray, GenericStringArray, OffsetSizeTrait};
use arrow_data::ArrayData;
use arrow_schema::DataType;
/// A [`GenericByteArray`] for storing `[u8]`
pub type GenericBinaryArray<OffsetSize> = GenericByteArray<GenericBinaryType<OffsetSize>>;
impl<OffsetSize: OffsetSizeTrait> GenericBinaryArray<OffsetSize> {
/// Creates a [GenericBinaryArray] from a vector of byte slices
///
/// See also [`Self::from_iter_values`]
pub fn from_vec(v: Vec<&[u8]>) -> Self {
Self::from_iter_values(v)
}
/// Creates a [GenericBinaryArray] from a vector of Optional (null) byte slices
pub fn from_opt_vec(v: Vec<Option<&[u8]>>) -> Self {
v.into_iter().collect()
}
fn from_list(v: GenericListArray<OffsetSize>) -> Self {
let v = v.into_data();
assert_eq!(
v.child_data().len(),
1,
"BinaryArray can only be created from list array of u8 values \
(i.e. List<PrimitiveArray<u8>>)."
);
let child_data = &v.child_data()[0];
assert_eq!(
child_data.child_data().len(),
0,
"BinaryArray can only be created from list array of u8 values \
(i.e. List<PrimitiveArray<u8>>)."
);
assert_eq!(
child_data.data_type(),
&DataType::UInt8,
"BinaryArray can only be created from List<u8> arrays, mismatched data types."
);
assert_eq!(
child_data.null_count(),
0,
"The child array cannot contain null values."
);
let builder = ArrayData::builder(Self::DATA_TYPE)
.len(v.len())
.offset(v.offset())
.add_buffer(v.buffers()[0].clone())
.add_buffer(child_data.buffers()[0].slice(child_data.offset()))
.nulls(v.nulls().cloned());
let data = unsafe { builder.build_unchecked() };
Self::from(data)
}
/// Returns an iterator that returns the values of `array.value(i)` for an iterator with each element `i`
pub fn take_iter<'a>(
&'a self,
indexes: impl Iterator<Item = Option<usize>> + 'a,
) -> impl Iterator<Item = Option<&'a [u8]>> {
indexes.map(|opt_index| opt_index.map(|index| self.value(index)))
}
/// Returns an iterator that returns the values of `array.value(i)` for an iterator with each element `i`
/// # Safety
///
/// caller must ensure that the indexes in the iterator are less than the `array.len()`
pub unsafe fn take_iter_unchecked<'a>(
&'a self,
indexes: impl Iterator<Item = Option<usize>> + 'a,
) -> impl Iterator<Item = Option<&'a [u8]>> {
unsafe { indexes.map(|opt_index| opt_index.map(|index| self.value_unchecked(index))) }
}
}
impl<OffsetSize: OffsetSizeTrait> From<Vec<Option<&[u8]>>> for GenericBinaryArray<OffsetSize> {
fn from(v: Vec<Option<&[u8]>>) -> Self {
Self::from_opt_vec(v)
}
}
impl<OffsetSize: OffsetSizeTrait> From<Vec<&[u8]>> for GenericBinaryArray<OffsetSize> {
fn from(v: Vec<&[u8]>) -> Self {
Self::from_iter_values(v)
}
}
impl<T: OffsetSizeTrait> From<GenericListArray<T>> for GenericBinaryArray<T> {
fn from(v: GenericListArray<T>) -> Self {
Self::from_list(v)
}
}
impl<OffsetSize: OffsetSizeTrait> From<GenericStringArray<OffsetSize>>
for GenericBinaryArray<OffsetSize>
{
fn from(value: GenericStringArray<OffsetSize>) -> Self {
let builder = value
.into_data()
.into_builder()
.data_type(GenericBinaryType::<OffsetSize>::DATA_TYPE);
// Safety:
// A StringArray is a valid BinaryArray
Self::from(unsafe { builder.build_unchecked() })
}
}
/// A [`GenericBinaryArray`] of `[u8]` using `i32` offsets
///
/// The byte length of each element is represented by an i32.
///
/// # Examples
///
/// Create a BinaryArray from a vector of byte slices.
///
/// ```
/// use arrow_array::{Array, BinaryArray};
/// let values: Vec<&[u8]> =
/// vec![b"one", b"two", b"", b"three"];
/// let array = BinaryArray::from_vec(values);
/// assert_eq!(4, array.len());
/// assert_eq!(b"one", array.value(0));
/// assert_eq!(b"two", array.value(1));
/// assert_eq!(b"", array.value(2));
/// assert_eq!(b"three", array.value(3));
/// ```
///
/// Create a BinaryArray from a vector of Optional (null) byte slices.
///
/// ```
/// use arrow_array::{Array, BinaryArray};
/// let values: Vec<Option<&[u8]>> =
/// vec![Some(b"one"), Some(b"two"), None, Some(b""), Some(b"three")];
/// let array = BinaryArray::from_opt_vec(values);
/// assert_eq!(5, array.len());
/// assert_eq!(b"one", array.value(0));
/// assert_eq!(b"two", array.value(1));
/// assert_eq!(b"", array.value(3));
/// assert_eq!(b"three", array.value(4));
/// assert!(!array.is_null(0));
/// assert!(!array.is_null(1));
/// assert!(array.is_null(2));
/// assert!(!array.is_null(3));
/// assert!(!array.is_null(4));
/// ```
///
/// See [`GenericByteArray`] for more information and examples
pub type BinaryArray = GenericBinaryArray<i32>;
/// A [`GenericBinaryArray`] of `[u8]` using `i64` offsets
///
/// # Examples
///
/// Create a LargeBinaryArray from a vector of byte slices.
///
/// ```
/// use arrow_array::{Array, LargeBinaryArray};
/// let values: Vec<&[u8]> =
/// vec![b"one", b"two", b"", b"three"];
/// let array = LargeBinaryArray::from_vec(values);
/// assert_eq!(4, array.len());
/// assert_eq!(b"one", array.value(0));
/// assert_eq!(b"two", array.value(1));
/// assert_eq!(b"", array.value(2));
/// assert_eq!(b"three", array.value(3));
/// ```
///
/// Create a LargeBinaryArray from a vector of Optional (null) byte slices.
///
/// ```
/// use arrow_array::{Array, LargeBinaryArray};
/// let values: Vec<Option<&[u8]>> =
/// vec![Some(b"one"), Some(b"two"), None, Some(b""), Some(b"three")];
/// let array = LargeBinaryArray::from_opt_vec(values);
/// assert_eq!(5, array.len());
/// assert_eq!(b"one", array.value(0));
/// assert_eq!(b"two", array.value(1));
/// assert_eq!(b"", array.value(3));
/// assert_eq!(b"three", array.value(4));
/// assert!(!array.is_null(0));
/// assert!(!array.is_null(1));
/// assert!(array.is_null(2));
/// assert!(!array.is_null(3));
/// assert!(!array.is_null(4));
/// ```
///
/// See [`GenericByteArray`] for more information and examples
pub type LargeBinaryArray = GenericBinaryArray<i64>;
#[cfg(test)]
mod tests {
use super::*;
use crate::{ListArray, StringArray};
use arrow_buffer::Buffer;
use arrow_schema::Field;
use std::sync::Arc;
#[test]
fn test_binary_array() {
let values: [u8; 12] = [
b'h', b'e', b'l', b'l', b'o', b'p', b'a', b'r', b'q', b'u', b'e', b't',
];
let offsets: [i32; 4] = [0, 5, 5, 12];
// Array data: ["hello", "", "parquet"]
let array_data = ArrayData::builder(DataType::Binary)
.len(3)
.add_buffer(Buffer::from_slice_ref(offsets))
.add_buffer(Buffer::from_slice_ref(values))
.build()
.unwrap();
let binary_array = BinaryArray::from(array_data);
assert_eq!(3, binary_array.len());
assert_eq!(0, binary_array.null_count());
assert_eq!([b'h', b'e', b'l', b'l', b'o'], binary_array.value(0));
assert_eq!([b'h', b'e', b'l', b'l', b'o'], unsafe {
binary_array.value_unchecked(0)
});
assert_eq!([] as [u8; 0], binary_array.value(1));
assert_eq!([] as [u8; 0], unsafe { binary_array.value_unchecked(1) });
assert_eq!(
[b'p', b'a', b'r', b'q', b'u', b'e', b't'],
binary_array.value(2)
);
assert_eq!([b'p', b'a', b'r', b'q', b'u', b'e', b't'], unsafe {
binary_array.value_unchecked(2)
});
assert_eq!(5, binary_array.value_offsets()[2]);
assert_eq!(7, binary_array.value_length(2));
for i in 0..3 {
assert!(binary_array.is_valid(i));
assert!(!binary_array.is_null(i));
}
}
#[test]
fn test_binary_array_with_offsets() {
let values: [u8; 12] = [
b'h', b'e', b'l', b'l', b'o', b'p', b'a', b'r', b'q', b'u', b'e', b't',
];
let offsets: [i32; 4] = [0, 5, 5, 12];
// Test binary array with offset
let array_data = ArrayData::builder(DataType::Binary)
.len(2)
.offset(1)
.add_buffer(Buffer::from_slice_ref(offsets))
.add_buffer(Buffer::from_slice_ref(values))
.build()
.unwrap();
let binary_array = BinaryArray::from(array_data);
assert_eq!(
[b'p', b'a', b'r', b'q', b'u', b'e', b't'],
binary_array.value(1)
);
assert_eq!(5, binary_array.value_offsets()[0]);
assert_eq!(0, binary_array.value_length(0));
assert_eq!(5, binary_array.value_offsets()[1]);
assert_eq!(7, binary_array.value_length(1));
}
#[test]
fn test_large_binary_array() {
let values: [u8; 12] = [
b'h', b'e', b'l', b'l', b'o', b'p', b'a', b'r', b'q', b'u', b'e', b't',
];
let offsets: [i64; 4] = [0, 5, 5, 12];
// Array data: ["hello", "", "parquet"]
let array_data = ArrayData::builder(DataType::LargeBinary)
.len(3)
.add_buffer(Buffer::from_slice_ref(offsets))
.add_buffer(Buffer::from_slice_ref(values))
.build()
.unwrap();
let binary_array = LargeBinaryArray::from(array_data);
assert_eq!(3, binary_array.len());
assert_eq!(0, binary_array.null_count());
assert_eq!([b'h', b'e', b'l', b'l', b'o'], binary_array.value(0));
assert_eq!([b'h', b'e', b'l', b'l', b'o'], unsafe {
binary_array.value_unchecked(0)
});
assert_eq!([] as [u8; 0], binary_array.value(1));
assert_eq!([] as [u8; 0], unsafe { binary_array.value_unchecked(1) });
assert_eq!(
[b'p', b'a', b'r', b'q', b'u', b'e', b't'],
binary_array.value(2)
);
assert_eq!([b'p', b'a', b'r', b'q', b'u', b'e', b't'], unsafe {
binary_array.value_unchecked(2)
});
assert_eq!(5, binary_array.value_offsets()[2]);
assert_eq!(7, binary_array.value_length(2));
for i in 0..3 {
assert!(binary_array.is_valid(i));
assert!(!binary_array.is_null(i));
}
}
#[test]
fn test_large_binary_array_with_offsets() {
let values: [u8; 12] = [
b'h', b'e', b'l', b'l', b'o', b'p', b'a', b'r', b'q', b'u', b'e', b't',
];
let offsets: [i64; 4] = [0, 5, 5, 12];
// Test binary array with offset
let array_data = ArrayData::builder(DataType::LargeBinary)
.len(2)
.offset(1)
.add_buffer(Buffer::from_slice_ref(offsets))
.add_buffer(Buffer::from_slice_ref(values))
.build()
.unwrap();
let binary_array = LargeBinaryArray::from(array_data);
assert_eq!(
[b'p', b'a', b'r', b'q', b'u', b'e', b't'],
binary_array.value(1)
);
assert_eq!([b'p', b'a', b'r', b'q', b'u', b'e', b't'], unsafe {
binary_array.value_unchecked(1)
});
assert_eq!(5, binary_array.value_offsets()[0]);
assert_eq!(0, binary_array.value_length(0));
assert_eq!(5, binary_array.value_offsets()[1]);
assert_eq!(7, binary_array.value_length(1));
}
fn _test_generic_binary_array_from_list_array<O: OffsetSizeTrait>() {
let values = b"helloparquet";
let child_data = ArrayData::builder(DataType::UInt8)
.len(12)
.add_buffer(Buffer::from(values))
.build()
.unwrap();
let offsets = [0, 5, 5, 12].map(|n| O::from_usize(n).unwrap());
// Array data: ["hello", "", "parquet"]
let array_data1 = ArrayData::builder(GenericBinaryArray::<O>::DATA_TYPE)
.len(3)
.add_buffer(Buffer::from_slice_ref(offsets))
.add_buffer(Buffer::from_slice_ref(values))
.build()
.unwrap();
let binary_array1 = GenericBinaryArray::<O>::from(array_data1);
let data_type = GenericListArray::<O>::DATA_TYPE_CONSTRUCTOR(Arc::new(
Field::new_list_field(DataType::UInt8, false),
));
let array_data2 = ArrayData::builder(data_type)
.len(3)
.add_buffer(Buffer::from_slice_ref(offsets))
.add_child_data(child_data)
.build()
.unwrap();
let list_array = GenericListArray::<O>::from(array_data2);
let binary_array2 = GenericBinaryArray::<O>::from(list_array);
assert_eq!(binary_array1.len(), binary_array2.len());
assert_eq!(binary_array1.null_count(), binary_array2.null_count());
assert_eq!(binary_array1.value_offsets(), binary_array2.value_offsets());
for i in 0..binary_array1.len() {
assert_eq!(binary_array1.value(i), binary_array2.value(i));
assert_eq!(binary_array1.value(i), unsafe {
binary_array2.value_unchecked(i)
});
assert_eq!(binary_array1.value_length(i), binary_array2.value_length(i));
}
}
#[test]
fn test_binary_array_from_list_array() {
_test_generic_binary_array_from_list_array::<i32>();
}
#[test]
fn test_large_binary_array_from_list_array() {
_test_generic_binary_array_from_list_array::<i64>();
}
fn _test_generic_binary_array_from_list_array_with_offset<O: OffsetSizeTrait>() {
let values = b"HelloArrowAndParquet";
// b"ArrowAndParquet"
let child_data = ArrayData::builder(DataType::UInt8)
.len(15)
.offset(5)
.add_buffer(Buffer::from(values))
.build()
.unwrap();
let offsets = [0, 5, 8, 15].map(|n| O::from_usize(n).unwrap());
let null_buffer = Buffer::from_slice_ref([0b101]);
let data_type = GenericListArray::<O>::DATA_TYPE_CONSTRUCTOR(Arc::new(
Field::new_list_field(DataType::UInt8, false),
));
// [None, Some(b"Parquet")]
let array_data = ArrayData::builder(data_type)
.len(2)
.offset(1)
.add_buffer(Buffer::from_slice_ref(offsets))
.null_bit_buffer(Some(null_buffer))
.add_child_data(child_data)
.build()
.unwrap();
let list_array = GenericListArray::<O>::from(array_data);
let binary_array = GenericBinaryArray::<O>::from(list_array);
assert_eq!(2, binary_array.len());
assert_eq!(1, binary_array.null_count());
assert!(binary_array.is_null(0));
assert!(binary_array.is_valid(1));
assert_eq!(b"Parquet", binary_array.value(1));
}
#[test]
fn test_binary_array_from_list_array_with_offset() {
_test_generic_binary_array_from_list_array_with_offset::<i32>();
}
#[test]
fn test_large_binary_array_from_list_array_with_offset() {
_test_generic_binary_array_from_list_array_with_offset::<i64>();
}
fn _test_generic_binary_array_from_list_array_with_child_nulls_failed<O: OffsetSizeTrait>() {
let values = b"HelloArrow";
let child_data = ArrayData::builder(DataType::UInt8)
.len(10)
.add_buffer(Buffer::from(values))
.null_bit_buffer(Some(Buffer::from_slice_ref([0b1010101010])))
.build()
.unwrap();
let offsets = [0, 5, 10].map(|n| O::from_usize(n).unwrap());
let data_type = GenericListArray::<O>::DATA_TYPE_CONSTRUCTOR(Arc::new(
Field::new_list_field(DataType::UInt8, true),
));
// [None, Some(b"Parquet")]
let array_data = ArrayData::builder(data_type)
.len(2)
.add_buffer(Buffer::from_slice_ref(offsets))
.add_child_data(child_data)
.build()
.unwrap();
let list_array = GenericListArray::<O>::from(array_data);
drop(GenericBinaryArray::<O>::from(list_array));
}
#[test]
#[should_panic(expected = "The child array cannot contain null values.")]
fn test_binary_array_from_list_array_with_child_nulls_failed() {
_test_generic_binary_array_from_list_array_with_child_nulls_failed::<i32>();
}
#[test]
#[should_panic(expected = "The child array cannot contain null values.")]
fn test_large_binary_array_from_list_array_with_child_nulls_failed() {
_test_generic_binary_array_from_list_array_with_child_nulls_failed::<i64>();
}
fn test_generic_binary_array_from_opt_vec<T: OffsetSizeTrait>() {
let values: Vec<Option<&[u8]>> =
vec![Some(b"one"), Some(b"two"), None, Some(b""), Some(b"three")];
let array = GenericBinaryArray::<T>::from_opt_vec(values);
assert_eq!(array.len(), 5);
assert_eq!(array.value(0), b"one");
assert_eq!(array.value(1), b"two");
assert_eq!(array.value(3), b"");
assert_eq!(array.value(4), b"three");
assert!(!array.is_null(0));
assert!(!array.is_null(1));
assert!(array.is_null(2));
assert!(!array.is_null(3));
assert!(!array.is_null(4));
}
#[test]
fn test_large_binary_array_from_opt_vec() {
test_generic_binary_array_from_opt_vec::<i64>()
}
#[test]
fn test_binary_array_from_opt_vec() {
test_generic_binary_array_from_opt_vec::<i32>()
}
#[test]
fn test_binary_array_from_unbound_iter() {
// iterator that doesn't declare (upper) size bound
let value_iter = (0..)
.scan(0usize, |pos, i| {
if *pos < 10 {
*pos += 1;
Some(Some(format!("value {i}")))
} else {
// actually returns up to 10 values
None
}
})
// limited using take()
.take(100);
let (_, upper_size_bound) = value_iter.size_hint();
// the upper bound, defined by take above, is 100
assert_eq!(upper_size_bound, Some(100));
let binary_array: BinaryArray = value_iter.collect();
// but the actual number of items in the array should be 10
assert_eq!(binary_array.len(), 10);
}
#[test]
#[should_panic(
expected = "BinaryArray can only be created from List<u8> arrays, mismatched data types."
)]
fn test_binary_array_from_incorrect_list_array() {
let values: [u32; 12] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11];
let values_data = ArrayData::builder(DataType::UInt32)
.len(12)
.add_buffer(Buffer::from_slice_ref(values))
.build()
.unwrap();
let offsets: [i32; 4] = [0, 5, 5, 12];
let data_type = DataType::List(Arc::new(Field::new_list_field(DataType::UInt32, false)));
let array_data = ArrayData::builder(data_type)
.len(3)
.add_buffer(Buffer::from_slice_ref(offsets))
.add_child_data(values_data)
.build()
.unwrap();
let list_array = ListArray::from(array_data);
drop(BinaryArray::from(list_array));
}
#[test]
#[should_panic(
expected = "Trying to access an element at index 4 from a BinaryArray of length 3"
)]
fn test_binary_array_get_value_index_out_of_bound() {
let values: [u8; 12] = [104, 101, 108, 108, 111, 112, 97, 114, 113, 117, 101, 116];
let offsets: [i32; 4] = [0, 5, 5, 12];
let array_data = ArrayData::builder(DataType::Binary)
.len(3)
.add_buffer(Buffer::from_slice_ref(offsets))
.add_buffer(Buffer::from_slice_ref(values))
.build()
.unwrap();
let binary_array = BinaryArray::from(array_data);
binary_array.value(4);
}
#[test]
#[should_panic(expected = "LargeBinaryArray expects DataType::LargeBinary")]
fn test_binary_array_validation() {
let array = BinaryArray::from_iter_values([&[1, 2]]);
let _ = LargeBinaryArray::from(array.into_data());
}
#[test]
fn test_binary_array_all_null() {
let data = vec![None];
let array = BinaryArray::from(data);
array
.into_data()
.validate_full()
.expect("All null array has valid array data");
}
#[test]
fn test_large_binary_array_all_null() {
let data = vec![None];
let array = LargeBinaryArray::from(data);
array
.into_data()
.validate_full()
.expect("All null array has valid array data");
}
#[test]
fn test_empty_offsets() {
let string = BinaryArray::from(
ArrayData::builder(DataType::Binary)
.buffers(vec![Buffer::from(&[]), Buffer::from(&[])])
.build()
.unwrap(),
);
assert_eq!(string.value_offsets(), &[0]);
let string = LargeBinaryArray::from(
ArrayData::builder(DataType::LargeBinary)
.buffers(vec![Buffer::from(&[]), Buffer::from(&[])])
.build()
.unwrap(),
);
assert_eq!(string.len(), 0);
assert_eq!(string.value_offsets(), &[0]);
}
#[test]
fn test_to_from_string() {
let s = StringArray::from_iter_values(["a", "b", "c", "d"]);
let b = BinaryArray::from(s.clone());
let sa = StringArray::from(b); // Performs UTF-8 validation again
assert_eq!(s, sa);
}
}
@@ -0,0 +1,861 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::array::print_long_array;
use crate::builder::BooleanBuilder;
use crate::iterator::BooleanIter;
use crate::{Array, ArrayAccessor, ArrayRef, Scalar};
use arrow_buffer::{BooleanBuffer, Buffer, MutableBuffer, NullBuffer, bit_util};
use arrow_data::{ArrayData, ArrayDataBuilder};
use arrow_schema::DataType;
use std::any::Any;
use std::sync::Arc;
/// An array of [boolean values](https://arrow.apache.org/docs/format/Columnar.html#fixed-size-primitive-layout)
///
/// # Example: From a Vec
///
/// ```
/// # use arrow_array::{Array, BooleanArray};
/// let arr: BooleanArray = vec![true, true, false].into();
/// ```
///
/// # Example: From an optional Vec
///
/// ```
/// # use arrow_array::{Array, BooleanArray};
/// let arr: BooleanArray = vec![Some(true), None, Some(false)].into();
/// ```
///
/// # Example: From an iterator
///
/// ```
/// # use arrow_array::{Array, BooleanArray};
/// let arr: BooleanArray = (0..5).map(|x| (x % 2 == 0).then(|| x % 3 == 0)).collect();
/// let values: Vec<_> = arr.iter().collect();
/// assert_eq!(&values, &[Some(true), None, Some(false), None, Some(false)])
/// ```
///
/// # Example: Using Builder
///
/// ```
/// # use arrow_array::Array;
/// # use arrow_array::builder::BooleanBuilder;
/// let mut builder = BooleanBuilder::new();
/// builder.append_value(true);
/// builder.append_null();
/// builder.append_value(false);
/// let array = builder.finish();
/// let values: Vec<_> = array.iter().collect();
/// assert_eq!(&values, &[Some(true), None, Some(false)])
/// ```
///
#[derive(Clone)]
pub struct BooleanArray {
values: BooleanBuffer,
nulls: Option<NullBuffer>,
}
impl std::fmt::Debug for BooleanArray {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(f, "BooleanArray\n[\n")?;
print_long_array(self, f, |array, index, f| {
std::fmt::Debug::fmt(&array.value(index), f)
})?;
write!(f, "]")
}
}
impl BooleanArray {
/// Create a new [`BooleanArray`] from the provided values and nulls
///
/// # Panics
///
/// Panics if `values.len() != nulls.len()`
pub fn new(values: BooleanBuffer, nulls: Option<NullBuffer>) -> Self {
if let Some(n) = nulls.as_ref() {
assert_eq!(values.len(), n.len());
}
Self { values, nulls }
}
/// Create a new [`BooleanArray`] with length `len` consisting only of nulls
pub fn new_null(len: usize) -> Self {
Self {
values: BooleanBuffer::new_unset(len),
nulls: Some(NullBuffer::new_null(len)),
}
}
/// Create a new [`Scalar`] from `value`
pub fn new_scalar(value: bool) -> Scalar<Self> {
let values = match value {
true => BooleanBuffer::new_set(1),
false => BooleanBuffer::new_unset(1),
};
Scalar::new(Self::new(values, None))
}
/// Create a new [`BooleanArray`] from a [`Buffer`] specified by `offset` and `len`, the `offset` and `len` in bits
/// Logically convert each bit in [`Buffer`] to boolean and use it to build [`BooleanArray`].
/// using this method will make the following points self-evident:
/// * there is no `null` in the constructed [`BooleanArray`];
/// * without considering `buffer.into()`, this method is efficient because there is no need to perform pack and unpack operations on boolean;
pub fn new_from_packed(buffer: impl Into<Buffer>, offset: usize, len: usize) -> Self {
BooleanBuffer::new(buffer.into(), offset, len).into()
}
/// Create a new [`BooleanArray`] from `&[u8]`
/// This method uses `new_from_packed` and constructs a [`Buffer`] using `value`, and offset is set to 0 and len is set to `value.len() * 8`
/// using this method will make the following points self-evident:
/// * there is no `null` in the constructed [`BooleanArray`];
/// * the length of the constructed [`BooleanArray`] is always a multiple of 8;
pub fn new_from_u8(value: &[u8]) -> Self {
BooleanBuffer::new(Buffer::from(value), 0, value.len() * 8).into()
}
/// Returns the length of this array.
pub fn len(&self) -> usize {
self.values.len()
}
/// Returns whether this array is empty.
pub fn is_empty(&self) -> bool {
self.values.is_empty()
}
/// Returns a zero-copy slice of this array with the indicated offset and length.
pub fn slice(&self, offset: usize, length: usize) -> Self {
Self {
values: self.values.slice(offset, length),
nulls: self.nulls.as_ref().map(|n| n.slice(offset, length)),
}
}
/// Returns a new boolean array builder
pub fn builder(capacity: usize) -> BooleanBuilder {
BooleanBuilder::with_capacity(capacity)
}
/// Returns the underlying [`BooleanBuffer`] holding all the values of this array
pub fn values(&self) -> &BooleanBuffer {
&self.values
}
/// Returns the number of non null, true values within this array
pub fn true_count(&self) -> usize {
match self.nulls() {
Some(nulls) => {
let null_chunks = nulls.inner().bit_chunks().iter_padded();
let value_chunks = self.values().bit_chunks().iter_padded();
null_chunks
.zip(value_chunks)
.map(|(a, b)| (a & b).count_ones() as usize)
.sum()
}
None => self.values().count_set_bits(),
}
}
/// Returns the number of non null, false values within this array
pub fn false_count(&self) -> usize {
self.len() - self.null_count() - self.true_count()
}
/// Returns the boolean value at index `i`.
///
/// Note: This method does not check for nulls and the value is arbitrary
/// if [`is_null`](Self::is_null) returns true for the index.
///
/// # Safety
/// This doesn't check bounds, the caller must ensure that index < self.len()
pub unsafe fn value_unchecked(&self, i: usize) -> bool {
unsafe { self.values.value_unchecked(i) }
}
/// Returns the boolean value at index `i`.
///
/// Note: This method does not check for nulls and the value is arbitrary
/// if [`is_null`](Self::is_null) returns true for the index.
///
/// # Panics
/// Panics if index `i` is out of bounds
pub fn value(&self, i: usize) -> bool {
assert!(
i < self.len(),
"Trying to access an element at index {} from a BooleanArray of length {}",
i,
self.len()
);
// Safety:
// `i < self.len()
unsafe { self.value_unchecked(i) }
}
/// Returns an iterator that returns the values of `array.value(i)` for an iterator with each element `i`
pub fn take_iter<'a>(
&'a self,
indexes: impl Iterator<Item = Option<usize>> + 'a,
) -> impl Iterator<Item = Option<bool>> + 'a {
indexes.map(|opt_index| opt_index.map(|index| self.value(index)))
}
/// Returns an iterator that returns the values of `array.value(i)` for an iterator with each element `i`
/// # Safety
///
/// caller must ensure that the offsets in the iterator are less than the array len()
pub unsafe fn take_iter_unchecked<'a>(
&'a self,
indexes: impl Iterator<Item = Option<usize>> + 'a,
) -> impl Iterator<Item = Option<bool>> + 'a {
indexes.map(|opt_index| opt_index.map(|index| unsafe { self.value_unchecked(index) }))
}
/// Create a [`BooleanArray`] by evaluating the operation for
/// each element of the provided array
///
/// ```
/// # use arrow_array::{BooleanArray, Int32Array};
///
/// let array = Int32Array::from(vec![1, 2, 3, 4, 5]);
/// let r = BooleanArray::from_unary(&array, |x| x > 2);
/// assert_eq!(&r, &BooleanArray::from(vec![false, false, true, true, true]));
/// ```
pub fn from_unary<T: ArrayAccessor, F>(left: T, mut op: F) -> Self
where
F: FnMut(T::Item) -> bool,
{
let nulls = left.logical_nulls();
let values = BooleanBuffer::collect_bool(left.len(), |i| unsafe {
// SAFETY: i in range 0..len
op(left.value_unchecked(i))
});
Self::new(values, nulls)
}
/// Create a [`BooleanArray`] by evaluating the binary operation for
/// each element of the provided arrays
///
/// ```
/// # use arrow_array::{BooleanArray, Int32Array};
///
/// let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
/// let b = Int32Array::from(vec![1, 2, 0, 2, 5]);
/// let r = BooleanArray::from_binary(&a, &b, |a, b| a == b);
/// assert_eq!(&r, &BooleanArray::from(vec![true, true, false, false, true]));
/// ```
///
/// # Panics
///
/// This function panics if left and right are not the same length
///
pub fn from_binary<T: ArrayAccessor, S: ArrayAccessor, F>(left: T, right: S, mut op: F) -> Self
where
F: FnMut(T::Item, S::Item) -> bool,
{
assert_eq!(left.len(), right.len());
let nulls = NullBuffer::union(
left.logical_nulls().as_ref(),
right.logical_nulls().as_ref(),
);
let values = BooleanBuffer::collect_bool(left.len(), |i| unsafe {
// SAFETY: i in range 0..len
op(left.value_unchecked(i), right.value_unchecked(i))
});
Self::new(values, nulls)
}
/// Deconstruct this array into its constituent parts
pub fn into_parts(self) -> (BooleanBuffer, Option<NullBuffer>) {
(self.values, self.nulls)
}
}
/// SAFETY: Correctly implements the contract of Arrow Arrays
unsafe impl Array for BooleanArray {
fn as_any(&self) -> &dyn Any {
self
}
fn to_data(&self) -> ArrayData {
self.clone().into()
}
fn into_data(self) -> ArrayData {
self.into()
}
fn data_type(&self) -> &DataType {
&DataType::Boolean
}
fn slice(&self, offset: usize, length: usize) -> ArrayRef {
Arc::new(self.slice(offset, length))
}
fn len(&self) -> usize {
self.values.len()
}
fn is_empty(&self) -> bool {
self.values.is_empty()
}
fn shrink_to_fit(&mut self) {
self.values.shrink_to_fit();
if let Some(nulls) = &mut self.nulls {
nulls.shrink_to_fit();
}
}
fn offset(&self) -> usize {
self.values.offset()
}
fn nulls(&self) -> Option<&NullBuffer> {
self.nulls.as_ref()
}
fn logical_null_count(&self) -> usize {
self.null_count()
}
fn get_buffer_memory_size(&self) -> usize {
let mut sum = self.values.inner().capacity();
if let Some(x) = &self.nulls {
sum += x.buffer().capacity()
}
sum
}
fn get_array_memory_size(&self) -> usize {
std::mem::size_of::<Self>() + self.get_buffer_memory_size()
}
}
impl ArrayAccessor for &BooleanArray {
type Item = bool;
fn value(&self, index: usize) -> Self::Item {
BooleanArray::value(self, index)
}
unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
unsafe { BooleanArray::value_unchecked(self, index) }
}
}
impl From<Vec<bool>> for BooleanArray {
fn from(data: Vec<bool>) -> Self {
let mut mut_buf = MutableBuffer::new_null(data.len());
{
let mut_slice = mut_buf.as_slice_mut();
for (i, b) in data.iter().enumerate() {
if *b {
bit_util::set_bit(mut_slice, i);
}
}
}
let array_data = ArrayData::builder(DataType::Boolean)
.len(data.len())
.add_buffer(mut_buf.into());
let array_data = unsafe { array_data.build_unchecked() };
BooleanArray::from(array_data)
}
}
impl From<Vec<Option<bool>>> for BooleanArray {
fn from(data: Vec<Option<bool>>) -> Self {
data.iter().collect()
}
}
impl From<ArrayData> for BooleanArray {
fn from(data: ArrayData) -> Self {
assert_eq!(
data.data_type(),
&DataType::Boolean,
"BooleanArray expected ArrayData with type {} got {}",
DataType::Boolean,
data.data_type()
);
assert_eq!(
data.buffers().len(),
1,
"BooleanArray data should contain a single buffer only (values buffer)"
);
let values = BooleanBuffer::new(data.buffers()[0].clone(), data.offset(), data.len());
Self {
values,
nulls: data.nulls().cloned(),
}
}
}
impl From<BooleanArray> for ArrayData {
fn from(array: BooleanArray) -> Self {
let builder = ArrayDataBuilder::new(DataType::Boolean)
.len(array.values.len())
.offset(array.values.offset())
.nulls(array.nulls)
.buffers(vec![array.values.into_inner()]);
unsafe { builder.build_unchecked() }
}
}
impl<'a> IntoIterator for &'a BooleanArray {
type Item = Option<bool>;
type IntoIter = BooleanIter<'a>;
fn into_iter(self) -> Self::IntoIter {
BooleanIter::<'a>::new(self)
}
}
impl<'a> BooleanArray {
/// constructs a new iterator
pub fn iter(&'a self) -> BooleanIter<'a> {
BooleanIter::<'a>::new(self)
}
}
/// An optional boolean value
///
/// This struct is used as an adapter when creating `BooleanArray` from an iterator.
/// `FromIterator` for `BooleanArray` takes an iterator where the elements can be `into`
/// this struct. So once implementing `From` or `Into` trait for a type, an iterator of
/// the type can be collected to `BooleanArray`.
///
/// See also [NativeAdapter](crate::array::NativeAdapter).
#[derive(Debug)]
struct BooleanAdapter {
/// Corresponding Rust native type if available
pub native: Option<bool>,
}
impl From<bool> for BooleanAdapter {
fn from(value: bool) -> Self {
BooleanAdapter {
native: Some(value),
}
}
}
impl From<&bool> for BooleanAdapter {
fn from(value: &bool) -> Self {
BooleanAdapter {
native: Some(*value),
}
}
}
impl From<Option<bool>> for BooleanAdapter {
fn from(value: Option<bool>) -> Self {
BooleanAdapter { native: value }
}
}
impl From<&Option<bool>> for BooleanAdapter {
fn from(value: &Option<bool>) -> Self {
BooleanAdapter { native: *value }
}
}
impl<Ptr: Into<BooleanAdapter>> FromIterator<Ptr> for BooleanArray {
fn from_iter<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
let iter = iter.into_iter();
let capacity = match iter.size_hint() {
(lower, Some(upper)) if lower == upper => lower,
_ => 0,
};
let mut builder = BooleanBuilder::with_capacity(capacity);
builder.extend(iter.map(|item| item.into().native));
builder.finish()
}
}
impl BooleanArray {
/// Creates a [`BooleanArray`] from an iterator of trusted length.
///
/// # Safety
///
/// The iterator must be [`TrustedLen`](https://doc.rust-lang.org/std/iter/trait.TrustedLen.html).
/// I.e. that `size_hint().1` correctly reports its length. Note that this is a stronger
/// guarantee that `ExactSizeIterator` provides which could still report a wrong length.
///
/// # Panics
///
/// Panics if the iterator does not report an upper bound on `size_hint()`.
#[inline]
#[allow(
private_bounds,
reason = "We will expose BooleanAdapter if there is a need"
)]
pub unsafe fn from_trusted_len_iter<I, P>(iter: I) -> Self
where
P: Into<BooleanAdapter>,
I: ExactSizeIterator<Item = P>,
{
let data_len = iter.len();
let num_bytes = bit_util::ceil(data_len, 8);
let mut null_builder = MutableBuffer::from_len_zeroed(num_bytes);
let mut val_builder = MutableBuffer::from_len_zeroed(num_bytes);
let data = val_builder.as_slice_mut();
let null_slice = null_builder.as_slice_mut();
iter.enumerate().for_each(|(i, item)| {
if let Some(a) = item.into().native {
unsafe {
// SAFETY: There will be enough space in the buffers due to the trusted len size
// hint
bit_util::set_bit_raw(null_slice.as_mut_ptr(), i);
if a {
bit_util::set_bit_raw(data.as_mut_ptr(), i);
}
}
}
});
let data = unsafe {
ArrayData::new_unchecked(
DataType::Boolean,
data_len,
None,
Some(null_builder.into()),
0,
vec![val_builder.into()],
vec![],
)
};
BooleanArray::from(data)
}
}
impl From<BooleanBuffer> for BooleanArray {
fn from(values: BooleanBuffer) -> Self {
Self {
values,
nulls: None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use arrow_buffer::Buffer;
use rand::{Rng, rng};
#[test]
fn test_boolean_fmt_debug() {
let arr = BooleanArray::from(vec![true, false, false]);
assert_eq!(
"BooleanArray\n[\n true,\n false,\n false,\n]",
format!("{arr:?}")
);
}
#[test]
fn test_boolean_with_null_fmt_debug() {
let mut builder = BooleanArray::builder(3);
builder.append_value(true);
builder.append_null();
builder.append_value(false);
let arr = builder.finish();
assert_eq!(
"BooleanArray\n[\n true,\n null,\n false,\n]",
format!("{arr:?}")
);
}
#[test]
fn test_boolean_array_from_vec() {
let buf = Buffer::from([10_u8]);
let arr = BooleanArray::from(vec![false, true, false, true]);
assert_eq!(&buf, arr.values().inner());
assert_eq!(4, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(0, arr.null_count());
for i in 0..4 {
assert!(!arr.is_null(i));
assert!(arr.is_valid(i));
assert_eq!(i == 1 || i == 3, arr.value(i), "failed at {i}")
}
}
#[test]
fn test_boolean_array_from_vec_option() {
let buf = Buffer::from([10_u8]);
let arr = BooleanArray::from(vec![Some(false), Some(true), None, Some(true)]);
assert_eq!(&buf, arr.values().inner());
assert_eq!(4, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(1, arr.null_count());
for i in 0..4 {
if i == 2 {
assert!(arr.is_null(i));
assert!(!arr.is_valid(i));
} else {
assert!(!arr.is_null(i));
assert!(arr.is_valid(i));
assert_eq!(i == 1 || i == 3, arr.value(i), "failed at {i}")
}
}
}
#[test]
fn test_boolean_array_from_packed() {
let v = [1_u8, 2_u8, 3_u8];
let arr = BooleanArray::new_from_packed(v, 0, 24);
assert_eq!(24, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(0, arr.null_count());
assert!(arr.nulls.is_none());
for i in 0..24 {
assert!(!arr.is_null(i));
assert!(arr.is_valid(i));
assert_eq!(
i == 0 || i == 9 || i == 16 || i == 17,
arr.value(i),
"failed t {i}"
)
}
}
#[test]
fn test_boolean_array_from_slice_u8() {
let v: Vec<u8> = vec![1, 2, 3];
let slice = &v[..];
let arr = BooleanArray::new_from_u8(slice);
assert_eq!(24, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(0, arr.null_count());
assert!(arr.nulls().is_none());
for i in 0..24 {
assert!(!arr.is_null(i));
assert!(arr.is_valid(i));
assert_eq!(
i == 0 || i == 9 || i == 16 || i == 17,
arr.value(i),
"failed t {i}"
)
}
}
#[test]
fn test_boolean_array_from_iter() {
let v = vec![Some(false), Some(true), Some(false), Some(true)];
let arr = v.into_iter().collect::<BooleanArray>();
assert_eq!(4, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(0, arr.null_count());
assert!(arr.nulls().is_none());
for i in 0..3 {
assert!(!arr.is_null(i));
assert!(arr.is_valid(i));
assert_eq!(i == 1 || i == 3, arr.value(i), "failed at {i}")
}
}
#[test]
fn test_boolean_array_from_non_nullable_iter() {
let v = vec![true, false, true];
let arr = v.into_iter().collect::<BooleanArray>();
assert_eq!(3, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(0, arr.null_count());
assert!(arr.nulls().is_none());
assert!(arr.value(0));
assert!(!arr.value(1));
assert!(arr.value(2));
}
#[test]
fn test_boolean_array_from_nullable_iter() {
let v = vec![Some(true), None, Some(false), None];
let arr = v.into_iter().collect::<BooleanArray>();
assert_eq!(4, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(2, arr.null_count());
assert!(arr.nulls().is_some());
assert!(arr.is_valid(0));
assert!(arr.is_null(1));
assert!(arr.is_valid(2));
assert!(arr.is_null(3));
assert!(arr.value(0));
assert!(!arr.value(2));
}
#[test]
fn test_boolean_array_from_nullable_trusted_len_iter() {
// Should exhibit the same behavior as `from_iter`, which is tested above.
let v = vec![Some(true), None, Some(false), None];
let expected = v.clone().into_iter().collect::<BooleanArray>();
let actual = unsafe {
// SAFETY: `v` has trusted length
BooleanArray::from_trusted_len_iter(v.into_iter())
};
assert_eq!(expected, actual);
}
#[test]
fn test_boolean_array_from_iter_with_larger_upper_bound() {
// See https://github.com/apache/arrow-rs/issues/8505
// This returns an upper size hint of 4
let iterator = vec![Some(true), None, Some(false), None]
.into_iter()
.filter(Option::is_some);
let arr = iterator.collect::<BooleanArray>();
assert_eq!(2, arr.len());
}
#[test]
fn test_boolean_array_builder() {
// Test building a boolean array with ArrayData builder and offset
// 000011011
let buf = Buffer::from([27_u8]);
let buf2 = buf.clone();
let data = ArrayData::builder(DataType::Boolean)
.len(5)
.offset(2)
.add_buffer(buf)
.build()
.unwrap();
let arr = BooleanArray::from(data);
assert_eq!(&buf2, arr.values().inner());
assert_eq!(5, arr.len());
assert_eq!(2, arr.offset());
assert_eq!(0, arr.null_count());
for i in 0..3 {
assert_eq!(i != 0, arr.value(i), "failed at {i}");
}
}
#[test]
#[should_panic(
expected = "Trying to access an element at index 4 from a BooleanArray of length 3"
)]
fn test_fixed_size_binary_array_get_value_index_out_of_bound() {
let v = vec![Some(true), None, Some(false)];
let array = v.into_iter().collect::<BooleanArray>();
array.value(4);
}
#[test]
#[should_panic(expected = "BooleanArray data should contain a single buffer only \
(values buffer)")]
// Different error messages, so skip for now
// https://github.com/apache/arrow-rs/issues/1545
#[cfg(not(feature = "force_validate"))]
fn test_boolean_array_invalid_buffer_len() {
let data = unsafe {
ArrayData::builder(DataType::Boolean)
.len(5)
.build_unchecked()
};
drop(BooleanArray::from(data));
}
#[test]
#[should_panic(expected = "BooleanArray expected ArrayData with type Boolean got Int32")]
fn test_from_array_data_validation() {
let _ = BooleanArray::from(ArrayData::new_empty(&DataType::Int32));
}
#[test]
#[cfg_attr(miri, ignore)] // Takes too long
fn test_true_false_count() {
let mut rng = rng();
for _ in 0..10 {
// No nulls
let d: Vec<_> = (0..2000).map(|_| rng.random_bool(0.5)).collect();
let b = BooleanArray::from(d.clone());
let expected_true = d.iter().filter(|x| **x).count();
assert_eq!(b.true_count(), expected_true);
assert_eq!(b.false_count(), d.len() - expected_true);
// With nulls
let d: Vec<_> = (0..2000)
.map(|_| rng.random_bool(0.5).then(|| rng.random_bool(0.5)))
.collect();
let b = BooleanArray::from(d.clone());
let expected_true = d.iter().filter(|x| matches!(x, Some(true))).count();
assert_eq!(b.true_count(), expected_true);
let expected_false = d.iter().filter(|x| matches!(x, Some(false))).count();
assert_eq!(b.false_count(), expected_false);
}
}
#[test]
fn test_into_parts() {
let boolean_array = [Some(true), None, Some(false)]
.into_iter()
.collect::<BooleanArray>();
let (values, nulls) = boolean_array.into_parts();
assert_eq!(values.values(), &[0b0000_0001]);
assert!(nulls.is_some());
assert_eq!(nulls.unwrap().buffer().as_slice(), &[0b0000_0101]);
let boolean_array =
BooleanArray::from(vec![false, false, false, false, false, false, false, true]);
let (values, nulls) = boolean_array.into_parts();
assert_eq!(values.values(), &[0b1000_0000]);
assert!(nulls.is_none());
}
#[test]
fn test_new_null_array() {
let arr = BooleanArray::new_null(5);
assert_eq!(arr.len(), 5);
assert_eq!(arr.null_count(), 5);
assert_eq!(arr.true_count(), 0);
assert_eq!(arr.false_count(), 0);
for i in 0..5 {
assert!(arr.is_null(i));
assert!(!arr.is_valid(i));
}
}
#[test]
fn test_slice_with_nulls() {
let arr = BooleanArray::from(vec![Some(true), None, Some(false)]);
let sliced = arr.slice(1, 2);
assert_eq!(sliced.len(), 2);
assert_eq!(sliced.null_count(), 1);
assert!(sliced.is_null(0));
assert!(sliced.is_valid(1));
assert!(!sliced.value(1));
}
}
+716
View File
@@ -0,0 +1,716 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::array::{get_offsets, print_long_array};
use crate::builder::GenericByteBuilder;
use crate::iterator::ArrayIter;
use crate::types::ByteArrayType;
use crate::types::bytes::ByteArrayNativeType;
use crate::{Array, ArrayAccessor, ArrayRef, OffsetSizeTrait, Scalar};
use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer};
use arrow_buffer::{NullBuffer, OffsetBuffer};
use arrow_data::{ArrayData, ArrayDataBuilder};
use arrow_schema::{ArrowError, DataType};
use std::any::Any;
use std::sync::Arc;
/// An array of [variable length byte arrays](https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-layout)
///
/// See [`StringArray`] and [`LargeStringArray`] for storing utf8 encoded string data
///
/// See [`BinaryArray`] and [`LargeBinaryArray`] for storing arbitrary bytes
///
/// # Example: From a Vec
///
/// ```
/// # use arrow_array::{Array, GenericByteArray, types::Utf8Type};
/// let arr: GenericByteArray<Utf8Type> = vec!["hello", "world", ""].into();
/// assert_eq!(arr.value_data(), b"helloworld");
/// assert_eq!(arr.value_offsets(), &[0, 5, 10, 10]);
/// let values: Vec<_> = arr.iter().collect();
/// assert_eq!(values, &[Some("hello"), Some("world"), Some("")]);
/// ```
///
/// # Example: From an optional Vec
///
/// ```
/// # use arrow_array::{Array, GenericByteArray, types::Utf8Type};
/// let arr: GenericByteArray<Utf8Type> = vec![Some("hello"), Some("world"), Some(""), None].into();
/// assert_eq!(arr.value_data(), b"helloworld");
/// assert_eq!(arr.value_offsets(), &[0, 5, 10, 10, 10]);
/// let values: Vec<_> = arr.iter().collect();
/// assert_eq!(values, &[Some("hello"), Some("world"), Some(""), None]);
/// ```
///
/// # Example: From an iterator of option
///
/// ```
/// # use arrow_array::{Array, GenericByteArray, types::Utf8Type};
/// let arr: GenericByteArray<Utf8Type> = (0..5).map(|x| (x % 2 == 0).then(|| x.to_string())).collect();
/// let values: Vec<_> = arr.iter().collect();
/// assert_eq!(values, &[Some("0"), None, Some("2"), None, Some("4")]);
/// ```
///
/// # Example: Using Builder
///
/// ```
/// # use arrow_array::Array;
/// # use arrow_array::builder::GenericByteBuilder;
/// # use arrow_array::types::Utf8Type;
/// let mut builder = GenericByteBuilder::<Utf8Type>::new();
/// builder.append_value("hello");
/// builder.append_null();
/// builder.append_value("world");
/// let array = builder.finish();
/// let values: Vec<_> = array.iter().collect();
/// assert_eq!(values, &[Some("hello"), None, Some("world")]);
/// ```
///
/// [`StringArray`]: crate::StringArray
/// [`LargeStringArray`]: crate::LargeStringArray
/// [`BinaryArray`]: crate::BinaryArray
/// [`LargeBinaryArray`]: crate::LargeBinaryArray
pub struct GenericByteArray<T: ByteArrayType> {
data_type: DataType,
value_offsets: OffsetBuffer<T::Offset>,
value_data: Buffer,
nulls: Option<NullBuffer>,
}
impl<T: ByteArrayType> Clone for GenericByteArray<T> {
fn clone(&self) -> Self {
Self {
data_type: T::DATA_TYPE,
value_offsets: self.value_offsets.clone(),
value_data: self.value_data.clone(),
nulls: self.nulls.clone(),
}
}
}
impl<T: ByteArrayType> GenericByteArray<T> {
/// Data type of the array.
pub const DATA_TYPE: DataType = T::DATA_TYPE;
/// Create a new [`GenericByteArray`] from the provided parts, panicking on failure
///
/// # Panics
///
/// Panics if [`GenericByteArray::try_new`] returns an error
pub fn new(
offsets: OffsetBuffer<T::Offset>,
values: Buffer,
nulls: Option<NullBuffer>,
) -> Self {
Self::try_new(offsets, values, nulls).unwrap()
}
/// Create a new [`GenericByteArray`] from the provided parts, returning an error on failure
///
/// # Errors
///
/// * `offsets.len() - 1 != nulls.len()`
/// * Any consecutive pair of `offsets` does not denote a valid slice of `values`
pub fn try_new(
offsets: OffsetBuffer<T::Offset>,
values: Buffer,
nulls: Option<NullBuffer>,
) -> Result<Self, ArrowError> {
let len = offsets.len() - 1;
// Verify that each pair of offsets is a valid slices of values
T::validate(&offsets, &values)?;
if let Some(n) = nulls.as_ref() {
if n.len() != len {
return Err(ArrowError::InvalidArgumentError(format!(
"Incorrect length of null buffer for {}{}Array, expected {len} got {}",
T::Offset::PREFIX,
T::PREFIX,
n.len(),
)));
}
}
Ok(Self {
data_type: T::DATA_TYPE,
value_offsets: offsets,
value_data: values,
nulls,
})
}
/// Create a new [`GenericByteArray`] from the provided parts, without validation
///
/// # Safety
///
/// Safe if [`Self::try_new`] would not error
pub unsafe fn new_unchecked(
offsets: OffsetBuffer<T::Offset>,
values: Buffer,
nulls: Option<NullBuffer>,
) -> Self {
if cfg!(feature = "force_validate") {
return Self::new(offsets, values, nulls);
}
Self {
data_type: T::DATA_TYPE,
value_offsets: offsets,
value_data: values,
nulls,
}
}
/// Create a new [`GenericByteArray`] of length `len` where all values are null
pub fn new_null(len: usize) -> Self {
Self {
data_type: T::DATA_TYPE,
value_offsets: OffsetBuffer::new_zeroed(len),
value_data: MutableBuffer::new(0).into(),
nulls: Some(NullBuffer::new_null(len)),
}
}
/// Create a new [`Scalar`] from `v`
pub fn new_scalar(value: impl AsRef<T::Native>) -> Scalar<Self> {
Scalar::new(Self::from_iter_values(std::iter::once(value)))
}
/// Create a new [`GenericByteArray`] where `value` is repeated `repeat_count` times.
///
/// # Panics
/// This will panic if value's length multiplied by `repeat_count` overflows usize.
///
pub fn new_repeated(value: impl AsRef<T::Native>, repeat_count: usize) -> Self {
let s: &[u8] = value.as_ref().as_ref();
let value_offsets = OffsetBuffer::from_repeated_length(s.len(), repeat_count);
let bytes: Buffer = {
let mut mutable_buffer = MutableBuffer::with_capacity(0);
mutable_buffer.repeat_slice_n_times(s, repeat_count);
mutable_buffer.into()
};
Self {
data_type: T::DATA_TYPE,
value_data: bytes,
value_offsets,
nulls: None,
}
}
/// Creates a [`GenericByteArray`] based on an iterator of values without nulls
pub fn from_iter_values<Ptr, I>(iter: I) -> Self
where
Ptr: AsRef<T::Native>,
I: IntoIterator<Item = Ptr>,
{
let iter = iter.into_iter();
let (_, data_len) = iter.size_hint();
let data_len = data_len.expect("Iterator must be sized"); // panic if no upper bound.
let mut offsets = MutableBuffer::new((data_len + 1) * std::mem::size_of::<T::Offset>());
offsets.push(T::Offset::usize_as(0));
let mut values = MutableBuffer::new(0);
for s in iter {
let s: &[u8] = s.as_ref().as_ref();
values.extend_from_slice(s);
offsets.push(T::Offset::usize_as(values.len()));
}
T::Offset::from_usize(values.len()).expect("offset overflow");
let offsets = Buffer::from(offsets);
// Safety: valid by construction
let value_offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
Self {
data_type: T::DATA_TYPE,
value_data: values.into(),
value_offsets,
nulls: None,
}
}
/// Deconstruct this array into its constituent parts
pub fn into_parts(self) -> (OffsetBuffer<T::Offset>, Buffer, Option<NullBuffer>) {
(self.value_offsets, self.value_data, self.nulls)
}
/// Returns the length for value at index `i`.
/// # Panics
/// Panics if index `i` is out of bounds.
#[inline]
pub fn value_length(&self, i: usize) -> T::Offset {
let offsets = self.value_offsets();
offsets[i + 1] - offsets[i]
}
/// Returns a reference to the offsets of this array
///
/// Unlike [`Self::value_offsets`] this returns the [`OffsetBuffer`]
/// allowing for zero-copy cloning
#[inline]
pub fn offsets(&self) -> &OffsetBuffer<T::Offset> {
&self.value_offsets
}
/// Returns the values of this array
///
/// Unlike [`Self::value_data`] this returns the [`Buffer`]
/// allowing for zero-copy cloning
#[inline]
pub fn values(&self) -> &Buffer {
&self.value_data
}
/// Returns the raw value data
pub fn value_data(&self) -> &[u8] {
self.value_data.as_slice()
}
/// Returns true if all data within this array is ASCII
pub fn is_ascii(&self) -> bool {
let offsets = self.value_offsets();
let start = offsets.first().unwrap();
let end = offsets.last().unwrap();
self.value_data()[start.as_usize()..end.as_usize()].is_ascii()
}
/// Returns the offset values in the offsets buffer
#[inline]
pub fn value_offsets(&self) -> &[T::Offset] {
&self.value_offsets
}
/// Returns the element at index `i`
///
/// Note: This method does not check for nulls and the value is arbitrary
/// if [`is_null`](Self::is_null) returns true for the index.
///
/// # Safety
/// Caller is responsible for ensuring that the index is within the bounds of the array
pub unsafe fn value_unchecked(&self, i: usize) -> &T::Native {
let end = *unsafe { self.value_offsets().get_unchecked(i + 1) };
let start = *unsafe { self.value_offsets().get_unchecked(i) };
// Soundness
// pointer alignment & location is ensured by RawPtrBox
// buffer bounds/offset is ensured by the value_offset invariants
// Safety of `to_isize().unwrap()`
// `start` and `end` are &OffsetSize, which is a generic type that implements the
// OffsetSizeTrait. Currently, only i32 and i64 implement OffsetSizeTrait,
// both of which should cleanly cast to isize on an architecture that supports
// 32/64-bit offsets
let b = unsafe {
std::slice::from_raw_parts(
self.value_data
.as_ptr()
.offset(start.to_isize().unwrap_unchecked()),
(end - start).to_usize().unwrap_unchecked(),
)
};
// SAFETY:
// ArrayData is valid
unsafe { T::Native::from_bytes_unchecked(b) }
}
/// Returns the element at index `i`
///
/// Note: This method does not check for nulls and the value is arbitrary
/// (but still well-defined) if [`is_null`](Self::is_null) returns true for the index.
///
/// # Panics
/// Panics if index `i` is out of bounds.
pub fn value(&self, i: usize) -> &T::Native {
assert!(
i < self.len(),
"Trying to access an element at index {} from a {}{}Array of length {}",
i,
T::Offset::PREFIX,
T::PREFIX,
self.len()
);
// SAFETY:
// Verified length above
unsafe { self.value_unchecked(i) }
}
/// constructs a new iterator
pub fn iter(&self) -> ArrayIter<&Self> {
ArrayIter::new(self)
}
/// Returns a zero-copy slice of this array with the indicated offset and length.
pub fn slice(&self, offset: usize, length: usize) -> Self {
Self {
data_type: T::DATA_TYPE,
value_offsets: self.value_offsets.slice(offset, length),
value_data: self.value_data.clone(),
nulls: self.nulls.as_ref().map(|n| n.slice(offset, length)),
}
}
/// Returns `GenericByteBuilder` of this byte array for mutating its values if the underlying
/// offset and data buffers are not shared by others.
pub fn into_builder(self) -> Result<GenericByteBuilder<T>, Self> {
let len = self.len();
let value_len = T::Offset::as_usize(self.value_offsets()[len] - self.value_offsets()[0]);
let data = self.into_data();
let null_bit_buffer = data.nulls().map(|b| b.inner().sliced());
let element_len = std::mem::size_of::<T::Offset>();
let offset_buffer = data.buffers()[0]
.slice_with_length(data.offset() * element_len, (len + 1) * element_len);
let element_len = std::mem::size_of::<u8>();
let value_buffer = data.buffers()[1]
.slice_with_length(data.offset() * element_len, value_len * element_len);
drop(data);
let try_mutable_null_buffer = match null_bit_buffer {
None => Ok(None),
Some(null_buffer) => {
// Null buffer exists, tries to make it mutable
null_buffer.into_mutable().map(Some)
}
};
let try_mutable_buffers = match try_mutable_null_buffer {
Ok(mutable_null_buffer) => {
// Got mutable null buffer, tries to get mutable value buffer
let try_mutable_offset_buffer = offset_buffer.into_mutable();
let try_mutable_value_buffer = value_buffer.into_mutable();
// try_mutable_offset_buffer.map(...).map_err(...) doesn't work as the compiler complains
// mutable_null_buffer is moved into map closure.
match (try_mutable_offset_buffer, try_mutable_value_buffer) {
(Ok(mutable_offset_buffer), Ok(mutable_value_buffer)) => unsafe {
Ok(GenericByteBuilder::<T>::new_from_buffer(
mutable_offset_buffer,
mutable_value_buffer,
mutable_null_buffer,
))
},
(Ok(mutable_offset_buffer), Err(value_buffer)) => Err((
mutable_offset_buffer.into(),
value_buffer,
mutable_null_buffer.map(|b| b.into()),
)),
(Err(offset_buffer), Ok(mutable_value_buffer)) => Err((
offset_buffer,
mutable_value_buffer.into(),
mutable_null_buffer.map(|b| b.into()),
)),
(Err(offset_buffer), Err(value_buffer)) => Err((
offset_buffer,
value_buffer,
mutable_null_buffer.map(|b| b.into()),
)),
}
}
Err(mutable_null_buffer) => {
// Unable to get mutable null buffer
Err((offset_buffer, value_buffer, Some(mutable_null_buffer)))
}
};
match try_mutable_buffers {
Ok(builder) => Ok(builder),
Err((offset_buffer, value_buffer, null_bit_buffer)) => {
let builder = ArrayData::builder(T::DATA_TYPE)
.len(len)
.add_buffer(offset_buffer)
.add_buffer(value_buffer)
.null_bit_buffer(null_bit_buffer);
let array_data = unsafe { builder.build_unchecked() };
let array = GenericByteArray::<T>::from(array_data);
Err(array)
}
}
}
}
impl<T: ByteArrayType> std::fmt::Debug for GenericByteArray<T> {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(f, "{}{}Array\n[\n", T::Offset::PREFIX, T::PREFIX)?;
print_long_array(self, f, |array, index, f| {
std::fmt::Debug::fmt(&array.value(index), f)
})?;
write!(f, "]")
}
}
/// SAFETY: Correctly implements the contract of Arrow Arrays
unsafe impl<T: ByteArrayType> Array for GenericByteArray<T> {
fn as_any(&self) -> &dyn Any {
self
}
fn to_data(&self) -> ArrayData {
self.clone().into()
}
fn into_data(self) -> ArrayData {
self.into()
}
fn data_type(&self) -> &DataType {
&self.data_type
}
fn slice(&self, offset: usize, length: usize) -> ArrayRef {
Arc::new(self.slice(offset, length))
}
fn len(&self) -> usize {
self.value_offsets.len() - 1
}
fn is_empty(&self) -> bool {
self.value_offsets.len() <= 1
}
fn shrink_to_fit(&mut self) {
self.value_offsets.shrink_to_fit();
self.value_data.shrink_to_fit();
if let Some(nulls) = &mut self.nulls {
nulls.shrink_to_fit();
}
}
fn offset(&self) -> usize {
0
}
fn nulls(&self) -> Option<&NullBuffer> {
self.nulls.as_ref()
}
fn logical_null_count(&self) -> usize {
// More efficient that the default implementation
self.null_count()
}
fn get_buffer_memory_size(&self) -> usize {
let mut sum = self.value_offsets.inner().inner().capacity();
sum += self.value_data.capacity();
if let Some(x) = &self.nulls {
sum += x.buffer().capacity()
}
sum
}
fn get_array_memory_size(&self) -> usize {
std::mem::size_of::<Self>() + self.get_buffer_memory_size()
}
}
impl<'a, T: ByteArrayType> ArrayAccessor for &'a GenericByteArray<T> {
type Item = &'a T::Native;
fn value(&self, index: usize) -> Self::Item {
GenericByteArray::value(self, index)
}
unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
unsafe { GenericByteArray::value_unchecked(self, index) }
}
}
impl<T: ByteArrayType> From<ArrayData> for GenericByteArray<T> {
fn from(data: ArrayData) -> Self {
assert_eq!(
data.data_type(),
&Self::DATA_TYPE,
"{}{}Array expects DataType::{}",
T::Offset::PREFIX,
T::PREFIX,
Self::DATA_TYPE
);
assert_eq!(
data.buffers().len(),
2,
"{}{}Array data should contain 2 buffers only (offsets and values)",
T::Offset::PREFIX,
T::PREFIX,
);
// SAFETY:
// ArrayData is valid, and verified type above
let value_offsets = unsafe { get_offsets(&data) };
let value_data = data.buffers()[1].clone();
Self {
value_offsets,
value_data,
data_type: T::DATA_TYPE,
nulls: data.nulls().cloned(),
}
}
}
impl<T: ByteArrayType> From<GenericByteArray<T>> for ArrayData {
fn from(array: GenericByteArray<T>) -> Self {
let len = array.len();
let offsets = array.value_offsets.into_inner().into_inner();
let builder = ArrayDataBuilder::new(array.data_type)
.len(len)
.buffers(vec![offsets, array.value_data])
.nulls(array.nulls);
unsafe { builder.build_unchecked() }
}
}
impl<'a, T: ByteArrayType> IntoIterator for &'a GenericByteArray<T> {
type Item = Option<&'a T::Native>;
type IntoIter = ArrayIter<Self>;
fn into_iter(self) -> Self::IntoIter {
ArrayIter::new(self)
}
}
impl<'a, Ptr, T: ByteArrayType> FromIterator<&'a Option<Ptr>> for GenericByteArray<T>
where
Ptr: AsRef<T::Native> + 'a,
{
fn from_iter<I: IntoIterator<Item = &'a Option<Ptr>>>(iter: I) -> Self {
iter.into_iter()
.map(|o| o.as_ref().map(|p| p.as_ref()))
.collect()
}
}
impl<Ptr, T: ByteArrayType> FromIterator<Option<Ptr>> for GenericByteArray<T>
where
Ptr: AsRef<T::Native>,
{
fn from_iter<I: IntoIterator<Item = Option<Ptr>>>(iter: I) -> Self {
let iter = iter.into_iter();
let mut builder = GenericByteBuilder::with_capacity(iter.size_hint().0, 1024);
builder.extend(iter);
builder.finish()
}
}
#[cfg(test)]
mod tests {
use crate::{Array, BinaryArray, StringArray};
use arrow_buffer::{Buffer, NullBuffer, OffsetBuffer};
#[test]
fn try_new() {
let data = Buffer::from_slice_ref("helloworld");
let offsets = OffsetBuffer::new(vec![0, 5, 10].into());
StringArray::new(offsets.clone(), data.clone(), None);
let nulls = NullBuffer::new_null(3);
let err =
StringArray::try_new(offsets.clone(), data.clone(), Some(nulls.clone())).unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: Incorrect length of null buffer for StringArray, expected 2 got 3"
);
let err = BinaryArray::try_new(offsets.clone(), data.clone(), Some(nulls)).unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: Incorrect length of null buffer for BinaryArray, expected 2 got 3"
);
let non_utf8_data = Buffer::from_slice_ref(b"he\xFFloworld");
let err = StringArray::try_new(offsets.clone(), non_utf8_data.clone(), None).unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: Encountered non UTF-8 data: invalid utf-8 sequence of 1 bytes from index 2"
);
BinaryArray::new(offsets, non_utf8_data, None);
let offsets = OffsetBuffer::new(vec![0, 5, 11].into());
let err = StringArray::try_new(offsets.clone(), data.clone(), None).unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: Offset of 11 exceeds length of values 10"
);
let err = BinaryArray::try_new(offsets.clone(), data, None).unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: Maximum offset of 11 is larger than values of length 10"
);
let non_ascii_data = Buffer::from_slice_ref("heìloworld");
StringArray::new(offsets.clone(), non_ascii_data.clone(), None);
BinaryArray::new(offsets, non_ascii_data.clone(), None);
let offsets = OffsetBuffer::new(vec![0, 3, 10].into());
let err = StringArray::try_new(offsets.clone(), non_ascii_data.clone(), None).unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: Split UTF-8 codepoint at offset 3"
);
BinaryArray::new(offsets, non_ascii_data, None);
}
#[test]
fn create_repeated() {
let arr = BinaryArray::new_repeated(b"hello", 3);
assert_eq!(arr.len(), 3);
assert_eq!(arr.value(0), b"hello");
assert_eq!(arr.value(1), b"hello");
assert_eq!(arr.value(2), b"hello");
let arr = StringArray::new_repeated("world", 2);
assert_eq!(arr.len(), 2);
assert_eq!(arr.value(0), "world");
assert_eq!(arr.value(1), "world");
}
#[test]
#[should_panic(expected = "usize overflow")]
fn create_repeated_usize_overflow_1() {
let _arr = BinaryArray::new_repeated(b"hello", (usize::MAX / "hello".len()) + 1);
}
#[test]
#[should_panic(expected = "usize overflow")]
fn create_repeated_usize_overflow_2() {
let _arr = BinaryArray::new_repeated(b"hello", usize::MAX);
}
#[test]
#[should_panic(expected = "offset overflow")]
fn create_repeated_i32_offset_overflow_1() {
let _arr = BinaryArray::new_repeated(b"hello", usize::MAX / "hello".len());
}
#[test]
#[should_panic(expected = "offset overflow")]
fn create_repeated_i32_offset_overflow_2() {
let _arr = BinaryArray::new_repeated(b"hello", ((i32::MAX as usize) / "hello".len()) + 1);
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,848 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::array::print_long_array;
use crate::builder::{FixedSizeListBuilder, PrimitiveBuilder};
use crate::iterator::FixedSizeListIter;
use crate::{Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType, make_array};
use arrow_buffer::ArrowNativeType;
use arrow_buffer::buffer::NullBuffer;
use arrow_data::{ArrayData, ArrayDataBuilder};
use arrow_schema::{ArrowError, DataType, FieldRef};
use std::any::Any;
use std::sync::Arc;
/// An array of [fixed length lists], similar to JSON arrays
/// (e.g. `["A", "B"]`).
///
/// Lists are represented using a `values` child
/// array where each list has a fixed size of `value_length`.
///
/// Use [`FixedSizeListBuilder`] to construct a [`FixedSizeListArray`].
///
/// # Representation
///
/// A [`FixedSizeListArray`] can represent a list of values of any other
/// supported Arrow type. Each element of the `FixedSizeListArray` itself is
/// a list which may contain NULL and non-null values,
/// or may itself be NULL.
///
/// For example, this `FixedSizeListArray` stores lists of strings:
///
/// ```text
/// ┌─────────────┐
/// │ [A,B] │
/// ├─────────────┤
/// │ NULL │
/// ├─────────────┤
/// │ [C,NULL] │
/// └─────────────┘
/// ```
///
/// The `values` of this `FixedSizeListArray`s are stored in a child
/// [`StringArray`] where logical null values take up `values_length` slots in the array
/// as shown in the following diagram. The logical values
/// are shown on the left, and the actual `FixedSizeListArray` encoding on the right
///
/// ```text
/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐
/// ┌ ─ ─ ─ ─ ─ ─ ─ ─┐
/// ┌─────────────┐ │ ┌───┐ ┌───┐ ┌──────┐ │
/// │ [A,B] │ │ 1 │ │ │ 1 │ │ A │ │ 0
/// ├─────────────┤ │ ├───┤ ├───┤ ├──────┤ │
/// │ NULL │ │ 0 │ │ │ 1 │ │ B │ │ 1
/// ├─────────────┤ │ ├───┤ ├───┤ ├──────┤ │
/// │ [C,NULL] │ │ 1 │ │ │ 0 │ │ ???? │ │ 2
/// └─────────────┘ │ └───┘ ├───┤ ├──────┤ │
/// | │ 0 │ │ ???? │ │ 3
/// Logical Values │ Validity ├───┤ ├──────┤ │
/// (nulls) │ │ 1 │ │ C │ │ 4
/// │ ├───┤ ├──────┤ │
/// │ │ 0 │ │ ???? │ │ 5
/// │ └───┘ └──────┘ │
/// │ Values │
/// │ FixedSizeListArray (Array) │
/// └ ─ ─ ─ ─ ─ ─ ─ ─┘
/// └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘
/// ```
///
/// # Example
///
/// ```
/// # use std::sync::Arc;
/// # use arrow_array::{Array, FixedSizeListArray, Int32Array};
/// # use arrow_data::ArrayData;
/// # use arrow_schema::{DataType, Field};
/// # use arrow_buffer::Buffer;
/// // Construct a value array
/// let value_data = ArrayData::builder(DataType::Int32)
/// .len(9)
/// .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7, 8]))
/// .build()
/// .unwrap();
/// let list_data_type = DataType::FixedSizeList(
/// Arc::new(Field::new_list_field(DataType::Int32, false)),
/// 3,
/// );
/// let list_data = ArrayData::builder(list_data_type.clone())
/// .len(3)
/// .add_child_data(value_data.clone())
/// .build()
/// .unwrap();
/// let list_array = FixedSizeListArray::from(list_data);
/// let list0 = list_array.value(0);
/// let list1 = list_array.value(1);
/// let list2 = list_array.value(2);
///
/// assert_eq!( &[0, 1, 2], list0.as_any().downcast_ref::<Int32Array>().unwrap().values());
/// assert_eq!( &[3, 4, 5], list1.as_any().downcast_ref::<Int32Array>().unwrap().values());
/// assert_eq!( &[6, 7, 8], list2.as_any().downcast_ref::<Int32Array>().unwrap().values());
/// ```
///
/// [`StringArray`]: crate::array::StringArray
/// [fixed size arrays](https://arrow.apache.org/docs/format/Columnar.html#fixed-size-list-layout)
#[derive(Clone)]
pub struct FixedSizeListArray {
data_type: DataType, // Must be DataType::FixedSizeList(value_length)
values: ArrayRef,
nulls: Option<NullBuffer>,
value_length: i32,
len: usize,
}
impl FixedSizeListArray {
/// Create a new [`FixedSizeListArray`] with `size` element size, panicking on failure.
///
/// Note that if `size == 0` and `nulls` is `None` (a degenerate, non-nullable
/// `FixedSizeListArray`), this function will set the length of the array to 0.
///
/// If you would like to have a degenerate, non-nullable `FixedSizeListArray` with arbitrary
/// length, use the [`try_new_with_length()`] constructor.
///
/// [`try_new_with_length()`]: Self::try_new_with_length
///
/// # Panics
///
/// Panics if [`Self::try_new`] returns an error
pub fn new(field: FieldRef, size: i32, values: ArrayRef, nulls: Option<NullBuffer>) -> Self {
Self::try_new(field, size, values, nulls).unwrap()
}
/// Create a new [`FixedSizeListArray`] from the provided parts, returning an error on failure.
///
/// Note that if `size == 0` and `nulls` is `None` (a degenerate, non-nullable
/// `FixedSizeListArray`), this function will set the length of the array to 0.
///
/// If you would like to have a degenerate, non-nullable `FixedSizeListArray` with arbitrary
/// length, use the [`try_new_with_length()`] constructor.
///
/// [`try_new_with_length()`]: Self::try_new_with_length
///
/// # Errors
///
/// * `size < 0`
/// * `values.len() != nulls.len() * size` if `nulls` is `Some`
/// * `values.data_type() != field.data_type()`
/// * `!field.is_nullable() && !nulls.expand(size).contains(values.logical_nulls())`
pub fn try_new(
field: FieldRef,
size: i32,
values: ArrayRef,
nulls: Option<NullBuffer>,
) -> Result<Self, ArrowError> {
let s = size.to_usize().ok_or_else(|| {
ArrowError::InvalidArgumentError(format!("Size cannot be negative, got {size}"))
})?;
if s == 0 {
// Note that for degenerate (`size == 0`) and non-nullable `FixedSizeList`s, we will set
// the length to 0 (`_or_default`).
let len = nulls.as_ref().map(|x| x.len()).unwrap_or_default();
Self::try_new_with_length(field, size, values, nulls, len)
} else {
if values.len() % s != 0 {
return Err(ArrowError::InvalidArgumentError(format!(
"Incorrect length of values buffer for FixedSizeListArray, \
expected a multiple of {s} got {}",
values.len(),
)));
}
let len = values.len() / s;
// Check that the null buffer length is correct (if it exists).
if let Some(null_buffer) = &nulls {
if s * null_buffer.len() != values.len() {
return Err(ArrowError::InvalidArgumentError(format!(
"Incorrect length of values buffer for FixedSizeListArray, \
expected {} got {}",
s * null_buffer.len(),
values.len(),
)));
}
}
Self::try_new_with_length(field, size, values, nulls, len)
}
}
/// Create a new [`FixedSizeListArray`] from the provided parts, returning an error on failure.
///
/// This method exists to allow the construction of arbitrary length degenerate (`size == 0`)
/// and non-nullable `FixedSizeListArray`s. If you want a nullable `FixedSizeListArray`, then
/// you can use [`try_new()`] instead.
///
/// [`try_new()`]: Self::try_new
///
/// # Errors
///
/// * `size < 0`
/// * `nulls.len() != len` if `nulls` is `Some`
/// * `values.len() != len * size`
/// * `values.data_type() != field.data_type()`
/// * `!field.is_nullable() && !nulls.expand(size).contains(values.logical_nulls())`
pub fn try_new_with_length(
field: FieldRef,
size: i32,
values: ArrayRef,
nulls: Option<NullBuffer>,
len: usize,
) -> Result<Self, ArrowError> {
let s = size.to_usize().ok_or_else(|| {
ArrowError::InvalidArgumentError(format!("Size cannot be negative, got {size}"))
})?;
if let Some(null_buffer) = &nulls {
if null_buffer.len() != len {
return Err(ArrowError::InvalidArgumentError(format!(
"Invalid null buffer for FixedSizeListArray, expected {len} found {}",
null_buffer.len()
)));
}
}
if s == 0 && !values.is_empty() {
return Err(ArrowError::InvalidArgumentError(format!(
"An degenerate FixedSizeListArray should have no underlying values, found {} values",
values.len()
)));
}
if values.len() != len * s {
return Err(ArrowError::InvalidArgumentError(format!(
"Incorrect length of values buffer for FixedSizeListArray, expected {} got {}",
len * s,
values.len(),
)));
}
if field.data_type() != values.data_type() {
return Err(ArrowError::InvalidArgumentError(format!(
"FixedSizeListArray expected data type {} got {} for {:?}",
field.data_type(),
values.data_type(),
field.name()
)));
}
if let Some(a) = values.logical_nulls() {
let nulls_valid = field.is_nullable()
|| nulls
.as_ref()
.map(|n| n.expand(size as _).contains(&a))
.unwrap_or_default()
|| (nulls.is_none() && a.null_count() == 0);
if !nulls_valid {
return Err(ArrowError::InvalidArgumentError(format!(
"Found unmasked nulls for non-nullable FixedSizeListArray field {:?}",
field.name()
)));
}
}
let data_type = DataType::FixedSizeList(field, size);
Ok(Self {
data_type,
values,
value_length: size,
nulls,
len,
})
}
/// Create a new [`FixedSizeListArray`] of length `len` where all values are null
///
/// # Panics
///
/// Panics if
///
/// * `size < 0`
/// * `size * len` would overflow `usize`
pub fn new_null(field: FieldRef, size: i32, len: usize) -> Self {
let capacity = size.to_usize().unwrap().checked_mul(len).unwrap();
Self {
values: make_array(ArrayData::new_null(field.data_type(), capacity)),
data_type: DataType::FixedSizeList(field, size),
nulls: Some(NullBuffer::new_null(len)),
value_length: size,
len,
}
}
/// Deconstruct this array into its constituent parts
pub fn into_parts(self) -> (FieldRef, i32, ArrayRef, Option<NullBuffer>) {
let f = match self.data_type {
DataType::FixedSizeList(f, _) => f,
_ => unreachable!(),
};
(f, self.value_length, self.values, self.nulls)
}
/// Returns a reference to the values of this list.
pub fn values(&self) -> &ArrayRef {
&self.values
}
/// Returns a clone of the value type of this list.
pub fn value_type(&self) -> DataType {
self.values.data_type().clone()
}
/// Returns ith value of this list array.
///
/// Note: This method does not check for nulls and the value is arbitrary
/// (but still well-defined) if [`is_null`](Self::is_null) returns true for the index.
///
/// # Panics
/// Panics if index `i` is out of bounds
pub fn value(&self, i: usize) -> ArrayRef {
self.values
.slice(self.value_offset_at(i), self.value_length() as usize)
}
/// Returns the offset for value at index `i`.
///
/// Note this doesn't do any bound checking, for performance reason.
#[inline]
pub fn value_offset(&self, i: usize) -> i32 {
self.value_offset_at(i) as i32
}
/// Returns the length for an element.
///
/// All elements have the same length as the array is a fixed size.
#[inline]
pub const fn value_length(&self) -> i32 {
self.value_length
}
#[inline]
const fn value_offset_at(&self, i: usize) -> usize {
i * self.value_length as usize
}
/// Returns a zero-copy slice of this array with the indicated offset and length.
pub fn slice(&self, offset: usize, len: usize) -> Self {
assert!(
offset.saturating_add(len) <= self.len,
"the length + offset of the sliced FixedSizeListArray cannot exceed the existing length"
);
let size = self.value_length as usize;
Self {
data_type: self.data_type.clone(),
values: self.values.slice(offset * size, len * size),
nulls: self.nulls.as_ref().map(|n| n.slice(offset, len)),
value_length: self.value_length,
len,
}
}
/// Creates a [`FixedSizeListArray`] from an iterator of primitive values
/// # Example
/// ```
/// # use arrow_array::FixedSizeListArray;
/// # use arrow_array::types::Int32Type;
///
/// let data = vec![
/// Some(vec![Some(0), Some(1), Some(2)]),
/// None,
/// Some(vec![Some(3), None, Some(5)]),
/// Some(vec![Some(6), Some(7), Some(45)]),
/// ];
/// let list_array = FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>(data, 3);
/// println!("{:?}", list_array);
/// ```
pub fn from_iter_primitive<T, P, I>(iter: I, length: i32) -> Self
where
T: ArrowPrimitiveType,
P: IntoIterator<Item = Option<<T as ArrowPrimitiveType>::Native>>,
I: IntoIterator<Item = Option<P>>,
{
let l = length as usize;
let iter = iter.into_iter();
let size_hint = iter.size_hint().0;
let mut builder = FixedSizeListBuilder::with_capacity(
PrimitiveBuilder::<T>::with_capacity(size_hint * l),
length,
size_hint,
);
for i in iter {
match i {
Some(p) => {
for t in p {
builder.values().append_option(t);
}
builder.append(true);
}
None => {
builder.values().append_nulls(l);
builder.append(false)
}
}
}
builder.finish()
}
/// constructs a new iterator
pub fn iter(&self) -> FixedSizeListIter<'_> {
FixedSizeListIter::new(self)
}
}
impl From<ArrayData> for FixedSizeListArray {
fn from(data: ArrayData) -> Self {
let value_length = match data.data_type() {
DataType::FixedSizeList(_, len) => *len,
data_type => {
panic!(
"FixedSizeListArray data should contain a FixedSizeList data type, got {data_type}"
)
}
};
let size = value_length as usize;
let values =
make_array(data.child_data()[0].slice(data.offset() * size, data.len() * size));
Self {
data_type: data.data_type().clone(),
values,
nulls: data.nulls().cloned(),
value_length,
len: data.len(),
}
}
}
impl From<FixedSizeListArray> for ArrayData {
fn from(array: FixedSizeListArray) -> Self {
let builder = ArrayDataBuilder::new(array.data_type)
.len(array.len)
.nulls(array.nulls)
.child_data(vec![array.values.to_data()]);
unsafe { builder.build_unchecked() }
}
}
/// SAFETY: Correctly implements the contract of Arrow Arrays
unsafe impl Array for FixedSizeListArray {
fn as_any(&self) -> &dyn Any {
self
}
fn to_data(&self) -> ArrayData {
self.clone().into()
}
fn into_data(self) -> ArrayData {
self.into()
}
fn data_type(&self) -> &DataType {
&self.data_type
}
fn slice(&self, offset: usize, length: usize) -> ArrayRef {
Arc::new(self.slice(offset, length))
}
fn len(&self) -> usize {
self.len
}
fn is_empty(&self) -> bool {
self.len == 0
}
fn shrink_to_fit(&mut self) {
self.values.shrink_to_fit();
if let Some(nulls) = &mut self.nulls {
nulls.shrink_to_fit();
}
}
fn offset(&self) -> usize {
0
}
fn nulls(&self) -> Option<&NullBuffer> {
self.nulls.as_ref()
}
fn logical_null_count(&self) -> usize {
// More efficient that the default implementation
self.null_count()
}
fn get_buffer_memory_size(&self) -> usize {
let mut size = self.values.get_buffer_memory_size();
if let Some(n) = self.nulls.as_ref() {
size += n.buffer().capacity();
}
size
}
fn get_array_memory_size(&self) -> usize {
let mut size = std::mem::size_of::<Self>() + self.values.get_array_memory_size();
if let Some(n) = self.nulls.as_ref() {
size += n.buffer().capacity();
}
size
}
}
impl ArrayAccessor for FixedSizeListArray {
type Item = ArrayRef;
fn value(&self, index: usize) -> Self::Item {
FixedSizeListArray::value(self, index)
}
unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
FixedSizeListArray::value(self, index)
}
}
impl std::fmt::Debug for FixedSizeListArray {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(f, "FixedSizeListArray<{}>\n[\n", self.value_length())?;
print_long_array(self, f, |array, index, f| {
std::fmt::Debug::fmt(&array.value(index), f)
})?;
write!(f, "]")
}
}
impl ArrayAccessor for &FixedSizeListArray {
type Item = ArrayRef;
fn value(&self, index: usize) -> Self::Item {
FixedSizeListArray::value(self, index)
}
unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
FixedSizeListArray::value(self, index)
}
}
#[cfg(test)]
mod tests {
use arrow_buffer::{BooleanBuffer, Buffer, bit_util};
use arrow_schema::Field;
use crate::cast::AsArray;
use crate::types::Int32Type;
use crate::{Int32Array, new_empty_array};
use super::*;
#[test]
fn test_fixed_size_list_array() {
// Construct a value array
let value_data = ArrayData::builder(DataType::Int32)
.len(9)
.add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8]))
.build()
.unwrap();
// Construct a list array from the above two
let list_data_type =
DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 3);
let list_data = ArrayData::builder(list_data_type.clone())
.len(3)
.add_child_data(value_data.clone())
.build()
.unwrap();
let list_array = FixedSizeListArray::from(list_data);
assert_eq!(value_data, list_array.values().to_data());
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(3, list_array.len());
assert_eq!(0, list_array.null_count());
assert_eq!(6, list_array.value_offset(2));
assert_eq!(3, list_array.value_length());
assert_eq!(0, list_array.value(0).as_primitive::<Int32Type>().value(0));
for i in 0..3 {
assert!(list_array.is_valid(i));
assert!(!list_array.is_null(i));
}
// Now test with a non-zero offset
let list_data = ArrayData::builder(list_data_type)
.len(2)
.offset(1)
.add_child_data(value_data.clone())
.build()
.unwrap();
let list_array = FixedSizeListArray::from(list_data);
assert_eq!(value_data.slice(3, 6), list_array.values().to_data());
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(2, list_array.len());
assert_eq!(0, list_array.null_count());
assert_eq!(3, list_array.value(0).as_primitive::<Int32Type>().value(0));
assert_eq!(3, list_array.value_offset(1));
assert_eq!(3, list_array.value_length());
}
#[test]
#[should_panic(expected = "assertion failed: end <= self.len()")]
// Different error messages, so skip for now
// https://github.com/apache/arrow-rs/issues/1545
#[cfg(not(feature = "force_validate"))]
fn test_fixed_size_list_array_unequal_children() {
// Construct a value array
let value_data = ArrayData::builder(DataType::Int32)
.len(8)
.add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7]))
.build()
.unwrap();
// Construct a list array from the above two
let list_data_type =
DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 3);
let list_data = unsafe {
ArrayData::builder(list_data_type)
.len(3)
.add_child_data(value_data)
.build_unchecked()
};
drop(FixedSizeListArray::from(list_data));
}
#[test]
fn test_fixed_size_list_array_slice() {
// Construct a value array
let value_data = ArrayData::builder(DataType::Int32)
.len(10)
.add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
.build()
.unwrap();
// Set null buts for the nested array:
// [[0, 1], null, null, [6, 7], [8, 9]]
// 01011001 00000001
let mut null_bits: [u8; 1] = [0; 1];
bit_util::set_bit(&mut null_bits, 0);
bit_util::set_bit(&mut null_bits, 3);
bit_util::set_bit(&mut null_bits, 4);
// Construct a fixed size list array from the above two
let list_data_type =
DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 2);
let list_data = ArrayData::builder(list_data_type)
.len(5)
.add_child_data(value_data.clone())
.null_bit_buffer(Some(Buffer::from(null_bits)))
.build()
.unwrap();
let list_array = FixedSizeListArray::from(list_data);
assert_eq!(value_data, list_array.values().to_data());
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(5, list_array.len());
assert_eq!(2, list_array.null_count());
assert_eq!(6, list_array.value_offset(3));
assert_eq!(2, list_array.value_length());
let sliced_array = list_array.slice(1, 4);
assert_eq!(4, sliced_array.len());
assert_eq!(2, sliced_array.null_count());
for i in 0..sliced_array.len() {
if bit_util::get_bit(&null_bits, 1 + i) {
assert!(sliced_array.is_valid(i));
} else {
assert!(sliced_array.is_null(i));
}
}
// Check offset and length for each non-null value.
let sliced_list_array = sliced_array
.as_any()
.downcast_ref::<FixedSizeListArray>()
.unwrap();
assert_eq!(2, sliced_list_array.value_length());
assert_eq!(4, sliced_list_array.value_offset(2));
assert_eq!(6, sliced_list_array.value_offset(3));
}
#[test]
#[should_panic(expected = "the offset of the new Buffer cannot exceed the existing length")]
fn test_fixed_size_list_array_index_out_of_bound() {
// Construct a value array
let value_data = ArrayData::builder(DataType::Int32)
.len(10)
.add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
.build()
.unwrap();
// Set null buts for the nested array:
// [[0, 1], null, null, [6, 7], [8, 9]]
// 01011001 00000001
let mut null_bits: [u8; 1] = [0; 1];
bit_util::set_bit(&mut null_bits, 0);
bit_util::set_bit(&mut null_bits, 3);
bit_util::set_bit(&mut null_bits, 4);
// Construct a fixed size list array from the above two
let list_data_type =
DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 2);
let list_data = ArrayData::builder(list_data_type)
.len(5)
.add_child_data(value_data)
.null_bit_buffer(Some(Buffer::from(null_bits)))
.build()
.unwrap();
let list_array = FixedSizeListArray::from(list_data);
list_array.value(10);
}
#[test]
fn test_fixed_size_list_constructors() {
let values = Arc::new(Int32Array::from_iter([
Some(1),
Some(2),
None,
None,
Some(3),
Some(4),
]));
let field = Arc::new(Field::new_list_field(DataType::Int32, true));
let list = FixedSizeListArray::new(field.clone(), 2, values.clone(), None);
assert_eq!(list.len(), 3);
let nulls = NullBuffer::new_null(3);
let list = FixedSizeListArray::new(field.clone(), 2, values.clone(), Some(nulls));
assert_eq!(list.len(), 3);
let list = FixedSizeListArray::new(field.clone(), 3, values.clone(), None);
assert_eq!(list.len(), 2);
let err = FixedSizeListArray::try_new(field.clone(), 4, values.clone(), None).unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: Incorrect length of values buffer for FixedSizeListArray, \
expected a multiple of 4 got 6",
);
let err =
FixedSizeListArray::try_new_with_length(field.clone(), 4, values.clone(), None, 1)
.unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: Incorrect length of values buffer for FixedSizeListArray, expected 4 got 6"
);
let err = FixedSizeListArray::try_new(field.clone(), -1, values.clone(), None).unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: Size cannot be negative, got -1"
);
let nulls = NullBuffer::new_null(2);
let err = FixedSizeListArray::try_new(field, 2, values.clone(), Some(nulls)).unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: Incorrect length of values buffer for FixedSizeListArray, expected 4 got 6"
);
let field = Arc::new(Field::new_list_field(DataType::Int32, false));
let err = FixedSizeListArray::try_new(field.clone(), 2, values.clone(), None).unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: Found unmasked nulls for non-nullable FixedSizeListArray field \"item\""
);
// Valid as nulls in child masked by parent
let nulls = NullBuffer::new(BooleanBuffer::new(Buffer::from([0b0000101]), 0, 3));
FixedSizeListArray::new(field, 2, values.clone(), Some(nulls));
let field = Arc::new(Field::new_list_field(DataType::Int64, true));
let err = FixedSizeListArray::try_new(field, 2, values, None).unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: FixedSizeListArray expected data type Int64 got Int32 for \"item\""
);
}
#[test]
fn degenerate_fixed_size_list() {
let field = Arc::new(Field::new_list_field(DataType::Int32, true));
let nulls = NullBuffer::new_null(2);
let values = new_empty_array(&DataType::Int32);
let list = FixedSizeListArray::new(field.clone(), 0, values.clone(), Some(nulls.clone()));
assert_eq!(list.len(), 2);
// Test invalid null buffer length.
let err = FixedSizeListArray::try_new_with_length(
field.clone(),
0,
values.clone(),
Some(nulls),
5,
)
.unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: Invalid null buffer for FixedSizeListArray, expected 5 found 2"
);
// Test non-empty values for degenerate list.
let non_empty_values = Arc::new(Int32Array::from(vec![1, 2, 3]));
let err =
FixedSizeListArray::try_new_with_length(field.clone(), 0, non_empty_values, None, 3)
.unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: An degenerate FixedSizeListArray should have no underlying values, found 3 values"
);
}
#[test]
fn test_fixed_size_list_new_null_len() {
let field = Arc::new(Field::new_list_field(DataType::Int32, true));
let array = FixedSizeListArray::new_null(field, 2, 5);
assert_eq!(array.len(), 5);
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+833
View File
@@ -0,0 +1,833 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::array::{get_offsets, print_long_array};
use crate::iterator::MapArrayIter;
use crate::{Array, ArrayAccessor, ArrayRef, ListArray, StringArray, StructArray, make_array};
use arrow_buffer::{ArrowNativeType, Buffer, NullBuffer, OffsetBuffer, ToByteSlice};
use arrow_data::{ArrayData, ArrayDataBuilder};
use arrow_schema::{ArrowError, DataType, Field, FieldRef};
use std::any::Any;
use std::sync::Arc;
/// An array of key-value maps
///
/// Keys should always be non-null, but values can be null.
///
/// [`MapArray`] is physically a [`ListArray`] of key values pairs stored as an `entries`
/// [`StructArray`] with 2 child fields.
///
/// See [`MapBuilder`](crate::builder::MapBuilder) for how to construct a [`MapArray`]
#[derive(Clone)]
pub struct MapArray {
data_type: DataType,
nulls: Option<NullBuffer>,
/// The [`StructArray`] that is the direct child of this array
entries: StructArray,
/// The start and end offsets of each entry
value_offsets: OffsetBuffer<i32>,
}
impl MapArray {
/// Create a new [`MapArray`] from the provided parts
///
/// See [`MapBuilder`](crate::builder::MapBuilder) for a higher-level interface
/// to construct a [`MapArray`]
///
/// # Errors
///
/// Errors if
///
/// * `offsets.len() - 1 != nulls.len()`
/// * `offsets.last() > entries.len()`
/// * `field.is_nullable()`
/// * `entries.null_count() != 0`
/// * `entries.columns().len() != 2`
/// * `field.data_type() != entries.data_type()`
pub fn try_new(
field: FieldRef,
offsets: OffsetBuffer<i32>,
entries: StructArray,
nulls: Option<NullBuffer>,
ordered: bool,
) -> Result<Self, ArrowError> {
let len = offsets.len() - 1; // Offsets guaranteed to not be empty
let end_offset = offsets.last().unwrap().as_usize();
// don't need to check other values of `offsets` because they are checked
// during construction of `OffsetBuffer`
if end_offset > entries.len() {
return Err(ArrowError::InvalidArgumentError(format!(
"Max offset of {end_offset} exceeds length of entries {}",
entries.len()
)));
}
if let Some(n) = nulls.as_ref() {
if n.len() != len {
return Err(ArrowError::InvalidArgumentError(format!(
"Incorrect length of null buffer for MapArray, expected {len} got {}",
n.len(),
)));
}
}
if field.is_nullable() || entries.null_count() != 0 {
return Err(ArrowError::InvalidArgumentError(
"MapArray entries cannot contain nulls".to_string(),
));
}
if field.data_type() != entries.data_type() {
return Err(ArrowError::InvalidArgumentError(format!(
"MapArray expected data type {} got {} for {:?}",
field.data_type(),
entries.data_type(),
field.name()
)));
}
if entries.columns().len() != 2 {
return Err(ArrowError::InvalidArgumentError(format!(
"MapArray entries must contain two children, got {}",
entries.columns().len()
)));
}
Ok(Self {
data_type: DataType::Map(field, ordered),
nulls,
entries,
value_offsets: offsets,
})
}
/// Create a new [`MapArray`] from the provided parts
///
/// See [`MapBuilder`](crate::builder::MapBuilder) for a higher-level interface
/// to construct a [`MapArray`]
///
/// # Panics
///
/// Panics if [`Self::try_new`] returns an error
pub fn new(
field: FieldRef,
offsets: OffsetBuffer<i32>,
entries: StructArray,
nulls: Option<NullBuffer>,
ordered: bool,
) -> Self {
Self::try_new(field, offsets, entries, nulls, ordered).unwrap()
}
/// Deconstruct this array into its constituent parts
pub fn into_parts(
self,
) -> (
FieldRef,
OffsetBuffer<i32>,
StructArray,
Option<NullBuffer>,
bool,
) {
let (f, ordered) = match self.data_type {
DataType::Map(f, ordered) => (f, ordered),
_ => unreachable!(),
};
(f, self.value_offsets, self.entries, self.nulls, ordered)
}
/// Returns a reference to the offsets of this map
///
/// Unlike [`Self::value_offsets`] this returns the [`OffsetBuffer`]
/// allowing for zero-copy cloning
#[inline]
pub fn offsets(&self) -> &OffsetBuffer<i32> {
&self.value_offsets
}
/// Returns a reference to the keys of this map
pub fn keys(&self) -> &ArrayRef {
self.entries.column(0)
}
/// Returns a reference to the values of this map
pub fn values(&self) -> &ArrayRef {
self.entries.column(1)
}
/// Returns a reference to the [`StructArray`] entries of this map
pub fn entries(&self) -> &StructArray {
&self.entries
}
/// Returns a reference to the fields of the [`StructArray`] that backs this map.
pub fn entries_fields(&self) -> (&Field, &Field) {
let fields = self.entries.fields().iter().collect::<Vec<_>>();
let fields = TryInto::<[&FieldRef; 2]>::try_into(fields)
.expect("Every map has a key and value field");
(fields[0].as_ref(), fields[1].as_ref())
}
/// Returns the data type of the map's keys.
pub fn key_type(&self) -> &DataType {
self.keys().data_type()
}
/// Returns the data type of the map's values.
pub fn value_type(&self) -> &DataType {
self.values().data_type()
}
/// Returns ith value of this map array.
///
/// Note: This method does not check for nulls and the value is arbitrary
/// if [`is_null`](Self::is_null) returns true for the index.
///
/// # Safety
/// Caller must ensure that the index is within the array bounds
pub unsafe fn value_unchecked(&self, i: usize) -> StructArray {
let end = *unsafe { self.value_offsets().get_unchecked(i + 1) };
let start = *unsafe { self.value_offsets().get_unchecked(i) };
self.entries
.slice(start.to_usize().unwrap(), (end - start).to_usize().unwrap())
}
/// Returns ith value of this map array.
///
/// This is a [`StructArray`] containing two fields
///
/// Note: This method does not check for nulls and the value is arbitrary
/// (but still well-defined) if [`is_null`](Self::is_null) returns true for the index.
///
/// # Panics
/// Panics if index `i` is out of bounds
pub fn value(&self, i: usize) -> StructArray {
let end = self.value_offsets()[i + 1] as usize;
let start = self.value_offsets()[i] as usize;
self.entries.slice(start, end - start)
}
/// Returns the offset values in the offsets buffer
#[inline]
pub fn value_offsets(&self) -> &[i32] {
&self.value_offsets
}
/// Returns the length for value at index `i`.
#[inline]
pub fn value_length(&self, i: usize) -> i32 {
let offsets = self.value_offsets();
offsets[i + 1] - offsets[i]
}
/// Returns a zero-copy slice of this array with the indicated offset and length.
pub fn slice(&self, offset: usize, length: usize) -> Self {
Self {
data_type: self.data_type.clone(),
nulls: self.nulls.as_ref().map(|n| n.slice(offset, length)),
entries: self.entries.clone(),
value_offsets: self.value_offsets.slice(offset, length),
}
}
/// constructs a new iterator
pub fn iter(&self) -> MapArrayIter<'_> {
MapArrayIter::new(self)
}
}
impl From<ArrayData> for MapArray {
fn from(data: ArrayData) -> Self {
Self::try_new_from_array_data(data)
.expect("Expected infallible creation of MapArray from ArrayData failed")
}
}
impl From<MapArray> for ArrayData {
fn from(array: MapArray) -> Self {
let len = array.len();
let builder = ArrayDataBuilder::new(array.data_type)
.len(len)
.nulls(array.nulls)
.buffers(vec![array.value_offsets.into_inner().into_inner()])
.child_data(vec![array.entries.to_data()]);
unsafe { builder.build_unchecked() }
}
}
impl MapArray {
fn try_new_from_array_data(data: ArrayData) -> Result<Self, ArrowError> {
if !matches!(data.data_type(), DataType::Map(_, _)) {
return Err(ArrowError::InvalidArgumentError(format!(
"MapArray expected ArrayData with DataType::Map got {}",
data.data_type()
)));
}
if data.buffers().len() != 1 {
return Err(ArrowError::InvalidArgumentError(format!(
"MapArray data should contain a single buffer only (value offsets), had {}",
data.len()
)));
}
if data.child_data().len() != 1 {
return Err(ArrowError::InvalidArgumentError(format!(
"MapArray should contain a single child array (values array), had {}",
data.child_data().len()
)));
}
let entries = data.child_data()[0].clone();
if let DataType::Struct(fields) = entries.data_type() {
if fields.len() != 2 {
return Err(ArrowError::InvalidArgumentError(format!(
"MapArray should contain a struct array with 2 fields, have {} fields",
fields.len()
)));
}
} else {
return Err(ArrowError::InvalidArgumentError(format!(
"MapArray should contain a struct array child, found {:?}",
entries.data_type()
)));
}
let entries = entries.into();
// SAFETY:
// ArrayData is valid, and verified type above
let value_offsets = unsafe { get_offsets(&data) };
Ok(Self {
data_type: data.data_type().clone(),
nulls: data.nulls().cloned(),
entries,
value_offsets,
})
}
/// Creates map array from provided keys, values and entry_offsets.
pub fn new_from_strings<'a>(
keys: impl Iterator<Item = &'a str>,
values: &dyn Array,
entry_offsets: &[u32],
) -> Result<Self, ArrowError> {
let entry_offsets_buffer = Buffer::from(entry_offsets.to_byte_slice());
let keys_data = StringArray::from_iter_values(keys);
let keys_field = Arc::new(Field::new("keys", DataType::Utf8, false));
let values_field = Arc::new(Field::new(
"values",
values.data_type().clone(),
values.null_count() > 0,
));
let entry_struct = StructArray::from(vec![
(keys_field, Arc::new(keys_data) as ArrayRef),
(values_field, make_array(values.to_data())),
]);
let map_data_type = DataType::Map(
Arc::new(Field::new(
"entries",
entry_struct.data_type().clone(),
false,
)),
false,
);
let map_data = ArrayData::builder(map_data_type)
.len(entry_offsets.len() - 1)
.add_buffer(entry_offsets_buffer)
.add_child_data(entry_struct.into_data())
.build()?;
Ok(MapArray::from(map_data))
}
}
/// SAFETY: Correctly implements the contract of Arrow Arrays
unsafe impl Array for MapArray {
fn as_any(&self) -> &dyn Any {
self
}
fn to_data(&self) -> ArrayData {
self.clone().into_data()
}
fn into_data(self) -> ArrayData {
self.into()
}
fn data_type(&self) -> &DataType {
&self.data_type
}
fn slice(&self, offset: usize, length: usize) -> ArrayRef {
Arc::new(self.slice(offset, length))
}
fn len(&self) -> usize {
self.value_offsets.len() - 1
}
fn is_empty(&self) -> bool {
self.value_offsets.len() <= 1
}
fn shrink_to_fit(&mut self) {
if let Some(nulls) = &mut self.nulls {
nulls.shrink_to_fit();
}
self.entries.shrink_to_fit();
self.value_offsets.shrink_to_fit();
}
fn offset(&self) -> usize {
0
}
fn nulls(&self) -> Option<&NullBuffer> {
self.nulls.as_ref()
}
fn logical_null_count(&self) -> usize {
// More efficient that the default implementation
self.null_count()
}
fn get_buffer_memory_size(&self) -> usize {
let mut size = self.entries.get_buffer_memory_size();
size += self.value_offsets.inner().inner().capacity();
if let Some(n) = self.nulls.as_ref() {
size += n.buffer().capacity();
}
size
}
fn get_array_memory_size(&self) -> usize {
let mut size = std::mem::size_of::<Self>() + self.entries.get_array_memory_size();
size += self.value_offsets.inner().inner().capacity();
if let Some(n) = self.nulls.as_ref() {
size += n.buffer().capacity();
}
size
}
}
impl ArrayAccessor for &MapArray {
type Item = StructArray;
fn value(&self, index: usize) -> Self::Item {
MapArray::value(self, index)
}
unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
MapArray::value(self, index)
}
}
impl std::fmt::Debug for MapArray {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(f, "MapArray\n[\n")?;
print_long_array(self, f, |array, index, f| {
std::fmt::Debug::fmt(&array.value(index), f)
})?;
write!(f, "]")
}
}
impl From<MapArray> for ListArray {
fn from(value: MapArray) -> Self {
let field = match value.data_type() {
DataType::Map(field, _) => field,
_ => unreachable!("This should be a map type."),
};
let data_type = DataType::List(field.clone());
let builder = value.into_data().into_builder().data_type(data_type);
let array_data = unsafe { builder.build_unchecked() };
ListArray::from(array_data)
}
}
#[cfg(test)]
mod tests {
use crate::cast::AsArray;
use crate::types::UInt32Type;
use crate::{Int32Array, UInt32Array};
use arrow_schema::Fields;
use super::*;
fn create_from_buffers() -> MapArray {
// Construct key and values
let keys_data = ArrayData::builder(DataType::Int32)
.len(8)
.add_buffer(Buffer::from([0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
.build()
.unwrap();
let values_data = ArrayData::builder(DataType::UInt32)
.len(8)
.add_buffer(Buffer::from(
[0u32, 10, 20, 30, 40, 50, 60, 70].to_byte_slice(),
))
.build()
.unwrap();
// Construct a buffer for value offsets, for the nested array:
// [[0, 1, 2], [3, 4, 5], [6, 7]]
let entry_offsets = Buffer::from([0, 3, 6, 8].to_byte_slice());
let keys = Arc::new(Field::new("keys", DataType::Int32, false));
let values = Arc::new(Field::new("values", DataType::UInt32, false));
let entry_struct = StructArray::from(vec![
(keys, make_array(keys_data)),
(values, make_array(values_data)),
]);
// Construct a map array from the above two
let map_data_type = DataType::Map(
Arc::new(Field::new(
"entries",
entry_struct.data_type().clone(),
false,
)),
false,
);
let map_data = ArrayData::builder(map_data_type)
.len(3)
.add_buffer(entry_offsets)
.add_child_data(entry_struct.into_data())
.build()
.unwrap();
MapArray::from(map_data)
}
#[test]
fn test_map_array() {
// Construct key and values
let key_data = ArrayData::builder(DataType::Int32)
.len(8)
.add_buffer(Buffer::from([0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
.build()
.unwrap();
let value_data = ArrayData::builder(DataType::UInt32)
.len(8)
.add_buffer(Buffer::from(
[0u32, 10, 20, 0, 40, 0, 60, 70].to_byte_slice(),
))
.null_bit_buffer(Some(Buffer::from(&[0b11010110])))
.build()
.unwrap();
// Construct a buffer for value offsets, for the nested array:
// [[0, 1, 2], [3, 4, 5], [6, 7]]
let entry_offsets = Buffer::from([0, 3, 6, 8].to_byte_slice());
let keys_field = Arc::new(Field::new("keys", DataType::Int32, false));
let values_field = Arc::new(Field::new("values", DataType::UInt32, true));
let entry_struct = StructArray::from(vec![
(keys_field.clone(), make_array(key_data)),
(values_field.clone(), make_array(value_data.clone())),
]);
// Construct a map array from the above two
let map_data_type = DataType::Map(
Arc::new(Field::new(
"entries",
entry_struct.data_type().clone(),
false,
)),
false,
);
let map_data = ArrayData::builder(map_data_type)
.len(3)
.add_buffer(entry_offsets)
.add_child_data(entry_struct.into_data())
.build()
.unwrap();
let map_array = MapArray::from(map_data);
assert_eq!(value_data, map_array.values().to_data());
assert_eq!(&DataType::UInt32, map_array.value_type());
assert_eq!(3, map_array.len());
assert_eq!(0, map_array.null_count());
assert_eq!(6, map_array.value_offsets()[2]);
assert_eq!(2, map_array.value_length(2));
let key_array = Arc::new(Int32Array::from(vec![0, 1, 2])) as ArrayRef;
let value_array =
Arc::new(UInt32Array::from(vec![None, Some(10u32), Some(20)])) as ArrayRef;
let struct_array = StructArray::from(vec![
(keys_field.clone(), key_array),
(values_field.clone(), value_array),
]);
assert_eq!(
struct_array,
StructArray::from(map_array.value(0).into_data())
);
assert_eq!(
&struct_array,
unsafe { map_array.value_unchecked(0) }
.as_any()
.downcast_ref::<StructArray>()
.unwrap()
);
for i in 0..3 {
assert!(map_array.is_valid(i));
assert!(!map_array.is_null(i));
}
// Now test with a non-zero offset
let map_array = map_array.slice(1, 2);
assert_eq!(value_data, map_array.values().to_data());
assert_eq!(&DataType::UInt32, map_array.value_type());
assert_eq!(2, map_array.len());
assert_eq!(0, map_array.null_count());
assert_eq!(6, map_array.value_offsets()[1]);
assert_eq!(2, map_array.value_length(1));
let key_array = Arc::new(Int32Array::from(vec![3, 4, 5])) as ArrayRef;
let value_array = Arc::new(UInt32Array::from(vec![None, Some(40), None])) as ArrayRef;
let struct_array =
StructArray::from(vec![(keys_field, key_array), (values_field, value_array)]);
assert_eq!(
&struct_array,
map_array
.value(0)
.as_any()
.downcast_ref::<StructArray>()
.unwrap()
);
assert_eq!(
&struct_array,
unsafe { map_array.value_unchecked(0) }
.as_any()
.downcast_ref::<StructArray>()
.unwrap()
);
}
#[test]
#[ignore = "Test fails because slice of <list<struct>> is still buggy"]
fn test_map_array_slice() {
let map_array = create_from_buffers();
let sliced_array = map_array.slice(1, 2);
assert_eq!(2, sliced_array.len());
assert_eq!(1, sliced_array.offset());
let sliced_array_data = sliced_array.to_data();
for array_data in sliced_array_data.child_data() {
assert_eq!(array_data.offset(), 1);
}
// Check offset and length for each non-null value.
let sliced_map_array = sliced_array.as_any().downcast_ref::<MapArray>().unwrap();
assert_eq!(3, sliced_map_array.value_offsets()[0]);
assert_eq!(3, sliced_map_array.value_length(0));
assert_eq!(6, sliced_map_array.value_offsets()[1]);
assert_eq!(2, sliced_map_array.value_length(1));
// Construct key and values
let keys_data = ArrayData::builder(DataType::Int32)
.len(5)
.add_buffer(Buffer::from([3, 4, 5, 6, 7].to_byte_slice()))
.build()
.unwrap();
let values_data = ArrayData::builder(DataType::UInt32)
.len(5)
.add_buffer(Buffer::from([30u32, 40, 50, 60, 70].to_byte_slice()))
.build()
.unwrap();
// Construct a buffer for value offsets, for the nested array:
// [[3, 4, 5], [6, 7]]
let entry_offsets = Buffer::from([0, 3, 5].to_byte_slice());
let keys = Arc::new(Field::new("keys", DataType::Int32, false));
let values = Arc::new(Field::new("values", DataType::UInt32, false));
let entry_struct = StructArray::from(vec![
(keys, make_array(keys_data)),
(values, make_array(values_data)),
]);
// Construct a map array from the above two
let map_data_type = DataType::Map(
Arc::new(Field::new(
"entries",
entry_struct.data_type().clone(),
false,
)),
false,
);
let expected_map_data = ArrayData::builder(map_data_type)
.len(2)
.add_buffer(entry_offsets)
.add_child_data(entry_struct.into_data())
.build()
.unwrap();
let expected_map_array = MapArray::from(expected_map_data);
assert_eq!(&expected_map_array, sliced_map_array)
}
#[test]
#[should_panic(expected = "index out of bounds: the len is ")]
fn test_map_array_index_out_of_bound() {
let map_array = create_from_buffers();
map_array.value(map_array.len());
}
#[test]
#[should_panic(expected = "MapArray expected ArrayData with DataType::Map got Dictionary")]
fn test_from_array_data_validation() {
// A DictionaryArray has similar buffer layout to a MapArray
// but the meaning of the values differs
let struct_t = DataType::Struct(Fields::from(vec![
Field::new("keys", DataType::Int32, true),
Field::new("values", DataType::UInt32, true),
]));
let dict_t = DataType::Dictionary(Box::new(DataType::Int32), Box::new(struct_t));
let _ = MapArray::from(ArrayData::new_empty(&dict_t));
}
#[test]
fn test_new_from_strings() {
let keys = vec!["a", "b", "c", "d", "e", "f", "g", "h"];
let values_data = UInt32Array::from(vec![0u32, 10, 20, 30, 40, 50, 60, 70]);
// Construct a buffer for value offsets, for the nested array:
// [[a, b, c], [d, e, f], [g, h]]
let entry_offsets = [0, 3, 6, 8];
let map_array =
MapArray::new_from_strings(keys.clone().into_iter(), &values_data, &entry_offsets)
.unwrap();
assert_eq!(
&values_data,
map_array.values().as_primitive::<UInt32Type>()
);
assert_eq!(&DataType::UInt32, map_array.value_type());
assert_eq!(3, map_array.len());
assert_eq!(0, map_array.null_count());
assert_eq!(6, map_array.value_offsets()[2]);
assert_eq!(2, map_array.value_length(2));
let key_array = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
let value_array = Arc::new(UInt32Array::from(vec![0u32, 10, 20])) as ArrayRef;
let keys_field = Arc::new(Field::new("keys", DataType::Utf8, false));
let values_field = Arc::new(Field::new("values", DataType::UInt32, false));
let struct_array =
StructArray::from(vec![(keys_field, key_array), (values_field, value_array)]);
assert_eq!(
struct_array,
StructArray::from(map_array.value(0).into_data())
);
assert_eq!(
&struct_array,
unsafe { map_array.value_unchecked(0) }
.as_any()
.downcast_ref::<StructArray>()
.unwrap()
);
for i in 0..3 {
assert!(map_array.is_valid(i));
assert!(!map_array.is_null(i));
}
}
#[test]
fn test_try_new() {
let offsets = OffsetBuffer::new(vec![0, 1, 4, 5].into());
let fields = Fields::from(vec![
Field::new("key", DataType::Int32, false),
Field::new("values", DataType::Int32, false),
]);
let columns = vec![
Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
];
let entries = StructArray::new(fields.clone(), columns, None);
let field = Arc::new(Field::new("entries", DataType::Struct(fields), false));
MapArray::new(field.clone(), offsets.clone(), entries.clone(), None, false);
let nulls = NullBuffer::new_null(3);
MapArray::new(field.clone(), offsets, entries.clone(), Some(nulls), false);
let nulls = NullBuffer::new_null(3);
let offsets = OffsetBuffer::new(vec![0, 1, 2, 4, 5].into());
let err = MapArray::try_new(
field.clone(),
offsets.clone(),
entries.clone(),
Some(nulls),
false,
)
.unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: Incorrect length of null buffer for MapArray, expected 4 got 3"
);
let err = MapArray::try_new(field, offsets.clone(), entries.slice(0, 2), None, false)
.unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: Max offset of 5 exceeds length of entries 2"
);
let field = Arc::new(Field::new("element", DataType::Int64, false));
let err = MapArray::try_new(field, offsets.clone(), entries, None, false)
.unwrap_err()
.to_string();
assert!(
err.starts_with("Invalid argument error: MapArray expected data type Int64 got Struct"),
"{err}"
);
let fields = Fields::from(vec![
Field::new("a", DataType::Int32, false),
Field::new("b", DataType::Int32, false),
Field::new("c", DataType::Int32, false),
]);
let columns = vec![
Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
];
let s = StructArray::new(fields.clone(), columns, None);
let field = Arc::new(Field::new("entries", DataType::Struct(fields), false));
let err = MapArray::try_new(field, offsets, s, None, false).unwrap_err();
assert_eq!(
err.to_string(),
"Invalid argument error: MapArray entries must contain two children, got 3"
);
}
}
File diff suppressed because it is too large Load Diff
+236
View File
@@ -0,0 +1,236 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Contains the `NullArray` type.
use crate::builder::NullBuilder;
use crate::{Array, ArrayRef};
use arrow_buffer::buffer::NullBuffer;
use arrow_data::{ArrayData, ArrayDataBuilder};
use arrow_schema::DataType;
use std::any::Any;
use std::sync::Arc;
/// An array of [null values](https://arrow.apache.org/docs/format/Columnar.html#null-layout)
///
/// A `NullArray` is a simplified array where all values are null.
///
/// # Example: Create an array
///
/// ```
/// use arrow_array::{Array, NullArray};
///
/// let array = NullArray::new(10);
///
/// assert!(array.is_nullable());
/// assert_eq!(array.len(), 10);
/// assert_eq!(array.null_count(), 0);
/// assert_eq!(array.logical_null_count(), 10);
/// assert_eq!(array.logical_nulls().unwrap().null_count(), 10);
/// ```
#[derive(Clone)]
pub struct NullArray {
len: usize,
}
impl NullArray {
/// Create a new [`NullArray`] of the specified length
///
/// *Note*: Use [`crate::array::new_null_array`] if you need an array of some
/// other [`DataType`].
///
pub fn new(length: usize) -> Self {
Self { len: length }
}
/// Returns a zero-copy slice of this array with the indicated offset and length.
pub fn slice(&self, offset: usize, len: usize) -> Self {
assert!(
offset.saturating_add(len) <= self.len,
"the length + offset of the sliced BooleanBuffer cannot exceed the existing length"
);
Self { len }
}
/// Returns a new null array builder
///
/// Note that the `capacity` parameter to this function is _deprecated_. It
/// now does nothing, and will be removed in a future version.
pub fn builder(_capacity: usize) -> NullBuilder {
NullBuilder::new()
}
}
/// SAFETY: Correctly implements the contract of Arrow Arrays
unsafe impl Array for NullArray {
fn as_any(&self) -> &dyn Any {
self
}
fn to_data(&self) -> ArrayData {
self.clone().into()
}
fn into_data(self) -> ArrayData {
self.into()
}
fn data_type(&self) -> &DataType {
&DataType::Null
}
fn slice(&self, offset: usize, length: usize) -> ArrayRef {
Arc::new(self.slice(offset, length))
}
fn len(&self) -> usize {
self.len
}
fn is_empty(&self) -> bool {
self.len == 0
}
fn offset(&self) -> usize {
0
}
fn nulls(&self) -> Option<&NullBuffer> {
None
}
fn logical_nulls(&self) -> Option<NullBuffer> {
(self.len != 0).then(|| NullBuffer::new_null(self.len))
}
fn is_nullable(&self) -> bool {
!self.is_empty()
}
fn logical_null_count(&self) -> usize {
self.len
}
fn get_buffer_memory_size(&self) -> usize {
0
}
fn get_array_memory_size(&self) -> usize {
std::mem::size_of::<Self>()
}
}
impl From<ArrayData> for NullArray {
fn from(data: ArrayData) -> Self {
assert_eq!(
data.data_type(),
&DataType::Null,
"NullArray data type should be Null"
);
assert_eq!(
data.buffers().len(),
0,
"NullArray data should contain 0 buffers"
);
assert!(
data.nulls().is_none(),
"NullArray data should not contain a null buffer, as no buffers are required"
);
Self { len: data.len() }
}
}
impl From<NullArray> for ArrayData {
fn from(array: NullArray) -> Self {
let builder = ArrayDataBuilder::new(DataType::Null).len(array.len);
unsafe { builder.build_unchecked() }
}
}
impl std::fmt::Debug for NullArray {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(f, "NullArray({})", self.len())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{Int64Array, StructArray, make_array};
use arrow_data::transform::MutableArrayData;
use arrow_schema::Field;
#[test]
fn test_null_array() {
let null_arr = NullArray::new(32);
assert_eq!(null_arr.len(), 32);
assert_eq!(null_arr.null_count(), 0);
assert_eq!(null_arr.logical_null_count(), 32);
assert_eq!(null_arr.logical_nulls().unwrap().null_count(), 32);
assert!(null_arr.is_valid(0));
assert!(null_arr.is_nullable());
}
#[test]
fn test_null_array_slice() {
let array1 = NullArray::new(32);
let array2 = array1.slice(8, 16);
assert_eq!(array2.len(), 16);
assert_eq!(array2.null_count(), 0);
assert_eq!(array2.logical_null_count(), 16);
assert_eq!(array2.logical_nulls().unwrap().null_count(), 16);
assert!(array2.is_valid(0));
assert!(array2.is_nullable());
}
#[test]
fn test_debug_null_array() {
let array = NullArray::new(1024 * 1024);
assert_eq!(format!("{array:?}"), "NullArray(1048576)");
}
#[test]
fn test_null_array_with_parent_null_buffer() {
let null_array = NullArray::new(1);
let int_array = Int64Array::from(vec![42]);
let fields = vec![
Field::new("a", DataType::Int64, true),
Field::new("b", DataType::Null, true),
];
let struct_array_data = ArrayData::builder(DataType::Struct(fields.into()))
.len(1)
.add_child_data(int_array.to_data())
.add_child_data(null_array.to_data())
.build()
.unwrap();
let mut mutable = MutableArrayData::new(vec![&struct_array_data], true, 1);
// Simulate a NULL value in the parent array, for instance, if array being queried by
// invalid index
mutable.extend_nulls(1);
let data = mutable.freeze();
let struct_array = Arc::new(StructArray::from(data.clone()));
assert!(make_array(data) == struct_array);
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,562 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::types::GenericStringType;
use crate::{GenericBinaryArray, GenericByteArray, GenericListArray, OffsetSizeTrait};
use arrow_schema::ArrowError;
/// A [`GenericByteArray`] for storing `str`
pub type GenericStringArray<OffsetSize> = GenericByteArray<GenericStringType<OffsetSize>>;
impl<OffsetSize: OffsetSizeTrait> GenericStringArray<OffsetSize> {
/// Returns the number of `Unicode Scalar Value` in the string at index `i`.
/// # Performance
/// This function has `O(n)` time complexity where `n` is the string length.
/// If you can make sure that all chars in the string are in the range `U+0x0000` ~ `U+0x007F`,
/// please use the function [`value_length`](#method.value_length) which has O(1) time complexity.
pub fn num_chars(&self, i: usize) -> usize {
self.value(i).chars().count()
}
/// Returns an iterator that returns the values of `array.value(i)` for an iterator with each element `i`
pub fn take_iter<'a>(
&'a self,
indexes: impl Iterator<Item = Option<usize>> + 'a,
) -> impl Iterator<Item = Option<&'a str>> {
indexes.map(|opt_index| opt_index.map(|index| self.value(index)))
}
/// Returns an iterator that returns the values of `array.value(i)` for an iterator with each element `i`
/// # Safety
///
/// caller must ensure that the indexes in the iterator are less than the `array.len()`
pub unsafe fn take_iter_unchecked<'a>(
&'a self,
indexes: impl Iterator<Item = Option<usize>> + 'a,
) -> impl Iterator<Item = Option<&'a str>> {
indexes.map(|opt_index| opt_index.map(|index| unsafe { self.value_unchecked(index) }))
}
/// Fallibly creates a [`GenericStringArray`] from a [`GenericBinaryArray`] returning
/// an error if [`GenericBinaryArray`] contains invalid UTF-8 data
pub fn try_from_binary(v: GenericBinaryArray<OffsetSize>) -> Result<Self, ArrowError> {
let (offsets, values, nulls) = v.into_parts();
Self::try_new(offsets, values, nulls)
}
}
impl<OffsetSize: OffsetSizeTrait> From<GenericListArray<OffsetSize>>
for GenericStringArray<OffsetSize>
{
fn from(v: GenericListArray<OffsetSize>) -> Self {
GenericBinaryArray::<OffsetSize>::from(v).into()
}
}
impl<OffsetSize: OffsetSizeTrait> From<GenericBinaryArray<OffsetSize>>
for GenericStringArray<OffsetSize>
{
fn from(v: GenericBinaryArray<OffsetSize>) -> Self {
Self::try_from_binary(v).unwrap()
}
}
impl<OffsetSize: OffsetSizeTrait> From<Vec<Option<&str>>> for GenericStringArray<OffsetSize> {
fn from(v: Vec<Option<&str>>) -> Self {
v.into_iter().collect()
}
}
impl<OffsetSize: OffsetSizeTrait> From<Vec<&str>> for GenericStringArray<OffsetSize> {
fn from(v: Vec<&str>) -> Self {
Self::from_iter_values(v)
}
}
impl<OffsetSize: OffsetSizeTrait> From<Vec<Option<String>>> for GenericStringArray<OffsetSize> {
fn from(v: Vec<Option<String>>) -> Self {
v.into_iter().collect()
}
}
impl<OffsetSize: OffsetSizeTrait> From<Vec<String>> for GenericStringArray<OffsetSize> {
fn from(v: Vec<String>) -> Self {
Self::from_iter_values(v)
}
}
/// A [`GenericStringArray`] of `str` using `i32` offsets
///
/// # Examples
///
/// Construction
///
/// ```
/// # use arrow_array::StringArray;
/// // Create from Vec<Option<&str>>
/// let arr = StringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]);
/// // Create from Vec<&str>
/// let arr = StringArray::from(vec!["foo", "bar", "baz"]);
/// // Create from iter/collect (requires Option<&str>)
/// let arr: StringArray = std::iter::repeat(Some("foo")).take(10).collect();
/// ```
///
/// Construction and Access
///
/// ```
/// # use arrow_array::StringArray;
/// let array = StringArray::from(vec![Some("foo"), None, Some("bar")]);
/// assert_eq!(array.value(0), "foo");
/// ```
///
/// See [`GenericByteArray`] for more information and examples
pub type StringArray = GenericStringArray<i32>;
/// A [`GenericStringArray`] of `str` using `i64` offsets
///
/// # Examples
///
/// Construction
///
/// ```
/// # use arrow_array::LargeStringArray;
/// // Create from Vec<Option<&str>>
/// let arr = LargeStringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]);
/// // Create from Vec<&str>
/// let arr = LargeStringArray::from(vec!["foo", "bar", "baz"]);
/// // Create from iter/collect (requires Option<&str>)
/// let arr: LargeStringArray = std::iter::repeat(Some("foo")).take(10).collect();
/// ```
///
/// Construction and Access
///
/// ```
/// use arrow_array::LargeStringArray;
/// let array = LargeStringArray::from(vec![Some("foo"), None, Some("bar")]);
/// assert_eq!(array.value(2), "bar");
/// ```
///
/// See [`GenericByteArray`] for more information and examples
pub type LargeStringArray = GenericStringArray<i64>;
#[cfg(test)]
mod tests {
use super::*;
use crate::Array;
use crate::builder::{ListBuilder, PrimitiveBuilder, StringBuilder};
use crate::types::UInt8Type;
use arrow_buffer::Buffer;
use arrow_data::ArrayData;
use arrow_schema::{DataType, Field};
use std::sync::Arc;
#[test]
fn test_string_array_from_u8_slice() {
let values: Vec<&str> = vec!["hello", "", "A£ऀ𖼚𝌆৩ƐZ"];
// Array data: ["hello", "", "A£ऀ𖼚𝌆৩ƐZ"]
let string_array = StringArray::from(values);
assert_eq!(3, string_array.len());
assert_eq!(0, string_array.null_count());
assert_eq!("hello", string_array.value(0));
assert_eq!("hello", unsafe { string_array.value_unchecked(0) });
assert_eq!("", string_array.value(1));
assert_eq!("", unsafe { string_array.value_unchecked(1) });
assert_eq!("A£ऀ𖼚𝌆৩ƐZ", string_array.value(2));
assert_eq!("A£ऀ𖼚𝌆৩ƐZ", unsafe {
string_array.value_unchecked(2)
});
assert_eq!(20, string_array.value_length(2)); // 1 + 2 + 3 + 4 + 4 + 3 + 2 + 1
assert_eq!(8, string_array.num_chars(2));
for i in 0..3 {
assert!(string_array.is_valid(i));
assert!(!string_array.is_null(i));
}
}
#[test]
#[should_panic(expected = "StringArray expects DataType::Utf8")]
fn test_string_array_from_int() {
let array = LargeStringArray::from(vec!["a", "b"]);
drop(StringArray::from(array.into_data()));
}
#[test]
fn test_large_string_array_from_u8_slice() {
let values: Vec<&str> = vec!["hello", "", "A£ऀ𖼚𝌆৩ƐZ"];
// Array data: ["hello", "", "A£ऀ𖼚𝌆৩ƐZ"]
let string_array = LargeStringArray::from(values);
assert_eq!(3, string_array.len());
assert_eq!(0, string_array.null_count());
assert_eq!("hello", string_array.value(0));
assert_eq!("hello", unsafe { string_array.value_unchecked(0) });
assert_eq!("", string_array.value(1));
assert_eq!("", unsafe { string_array.value_unchecked(1) });
assert_eq!("A£ऀ𖼚𝌆৩ƐZ", string_array.value(2));
assert_eq!("A£ऀ𖼚𝌆৩ƐZ", unsafe {
string_array.value_unchecked(2)
});
assert_eq!(5, string_array.value_offsets()[2]);
assert_eq!(20, string_array.value_length(2)); // 1 + 2 + 3 + 4 + 4 + 3 + 2 + 1
assert_eq!(8, string_array.num_chars(2));
for i in 0..3 {
assert!(string_array.is_valid(i));
assert!(!string_array.is_null(i));
}
}
#[test]
fn test_nested_string_array() {
let string_builder = StringBuilder::with_capacity(3, 10);
let mut list_of_string_builder = ListBuilder::new(string_builder);
list_of_string_builder.values().append_value("foo");
list_of_string_builder.values().append_value("bar");
list_of_string_builder.append(true);
list_of_string_builder.values().append_value("foobar");
list_of_string_builder.append(true);
let list_of_strings = list_of_string_builder.finish();
assert_eq!(list_of_strings.len(), 2);
let first_slot = list_of_strings.value(0);
let first_list = first_slot.as_any().downcast_ref::<StringArray>().unwrap();
assert_eq!(first_list.len(), 2);
assert_eq!(first_list.value(0), "foo");
assert_eq!(unsafe { first_list.value_unchecked(0) }, "foo");
assert_eq!(first_list.value(1), "bar");
assert_eq!(unsafe { first_list.value_unchecked(1) }, "bar");
let second_slot = list_of_strings.value(1);
let second_list = second_slot.as_any().downcast_ref::<StringArray>().unwrap();
assert_eq!(second_list.len(), 1);
assert_eq!(second_list.value(0), "foobar");
assert_eq!(unsafe { second_list.value_unchecked(0) }, "foobar");
}
#[test]
#[should_panic(
expected = "Trying to access an element at index 4 from a StringArray of length 3"
)]
fn test_string_array_get_value_index_out_of_bound() {
let values: [u8; 12] = [
b'h', b'e', b'l', b'l', b'o', b'p', b'a', b'r', b'q', b'u', b'e', b't',
];
let offsets: [i32; 4] = [0, 5, 5, 12];
let array_data = ArrayData::builder(DataType::Utf8)
.len(3)
.add_buffer(Buffer::from_slice_ref(offsets))
.add_buffer(Buffer::from_slice_ref(values))
.build()
.unwrap();
let string_array = StringArray::from(array_data);
string_array.value(4);
}
#[test]
fn test_string_array_fmt_debug() {
let arr: StringArray = vec!["hello", "arrow"].into();
assert_eq!(
"StringArray\n[\n \"hello\",\n \"arrow\",\n]",
format!("{arr:?}")
);
}
#[test]
fn test_large_string_array_fmt_debug() {
let arr: LargeStringArray = vec!["hello", "arrow"].into();
assert_eq!(
"LargeStringArray\n[\n \"hello\",\n \"arrow\",\n]",
format!("{arr:?}")
);
}
#[test]
fn test_string_array_from_iter() {
let data = [Some("hello"), None, Some("arrow")];
let data_vec = data.to_vec();
// from Vec<Option<&str>>
let array1 = StringArray::from(data_vec.clone());
// from Iterator<Option<&str>>
let array2: StringArray = data_vec.clone().into_iter().collect();
// from Iterator<Option<String>>
let array3: StringArray = data_vec
.into_iter()
.map(|x| x.map(|s| s.to_string()))
.collect();
// from Iterator<&Option<&str>>
let array4: StringArray = data.iter().collect::<StringArray>();
assert_eq!(array1, array2);
assert_eq!(array2, array3);
assert_eq!(array3, array4);
}
#[test]
fn test_string_array_from_iter_values() {
let data = ["hello", "hello2"];
let array1 = StringArray::from_iter_values(data.iter());
assert_eq!(array1.value(0), "hello");
assert_eq!(array1.value(1), "hello2");
// Also works with String types.
let data2 = ["goodbye".to_string(), "goodbye2".to_string()];
let array2 = StringArray::from_iter_values(data2.iter());
assert_eq!(array2.value(0), "goodbye");
assert_eq!(array2.value(1), "goodbye2");
}
#[test]
fn test_string_array_from_unbound_iter() {
// iterator that doesn't declare (upper) size bound
let string_iter = (0..)
.scan(0usize, |pos, i| {
if *pos < 10 {
*pos += 1;
Some(Some(format!("value {i}")))
} else {
// actually returns up to 10 values
None
}
})
// limited using take()
.take(100);
let (_, upper_size_bound) = string_iter.size_hint();
// the upper bound, defined by take above, is 100
assert_eq!(upper_size_bound, Some(100));
let string_array: StringArray = string_iter.collect();
// but the actual number of items in the array should be 10
assert_eq!(string_array.len(), 10);
}
#[test]
fn test_string_array_all_null() {
let data: Vec<Option<&str>> = vec![None];
let array = StringArray::from(data);
array
.into_data()
.validate_full()
.expect("All null array has valid array data");
}
#[test]
fn test_large_string_array_all_null() {
let data: Vec<Option<&str>> = vec![None];
let array = LargeStringArray::from(data);
array
.into_data()
.validate_full()
.expect("All null array has valid array data");
}
fn _test_generic_string_array_from_list_array<O: OffsetSizeTrait>() {
let values = b"HelloArrowAndParquet";
// "ArrowAndParquet"
let child_data = ArrayData::builder(DataType::UInt8)
.len(15)
.offset(5)
.add_buffer(Buffer::from(values))
.build()
.unwrap();
let offsets = [0, 5, 8, 15].map(|n| O::from_usize(n).unwrap());
let null_buffer = Buffer::from_slice_ref([0b101]);
let data_type = GenericListArray::<O>::DATA_TYPE_CONSTRUCTOR(Arc::new(
Field::new_list_field(DataType::UInt8, false),
));
// [None, Some("Parquet")]
let array_data = ArrayData::builder(data_type)
.len(2)
.offset(1)
.add_buffer(Buffer::from_slice_ref(offsets))
.null_bit_buffer(Some(null_buffer))
.add_child_data(child_data)
.build()
.unwrap();
let list_array = GenericListArray::<O>::from(array_data);
let string_array = GenericStringArray::<O>::from(list_array);
assert_eq!(2, string_array.len());
assert_eq!(1, string_array.null_count());
assert!(string_array.is_null(0));
assert!(string_array.is_valid(1));
assert_eq!("Parquet", string_array.value(1));
}
#[test]
fn test_string_array_from_list_array() {
_test_generic_string_array_from_list_array::<i32>();
}
#[test]
fn test_large_string_array_from_list_array() {
_test_generic_string_array_from_list_array::<i64>();
}
fn _test_generic_string_array_from_list_array_with_child_nulls_failed<O: OffsetSizeTrait>() {
let values = b"HelloArrow";
let child_data = ArrayData::builder(DataType::UInt8)
.len(10)
.add_buffer(Buffer::from(values))
.null_bit_buffer(Some(Buffer::from_slice_ref([0b1010101010])))
.build()
.unwrap();
let offsets = [0, 5, 10].map(|n| O::from_usize(n).unwrap());
// It is possible to create a null struct containing a non-nullable child
// see https://github.com/apache/arrow-rs/pull/3244 for details
let data_type = GenericListArray::<O>::DATA_TYPE_CONSTRUCTOR(Arc::new(
Field::new_list_field(DataType::UInt8, true),
));
// [None, Some(b"Parquet")]
let array_data = ArrayData::builder(data_type)
.len(2)
.add_buffer(Buffer::from_slice_ref(offsets))
.add_child_data(child_data)
.build()
.unwrap();
let list_array = GenericListArray::<O>::from(array_data);
drop(GenericStringArray::<O>::from(list_array));
}
#[test]
#[should_panic(expected = "The child array cannot contain null values.")]
fn test_string_array_from_list_array_with_child_nulls_failed() {
_test_generic_string_array_from_list_array_with_child_nulls_failed::<i32>();
}
#[test]
#[should_panic(expected = "The child array cannot contain null values.")]
fn test_large_string_array_from_list_array_with_child_nulls_failed() {
_test_generic_string_array_from_list_array_with_child_nulls_failed::<i64>();
}
fn _test_generic_string_array_from_list_array_wrong_type<O: OffsetSizeTrait>() {
let values = b"HelloArrow";
let child_data = ArrayData::builder(DataType::UInt16)
.len(5)
.add_buffer(Buffer::from(values))
.build()
.unwrap();
let offsets = [0, 2, 3].map(|n| O::from_usize(n).unwrap());
let data_type = GenericListArray::<O>::DATA_TYPE_CONSTRUCTOR(Arc::new(
Field::new_list_field(DataType::UInt16, false),
));
let array_data = ArrayData::builder(data_type)
.len(2)
.add_buffer(Buffer::from_slice_ref(offsets))
.add_child_data(child_data)
.build()
.unwrap();
let list_array = GenericListArray::<O>::from(array_data);
drop(GenericStringArray::<O>::from(list_array));
}
#[test]
#[should_panic(
expected = "BinaryArray can only be created from List<u8> arrays, mismatched data types."
)]
fn test_string_array_from_list_array_wrong_type() {
_test_generic_string_array_from_list_array_wrong_type::<i32>();
}
#[test]
#[should_panic(
expected = "BinaryArray can only be created from List<u8> arrays, mismatched data types."
)]
fn test_large_string_array_from_list_array_wrong_type() {
_test_generic_string_array_from_list_array_wrong_type::<i64>();
}
#[test]
#[should_panic(
expected = "Encountered non UTF-8 data: invalid utf-8 sequence of 1 bytes from index 0"
)]
fn test_list_array_utf8_validation() {
let mut builder = ListBuilder::new(PrimitiveBuilder::<UInt8Type>::new());
builder.values().append_value(0xFF);
builder.append(true);
let list = builder.finish();
let _ = StringArray::from(list);
}
#[test]
fn test_empty_offsets() {
let string = StringArray::from(
ArrayData::builder(DataType::Utf8)
.buffers(vec![Buffer::from(&[]), Buffer::from(&[])])
.build()
.unwrap(),
);
assert_eq!(string.len(), 0);
assert_eq!(string.value_offsets(), &[0]);
let string = LargeStringArray::from(
ArrayData::builder(DataType::LargeUtf8)
.buffers(vec![Buffer::from(&[]), Buffer::from(&[])])
.build()
.unwrap(),
);
assert_eq!(string.len(), 0);
assert_eq!(string.value_offsets(), &[0]);
}
#[test]
fn test_into_builder() {
let array: StringArray = vec!["hello", "arrow"].into();
// Append values
let mut builder = array.into_builder().unwrap();
builder.append_value("rust");
let expected: StringArray = vec!["hello", "arrow", "rust"].into();
let array = builder.finish();
assert_eq!(expected, array);
}
#[test]
fn test_into_builder_err() {
let array: StringArray = vec!["hello", "arrow"].into();
// Clone it, so we cannot get a mutable builder back
let shared_array = array.clone();
let err_return = array.into_builder().unwrap_err();
assert_eq!(&err_return, &shared_array);
}
#[test]
fn test_non_null_string_array_equal() {
let a = StringArray::from(vec![Some("ab"), Some("c")]);
let b = StringArray::from(vec![Some("a"), Some("bc")]);
assert_ne!(a, b);
}
}
@@ -0,0 +1,950 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::array::print_long_array;
use crate::{Array, ArrayRef, RecordBatch, make_array, new_null_array};
use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer};
use arrow_data::{ArrayData, ArrayDataBuilder};
use arrow_schema::{ArrowError, DataType, Field, FieldRef, Fields};
use std::sync::Arc;
use std::{any::Any, ops::Index};
/// An array of [structs](https://arrow.apache.org/docs/format/Columnar.html#struct-layout)
///
/// Each child (called *field*) is represented by a separate array.
///
/// # Comparison with [RecordBatch]
///
/// Both [`RecordBatch`] and [`StructArray`] represent a collection of columns / arrays with the
/// same length.
///
/// However, there are a couple of key differences:
///
/// * [`StructArray`] can be nested within other [`Array`], including itself
/// * [`RecordBatch`] can contain top-level metadata on its associated [`Schema`][arrow_schema::Schema]
/// * [`StructArray`] can contain top-level nulls, i.e. `null`
/// * [`RecordBatch`] can only represent nulls in its child columns, i.e. `{"field": null}`
///
/// [`StructArray`] is therefore a more general data container than [`RecordBatch`], and as such
/// code that needs to handle both will typically share an implementation in terms of
/// [`StructArray`] and convert to/from [`RecordBatch`] as necessary.
///
/// [`From`] implementations are provided to facilitate this conversion, however, converting
/// from a [`StructArray`] containing top-level nulls to a [`RecordBatch`] will panic, as there
/// is no way to preserve them.
///
/// # Example: Create an array from a vector of fields
///
/// ```
/// use std::sync::Arc;
/// use arrow_array::{Array, ArrayRef, BooleanArray, Int32Array, StructArray};
/// use arrow_schema::{DataType, Field};
///
/// let boolean = Arc::new(BooleanArray::from(vec![false, false, true, true]));
/// let int = Arc::new(Int32Array::from(vec![42, 28, 19, 31]));
///
/// let struct_array = StructArray::from(vec![
/// (
/// Arc::new(Field::new("b", DataType::Boolean, false)),
/// boolean.clone() as ArrayRef,
/// ),
/// (
/// Arc::new(Field::new("c", DataType::Int32, false)),
/// int.clone() as ArrayRef,
/// ),
/// ]);
/// assert_eq!(struct_array.column(0).as_ref(), boolean.as_ref());
/// assert_eq!(struct_array.column(1).as_ref(), int.as_ref());
/// assert_eq!(4, struct_array.len());
/// assert_eq!(0, struct_array.null_count());
/// assert_eq!(0, struct_array.offset());
/// ```
#[derive(Clone)]
pub struct StructArray {
len: usize,
data_type: DataType,
nulls: Option<NullBuffer>,
fields: Vec<ArrayRef>,
}
impl StructArray {
/// Create a new [`StructArray`] from the provided parts, panicking on failure
///
/// # Panics
///
/// Panics if [`Self::try_new`] returns an error
pub fn new(fields: Fields, arrays: Vec<ArrayRef>, nulls: Option<NullBuffer>) -> Self {
Self::try_new(fields, arrays, nulls).unwrap()
}
/// Create a new [`StructArray`] from the provided parts, returning an error on failure
///
/// The length will be inferred from the length of the child arrays. Returns an error if
/// there are no child arrays. Consider using [`Self::try_new_with_length`] if the length
/// is known to avoid this.
///
/// # Errors
///
/// Errors if
///
/// * `fields.len() == 0`
/// * Any reason that [`Self::try_new_with_length`] would error
pub fn try_new(
fields: Fields,
arrays: Vec<ArrayRef>,
nulls: Option<NullBuffer>,
) -> Result<Self, ArrowError> {
let len = arrays.first().map(|x| x.len()).ok_or_else(||ArrowError::InvalidArgumentError("use StructArray::try_new_with_length or StructArray::new_empty_fields to create a struct array with no fields so that the length can be set correctly".to_string()))?;
Self::try_new_with_length(fields, arrays, nulls, len)
}
/// Create a new [`StructArray`] from the provided parts, returning an error on failure
///
/// # Errors
///
/// Errors if
///
/// * `fields.len() != arrays.len()`
/// * `fields[i].data_type() != arrays[i].data_type()`
/// * `arrays[i].len() != arrays[j].len()`
/// * `arrays[i].len() != nulls.len()`
/// * `!fields[i].is_nullable() && !nulls.contains(arrays[i].nulls())`
pub fn try_new_with_length(
fields: Fields,
arrays: Vec<ArrayRef>,
nulls: Option<NullBuffer>,
len: usize,
) -> Result<Self, ArrowError> {
if fields.len() != arrays.len() {
return Err(ArrowError::InvalidArgumentError(format!(
"Incorrect number of arrays for StructArray fields, expected {} got {}",
fields.len(),
arrays.len()
)));
}
if let Some(n) = nulls.as_ref() {
if n.len() != len {
return Err(ArrowError::InvalidArgumentError(format!(
"Incorrect number of nulls for StructArray, expected {len} got {}",
n.len(),
)));
}
}
for (f, a) in fields.iter().zip(&arrays) {
if f.data_type() != a.data_type() {
return Err(ArrowError::InvalidArgumentError(format!(
"Incorrect datatype for StructArray field {:?}, expected {} got {}",
f.name(),
f.data_type(),
a.data_type()
)));
}
if a.len() != len {
return Err(ArrowError::InvalidArgumentError(format!(
"Incorrect array length for StructArray field {:?}, expected {} got {}",
f.name(),
len,
a.len()
)));
}
if !f.is_nullable() {
if let Some(a) = a.logical_nulls() {
if !nulls.as_ref().map(|n| n.contains(&a)).unwrap_or_default()
&& a.null_count() > 0
{
return Err(ArrowError::InvalidArgumentError(format!(
"Found unmasked nulls for non-nullable StructArray field {:?}",
f.name()
)));
}
}
}
}
Ok(Self {
len,
data_type: DataType::Struct(fields),
nulls: nulls.filter(|n| n.null_count() > 0),
fields: arrays,
})
}
/// Create a new [`StructArray`] of length `len` where all values are null
pub fn new_null(fields: Fields, len: usize) -> Self {
let arrays = fields
.iter()
.map(|f| new_null_array(f.data_type(), len))
.collect();
Self {
len,
data_type: DataType::Struct(fields),
nulls: Some(NullBuffer::new_null(len)),
fields: arrays,
}
}
/// Create a new [`StructArray`] from the provided parts without validation
///
/// The length will be inferred from the length of the child arrays. Panics if there are no
/// child arrays. Consider using [`Self::new_unchecked_with_length`] if the length is known
/// to avoid this.
///
/// # Safety
///
/// Safe if [`Self::new`] would not panic with the given arguments
pub unsafe fn new_unchecked(
fields: Fields,
arrays: Vec<ArrayRef>,
nulls: Option<NullBuffer>,
) -> Self {
if cfg!(feature = "force_validate") {
return Self::new(fields, arrays, nulls);
}
let len = arrays.first().map(|x| x.len()).expect(
"cannot use StructArray::new_unchecked if there are no fields, length is unknown",
);
Self {
len,
data_type: DataType::Struct(fields),
nulls,
fields: arrays,
}
}
/// Create a new [`StructArray`] from the provided parts without validation
///
/// # Safety
///
/// Safe if [`Self::new`] would not panic with the given arguments
pub unsafe fn new_unchecked_with_length(
fields: Fields,
arrays: Vec<ArrayRef>,
nulls: Option<NullBuffer>,
len: usize,
) -> Self {
if cfg!(feature = "force_validate") {
return Self::try_new_with_length(fields, arrays, nulls, len).unwrap();
}
Self {
len,
data_type: DataType::Struct(fields),
nulls,
fields: arrays,
}
}
/// Create a new [`StructArray`] containing no fields
///
/// # Panics
///
/// If `len != nulls.len()`
pub fn new_empty_fields(len: usize, nulls: Option<NullBuffer>) -> Self {
if let Some(n) = &nulls {
assert_eq!(len, n.len())
}
Self {
len,
data_type: DataType::Struct(Fields::empty()),
fields: vec![],
nulls,
}
}
/// Deconstruct this array into its constituent parts
pub fn into_parts(self) -> (Fields, Vec<ArrayRef>, Option<NullBuffer>) {
let f = match self.data_type {
DataType::Struct(f) => f,
_ => unreachable!(),
};
(f, self.fields, self.nulls)
}
/// Returns the field at `pos`.
pub fn column(&self, pos: usize) -> &ArrayRef {
&self.fields[pos]
}
/// Return the number of fields in this struct array
pub fn num_columns(&self) -> usize {
self.fields.len()
}
/// Returns the fields of the struct array
pub fn columns(&self) -> &[ArrayRef] {
&self.fields
}
/// Return field names in this struct array
pub fn column_names(&self) -> Vec<&str> {
match self.data_type() {
DataType::Struct(fields) => fields
.iter()
.map(|f| f.name().as_str())
.collect::<Vec<&str>>(),
_ => unreachable!("Struct array's data type is not struct!"),
}
}
/// Returns the [`Fields`] of this [`StructArray`]
pub fn fields(&self) -> &Fields {
match self.data_type() {
DataType::Struct(f) => f,
_ => unreachable!(),
}
}
/// Return child array whose field name equals to column_name
///
/// Note: A schema can currently have duplicate field names, in which case
/// the first field will always be selected.
/// This issue will be addressed in [ARROW-11178](https://issues.apache.org/jira/browse/ARROW-11178)
pub fn column_by_name(&self, column_name: &str) -> Option<&ArrayRef> {
self.column_names()
.iter()
.position(|c| c == &column_name)
.map(|pos| self.column(pos))
}
/// Returns a zero-copy slice of this array with the indicated offset and length.
pub fn slice(&self, offset: usize, len: usize) -> Self {
assert!(
offset.saturating_add(len) <= self.len,
"the length + offset of the sliced StructArray cannot exceed the existing length"
);
let fields = self.fields.iter().map(|a| a.slice(offset, len)).collect();
Self {
len,
data_type: self.data_type.clone(),
nulls: self.nulls.as_ref().map(|n| n.slice(offset, len)),
fields,
}
}
}
impl From<ArrayData> for StructArray {
fn from(data: ArrayData) -> Self {
let parent_offset = data.offset();
let parent_len = data.len();
let fields = data
.child_data()
.iter()
.map(|cd| {
if parent_offset != 0 || parent_len != cd.len() {
make_array(cd.slice(parent_offset, parent_len))
} else {
make_array(cd.clone())
}
})
.collect();
Self {
len: data.len(),
data_type: data.data_type().clone(),
nulls: data.nulls().cloned(),
fields,
}
}
}
impl From<StructArray> for ArrayData {
fn from(array: StructArray) -> Self {
let builder = ArrayDataBuilder::new(array.data_type)
.len(array.len)
.nulls(array.nulls)
.child_data(array.fields.iter().map(|x| x.to_data()).collect());
unsafe { builder.build_unchecked() }
}
}
impl TryFrom<Vec<(&str, ArrayRef)>> for StructArray {
type Error = ArrowError;
/// builds a StructArray from a vector of names and arrays.
fn try_from(values: Vec<(&str, ArrayRef)>) -> Result<Self, ArrowError> {
let (fields, arrays): (Vec<_>, _) = values
.into_iter()
.map(|(name, array)| {
(
Field::new(name, array.data_type().clone(), array.is_nullable()),
array,
)
})
.unzip();
StructArray::try_new(fields.into(), arrays, None)
}
}
/// SAFETY: Correctly implements the contract of Arrow Arrays
unsafe impl Array for StructArray {
fn as_any(&self) -> &dyn Any {
self
}
fn to_data(&self) -> ArrayData {
self.clone().into()
}
fn into_data(self) -> ArrayData {
self.into()
}
fn data_type(&self) -> &DataType {
&self.data_type
}
fn slice(&self, offset: usize, length: usize) -> ArrayRef {
Arc::new(self.slice(offset, length))
}
fn len(&self) -> usize {
self.len
}
fn is_empty(&self) -> bool {
self.len == 0
}
fn shrink_to_fit(&mut self) {
if let Some(nulls) = &mut self.nulls {
nulls.shrink_to_fit();
}
self.fields.iter_mut().for_each(|n| n.shrink_to_fit());
}
fn offset(&self) -> usize {
0
}
fn nulls(&self) -> Option<&NullBuffer> {
self.nulls.as_ref()
}
fn logical_null_count(&self) -> usize {
// More efficient that the default implementation
self.null_count()
}
fn get_buffer_memory_size(&self) -> usize {
let mut size = self.fields.iter().map(|a| a.get_buffer_memory_size()).sum();
if let Some(n) = self.nulls.as_ref() {
size += n.buffer().capacity();
}
size
}
fn get_array_memory_size(&self) -> usize {
let mut size = self.fields.iter().map(|a| a.get_array_memory_size()).sum();
size += std::mem::size_of::<Self>();
if let Some(n) = self.nulls.as_ref() {
size += n.buffer().capacity();
}
size
}
}
impl From<Vec<(FieldRef, ArrayRef)>> for StructArray {
fn from(v: Vec<(FieldRef, ArrayRef)>) -> Self {
let (fields, arrays): (Vec<_>, _) = v.into_iter().unzip();
StructArray::new(fields.into(), arrays, None)
}
}
impl std::fmt::Debug for StructArray {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
writeln!(f, "StructArray")?;
writeln!(f, "-- validity:")?;
writeln!(f, "[")?;
print_long_array(self, f, |_array, _index, f| write!(f, "valid"))?;
writeln!(f, "]\n[")?;
for (child_index, name) in self.column_names().iter().enumerate() {
let column = self.column(child_index);
writeln!(
f,
"-- child {}: \"{}\" ({:?})",
child_index,
name,
column.data_type()
)?;
std::fmt::Debug::fmt(column, f)?;
writeln!(f)?;
}
write!(f, "]")
}
}
impl From<(Vec<(FieldRef, ArrayRef)>, Buffer)> for StructArray {
fn from(pair: (Vec<(FieldRef, ArrayRef)>, Buffer)) -> Self {
let len = pair.0.first().map(|x| x.1.len()).unwrap_or_default();
let (fields, arrays): (Vec<_>, Vec<_>) = pair.0.into_iter().unzip();
let nulls = NullBuffer::new(BooleanBuffer::new(pair.1, 0, len));
Self::new(fields.into(), arrays, Some(nulls))
}
}
impl From<RecordBatch> for StructArray {
fn from(value: RecordBatch) -> Self {
Self {
len: value.num_rows(),
data_type: DataType::Struct(value.schema().fields().clone()),
nulls: None,
fields: value.columns().to_vec(),
}
}
}
impl Index<&str> for StructArray {
type Output = ArrayRef;
/// Get a reference to a column's array by name.
///
/// Note: A schema can currently have duplicate field names, in which case
/// the first field will always be selected.
/// This issue will be addressed in [ARROW-11178](https://issues.apache.org/jira/browse/ARROW-11178)
///
/// # Panics
///
/// Panics if the name is not in the schema.
fn index(&self, name: &str) -> &Self::Output {
self.column_by_name(name).unwrap()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array, StringArray};
use arrow_buffer::ToByteSlice;
#[test]
fn test_struct_array_builder() {
let boolean_array = BooleanArray::from(vec![false, false, true, true]);
let int_array = Int64Array::from(vec![42, 28, 19, 31]);
let fields = vec![
Field::new("a", DataType::Boolean, false),
Field::new("b", DataType::Int64, false),
];
let struct_array_data = ArrayData::builder(DataType::Struct(fields.into()))
.len(4)
.add_child_data(boolean_array.to_data())
.add_child_data(int_array.to_data())
.build()
.unwrap();
let struct_array = StructArray::from(struct_array_data);
assert_eq!(struct_array.column(0).as_ref(), &boolean_array);
assert_eq!(struct_array.column(1).as_ref(), &int_array);
}
#[test]
fn test_struct_array_from() {
let boolean = Arc::new(BooleanArray::from(vec![false, false, true, true]));
let int = Arc::new(Int32Array::from(vec![42, 28, 19, 31]));
let struct_array = StructArray::from(vec![
(
Arc::new(Field::new("b", DataType::Boolean, false)),
boolean.clone() as ArrayRef,
),
(
Arc::new(Field::new("c", DataType::Int32, false)),
int.clone() as ArrayRef,
),
]);
assert_eq!(struct_array.column(0).as_ref(), boolean.as_ref());
assert_eq!(struct_array.column(1).as_ref(), int.as_ref());
assert_eq!(4, struct_array.len());
assert_eq!(0, struct_array.null_count());
assert_eq!(0, struct_array.offset());
}
#[test]
fn test_struct_array_from_data_with_offset_and_length() {
// Various ways to make the struct array:
//
// [{x: 2}, {x: 3}, None]
//
// from slicing larger buffers/arrays with offsets and lengths
let int_arr = Int32Array::from(vec![1, 2, 3, 4, 5]);
let int_field = Field::new("x", DataType::Int32, false);
let struct_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, true, false]));
let int_data = int_arr.to_data();
// Case 1: Offset + length, nulls are not sliced
let case1 = ArrayData::builder(DataType::Struct(Fields::from(vec![int_field.clone()])))
.len(3)
.offset(1)
.nulls(Some(struct_nulls))
.add_child_data(int_data.clone())
.build()
.unwrap();
// Case 2: Offset + length, nulls are sliced
let struct_nulls =
NullBuffer::new(BooleanBuffer::from(vec![true, true, true, false, true]).slice(1, 3));
let case2 = ArrayData::builder(DataType::Struct(Fields::from(vec![int_field.clone()])))
.len(3)
.offset(1)
.nulls(Some(struct_nulls.clone()))
.add_child_data(int_data.clone())
.build()
.unwrap();
// Case 3: struct length is smaller than child length but no offset
let offset_int_data = int_data.slice(1, 4);
let case3 = ArrayData::builder(DataType::Struct(Fields::from(vec![int_field.clone()])))
.len(3)
.nulls(Some(struct_nulls))
.add_child_data(offset_int_data)
.build()
.unwrap();
let expected = StructArray::new(
Fields::from(vec![int_field.clone()]),
vec![Arc::new(int_arr)],
Some(NullBuffer::new(BooleanBuffer::from(vec![
true, true, true, false, true,
]))),
)
.slice(1, 3);
for case in [case1, case2, case3] {
let struct_arr_from_data = StructArray::from(case);
assert_eq!(struct_arr_from_data, expected);
assert_eq!(struct_arr_from_data.column(0), expected.column(0));
}
}
#[test]
#[should_panic(expected = "assertion failed: end <= self.len()")]
fn test_struct_array_from_data_with_offset_and_length_error() {
let int_arr = Int32Array::from(vec![1, 2, 3, 4, 5]);
let int_field = Field::new("x", DataType::Int32, false);
let struct_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, true, false]));
let int_data = int_arr.to_data();
// If parent offset is 3 and len is 3 then child must have 6 items
let struct_data =
ArrayData::builder(DataType::Struct(Fields::from(vec![int_field.clone()])))
.len(3)
.offset(3)
.nulls(Some(struct_nulls))
.add_child_data(int_data)
.build()
.unwrap();
let _ = StructArray::from(struct_data);
}
/// validates that struct can be accessed using `column_name` as index i.e. `struct_array["column_name"]`.
#[test]
fn test_struct_array_index_access() {
let boolean = Arc::new(BooleanArray::from(vec![false, false, true, true]));
let int = Arc::new(Int32Array::from(vec![42, 28, 19, 31]));
let struct_array = StructArray::from(vec![
(
Arc::new(Field::new("b", DataType::Boolean, false)),
boolean.clone() as ArrayRef,
),
(
Arc::new(Field::new("c", DataType::Int32, false)),
int.clone() as ArrayRef,
),
]);
assert_eq!(struct_array["b"].as_ref(), boolean.as_ref());
assert_eq!(struct_array["c"].as_ref(), int.as_ref());
}
/// validates that the in-memory representation follows [the spec](https://arrow.apache.org/docs/format/Columnar.html#struct-layout)
#[test]
fn test_struct_array_from_vec() {
let strings: ArrayRef = Arc::new(StringArray::from(vec![
Some("joe"),
None,
None,
Some("mark"),
]));
let ints: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)]));
let arr =
StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())]).unwrap();
let struct_data = arr.into_data();
assert_eq!(4, struct_data.len());
assert_eq!(0, struct_data.null_count());
let expected_string_data = ArrayData::builder(DataType::Utf8)
.len(4)
.null_bit_buffer(Some(Buffer::from(&[9_u8])))
.add_buffer(Buffer::from([0, 3, 3, 3, 7].to_byte_slice()))
.add_buffer(Buffer::from(b"joemark"))
.build()
.unwrap();
let expected_int_data = ArrayData::builder(DataType::Int32)
.len(4)
.null_bit_buffer(Some(Buffer::from(&[11_u8])))
.add_buffer(Buffer::from([1, 2, 0, 4].to_byte_slice()))
.build()
.unwrap();
assert_eq!(expected_string_data, struct_data.child_data()[0]);
assert_eq!(expected_int_data, struct_data.child_data()[1]);
}
#[test]
fn test_struct_array_from_vec_error() {
let strings: ArrayRef = Arc::new(StringArray::from(vec![
Some("joe"),
None,
None,
// 3 elements, not 4
]));
let ints: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)]));
let err = StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())])
.unwrap_err()
.to_string();
assert_eq!(
err,
"Invalid argument error: Incorrect array length for StructArray field \"f2\", expected 3 got 4"
)
}
#[test]
#[should_panic(
expected = "Incorrect datatype for StructArray field \\\"b\\\", expected Int16 got Boolean"
)]
fn test_struct_array_from_mismatched_types_single() {
drop(StructArray::from(vec![(
Arc::new(Field::new("b", DataType::Int16, false)),
Arc::new(BooleanArray::from(vec![false, false, true, true])) as Arc<dyn Array>,
)]));
}
#[test]
#[should_panic(
expected = "Incorrect datatype for StructArray field \\\"b\\\", expected Int16 got Boolean"
)]
fn test_struct_array_from_mismatched_types_multiple() {
drop(StructArray::from(vec![
(
Arc::new(Field::new("b", DataType::Int16, false)),
Arc::new(BooleanArray::from(vec![false, false, true, true])) as Arc<dyn Array>,
),
(
Arc::new(Field::new("c", DataType::Utf8, false)),
Arc::new(Int32Array::from(vec![42, 28, 19, 31])),
),
]));
}
#[test]
fn test_struct_array_slice() {
let boolean_data = ArrayData::builder(DataType::Boolean)
.len(5)
.add_buffer(Buffer::from([0b00010000]))
.null_bit_buffer(Some(Buffer::from([0b00010001])))
.build()
.unwrap();
let int_data = ArrayData::builder(DataType::Int32)
.len(5)
.add_buffer(Buffer::from([0, 28, 42, 0, 0].to_byte_slice()))
.null_bit_buffer(Some(Buffer::from([0b00000110])))
.build()
.unwrap();
let field_types = vec![
Field::new("a", DataType::Boolean, true),
Field::new("b", DataType::Int32, true),
];
let struct_array_data = ArrayData::builder(DataType::Struct(field_types.into()))
.len(5)
.add_child_data(boolean_data.clone())
.add_child_data(int_data.clone())
.null_bit_buffer(Some(Buffer::from([0b00010111])))
.build()
.unwrap();
let struct_array = StructArray::from(struct_array_data);
assert_eq!(5, struct_array.len());
assert_eq!(1, struct_array.null_count());
assert!(struct_array.is_valid(0));
assert!(struct_array.is_valid(1));
assert!(struct_array.is_valid(2));
assert!(struct_array.is_null(3));
assert!(struct_array.is_valid(4));
assert_eq!(boolean_data, struct_array.column(0).to_data());
assert_eq!(int_data, struct_array.column(1).to_data());
let c0 = struct_array.column(0);
let c0 = c0.as_any().downcast_ref::<BooleanArray>().unwrap();
assert_eq!(5, c0.len());
assert_eq!(3, c0.null_count());
assert!(c0.is_valid(0));
assert!(!c0.value(0));
assert!(c0.is_null(1));
assert!(c0.is_null(2));
assert!(c0.is_null(3));
assert!(c0.is_valid(4));
assert!(c0.value(4));
let c1 = struct_array.column(1);
let c1 = c1.as_any().downcast_ref::<Int32Array>().unwrap();
assert_eq!(5, c1.len());
assert_eq!(3, c1.null_count());
assert!(c1.is_null(0));
assert!(c1.is_valid(1));
assert_eq!(28, c1.value(1));
assert!(c1.is_valid(2));
assert_eq!(42, c1.value(2));
assert!(c1.is_null(3));
assert!(c1.is_null(4));
let sliced_array = struct_array.slice(2, 3);
let sliced_array = sliced_array.as_any().downcast_ref::<StructArray>().unwrap();
assert_eq!(3, sliced_array.len());
assert_eq!(1, sliced_array.null_count());
assert!(sliced_array.is_valid(0));
assert!(sliced_array.is_null(1));
assert!(sliced_array.is_valid(2));
let sliced_c0 = sliced_array.column(0);
let sliced_c0 = sliced_c0.as_any().downcast_ref::<BooleanArray>().unwrap();
assert_eq!(3, sliced_c0.len());
assert!(sliced_c0.is_null(0));
assert!(sliced_c0.is_null(1));
assert!(sliced_c0.is_valid(2));
assert!(sliced_c0.value(2));
let sliced_c1 = sliced_array.column(1);
let sliced_c1 = sliced_c1.as_any().downcast_ref::<Int32Array>().unwrap();
assert_eq!(3, sliced_c1.len());
assert!(sliced_c1.is_valid(0));
assert_eq!(42, sliced_c1.value(0));
assert!(sliced_c1.is_null(1));
assert!(sliced_c1.is_null(2));
}
#[test]
#[should_panic(
expected = "Incorrect array length for StructArray field \\\"c\\\", expected 1 got 2"
)]
fn test_invalid_struct_child_array_lengths() {
drop(StructArray::from(vec![
(
Arc::new(Field::new("b", DataType::Float32, false)),
Arc::new(Float32Array::from(vec![1.1])) as Arc<dyn Array>,
),
(
Arc::new(Field::new("c", DataType::Float64, false)),
Arc::new(Float64Array::from(vec![2.2, 3.3])),
),
]));
}
#[test]
#[should_panic(expected = "use StructArray::try_new_with_length")]
fn test_struct_array_from_empty() {
// This can't work because we don't know how many rows the array should have. Previously we inferred 0 but
// that often led to bugs.
let _ = StructArray::from(vec![]);
}
#[test]
fn test_empty_struct_array() {
assert!(StructArray::try_new(Fields::empty(), vec![], None).is_err());
let arr = StructArray::new_empty_fields(10, None);
assert_eq!(arr.len(), 10);
assert_eq!(arr.null_count(), 0);
assert_eq!(arr.num_columns(), 0);
let arr2 = StructArray::try_new_with_length(Fields::empty(), vec![], None, 10).unwrap();
assert_eq!(arr2.len(), 10);
let arr = StructArray::new_empty_fields(10, Some(NullBuffer::new_null(10)));
assert_eq!(arr.len(), 10);
assert_eq!(arr.null_count(), 10);
assert_eq!(arr.num_columns(), 0);
let arr2 = StructArray::try_new_with_length(
Fields::empty(),
vec![],
Some(NullBuffer::new_null(10)),
10,
)
.unwrap();
assert_eq!(arr2.len(), 10);
}
#[test]
#[should_panic(expected = "Found unmasked nulls for non-nullable StructArray field \\\"c\\\"")]
fn test_struct_array_from_mismatched_nullability() {
drop(StructArray::from(vec![(
Arc::new(Field::new("c", DataType::Int32, false)),
Arc::new(Int32Array::from(vec![Some(42), None, Some(19)])) as ArrayRef,
)]));
}
#[test]
fn test_struct_array_fmt_debug() {
let arr: StructArray = StructArray::new(
vec![Arc::new(Field::new("c", DataType::Int32, true))].into(),
vec![Arc::new(Int32Array::from((0..30).collect::<Vec<_>>())) as ArrayRef],
Some(NullBuffer::new(BooleanBuffer::from(
(0..30).map(|i| i % 2 == 0).collect::<Vec<_>>(),
))),
);
assert_eq!(
format!("{arr:?}"),
"StructArray\n-- validity:\n[\n valid,\n null,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n ...10 elements...,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n]\n[\n-- child 0: \"c\" (Int32)\nPrimitiveArray<Int32>\n[\n 0,\n 1,\n 2,\n 3,\n 4,\n 5,\n 6,\n 7,\n 8,\n 9,\n ...10 elements...,\n 20,\n 21,\n 22,\n 23,\n 24,\n 25,\n 26,\n 27,\n 28,\n 29,\n]\n]"
)
}
#[test]
fn test_struct_array_logical_nulls() {
// Field is non-nullable
let field = Field::new("a", DataType::Int32, false);
let values = vec![1, 2, 3];
// Create a NullBuffer with all bits set to valid (true)
let nulls = NullBuffer::from(vec![true, true, true]);
let array = Int32Array::new(values.into(), Some(nulls));
let child = Arc::new(array) as ArrayRef;
assert!(child.logical_nulls().is_some());
assert_eq!(child.logical_nulls().unwrap().null_count(), 0);
let fields = Fields::from(vec![field]);
let arrays = vec![child];
let nulls = None;
StructArray::try_new(fields, arrays, nulls).expect("should not error");
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,411 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::builder::{ArrayBuilder, BooleanBufferBuilder};
use crate::{Array, ArrayRef, BooleanArray};
use arrow_buffer::Buffer;
use arrow_buffer::NullBufferBuilder;
use arrow_data::ArrayData;
use arrow_schema::{ArrowError, DataType};
use std::any::Any;
use std::sync::Arc;
/// Builder for [`BooleanArray`]
///
/// # Example
///
/// Create a `BooleanArray` from a `BooleanBuilder`
///
/// ```
///
/// # use arrow_array::{Array, BooleanArray, builder::BooleanBuilder};
///
/// let mut b = BooleanBuilder::new();
/// b.append_value(true);
/// b.append_null();
/// b.append_value(false);
/// b.append_value(true);
/// let arr = b.finish();
///
/// assert_eq!(4, arr.len());
/// assert_eq!(1, arr.null_count());
/// assert_eq!(true, arr.value(0));
/// assert!(arr.is_valid(0));
/// assert!(!arr.is_null(0));
/// assert!(!arr.is_valid(1));
/// assert!(arr.is_null(1));
/// assert_eq!(false, arr.value(2));
/// assert!(arr.is_valid(2));
/// assert!(!arr.is_null(2));
/// assert_eq!(true, arr.value(3));
/// assert!(arr.is_valid(3));
/// assert!(!arr.is_null(3));
/// ```
#[derive(Debug)]
pub struct BooleanBuilder {
values_builder: BooleanBufferBuilder,
null_buffer_builder: NullBufferBuilder,
}
impl Default for BooleanBuilder {
fn default() -> Self {
Self::new()
}
}
impl BooleanBuilder {
/// Creates a new boolean builder
pub fn new() -> Self {
Self::with_capacity(1024)
}
/// Creates a new boolean builder with space for `capacity` elements without re-allocating
pub fn with_capacity(capacity: usize) -> Self {
Self {
values_builder: BooleanBufferBuilder::new(capacity),
null_buffer_builder: NullBufferBuilder::new(capacity),
}
}
/// Returns the capacity of this builder measured in slots of type `T`
pub fn capacity(&self) -> usize {
self.values_builder.capacity()
}
/// Appends a value of type `T` into the builder
#[inline]
pub fn append_value(&mut self, v: bool) {
self.values_builder.append(v);
self.null_buffer_builder.append_non_null();
}
/// Appends a null slot into the builder
#[inline]
pub fn append_null(&mut self) {
self.null_buffer_builder.append_null();
self.values_builder.advance(1);
}
/// Appends `n` `null`s into the builder.
#[inline]
pub fn append_nulls(&mut self, n: usize) {
self.null_buffer_builder.append_n_nulls(n);
self.values_builder.advance(n);
}
/// Appends an `Option<T>` into the builder
#[inline]
pub fn append_option(&mut self, v: Option<bool>) {
match v {
None => self.append_null(),
Some(v) => self.append_value(v),
};
}
/// Appends a slice of type `T` into the builder
#[inline]
pub fn append_slice(&mut self, v: &[bool]) {
self.values_builder.append_slice(v);
self.null_buffer_builder.append_n_non_nulls(v.len());
}
/// Appends n `additional` bits of value `v` into the buffer
#[inline]
pub fn append_n(&mut self, additional: usize, v: bool) {
self.values_builder.append_n(additional, v);
self.null_buffer_builder.append_n_non_nulls(additional);
}
/// Appends values from a slice of type `T` and a validity boolean slice.
///
/// Returns an error if the slices are of different lengths
#[inline]
pub fn append_values(&mut self, values: &[bool], is_valid: &[bool]) -> Result<(), ArrowError> {
if values.len() != is_valid.len() {
Err(ArrowError::InvalidArgumentError(
"Value and validity lengths must be equal".to_string(),
))
} else {
self.null_buffer_builder.append_slice(is_valid);
self.values_builder.append_slice(values);
Ok(())
}
}
/// Appends array values and null to this builder as is
/// (this means that underlying null values are copied as is).
#[inline]
pub fn append_array(&mut self, array: &BooleanArray) {
self.values_builder.append_buffer(array.values());
if let Some(null_buffer) = array.nulls() {
self.null_buffer_builder.append_buffer(null_buffer);
} else {
self.null_buffer_builder.append_n_non_nulls(array.len());
}
}
/// Builds the [BooleanArray] and reset this builder.
pub fn finish(&mut self) -> BooleanArray {
let len = self.len();
let null_bit_buffer = self.null_buffer_builder.finish();
let builder = ArrayData::builder(DataType::Boolean)
.len(len)
.add_buffer(self.values_builder.finish().into_inner())
.nulls(null_bit_buffer);
let array_data = unsafe { builder.build_unchecked() };
BooleanArray::from(array_data)
}
/// Builds the [BooleanArray] without resetting the builder.
pub fn finish_cloned(&self) -> BooleanArray {
let len = self.len();
let nulls = self.null_buffer_builder.finish_cloned();
let value_buffer = Buffer::from_slice_ref(self.values_builder.as_slice());
let builder = ArrayData::builder(DataType::Boolean)
.len(len)
.add_buffer(value_buffer)
.nulls(nulls);
let array_data = unsafe { builder.build_unchecked() };
BooleanArray::from(array_data)
}
/// Returns the current values buffer as a slice
///
/// Boolean values are bit-packed into bytes. To extract the i-th boolean
/// from the bytes, you can use `arrow_buffer::bit_util::get_bit()`.
pub fn values_slice(&self) -> &[u8] {
self.values_builder.as_slice()
}
/// Returns the current null buffer as a slice
pub fn validity_slice(&self) -> Option<&[u8]> {
self.null_buffer_builder.as_slice()
}
}
impl ArrayBuilder for BooleanBuilder {
/// Returns the builder as a non-mutable `Any` reference.
fn as_any(&self) -> &dyn Any {
self
}
/// Returns the builder as a mutable `Any` reference.
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
/// Returns the boxed builder as a box of `Any`.
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
/// Returns the number of array slots in the builder
fn len(&self) -> usize {
self.values_builder.len()
}
/// Builds the array and reset this builder.
fn finish(&mut self) -> ArrayRef {
Arc::new(self.finish())
}
/// Builds the array without resetting the builder.
fn finish_cloned(&self) -> ArrayRef {
Arc::new(self.finish_cloned())
}
}
impl Extend<Option<bool>> for BooleanBuilder {
#[inline]
fn extend<T: IntoIterator<Item = Option<bool>>>(&mut self, iter: T) {
let buffered = iter.into_iter().collect::<Vec<_>>();
let array = unsafe {
// SAFETY: std::vec::IntoIter implements TrustedLen
BooleanArray::from_trusted_len_iter(buffered.into_iter())
};
self.append_array(&array)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Array;
use arrow_buffer::{BooleanBuffer, NullBuffer};
#[test]
fn test_boolean_array_builder() {
// 00000010 01001000
let buf = Buffer::from([72_u8, 2_u8]);
let mut builder = BooleanArray::builder(10);
for i in 0..10 {
if i == 3 || i == 6 || i == 9 {
builder.append_value(true);
} else {
builder.append_value(false);
}
}
let arr = builder.finish();
assert_eq!(&buf, arr.values().inner());
assert_eq!(10, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(0, arr.null_count());
for i in 0..10 {
assert!(!arr.is_null(i));
assert!(arr.is_valid(i));
assert_eq!(i == 3 || i == 6 || i == 9, arr.value(i), "failed at {i}")
}
}
#[test]
fn test_boolean_array_builder_append_slice() {
let arr1 = BooleanArray::from(vec![Some(true), Some(false), None, None, Some(false)]);
let mut builder = BooleanArray::builder(0);
builder.append_slice(&[true, false]);
builder.append_null();
builder.append_null();
builder.append_value(false);
let arr2 = builder.finish();
assert_eq!(arr1, arr2);
}
#[test]
fn test_boolean_array_builder_append_slice_large() {
let arr1 = BooleanArray::from(vec![true; 513]);
let mut builder = BooleanArray::builder(512);
builder.append_slice(&[true; 513]);
let arr2 = builder.finish();
assert_eq!(arr1, arr2);
}
#[test]
fn test_boolean_array_builder_no_null() {
let mut builder = BooleanArray::builder(0);
builder.append_option(Some(true));
builder.append_value(false);
builder.append_slice(&[true, false, true]);
builder
.append_values(&[false, false, true], &[true, true, true])
.unwrap();
let array = builder.finish();
assert_eq!(0, array.null_count());
assert!(array.nulls().is_none());
}
#[test]
fn test_boolean_array_builder_finish_cloned() {
let mut builder = BooleanArray::builder(16);
builder.append_option(Some(true));
builder.append_value(false);
builder.append_slice(&[true, false, true]);
let mut array = builder.finish_cloned();
assert_eq!(3, array.true_count());
assert_eq!(2, array.false_count());
builder
.append_values(&[false, false, true], &[true, true, true])
.unwrap();
array = builder.finish();
assert_eq!(4, array.true_count());
assert_eq!(4, array.false_count());
assert_eq!(0, array.null_count());
assert!(array.nulls().is_none());
}
#[test]
fn test_extend() {
let mut builder = BooleanBuilder::new();
builder.extend([false, false, true, false, false].into_iter().map(Some));
builder.extend([true, true, false].into_iter().map(Some));
let array = builder.finish();
let values = array.iter().map(|x| x.unwrap()).collect::<Vec<_>>();
assert_eq!(
&values,
&[false, false, true, false, false, true, true, false]
)
}
#[test]
fn test_boolean_array_builder_append_n() {
let mut builder = BooleanBuilder::new();
builder.append_n(3, true);
builder.append_n(2, false);
let array = builder.finish();
assert_eq!(3, array.true_count());
assert_eq!(2, array.false_count());
assert_eq!(0, array.null_count());
let values = array.iter().map(|x| x.unwrap()).collect::<Vec<_>>();
assert_eq!(&values, &[true, true, true, false, false])
}
#[test]
fn test_append_array() {
let input = vec![
Some(true),
None,
Some(true),
None,
Some(false),
None,
None,
None,
Some(false),
Some(false),
Some(false),
Some(true),
Some(false),
];
let arr1 = BooleanArray::from(input[..5].to_vec());
let arr2 = BooleanArray::from(input[5..8].to_vec());
let arr3 = BooleanArray::from(input[8..].to_vec());
let mut builder = BooleanBuilder::new();
builder.append_array(&arr1);
builder.append_array(&arr2);
builder.append_array(&arr3);
let actual = builder.finish();
let expected = BooleanArray::from(input);
assert_eq!(actual, expected);
}
#[test]
fn test_append_array_add_underlying_null_values() {
let array = BooleanArray::new(
BooleanBuffer::from(vec![true, false, true, false]),
Some(NullBuffer::from(&[true, true, false, false])),
);
let mut builder = BooleanBuilder::new();
builder.append_array(&array);
let actual = builder.finish();
assert_eq!(actual, array);
assert_eq!(actual.values(), array.values())
}
}
@@ -0,0 +1,231 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
pub use arrow_buffer::BufferBuilder;
pub use arrow_buffer::OffsetBufferBuilder;
use half::f16;
use crate::types::*;
/// Buffer builder for signed 8-bit integer type.
pub type Int8BufferBuilder = BufferBuilder<i8>;
/// Buffer builder for signed 16-bit integer type.
pub type Int16BufferBuilder = BufferBuilder<i16>;
/// Buffer builder for signed 32-bit integer type.
pub type Int32BufferBuilder = BufferBuilder<i32>;
/// Buffer builder for signed 64-bit integer type.
pub type Int64BufferBuilder = BufferBuilder<i64>;
/// Buffer builder for usigned 8-bit integer type.
pub type UInt8BufferBuilder = BufferBuilder<u8>;
/// Buffer builder for usigned 16-bit integer type.
pub type UInt16BufferBuilder = BufferBuilder<u16>;
/// Buffer builder for usigned 32-bit integer type.
pub type UInt32BufferBuilder = BufferBuilder<u32>;
/// Buffer builder for usigned 64-bit integer type.
pub type UInt64BufferBuilder = BufferBuilder<u64>;
/// Buffer builder for 16-bit floating point type.
pub type Float16BufferBuilder = BufferBuilder<f16>;
/// Buffer builder for 32-bit floating point type.
pub type Float32BufferBuilder = BufferBuilder<f32>;
/// Buffer builder for 64-bit floating point type.
pub type Float64BufferBuilder = BufferBuilder<f64>;
/// Buffer builder for 32-bit decimal type.
pub type Decimal32BufferBuilder = BufferBuilder<<Decimal32Type as ArrowPrimitiveType>::Native>;
/// Buffer builder for 64-bit decimal type.
pub type Decimal64BufferBuilder = BufferBuilder<<Decimal64Type as ArrowPrimitiveType>::Native>;
/// Buffer builder for 128-bit decimal type.
pub type Decimal128BufferBuilder = BufferBuilder<<Decimal128Type as ArrowPrimitiveType>::Native>;
/// Buffer builder for 256-bit decimal type.
pub type Decimal256BufferBuilder = BufferBuilder<<Decimal256Type as ArrowPrimitiveType>::Native>;
/// Buffer builder for timestamp type of second unit.
pub type TimestampSecondBufferBuilder =
BufferBuilder<<TimestampSecondType as ArrowPrimitiveType>::Native>;
/// Buffer builder for timestamp type of millisecond unit.
pub type TimestampMillisecondBufferBuilder =
BufferBuilder<<TimestampMillisecondType as ArrowPrimitiveType>::Native>;
/// Buffer builder for timestamp type of microsecond unit.
pub type TimestampMicrosecondBufferBuilder =
BufferBuilder<<TimestampMicrosecondType as ArrowPrimitiveType>::Native>;
/// Buffer builder for timestamp type of nanosecond unit.
pub type TimestampNanosecondBufferBuilder =
BufferBuilder<<TimestampNanosecondType as ArrowPrimitiveType>::Native>;
/// Buffer builder for 32-bit date type.
pub type Date32BufferBuilder = BufferBuilder<<Date32Type as ArrowPrimitiveType>::Native>;
/// Buffer builder for 64-bit date type.
pub type Date64BufferBuilder = BufferBuilder<<Date64Type as ArrowPrimitiveType>::Native>;
/// Buffer builder for 32-bit elaspsed time since midnight of second unit.
pub type Time32SecondBufferBuilder =
BufferBuilder<<Time32SecondType as ArrowPrimitiveType>::Native>;
/// Buffer builder for 32-bit elaspsed time since midnight of millisecond unit.
pub type Time32MillisecondBufferBuilder =
BufferBuilder<<Time32MillisecondType as ArrowPrimitiveType>::Native>;
/// Buffer builder for 64-bit elaspsed time since midnight of microsecond unit.
pub type Time64MicrosecondBufferBuilder =
BufferBuilder<<Time64MicrosecondType as ArrowPrimitiveType>::Native>;
/// Buffer builder for 64-bit elaspsed time since midnight of nanosecond unit.
pub type Time64NanosecondBufferBuilder =
BufferBuilder<<Time64NanosecondType as ArrowPrimitiveType>::Native>;
/// Buffer builder for “calendar” interval in months.
pub type IntervalYearMonthBufferBuilder =
BufferBuilder<<IntervalYearMonthType as ArrowPrimitiveType>::Native>;
/// Buffer builder for “calendar” interval in days and milliseconds.
pub type IntervalDayTimeBufferBuilder =
BufferBuilder<<IntervalDayTimeType as ArrowPrimitiveType>::Native>;
/// Buffer builder “calendar” interval in months, days, and nanoseconds.
pub type IntervalMonthDayNanoBufferBuilder =
BufferBuilder<<IntervalMonthDayNanoType as ArrowPrimitiveType>::Native>;
/// Buffer builder for elaspsed time of second unit.
pub type DurationSecondBufferBuilder =
BufferBuilder<<DurationSecondType as ArrowPrimitiveType>::Native>;
/// Buffer builder for elaspsed time of milliseconds unit.
pub type DurationMillisecondBufferBuilder =
BufferBuilder<<DurationMillisecondType as ArrowPrimitiveType>::Native>;
/// Buffer builder for elaspsed time of microseconds unit.
pub type DurationMicrosecondBufferBuilder =
BufferBuilder<<DurationMicrosecondType as ArrowPrimitiveType>::Native>;
/// Buffer builder for elaspsed time of nanoseconds unit.
pub type DurationNanosecondBufferBuilder =
BufferBuilder<<DurationNanosecondType as ArrowPrimitiveType>::Native>;
#[cfg(test)]
mod tests {
use crate::Array;
use crate::builder::{ArrayBuilder, Int8Builder, Int32BufferBuilder, UInt8BufferBuilder};
#[test]
fn test_builder_i32_empty() {
let mut b = Int32BufferBuilder::new(5);
assert_eq!(0, b.len());
assert_eq!(16, b.capacity());
let a = b.finish();
assert_eq!(0, a.len());
}
#[test]
fn test_builder_i32_alloc_zero_bytes() {
let mut b = Int32BufferBuilder::new(0);
b.append(123);
let a = b.finish();
assert_eq!(4, a.len());
}
#[test]
fn test_builder_i32() {
let mut b = Int32BufferBuilder::new(5);
for i in 0..5 {
b.append(i);
}
assert_eq!(16, b.capacity());
let a = b.finish();
assert_eq!(20, a.len());
}
#[test]
fn test_builder_i32_grow_buffer() {
let mut b = Int32BufferBuilder::new(2);
assert_eq!(16, b.capacity());
for i in 0..20 {
b.append(i);
}
assert_eq!(32, b.capacity());
let a = b.finish();
assert_eq!(80, a.len());
}
#[test]
fn test_builder_finish() {
let mut b = Int32BufferBuilder::new(5);
assert_eq!(16, b.capacity());
for i in 0..10 {
b.append(i);
}
let mut a = b.finish();
assert_eq!(40, a.len());
assert_eq!(0, b.len());
assert_eq!(0, b.capacity());
// Try build another buffer after cleaning up.
for i in 0..20 {
b.append(i)
}
assert_eq!(32, b.capacity());
a = b.finish();
assert_eq!(80, a.len());
}
#[test]
fn test_reserve() {
let mut b = UInt8BufferBuilder::new(2);
assert_eq!(64, b.capacity());
b.reserve(64);
assert_eq!(64, b.capacity());
b.reserve(65);
assert_eq!(128, b.capacity());
let mut b = Int32BufferBuilder::new(2);
assert_eq!(16, b.capacity());
b.reserve(16);
assert_eq!(16, b.capacity());
b.reserve(17);
assert_eq!(32, b.capacity());
}
#[test]
fn test_append_slice() {
let mut b = UInt8BufferBuilder::new(0);
b.append_slice(b"Hello, ");
b.append_slice(b"World!");
let buffer = b.finish();
assert_eq!(13, buffer.len());
let mut b = Int32BufferBuilder::new(0);
b.append_slice(&[32, 54]);
let buffer = b.finish();
assert_eq!(8, buffer.len());
}
#[test]
fn test_append_values() {
let mut a = Int8Builder::new();
a.append_value(1);
a.append_null();
a.append_value(-2);
assert_eq!(a.len(), 3);
// append values
let values = &[1, 2, 3, 4];
let is_valid = &[true, true, false, true];
a.append_values(values, is_valid);
assert_eq!(a.len(), 7);
let array = a.finish();
assert_eq!(array.value(0), 1);
assert!(array.is_null(1));
assert_eq!(array.value(2), -2);
assert_eq!(array.value(3), 1);
assert_eq!(array.value(4), 2);
assert!(array.is_null(5));
assert_eq!(array.value(6), 4);
}
}
@@ -0,0 +1,332 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::array::Array;
use crate::builder::ArrayBuilder;
use crate::{ArrayRef, FixedSizeBinaryArray};
use arrow_buffer::Buffer;
use arrow_buffer::NullBufferBuilder;
use arrow_data::ArrayData;
use arrow_schema::{ArrowError, DataType};
use std::any::Any;
use std::sync::Arc;
/// Builder for [`FixedSizeBinaryArray`]
/// ```
/// # use arrow_array::builder::FixedSizeBinaryBuilder;
/// # use arrow_array::Array;
/// #
/// let mut builder = FixedSizeBinaryBuilder::with_capacity(3, 5);
/// // [b"hello", null, b"arrow"]
/// builder.append_value(b"hello").unwrap();
/// builder.append_null();
/// builder.append_value(b"arrow").unwrap();
///
/// let array = builder.finish();
/// assert_eq!(array.value(0), b"hello");
/// assert!(array.is_null(1));
/// assert_eq!(array.value(2), b"arrow");
/// ```
#[derive(Debug)]
pub struct FixedSizeBinaryBuilder {
values_builder: Vec<u8>,
null_buffer_builder: NullBufferBuilder,
value_length: i32,
}
impl FixedSizeBinaryBuilder {
/// Creates a new [`FixedSizeBinaryBuilder`]
pub fn new(byte_width: i32) -> Self {
Self::with_capacity(1024, byte_width)
}
/// Creates a new [`FixedSizeBinaryBuilder`], `capacity` is the number of byte slices
/// that can be appended without reallocating
pub fn with_capacity(capacity: usize, byte_width: i32) -> Self {
assert!(
byte_width >= 0,
"value length ({byte_width}) of the array must >= 0"
);
Self {
values_builder: Vec::with_capacity(capacity * byte_width as usize),
null_buffer_builder: NullBufferBuilder::new(capacity),
value_length: byte_width,
}
}
/// Appends a byte slice into the builder.
///
/// Automatically update the null buffer to delimit the slice appended in as a
/// distinct value element.
#[inline]
pub fn append_value(&mut self, value: impl AsRef<[u8]>) -> Result<(), ArrowError> {
if self.value_length != value.as_ref().len() as i32 {
Err(ArrowError::InvalidArgumentError(
"Byte slice does not have the same length as FixedSizeBinaryBuilder value lengths"
.to_string(),
))
} else {
self.values_builder.extend_from_slice(value.as_ref());
self.null_buffer_builder.append_non_null();
Ok(())
}
}
/// Append a null value to the array.
#[inline]
pub fn append_null(&mut self) {
self.values_builder
.extend(std::iter::repeat_n(0u8, self.value_length as usize));
self.null_buffer_builder.append_null();
}
/// Appends `n` `null`s into the builder.
#[inline]
pub fn append_nulls(&mut self, n: usize) {
self.values_builder
.extend(std::iter::repeat_n(0u8, self.value_length as usize * n));
self.null_buffer_builder.append_n_nulls(n);
}
/// Appends all elements in array into the builder.
pub fn append_array(&mut self, array: &FixedSizeBinaryArray) -> Result<(), ArrowError> {
if self.value_length != array.value_length() {
return Err(ArrowError::InvalidArgumentError(
"Cannot append FixedSizeBinaryArray with different value length".to_string(),
));
}
let buffer = array.value_data();
self.values_builder.extend_from_slice(buffer);
if let Some(validity) = array.nulls() {
self.null_buffer_builder.append_buffer(validity);
} else {
self.null_buffer_builder.append_n_non_nulls(array.len());
}
Ok(())
}
/// Returns the current values buffer as a slice
pub fn values_slice(&self) -> &[u8] {
self.values_builder.as_slice()
}
/// Builds the [`FixedSizeBinaryArray`] and reset this builder.
pub fn finish(&mut self) -> FixedSizeBinaryArray {
let array_length = self.len();
let array_data_builder = ArrayData::builder(DataType::FixedSizeBinary(self.value_length))
.add_buffer(std::mem::take(&mut self.values_builder).into())
.nulls(self.null_buffer_builder.finish())
.len(array_length);
let array_data = unsafe { array_data_builder.build_unchecked() };
FixedSizeBinaryArray::from(array_data)
}
/// Builds the [`FixedSizeBinaryArray`] without resetting the builder.
pub fn finish_cloned(&self) -> FixedSizeBinaryArray {
let array_length = self.len();
let values_buffer = Buffer::from_slice_ref(self.values_builder.as_slice());
let array_data_builder = ArrayData::builder(DataType::FixedSizeBinary(self.value_length))
.add_buffer(values_buffer)
.nulls(self.null_buffer_builder.finish_cloned())
.len(array_length);
let array_data = unsafe { array_data_builder.build_unchecked() };
FixedSizeBinaryArray::from(array_data)
}
/// Returns the current null buffer as a slice
pub fn validity_slice(&self) -> Option<&[u8]> {
self.null_buffer_builder.as_slice()
}
}
impl ArrayBuilder for FixedSizeBinaryBuilder {
/// Returns the builder as a non-mutable `Any` reference.
fn as_any(&self) -> &dyn Any {
self
}
/// Returns the builder as a mutable `Any` reference.
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
/// Returns the boxed builder as a box of `Any`.
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
/// Returns the number of array slots in the builder
fn len(&self) -> usize {
self.null_buffer_builder.len()
}
/// Builds the array and reset this builder.
fn finish(&mut self) -> ArrayRef {
Arc::new(self.finish())
}
/// Builds the array without resetting the builder.
fn finish_cloned(&self) -> ArrayRef {
Arc::new(self.finish_cloned())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Array;
#[test]
fn test_fixed_size_binary_builder() {
let mut builder = FixedSizeBinaryBuilder::with_capacity(3, 5);
// [b"hello", null, "arrow", null, null, "world"]
builder.append_value(b"hello").unwrap();
builder.append_null();
builder.append_value(b"arrow").unwrap();
builder.append_nulls(2);
builder.append_value(b"world").unwrap();
let array: FixedSizeBinaryArray = builder.finish();
assert_eq!(&DataType::FixedSizeBinary(5), array.data_type());
assert_eq!(6, array.len());
assert_eq!(3, array.null_count());
assert_eq!(10, array.value_offset(2));
assert_eq!(15, array.value_offset(3));
assert_eq!(5, array.value_length());
assert!(array.is_null(3));
assert!(array.is_null(4));
}
#[test]
fn test_fixed_size_binary_builder_finish_cloned() {
let mut builder = FixedSizeBinaryBuilder::with_capacity(3, 5);
// [b"hello", null, "arrow"]
builder.append_value(b"hello").unwrap();
builder.append_null();
builder.append_value(b"arrow").unwrap();
let mut array: FixedSizeBinaryArray = builder.finish_cloned();
assert_eq!(&DataType::FixedSizeBinary(5), array.data_type());
assert_eq!(3, array.len());
assert_eq!(1, array.null_count());
assert_eq!(10, array.value_offset(2));
assert_eq!(5, array.value_length());
// [b"finis", null, "clone"]
builder.append_value(b"finis").unwrap();
builder.append_null();
builder.append_value(b"clone").unwrap();
array = builder.finish();
assert_eq!(&DataType::FixedSizeBinary(5), array.data_type());
assert_eq!(6, array.len());
assert_eq!(2, array.null_count());
assert_eq!(25, array.value_offset(5));
assert_eq!(5, array.value_length());
}
#[test]
fn test_fixed_size_binary_builder_with_zero_value_length() {
let mut builder = FixedSizeBinaryBuilder::new(0);
builder.append_value(b"").unwrap();
builder.append_null();
builder.append_value(b"").unwrap();
assert!(!builder.is_empty());
let array: FixedSizeBinaryArray = builder.finish();
assert_eq!(&DataType::FixedSizeBinary(0), array.data_type());
assert_eq!(3, array.len());
assert_eq!(1, array.null_count());
assert_eq!(0, array.value_offset(2));
assert_eq!(0, array.value_length());
assert_eq!(b"", array.value(0));
assert_eq!(b"", array.value(2));
}
#[test]
#[should_panic(
expected = "Byte slice does not have the same length as FixedSizeBinaryBuilder value lengths"
)]
fn test_fixed_size_binary_builder_with_inconsistent_value_length() {
let mut builder = FixedSizeBinaryBuilder::with_capacity(1, 4);
builder.append_value(b"hello").unwrap();
}
#[test]
fn test_fixed_size_binary_builder_empty() {
let mut builder = FixedSizeBinaryBuilder::new(5);
assert!(builder.is_empty());
let fixed_size_binary_array = builder.finish();
assert_eq!(
&DataType::FixedSizeBinary(5),
fixed_size_binary_array.data_type()
);
assert_eq!(0, fixed_size_binary_array.len());
}
#[test]
#[should_panic(expected = "value length (-1) of the array must >= 0")]
fn test_fixed_size_binary_builder_invalid_value_length() {
let _ = FixedSizeBinaryBuilder::with_capacity(15, -1);
}
#[test]
fn test_fixed_size_binary_builder_append_array() {
let mut other_builder = FixedSizeBinaryBuilder::with_capacity(3, 5);
other_builder.append_value(b"hello").unwrap();
other_builder.append_null();
other_builder.append_value(b"arrow").unwrap();
let other_array = other_builder.finish();
let mut builder = FixedSizeBinaryBuilder::with_capacity(6, 5);
builder.append_array(&other_array).unwrap();
// Append again to test if breaks when appending multiple times
builder.append_array(&other_array).unwrap();
let array = builder.finish();
assert_eq!(array.value_length(), other_array.value_length());
assert_eq!(&DataType::FixedSizeBinary(5), array.data_type());
assert_eq!(6, array.len());
assert_eq!(2, array.null_count());
for i in 0..6 {
assert_eq!(i * 5, array.value_offset(i as usize));
}
assert_eq!(b"hello", array.value(0));
assert!(array.is_null(1));
assert_eq!(b"arrow", array.value(2));
assert_eq!(b"hello", array.value(3));
assert!(array.is_null(4));
assert_eq!(b"arrow", array.value(5));
}
#[test]
#[should_panic(expected = "Cannot append FixedSizeBinaryArray with different value length")]
fn test_fixed_size_binary_builder_append_array_invalid_value_length() {
let mut other_builder = FixedSizeBinaryBuilder::with_capacity(3, 4);
other_builder.append_value(b"test").unwrap();
let other_array = other_builder.finish();
let mut builder = FixedSizeBinaryBuilder::with_capacity(3, 5);
builder.append_array(&other_array).unwrap();
}
}
@@ -0,0 +1,670 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::builder::{ArrayBuilder, FixedSizeBinaryBuilder, PrimitiveBuilder};
use crate::types::ArrowDictionaryKeyType;
use crate::{Array, ArrayRef, DictionaryArray, PrimitiveArray};
use arrow_buffer::ArrowNativeType;
use arrow_schema::DataType::FixedSizeBinary;
use arrow_schema::{ArrowError, DataType};
use hashbrown::HashTable;
use num_traits::NumCast;
use std::any::Any;
use std::sync::Arc;
/// Builder for [`DictionaryArray`] of [`FixedSizeBinaryArray`]
///
/// The output array has a dictionary of unique, fixed-size binary values. The
/// builder handles deduplication.
///
/// # Example
/// ```
/// # use arrow_array::builder::{FixedSizeBinaryDictionaryBuilder};
/// # use arrow_array::array::{Array, FixedSizeBinaryArray};
/// # use arrow_array::DictionaryArray;
/// # use arrow_array::types::Int8Type;
/// // Build 3 byte FixedBinaryArrays
/// let byte_width = 3;
/// let mut builder = FixedSizeBinaryDictionaryBuilder::<Int8Type>::new(3);
/// builder.append("abc").unwrap();
/// builder.append_null();
/// builder.append(b"def").unwrap();
/// builder.append(b"def").unwrap(); // duplicate value
/// // Result is a Dictionary Array
/// let array = builder.finish();
/// let dict_array = array.as_any().downcast_ref::<DictionaryArray<Int8Type>>().unwrap();
/// // The array represents "abc", null, "def", "def"
/// assert_eq!(array.keys().len(), 4);
/// // but there are only 2 unique values
/// assert_eq!(array.values().len(), 2);
/// let values = dict_array.values().as_any().downcast_ref::<FixedSizeBinaryArray>().unwrap();
/// assert_eq!(values.value(0), "abc".as_bytes());
/// assert_eq!(values.value(1), "def".as_bytes());
/// ```
///
/// [`FixedSizeBinaryArray`]: crate::FixedSizeBinaryArray
#[derive(Debug)]
pub struct FixedSizeBinaryDictionaryBuilder<K>
where
K: ArrowDictionaryKeyType,
{
state: ahash::RandomState,
dedup: HashTable<usize>,
keys_builder: PrimitiveBuilder<K>,
values_builder: FixedSizeBinaryBuilder,
byte_width: i32,
}
impl<K> FixedSizeBinaryDictionaryBuilder<K>
where
K: ArrowDictionaryKeyType,
{
/// Creates a new `FixedSizeBinaryDictionaryBuilder`
pub fn new(byte_width: i32) -> Self {
let keys_builder = PrimitiveBuilder::new();
let values_builder = FixedSizeBinaryBuilder::new(byte_width);
Self {
state: Default::default(),
dedup: HashTable::with_capacity(keys_builder.capacity()),
keys_builder,
values_builder,
byte_width,
}
}
/// Creates a new `FixedSizeBinaryDictionaryBuilder` with the provided capacities
///
/// `keys_capacity`: the number of keys, i.e. length of array to build
/// `value_capacity`: the number of distinct dictionary values, i.e. size of dictionary
/// `byte_width`: the byte width for individual values in the values array
pub fn with_capacity(keys_capacity: usize, value_capacity: usize, byte_width: i32) -> Self {
Self {
state: Default::default(),
dedup: Default::default(),
keys_builder: PrimitiveBuilder::with_capacity(keys_capacity),
values_builder: FixedSizeBinaryBuilder::with_capacity(value_capacity, byte_width),
byte_width,
}
}
/// Creates a new `FixedSizeBinaryDictionaryBuilder` from the existing builder with the same
/// keys and values, but with a new data type for the keys.
///
/// # Example
/// ```
/// # use arrow_array::builder::FixedSizeBinaryDictionaryBuilder;
/// # use arrow_array::types::{UInt8Type, UInt16Type, UInt64Type};
/// # use arrow_array::UInt16Array;
/// # use arrow_schema::ArrowError;
///
/// let mut u8_keyed_builder = FixedSizeBinaryDictionaryBuilder::<UInt8Type>::new(2);
/// // appending too many values causes the dictionary to overflow
/// for i in 0..=255 {
/// u8_keyed_builder.append_value(vec![0, i]);
/// }
/// let result = u8_keyed_builder.append(vec![1, 0]);
/// assert!(matches!(result, Err(ArrowError::DictionaryKeyOverflowError{})));
///
/// // we need to upgrade to a larger key type
/// let mut u16_keyed_builder = FixedSizeBinaryDictionaryBuilder::<UInt16Type>::try_new_from_builder(u8_keyed_builder).unwrap();
/// let dictionary_array = u16_keyed_builder.finish();
/// let keys = dictionary_array.keys();
///
/// assert_eq!(keys, &UInt16Array::from_iter(0..256));
/// ```
pub fn try_new_from_builder<K2>(
mut source: FixedSizeBinaryDictionaryBuilder<K2>,
) -> Result<Self, ArrowError>
where
K::Native: NumCast,
K2: ArrowDictionaryKeyType,
K2::Native: NumCast,
{
let state = source.state;
let dedup = source.dedup;
let values_builder = source.values_builder;
let byte_width = source.byte_width;
let source_keys = source.keys_builder.finish();
let new_keys: PrimitiveArray<K> = source_keys.try_unary(|value| {
num_traits::cast::cast::<K2::Native, K::Native>(value).ok_or_else(|| {
ArrowError::CastError(format!(
"Can't cast dictionary keys from source type {:?} to type {:?}",
K2::DATA_TYPE,
K::DATA_TYPE
))
})
})?;
// drop source key here because currently source_keys and new_keys are holding reference to
// the same underlying null_buffer. Below we want to call new_keys.into_builder() it must
// be the only reference holder.
drop(source_keys);
Ok(Self {
state,
dedup,
keys_builder: new_keys
.into_builder()
.expect("underlying buffer has no references"),
values_builder,
byte_width,
})
}
}
impl<K> ArrayBuilder for FixedSizeBinaryDictionaryBuilder<K>
where
K: ArrowDictionaryKeyType,
{
/// Returns the builder as an non-mutable `Any` reference.
fn as_any(&self) -> &dyn Any {
self
}
/// Returns the builder as an mutable `Any` reference.
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
/// Returns the boxed builder as a box of `Any`.
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
/// Returns the number of array slots in the builder
fn len(&self) -> usize {
self.keys_builder.len()
}
/// Builds the array and reset this builder.
fn finish(&mut self) -> ArrayRef {
Arc::new(self.finish())
}
/// Builds the array without resetting the builder.
fn finish_cloned(&self) -> ArrayRef {
Arc::new(self.finish_cloned())
}
}
impl<K> FixedSizeBinaryDictionaryBuilder<K>
where
K: ArrowDictionaryKeyType,
{
fn get_or_insert_key(&mut self, value: impl AsRef<[u8]>) -> Result<K::Native, ArrowError> {
let value_bytes: &[u8] = value.as_ref();
let state = &self.state;
let storage = &mut self.values_builder;
let hash = state.hash_one(value_bytes);
let idx = *self
.dedup
.entry(
hash,
|idx| value_bytes == get_bytes(storage, self.byte_width, *idx),
|idx| state.hash_one(get_bytes(storage, self.byte_width, *idx)),
)
.or_insert_with(|| {
let idx = storage.len();
let _ = storage.append_value(value);
idx
})
.get();
let key = K::Native::from_usize(idx).ok_or(ArrowError::DictionaryKeyOverflowError)?;
Ok(key)
}
/// Append a value to the array. Return an existing index
/// if already present in the values array or a new index if the
/// value is appended to the values array.
///
/// Returns an error if the new index would overflow the key type.
pub fn append(&mut self, value: impl AsRef<[u8]>) -> Result<K::Native, ArrowError> {
if self.byte_width != value.as_ref().len() as i32 {
Err(ArrowError::InvalidArgumentError(format!(
"Invalid input length passed to FixedSizeBinaryBuilder. Expected {} got {}",
self.byte_width,
value.as_ref().len()
)))
} else {
let key = self.get_or_insert_key(value)?;
self.keys_builder.append_value(key);
Ok(key)
}
}
/// Append a value multiple times to the array.
/// This is the same as [`Self::append`] but allows to append the same value multiple times without doing multiple lookups.
///
/// Returns an error if the new index would overflow the key type.
pub fn append_n(
&mut self,
value: impl AsRef<[u8]>,
count: usize,
) -> Result<K::Native, ArrowError> {
if self.byte_width != value.as_ref().len() as i32 {
Err(ArrowError::InvalidArgumentError(format!(
"Invalid input length passed to FixedSizeBinaryBuilder. Expected {} got {}",
self.byte_width,
value.as_ref().len()
)))
} else {
let key = self.get_or_insert_key(value)?;
self.keys_builder.append_value_n(key, count);
Ok(key)
}
}
/// Appends a null slot into the builder
#[inline]
pub fn append_null(&mut self) {
self.keys_builder.append_null()
}
/// Appends `n` `null`s into the builder.
#[inline]
pub fn append_nulls(&mut self, n: usize) {
self.keys_builder.append_nulls(n);
}
/// Infallibly append a value to this builder
///
/// # Panics
///
/// Panics if the resulting length of the dictionary values array would exceed `T::Native::MAX`
pub fn append_value(&mut self, value: impl AsRef<[u8]>) {
self.append(value).expect("dictionary key overflow");
}
/// Builds the `DictionaryArray` and reset this builder.
pub fn finish(&mut self) -> DictionaryArray<K> {
self.dedup.clear();
let values = self.values_builder.finish();
let keys = self.keys_builder.finish();
let data_type = DataType::Dictionary(
Box::new(K::DATA_TYPE),
Box::new(FixedSizeBinary(self.byte_width)),
);
let builder = keys
.into_data()
.into_builder()
.data_type(data_type)
.child_data(vec![values.into_data()]);
DictionaryArray::from(unsafe { builder.build_unchecked() })
}
/// Builds the `DictionaryArray` without resetting the builder.
pub fn finish_cloned(&self) -> DictionaryArray<K> {
let values = self.values_builder.finish_cloned();
let keys = self.keys_builder.finish_cloned();
let data_type = DataType::Dictionary(
Box::new(K::DATA_TYPE),
Box::new(FixedSizeBinary(self.byte_width)),
);
let builder = keys
.into_data()
.into_builder()
.data_type(data_type)
.child_data(vec![values.into_data()]);
DictionaryArray::from(unsafe { builder.build_unchecked() })
}
/// Builds the `DictionaryArray` without resetting the values builder or
/// the internal de-duplication map.
///
/// The advantage of doing this is that the values will represent the entire
/// set of what has been built so-far by this builder and ensures
/// consistency in the assignment of keys to values across multiple calls
/// to `finish_preserve_values`. This enables ipc writers to efficiently
/// emit delta dictionaries.
///
/// The downside to this is that building the record requires creating a
/// copy of the values, which can become slowly more expensive if the
/// dictionary grows.
///
/// Additionally, if record batches from multiple different dictionary
/// builders for the same column are fed into a single ipc writer, beware
/// that entire dictionaries are likely to be re-sent frequently even when
/// the majority of the values are not used by the current record batch.
pub fn finish_preserve_values(&mut self) -> DictionaryArray<K> {
let values = self.values_builder.finish_cloned();
let keys = self.keys_builder.finish();
let data_type = DataType::Dictionary(
Box::new(K::DATA_TYPE),
Box::new(FixedSizeBinary(self.byte_width)),
);
let builder = keys
.into_data()
.into_builder()
.data_type(data_type)
.child_data(vec![values.into_data()]);
DictionaryArray::from(unsafe { builder.build_unchecked() })
}
}
fn get_bytes(values: &FixedSizeBinaryBuilder, byte_width: i32, idx: usize) -> &[u8] {
let values = values.values_slice();
let start = idx * byte_width.as_usize();
let end = idx * byte_width.as_usize() + byte_width.as_usize();
&values[start..end]
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::{Int8Type, Int16Type, Int32Type, UInt8Type, UInt16Type};
use crate::{ArrowPrimitiveType, FixedSizeBinaryArray, Int8Array};
#[test]
fn test_fixed_size_dictionary_builder() {
let values = ["abc", "def"];
let mut b = FixedSizeBinaryDictionaryBuilder::<Int8Type>::new(3);
assert_eq!(b.append(values[0]).unwrap(), 0);
b.append_null();
assert_eq!(b.append(values[1]).unwrap(), 1);
assert_eq!(b.append(values[1]).unwrap(), 1);
assert_eq!(b.append(values[0]).unwrap(), 0);
b.append_nulls(2);
assert_eq!(b.append(values[0]).unwrap(), 0);
let array = b.finish();
assert_eq!(
array.keys(),
&Int8Array::from(vec![
Some(0),
None,
Some(1),
Some(1),
Some(0),
None,
None,
Some(0)
]),
);
// Values are polymorphic and so require a downcast.
let ava = array
.values()
.as_any()
.downcast_ref::<FixedSizeBinaryArray>()
.unwrap();
assert_eq!(ava.value(0), values[0].as_bytes());
assert_eq!(ava.value(1), values[1].as_bytes());
}
#[test]
fn test_fixed_size_dictionary_builder_append_n() {
let values = ["abc", "def"];
let mut b = FixedSizeBinaryDictionaryBuilder::<Int8Type>::new(3);
assert_eq!(b.append_n(values[0], 2).unwrap(), 0);
assert_eq!(b.append_n(values[1], 3).unwrap(), 1);
assert_eq!(b.append_n(values[0], 2).unwrap(), 0);
let array = b.finish();
assert_eq!(
array.keys(),
&Int8Array::from(vec![
Some(0),
Some(0),
Some(1),
Some(1),
Some(1),
Some(0),
Some(0),
]),
);
// Values are polymorphic and so require a downcast.
let ava = array
.values()
.as_any()
.downcast_ref::<FixedSizeBinaryArray>()
.unwrap();
assert_eq!(ava.value(0), values[0].as_bytes());
assert_eq!(ava.value(1), values[1].as_bytes());
}
#[test]
fn test_fixed_size_dictionary_builder_wrong_size() {
let mut b = FixedSizeBinaryDictionaryBuilder::<Int8Type>::new(3);
let err = b.append(b"too long").unwrap_err().to_string();
assert_eq!(
err,
"Invalid argument error: Invalid input length passed to FixedSizeBinaryBuilder. Expected 3 got 8"
);
let err = b.append("").unwrap_err().to_string();
assert_eq!(
err,
"Invalid argument error: Invalid input length passed to FixedSizeBinaryBuilder. Expected 3 got 0"
);
let err = b.append_n("a", 3).unwrap_err().to_string();
assert_eq!(
err,
"Invalid argument error: Invalid input length passed to FixedSizeBinaryBuilder. Expected 3 got 1"
);
}
#[test]
fn test_fixed_size_dictionary_builder_finish_cloned() {
let values = ["abc", "def", "ghi"];
let mut builder = FixedSizeBinaryDictionaryBuilder::<Int8Type>::new(3);
builder.append(values[0]).unwrap();
builder.append_null();
builder.append(values[1]).unwrap();
builder.append(values[1]).unwrap();
builder.append(values[0]).unwrap();
let mut array = builder.finish_cloned();
assert_eq!(
array.keys(),
&Int8Array::from(vec![Some(0), None, Some(1), Some(1), Some(0)])
);
// Values are polymorphic and so require a downcast.
let ava = array
.values()
.as_any()
.downcast_ref::<FixedSizeBinaryArray>()
.unwrap();
assert_eq!(ava.value(0), values[0].as_bytes());
assert_eq!(ava.value(1), values[1].as_bytes());
builder.append(values[0]).unwrap();
builder.append(values[2]).unwrap();
builder.append(values[1]).unwrap();
array = builder.finish();
assert_eq!(
array.keys(),
&Int8Array::from(vec![
Some(0),
None,
Some(1),
Some(1),
Some(0),
Some(0),
Some(2),
Some(1)
])
);
// Values are polymorphic and so require a downcast.
let ava2 = array
.values()
.as_any()
.downcast_ref::<FixedSizeBinaryArray>()
.unwrap();
assert_eq!(ava2.value(0), values[0].as_bytes());
assert_eq!(ava2.value(1), values[1].as_bytes());
assert_eq!(ava2.value(2), values[2].as_bytes());
}
fn _test_try_new_from_builder_generic_for_key_types<K1, K2>(values: Vec<[u8; 3]>)
where
K1: ArrowDictionaryKeyType,
K1::Native: NumCast,
K2: ArrowDictionaryKeyType,
K2::Native: NumCast + From<u8>,
{
let mut source = FixedSizeBinaryDictionaryBuilder::<K1>::new(3);
source.append_value(values[0]);
source.append_null();
source.append_value(values[1]);
source.append_value(values[2]);
let mut result =
FixedSizeBinaryDictionaryBuilder::<K2>::try_new_from_builder(source).unwrap();
let array = result.finish();
let mut expected_keys_builder = PrimitiveBuilder::<K2>::new();
expected_keys_builder
.append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(0u8));
expected_keys_builder.append_null();
expected_keys_builder
.append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(1u8));
expected_keys_builder
.append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(2u8));
let expected_keys = expected_keys_builder.finish();
assert_eq!(array.keys(), &expected_keys);
let av = array.values();
let ava = av.as_any().downcast_ref::<FixedSizeBinaryArray>().unwrap();
assert_eq!(ava.value(0), values[0]);
assert_eq!(ava.value(1), values[1]);
assert_eq!(ava.value(2), values[2]);
}
#[test]
fn test_try_new_from_builder() {
let values = vec![[1, 2, 3], [5, 6, 7], [6, 7, 8]];
// test cast to bigger size unsigned
_test_try_new_from_builder_generic_for_key_types::<UInt8Type, UInt16Type>(values.clone());
// test cast going to smaller size unsigned
_test_try_new_from_builder_generic_for_key_types::<UInt16Type, UInt8Type>(values.clone());
// test cast going to bigger size signed
_test_try_new_from_builder_generic_for_key_types::<Int8Type, Int16Type>(values.clone());
// test cast going to smaller size signed
_test_try_new_from_builder_generic_for_key_types::<Int32Type, Int16Type>(values.clone());
// test going from signed to signed for different size changes
_test_try_new_from_builder_generic_for_key_types::<UInt8Type, Int16Type>(values.clone());
_test_try_new_from_builder_generic_for_key_types::<Int8Type, UInt8Type>(values.clone());
_test_try_new_from_builder_generic_for_key_types::<Int8Type, UInt16Type>(values.clone());
_test_try_new_from_builder_generic_for_key_types::<Int32Type, Int16Type>(values.clone());
}
#[test]
fn test_try_new_from_builder_cast_fails() {
let mut source_builder = FixedSizeBinaryDictionaryBuilder::<UInt16Type>::new(2);
for i in 0u16..257u16 {
source_builder.append_value(vec![(i >> 8) as u8, i as u8]);
}
// there should be too many values that we can't downcast to the underlying type
// we have keys that wouldn't fit into UInt8Type
let result =
FixedSizeBinaryDictionaryBuilder::<UInt8Type>::try_new_from_builder(source_builder);
assert!(result.is_err());
if let Err(e) = result {
assert!(matches!(e, ArrowError::CastError(_)));
assert_eq!(
e.to_string(),
"Cast error: Can't cast dictionary keys from source type UInt16 to type UInt8"
);
}
}
#[test]
fn test_finish_preserve_values() {
// Create the first dictionary
let mut builder = FixedSizeBinaryDictionaryBuilder::<Int32Type>::new(3);
builder.append_value("aaa");
builder.append_value("bbb");
builder.append_value("ccc");
let dict = builder.finish_preserve_values();
assert_eq!(dict.keys().values(), &[0, 1, 2]);
let values = dict
.downcast_dict::<FixedSizeBinaryArray>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(
values,
vec![
Some("aaa".as_bytes()),
Some("bbb".as_bytes()),
Some("ccc".as_bytes())
]
);
// Create a new dictionary
builder.append_value("ddd");
builder.append_value("eee");
let dict2 = builder.finish_preserve_values();
// Make sure the keys are assigned after the old ones and we have the
// right values
assert_eq!(dict2.keys().values(), &[3, 4]);
let values = dict2
.downcast_dict::<FixedSizeBinaryArray>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(values, [Some("ddd".as_bytes()), Some("eee".as_bytes())]);
// Check that we have all of the expected values
let all_values = dict2
.values()
.as_any()
.downcast_ref::<FixedSizeBinaryArray>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(
all_values,
[
Some("aaa".as_bytes()),
Some("bbb".as_bytes()),
Some("ccc".as_bytes()),
Some("ddd".as_bytes()),
Some("eee".as_bytes())
]
);
}
}
@@ -0,0 +1,494 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::builder::ArrayBuilder;
use crate::{ArrayRef, FixedSizeListArray};
use arrow_buffer::NullBufferBuilder;
use arrow_schema::{Field, FieldRef};
use std::any::Any;
use std::sync::Arc;
/// Builder for [`FixedSizeListArray`]
/// ```
/// use arrow_array::{builder::{Int32Builder, FixedSizeListBuilder}, Array, Int32Array};
/// let values_builder = Int32Builder::new();
/// let mut builder = FixedSizeListBuilder::new(values_builder, 3);
///
/// // [[0, 1, 2], null, [3, null, 5], [6, 7, null]]
/// builder.values().append_value(0);
/// builder.values().append_value(1);
/// builder.values().append_value(2);
/// builder.append(true);
/// builder.values().append_null();
/// builder.values().append_null();
/// builder.values().append_null();
/// builder.append(false);
/// builder.values().append_value(3);
/// builder.values().append_null();
/// builder.values().append_value(5);
/// builder.append(true);
/// builder.values().append_value(6);
/// builder.values().append_value(7);
/// builder.values().append_null();
/// builder.append(true);
/// let list_array = builder.finish();
/// assert_eq!(
/// *list_array.value(0),
/// Int32Array::from(vec![Some(0), Some(1), Some(2)])
/// );
/// assert!(list_array.is_null(1));
/// assert_eq!(
/// *list_array.value(2),
/// Int32Array::from(vec![Some(3), None, Some(5)])
/// );
/// assert_eq!(
/// *list_array.value(3),
/// Int32Array::from(vec![Some(6), Some(7), None])
/// )
/// ```
///
#[derive(Debug)]
pub struct FixedSizeListBuilder<T: ArrayBuilder> {
null_buffer_builder: NullBufferBuilder,
values_builder: T,
list_len: i32,
field: Option<FieldRef>,
}
impl<T: ArrayBuilder> FixedSizeListBuilder<T> {
/// Creates a new [`FixedSizeListBuilder`] from a given values array builder
/// `value_length` is the number of values within each array
pub fn new(values_builder: T, value_length: i32) -> Self {
let capacity = values_builder
.len()
.checked_div(value_length as _)
.unwrap_or_default();
Self::with_capacity(values_builder, value_length, capacity)
}
/// Creates a new [`FixedSizeListBuilder`] from a given values array builder
/// `value_length` is the number of values within each array
/// `capacity` is the number of items to pre-allocate space for in this builder
pub fn with_capacity(values_builder: T, value_length: i32, capacity: usize) -> Self {
Self {
null_buffer_builder: NullBufferBuilder::new(capacity),
values_builder,
list_len: value_length,
field: None,
}
}
/// Override the field passed to [`FixedSizeListArray::new`]
///
/// By default, a nullable field is created with the name `item`
///
/// Note: [`Self::finish`] and [`Self::finish_cloned`] will panic if the
/// field's data type does not match that of `T`
pub fn with_field(self, field: impl Into<FieldRef>) -> Self {
Self {
field: Some(field.into()),
..self
}
}
}
impl<T: ArrayBuilder> ArrayBuilder for FixedSizeListBuilder<T>
where
T: 'static,
{
/// Returns the builder as a non-mutable `Any` reference.
fn as_any(&self) -> &dyn Any {
self
}
/// Returns the builder as a mutable `Any` reference.
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
/// Returns the boxed builder as a box of `Any`.
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
/// Returns the number of array slots in the builder
fn len(&self) -> usize {
self.null_buffer_builder.len()
}
/// Builds the array and reset this builder.
fn finish(&mut self) -> ArrayRef {
Arc::new(self.finish())
}
/// Builds the array without resetting the builder.
fn finish_cloned(&self) -> ArrayRef {
Arc::new(self.finish_cloned())
}
}
impl<T: ArrayBuilder> FixedSizeListBuilder<T>
where
T: 'static,
{
/// Returns the child array builder as a mutable reference.
///
/// This mutable reference can be used to append values into the child array builder,
/// but you must call [`append`](#method.append) to delimit each distinct list value.
pub fn values(&mut self) -> &mut T {
&mut self.values_builder
}
/// Returns the length of the list
pub fn value_length(&self) -> i32 {
self.list_len
}
/// Finish the current fixed-length list array slot
#[inline]
pub fn append(&mut self, is_valid: bool) {
self.null_buffer_builder.append(is_valid);
}
/// Builds the [`FixedSizeListBuilder`] and reset this builder.
pub fn finish(&mut self) -> FixedSizeListArray {
let len = self.len();
let values = self.values_builder.finish();
let nulls = self.null_buffer_builder.finish();
assert_eq!(
values.len(),
len * self.list_len as usize,
"Length of the child array ({}) must be the multiple of the value length ({}) and the array length ({}).",
values.len(),
self.list_len,
len,
);
let field = self
.field
.clone()
.unwrap_or_else(|| Arc::new(Field::new_list_field(values.data_type().clone(), true)));
FixedSizeListArray::new(field, self.list_len, values, nulls)
}
/// Builds the [`FixedSizeListBuilder`] without resetting the builder.
pub fn finish_cloned(&self) -> FixedSizeListArray {
let len = self.len();
let values = self.values_builder.finish_cloned();
let nulls = self.null_buffer_builder.finish_cloned();
assert_eq!(
values.len(),
len * self.list_len as usize,
"Length of the child array ({}) must be the multiple of the value length ({}) and the array length ({}).",
values.len(),
self.list_len,
len,
);
let field = self
.field
.clone()
.unwrap_or_else(|| Arc::new(Field::new_list_field(values.data_type().clone(), true)));
FixedSizeListArray::new(field, self.list_len, values, nulls)
}
/// Returns the current null buffer as a slice
pub fn validity_slice(&self) -> Option<&[u8]> {
self.null_buffer_builder.as_slice()
}
}
#[cfg(test)]
mod tests {
use super::*;
use arrow_schema::DataType;
use crate::Array;
use crate::Int32Array;
use crate::builder::Int32Builder;
fn make_list_builder(
include_null_element: bool,
include_null_in_values: bool,
) -> FixedSizeListBuilder<crate::builder::PrimitiveBuilder<crate::types::Int32Type>> {
let values_builder = Int32Builder::new();
let mut builder = FixedSizeListBuilder::new(values_builder, 3);
builder.values().append_value(0);
builder.values().append_value(1);
builder.values().append_value(2);
builder.append(true);
builder.values().append_value(2);
builder.values().append_value(3);
builder.values().append_value(4);
builder.append(true);
if include_null_element {
builder.values().append_null();
builder.values().append_null();
builder.values().append_null();
builder.append(false);
} else {
builder.values().append_value(2);
builder.values().append_value(3);
builder.values().append_value(4);
builder.append(true);
}
if include_null_in_values {
builder.values().append_value(3);
builder.values().append_null();
builder.values().append_value(5);
builder.append(true);
} else {
builder.values().append_value(3);
builder.values().append_value(4);
builder.values().append_value(5);
builder.append(true);
}
builder
}
#[test]
fn test_fixed_size_list_array_builder() {
let mut builder = make_list_builder(true, true);
let list_array = builder.finish();
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(4, list_array.len());
assert_eq!(1, list_array.null_count());
assert_eq!(6, list_array.value_offset(2));
assert_eq!(3, list_array.value_length());
}
#[test]
fn test_fixed_size_list_array_builder_with_field() {
let builder = make_list_builder(false, false);
let mut builder = builder.with_field(Field::new("list_element", DataType::Int32, false));
let list_array = builder.finish();
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(4, list_array.len());
assert_eq!(0, list_array.null_count());
assert_eq!(6, list_array.value_offset(2));
assert_eq!(3, list_array.value_length());
}
#[test]
fn test_fixed_size_list_array_builder_with_field_and_null() {
let builder = make_list_builder(true, false);
let mut builder = builder.with_field(Field::new("list_element", DataType::Int32, false));
let list_array = builder.finish();
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(4, list_array.len());
assert_eq!(1, list_array.null_count());
assert_eq!(6, list_array.value_offset(2));
assert_eq!(3, list_array.value_length());
}
#[test]
#[should_panic(expected = "Found unmasked nulls for non-nullable FixedSizeListArray")]
fn test_fixed_size_list_array_builder_with_field_null_panic() {
let builder = make_list_builder(true, true);
let mut builder = builder.with_field(Field::new("list_item", DataType::Int32, false));
builder.finish();
}
#[test]
#[should_panic(expected = "FixedSizeListArray expected data type Int64 got Int32")]
fn test_fixed_size_list_array_builder_with_field_type_panic() {
let values_builder = Int32Builder::new();
let builder = FixedSizeListBuilder::new(values_builder, 3);
let mut builder = builder.with_field(Field::new("list_item", DataType::Int64, true));
// [[0, 1, 2], null, [3, null, 5], [6, 7, null]]
builder.values().append_value(0);
builder.values().append_value(1);
builder.values().append_value(2);
builder.append(true);
builder.values().append_null();
builder.values().append_null();
builder.values().append_null();
builder.append(false);
builder.values().append_value(3);
builder.values().append_value(4);
builder.values().append_value(5);
builder.append(true);
builder.finish();
}
#[test]
fn test_fixed_size_list_array_builder_cloned_with_field() {
let builder = make_list_builder(true, true);
let builder = builder.with_field(Field::new("list_element", DataType::Int32, true));
let list_array = builder.finish_cloned();
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(4, list_array.len());
assert_eq!(1, list_array.null_count());
assert_eq!(6, list_array.value_offset(2));
assert_eq!(3, list_array.value_length());
}
#[test]
#[should_panic(expected = "Found unmasked nulls for non-nullable FixedSizeListArray")]
fn test_fixed_size_list_array_builder_cloned_with_field_null_panic() {
let builder = make_list_builder(true, true);
let builder = builder.with_field(Field::new("list_item", DataType::Int32, false));
builder.finish_cloned();
}
#[test]
fn test_fixed_size_list_array_builder_cloned_with_field_and_null() {
let builder = make_list_builder(true, false);
let mut builder = builder.with_field(Field::new("list_element", DataType::Int32, false));
let list_array = builder.finish();
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(4, list_array.len());
assert_eq!(1, list_array.null_count());
assert_eq!(6, list_array.value_offset(2));
assert_eq!(3, list_array.value_length());
}
#[test]
#[should_panic(expected = "FixedSizeListArray expected data type Int64 got Int32")]
fn test_fixed_size_list_array_builder_cloned_with_field_type_panic() {
let builder = make_list_builder(false, false);
let builder = builder.with_field(Field::new("list_item", DataType::Int64, true));
builder.finish_cloned();
}
#[test]
fn test_fixed_size_list_array_builder_finish_cloned() {
let mut builder = make_list_builder(true, true);
let mut list_array = builder.finish_cloned();
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(4, list_array.len());
assert_eq!(1, list_array.null_count());
assert_eq!(3, list_array.value_length());
builder.values().append_value(6);
builder.values().append_value(7);
builder.values().append_null();
builder.append(true);
builder.values().append_null();
builder.values().append_null();
builder.values().append_null();
builder.append(false);
list_array = builder.finish();
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(6, list_array.len());
assert_eq!(2, list_array.null_count());
assert_eq!(6, list_array.value_offset(2));
assert_eq!(3, list_array.value_length());
}
#[test]
fn test_fixed_size_list_array_builder_with_field_empty() {
let values_builder = Int32Array::builder(0);
let mut builder = FixedSizeListBuilder::new(values_builder, 3).with_field(Field::new(
"list_item",
DataType::Int32,
false,
));
assert!(builder.is_empty());
let arr = builder.finish();
assert_eq!(0, arr.len());
assert_eq!(0, builder.len());
}
#[test]
fn test_fixed_size_list_array_builder_cloned_with_field_empty() {
let values_builder = Int32Array::builder(0);
let builder = FixedSizeListBuilder::new(values_builder, 3).with_field(Field::new(
"list_item",
DataType::Int32,
false,
));
assert!(builder.is_empty());
let arr = builder.finish_cloned();
assert_eq!(0, arr.len());
assert_eq!(0, builder.len());
}
#[test]
fn test_fixed_size_list_array_builder_empty() {
let values_builder = Int32Array::builder(5);
let mut builder = FixedSizeListBuilder::new(values_builder, 3);
assert!(builder.is_empty());
let arr = builder.finish();
assert_eq!(0, arr.len());
assert_eq!(0, builder.len());
}
#[test]
fn test_fixed_size_list_array_builder_finish() {
let values_builder = Int32Array::builder(5);
let mut builder = FixedSizeListBuilder::new(values_builder, 3);
builder.values().append_slice(&[1, 2, 3]);
builder.append(true);
builder.values().append_slice(&[4, 5, 6]);
builder.append(true);
let mut arr = builder.finish();
assert_eq!(2, arr.len());
assert_eq!(0, builder.len());
builder.values().append_slice(&[7, 8, 9]);
builder.append(true);
arr = builder.finish();
assert_eq!(1, arr.len());
assert_eq!(0, builder.len());
}
#[test]
#[should_panic(
expected = "Length of the child array (10) must be the multiple of the value length (3) and the array length (3)."
)]
fn test_fixed_size_list_array_builder_fail() {
let values_builder = Int32Array::builder(5);
let mut builder = FixedSizeListBuilder::new(values_builder, 3);
builder.values().append_slice(&[1, 2, 3]);
builder.append(true);
builder.values().append_slice(&[4, 5, 6]);
builder.append(true);
builder.values().append_slice(&[7, 8, 9, 10]);
builder.append(true);
builder.finish();
}
}
@@ -0,0 +1,517 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::types::bytes::ByteArrayNativeType;
use std::{any::Any, sync::Arc};
use crate::{
ArrayRef, ArrowPrimitiveType, RunArray,
types::{BinaryType, ByteArrayType, LargeBinaryType, LargeUtf8Type, RunEndIndexType, Utf8Type},
};
use super::{ArrayBuilder, GenericByteBuilder, PrimitiveBuilder};
use arrow_buffer::ArrowNativeType;
/// Builder for [`RunArray`] of [`GenericByteArray`](crate::array::GenericByteArray)
///
/// # Example:
///
/// ```
///
/// # use arrow_array::builder::GenericByteRunBuilder;
/// # use arrow_array::{GenericByteArray, BinaryArray};
/// # use arrow_array::types::{BinaryType, Int16Type};
/// # use arrow_array::{Array, Int16Array};
/// # use arrow_array::cast::AsArray;
///
/// let mut builder =
/// GenericByteRunBuilder::<Int16Type, BinaryType>::new();
/// builder.extend([Some(b"abc"), Some(b"abc"), None, Some(b"def")].into_iter());
/// builder.append_value(b"def");
/// builder.append_null();
/// let array = builder.finish();
///
/// assert_eq!(array.run_ends().values(), &[2, 3, 5, 6]);
///
/// let av = array.values();
///
/// assert!(!av.is_null(0));
/// assert!(av.is_null(1));
/// assert!(!av.is_null(2));
/// assert!(av.is_null(3));
///
/// // Values are polymorphic and so require a downcast.
/// let ava: &BinaryArray = av.as_binary();
///
/// assert_eq!(ava.value(0), b"abc");
/// assert_eq!(ava.value(2), b"def");
/// ```
#[derive(Debug)]
pub struct GenericByteRunBuilder<R, V>
where
R: ArrowPrimitiveType,
V: ByteArrayType,
{
run_ends_builder: PrimitiveBuilder<R>,
values_builder: GenericByteBuilder<V>,
current_value: Vec<u8>,
has_current_value: bool,
current_run_end_index: usize,
prev_run_end_index: usize,
}
impl<R, V> Default for GenericByteRunBuilder<R, V>
where
R: ArrowPrimitiveType,
V: ByteArrayType,
{
fn default() -> Self {
Self::new()
}
}
impl<R, V> GenericByteRunBuilder<R, V>
where
R: ArrowPrimitiveType,
V: ByteArrayType,
{
/// Creates a new `GenericByteRunBuilder`
pub fn new() -> Self {
Self {
run_ends_builder: PrimitiveBuilder::new(),
values_builder: GenericByteBuilder::<V>::new(),
current_value: Vec::new(),
has_current_value: false,
current_run_end_index: 0,
prev_run_end_index: 0,
}
}
/// Creates a new `GenericByteRunBuilder` with the provided capacity
///
/// `capacity`: the expected number of run-end encoded values.
/// `data_capacity`: the expected number of bytes of run end encoded values
pub fn with_capacity(capacity: usize, data_capacity: usize) -> Self {
Self {
run_ends_builder: PrimitiveBuilder::with_capacity(capacity),
values_builder: GenericByteBuilder::<V>::with_capacity(capacity, data_capacity),
current_value: Vec::new(),
has_current_value: false,
current_run_end_index: 0,
prev_run_end_index: 0,
}
}
}
impl<R, V> ArrayBuilder for GenericByteRunBuilder<R, V>
where
R: RunEndIndexType,
V: ByteArrayType,
{
/// Returns the builder as a non-mutable `Any` reference.
fn as_any(&self) -> &dyn Any {
self
}
/// Returns the builder as a mutable `Any` reference.
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
/// Returns the boxed builder as a box of `Any`.
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
/// Returns the length of logical array encoded by
/// the eventual runs array.
fn len(&self) -> usize {
self.current_run_end_index
}
/// Builds the array and reset this builder.
fn finish(&mut self) -> ArrayRef {
Arc::new(self.finish())
}
/// Builds the array without resetting the builder.
fn finish_cloned(&self) -> ArrayRef {
Arc::new(self.finish_cloned())
}
}
impl<R, V> GenericByteRunBuilder<R, V>
where
R: RunEndIndexType,
V: ByteArrayType,
{
/// Appends optional value to the logical array encoded by the RunArray.
pub fn append_option(&mut self, input_value: Option<impl AsRef<V::Native>>) {
match input_value {
Some(value) => self.append_value(value),
None => self.append_null(),
}
}
/// Appends value to the logical array encoded by the RunArray.
pub fn append_value(&mut self, input_value: impl AsRef<V::Native>) {
let value: &[u8] = input_value.as_ref().as_ref();
if !self.has_current_value {
self.append_run_end();
self.current_value.extend_from_slice(value);
self.has_current_value = true;
} else if self.current_value.as_slice() != value {
self.append_run_end();
self.current_value.clear();
self.current_value.extend_from_slice(value);
}
self.current_run_end_index += 1;
}
/// Appends null to the logical array encoded by the RunArray.
pub fn append_null(&mut self) {
if self.has_current_value {
self.append_run_end();
self.current_value.clear();
self.has_current_value = false;
}
self.current_run_end_index += 1;
}
/// Creates the RunArray and resets the builder.
/// Panics if RunArray cannot be built.
pub fn finish(&mut self) -> RunArray<R> {
// write the last run end to the array.
self.append_run_end();
// reset the run end index to zero.
self.current_value.clear();
self.has_current_value = false;
self.current_run_end_index = 0;
self.prev_run_end_index = 0;
// build the run encoded array by adding run_ends and values array as its children.
let run_ends_array = self.run_ends_builder.finish();
let values_array = self.values_builder.finish();
RunArray::<R>::try_new(&run_ends_array, &values_array).unwrap()
}
/// Creates the RunArray and without resetting the builder.
/// Panics if RunArray cannot be built.
pub fn finish_cloned(&self) -> RunArray<R> {
let mut run_ends_array = self.run_ends_builder.finish_cloned();
let mut values_array = self.values_builder.finish_cloned();
// Add current run if one exists
if self.prev_run_end_index != self.current_run_end_index {
let mut run_end_builder = run_ends_array.into_builder().unwrap();
let mut values_builder = values_array.into_builder().unwrap();
self.append_run_end_with_builders(&mut run_end_builder, &mut values_builder);
run_ends_array = run_end_builder.finish();
values_array = values_builder.finish();
}
RunArray::<R>::try_new(&run_ends_array, &values_array).unwrap()
}
// Appends the current run to the array.
fn append_run_end(&mut self) {
// empty array or the function called without appending any value.
if self.prev_run_end_index == self.current_run_end_index {
return;
}
let run_end_index = self.run_end_index_as_native();
self.run_ends_builder.append_value(run_end_index);
if self.has_current_value {
let slice = self.current_value.as_slice();
let native = unsafe {
// Safety:
// As self.current_value is created from V::Native. The value V::Native can be
// built back from the bytes without validations
V::Native::from_bytes_unchecked(slice)
};
self.values_builder.append_value(native);
} else {
self.values_builder.append_null();
}
self.prev_run_end_index = self.current_run_end_index;
}
// Similar to `append_run_end` but on custom builders.
// Used in `finish_cloned` which is not suppose to mutate `self`.
fn append_run_end_with_builders(
&self,
run_ends_builder: &mut PrimitiveBuilder<R>,
values_builder: &mut GenericByteBuilder<V>,
) {
let run_end_index = self.run_end_index_as_native();
run_ends_builder.append_value(run_end_index);
if self.has_current_value {
let slice = self.current_value.as_slice();
let native = unsafe {
// Safety:
// As self.current_value is created from V::Native. The value V::Native can be
// built back from the bytes without validations
V::Native::from_bytes_unchecked(slice)
};
values_builder.append_value(native);
} else {
values_builder.append_null();
}
}
fn run_end_index_as_native(&self) -> R::Native {
R::Native::from_usize(self.current_run_end_index).unwrap_or_else(|| {
panic!(
"Cannot convert the value {} from `usize` to native form of arrow datatype {}",
self.current_run_end_index,
R::DATA_TYPE
)
})
}
}
impl<R, V, S> Extend<Option<S>> for GenericByteRunBuilder<R, V>
where
R: RunEndIndexType,
V: ByteArrayType,
S: AsRef<V::Native>,
{
fn extend<T: IntoIterator<Item = Option<S>>>(&mut self, iter: T) {
for elem in iter {
self.append_option(elem);
}
}
}
/// Builder for [`RunArray`] of [`StringArray`](crate::array::StringArray)
///
/// ```
/// // Create a run-end encoded array with run-end indexes data type as `i16`.
/// // The encoded values are Strings.
///
/// # use arrow_array::builder::StringRunBuilder;
/// # use arrow_array::{Int16Array, StringArray};
/// # use arrow_array::types::Int16Type;
/// # use arrow_array::cast::AsArray;
/// #
/// let mut builder = StringRunBuilder::<Int16Type>::new();
///
/// // The builder builds the dictionary value by value
/// builder.append_value("abc");
/// builder.append_null();
/// builder.extend([Some("def"), Some("def"), Some("abc")]);
/// let array = builder.finish();
///
/// assert_eq!(array.run_ends().values(), &[1, 2, 4, 5]);
///
/// // Values are polymorphic and so require a downcast.
/// let av = array.values();
/// let ava: &StringArray = av.as_string::<i32>();
///
/// assert_eq!(ava.value(0), "abc");
/// assert!(av.is_null(1));
/// assert_eq!(ava.value(2), "def");
/// assert_eq!(ava.value(3), "abc");
///
/// ```
pub type StringRunBuilder<K> = GenericByteRunBuilder<K, Utf8Type>;
/// Builder for [`RunArray`] of [`LargeStringArray`](crate::array::LargeStringArray)
pub type LargeStringRunBuilder<K> = GenericByteRunBuilder<K, LargeUtf8Type>;
/// Builder for [`RunArray`] of [`BinaryArray`](crate::array::BinaryArray)
///
/// ```
/// // Create a run-end encoded array with run-end indexes data type as `i16`.
/// // The encoded data is binary values.
///
/// # use arrow_array::builder::BinaryRunBuilder;
/// # use arrow_array::{BinaryArray, Int16Array};
/// # use arrow_array::cast::AsArray;
/// # use arrow_array::types::Int16Type;
///
/// let mut builder = BinaryRunBuilder::<Int16Type>::new();
///
/// // The builder builds the dictionary value by value
/// builder.append_value(b"abc");
/// builder.append_null();
/// builder.extend([Some(b"def"), Some(b"def"), Some(b"abc")]);
/// let array = builder.finish();
///
/// assert_eq!(array.run_ends().values(), &[1, 2, 4, 5]);
///
/// // Values are polymorphic and so require a downcast.
/// let av = array.values();
/// let ava: &BinaryArray = av.as_binary();
///
/// assert_eq!(ava.value(0), b"abc");
/// assert!(av.is_null(1));
/// assert_eq!(ava.value(2), b"def");
/// assert_eq!(ava.value(3), b"abc");
///
/// ```
pub type BinaryRunBuilder<K> = GenericByteRunBuilder<K, BinaryType>;
/// Builder for [`RunArray`] of [`LargeBinaryArray`](crate::array::LargeBinaryArray)
pub type LargeBinaryRunBuilder<K> = GenericByteRunBuilder<K, LargeBinaryType>;
#[cfg(test)]
mod tests {
use super::*;
use crate::GenericByteArray;
use crate::Int16RunArray;
use crate::array::Array;
use crate::cast::AsArray;
use crate::types::{Int16Type, Int32Type};
fn test_bytes_run_builder<T>(values: Vec<&T::Native>)
where
T: ByteArrayType,
<T as ByteArrayType>::Native: PartialEq,
<T as ByteArrayType>::Native: AsRef<<T as ByteArrayType>::Native>,
{
let mut builder = GenericByteRunBuilder::<Int16Type, T>::new();
builder.append_value(values[0]);
builder.append_value(values[0]);
builder.append_value(values[0]);
builder.append_null();
builder.append_null();
builder.append_value(values[1]);
builder.append_value(values[1]);
builder.append_value(values[2]);
builder.append_value(values[2]);
builder.append_value(values[2]);
builder.append_value(values[2]);
let array = builder.finish();
assert_eq!(array.len(), 11);
assert_eq!(array.null_count(), 0);
assert_eq!(array.logical_null_count(), 2);
assert_eq!(array.run_ends().values(), &[3, 5, 7, 11]);
// Values are polymorphic and so require a downcast.
let av = array.values();
let ava: &GenericByteArray<T> = av.as_any().downcast_ref::<GenericByteArray<T>>().unwrap();
assert_eq!(*ava.value(0), *values[0]);
assert!(ava.is_null(1));
assert_eq!(*ava.value(2), *values[1]);
assert_eq!(*ava.value(3), *values[2]);
}
#[test]
fn test_string_run_builder() {
test_bytes_run_builder::<Utf8Type>(vec!["abc", "def", "ghi"]);
}
#[test]
fn test_string_run_builder_with_empty_strings() {
test_bytes_run_builder::<Utf8Type>(vec!["abc", "", "ghi"]);
}
#[test]
fn test_binary_run_builder() {
test_bytes_run_builder::<BinaryType>(vec![b"abc", b"def", b"ghi"]);
}
fn test_bytes_run_builder_finish_cloned<T>(values: Vec<&T::Native>)
where
T: ByteArrayType,
<T as ByteArrayType>::Native: PartialEq,
<T as ByteArrayType>::Native: AsRef<<T as ByteArrayType>::Native>,
{
let mut builder = GenericByteRunBuilder::<Int16Type, T>::new();
builder.append_value(values[0]);
builder.append_null();
builder.append_value(values[1]);
builder.append_value(values[1]);
builder.append_value(values[0]);
let mut array: Int16RunArray = builder.finish_cloned();
assert_eq!(array.len(), 5);
assert_eq!(array.null_count(), 0);
assert_eq!(array.logical_null_count(), 1);
assert_eq!(array.run_ends().values(), &[1, 2, 4, 5]);
// Values are polymorphic and so require a downcast.
let av = array.values();
let ava: &GenericByteArray<T> = av.as_any().downcast_ref::<GenericByteArray<T>>().unwrap();
assert_eq!(ava.value(0), values[0]);
assert!(ava.is_null(1));
assert_eq!(ava.value(2), values[1]);
assert_eq!(ava.value(3), values[0]);
// Append last value before `finish_cloned` (`value[0]`) again and ensure it has only
// one entry in final output.
builder.append_value(values[0]);
builder.append_value(values[0]);
builder.append_value(values[1]);
array = builder.finish();
assert_eq!(array.len(), 8);
assert_eq!(array.null_count(), 0);
assert_eq!(array.logical_null_count(), 1);
assert_eq!(array.run_ends().values(), &[1, 2, 4, 7, 8]);
// Values are polymorphic and so require a downcast.
let av2 = array.values();
let ava2: &GenericByteArray<T> =
av2.as_any().downcast_ref::<GenericByteArray<T>>().unwrap();
assert_eq!(ava2.value(0), values[0]);
assert!(ava2.is_null(1));
assert_eq!(ava2.value(2), values[1]);
// The value appended before and after `finish_cloned` has only one entry.
assert_eq!(ava2.value(3), values[0]);
assert_eq!(ava2.value(4), values[1]);
}
#[test]
fn test_string_run_builder_finish_cloned() {
test_bytes_run_builder_finish_cloned::<Utf8Type>(vec!["abc", "def", "ghi"]);
}
#[test]
fn test_binary_run_builder_finish_cloned() {
test_bytes_run_builder_finish_cloned::<BinaryType>(vec![b"abc", b"def", b"ghi"]);
}
#[test]
fn test_extend() {
let mut builder = StringRunBuilder::<Int32Type>::new();
builder.extend(["a", "a", "a", "", "", "b", "b"].into_iter().map(Some));
builder.extend(["b", "cupcakes", "cupcakes"].into_iter().map(Some));
let array = builder.finish();
assert_eq!(array.len(), 10);
assert_eq!(array.run_ends().values(), &[3, 5, 8, 10]);
let str_array = array.values().as_string::<i32>();
assert_eq!(str_array.value(0), "a");
assert_eq!(str_array.value(1), "");
assert_eq!(str_array.value(2), "b");
assert_eq!(str_array.value(3), "cupcakes");
}
}
@@ -0,0 +1,942 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::builder::ArrayBuilder;
use crate::types::{ByteArrayType, GenericBinaryType, GenericStringType};
use crate::{Array, ArrayRef, GenericByteArray, OffsetSizeTrait};
use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer, NullBufferBuilder, ScalarBuffer};
use arrow_data::ArrayDataBuilder;
use arrow_schema::ArrowError;
use std::any::Any;
use std::sync::Arc;
/// Builder for [`GenericByteArray`]
///
/// For building strings, see docs on [`GenericStringBuilder`].
/// For building binary, see docs on [`GenericBinaryBuilder`].
pub struct GenericByteBuilder<T: ByteArrayType> {
value_builder: Vec<u8>,
offsets_builder: Vec<T::Offset>,
null_buffer_builder: NullBufferBuilder,
}
impl<T: ByteArrayType> GenericByteBuilder<T> {
/// Creates a new [`GenericByteBuilder`].
pub fn new() -> Self {
Self::with_capacity(1024, 1024)
}
/// Creates a new [`GenericByteBuilder`].
///
/// - `item_capacity` is the number of items to pre-allocate.
/// The size of the preallocated buffer of offsets is the number of items plus one.
/// - `data_capacity` is the total number of bytes of data to pre-allocate
/// (for all items, not per item).
pub fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
let mut offsets_builder = Vec::with_capacity(item_capacity + 1);
offsets_builder.push(T::Offset::from_usize(0).unwrap());
Self {
value_builder: Vec::with_capacity(data_capacity),
offsets_builder,
null_buffer_builder: NullBufferBuilder::new(item_capacity),
}
}
/// Creates a new [`GenericByteBuilder`] from buffers.
///
/// # Safety
///
/// This doesn't verify buffer contents as it assumes the buffers are from
/// existing and valid [`GenericByteArray`].
pub unsafe fn new_from_buffer(
offsets_buffer: MutableBuffer,
value_buffer: MutableBuffer,
null_buffer: Option<MutableBuffer>,
) -> Self {
let offsets_builder: Vec<T::Offset> =
ScalarBuffer::<T::Offset>::from(offsets_buffer).into();
let value_builder: Vec<u8> = ScalarBuffer::<u8>::from(value_buffer).into();
let null_buffer_builder = null_buffer
.map(|buffer| NullBufferBuilder::new_from_buffer(buffer, offsets_builder.len() - 1))
.unwrap_or_else(|| NullBufferBuilder::new_with_len(offsets_builder.len() - 1));
Self {
offsets_builder,
value_builder,
null_buffer_builder,
}
}
#[inline]
fn next_offset(&self) -> T::Offset {
T::Offset::from_usize(self.value_builder.len()).expect("byte array offset overflow")
}
/// Appends a value into the builder.
///
/// See the [GenericStringBuilder] documentation for examples of
/// incrementally building string values with multiple `write!` calls.
///
/// # Panics
///
/// Panics if the resulting length of [`Self::values_slice`] would exceed
/// `T::Offset::MAX` bytes.
///
/// For example, this can happen with [`StringArray`] or [`BinaryArray`]
/// where the total length of all values exceeds 2GB
///
/// [`StringArray`]: crate::StringArray
/// [`BinaryArray`]: crate::BinaryArray
#[inline]
pub fn append_value(&mut self, value: impl AsRef<T::Native>) {
self.value_builder
.extend_from_slice(value.as_ref().as_ref());
self.null_buffer_builder.append(true);
self.offsets_builder.push(self.next_offset());
}
/// Append an `Option` value into the builder.
///
/// - A `None` value will append a null value.
/// - A `Some` value will append the value.
///
/// See [`Self::append_value`] for more panic information.
#[inline]
pub fn append_option(&mut self, value: Option<impl AsRef<T::Native>>) {
match value {
None => self.append_null(),
Some(v) => self.append_value(v),
};
}
/// Append a null value into the builder.
#[inline]
pub fn append_null(&mut self) {
self.null_buffer_builder.append(false);
self.offsets_builder.push(self.next_offset());
}
/// Appends `n` `null`s into the builder.
#[inline]
pub fn append_nulls(&mut self, n: usize) {
self.null_buffer_builder.append_n_nulls(n);
let next_offset = self.next_offset();
self.offsets_builder
.extend(std::iter::repeat_n(next_offset, n));
}
/// Appends array values and null to this builder as is
/// (this means that underlying null values are copied as is).
#[inline]
pub fn append_array(&mut self, array: &GenericByteArray<T>) -> Result<(), ArrowError> {
use num_traits::CheckedAdd;
if array.len() == 0 {
return Ok(());
}
let offsets = array.offsets();
// If the offsets are contiguous, we can append them directly avoiding the need to align
// for example, when the first appended array is not sliced (starts at offset 0)
if self.next_offset() == offsets[0] {
self.offsets_builder.extend_from_slice(&offsets[1..]);
} else {
// Shifting all the offsets
let shift: T::Offset = self.next_offset() - offsets[0];
if shift.checked_add(&offsets[offsets.len() - 1]).is_none() {
return Err(ArrowError::OffsetOverflowError(
shift.as_usize() + offsets[offsets.len() - 1].as_usize(),
));
}
self.offsets_builder
.extend(offsets[1..].iter().map(|&offset| offset + shift));
}
// Append underlying values, starting from the first offset and ending at the last offset
self.value_builder.extend_from_slice(
&array.values().as_slice()[offsets[0].as_usize()..offsets[array.len()].as_usize()],
);
if let Some(null_buffer) = array.nulls() {
self.null_buffer_builder.append_buffer(null_buffer);
} else {
self.null_buffer_builder.append_n_non_nulls(array.len());
}
Ok(())
}
/// Builds the [`GenericByteArray`] and reset this builder.
pub fn finish(&mut self) -> GenericByteArray<T> {
let array_type = T::DATA_TYPE;
let array_builder = ArrayDataBuilder::new(array_type)
.len(self.len())
.add_buffer(std::mem::take(&mut self.offsets_builder).into())
.add_buffer(std::mem::take(&mut self.value_builder).into())
.nulls(self.null_buffer_builder.finish());
self.offsets_builder.push(self.next_offset());
let array_data = unsafe { array_builder.build_unchecked() };
GenericByteArray::from(array_data)
}
/// Builds the [`GenericByteArray`] without resetting the builder.
pub fn finish_cloned(&self) -> GenericByteArray<T> {
let array_type = T::DATA_TYPE;
let offset_buffer = Buffer::from_slice_ref(self.offsets_builder.as_slice());
let value_buffer = Buffer::from_slice_ref(self.value_builder.as_slice());
let array_builder = ArrayDataBuilder::new(array_type)
.len(self.len())
.add_buffer(offset_buffer)
.add_buffer(value_buffer)
.nulls(self.null_buffer_builder.finish_cloned());
let array_data = unsafe { array_builder.build_unchecked() };
GenericByteArray::from(array_data)
}
/// Returns the current values buffer as a slice
pub fn values_slice(&self) -> &[u8] {
self.value_builder.as_slice()
}
/// Returns the current offsets buffer as a slice
pub fn offsets_slice(&self) -> &[T::Offset] {
self.offsets_builder.as_slice()
}
/// Returns the current null buffer as a slice
pub fn validity_slice(&self) -> Option<&[u8]> {
self.null_buffer_builder.as_slice()
}
/// Returns the current null buffer as a mutable slice
pub fn validity_slice_mut(&mut self) -> Option<&mut [u8]> {
self.null_buffer_builder.as_slice_mut()
}
}
impl<T: ByteArrayType> std::fmt::Debug for GenericByteBuilder<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}{}Builder", T::Offset::PREFIX, T::PREFIX)?;
f.debug_struct("")
.field("value_builder", &self.value_builder)
.field("offsets_builder", &self.offsets_builder)
.field("null_buffer_builder", &self.null_buffer_builder)
.finish()
}
}
impl<T: ByteArrayType> Default for GenericByteBuilder<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: ByteArrayType> ArrayBuilder for GenericByteBuilder<T> {
/// Returns the number of binary slots in the builder
fn len(&self) -> usize {
self.null_buffer_builder.len()
}
/// Builds the array and reset this builder.
fn finish(&mut self) -> ArrayRef {
Arc::new(self.finish())
}
/// Builds the array without resetting the builder.
fn finish_cloned(&self) -> ArrayRef {
Arc::new(self.finish_cloned())
}
/// Returns the builder as a non-mutable `Any` reference.
fn as_any(&self) -> &dyn Any {
self
}
/// Returns the builder as a mutable `Any` reference.
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
/// Returns the boxed builder as a box of `Any`.
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
}
impl<T: ByteArrayType, V: AsRef<T::Native>> Extend<Option<V>> for GenericByteBuilder<T> {
#[inline]
fn extend<I: IntoIterator<Item = Option<V>>>(&mut self, iter: I) {
for v in iter {
self.append_option(v)
}
}
}
/// Array builder for [`GenericStringArray`][crate::GenericStringArray]
///
/// Values can be appended using [`GenericByteBuilder::append_value`], and nulls with
/// [`GenericByteBuilder::append_null`].
///
/// This builder also implements [`std::fmt::Write`] with any written data
/// included in the next appended value. This allows using [`std::fmt::Display`]
/// with standard Rust idioms like `write!` and `writeln!` to write data
/// directly to the builder without intermediate allocations.
///
/// # Example writing strings with `append_value`
/// ```
/// # use arrow_array::builder::GenericStringBuilder;
/// let mut builder = GenericStringBuilder::<i32>::new();
///
/// // Write one string value
/// builder.append_value("foobarbaz");
///
/// // Write a second string
/// builder.append_value("v2");
///
/// let array = builder.finish();
/// assert_eq!(array.value(0), "foobarbaz");
/// assert_eq!(array.value(1), "v2");
/// ```
///
/// # Example incrementally writing strings with `std::fmt::Write`
///
/// ```
/// # use std::fmt::Write;
/// # use arrow_array::builder::GenericStringBuilder;
/// let mut builder = GenericStringBuilder::<i32>::new();
///
/// // Write data in multiple `write!` calls
/// write!(builder, "foo").unwrap();
/// write!(builder, "bar").unwrap();
/// // The next call to append_value finishes the current string
/// // including all previously written strings.
/// builder.append_value("baz");
///
/// // Write second value with a single write call
/// write!(builder, "v2").unwrap();
/// // finish the value by calling append_value with an empty string
/// builder.append_value("");
///
/// let array = builder.finish();
/// assert_eq!(array.value(0), "foobarbaz");
/// assert_eq!(array.value(1), "v2");
/// ```
pub type GenericStringBuilder<O> = GenericByteBuilder<GenericStringType<O>>;
impl<O: OffsetSizeTrait> std::fmt::Write for GenericStringBuilder<O> {
fn write_str(&mut self, s: &str) -> std::fmt::Result {
self.value_builder.extend_from_slice(s.as_bytes());
Ok(())
}
}
/// A byte size value representing the number of bytes to allocate per string in [`GenericStringBuilder`]
///
/// To create a [`GenericStringBuilder`] using `.with_capacity` we are required to provide: \
/// - `item_capacity` - the row count \
/// - `data_capacity` - total string byte count \
///
/// We will use the `AVERAGE_STRING_LENGTH` * row_count for `data_capacity`. \
///
/// These capacities are preallocation hints used to improve performance,
/// but consequences of passing a hint too large or too small should be negligible.
const AVERAGE_STRING_LENGTH: usize = 16;
/// Trait for string-like array builders
///
/// This trait provides unified interface for builders that append string-like data
/// such as [`GenericStringBuilder<O>`] and [`crate::builder::StringViewBuilder`]
pub trait StringLikeArrayBuilder: ArrayBuilder {
/// Returns a human-readable type name for the builder.
fn type_name() -> &'static str;
/// Creates a new builder with the given row capacity.
fn with_capacity(capacity: usize) -> Self;
/// Appends a non-null string value to the builder.
fn append_value(&mut self, value: &str);
/// Appends a null value to the builder.
fn append_null(&mut self);
}
impl<O: OffsetSizeTrait> StringLikeArrayBuilder for GenericStringBuilder<O> {
fn type_name() -> &'static str {
std::any::type_name::<Self>()
}
fn with_capacity(capacity: usize) -> Self {
Self::with_capacity(capacity, capacity * AVERAGE_STRING_LENGTH)
}
fn append_value(&mut self, value: &str) {
Self::append_value(self, value);
}
fn append_null(&mut self) {
Self::append_null(self);
}
}
/// A byte size value representing the number of bytes to allocate per binary in [`GenericBinaryBuilder`]
///
/// To create a [`GenericBinaryBuilder`] using `.with_capacity` we are required to provide: \
/// - `item_capacity` - the row count \
/// - `data_capacity` - total binary byte count \
///
/// We will use the `AVERAGE_BINARY_LENGTH` * row_count for `data_capacity`. \
///
/// These capacities are preallocation hints used to improve performance,
/// but consequences of passing a hint too large or too small should be negligible.
const AVERAGE_BINARY_LENGTH: usize = 128;
/// Trait for binary-like array builders
///
/// This trait provides unified interface for builders that append binary-like data
/// such as [`GenericBinaryBuilder<O>`] and [`crate::builder::BinaryViewBuilder`]
pub trait BinaryLikeArrayBuilder: ArrayBuilder {
/// Returns a human-readable type name for the builder.
fn type_name() -> &'static str;
/// Creates a new builder with the given row capacity.
fn with_capacity(capacity: usize) -> Self;
/// Appends a non-null string value to the builder.
fn append_value(&mut self, value: &[u8]);
/// Appends a null value to the builder.
fn append_null(&mut self);
}
impl<O: OffsetSizeTrait> BinaryLikeArrayBuilder for GenericBinaryBuilder<O> {
fn type_name() -> &'static str {
std::any::type_name::<Self>()
}
fn with_capacity(capacity: usize) -> Self {
Self::with_capacity(capacity, capacity * AVERAGE_BINARY_LENGTH)
}
fn append_value(&mut self, value: &[u8]) {
Self::append_value(self, value);
}
fn append_null(&mut self) {
Self::append_null(self);
}
}
/// Array builder for [`GenericBinaryArray`][crate::GenericBinaryArray]
///
/// Values can be appended using [`GenericByteBuilder::append_value`], and nulls with
/// [`GenericByteBuilder::append_null`].
///
/// # Example
/// ```
/// # use arrow_array::builder::GenericBinaryBuilder;
/// let mut builder = GenericBinaryBuilder::<i32>::new();
///
/// // Write data
/// builder.append_value("foo");
///
/// // Write second value
/// builder.append_value(&[0,1,2]);
///
/// let array = builder.finish();
/// // binary values
/// assert_eq!(array.value(0), b"foo");
/// assert_eq!(array.value(1), b"\x00\x01\x02");
/// ```
///
/// # Example incrementally writing bytes with `write_bytes`
///
/// ```
/// # use std::io::Write;
/// # use arrow_array::builder::GenericBinaryBuilder;
/// let mut builder = GenericBinaryBuilder::<i32>::new();
///
/// // Write data in multiple `write_bytes` calls
/// write!(builder, "foo").unwrap();
/// write!(builder, "bar").unwrap();
/// // The next call to append_value finishes the current string
/// // including all previously written strings.
/// builder.append_value("baz");
///
/// // Write second value with a single write call
/// write!(builder, "v2").unwrap();
/// // finish the value by calling append_value with an empty string
/// builder.append_value("");
///
/// let array = builder.finish();
/// assert_eq!(array.value(0), "foobarbaz".as_bytes());
/// assert_eq!(array.value(1), "v2".as_bytes());
/// ```
pub type GenericBinaryBuilder<O> = GenericByteBuilder<GenericBinaryType<O>>;
impl<O: OffsetSizeTrait> std::io::Write for GenericBinaryBuilder<O> {
fn write(&mut self, bs: &[u8]) -> std::io::Result<usize> {
self.value_builder.extend_from_slice(bs);
Ok(bs.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::GenericStringArray;
use crate::array::Array;
use arrow_buffer::NullBuffer;
use std::fmt::Write as _;
use std::io::Write as _;
fn _test_generic_binary_builder<O: OffsetSizeTrait>() {
let mut builder = GenericBinaryBuilder::<O>::new();
builder.append_value(b"hello");
builder.append_value(b"");
builder.append_null();
builder.append_value(b"rust");
let array = builder.finish();
assert_eq!(4, array.len());
assert_eq!(1, array.null_count());
assert_eq!(b"hello", array.value(0));
assert_eq!([] as [u8; 0], array.value(1));
assert!(array.is_null(2));
assert_eq!(b"rust", array.value(3));
assert_eq!(O::from_usize(5).unwrap(), array.value_offsets()[2]);
assert_eq!(O::from_usize(4).unwrap(), array.value_length(3));
}
#[test]
fn test_binary_builder() {
_test_generic_binary_builder::<i32>()
}
#[test]
fn test_large_binary_builder() {
_test_generic_binary_builder::<i64>()
}
fn _test_generic_binary_builder_all_nulls<O: OffsetSizeTrait>() {
let mut builder = GenericBinaryBuilder::<O>::new();
builder.append_null();
builder.append_null();
builder.append_null();
builder.append_nulls(2);
assert_eq!(5, builder.len());
assert!(!builder.is_empty());
let array = builder.finish();
assert_eq!(5, array.null_count());
assert_eq!(5, array.len());
assert!(array.is_null(0));
assert!(array.is_null(1));
assert!(array.is_null(2));
assert!(array.is_null(3));
assert!(array.is_null(4));
}
#[test]
fn test_binary_builder_all_nulls() {
_test_generic_binary_builder_all_nulls::<i32>()
}
#[test]
fn test_large_binary_builder_all_nulls() {
_test_generic_binary_builder_all_nulls::<i64>()
}
fn _test_generic_binary_builder_reset<O: OffsetSizeTrait>() {
let mut builder = GenericBinaryBuilder::<O>::new();
builder.append_value(b"hello");
builder.append_value(b"");
builder.append_null();
builder.append_value(b"rust");
builder.finish();
assert!(builder.is_empty());
builder.append_value(b"parquet");
builder.append_null();
builder.append_value(b"arrow");
builder.append_value(b"");
builder.append_nulls(2);
builder.append_value(b"hi");
let array = builder.finish();
assert_eq!(7, array.len());
assert_eq!(3, array.null_count());
assert_eq!(b"parquet", array.value(0));
assert!(array.is_null(1));
assert!(array.is_null(4));
assert!(array.is_null(5));
assert_eq!(b"arrow", array.value(2));
assert_eq!(b"", array.value(1));
assert_eq!(b"hi", array.value(6));
assert_eq!(O::zero(), array.value_offsets()[0]);
assert_eq!(O::from_usize(7).unwrap(), array.value_offsets()[2]);
assert_eq!(O::from_usize(14).unwrap(), array.value_offsets()[7]);
assert_eq!(O::from_usize(5).unwrap(), array.value_length(2));
}
#[test]
fn test_binary_builder_reset() {
_test_generic_binary_builder_reset::<i32>()
}
#[test]
fn test_large_binary_builder_reset() {
_test_generic_binary_builder_reset::<i64>()
}
fn _test_generic_string_array_builder<O: OffsetSizeTrait>() {
let mut builder = GenericStringBuilder::<O>::new();
let owned = "arrow".to_owned();
builder.append_value("hello");
builder.append_value("");
builder.append_value(&owned);
builder.append_null();
builder.append_option(Some("rust"));
builder.append_option(None::<&str>);
builder.append_option(None::<String>);
builder.append_nulls(2);
builder.append_value("parquet");
assert_eq!(10, builder.len());
assert_eq!(
GenericStringArray::<O>::from(vec![
Some("hello"),
Some(""),
Some("arrow"),
None,
Some("rust"),
None,
None,
None,
None,
Some("parquet")
]),
builder.finish()
);
}
#[test]
fn test_string_array_builder() {
_test_generic_string_array_builder::<i32>()
}
#[test]
fn test_large_string_array_builder() {
_test_generic_string_array_builder::<i64>()
}
fn _test_generic_string_array_builder_finish<O: OffsetSizeTrait>() {
let mut builder = GenericStringBuilder::<O>::with_capacity(3, 11);
builder.append_value("hello");
builder.append_value("rust");
builder.append_null();
builder.finish();
assert!(builder.is_empty());
assert_eq!(&[O::zero()], builder.offsets_slice());
builder.append_value("arrow");
builder.append_value("parquet");
let arr = builder.finish();
// array should not have null buffer because there is not `null` value.
assert!(arr.nulls().is_none());
assert_eq!(GenericStringArray::<O>::from(vec!["arrow", "parquet"]), arr,)
}
#[test]
fn test_string_array_builder_finish() {
_test_generic_string_array_builder_finish::<i32>()
}
#[test]
fn test_large_string_array_builder_finish() {
_test_generic_string_array_builder_finish::<i64>()
}
fn _test_generic_string_array_builder_finish_cloned<O: OffsetSizeTrait>() {
let mut builder = GenericStringBuilder::<O>::with_capacity(3, 11);
builder.append_value("hello");
builder.append_value("rust");
builder.append_null();
let mut arr = builder.finish_cloned();
assert!(!builder.is_empty());
assert_eq!(3, arr.len());
builder.append_value("arrow");
builder.append_value("parquet");
arr = builder.finish();
assert!(arr.nulls().is_some());
assert_eq!(&[O::zero()], builder.offsets_slice());
assert_eq!(5, arr.len());
}
#[test]
fn test_string_array_builder_finish_cloned() {
_test_generic_string_array_builder_finish_cloned::<i32>()
}
#[test]
fn test_large_string_array_builder_finish_cloned() {
_test_generic_string_array_builder_finish_cloned::<i64>()
}
#[test]
fn test_extend() {
let mut builder = GenericStringBuilder::<i32>::new();
builder.extend(["a", "b", "c", "", "a", "b", "c"].into_iter().map(Some));
builder.extend(["d", "cupcakes", "hello"].into_iter().map(Some));
let array = builder.finish();
assert_eq!(array.value_offsets(), &[0, 1, 2, 3, 3, 4, 5, 6, 7, 15, 20]);
assert_eq!(array.value_data(), b"abcabcdcupcakeshello");
}
#[test]
fn test_write_str() {
let mut builder = GenericStringBuilder::<i32>::new();
write!(builder, "foo").unwrap();
builder.append_value("");
writeln!(builder, "bar").unwrap();
builder.append_value("");
write!(builder, "fiz").unwrap();
write!(builder, "buz").unwrap();
builder.append_value("");
let a = builder.finish();
let r: Vec<_> = a.iter().flatten().collect();
assert_eq!(r, &["foo", "bar\n", "fizbuz"])
}
#[test]
fn test_write_bytes() {
let mut builder = GenericBinaryBuilder::<i32>::new();
write!(builder, "foo").unwrap();
builder.append_value("");
writeln!(builder, "bar").unwrap();
builder.append_value("");
write!(builder, "fiz").unwrap();
write!(builder, "buz").unwrap();
builder.append_value("");
let a = builder.finish();
let r: Vec<_> = a.iter().flatten().collect();
assert_eq!(
r,
&["foo".as_bytes(), "bar\n".as_bytes(), "fizbuz".as_bytes()]
)
}
#[test]
fn test_append_array_without_nulls() {
let input = vec![
"hello", "world", "how", "are", "you", "doing", "today", "I", "am", "doing", "well",
"thank", "you", "for", "asking",
];
let arr1 = GenericStringArray::<i32>::from(input[..3].to_vec());
let arr2 = GenericStringArray::<i32>::from(input[3..7].to_vec());
let arr3 = GenericStringArray::<i32>::from(input[7..].to_vec());
let mut builder = GenericStringBuilder::<i32>::new();
builder.append_array(&arr1).unwrap();
builder.append_array(&arr2).unwrap();
builder.append_array(&arr3).unwrap();
let actual = builder.finish();
let expected = GenericStringArray::<i32>::from(input);
assert_eq!(actual, expected);
}
#[test]
fn test_append_array_with_nulls() {
let input = vec![
Some("hello"),
None,
Some("how"),
None,
None,
None,
None,
Some("I"),
Some("am"),
Some("doing"),
Some("well"),
];
let arr1 = GenericStringArray::<i32>::from(input[..3].to_vec());
let arr2 = GenericStringArray::<i32>::from(input[3..7].to_vec());
let arr3 = GenericStringArray::<i32>::from(input[7..].to_vec());
let mut builder = GenericStringBuilder::<i32>::new();
builder.append_array(&arr1).unwrap();
builder.append_array(&arr2).unwrap();
builder.append_array(&arr3).unwrap();
let actual = builder.finish();
let expected = GenericStringArray::<i32>::from(input);
assert_eq!(actual, expected);
}
#[test]
fn test_append_empty_array() {
let arr = GenericStringArray::<i32>::from(Vec::<&str>::new());
let mut builder = GenericStringBuilder::<i32>::new();
builder.append_array(&arr).unwrap();
let result = builder.finish();
assert_eq!(result.len(), 0);
}
#[test]
fn test_append_array_with_offset_not_starting_at_0() {
let input = vec![
Some("hello"),
None,
Some("how"),
None,
None,
None,
None,
Some("I"),
Some("am"),
Some("doing"),
Some("well"),
];
let full_array = GenericStringArray::<i32>::from(input);
let sliced = full_array.slice(1, 4);
assert_ne!(sliced.offsets()[0].as_usize(), 0);
assert_ne!(sliced.offsets().last(), full_array.offsets().last());
let mut builder = GenericStringBuilder::<i32>::new();
builder.append_array(&sliced).unwrap();
let actual = builder.finish();
let expected = GenericStringArray::<i32>::from(vec![None, Some("how"), None, None]);
assert_eq!(actual, expected);
}
#[test]
fn test_append_underlying_null_values_added_as_is() {
let input_1_array_with_nulls = {
let input = vec![
"hello", "world", "how", "are", "you", "doing", "today", "I", "am",
];
let (offsets, buffer, _) = GenericStringArray::<i32>::from(input).into_parts();
GenericStringArray::<i32>::new(
offsets,
buffer,
Some(NullBuffer::from(&[
true, false, true, false, false, true, true, true, false,
])),
)
};
let input_2_array_with_nulls = {
let input = vec!["doing", "well", "thank", "you", "for", "asking"];
let (offsets, buffer, _) = GenericStringArray::<i32>::from(input).into_parts();
GenericStringArray::<i32>::new(
offsets,
buffer,
Some(NullBuffer::from(&[false, false, true, false, true, true])),
)
};
let mut builder = GenericStringBuilder::<i32>::new();
builder.append_array(&input_1_array_with_nulls).unwrap();
builder.append_array(&input_2_array_with_nulls).unwrap();
let actual = builder.finish();
let expected = GenericStringArray::<i32>::from(vec![
Some("hello"),
None, // world
Some("how"),
None, // are
None, // you
Some("doing"),
Some("today"),
Some("I"),
None, // am
None, // doing
None, // well
Some("thank"),
None, // "you",
Some("for"),
Some("asking"),
]);
assert_eq!(actual, expected);
let expected_underlying_buffer = Buffer::from(
[
"hello", "world", "how", "are", "you", "doing", "today", "I", "am", "doing",
"well", "thank", "you", "for", "asking",
]
.join("")
.as_bytes(),
);
assert_eq!(actual.values(), &expected_underlying_buffer);
}
#[test]
fn append_array_with_continues_indices() {
let input = vec![
"hello", "world", "how", "are", "you", "doing", "today", "I", "am", "doing", "well",
"thank", "you", "for", "asking",
];
let full_array = GenericStringArray::<i32>::from(input);
let slice1 = full_array.slice(0, 3);
let slice2 = full_array.slice(3, 4);
let slice3 = full_array.slice(7, full_array.len() - 7);
let mut builder = GenericStringBuilder::<i32>::new();
builder.append_array(&slice1).unwrap();
builder.append_array(&slice2).unwrap();
builder.append_array(&slice3).unwrap();
let actual = builder.finish();
assert_eq!(actual, full_array);
}
#[test]
fn test_append_array_offset_overflow_precise() {
let mut builder = GenericStringBuilder::<i32>::new();
let initial_string = "x".repeat(i32::MAX as usize - 100);
builder.append_value(&initial_string);
let overflow_string = "y".repeat(200);
let overflow_array = GenericStringArray::<i32>::from(vec![overflow_string.as_str()]);
let result = builder.append_array(&overflow_array);
assert!(matches!(result, Err(ArrowError::OffsetOverflowError(_))));
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,821 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::builder::ArrayBuilder;
use crate::{Array, ArrayRef, GenericListArray, OffsetSizeTrait};
use arrow_buffer::NullBufferBuilder;
use arrow_buffer::{Buffer, OffsetBuffer};
use arrow_schema::{Field, FieldRef};
use std::any::Any;
use std::sync::Arc;
/// Builder for [`GenericListArray`]
///
/// Use [`ListBuilder`] to build [`ListArray`]s and [`LargeListBuilder`] to build [`LargeListArray`]s.
///
/// # Example
///
/// Here is code that constructs a ListArray with the contents:
/// `[[A,B,C], [], NULL, [D], [NULL, F]]`
///
/// ```
/// # use std::sync::Arc;
/// # use arrow_array::{builder::ListBuilder, builder::StringBuilder, ArrayRef, StringArray, Array};
/// #
/// let values_builder = StringBuilder::new();
/// let mut builder = ListBuilder::new(values_builder);
///
/// // [A, B, C]
/// builder.values().append_value("A");
/// builder.values().append_value("B");
/// builder.values().append_value("C");
/// builder.append(true);
///
/// // [ ] (empty list)
/// builder.append(true);
///
/// // Null
/// builder.append(false);
///
/// // [D]
/// builder.values().append_value("D");
/// builder.append(true);
///
/// // [NULL, F]
/// builder.values().append_null();
/// builder.values().append_value("F");
/// builder.append(true);
///
/// // Build the array
/// let array = builder.finish();
///
/// // Values is a string array
/// // "A", "B" "C", "?", "D", NULL, "F"
/// assert_eq!(
/// array.values().as_ref(),
/// &StringArray::from(vec![
/// Some("A"), Some("B"), Some("C"),
/// Some("D"), None, Some("F")
/// ])
/// );
///
/// // Offsets are indexes into the values array
/// assert_eq!(
/// array.value_offsets(),
/// &[0, 3, 3, 3, 4, 6]
/// );
/// ```
///
/// [`ListBuilder`]: crate::builder::ListBuilder
/// [`ListArray`]: crate::array::ListArray
/// [`LargeListBuilder`]: crate::builder::LargeListBuilder
/// [`LargeListArray`]: crate::array::LargeListArray
#[derive(Debug)]
pub struct GenericListBuilder<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> {
offsets_builder: Vec<OffsetSize>,
null_buffer_builder: NullBufferBuilder,
values_builder: T,
field: Option<FieldRef>,
}
impl<O: OffsetSizeTrait, T: ArrayBuilder + Default> Default for GenericListBuilder<O, T> {
fn default() -> Self {
Self::new(T::default())
}
}
impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> GenericListBuilder<OffsetSize, T> {
/// Creates a new [`GenericListBuilder`] from a given values array builder
pub fn new(values_builder: T) -> Self {
let capacity = values_builder.len();
Self::with_capacity(values_builder, capacity)
}
/// Creates a new [`GenericListBuilder`] from a given values array builder
/// `capacity` is the number of items to pre-allocate space for in this builder
pub fn with_capacity(values_builder: T, capacity: usize) -> Self {
let mut offsets_builder = Vec::with_capacity(capacity + 1);
offsets_builder.push(OffsetSize::zero());
Self {
offsets_builder,
null_buffer_builder: NullBufferBuilder::new(capacity),
values_builder,
field: None,
}
}
/// Override the field passed to [`GenericListArray::new`]
///
/// By default a nullable field is created with the name `item`
///
/// Note: [`Self::finish`] and [`Self::finish_cloned`] will panic if the
/// field's data type does not match that of `T`
pub fn with_field(self, field: impl Into<FieldRef>) -> Self {
Self {
field: Some(field.into()),
..self
}
}
}
impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> ArrayBuilder
for GenericListBuilder<OffsetSize, T>
where
T: 'static,
{
/// Returns the builder as a non-mutable `Any` reference.
fn as_any(&self) -> &dyn Any {
self
}
/// Returns the builder as a mutable `Any` reference.
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
/// Returns the boxed builder as a box of `Any`.
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
/// Returns the number of array slots in the builder
fn len(&self) -> usize {
self.null_buffer_builder.len()
}
/// Builds the array and reset this builder.
fn finish(&mut self) -> ArrayRef {
Arc::new(self.finish())
}
/// Builds the array without resetting the builder.
fn finish_cloned(&self) -> ArrayRef {
Arc::new(self.finish_cloned())
}
}
impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> GenericListBuilder<OffsetSize, T>
where
T: 'static,
{
/// Returns the child array builder as a mutable reference.
///
/// This mutable reference can be used to append values into the child array builder,
/// but you must call [`append`](#method.append) to delimit each distinct list value.
pub fn values(&mut self) -> &mut T {
&mut self.values_builder
}
/// Returns the child array builder as an immutable reference
pub fn values_ref(&self) -> &T {
&self.values_builder
}
/// Finish the current variable-length list array slot
///
/// # Panics
///
/// Panics if the length of [`Self::values`] exceeds `OffsetSize::MAX`
#[inline]
pub fn append(&mut self, is_valid: bool) {
self.offsets_builder.push(self.next_offset());
self.null_buffer_builder.append(is_valid);
}
/// Returns the next offset
///
/// # Panics
///
/// Panics if the length of [`Self::values`] exceeds `OffsetSize::MAX`
#[inline]
fn next_offset(&self) -> OffsetSize {
OffsetSize::from_usize(self.values_builder.len()).unwrap()
}
/// Append a value to this [`GenericListBuilder`]
///
/// ```
/// # use arrow_array::builder::{Int32Builder, ListBuilder};
/// # use arrow_array::cast::AsArray;
/// # use arrow_array::{Array, Int32Array};
/// # use arrow_array::types::Int32Type;
/// let mut builder = ListBuilder::new(Int32Builder::new());
///
/// builder.append_value([Some(1), Some(2), Some(3)]);
/// builder.append_value([]);
/// builder.append_value([None]);
///
/// let array = builder.finish();
/// assert_eq!(array.len(), 3);
///
/// assert_eq!(array.value_offsets(), &[0, 3, 3, 4]);
/// let values = array.values().as_primitive::<Int32Type>();
/// assert_eq!(values, &Int32Array::from(vec![Some(1), Some(2), Some(3), None]));
/// ```
///
/// This is an alternative API to appending directly to [`Self::values`] and
/// delimiting the result with [`Self::append`]
///
/// ```
/// # use arrow_array::builder::{Int32Builder, ListBuilder};
/// # use arrow_array::cast::AsArray;
/// # use arrow_array::{Array, Int32Array};
/// # use arrow_array::types::Int32Type;
/// let mut builder = ListBuilder::new(Int32Builder::new());
///
/// builder.values().append_value(1);
/// builder.values().append_value(2);
/// builder.values().append_value(3);
/// builder.append(true);
/// builder.append(true);
/// builder.values().append_null();
/// builder.append(true);
///
/// let array = builder.finish();
/// assert_eq!(array.len(), 3);
///
/// assert_eq!(array.value_offsets(), &[0, 3, 3, 4]);
/// let values = array.values().as_primitive::<Int32Type>();
/// assert_eq!(values, &Int32Array::from(vec![Some(1), Some(2), Some(3), None]));
/// ```
#[inline]
pub fn append_value<I, V>(&mut self, i: I)
where
T: Extend<Option<V>>,
I: IntoIterator<Item = Option<V>>,
{
self.extend(std::iter::once(Some(i)))
}
/// Append a null to this [`GenericListBuilder`]
///
/// See [`Self::append_value`] for an example use.
#[inline]
pub fn append_null(&mut self) {
self.offsets_builder.push(self.next_offset());
self.null_buffer_builder.append_null();
}
/// Appends `n` `null`s into the builder.
#[inline]
pub fn append_nulls(&mut self, n: usize) {
let next_offset = self.next_offset();
self.offsets_builder
.extend(std::iter::repeat_n(next_offset, n));
self.null_buffer_builder.append_n_nulls(n);
}
/// Appends an optional value into this [`GenericListBuilder`]
///
/// If `Some` calls [`Self::append_value`] otherwise calls [`Self::append_null`]
#[inline]
pub fn append_option<I, V>(&mut self, i: Option<I>)
where
T: Extend<Option<V>>,
I: IntoIterator<Item = Option<V>>,
{
match i {
Some(i) => self.append_value(i),
None => self.append_null(),
}
}
/// Builds the [`GenericListArray`] and reset this builder.
pub fn finish(&mut self) -> GenericListArray<OffsetSize> {
let values = self.values_builder.finish();
let nulls = self.null_buffer_builder.finish();
let offsets = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
// Safety: Safe by construction
let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
self.offsets_builder.push(OffsetSize::zero());
let field = match &self.field {
Some(f) => f.clone(),
None => Arc::new(Field::new_list_field(values.data_type().clone(), true)),
};
GenericListArray::new(field, offsets, values, nulls)
}
/// Builds the [`GenericListArray`] without resetting the builder.
pub fn finish_cloned(&self) -> GenericListArray<OffsetSize> {
let values = self.values_builder.finish_cloned();
let nulls = self.null_buffer_builder.finish_cloned();
let offsets = Buffer::from_slice_ref(self.offsets_builder.as_slice());
// Safety: safe by construction
let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
let field = match &self.field {
Some(f) => f.clone(),
None => Arc::new(Field::new_list_field(values.data_type().clone(), true)),
};
GenericListArray::new(field, offsets, values, nulls)
}
/// Returns the current offsets buffer as a slice
pub fn offsets_slice(&self) -> &[OffsetSize] {
self.offsets_builder.as_slice()
}
/// Returns the current null buffer as a slice
pub fn validity_slice(&self) -> Option<&[u8]> {
self.null_buffer_builder.as_slice()
}
}
impl<O, B, V, E> Extend<Option<V>> for GenericListBuilder<O, B>
where
O: OffsetSizeTrait,
B: ArrayBuilder + Extend<E>,
V: IntoIterator<Item = E>,
{
#[inline]
fn extend<T: IntoIterator<Item = Option<V>>>(&mut self, iter: T) {
for v in iter {
match v {
Some(elements) => {
self.values_builder.extend(elements);
self.append(true);
}
None => self.append(false),
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Int32Array;
use crate::builder::{Int32Builder, ListBuilder, make_builder};
use crate::cast::AsArray;
use crate::types::Int32Type;
use arrow_schema::DataType;
fn _test_generic_list_array_builder<O: OffsetSizeTrait>() {
let values_builder = Int32Builder::with_capacity(10);
let mut builder = GenericListBuilder::<O, _>::new(values_builder);
// [[0, 1, 2], [3, 4, 5], [6, 7]]
builder.values().append_value(0);
builder.values().append_value(1);
builder.values().append_value(2);
builder.append(true);
builder.values().append_value(3);
builder.values().append_value(4);
builder.values().append_value(5);
builder.append(true);
builder.values().append_value(6);
builder.values().append_value(7);
builder.append(true);
let list_array = builder.finish();
let list_values = list_array.values().as_primitive::<Int32Type>();
assert_eq!(list_values.values(), &[0, 1, 2, 3, 4, 5, 6, 7]);
assert_eq!(list_array.value_offsets(), [0, 3, 6, 8].map(O::usize_as));
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(3, list_array.len());
assert_eq!(0, list_array.null_count());
assert_eq!(O::from_usize(6).unwrap(), list_array.value_offsets()[2]);
assert_eq!(O::from_usize(2).unwrap(), list_array.value_length(2));
for i in 0..3 {
assert!(list_array.is_valid(i));
assert!(!list_array.is_null(i));
}
}
#[test]
fn test_list_array_builder() {
_test_generic_list_array_builder::<i32>()
}
#[test]
fn test_large_list_array_builder() {
_test_generic_list_array_builder::<i64>()
}
fn _test_generic_list_array_builder_nulls<O: OffsetSizeTrait>() {
let values_builder = Int32Builder::with_capacity(10);
let mut builder = GenericListBuilder::<O, _>::new(values_builder);
// [[0, 1, 2], null, [3, null, 5], [6, 7], null, null, [8]]
builder.values().append_value(0);
builder.values().append_value(1);
builder.values().append_value(2);
builder.append(true);
builder.append(false);
builder.values().append_value(3);
builder.values().append_null();
builder.values().append_value(5);
builder.append(true);
builder.values().append_value(6);
builder.values().append_value(7);
builder.append(true);
builder.append_nulls(2);
builder.values().append_value(8);
builder.append(true);
let list_array = builder.finish();
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(7, list_array.len());
assert_eq!(3, list_array.null_count());
assert_eq!(O::from_usize(3).unwrap(), list_array.value_offsets()[2]);
assert_eq!(O::from_usize(9).unwrap(), list_array.value_offsets()[7]);
assert_eq!(O::from_usize(3).unwrap(), list_array.value_length(2));
assert!(list_array.is_null(4));
assert!(list_array.is_null(5));
}
#[test]
fn test_list_array_builder_nulls() {
_test_generic_list_array_builder_nulls::<i32>()
}
#[test]
fn test_large_list_array_builder_nulls() {
_test_generic_list_array_builder_nulls::<i64>()
}
#[test]
fn test_list_array_builder_finish() {
let values_builder = Int32Array::builder(5);
let mut builder = ListBuilder::new(values_builder);
builder.values().append_slice(&[1, 2, 3]);
builder.append(true);
builder.values().append_slice(&[4, 5, 6]);
builder.append(true);
let mut arr = builder.finish();
assert_eq!(2, arr.len());
assert!(builder.is_empty());
builder.values().append_slice(&[7, 8, 9]);
builder.append(true);
arr = builder.finish();
assert_eq!(1, arr.len());
assert!(builder.is_empty());
}
#[test]
fn test_list_array_builder_finish_cloned() {
let values_builder = Int32Array::builder(5);
let mut builder = ListBuilder::new(values_builder);
builder.values().append_slice(&[1, 2, 3]);
builder.append(true);
builder.values().append_slice(&[4, 5, 6]);
builder.append(true);
let mut arr = builder.finish_cloned();
assert_eq!(2, arr.len());
assert!(!builder.is_empty());
builder.values().append_slice(&[7, 8, 9]);
builder.append(true);
arr = builder.finish();
assert_eq!(3, arr.len());
assert!(builder.is_empty());
}
#[test]
fn test_list_list_array_builder() {
let primitive_builder = Int32Builder::with_capacity(10);
let values_builder = ListBuilder::new(primitive_builder);
let mut builder = ListBuilder::new(values_builder);
// [[[1, 2], [3, 4]], [[5, 6, 7], null, [8]], null, [[9, 10]]]
builder.values().values().append_value(1);
builder.values().values().append_value(2);
builder.values().append(true);
builder.values().values().append_value(3);
builder.values().values().append_value(4);
builder.values().append(true);
builder.append(true);
builder.values().values().append_value(5);
builder.values().values().append_value(6);
builder.values().values().append_value(7);
builder.values().append(true);
builder.values().append(false);
builder.values().values().append_value(8);
builder.values().append(true);
builder.append(true);
builder.append(false);
builder.values().values().append_value(9);
builder.values().values().append_value(10);
builder.values().append(true);
builder.append(true);
let l1 = builder.finish();
assert_eq!(4, l1.len());
assert_eq!(1, l1.null_count());
assert_eq!(l1.value_offsets(), &[0, 2, 5, 5, 6]);
let l2 = l1.values().as_list::<i32>();
assert_eq!(6, l2.len());
assert_eq!(1, l2.null_count());
assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8, 10]);
let i1 = l2.values().as_primitive::<Int32Type>();
assert_eq!(10, i1.len());
assert_eq!(0, i1.null_count());
assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
}
#[test]
fn test_extend() {
let mut builder = ListBuilder::new(Int32Builder::new());
builder.extend([
Some(vec![Some(1), Some(2), Some(7), None]),
Some(vec![]),
Some(vec![Some(4), Some(5)]),
None,
]);
let array = builder.finish();
assert_eq!(array.value_offsets(), [0, 4, 4, 6, 6]);
assert_eq!(array.null_count(), 1);
assert_eq!(array.logical_null_count(), 1);
assert!(array.is_null(3));
let elements = array.values().as_primitive::<Int32Type>();
assert_eq!(elements.values(), &[1, 2, 7, 0, 4, 5]);
assert_eq!(elements.null_count(), 1);
assert_eq!(elements.logical_null_count(), 1);
assert!(elements.is_null(3));
}
#[test]
fn test_boxed_primitive_array_builder() {
let values_builder = make_builder(&DataType::Int32, 5);
let mut builder = ListBuilder::new(values_builder);
builder
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_slice(&[1, 2, 3]);
builder.append(true);
builder
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_slice(&[4, 5, 6]);
builder.append(true);
let arr = builder.finish();
assert_eq!(2, arr.len());
let elements = arr.values().as_primitive::<Int32Type>();
assert_eq!(elements.values(), &[1, 2, 3, 4, 5, 6]);
}
#[test]
fn test_boxed_list_list_array_builder() {
// This test is same as `test_list_list_array_builder` but uses boxed builders.
let values_builder = make_builder(
&DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))),
10,
);
test_boxed_generic_list_generic_list_array_builder::<i32>(values_builder);
}
#[test]
fn test_boxed_large_list_large_list_array_builder() {
// This test is same as `test_list_list_array_builder` but uses boxed builders.
let values_builder = make_builder(
&DataType::LargeList(Arc::new(Field::new_list_field(DataType::Int32, true))),
10,
);
test_boxed_generic_list_generic_list_array_builder::<i64>(values_builder);
}
fn test_boxed_generic_list_generic_list_array_builder<O: OffsetSizeTrait + PartialEq>(
values_builder: Box<dyn ArrayBuilder>,
) {
let mut builder: GenericListBuilder<O, Box<dyn ArrayBuilder>> =
GenericListBuilder::<O, Box<dyn ArrayBuilder>>::new(values_builder);
// [[[1, 2], [3, 4]], [[5, 6, 7], null, [8]], null, [[9, 10]]]
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(1);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(2);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.append(true);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(3);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(4);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.append(true);
builder.append(true);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(5);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(6);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an (Large)ListBuilder")
.append_value(7);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.append(true);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.append(false);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(8);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.append(true);
builder.append(true);
builder.append(false);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(9);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(10);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.append(true);
builder.append(true);
let l1 = builder.finish();
assert_eq!(4, l1.len());
assert_eq!(1, l1.null_count());
assert_eq!(l1.value_offsets(), &[0, 2, 5, 5, 6].map(O::usize_as));
let l2 = l1.values().as_list::<O>();
assert_eq!(6, l2.len());
assert_eq!(1, l2.null_count());
assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8, 10].map(O::usize_as));
let i1 = l2.values().as_primitive::<Int32Type>();
assert_eq!(10, i1.len());
assert_eq!(0, i1.null_count());
assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
}
#[test]
fn test_with_field() {
let field = Arc::new(Field::new("bar", DataType::Int32, false));
let mut builder = ListBuilder::new(Int32Builder::new()).with_field(field.clone());
builder.append_value([Some(1), Some(2), Some(3)]);
builder.append_null(); // This is fine as nullability refers to nullability of values
builder.append_value([Some(4)]);
let array = builder.finish();
assert_eq!(array.len(), 3);
assert_eq!(array.data_type(), &DataType::List(field.clone()));
builder.append_value([Some(4), Some(5)]);
let array = builder.finish();
assert_eq!(array.data_type(), &DataType::List(field));
assert_eq!(array.len(), 1);
}
#[test]
#[should_panic(expected = "Non-nullable field of ListArray \\\"item\\\" cannot contain nulls")]
fn test_checks_nullability() {
let field = Arc::new(Field::new_list_field(DataType::Int32, false));
let mut builder = ListBuilder::new(Int32Builder::new()).with_field(field.clone());
builder.append_value([Some(1), None]);
builder.finish();
}
#[test]
#[should_panic(expected = "ListArray expected data type Int64 got Int32")]
fn test_checks_data_type() {
let field = Arc::new(Field::new_list_field(DataType::Int64, false));
let mut builder = ListBuilder::new(Int32Builder::new()).with_field(field.clone());
builder.append_value([Some(1)]);
builder.finish();
}
}
@@ -0,0 +1,706 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::builder::ArrayBuilder;
use crate::{ArrayRef, GenericListViewArray, OffsetSizeTrait};
use arrow_buffer::{Buffer, NullBufferBuilder, ScalarBuffer};
use arrow_schema::{Field, FieldRef};
use std::any::Any;
use std::sync::Arc;
/// Builder for [`GenericListViewArray`]
#[derive(Debug)]
pub struct GenericListViewBuilder<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> {
offsets_builder: Vec<OffsetSize>,
sizes_builder: Vec<OffsetSize>,
null_buffer_builder: NullBufferBuilder,
values_builder: T,
field: Option<FieldRef>,
current_offset: OffsetSize,
}
impl<O: OffsetSizeTrait, T: ArrayBuilder + Default> Default for GenericListViewBuilder<O, T> {
fn default() -> Self {
Self::new(T::default())
}
}
impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> ArrayBuilder
for GenericListViewBuilder<OffsetSize, T>
{
/// Returns the builder as a non-mutable `Any` reference.
fn as_any(&self) -> &dyn Any {
self
}
/// Returns the builder as a mutable `Any` reference.
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
/// Returns the boxed builder as a box of `Any`.
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
/// Returns the number of array slots in the builder
fn len(&self) -> usize {
self.null_buffer_builder.len()
}
/// Builds the array and reset this builder.
fn finish(&mut self) -> ArrayRef {
Arc::new(self.finish())
}
/// Builds the array without resetting the builder.
fn finish_cloned(&self) -> ArrayRef {
Arc::new(self.finish_cloned())
}
}
impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> GenericListViewBuilder<OffsetSize, T> {
/// Creates a new [`GenericListViewBuilder`] from a given values array builder
pub fn new(values_builder: T) -> Self {
let capacity = values_builder.len();
Self::with_capacity(values_builder, capacity)
}
/// Creates a new [`GenericListViewBuilder`] from a given values array builder
/// `capacity` is the number of items to pre-allocate space for in this builder
pub fn with_capacity(values_builder: T, capacity: usize) -> Self {
let offsets_builder = Vec::with_capacity(capacity);
let sizes_builder = Vec::with_capacity(capacity);
Self {
offsets_builder,
null_buffer_builder: NullBufferBuilder::new(capacity),
values_builder,
sizes_builder,
field: None,
current_offset: OffsetSize::zero(),
}
}
///
/// By default a nullable field is created with the name `item`
///
/// Note: [`Self::finish`] and [`Self::finish_cloned`] will panic if the
/// field's data type does not match that of `T`
pub fn with_field(self, field: impl Into<FieldRef>) -> Self {
Self {
field: Some(field.into()),
..self
}
}
}
impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> GenericListViewBuilder<OffsetSize, T>
where
T: 'static,
{
/// Returns the child array builder as a mutable reference.
///
/// This mutable reference can be used to append values into the child array builder,
/// but you must call [`append`](#method.append) to delimit each distinct list value.
pub fn values(&mut self) -> &mut T {
&mut self.values_builder
}
/// Returns the child array builder as an immutable reference
pub fn values_ref(&self) -> &T {
&self.values_builder
}
/// Finish the current variable-length list array slot
///
/// # Panics
///
/// Panics if the length of [`Self::values`] exceeds `OffsetSize::MAX`
#[inline]
pub fn append(&mut self, is_valid: bool) {
self.offsets_builder.push(self.current_offset);
self.sizes_builder.push(
OffsetSize::from_usize(
self.values_builder.len() - self.current_offset.to_usize().unwrap(),
)
.unwrap(),
);
self.null_buffer_builder.append(is_valid);
self.current_offset = OffsetSize::from_usize(self.values_builder.len()).unwrap();
}
/// Append value into this [`GenericListViewBuilder`]
#[inline]
pub fn append_value<I, V>(&mut self, i: I)
where
T: Extend<Option<V>>,
I: IntoIterator<Item = Option<V>>,
{
self.extend(std::iter::once(Some(i)))
}
/// Append a null to this [`GenericListViewBuilder`]
///
/// See [`Self::append_value`] for an example use.
#[inline]
pub fn append_null(&mut self) {
self.offsets_builder.push(self.current_offset);
self.sizes_builder.push(OffsetSize::from_usize(0).unwrap());
self.null_buffer_builder.append_null();
}
/// Appends an optional value into this [`GenericListViewBuilder`]
///
/// If `Some` calls [`Self::append_value`] otherwise calls [`Self::append_null`]
#[inline]
pub fn append_option<I, V>(&mut self, i: Option<I>)
where
T: Extend<Option<V>>,
I: IntoIterator<Item = Option<V>>,
{
match i {
Some(i) => self.append_value(i),
None => self.append_null(),
}
}
/// Builds the [`GenericListViewArray`] and reset this builder.
pub fn finish(&mut self) -> GenericListViewArray<OffsetSize> {
let values = self.values_builder.finish();
let nulls = self.null_buffer_builder.finish();
let offsets = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
self.current_offset = OffsetSize::zero();
// Safety: Safe by construction
let offsets = ScalarBuffer::from(offsets);
let sizes = Buffer::from_vec(std::mem::take(&mut self.sizes_builder));
let sizes = ScalarBuffer::from(sizes);
let field = match &self.field {
Some(f) => f.clone(),
None => Arc::new(Field::new("item", values.data_type().clone(), true)),
};
GenericListViewArray::new(field, offsets, sizes, values, nulls)
}
/// Builds the [`GenericListViewArray`] without resetting the builder.
pub fn finish_cloned(&self) -> GenericListViewArray<OffsetSize> {
let values = self.values_builder.finish_cloned();
let nulls = self.null_buffer_builder.finish_cloned();
let offsets = Buffer::from_slice_ref(self.offsets_builder.as_slice());
// Safety: safe by construction
let offsets = ScalarBuffer::from(offsets);
let sizes = Buffer::from_slice_ref(self.sizes_builder.as_slice());
let sizes = ScalarBuffer::from(sizes);
let field = match &self.field {
Some(f) => f.clone(),
None => Arc::new(Field::new("item", values.data_type().clone(), true)),
};
GenericListViewArray::new(field, offsets, sizes, values, nulls)
}
/// Returns the current offsets buffer as a slice
pub fn offsets_slice(&self) -> &[OffsetSize] {
self.offsets_builder.as_slice()
}
}
impl<O, B, V, E> Extend<Option<V>> for GenericListViewBuilder<O, B>
where
O: OffsetSizeTrait,
B: ArrayBuilder + Extend<E>,
V: IntoIterator<Item = E>,
{
#[inline]
fn extend<T: IntoIterator<Item = Option<V>>>(&mut self, iter: T) {
for v in iter {
match v {
Some(elements) => {
self.values_builder.extend(elements);
self.append(true);
}
None => self.append(false),
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::builder::{Int32Builder, ListViewBuilder, make_builder};
use crate::cast::AsArray;
use crate::types::Int32Type;
use crate::{Array, Int32Array};
use arrow_schema::DataType;
fn test_generic_list_view_array_builder_impl<O: OffsetSizeTrait>() {
let values_builder = Int32Builder::with_capacity(10);
let mut builder = GenericListViewBuilder::<O, _>::new(values_builder);
// [[0, 1, 2], [3, 4, 5], [6, 7]]
builder.values().append_value(0);
builder.values().append_value(1);
builder.values().append_value(2);
builder.append(true);
builder.values().append_value(3);
builder.values().append_value(4);
builder.values().append_value(5);
builder.append(true);
builder.values().append_value(6);
builder.values().append_value(7);
builder.append(true);
let list_array = builder.finish();
let list_values = list_array.values().as_primitive::<Int32Type>();
assert_eq!(list_values.values(), &[0, 1, 2, 3, 4, 5, 6, 7]);
assert_eq!(list_array.value_offsets(), [0, 3, 6].map(O::usize_as));
assert_eq!(list_array.value_sizes(), [3, 3, 2].map(O::usize_as));
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(3, list_array.len());
assert_eq!(0, list_array.null_count());
assert_eq!(O::from_usize(6).unwrap(), list_array.value_offsets()[2]);
assert_eq!(O::from_usize(2).unwrap(), list_array.value_sizes()[2]);
for i in 0..2 {
assert!(list_array.is_valid(i));
assert!(!list_array.is_null(i));
}
}
#[test]
fn test_list_view_array_builder() {
test_generic_list_view_array_builder_impl::<i32>()
}
#[test]
fn test_large_list_view_array_builder() {
test_generic_list_view_array_builder_impl::<i64>()
}
fn test_generic_list_view_array_builder_nulls_impl<O: OffsetSizeTrait>() {
let values_builder = Int32Builder::with_capacity(10);
let mut builder = GenericListViewBuilder::<O, _>::new(values_builder);
// [[0, 1, 2], null, [3, null, 5], [6, 7]]
builder.values().append_value(0);
builder.values().append_value(1);
builder.values().append_value(2);
builder.append(true);
builder.append(false);
builder.values().append_value(3);
builder.values().append_null();
builder.values().append_value(5);
builder.append(true);
builder.values().append_value(6);
builder.values().append_value(7);
builder.append(true);
let list_array = builder.finish();
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(4, list_array.len());
assert_eq!(1, list_array.null_count());
assert_eq!(O::from_usize(3).unwrap(), list_array.value_offsets()[2]);
assert_eq!(O::from_usize(3).unwrap(), list_array.value_sizes()[2]);
}
#[test]
fn test_list_view_array_builder_nulls() {
test_generic_list_view_array_builder_nulls_impl::<i32>()
}
#[test]
fn test_large_list_view_array_builder_nulls() {
test_generic_list_view_array_builder_nulls_impl::<i64>()
}
#[test]
fn test_list_view_array_builder_finish() {
let values_builder = Int32Array::builder(5);
let mut builder = ListViewBuilder::new(values_builder);
builder.values().append_slice(&[1, 2, 3]);
builder.append(true);
builder.values().append_slice(&[4, 5, 6]);
builder.append(true);
let mut arr = builder.finish();
assert_eq!(2, arr.len());
assert!(builder.is_empty());
builder.values().append_slice(&[7, 8, 9]);
builder.append(true);
arr = builder.finish();
assert_eq!(1, arr.len());
assert!(builder.is_empty());
}
#[test]
fn test_list_view_array_builder_finish_cloned() {
let values_builder = Int32Array::builder(5);
let mut builder = ListViewBuilder::new(values_builder);
builder.values().append_slice(&[1, 2, 3]);
builder.append(true);
builder.values().append_slice(&[4, 5, 6]);
builder.append(true);
let mut arr = builder.finish_cloned();
assert_eq!(2, arr.len());
assert!(!builder.is_empty());
builder.values().append_slice(&[7, 8, 9]);
builder.append(true);
arr = builder.finish();
assert_eq!(3, arr.len());
assert!(builder.is_empty());
}
#[test]
fn test_list_view_list_view_array_builder() {
let primitive_builder = Int32Builder::with_capacity(10);
let values_builder = ListViewBuilder::new(primitive_builder);
let mut builder = ListViewBuilder::new(values_builder);
// [[[1, 2], [3, 4]], [[5, 6, 7], null, [8]], null, [[9, 10]]]
builder.values().values().append_value(1);
builder.values().values().append_value(2);
builder.values().append(true);
builder.values().values().append_value(3);
builder.values().values().append_value(4);
builder.values().append(true);
builder.append(true);
builder.values().values().append_value(5);
builder.values().values().append_value(6);
builder.values().values().append_value(7);
builder.values().append(true);
builder.values().append(false);
builder.values().values().append_value(8);
builder.values().append(true);
builder.append(true);
builder.append(false);
builder.values().values().append_value(9);
builder.values().values().append_value(10);
builder.values().append(true);
builder.append(true);
let l1 = builder.finish();
assert_eq!(4, l1.len());
assert_eq!(1, l1.null_count());
assert_eq!(l1.value_offsets(), &[0, 2, 5, 5]);
assert_eq!(l1.value_sizes(), &[2, 3, 0, 1]);
let l2 = l1.values().as_list_view::<i32>();
assert_eq!(6, l2.len());
assert_eq!(1, l2.null_count());
assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8]);
assert_eq!(l2.value_sizes(), &[2, 2, 3, 0, 1, 2]);
let i1 = l2.values().as_primitive::<Int32Type>();
assert_eq!(10, i1.len());
assert_eq!(0, i1.null_count());
assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
}
#[test]
fn test_extend() {
let mut builder = ListViewBuilder::new(Int32Builder::new());
builder.extend([
Some(vec![Some(1), Some(2), Some(7), None]),
Some(vec![]),
Some(vec![Some(4), Some(5)]),
None,
]);
let array = builder.finish();
assert_eq!(array.value_offsets(), [0, 4, 4, 6]);
assert_eq!(array.value_sizes(), [4, 0, 2, 0]);
assert_eq!(array.null_count(), 1);
assert!(array.is_null(3));
let elements = array.values().as_primitive::<Int32Type>();
assert_eq!(elements.values(), &[1, 2, 7, 0, 4, 5]);
assert_eq!(elements.null_count(), 1);
assert!(elements.is_null(3));
}
#[test]
fn test_boxed_primitive_array_builder() {
let values_builder = make_builder(&DataType::Int32, 5);
let mut builder = ListViewBuilder::new(values_builder);
builder
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_slice(&[1, 2, 3]);
builder.append(true);
builder
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_slice(&[4, 5, 6]);
builder.append(true);
let arr = builder.finish();
assert_eq!(2, arr.len());
let elements = arr.values().as_primitive::<Int32Type>();
assert_eq!(elements.values(), &[1, 2, 3, 4, 5, 6]);
}
#[test]
fn test_boxed_list_view_list_view_array_builder() {
// This test is same as `test_list_list_array_builder` but uses boxed builders.
let values_builder = make_builder(
&DataType::ListView(Arc::new(Field::new("item", DataType::Int32, true))),
10,
);
test_boxed_generic_list_view_generic_list_view_array_builder::<i32>(values_builder);
}
#[test]
fn test_boxed_large_list_view_large_list_view_array_builder() {
// This test is same as `test_list_list_array_builder` but uses boxed builders.
let values_builder = make_builder(
&DataType::LargeListView(Arc::new(Field::new("item", DataType::Int32, true))),
10,
);
test_boxed_generic_list_view_generic_list_view_array_builder::<i64>(values_builder);
}
fn test_boxed_generic_list_view_generic_list_view_array_builder<O>(
values_builder: Box<dyn ArrayBuilder>,
) where
O: OffsetSizeTrait + PartialEq,
{
let mut builder: GenericListViewBuilder<O, Box<dyn ArrayBuilder>> =
GenericListViewBuilder::<O, Box<dyn ArrayBuilder>>::new(values_builder);
// [[[1, 2], [3, 4]], [[5, 6, 7], null, [8]], null, [[9, 10]]]
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(1);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(2);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.append(true);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(3);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(4);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.append(true);
builder.append(true);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(5);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(6);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an (Large)ListViewBuilder")
.append_value(7);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.append(true);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.append(false);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(8);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.append(true);
builder.append(true);
builder.append(false);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(9);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(10);
builder
.values()
.as_any_mut()
.downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListViewBuilder")
.append(true);
builder.append(true);
let l1 = builder.finish();
assert_eq!(4, l1.len());
assert_eq!(1, l1.null_count());
assert_eq!(l1.value_offsets(), &[0, 2, 5, 5].map(O::usize_as));
assert_eq!(l1.value_sizes(), &[2, 3, 0, 1].map(O::usize_as));
let l2 = l1.values().as_list_view::<O>();
assert_eq!(6, l2.len());
assert_eq!(1, l2.null_count());
assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8].map(O::usize_as));
assert_eq!(l2.value_sizes(), &[2, 2, 3, 0, 1, 2].map(O::usize_as));
let i1 = l2.values().as_primitive::<Int32Type>();
assert_eq!(10, i1.len());
assert_eq!(0, i1.null_count());
assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
}
#[test]
fn test_with_field() {
let field = Arc::new(Field::new("bar", DataType::Int32, false));
let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone());
builder.append_value([Some(1), Some(2), Some(3)]);
builder.append_null(); // This is fine as nullability refers to nullability of values
builder.append_value([Some(4)]);
let array = builder.finish();
assert_eq!(array.len(), 3);
assert_eq!(array.data_type(), &DataType::ListView(field.clone()));
builder.append_value([Some(4), Some(5)]);
let array = builder.finish();
assert_eq!(array.data_type(), &DataType::ListView(field));
assert_eq!(array.len(), 1);
}
#[test]
#[should_panic(
expected = r#"Non-nullable field of ListViewArray \"item\" cannot contain nulls"#
)]
// If a non-nullable type is declared but a null value is used, it will be intercepted by the null check.
fn test_checks_nullability() {
let field = Arc::new(Field::new("item", DataType::Int32, false));
let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone());
builder.append_value([Some(1), None]);
builder.finish();
}
#[test]
#[should_panic(expected = "ListViewArray expected data type Int64 got Int32")]
// If the declared type does not match the actual appended type, it will be intercepted by type checking in the finish function.
fn test_checks_data_type() {
let field = Arc::new(Field::new("item", DataType::Int64, false));
let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone());
builder.append_value([Some(1)]);
builder.finish();
}
}
@@ -0,0 +1,464 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::builder::ArrayBuilder;
use crate::{Array, ArrayRef, MapArray, StructArray};
use arrow_buffer::Buffer;
use arrow_buffer::{NullBuffer, NullBufferBuilder};
use arrow_data::ArrayData;
use arrow_schema::{ArrowError, DataType, Field, FieldRef};
use std::any::Any;
use std::sync::Arc;
/// Builder for [`MapArray`]
///
/// ```
/// # use arrow_array::builder::{Int32Builder, MapBuilder, StringBuilder};
/// # use arrow_array::{Int32Array, StringArray};
///
/// let string_builder = StringBuilder::new();
/// let int_builder = Int32Builder::with_capacity(4);
///
/// // Construct `[{"joe": 1}, {"blogs": 2, "foo": 4}, {}, null]`
/// let mut builder = MapBuilder::new(None, string_builder, int_builder);
///
/// builder.keys().append_value("joe");
/// builder.values().append_value(1);
/// builder.append(true).unwrap();
///
/// builder.keys().append_value("blogs");
/// builder.values().append_value(2);
/// builder.keys().append_value("foo");
/// builder.values().append_value(4);
/// builder.append(true).unwrap();
/// builder.append(true).unwrap();
/// builder.append(false).unwrap();
///
/// let array = builder.finish();
/// assert_eq!(array.value_offsets(), &[0, 1, 3, 3, 3]);
/// assert_eq!(array.values().as_ref(), &Int32Array::from(vec![1, 2, 4]));
/// assert_eq!(array.keys().as_ref(), &StringArray::from(vec!["joe", "blogs", "foo"]));
///
/// ```
#[derive(Debug)]
pub struct MapBuilder<K: ArrayBuilder, V: ArrayBuilder> {
offsets_builder: Vec<i32>,
null_buffer_builder: NullBufferBuilder,
field_names: MapFieldNames,
key_builder: K,
value_builder: V,
key_field: Option<FieldRef>,
value_field: Option<FieldRef>,
}
/// The [`Field`] names for a [`MapArray`]
#[derive(Debug, Clone)]
pub struct MapFieldNames {
/// [`Field`] name for map entries
pub entry: String,
/// [`Field`] name for map key
pub key: String,
/// [`Field`] name for map value
pub value: String,
}
impl Default for MapFieldNames {
fn default() -> Self {
Self {
entry: "entries".to_string(),
key: "keys".to_string(),
value: "values".to_string(),
}
}
}
impl<K: ArrayBuilder, V: ArrayBuilder> MapBuilder<K, V> {
/// Creates a new `MapBuilder`
pub fn new(field_names: Option<MapFieldNames>, key_builder: K, value_builder: V) -> Self {
let capacity = key_builder.len();
Self::with_capacity(field_names, key_builder, value_builder, capacity)
}
/// Creates a new `MapBuilder` with capacity
pub fn with_capacity(
field_names: Option<MapFieldNames>,
key_builder: K,
value_builder: V,
capacity: usize,
) -> Self {
let mut offsets_builder = Vec::with_capacity(capacity + 1);
offsets_builder.push(0);
Self {
offsets_builder,
null_buffer_builder: NullBufferBuilder::new(capacity),
field_names: field_names.unwrap_or_default(),
key_builder,
value_builder,
key_field: None,
value_field: None,
}
}
/// Override the field passed to [`MapBuilder::new`]
///
/// By default, a non-nullable field is created with the name `keys`
///
/// Note: [`Self::finish`] and [`Self::finish_cloned`] will panic if the
/// field's data type does not match that of `K` or the field is nullable
pub fn with_keys_field(self, field: impl Into<FieldRef>) -> Self {
Self {
key_field: Some(field.into()),
..self
}
}
/// Override the field passed to [`MapBuilder::new`]
///
/// By default, a nullable field is created with the name `values`
///
/// Note: [`Self::finish`] and [`Self::finish_cloned`] will panic if the
/// field's data type does not match that of `V`
pub fn with_values_field(self, field: impl Into<FieldRef>) -> Self {
Self {
value_field: Some(field.into()),
..self
}
}
/// Returns the key array builder of the map
pub fn keys(&mut self) -> &mut K {
&mut self.key_builder
}
/// Returns the value array builder of the map
pub fn values(&mut self) -> &mut V {
&mut self.value_builder
}
/// Returns both the key and value array builders of the map
pub fn entries(&mut self) -> (&mut K, &mut V) {
(&mut self.key_builder, &mut self.value_builder)
}
/// Finish the current map array slot
///
/// Returns an error if the key and values builders are in an inconsistent state.
#[inline]
pub fn append(&mut self, is_valid: bool) -> Result<(), ArrowError> {
if self.key_builder.len() != self.value_builder.len() {
return Err(ArrowError::InvalidArgumentError(format!(
"Cannot append to a map builder when its keys and values have unequal lengths of {} and {}",
self.key_builder.len(),
self.value_builder.len()
)));
}
self.offsets_builder.push(self.key_builder.len() as i32);
self.null_buffer_builder.append(is_valid);
Ok(())
}
/// Builds the [`MapArray`]
pub fn finish(&mut self) -> MapArray {
let len = self.len();
// Build the keys
let keys_arr = self.key_builder.finish();
let values_arr = self.value_builder.finish();
let offset_buffer = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
self.offsets_builder.push(0);
let null_bit_buffer = self.null_buffer_builder.finish();
self.finish_helper(keys_arr, values_arr, offset_buffer, null_bit_buffer, len)
}
/// Builds the [`MapArray`] without resetting the builder.
pub fn finish_cloned(&self) -> MapArray {
let len = self.len();
// Build the keys
let keys_arr = self.key_builder.finish_cloned();
let values_arr = self.value_builder.finish_cloned();
let offset_buffer = Buffer::from_slice_ref(self.offsets_builder.as_slice());
let nulls = self.null_buffer_builder.finish_cloned();
self.finish_helper(keys_arr, values_arr, offset_buffer, nulls, len)
}
fn finish_helper(
&self,
keys_arr: Arc<dyn Array>,
values_arr: Arc<dyn Array>,
offset_buffer: Buffer,
nulls: Option<NullBuffer>,
len: usize,
) -> MapArray {
assert!(
keys_arr.null_count() == 0,
"Keys array must have no null values, found {} null value(s)",
keys_arr.null_count()
);
let keys_field = match &self.key_field {
Some(f) => {
assert!(!f.is_nullable(), "Keys field must not be nullable");
f.clone()
}
None => Arc::new(Field::new(
self.field_names.key.as_str(),
keys_arr.data_type().clone(),
false, // always non-nullable
)),
};
let values_field = match &self.value_field {
Some(f) => f.clone(),
None => Arc::new(Field::new(
self.field_names.value.as_str(),
values_arr.data_type().clone(),
true,
)),
};
let struct_array =
StructArray::from(vec![(keys_field, keys_arr), (values_field, values_arr)]);
let map_field = Arc::new(Field::new(
self.field_names.entry.as_str(),
struct_array.data_type().clone(),
false, // always non-nullable
));
let array_data = ArrayData::builder(DataType::Map(map_field, false)) // TODO: support sorted keys
.len(len)
.add_buffer(offset_buffer)
.add_child_data(struct_array.into_data())
.nulls(nulls);
let array_data = unsafe { array_data.build_unchecked() };
MapArray::from(array_data)
}
/// Returns the current null buffer as a slice
pub fn validity_slice(&self) -> Option<&[u8]> {
self.null_buffer_builder.as_slice()
}
}
impl<K: ArrayBuilder, V: ArrayBuilder> ArrayBuilder for MapBuilder<K, V> {
fn len(&self) -> usize {
self.null_buffer_builder.len()
}
fn finish(&mut self) -> ArrayRef {
Arc::new(self.finish())
}
/// Builds the array without resetting the builder.
fn finish_cloned(&self) -> ArrayRef {
Arc::new(self.finish_cloned())
}
fn as_any(&self) -> &dyn Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::builder::{Int32Builder, StringBuilder, make_builder};
use crate::{Int32Array, StringArray};
use std::collections::HashMap;
#[test]
#[should_panic(expected = "Keys array must have no null values, found 1 null value(s)")]
fn test_map_builder_with_null_keys_panics() {
let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
builder.keys().append_null();
builder.values().append_value(42);
builder.append(true).unwrap();
builder.finish();
}
#[test]
fn test_boxed_map_builder() {
let keys_builder = make_builder(&DataType::Utf8, 5);
let values_builder = make_builder(&DataType::Int32, 5);
let mut builder = MapBuilder::new(None, keys_builder, values_builder);
builder
.keys()
.as_any_mut()
.downcast_mut::<StringBuilder>()
.expect("should be an StringBuilder")
.append_value("1");
builder
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(42);
builder.append(true).unwrap();
let map_array = builder.finish();
assert_eq!(
map_array
.keys()
.as_any()
.downcast_ref::<StringArray>()
.expect("should be an StringArray")
.value(0),
"1"
);
assert_eq!(
map_array
.values()
.as_any()
.downcast_ref::<Int32Array>()
.expect("should be an Int32Array")
.value(0),
42
);
}
#[test]
fn test_with_values_field() {
let value_field = Arc::new(Field::new("bars", DataType::Int32, false));
let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new())
.with_values_field(value_field.clone());
builder.keys().append_value(1);
builder.values().append_value(2);
builder.append(true).unwrap();
builder.append(false).unwrap(); // This is fine as nullability refers to nullability of values
builder.keys().append_value(3);
builder.values().append_value(4);
builder.append(true).unwrap();
let map = builder.finish();
assert_eq!(map.len(), 3);
assert_eq!(
map.data_type(),
&DataType::Map(
Arc::new(Field::new(
"entries",
DataType::Struct(
vec![
Arc::new(Field::new("keys", DataType::Int32, false)),
value_field.clone()
]
.into()
),
false,
)),
false
)
);
builder.keys().append_value(5);
builder.values().append_value(6);
builder.append(true).unwrap();
let map = builder.finish();
assert_eq!(map.len(), 1);
assert_eq!(
map.data_type(),
&DataType::Map(
Arc::new(Field::new(
"entries",
DataType::Struct(
vec![
Arc::new(Field::new("keys", DataType::Int32, false)),
value_field
]
.into()
),
false,
)),
false
)
);
}
#[test]
fn test_with_keys_field() {
let mut key_metadata = HashMap::new();
key_metadata.insert("foo".to_string(), "bar".to_string());
let key_field = Arc::new(
Field::new("keys", DataType::Int32, false).with_metadata(key_metadata.clone()),
);
let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new())
.with_keys_field(key_field.clone());
builder.keys().append_value(1);
builder.values().append_value(2);
builder.append(true).unwrap();
let map = builder.finish();
assert_eq!(map.len(), 1);
assert_eq!(
map.data_type(),
&DataType::Map(
Arc::new(Field::new(
"entries",
DataType::Struct(
vec![
Arc::new(
Field::new("keys", DataType::Int32, false)
.with_metadata(key_metadata)
),
Arc::new(Field::new("values", DataType::Int32, true))
]
.into()
),
false,
)),
false
)
);
}
#[test]
#[should_panic(expected = "Keys field must not be nullable")]
fn test_with_nullable_keys_field() {
let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new())
.with_keys_field(Arc::new(Field::new("keys", DataType::Int32, true)));
builder.keys().append_value(1);
builder.values().append_value(2);
builder.append(true).unwrap();
builder.finish();
}
#[test]
#[should_panic(expected = "Incorrect datatype")]
fn test_keys_field_type_mismatch() {
let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new())
.with_keys_field(Arc::new(Field::new("keys", DataType::Utf8, false)));
builder.keys().append_value(1);
builder.values().append_value(2);
builder.append(true).unwrap();
builder.finish();
}
}
+615
View File
@@ -0,0 +1,615 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Defines push-based APIs for constructing arrays
//!
//! # Basic Usage
//!
//! Builders can be used to build simple, non-nested arrays
//!
//! ```
//! # use arrow_array::builder::Int32Builder;
//! # use arrow_array::PrimitiveArray;
//! let mut a = Int32Builder::new();
//! a.append_value(1);
//! a.append_null();
//! a.append_value(2);
//! let a = a.finish();
//!
//! assert_eq!(a, PrimitiveArray::from(vec![Some(1), None, Some(2)]));
//! ```
//!
//! ```
//! # use arrow_array::builder::StringBuilder;
//! # use arrow_array::{Array, StringArray};
//! let mut a = StringBuilder::new();
//! a.append_value("foo");
//! a.append_value("bar");
//! a.append_null();
//! let a = a.finish();
//!
//! assert_eq!(a, StringArray::from_iter([Some("foo"), Some("bar"), None]));
//! ```
//!
//! # Nested Usage
//!
//! Builders can also be used to build more complex nested arrays, such as lists
//!
//! ```
//! # use arrow_array::builder::{Int32Builder, ListBuilder};
//! # use arrow_array::ListArray;
//! # use arrow_array::types::Int32Type;
//! let mut a = ListBuilder::new(Int32Builder::new());
//! // [1, 2]
//! a.values().append_value(1);
//! a.values().append_value(2);
//! a.append(true);
//! // null
//! a.append(false);
//! // []
//! a.append(true);
//! // [3, null]
//! a.values().append_value(3);
//! a.values().append_null();
//! a.append(true);
//!
//! // [[1, 2], null, [], [3, null]]
//! let a = a.finish();
//!
//! assert_eq!(a, ListArray::from_iter_primitive::<Int32Type, _, _>([
//! Some(vec![Some(1), Some(2)]),
//! None,
//! Some(vec![]),
//! Some(vec![Some(3), None])]
//! ))
//! ```
//!
//! # Using the [`Extend`] trait to append values from an iterable:
//!
//! ```
//! # use arrow_array::{Array};
//! # use arrow_array::builder::{ArrayBuilder, StringBuilder};
//!
//! let mut builder = StringBuilder::new();
//! builder.extend(vec![Some("🍐"), Some("🍎"), None]);
//! assert_eq!(builder.finish().len(), 3);
//! ```
//!
//! # Using the [`Extend`] trait to write generic functions:
//!
//! ```
//! # use arrow_array::{Array, ArrayRef, StringArray};
//! # use arrow_array::builder::{ArrayBuilder, Int32Builder, ListBuilder, StringBuilder};
//!
//! // For generic methods that fill a list of values for an [`ArrayBuilder`], use the [`Extend`] trait.
//! fn filter_and_fill<V, I: IntoIterator<Item = V>>(builder: &mut impl Extend<V>, values: I, filter: V)
//! where V: PartialEq
//! {
//! builder.extend(values.into_iter().filter(|v| *v == filter));
//! }
//! let mut string_builder = StringBuilder::new();
//! filter_and_fill(
//! &mut string_builder,
//! vec![Some("🍐"), Some("🍎"), None],
//! Some("🍎"),
//! );
//! assert_eq!(string_builder.finish().len(), 1);
//!
//! let mut int_builder = Int32Builder::new();
//! filter_and_fill(
//! &mut int_builder,
//! vec![Some(11), Some(42), None],
//! Some(42),
//! );
//! assert_eq!(int_builder.finish().len(), 1);
//!
//! // For generic methods that fill lists-of-lists for an [`ArrayBuilder`], use the [`Extend`] trait.
//! fn filter_and_fill_if_contains<T, V, I: IntoIterator<Item = Option<V>>>(
//! list_builder: &mut impl Extend<Option<V>>,
//! values: I,
//! filter: Option<T>,
//! ) where
//! T: PartialEq,
//! for<'a> &'a V: IntoIterator<Item = &'a Option<T>>,
//! {
//! list_builder.extend(values.into_iter().filter(|string: &Option<V>| {
//! string
//! .as_ref()
//! .map(|str: &V| str.into_iter().any(|ch: &Option<T>| ch == &filter))
//! .unwrap_or(false)
//! }));
//! }
//! let builder = StringBuilder::new();
//! let mut list_builder = ListBuilder::new(builder);
//! let pear_pear = vec![Some("🍐"),Some("🍐")];
//! let pear_app = vec![Some("🍐"),Some("🍎")];
//! filter_and_fill_if_contains(
//! &mut list_builder,
//! vec![Some(pear_pear), Some(pear_app), None],
//! Some("🍎"),
//! );
//! assert_eq!(list_builder.finish().len(), 1);
//! ```
//!
//! # Custom Builders
//!
//! It is common to have a collection of statically defined Rust types that
//! you want to convert to Arrow arrays.
//!
//! An example of doing so is below
//!
//! ```
//! # use std::any::Any;
//! # use arrow_array::builder::{ArrayBuilder, Int32Builder, ListBuilder, StringBuilder};
//! # use arrow_array::{ArrayRef, RecordBatch, StructArray};
//! # use arrow_schema::{DataType, Field};
//! # use std::sync::Arc;
//! /// A custom row representation
//! struct MyRow {
//! i32: i32,
//! optional_i32: Option<i32>,
//! string: Option<String>,
//! i32_list: Option<Vec<Option<i32>>>,
//! }
//!
//! /// Converts `Vec<Row>` into `StructArray`
//! #[derive(Debug, Default)]
//! struct MyRowBuilder {
//! i32: Int32Builder,
//! string: StringBuilder,
//! i32_list: ListBuilder<Int32Builder>,
//! }
//!
//! impl MyRowBuilder {
//! fn append(&mut self, row: &MyRow) {
//! self.i32.append_value(row.i32);
//! self.string.append_option(row.string.as_ref());
//! self.i32_list.append_option(row.i32_list.as_ref().map(|x| x.iter().copied()));
//! }
//!
//! /// Note: returns StructArray to allow nesting within another array if desired
//! fn finish(&mut self) -> StructArray {
//! let i32 = Arc::new(self.i32.finish()) as ArrayRef;
//! let i32_field = Arc::new(Field::new("i32", DataType::Int32, false));
//!
//! let string = Arc::new(self.string.finish()) as ArrayRef;
//! let string_field = Arc::new(Field::new("i32", DataType::Utf8, false));
//!
//! let i32_list = Arc::new(self.i32_list.finish()) as ArrayRef;
//! let value_field = Arc::new(Field::new_list_field(DataType::Int32, true));
//! let i32_list_field = Arc::new(Field::new("i32_list", DataType::List(value_field), true));
//!
//! StructArray::from(vec![
//! (i32_field, i32),
//! (string_field, string),
//! (i32_list_field, i32_list),
//! ])
//! }
//! }
//!
//! /// For building arrays in generic code, use Extend instead of the append_* methods
//! /// e.g. append_value, append_option, append_null
//! impl<'a> Extend<&'a MyRow> for MyRowBuilder {
//! fn extend<T: IntoIterator<Item = &'a MyRow>>(&mut self, iter: T) {
//! iter.into_iter().for_each(|row| self.append(row));
//! }
//! }
//!
//! /// Converts a slice of [`MyRow`] to a [`RecordBatch`]
//! fn rows_to_batch(rows: &[MyRow]) -> RecordBatch {
//! let mut builder = MyRowBuilder::default();
//! builder.extend(rows);
//! RecordBatch::from(&builder.finish())
//! }
//! ```
//!
//! # Null / Validity Masks
//!
//! The [`NullBufferBuilder`] is optimized for creating the null mask for an array.
//!
//! ```
//! # use arrow_array::builder::NullBufferBuilder;
//! let mut builder = NullBufferBuilder::new(8);
//! let mut builder = NullBufferBuilder::new(8);
//! builder.append_n_non_nulls(7);
//! builder.append_null();
//! let buffer = builder.finish().unwrap();
//! assert_eq!(buffer.len(), 8);
//! assert_eq!(buffer.iter().collect::<Vec<_>>(), vec![true, true, true, true, true, true, true, false]);
//! ```
pub use arrow_buffer::BooleanBufferBuilder;
pub use arrow_buffer::NullBufferBuilder;
mod boolean_builder;
pub use boolean_builder::*;
mod buffer_builder;
pub use buffer_builder::*;
mod fixed_size_binary_builder;
pub use fixed_size_binary_builder::*;
mod fixed_size_list_builder;
pub use fixed_size_list_builder::*;
mod fixed_size_binary_dictionary_builder;
pub use fixed_size_binary_dictionary_builder::*;
mod generic_bytes_builder;
pub use generic_bytes_builder::*;
mod generic_list_builder;
pub use generic_list_builder::*;
mod map_builder;
pub use map_builder::*;
mod null_builder;
pub use null_builder::*;
mod primitive_builder;
pub use primitive_builder::*;
mod primitive_dictionary_builder;
pub use primitive_dictionary_builder::*;
mod primitive_run_builder;
pub use primitive_run_builder::*;
mod struct_builder;
pub use struct_builder::*;
mod generic_bytes_dictionary_builder;
pub use generic_bytes_dictionary_builder::*;
mod generic_byte_run_builder;
pub use generic_byte_run_builder::*;
mod generic_bytes_view_builder;
pub use generic_bytes_view_builder::*;
mod generic_list_view_builder;
pub use generic_list_view_builder::*;
mod union_builder;
pub use union_builder::*;
use crate::ArrayRef;
use crate::types::{Int8Type, Int16Type, Int32Type, Int64Type};
use arrow_schema::{DataType, IntervalUnit, TimeUnit};
use std::any::Any;
/// Trait for dealing with different array builders at runtime
///
/// # Example
///
/// ```
/// // Create
/// # use arrow_array::{ArrayRef, StringArray};
/// # use arrow_array::builder::{ArrayBuilder, Float64Builder, Int64Builder, StringBuilder};
///
/// let mut data_builders: Vec<Box<dyn ArrayBuilder>> = vec![
/// Box::new(Float64Builder::new()),
/// Box::new(Int64Builder::new()),
/// Box::new(StringBuilder::new()),
/// ];
///
/// // Fill
/// data_builders[0]
/// .as_any_mut()
/// .downcast_mut::<Float64Builder>()
/// .unwrap()
/// .append_value(3.14);
/// data_builders[1]
/// .as_any_mut()
/// .downcast_mut::<Int64Builder>()
/// .unwrap()
/// .append_value(-1);
/// data_builders[2]
/// .as_any_mut()
/// .downcast_mut::<StringBuilder>()
/// .unwrap()
/// .append_value("🍎");
///
/// // Finish
/// let array_refs: Vec<ArrayRef> = data_builders
/// .iter_mut()
/// .map(|builder| builder.finish())
/// .collect();
/// assert_eq!(array_refs[0].len(), 1);
/// assert_eq!(array_refs[1].is_null(0), false);
/// assert_eq!(
/// array_refs[2]
/// .as_any()
/// .downcast_ref::<StringArray>()
/// .unwrap()
/// .value(0),
/// "🍎"
/// );
/// ```
pub trait ArrayBuilder: Any + Send + Sync {
/// Returns the number of array slots in the builder
fn len(&self) -> usize;
/// Returns whether number of array slots is zero
fn is_empty(&self) -> bool {
self.len() == 0
}
/// Builds the array
fn finish(&mut self) -> ArrayRef;
/// Builds the array without resetting the underlying builder.
fn finish_cloned(&self) -> ArrayRef;
/// Returns the builder as a non-mutable `Any` reference.
///
/// This is most useful when one wants to call non-mutable APIs on a specific builder
/// type. In this case, one can first cast this into a `Any`, and then use
/// `downcast_ref` to get a reference on the specific builder.
fn as_any(&self) -> &dyn Any;
/// Returns the builder as a mutable `Any` reference.
///
/// This is most useful when one wants to call mutable APIs on a specific builder
/// type. In this case, one can first cast this into a `Any`, and then use
/// `downcast_mut` to get a reference on the specific builder.
fn as_any_mut(&mut self) -> &mut dyn Any;
/// Returns the boxed builder as a box of `Any`.
fn into_box_any(self: Box<Self>) -> Box<dyn Any>;
}
impl ArrayBuilder for Box<dyn ArrayBuilder> {
fn len(&self) -> usize {
(**self).len()
}
fn is_empty(&self) -> bool {
(**self).is_empty()
}
fn finish(&mut self) -> ArrayRef {
(**self).finish()
}
fn finish_cloned(&self) -> ArrayRef {
(**self).finish_cloned()
}
fn as_any(&self) -> &dyn Any {
(**self).as_any()
}
fn as_any_mut(&mut self) -> &mut dyn Any {
(**self).as_any_mut()
}
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
}
/// Builder for [`ListArray`](crate::array::ListArray)
pub type ListBuilder<T> = GenericListBuilder<i32, T>;
/// Builder for [`LargeListArray`](crate::array::LargeListArray)
pub type LargeListBuilder<T> = GenericListBuilder<i64, T>;
/// Builder for [`ListViewArray`](crate::array::ListViewArray)
pub type ListViewBuilder<T> = GenericListViewBuilder<i32, T>;
/// Builder for [`LargeListViewArray`](crate::array::LargeListViewArray)
pub type LargeListViewBuilder<T> = GenericListViewBuilder<i64, T>;
/// Builder for [`BinaryArray`](crate::array::BinaryArray)
///
/// See examples on [`GenericBinaryBuilder`]
pub type BinaryBuilder = GenericBinaryBuilder<i32>;
/// Builder for [`LargeBinaryArray`](crate::array::LargeBinaryArray)
///
/// See examples on [`GenericBinaryBuilder`]
pub type LargeBinaryBuilder = GenericBinaryBuilder<i64>;
/// Builder for [`StringArray`](crate::array::StringArray)
///
/// See examples on [`GenericStringBuilder`]
pub type StringBuilder = GenericStringBuilder<i32>;
/// Builder for [`LargeStringArray`](crate::array::LargeStringArray)
///
/// See examples on [`GenericStringBuilder`]
pub type LargeStringBuilder = GenericStringBuilder<i64>;
/// Returns a builder with capacity for `capacity` elements of datatype
/// `DataType`.
///
/// This function is useful to construct arrays from an arbitrary vectors with
/// known/expected schema.
///
/// See comments on [StructBuilder] for retrieving collection builders built by
/// make_builder.
pub fn make_builder(datatype: &DataType, capacity: usize) -> Box<dyn ArrayBuilder> {
use crate::builder::*;
match datatype {
DataType::Null => Box::new(NullBuilder::new()),
DataType::Boolean => Box::new(BooleanBuilder::with_capacity(capacity)),
DataType::Int8 => Box::new(Int8Builder::with_capacity(capacity)),
DataType::Int16 => Box::new(Int16Builder::with_capacity(capacity)),
DataType::Int32 => Box::new(Int32Builder::with_capacity(capacity)),
DataType::Int64 => Box::new(Int64Builder::with_capacity(capacity)),
DataType::UInt8 => Box::new(UInt8Builder::with_capacity(capacity)),
DataType::UInt16 => Box::new(UInt16Builder::with_capacity(capacity)),
DataType::UInt32 => Box::new(UInt32Builder::with_capacity(capacity)),
DataType::UInt64 => Box::new(UInt64Builder::with_capacity(capacity)),
DataType::Float16 => Box::new(Float16Builder::with_capacity(capacity)),
DataType::Float32 => Box::new(Float32Builder::with_capacity(capacity)),
DataType::Float64 => Box::new(Float64Builder::with_capacity(capacity)),
DataType::Binary => Box::new(BinaryBuilder::with_capacity(capacity, 1024)),
DataType::LargeBinary => Box::new(LargeBinaryBuilder::with_capacity(capacity, 1024)),
DataType::BinaryView => Box::new(BinaryViewBuilder::with_capacity(capacity)),
DataType::FixedSizeBinary(len) => {
Box::new(FixedSizeBinaryBuilder::with_capacity(capacity, *len))
}
DataType::Decimal32(p, s) => Box::new(
Decimal32Builder::with_capacity(capacity).with_data_type(DataType::Decimal32(*p, *s)),
),
DataType::Decimal64(p, s) => Box::new(
Decimal64Builder::with_capacity(capacity).with_data_type(DataType::Decimal64(*p, *s)),
),
DataType::Decimal128(p, s) => Box::new(
Decimal128Builder::with_capacity(capacity).with_data_type(DataType::Decimal128(*p, *s)),
),
DataType::Decimal256(p, s) => Box::new(
Decimal256Builder::with_capacity(capacity).with_data_type(DataType::Decimal256(*p, *s)),
),
DataType::Utf8 => Box::new(StringBuilder::with_capacity(capacity, 1024)),
DataType::LargeUtf8 => Box::new(LargeStringBuilder::with_capacity(capacity, 1024)),
DataType::Utf8View => Box::new(StringViewBuilder::with_capacity(capacity)),
DataType::Date32 => Box::new(Date32Builder::with_capacity(capacity)),
DataType::Date64 => Box::new(Date64Builder::with_capacity(capacity)),
DataType::Time32(TimeUnit::Second) => {
Box::new(Time32SecondBuilder::with_capacity(capacity))
}
DataType::Time32(TimeUnit::Millisecond) => {
Box::new(Time32MillisecondBuilder::with_capacity(capacity))
}
DataType::Time64(TimeUnit::Microsecond) => {
Box::new(Time64MicrosecondBuilder::with_capacity(capacity))
}
DataType::Time64(TimeUnit::Nanosecond) => {
Box::new(Time64NanosecondBuilder::with_capacity(capacity))
}
DataType::Timestamp(TimeUnit::Second, tz) => Box::new(
TimestampSecondBuilder::with_capacity(capacity)
.with_data_type(DataType::Timestamp(TimeUnit::Second, tz.clone())),
),
DataType::Timestamp(TimeUnit::Millisecond, tz) => Box::new(
TimestampMillisecondBuilder::with_capacity(capacity)
.with_data_type(DataType::Timestamp(TimeUnit::Millisecond, tz.clone())),
),
DataType::Timestamp(TimeUnit::Microsecond, tz) => Box::new(
TimestampMicrosecondBuilder::with_capacity(capacity)
.with_data_type(DataType::Timestamp(TimeUnit::Microsecond, tz.clone())),
),
DataType::Timestamp(TimeUnit::Nanosecond, tz) => Box::new(
TimestampNanosecondBuilder::with_capacity(capacity)
.with_data_type(DataType::Timestamp(TimeUnit::Nanosecond, tz.clone())),
),
DataType::Interval(IntervalUnit::YearMonth) => {
Box::new(IntervalYearMonthBuilder::with_capacity(capacity))
}
DataType::Interval(IntervalUnit::DayTime) => {
Box::new(IntervalDayTimeBuilder::with_capacity(capacity))
}
DataType::Interval(IntervalUnit::MonthDayNano) => {
Box::new(IntervalMonthDayNanoBuilder::with_capacity(capacity))
}
DataType::Duration(TimeUnit::Second) => {
Box::new(DurationSecondBuilder::with_capacity(capacity))
}
DataType::Duration(TimeUnit::Millisecond) => {
Box::new(DurationMillisecondBuilder::with_capacity(capacity))
}
DataType::Duration(TimeUnit::Microsecond) => {
Box::new(DurationMicrosecondBuilder::with_capacity(capacity))
}
DataType::Duration(TimeUnit::Nanosecond) => {
Box::new(DurationNanosecondBuilder::with_capacity(capacity))
}
DataType::List(field) => {
let builder = make_builder(field.data_type(), capacity);
Box::new(ListBuilder::with_capacity(builder, capacity).with_field(field.clone()))
}
DataType::LargeList(field) => {
let builder = make_builder(field.data_type(), capacity);
Box::new(LargeListBuilder::with_capacity(builder, capacity).with_field(field.clone()))
}
DataType::FixedSizeList(field, size) => {
let size = *size;
let values_builder_capacity = {
let size: usize = size.try_into().unwrap();
capacity * size
};
let builder = make_builder(field.data_type(), values_builder_capacity);
Box::new(
FixedSizeListBuilder::with_capacity(builder, size, capacity)
.with_field(field.clone()),
)
}
DataType::ListView(field) => {
let builder = make_builder(field.data_type(), capacity);
Box::new(ListViewBuilder::with_capacity(builder, capacity).with_field(field.clone()))
}
DataType::LargeListView(field) => {
let builder = make_builder(field.data_type(), capacity);
Box::new(
LargeListViewBuilder::with_capacity(builder, capacity).with_field(field.clone()),
)
}
DataType::Map(field, _) => match field.data_type() {
DataType::Struct(fields) => {
let map_field_names = MapFieldNames {
key: fields[0].name().clone(),
value: fields[1].name().clone(),
entry: field.name().clone(),
};
let key_builder = make_builder(fields[0].data_type(), capacity);
let value_builder = make_builder(fields[1].data_type(), capacity);
Box::new(
MapBuilder::with_capacity(
Some(map_field_names),
key_builder,
value_builder,
capacity,
)
.with_keys_field(fields[0].clone())
.with_values_field(fields[1].clone()),
)
}
t => panic!("The field of Map data type {t} should have a child Struct field"),
},
DataType::Struct(fields) => Box::new(StructBuilder::from_fields(fields.clone(), capacity)),
t @ DataType::Dictionary(key_type, value_type) => {
macro_rules! dict_builder {
($key_type:ty) => {
match &**value_type {
DataType::Utf8 => {
let dict_builder: StringDictionaryBuilder<$key_type> =
StringDictionaryBuilder::with_capacity(capacity, 256, 1024);
Box::new(dict_builder)
}
DataType::LargeUtf8 => {
let dict_builder: LargeStringDictionaryBuilder<$key_type> =
LargeStringDictionaryBuilder::with_capacity(capacity, 256, 1024);
Box::new(dict_builder)
}
DataType::Binary => {
let dict_builder: BinaryDictionaryBuilder<$key_type> =
BinaryDictionaryBuilder::with_capacity(capacity, 256, 1024);
Box::new(dict_builder)
}
DataType::LargeBinary => {
let dict_builder: LargeBinaryDictionaryBuilder<$key_type> =
LargeBinaryDictionaryBuilder::with_capacity(capacity, 256, 1024);
Box::new(dict_builder)
}
t => unimplemented!("Dictionary value type {t} is not currently supported"),
}
};
}
match &**key_type {
DataType::Int8 => dict_builder!(Int8Type),
DataType::Int16 => dict_builder!(Int16Type),
DataType::Int32 => dict_builder!(Int32Type),
DataType::Int64 => dict_builder!(Int64Type),
_ => {
unimplemented!(
"Data type {t} with key type {key_type} is not currently supported"
)
}
}
}
t => unimplemented!("Data type {t} is not currently supported"),
}
}
@@ -0,0 +1,170 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::builder::ArrayBuilder;
use crate::{ArrayRef, NullArray};
use arrow_data::ArrayData;
use arrow_schema::DataType;
use std::any::Any;
use std::sync::Arc;
/// Builder for [`NullArray`]
///
/// # Example
///
/// Create a `NullArray` from a `NullBuilder`
///
/// ```
///
/// # use arrow_array::{Array, NullArray, builder::NullBuilder};
///
/// let mut b = NullBuilder::new();
/// b.append_empty_value();
/// b.append_null();
/// b.append_nulls(3);
/// b.append_empty_values(3);
/// let arr = b.finish();
///
/// assert_eq!(8, arr.len());
/// assert_eq!(0, arr.null_count());
/// ```
#[derive(Debug)]
pub struct NullBuilder {
len: usize,
}
impl Default for NullBuilder {
fn default() -> Self {
Self::new()
}
}
impl NullBuilder {
/// Creates a new null builder
pub fn new() -> Self {
Self { len: 0 }
}
/// Appends a null slot into the builder
#[inline]
pub fn append_null(&mut self) {
self.len += 1;
}
/// Appends `n` `null`s into the builder.
#[inline]
pub fn append_nulls(&mut self, n: usize) {
self.len += n;
}
/// Appends a null slot into the builder
#[inline]
pub fn append_empty_value(&mut self) {
self.append_null();
}
/// Appends `n` `null`s into the builder.
#[inline]
pub fn append_empty_values(&mut self, n: usize) {
self.append_nulls(n);
}
/// Builds the [NullArray] and reset this builder.
pub fn finish(&mut self) -> NullArray {
let len = self.len();
let builder = ArrayData::new_null(&DataType::Null, len).into_builder();
let array_data = unsafe { builder.build_unchecked() };
NullArray::from(array_data)
}
/// Builds the [NullArray] without resetting the builder.
pub fn finish_cloned(&self) -> NullArray {
let len = self.len();
let builder = ArrayData::new_null(&DataType::Null, len).into_builder();
let array_data = unsafe { builder.build_unchecked() };
NullArray::from(array_data)
}
}
impl ArrayBuilder for NullBuilder {
/// Returns the builder as a non-mutable `Any` reference.
fn as_any(&self) -> &dyn Any {
self
}
/// Returns the builder as a mutable `Any` reference.
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
/// Returns the boxed builder as a box of `Any`.
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
/// Returns the number of array slots in the builder
fn len(&self) -> usize {
self.len
}
/// Builds the array and reset this builder.
fn finish(&mut self) -> ArrayRef {
Arc::new(self.finish())
}
/// Builds the array without resetting the builder.
fn finish_cloned(&self) -> ArrayRef {
Arc::new(self.finish_cloned())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Array;
#[test]
fn test_null_array_builder() {
let mut builder = NullArray::builder(10);
builder.append_null();
builder.append_nulls(4);
builder.append_empty_value();
builder.append_empty_values(4);
let arr = builder.finish();
assert_eq!(10, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(0, arr.null_count());
assert!(arr.is_nullable());
}
#[test]
fn test_null_array_builder_finish_cloned() {
let mut builder = NullArray::builder(16);
builder.append_null();
builder.append_empty_value();
builder.append_empty_values(3);
let mut array = builder.finish_cloned();
assert_eq!(5, array.len());
builder.append_empty_values(5);
array = builder.finish();
assert_eq!(10, array.len());
}
}
@@ -0,0 +1,705 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::builder::ArrayBuilder;
use crate::types::*;
use crate::{Array, ArrayRef, PrimitiveArray};
use arrow_buffer::{Buffer, MutableBuffer, NullBufferBuilder, ScalarBuffer};
use arrow_data::ArrayData;
use arrow_schema::{ArrowError, DataType};
use std::any::Any;
use std::sync::Arc;
/// A signed 8-bit integer array builder.
pub type Int8Builder = PrimitiveBuilder<Int8Type>;
/// A signed 16-bit integer array builder.
pub type Int16Builder = PrimitiveBuilder<Int16Type>;
/// A signed 32-bit integer array builder.
pub type Int32Builder = PrimitiveBuilder<Int32Type>;
/// A signed 64-bit integer array builder.
pub type Int64Builder = PrimitiveBuilder<Int64Type>;
/// An usigned 8-bit integer array builder.
pub type UInt8Builder = PrimitiveBuilder<UInt8Type>;
/// An usigned 16-bit integer array builder.
pub type UInt16Builder = PrimitiveBuilder<UInt16Type>;
/// An usigned 32-bit integer array builder.
pub type UInt32Builder = PrimitiveBuilder<UInt32Type>;
/// An usigned 64-bit integer array builder.
pub type UInt64Builder = PrimitiveBuilder<UInt64Type>;
/// A 16-bit floating point array builder.
pub type Float16Builder = PrimitiveBuilder<Float16Type>;
/// A 32-bit floating point array builder.
pub type Float32Builder = PrimitiveBuilder<Float32Type>;
/// A 64-bit floating point array builder.
pub type Float64Builder = PrimitiveBuilder<Float64Type>;
/// A timestamp second array builder.
pub type TimestampSecondBuilder = PrimitiveBuilder<TimestampSecondType>;
/// A timestamp millisecond array builder.
pub type TimestampMillisecondBuilder = PrimitiveBuilder<TimestampMillisecondType>;
/// A timestamp microsecond array builder.
pub type TimestampMicrosecondBuilder = PrimitiveBuilder<TimestampMicrosecondType>;
/// A timestamp nanosecond array builder.
pub type TimestampNanosecondBuilder = PrimitiveBuilder<TimestampNanosecondType>;
/// A 32-bit date array builder.
pub type Date32Builder = PrimitiveBuilder<Date32Type>;
/// A 64-bit date array builder.
pub type Date64Builder = PrimitiveBuilder<Date64Type>;
/// A 32-bit elaspsed time in seconds array builder.
pub type Time32SecondBuilder = PrimitiveBuilder<Time32SecondType>;
/// A 32-bit elaspsed time in milliseconds array builder.
pub type Time32MillisecondBuilder = PrimitiveBuilder<Time32MillisecondType>;
/// A 64-bit elaspsed time in microseconds array builder.
pub type Time64MicrosecondBuilder = PrimitiveBuilder<Time64MicrosecondType>;
/// A 64-bit elaspsed time in nanoseconds array builder.
pub type Time64NanosecondBuilder = PrimitiveBuilder<Time64NanosecondType>;
/// A “calendar” interval in months array builder.
pub type IntervalYearMonthBuilder = PrimitiveBuilder<IntervalYearMonthType>;
/// A “calendar” interval in days and milliseconds array builder.
pub type IntervalDayTimeBuilder = PrimitiveBuilder<IntervalDayTimeType>;
/// A “calendar” interval in months, days, and nanoseconds array builder.
pub type IntervalMonthDayNanoBuilder = PrimitiveBuilder<IntervalMonthDayNanoType>;
/// An elapsed time in seconds array builder.
pub type DurationSecondBuilder = PrimitiveBuilder<DurationSecondType>;
/// An elapsed time in milliseconds array builder.
pub type DurationMillisecondBuilder = PrimitiveBuilder<DurationMillisecondType>;
/// An elapsed time in microseconds array builder.
pub type DurationMicrosecondBuilder = PrimitiveBuilder<DurationMicrosecondType>;
/// An elapsed time in nanoseconds array builder.
pub type DurationNanosecondBuilder = PrimitiveBuilder<DurationNanosecondType>;
/// A decimal 32 array builder
pub type Decimal32Builder = PrimitiveBuilder<Decimal32Type>;
/// A decimal 64 array builder
pub type Decimal64Builder = PrimitiveBuilder<Decimal64Type>;
/// A decimal 128 array builder
pub type Decimal128Builder = PrimitiveBuilder<Decimal128Type>;
/// A decimal 256 array builder
pub type Decimal256Builder = PrimitiveBuilder<Decimal256Type>;
/// Builder for [`PrimitiveArray`]
#[derive(Debug)]
pub struct PrimitiveBuilder<T: ArrowPrimitiveType> {
values_builder: Vec<T::Native>,
null_buffer_builder: NullBufferBuilder,
data_type: DataType,
}
impl<T: ArrowPrimitiveType> ArrayBuilder for PrimitiveBuilder<T> {
/// Returns the builder as a non-mutable `Any` reference.
fn as_any(&self) -> &dyn Any {
self
}
/// Returns the builder as a mutable `Any` reference.
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
/// Returns the boxed builder as a box of `Any`.
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
/// Returns the number of array slots in the builder
fn len(&self) -> usize {
self.values_builder.len()
}
/// Builds the array and reset this builder.
fn finish(&mut self) -> ArrayRef {
Arc::new(self.finish())
}
/// Builds the array without resetting the builder.
fn finish_cloned(&self) -> ArrayRef {
Arc::new(self.finish_cloned())
}
}
impl<T: ArrowPrimitiveType> Default for PrimitiveBuilder<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
/// Creates a new primitive array builder
pub fn new() -> Self {
Self::with_capacity(1024)
}
/// Creates a new primitive array builder with capacity no of items
pub fn with_capacity(capacity: usize) -> Self {
Self {
values_builder: Vec::with_capacity(capacity),
null_buffer_builder: NullBufferBuilder::new(capacity),
data_type: T::DATA_TYPE,
}
}
/// Creates a new primitive array builder from buffers
pub fn new_from_buffer(
values_buffer: MutableBuffer,
null_buffer: Option<MutableBuffer>,
) -> Self {
let values_builder: Vec<T::Native> = ScalarBuffer::<T::Native>::from(values_buffer).into();
let null_buffer_builder = null_buffer
.map(|buffer| NullBufferBuilder::new_from_buffer(buffer, values_builder.len()))
.unwrap_or_else(|| NullBufferBuilder::new_with_len(values_builder.len()));
Self {
values_builder,
null_buffer_builder,
data_type: T::DATA_TYPE,
}
}
/// By default [`PrimitiveBuilder`] uses [`ArrowPrimitiveType::DATA_TYPE`] as the
/// data type of the generated array.
///
/// This method allows overriding the data type, to allow specifying timezones
/// for [`DataType::Timestamp`] or precision and scale for [`DataType::Decimal32`],
/// [`DataType::Decimal64`], [`DataType::Decimal128`] and [`DataType::Decimal256`]
///
/// # Panics
///
/// This method panics if `data_type` is not [PrimitiveArray::is_compatible]
pub fn with_data_type(self, data_type: DataType) -> Self {
assert!(
PrimitiveArray::<T>::is_compatible(&data_type),
"incompatible data type for builder, expected {} got {}",
T::DATA_TYPE,
data_type
);
Self { data_type, ..self }
}
/// Returns the capacity of this builder measured in slots of type `T`
pub fn capacity(&self) -> usize {
self.values_builder.capacity()
}
/// Appends a value of type `T` into the builder
#[inline]
pub fn append_value(&mut self, v: T::Native) {
self.null_buffer_builder.append_non_null();
self.values_builder.push(v);
}
/// Appends a value of type `T` into the builder `n` times
#[inline]
pub fn append_value_n(&mut self, v: T::Native, n: usize) {
self.null_buffer_builder.append_n_non_nulls(n);
self.values_builder.extend(std::iter::repeat_n(v, n));
}
/// Appends a null slot into the builder
#[inline]
pub fn append_null(&mut self) {
self.null_buffer_builder.append_null();
self.values_builder.push(T::Native::default());
}
/// Appends `n` no. of null's into the builder
#[inline]
pub fn append_nulls(&mut self, n: usize) {
self.null_buffer_builder.append_n_nulls(n);
self.values_builder
.extend(std::iter::repeat_n(T::Native::default(), n));
}
/// Appends an `Option<T>` into the builder
#[inline]
pub fn append_option(&mut self, v: Option<T::Native>) {
match v {
None => self.append_null(),
Some(v) => self.append_value(v),
};
}
/// Appends a slice of type `T` into the builder
#[inline]
pub fn append_slice(&mut self, v: &[T::Native]) {
self.null_buffer_builder.append_n_non_nulls(v.len());
self.values_builder.extend_from_slice(v);
}
/// Appends values from a slice of type `T` and a validity boolean slice
///
/// # Panics
///
/// Panics if `values` and `is_valid` have different lengths
#[inline]
pub fn append_values(&mut self, values: &[T::Native], is_valid: &[bool]) {
assert_eq!(
values.len(),
is_valid.len(),
"Value and validity lengths must be equal"
);
self.null_buffer_builder.append_slice(is_valid);
self.values_builder.extend_from_slice(values);
}
/// Appends array values and null to this builder as is
/// (this means that underlying null values are copied as is).
///
/// # Panics
///
/// Panics if `array` and `self` data types are different
#[inline]
pub fn append_array(&mut self, array: &PrimitiveArray<T>) {
assert_eq!(
&self.data_type,
array.data_type(),
"array data type mismatch"
);
self.values_builder.extend_from_slice(array.values());
if let Some(null_buffer) = array.nulls() {
self.null_buffer_builder.append_buffer(null_buffer);
} else {
self.null_buffer_builder.append_n_non_nulls(array.len());
}
}
/// Appends values from a trusted length iterator.
///
/// # Safety
/// This requires the iterator be a trusted length. This could instead require
/// the iterator implement `TrustedLen` once that is stabilized.
#[inline]
pub unsafe fn append_trusted_len_iter(&mut self, iter: impl IntoIterator<Item = T::Native>) {
let iter = iter.into_iter();
let len = iter
.size_hint()
.1
.expect("append_trusted_len_iter requires an upper bound");
self.null_buffer_builder.append_n_non_nulls(len);
self.values_builder.extend(iter);
}
/// Builds the [`PrimitiveArray`] and reset this builder.
pub fn finish(&mut self) -> PrimitiveArray<T> {
let len = self.len();
let nulls = self.null_buffer_builder.finish();
let builder = ArrayData::builder(self.data_type.clone())
.len(len)
.add_buffer(std::mem::take(&mut self.values_builder).into())
.nulls(nulls);
let array_data = unsafe { builder.build_unchecked() };
PrimitiveArray::<T>::from(array_data)
}
/// Builds the [`PrimitiveArray`] without resetting the builder.
pub fn finish_cloned(&self) -> PrimitiveArray<T> {
let len = self.len();
let nulls = self.null_buffer_builder.finish_cloned();
let values_buffer = Buffer::from_slice_ref(self.values_builder.as_slice());
let builder = ArrayData::builder(self.data_type.clone())
.len(len)
.add_buffer(values_buffer)
.nulls(nulls);
let array_data = unsafe { builder.build_unchecked() };
PrimitiveArray::<T>::from(array_data)
}
/// Returns the current values buffer as a slice
pub fn values_slice(&self) -> &[T::Native] {
self.values_builder.as_slice()
}
/// Returns the current values buffer as a mutable slice
pub fn values_slice_mut(&mut self) -> &mut [T::Native] {
self.values_builder.as_mut_slice()
}
/// Returns the current null buffer as a slice
pub fn validity_slice(&self) -> Option<&[u8]> {
self.null_buffer_builder.as_slice()
}
/// Returns the current null buffer as a mutable slice
pub fn validity_slice_mut(&mut self) -> Option<&mut [u8]> {
self.null_buffer_builder.as_slice_mut()
}
/// Returns the current values buffer and null buffer as a slice
pub fn slices_mut(&mut self) -> (&mut [T::Native], Option<&mut [u8]>) {
(
self.values_builder.as_mut_slice(),
self.null_buffer_builder.as_slice_mut(),
)
}
}
impl<P: DecimalType> PrimitiveBuilder<P> {
/// Sets the precision and scale
pub fn with_precision_and_scale(self, precision: u8, scale: i8) -> Result<Self, ArrowError> {
validate_decimal_precision_and_scale::<P>(precision, scale)?;
Ok(Self {
data_type: P::TYPE_CONSTRUCTOR(precision, scale),
..self
})
}
}
impl<P: ArrowTimestampType> PrimitiveBuilder<P> {
/// Sets the timezone
pub fn with_timezone(self, timezone: impl Into<Arc<str>>) -> Self {
self.with_timezone_opt(Some(timezone.into()))
}
/// Sets an optional timezone
pub fn with_timezone_opt<S: Into<Arc<str>>>(self, timezone: Option<S>) -> Self {
Self {
data_type: DataType::Timestamp(P::UNIT, timezone.map(Into::into)),
..self
}
}
}
impl<P: ArrowPrimitiveType> Extend<Option<P::Native>> for PrimitiveBuilder<P> {
#[inline]
fn extend<T: IntoIterator<Item = Option<P::Native>>>(&mut self, iter: T) {
for v in iter {
self.append_option(v)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use arrow_buffer::{NullBuffer, ScalarBuffer};
use arrow_schema::TimeUnit;
use crate::array::Array;
use crate::array::BooleanArray;
use crate::array::Date32Array;
use crate::array::Int32Array;
use crate::array::TimestampSecondArray;
#[test]
fn test_primitive_array_builder_i32() {
let mut builder = Int32Array::builder(5);
for i in 0..5 {
builder.append_value(i);
}
let arr = builder.finish();
assert_eq!(5, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(0, arr.null_count());
for i in 0..5 {
assert!(!arr.is_null(i));
assert!(arr.is_valid(i));
assert_eq!(i as i32, arr.value(i));
}
}
#[test]
fn test_primitive_array_builder_i32_append_iter() {
let mut builder = Int32Array::builder(5);
unsafe { builder.append_trusted_len_iter(0..5) };
let arr = builder.finish();
assert_eq!(5, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(0, arr.null_count());
for i in 0..5 {
assert!(!arr.is_null(i));
assert!(arr.is_valid(i));
assert_eq!(i as i32, arr.value(i));
}
}
#[test]
fn test_primitive_array_builder_i32_append_nulls() {
let mut builder = Int32Array::builder(5);
builder.append_nulls(5);
let arr = builder.finish();
assert_eq!(5, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(5, arr.null_count());
for i in 0..5 {
assert!(arr.is_null(i));
assert!(!arr.is_valid(i));
}
}
#[test]
fn test_primitive_array_builder_date32() {
let mut builder = Date32Array::builder(5);
for i in 0..5 {
builder.append_value(i);
}
let arr = builder.finish();
assert_eq!(5, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(0, arr.null_count());
for i in 0..5 {
assert!(!arr.is_null(i));
assert!(arr.is_valid(i));
assert_eq!(i as i32, arr.value(i));
}
}
#[test]
fn test_primitive_array_builder_timestamp_second() {
let mut builder = TimestampSecondArray::builder(5);
for i in 0..5 {
builder.append_value(i);
}
let arr = builder.finish();
assert_eq!(5, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(0, arr.null_count());
for i in 0..5 {
assert!(!arr.is_null(i));
assert!(arr.is_valid(i));
assert_eq!(i as i64, arr.value(i));
}
}
#[test]
fn test_primitive_array_builder_bool() {
// 00000010 01001000
let buf = Buffer::from([72_u8, 2_u8]);
let mut builder = BooleanArray::builder(10);
for i in 0..10 {
if i == 3 || i == 6 || i == 9 {
builder.append_value(true);
} else {
builder.append_value(false);
}
}
let arr = builder.finish();
assert_eq!(&buf, arr.values().inner());
assert_eq!(10, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(0, arr.null_count());
for i in 0..10 {
assert!(!arr.is_null(i));
assert!(arr.is_valid(i));
assert_eq!(i == 3 || i == 6 || i == 9, arr.value(i), "failed at {i}")
}
}
#[test]
fn test_primitive_array_builder_append_option() {
let arr1 = Int32Array::from(vec![Some(0), None, Some(2), None, Some(4)]);
let mut builder = Int32Array::builder(5);
builder.append_option(Some(0));
builder.append_option(None);
builder.append_option(Some(2));
builder.append_option(None);
builder.append_option(Some(4));
let arr2 = builder.finish();
assert_eq!(arr1.len(), arr2.len());
assert_eq!(arr1.offset(), arr2.offset());
assert_eq!(arr1.null_count(), arr2.null_count());
for i in 0..5 {
assert_eq!(arr1.is_null(i), arr2.is_null(i));
assert_eq!(arr1.is_valid(i), arr2.is_valid(i));
if arr1.is_valid(i) {
assert_eq!(arr1.value(i), arr2.value(i));
}
}
}
#[test]
fn test_primitive_array_builder_append_null() {
let arr1 = Int32Array::from(vec![Some(0), Some(2), None, None, Some(4)]);
let mut builder = Int32Array::builder(5);
builder.append_value(0);
builder.append_value(2);
builder.append_null();
builder.append_null();
builder.append_value(4);
let arr2 = builder.finish();
assert_eq!(arr1.len(), arr2.len());
assert_eq!(arr1.offset(), arr2.offset());
assert_eq!(arr1.null_count(), arr2.null_count());
for i in 0..5 {
assert_eq!(arr1.is_null(i), arr2.is_null(i));
assert_eq!(arr1.is_valid(i), arr2.is_valid(i));
if arr1.is_valid(i) {
assert_eq!(arr1.value(i), arr2.value(i));
}
}
}
#[test]
fn test_primitive_array_builder_append_slice() {
let arr1 = Int32Array::from(vec![Some(0), Some(2), None, None, Some(4)]);
let mut builder = Int32Array::builder(5);
builder.append_slice(&[0, 2]);
builder.append_null();
builder.append_null();
builder.append_value(4);
let arr2 = builder.finish();
assert_eq!(arr1.len(), arr2.len());
assert_eq!(arr1.offset(), arr2.offset());
assert_eq!(arr1.null_count(), arr2.null_count());
for i in 0..5 {
assert_eq!(arr1.is_null(i), arr2.is_null(i));
assert_eq!(arr1.is_valid(i), arr2.is_valid(i));
if arr1.is_valid(i) {
assert_eq!(arr1.value(i), arr2.value(i));
}
}
}
#[test]
fn test_primitive_array_builder_finish() {
let mut builder = Int32Builder::new();
builder.append_slice(&[2, 4, 6, 8]);
let mut arr = builder.finish();
assert_eq!(4, arr.len());
assert_eq!(0, builder.len());
builder.append_slice(&[1, 3, 5, 7, 9]);
arr = builder.finish();
assert_eq!(5, arr.len());
assert_eq!(0, builder.len());
}
#[test]
fn test_primitive_array_builder_finish_cloned() {
let mut builder = Int32Builder::new();
builder.append_value(23);
builder.append_value(45);
let result = builder.finish_cloned();
assert_eq!(result, Int32Array::from(vec![23, 45]));
builder.append_value(56);
assert_eq!(builder.finish_cloned(), Int32Array::from(vec![23, 45, 56]));
builder.append_slice(&[2, 4, 6, 8]);
let mut arr = builder.finish();
assert_eq!(7, arr.len());
assert_eq!(arr, Int32Array::from(vec![23, 45, 56, 2, 4, 6, 8]));
assert_eq!(0, builder.len());
builder.append_slice(&[1, 3, 5, 7, 9]);
arr = builder.finish();
assert_eq!(5, arr.len());
assert_eq!(0, builder.len());
}
#[test]
fn test_primitive_array_builder_with_data_type() {
let mut builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(1, 2));
builder.append_value(1);
let array = builder.finish();
assert_eq!(array.precision(), 1);
assert_eq!(array.scale(), 2);
let data_type = DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into()));
let mut builder = TimestampNanosecondBuilder::new().with_data_type(data_type.clone());
builder.append_value(1);
let array = builder.finish();
assert_eq!(array.data_type(), &data_type);
}
#[test]
#[should_panic(expected = "incompatible data type for builder, expected Int32 got Int64")]
fn test_invalid_with_data_type() {
Int32Builder::new().with_data_type(DataType::Int64);
}
#[test]
fn test_extend() {
let mut builder = PrimitiveBuilder::<Int16Type>::new();
builder.extend([1, 2, 3, 5, 2, 4, 4].into_iter().map(Some));
builder.extend([2, 4, 6, 2].into_iter().map(Some));
let array = builder.finish();
assert_eq!(array.values(), &[1, 2, 3, 5, 2, 4, 4, 2, 4, 6, 2]);
}
#[test]
fn test_primitive_array_append_array() {
let input = vec![
Some(1),
None,
Some(3),
None,
Some(5),
None,
None,
None,
Some(7),
Some(9),
Some(8),
Some(6),
Some(4),
];
let arr1 = Int32Array::from(input[..5].to_vec());
let arr2 = Int32Array::from(input[5..8].to_vec());
let arr3 = Int32Array::from(input[8..].to_vec());
let mut builder = Int32Array::builder(5);
builder.append_array(&arr1);
builder.append_array(&arr2);
builder.append_array(&arr3);
let actual = builder.finish();
let expected = Int32Array::from(input);
assert_eq!(actual, expected);
}
#[test]
fn test_append_array_add_underlying_null_values() {
let array = Int32Array::new(
ScalarBuffer::from(vec![2, 3, 4, 5]),
Some(NullBuffer::from(&[true, true, false, false])),
);
let mut builder = Int32Array::builder(5);
builder.append_array(&array);
let actual = builder.finish();
assert_eq!(actual, array);
assert_eq!(actual.values(), array.values())
}
#[test]
#[should_panic(expected = "array data type mismatch")]
fn test_invalid_with_data_type_in_append_array() {
let array = {
let mut builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(1, 2));
builder.append_value(1);
builder.finish()
};
let mut builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(2, 3));
builder.append_array(&array)
}
}
@@ -0,0 +1,893 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::builder::{ArrayBuilder, PrimitiveBuilder};
use crate::types::ArrowDictionaryKeyType;
use crate::{
Array, ArrayRef, ArrowPrimitiveType, DictionaryArray, PrimitiveArray, TypedDictionaryArray,
};
use arrow_buffer::{ArrowNativeType, ToByteSlice};
use arrow_schema::{ArrowError, DataType};
use num_traits::NumCast;
use std::any::Any;
use std::collections::HashMap;
use std::sync::Arc;
/// Wraps a type implementing `ToByteSlice` implementing `Hash` and `Eq` for it
///
/// This is necessary to handle types such as f32, which don't natively implement these
#[derive(Debug)]
struct Value<T>(T);
impl<T: ToByteSlice> std::hash::Hash for Value<T> {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.0.to_byte_slice().hash(state)
}
}
impl<T: ToByteSlice> PartialEq for Value<T> {
fn eq(&self, other: &Self) -> bool {
self.0.to_byte_slice().eq(other.0.to_byte_slice())
}
}
impl<T: ToByteSlice> Eq for Value<T> {}
/// Builder for [`DictionaryArray`] of [`PrimitiveArray`]
///
/// # Example:
///
/// ```
///
/// # use arrow_array::builder::PrimitiveDictionaryBuilder;
/// # use arrow_array::types::{UInt32Type, UInt8Type};
/// # use arrow_array::{Array, UInt32Array, UInt8Array};
///
/// let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::new();
/// builder.append(12345678).unwrap();
/// builder.append_null();
/// builder.append(22345678).unwrap();
/// let array = builder.finish();
///
/// assert_eq!(
/// array.keys(),
/// &UInt8Array::from(vec![Some(0), None, Some(1)])
/// );
///
/// // Values are polymorphic and so require a downcast.
/// let av = array.values();
/// let ava: &UInt32Array = av.as_any().downcast_ref::<UInt32Array>().unwrap();
/// let avs: &[u32] = ava.values();
///
/// assert!(!array.is_null(0));
/// assert!(array.is_null(1));
/// assert!(!array.is_null(2));
///
/// assert_eq!(avs, &[12345678, 22345678]);
/// ```
#[derive(Debug)]
pub struct PrimitiveDictionaryBuilder<K, V>
where
K: ArrowPrimitiveType,
V: ArrowPrimitiveType,
{
keys_builder: PrimitiveBuilder<K>,
values_builder: PrimitiveBuilder<V>,
map: HashMap<Value<V::Native>, usize>,
}
impl<K, V> Default for PrimitiveDictionaryBuilder<K, V>
where
K: ArrowPrimitiveType,
V: ArrowPrimitiveType,
{
fn default() -> Self {
Self::new()
}
}
impl<K, V> PrimitiveDictionaryBuilder<K, V>
where
K: ArrowPrimitiveType,
V: ArrowPrimitiveType,
{
/// Creates a new `PrimitiveDictionaryBuilder`.
pub fn new() -> Self {
Self {
keys_builder: PrimitiveBuilder::new(),
values_builder: PrimitiveBuilder::new(),
map: HashMap::new(),
}
}
/// Creates a new `PrimitiveDictionaryBuilder` from the provided keys and values builders.
///
/// # Panics
///
/// This method panics if `keys_builder` or `values_builder` is not empty.
pub fn new_from_empty_builders(
keys_builder: PrimitiveBuilder<K>,
values_builder: PrimitiveBuilder<V>,
) -> Self {
assert!(
keys_builder.is_empty() && values_builder.is_empty(),
"keys and values builders must be empty"
);
let values_capacity = values_builder.capacity();
Self {
keys_builder,
values_builder,
map: HashMap::with_capacity(values_capacity),
}
}
/// Creates a new `PrimitiveDictionaryBuilder` from existing `PrimitiveBuilder`s of keys and values.
///
/// # Safety
///
/// caller must ensure that the passed in builders are valid for DictionaryArray.
pub unsafe fn new_from_builders(
keys_builder: PrimitiveBuilder<K>,
values_builder: PrimitiveBuilder<V>,
) -> Self {
let keys = keys_builder.values_slice();
let values = values_builder.values_slice();
let mut map = HashMap::with_capacity(values.len());
keys.iter().zip(values.iter()).for_each(|(key, value)| {
map.insert(Value(*value), K::Native::to_usize(*key).unwrap());
});
Self {
keys_builder,
values_builder,
map,
}
}
/// Creates a new `PrimitiveDictionaryBuilder` with the provided capacities
///
/// `keys_capacity`: the number of keys, i.e. length of array to build
/// `values_capacity`: the number of distinct dictionary values, i.e. size of dictionary
pub fn with_capacity(keys_capacity: usize, values_capacity: usize) -> Self {
Self {
keys_builder: PrimitiveBuilder::with_capacity(keys_capacity),
values_builder: PrimitiveBuilder::with_capacity(values_capacity),
map: HashMap::with_capacity(values_capacity),
}
}
/// Creates a new `PrimitiveDictionaryBuilder` from the existing builder with the same
/// keys and values, but with a new data type for the keys.
///
/// # Example
/// ```
/// #
/// # use arrow_array::builder::PrimitiveDictionaryBuilder;
/// # use arrow_array::types::{UInt8Type, UInt16Type, UInt64Type};
/// # use arrow_array::UInt16Array;
/// # use arrow_schema::ArrowError;
///
/// let mut u8_keyed_builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt64Type>::new();
///
/// // appending too many values causes the dictionary to overflow
/// for i in 0..256 {
/// u8_keyed_builder.append_value(i);
/// }
/// let result = u8_keyed_builder.append(256);
/// assert!(matches!(result, Err(ArrowError::DictionaryKeyOverflowError{})));
///
/// // we need to upgrade to a larger key type
/// let mut u16_keyed_builder = PrimitiveDictionaryBuilder::<UInt16Type, UInt64Type>::try_new_from_builder(u8_keyed_builder).unwrap();
/// let dictionary_array = u16_keyed_builder.finish();
/// let keys = dictionary_array.keys();
///
/// assert_eq!(keys, &UInt16Array::from_iter(0..256));
pub fn try_new_from_builder<K2>(
mut source: PrimitiveDictionaryBuilder<K2, V>,
) -> Result<Self, ArrowError>
where
K::Native: NumCast,
K2: ArrowDictionaryKeyType,
K2::Native: NumCast,
{
let map = source.map;
let values_builder = source.values_builder;
let source_keys = source.keys_builder.finish();
let new_keys: PrimitiveArray<K> = source_keys.try_unary(|value| {
num_traits::cast::cast::<K2::Native, K::Native>(value).ok_or_else(|| {
ArrowError::CastError(format!(
"Can't cast dictionary keys from source type {:?} to type {:?}",
K2::DATA_TYPE,
K::DATA_TYPE
))
})
})?;
// drop source key here because currently source_keys and new_keys are holding reference to
// the same underlying null_buffer. Below we want to call new_keys.into_builder() it must
// be the only reference holder.
drop(source_keys);
Ok(Self {
map,
keys_builder: new_keys
.into_builder()
.expect("underlying buffer has no references"),
values_builder,
})
}
}
impl<K, V> ArrayBuilder for PrimitiveDictionaryBuilder<K, V>
where
K: ArrowDictionaryKeyType,
V: ArrowPrimitiveType,
{
/// Returns the builder as an non-mutable `Any` reference.
fn as_any(&self) -> &dyn Any {
self
}
/// Returns the builder as an mutable `Any` reference.
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
/// Returns the boxed builder as a box of `Any`.
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
/// Returns the number of array slots in the builder
fn len(&self) -> usize {
self.keys_builder.len()
}
/// Builds the array and reset this builder.
fn finish(&mut self) -> ArrayRef {
Arc::new(self.finish())
}
/// Builds the array without resetting the builder.
fn finish_cloned(&self) -> ArrayRef {
Arc::new(self.finish_cloned())
}
}
impl<K, V> PrimitiveDictionaryBuilder<K, V>
where
K: ArrowDictionaryKeyType,
V: ArrowPrimitiveType,
{
#[inline]
fn get_or_insert_key(&mut self, value: V::Native) -> Result<K::Native, ArrowError> {
match self.map.get(&Value(value)) {
Some(&key) => {
Ok(K::Native::from_usize(key).ok_or(ArrowError::DictionaryKeyOverflowError)?)
}
None => {
let key = self.values_builder.len();
self.values_builder.append_value(value);
self.map.insert(Value(value), key);
Ok(K::Native::from_usize(key).ok_or(ArrowError::DictionaryKeyOverflowError)?)
}
}
}
/// Append a primitive value to the array. Return an existing index
/// if already present in the values array or a new index if the
/// value is appended to the values array.
#[inline]
pub fn append(&mut self, value: V::Native) -> Result<K::Native, ArrowError> {
let key = self.get_or_insert_key(value)?;
self.keys_builder.append_value(key);
Ok(key)
}
/// Append a value multiple times to the array.
/// This is the same as `append` but allows to append the same value multiple times without doing multiple lookups.
///
/// Returns an error if the new index would overflow the key type.
pub fn append_n(&mut self, value: V::Native, count: usize) -> Result<K::Native, ArrowError> {
let key = self.get_or_insert_key(value)?;
self.keys_builder.append_value_n(key, count);
Ok(key)
}
/// Infallibly append a value to this builder
///
/// # Panics
///
/// Panics if the resulting length of the dictionary values array would exceed `T::Native::MAX`
#[inline]
pub fn append_value(&mut self, value: V::Native) {
self.append(value).expect("dictionary key overflow");
}
/// Infallibly append a value to this builder repeatedly `count` times.
/// This is the same as `append_value` but allows to append the same value multiple times without doing multiple lookups.
///
/// # Panics
///
/// Panics if the resulting length of the dictionary values array would exceed `T::Native::MAX`
pub fn append_values(&mut self, value: V::Native, count: usize) {
self.append_n(value, count)
.expect("dictionary key overflow");
}
/// Appends a null slot into the builder
#[inline]
pub fn append_null(&mut self) {
self.keys_builder.append_null()
}
/// Append `n` null slots into the builder
#[inline]
pub fn append_nulls(&mut self, n: usize) {
self.keys_builder.append_nulls(n)
}
/// Append an `Option` value into the builder
///
/// # Panics
///
/// Panics if the resulting length of the dictionary values array would exceed `T::Native::MAX`
#[inline]
pub fn append_option(&mut self, value: Option<V::Native>) {
match value {
None => self.append_null(),
Some(v) => self.append_value(v),
};
}
/// Append an `Option` value into the builder repeatedly `count` times.
/// This is the same as `append_option` but allows to append the same value multiple times without doing multiple lookups.
///
/// # Panics
///
/// Panics if the resulting length of the dictionary values array would exceed `T::Native::MAX`
pub fn append_options(&mut self, value: Option<V::Native>, count: usize) {
match value {
None => self.keys_builder.append_nulls(count),
Some(v) => self.append_values(v, count),
};
}
/// Extends builder with dictionary
///
/// This is the same as [`Self::extend`] but is faster as it translates
/// the dictionary values once rather than doing a lookup for each item in the iterator
///
/// when dictionary values are null (the actual mapped values) the keys are null
///
pub fn extend_dictionary(
&mut self,
dictionary: &TypedDictionaryArray<K, PrimitiveArray<V>>,
) -> Result<(), ArrowError> {
let values = dictionary.values();
let v_len = values.len();
let k_len = dictionary.keys().len();
if v_len == 0 && k_len == 0 {
return Ok(());
}
// All nulls
if v_len == 0 {
self.append_nulls(k_len);
return Ok(());
}
if k_len == 0 {
return Err(ArrowError::InvalidArgumentError(
"Dictionary keys should not be empty when values are not empty".to_string(),
));
}
// Orphan values will be carried over to the new dictionary
let mapped_values = values
.iter()
// Dictionary values can technically be null, so we need to handle that
.map(|dict_value| {
dict_value
.map(|dict_value| self.get_or_insert_key(dict_value))
.transpose()
})
.collect::<Result<Vec<_>, _>>()?;
// Just insert the keys without additional lookups
dictionary.keys().iter().for_each(|key| match key {
None => self.append_null(),
Some(original_dict_index) => {
let index = original_dict_index.as_usize().min(v_len - 1);
match mapped_values[index] {
None => self.append_null(),
Some(mapped_value) => self.keys_builder.append_value(mapped_value),
}
}
});
Ok(())
}
/// Builds the `DictionaryArray` and reset this builder.
pub fn finish(&mut self) -> DictionaryArray<K> {
self.map.clear();
let values = self.values_builder.finish();
let keys = self.keys_builder.finish();
let data_type =
DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(values.data_type().clone()));
let builder = keys
.into_data()
.into_builder()
.data_type(data_type)
.child_data(vec![values.into_data()]);
DictionaryArray::from(unsafe { builder.build_unchecked() })
}
/// Builds the `DictionaryArray` without resetting the builder.
pub fn finish_cloned(&self) -> DictionaryArray<K> {
let values = self.values_builder.finish_cloned();
let keys = self.keys_builder.finish_cloned();
let data_type = DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(V::DATA_TYPE));
let builder = keys
.into_data()
.into_builder()
.data_type(data_type)
.child_data(vec![values.into_data()]);
DictionaryArray::from(unsafe { builder.build_unchecked() })
}
/// Builds the `DictionaryArray` without resetting the values builder or
/// the internal de-duplication map.
///
/// The advantage of doing this is that the values will represent the entire
/// set of what has been built so-far by this builder and ensures
/// consistency in the assignment of keys to values across multiple calls
/// to `finish_preserve_values`. This enables ipc writers to efficiently
/// emit delta dictionaries.
///
/// The downside to this is that building the record requires creating a
/// copy of the values, which can become slowly more expensive if the
/// dictionary grows.
///
/// Additionally, if record batches from multiple different dictionary
/// builders for the same column are fed into a single ipc writer, beware
/// that entire dictionaries are likely to be re-sent frequently even when
/// the majority of the values are not used by the current record batch.
pub fn finish_preserve_values(&mut self) -> DictionaryArray<K> {
let values = self.values_builder.finish_cloned();
let keys = self.keys_builder.finish();
let data_type = DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(V::DATA_TYPE));
let builder = keys
.into_data()
.into_builder()
.data_type(data_type)
.child_data(vec![values.into_data()]);
DictionaryArray::from(unsafe { builder.build_unchecked() })
}
/// Returns the current dictionary values buffer as a slice
pub fn values_slice(&self) -> &[V::Native] {
self.values_builder.values_slice()
}
/// Returns the current dictionary values buffer as a mutable slice
pub fn values_slice_mut(&mut self) -> &mut [V::Native] {
self.values_builder.values_slice_mut()
}
/// Returns the current null buffer as a slice
pub fn validity_slice(&self) -> Option<&[u8]> {
self.keys_builder.validity_slice()
}
}
impl<K: ArrowDictionaryKeyType, P: ArrowPrimitiveType> Extend<Option<P::Native>>
for PrimitiveDictionaryBuilder<K, P>
{
#[inline]
fn extend<T: IntoIterator<Item = Option<P::Native>>>(&mut self, iter: T) {
for v in iter {
self.append_option(v)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::array::{Int32Array, UInt8Array, UInt32Array};
use crate::builder::Decimal128Builder;
use crate::cast::AsArray;
use crate::types::{
Date32Type, Decimal128Type, DurationNanosecondType, Float32Type, Float64Type, Int8Type,
Int16Type, Int32Type, Int64Type, TimestampNanosecondType, UInt8Type, UInt16Type,
UInt32Type, UInt64Type,
};
#[test]
fn test_primitive_dictionary_builder() {
let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::with_capacity(3, 2);
builder.append(12345678).unwrap();
builder.append_null();
builder.append(22345678).unwrap();
let array = builder.finish();
assert_eq!(
array.keys(),
&UInt8Array::from(vec![Some(0), None, Some(1)])
);
// Values are polymorphic and so require a downcast.
let av = array.values();
let ava: &UInt32Array = av.as_any().downcast_ref::<UInt32Array>().unwrap();
let avs: &[u32] = ava.values();
assert!(!array.is_null(0));
assert!(array.is_null(1));
assert!(!array.is_null(2));
assert_eq!(avs, &[12345678, 22345678]);
}
#[test]
fn test_extend() {
let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
builder.extend([1, 2, 3, 1, 2, 3, 1, 2, 3].into_iter().map(Some));
builder.extend([4, 5, 1, 3, 1].into_iter().map(Some));
let dict = builder.finish();
assert_eq!(
dict.keys().values(),
&[0, 1, 2, 0, 1, 2, 0, 1, 2, 3, 4, 0, 2, 0]
);
assert_eq!(dict.values().len(), 5);
}
#[test]
#[should_panic(expected = "DictionaryKeyOverflowError")]
fn test_primitive_dictionary_overflow() {
let mut builder =
PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::with_capacity(257, 257);
// 256 unique keys.
for i in 0..256 {
builder.append(i + 1000).unwrap();
}
// Special error if the key overflows (256th entry)
builder.append(1257).unwrap();
}
#[test]
fn test_primitive_dictionary_with_builders() {
let keys_builder = PrimitiveBuilder::<Int32Type>::new();
let values_builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(1, 2));
let mut builder =
PrimitiveDictionaryBuilder::<Int32Type, Decimal128Type>::new_from_empty_builders(
keys_builder,
values_builder,
);
let dict_array = builder.finish();
assert_eq!(dict_array.value_type(), DataType::Decimal128(1, 2));
assert_eq!(
dict_array.data_type(),
&DataType::Dictionary(
Box::new(DataType::Int32),
Box::new(DataType::Decimal128(1, 2)),
)
);
}
#[test]
fn test_extend_dictionary() {
let some_dict = {
let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
builder.extend([1, 2, 3, 1, 2, 3, 1, 2, 3].into_iter().map(Some));
builder.extend([None::<i32>]);
builder.extend([4, 5, 1, 3, 1].into_iter().map(Some));
builder.append_null();
builder.finish()
};
let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
builder.extend([6, 6, 7, 6, 5].into_iter().map(Some));
builder
.extend_dictionary(&some_dict.downcast_dict().unwrap())
.unwrap();
let dict = builder.finish();
assert_eq!(dict.values().len(), 7);
let values = dict
.downcast_dict::<Int32Array>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(
values,
[
Some(6),
Some(6),
Some(7),
Some(6),
Some(5),
Some(1),
Some(2),
Some(3),
Some(1),
Some(2),
Some(3),
Some(1),
Some(2),
Some(3),
None,
Some(4),
Some(5),
Some(1),
Some(3),
Some(1),
None
]
);
}
#[test]
fn test_extend_dictionary_with_null_in_mapped_value() {
let some_dict = {
let mut values_builder = PrimitiveBuilder::<Int32Type>::new();
let mut keys_builder = PrimitiveBuilder::<Int32Type>::new();
// Manually build a dictionary values that the mapped values have null
values_builder.append_null();
keys_builder.append_value(0);
values_builder.append_value(42);
keys_builder.append_value(1);
let values = values_builder.finish();
let keys = keys_builder.finish();
let data_type = DataType::Dictionary(
Box::new(Int32Type::DATA_TYPE),
Box::new(values.data_type().clone()),
);
let builder = keys
.into_data()
.into_builder()
.data_type(data_type)
.child_data(vec![values.into_data()]);
DictionaryArray::from(unsafe { builder.build_unchecked() })
};
let some_dict_values = some_dict.values().as_primitive::<Int32Type>();
assert_eq!(
some_dict_values.into_iter().collect::<Vec<_>>(),
&[None, Some(42)]
);
let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
builder
.extend_dictionary(&some_dict.downcast_dict().unwrap())
.unwrap();
let dict = builder.finish();
assert_eq!(dict.values().len(), 1);
let values = dict
.downcast_dict::<Int32Array>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(values, [None, Some(42)]);
}
#[test]
fn test_extend_all_null_dictionary() {
let some_dict = {
let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
builder.append_nulls(2);
builder.finish()
};
let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
builder
.extend_dictionary(&some_dict.downcast_dict().unwrap())
.unwrap();
let dict = builder.finish();
assert_eq!(dict.values().len(), 0);
let values = dict
.downcast_dict::<Int32Array>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(values, [None, None]);
}
#[test]
fn creating_dictionary_from_builders_should_use_values_capacity_for_the_map() {
let builder = PrimitiveDictionaryBuilder::<Int32Type, crate::types::TimestampMicrosecondType>::new_from_empty_builders(
PrimitiveBuilder::with_capacity(1).with_data_type(DataType::Int32),
PrimitiveBuilder::with_capacity(2).with_data_type(DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some("+08:00".into()))),
);
assert!(
builder.map.capacity() >= builder.values_builder.capacity(),
"map capacity {} should be at least the values capacity {}",
builder.map.capacity(),
builder.values_builder.capacity()
)
}
fn _test_try_new_from_builder_generic_for_key_types<K1, K2, V>(values: Vec<V::Native>)
where
K1: ArrowDictionaryKeyType,
K1::Native: NumCast,
K2: ArrowDictionaryKeyType,
K2::Native: NumCast + From<u8>,
V: ArrowPrimitiveType,
{
let mut source = PrimitiveDictionaryBuilder::<K1, V>::new();
source.append(values[0]).unwrap();
source.append_null();
source.append(values[1]).unwrap();
source.append(values[2]).unwrap();
let mut result = PrimitiveDictionaryBuilder::<K2, V>::try_new_from_builder(source).unwrap();
let array = result.finish();
let mut expected_keys_builder = PrimitiveBuilder::<K2>::new();
expected_keys_builder
.append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(0u8));
expected_keys_builder.append_null();
expected_keys_builder
.append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(1u8));
expected_keys_builder
.append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(2u8));
let expected_keys = expected_keys_builder.finish();
assert_eq!(array.keys(), &expected_keys);
let av = array.values();
let ava = av.as_any().downcast_ref::<PrimitiveArray<V>>().unwrap();
assert_eq!(ava.value(0), values[0]);
assert_eq!(ava.value(1), values[1]);
assert_eq!(ava.value(2), values[2]);
}
fn _test_try_new_from_builder_generic_for_value<T>(values: Vec<T::Native>)
where
T: ArrowPrimitiveType,
{
// test cast to bigger size unsigned
_test_try_new_from_builder_generic_for_key_types::<UInt8Type, UInt16Type, T>(
values.clone(),
);
// test cast going to smaller size unsigned
_test_try_new_from_builder_generic_for_key_types::<UInt16Type, UInt8Type, T>(
values.clone(),
);
// test cast going to bigger size signed
_test_try_new_from_builder_generic_for_key_types::<Int8Type, Int16Type, T>(values.clone());
// test cast going to smaller size signed
_test_try_new_from_builder_generic_for_key_types::<Int32Type, Int16Type, T>(values.clone());
// test going from signed to signed for different size changes
_test_try_new_from_builder_generic_for_key_types::<UInt8Type, Int16Type, T>(values.clone());
_test_try_new_from_builder_generic_for_key_types::<Int8Type, UInt8Type, T>(values.clone());
_test_try_new_from_builder_generic_for_key_types::<Int8Type, UInt16Type, T>(values.clone());
_test_try_new_from_builder_generic_for_key_types::<Int32Type, Int16Type, T>(values.clone());
}
#[test]
fn test_try_new_from_builder() {
// test unsigned types
_test_try_new_from_builder_generic_for_value::<UInt8Type>(vec![1, 2, 3]);
_test_try_new_from_builder_generic_for_value::<UInt16Type>(vec![1, 2, 3]);
_test_try_new_from_builder_generic_for_value::<UInt32Type>(vec![1, 2, 3]);
_test_try_new_from_builder_generic_for_value::<UInt64Type>(vec![1, 2, 3]);
// test signed types
_test_try_new_from_builder_generic_for_value::<Int8Type>(vec![-1, 0, 1]);
_test_try_new_from_builder_generic_for_value::<Int16Type>(vec![-1, 0, 1]);
_test_try_new_from_builder_generic_for_value::<Int32Type>(vec![-1, 0, 1]);
_test_try_new_from_builder_generic_for_value::<Int64Type>(vec![-1, 0, 1]);
// test some date types
_test_try_new_from_builder_generic_for_value::<Date32Type>(vec![5, 6, 7]);
_test_try_new_from_builder_generic_for_value::<DurationNanosecondType>(vec![1, 2, 3]);
_test_try_new_from_builder_generic_for_value::<TimestampNanosecondType>(vec![1, 2, 3]);
// test some floating point types
_test_try_new_from_builder_generic_for_value::<Float32Type>(vec![0.1, 0.2, 0.3]);
_test_try_new_from_builder_generic_for_value::<Float64Type>(vec![-0.1, 0.2, 0.3]);
}
#[test]
fn test_try_new_from_builder_cast_fails() {
let mut source_builder = PrimitiveDictionaryBuilder::<UInt16Type, UInt64Type>::new();
for i in 0..257 {
source_builder.append_value(i);
}
// there should be too many values that we can't downcast to the underlying type
// we have keys that wouldn't fit into UInt8Type
let result = PrimitiveDictionaryBuilder::<UInt8Type, UInt64Type>::try_new_from_builder(
source_builder,
);
assert!(result.is_err());
if let Err(e) = result {
assert!(matches!(e, ArrowError::CastError(_)));
assert_eq!(
e.to_string(),
"Cast error: Can't cast dictionary keys from source type UInt16 to type UInt8"
);
}
}
#[test]
fn test_finish_preserve_values() {
// Create the first dictionary
let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::new();
builder.append(10).unwrap();
builder.append(20).unwrap();
let array = builder.finish_preserve_values();
assert_eq!(array.keys(), &UInt8Array::from(vec![Some(0), Some(1)]));
let values: &[u32] = array
.values()
.as_any()
.downcast_ref::<UInt32Array>()
.unwrap()
.values();
assert_eq!(values, &[10, 20]);
// Create a new dictionary
builder.append(30).unwrap();
builder.append(40).unwrap();
let array2 = builder.finish_preserve_values();
// Make sure the keys are assigned after the old ones
// and that we have the right values
assert_eq!(array2.keys(), &UInt8Array::from(vec![Some(2), Some(3)]));
let values = array2
.downcast_dict::<UInt32Array>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(values, vec![Some(30), Some(40)]);
// Check that we have all of the expected values
let all_values: &[u32] = array2
.values()
.as_any()
.downcast_ref::<UInt32Array>()
.unwrap()
.values();
assert_eq!(all_values, &[10, 20, 30, 40]);
}
}
@@ -0,0 +1,313 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use std::{any::Any, sync::Arc};
use crate::{ArrayRef, ArrowPrimitiveType, RunArray, types::RunEndIndexType};
use super::{ArrayBuilder, PrimitiveBuilder};
use arrow_buffer::ArrowNativeType;
/// Builder for [`RunArray`] of [`PrimitiveArray`](crate::array::PrimitiveArray)
///
/// # Example:
///
/// ```
///
/// # use arrow_array::builder::PrimitiveRunBuilder;
/// # use arrow_array::cast::AsArray;
/// # use arrow_array::types::{UInt32Type, Int16Type};
/// # use arrow_array::{Array, UInt32Array, Int16Array};
///
/// let mut builder =
/// PrimitiveRunBuilder::<Int16Type, UInt32Type>::new();
/// builder.append_value(1234);
/// builder.append_value(1234);
/// builder.append_value(1234);
/// builder.append_null();
/// builder.append_value(5678);
/// builder.append_value(5678);
/// let array = builder.finish();
///
/// assert_eq!(array.run_ends().values(), &[3, 4, 6]);
///
/// let av = array.values();
///
/// assert!(!av.is_null(0));
/// assert!(av.is_null(1));
/// assert!(!av.is_null(2));
///
/// // Values are polymorphic and so require a downcast.
/// let ava: &UInt32Array = av.as_primitive::<UInt32Type>();
///
/// assert_eq!(ava, &UInt32Array::from(vec![Some(1234), None, Some(5678)]));
/// ```
#[derive(Debug)]
pub struct PrimitiveRunBuilder<R, V>
where
R: RunEndIndexType,
V: ArrowPrimitiveType,
{
run_ends_builder: PrimitiveBuilder<R>,
values_builder: PrimitiveBuilder<V>,
current_value: Option<V::Native>,
current_run_end_index: usize,
prev_run_end_index: usize,
}
impl<R, V> Default for PrimitiveRunBuilder<R, V>
where
R: RunEndIndexType,
V: ArrowPrimitiveType,
{
fn default() -> Self {
Self::new()
}
}
impl<R, V> PrimitiveRunBuilder<R, V>
where
R: RunEndIndexType,
V: ArrowPrimitiveType,
{
/// Creates a new `PrimitiveRunBuilder`
pub fn new() -> Self {
Self {
run_ends_builder: PrimitiveBuilder::new(),
values_builder: PrimitiveBuilder::new(),
current_value: None,
current_run_end_index: 0,
prev_run_end_index: 0,
}
}
/// Creates a new `PrimitiveRunBuilder` with the provided capacity
///
/// `capacity`: the expected number of run-end encoded values.
pub fn with_capacity(capacity: usize) -> Self {
Self {
run_ends_builder: PrimitiveBuilder::with_capacity(capacity),
values_builder: PrimitiveBuilder::with_capacity(capacity),
current_value: None,
current_run_end_index: 0,
prev_run_end_index: 0,
}
}
}
impl<R, V> ArrayBuilder for PrimitiveRunBuilder<R, V>
where
R: RunEndIndexType,
V: ArrowPrimitiveType,
{
/// Returns the builder as a non-mutable `Any` reference.
fn as_any(&self) -> &dyn Any {
self
}
/// Returns the builder as a mutable `Any` reference.
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
/// Returns the boxed builder as a box of `Any`.
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
/// Returns the length of logical array encoded by
/// the eventual runs array.
fn len(&self) -> usize {
self.current_run_end_index
}
/// Builds the array and reset this builder.
fn finish(&mut self) -> ArrayRef {
Arc::new(self.finish())
}
/// Builds the array without resetting the builder.
fn finish_cloned(&self) -> ArrayRef {
Arc::new(self.finish_cloned())
}
}
impl<R, V> PrimitiveRunBuilder<R, V>
where
R: RunEndIndexType,
V: ArrowPrimitiveType,
{
/// Appends optional value to the logical array encoded by the RunArray.
pub fn append_option(&mut self, value: Option<V::Native>) {
if self.current_run_end_index == 0 {
self.current_run_end_index = 1;
self.current_value = value;
return;
}
if self.current_value != value {
self.append_run_end();
self.current_value = value;
}
self.current_run_end_index += 1;
}
/// Appends value to the logical array encoded by the run-ends array.
pub fn append_value(&mut self, value: V::Native) {
self.append_option(Some(value))
}
/// Appends null to the logical array encoded by the run-ends array.
pub fn append_null(&mut self) {
self.append_option(None)
}
/// Creates the RunArray and resets the builder.
/// Panics if RunArray cannot be built.
pub fn finish(&mut self) -> RunArray<R> {
// write the last run end to the array.
self.append_run_end();
// reset the run index to zero.
self.current_value = None;
self.current_run_end_index = 0;
// build the run encoded array by adding run_ends and values array as its children.
let run_ends_array = self.run_ends_builder.finish();
let values_array = self.values_builder.finish();
RunArray::<R>::try_new(&run_ends_array, &values_array).unwrap()
}
/// Creates the RunArray and without resetting the builder.
/// Panics if RunArray cannot be built.
pub fn finish_cloned(&self) -> RunArray<R> {
let mut run_ends_array = self.run_ends_builder.finish_cloned();
let mut values_array = self.values_builder.finish_cloned();
// Add current run if one exists
if self.prev_run_end_index != self.current_run_end_index {
let mut run_end_builder = run_ends_array.into_builder().unwrap();
let mut values_builder = values_array.into_builder().unwrap();
self.append_run_end_with_builders(&mut run_end_builder, &mut values_builder);
run_ends_array = run_end_builder.finish();
values_array = values_builder.finish();
}
RunArray::try_new(&run_ends_array, &values_array).unwrap()
}
// Appends the current run to the array.
fn append_run_end(&mut self) {
// empty array or the function called without appending any value.
if self.prev_run_end_index == self.current_run_end_index {
return;
}
let run_end_index = self.run_end_index_as_native();
self.run_ends_builder.append_value(run_end_index);
self.values_builder.append_option(self.current_value);
self.prev_run_end_index = self.current_run_end_index;
}
// Similar to `append_run_end` but on custom builders.
// Used in `finish_cloned` which is not suppose to mutate `self`.
fn append_run_end_with_builders(
&self,
run_ends_builder: &mut PrimitiveBuilder<R>,
values_builder: &mut PrimitiveBuilder<V>,
) {
let run_end_index = self.run_end_index_as_native();
run_ends_builder.append_value(run_end_index);
values_builder.append_option(self.current_value);
}
fn run_end_index_as_native(&self) -> R::Native {
R::Native::from_usize(self.current_run_end_index)
.unwrap_or_else(|| panic!(
"Cannot convert `current_run_end_index` {} from `usize` to native form of arrow datatype {}",
self.current_run_end_index,
R::DATA_TYPE
))
}
}
impl<R, V> Extend<Option<V::Native>> for PrimitiveRunBuilder<R, V>
where
R: RunEndIndexType,
V: ArrowPrimitiveType,
{
fn extend<T: IntoIterator<Item = Option<V::Native>>>(&mut self, iter: T) {
for elem in iter {
self.append_option(elem);
}
}
}
#[cfg(test)]
mod tests {
use crate::builder::PrimitiveRunBuilder;
use crate::cast::AsArray;
use crate::types::{Int16Type, UInt32Type};
use crate::{Array, UInt32Array};
#[test]
fn test_primitive_ree_array_builder() {
let mut builder = PrimitiveRunBuilder::<Int16Type, UInt32Type>::new();
builder.append_value(1234);
builder.append_value(1234);
builder.append_value(1234);
builder.append_null();
builder.append_value(5678);
builder.append_value(5678);
let array = builder.finish();
assert_eq!(array.null_count(), 0);
assert_eq!(array.logical_null_count(), 1);
assert_eq!(array.len(), 6);
assert_eq!(array.run_ends().values(), &[3, 4, 6]);
let av = array.values();
assert!(!av.is_null(0));
assert!(av.is_null(1));
assert!(!av.is_null(2));
// Values are polymorphic and so require a downcast.
let ava: &UInt32Array = av.as_primitive::<UInt32Type>();
assert_eq!(ava, &UInt32Array::from(vec![Some(1234), None, Some(5678)]));
}
#[test]
fn test_extend() {
let mut builder = PrimitiveRunBuilder::<Int16Type, Int16Type>::new();
builder.extend([1, 2, 2, 5, 5, 4, 4].into_iter().map(Some));
builder.extend([4, 4, 6, 2].into_iter().map(Some));
let array = builder.finish();
assert_eq!(array.len(), 11);
assert_eq!(array.null_count(), 0);
assert_eq!(array.logical_null_count(), 0);
assert_eq!(array.run_ends().values(), &[1, 3, 5, 9, 10, 11]);
assert_eq!(
array.values().as_primitive::<Int16Type>().values(),
&[1, 2, 5, 4, 6, 2]
);
}
}
@@ -0,0 +1,730 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::StructArray;
use crate::builder::*;
use arrow_buffer::NullBufferBuilder;
use arrow_schema::{Fields, SchemaBuilder};
use std::sync::Arc;
/// Builder for [`StructArray`]
///
/// Note that callers should make sure that methods of all the child field builders are
/// properly called to maintain the consistency of the data structure.
///
///
/// Handling arrays with complex layouts, such as `List<Struct<List<Struct>>>`, in Rust can be challenging due to its strong typing system.
/// To construct a collection builder ([`ListBuilder`], [`LargeListBuilder`], or [`MapBuilder`]) using [`make_builder`], multiple calls are required. This complexity arises from the recursive approach utilized by [`StructBuilder::from_fields`].
///
/// Initially, [`StructBuilder::from_fields`] invokes [`make_builder`], which returns a `Box<dyn ArrayBuilder>`. To obtain the specific collection builder, one must first use [`StructBuilder::field_builder`] to get a `Collection<[Box<dyn ArrayBuilder>]>`. Subsequently, the `values()` result from this operation can be downcast to the desired builder type.
///
/// For example, when working with [`ListBuilder`], you would first call [`StructBuilder::field_builder::<ListBuilder<Box<dyn ArrayBuilder>>>`] and then downcast the [`Box<dyn ArrayBuilder>`] to the specific [`StructBuilder`] you need.
///
/// For a practical example see the code below:
///
/// ```rust
/// use arrow_array::builder::{ArrayBuilder, ListBuilder, StringBuilder, StructBuilder};
/// use arrow_schema::{DataType, Field, Fields};
/// use std::sync::Arc;
///
/// // This is an example column that has a List<Struct<List<Struct>>> layout
/// let mut example_col = ListBuilder::new(StructBuilder::from_fields(
/// vec![Field::new(
/// "value_list",
/// DataType::List(Arc::new(Field::new_list_field(
/// DataType::Struct(Fields::from(vec![
/// Field::new("key", DataType::Utf8, true),
/// Field::new("value", DataType::Utf8, true),
/// ])), //In this example we are trying to get to this builder and insert key/value pairs
/// true,
/// ))),
/// true,
/// )],
/// 0,
/// ));
///
/// // We can obtain the StructBuilder without issues, because example_col was created with StructBuilder
/// let col_struct_builder: &mut StructBuilder = example_col.values();
///
/// // We can't obtain the ListBuilder<StructBuilder> with the expected generic types, because under the hood
/// // the StructBuilder was returned as a Box<dyn ArrayBuilder> and passed as such to the ListBuilder constructor
///
/// // This panics in runtime, even though we know that the builder is a ListBuilder<StructBuilder>.
/// // let sb = col_struct_builder
/// // .field_builder::<ListBuilder<StructBuilder>>(0)
/// // .as_mut()
/// // .unwrap();
///
/// //To keep in line with Rust's strong typing, we fetch a ListBuilder<Box<dyn ArrayBuilder>> from the column StructBuilder first...
/// let mut list_builder_option =
/// col_struct_builder.field_builder::<ListBuilder<Box<dyn ArrayBuilder>>>(0);
///
/// let list_builder = list_builder_option.as_mut().unwrap();
///
/// // ... and then downcast the key/value pair values to a StructBuilder
/// let struct_builder = list_builder
/// .values()
/// .as_any_mut()
/// .downcast_mut::<StructBuilder>()
/// .unwrap();
///
/// // We can now append values to the StructBuilder
/// let key_builder = struct_builder.field_builder::<StringBuilder>(0).unwrap();
/// key_builder.append_value("my key");
///
/// let value_builder = struct_builder.field_builder::<StringBuilder>(1).unwrap();
/// value_builder.append_value("my value");
///
/// struct_builder.append(true);
/// list_builder.append(true);
/// col_struct_builder.append(true);
/// example_col.append(true);
///
/// let array = example_col.finish();
///
/// println!("My array: {:?}", array);
/// ```
///
pub struct StructBuilder {
fields: Fields,
field_builders: Vec<Box<dyn ArrayBuilder>>,
null_buffer_builder: NullBufferBuilder,
}
impl std::fmt::Debug for StructBuilder {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("StructBuilder")
.field("fields", &self.fields)
.field("bitmap_builder", &self.null_buffer_builder)
.field("len", &self.len())
.finish()
}
}
impl ArrayBuilder for StructBuilder {
/// Returns the number of array slots in the builder.
///
/// Note that this always return the first child field builder's length, and it is
/// the caller's responsibility to maintain the consistency that all the child field
/// builder should have the equal number of elements.
fn len(&self) -> usize {
self.null_buffer_builder.len()
}
/// Builds the array.
fn finish(&mut self) -> ArrayRef {
Arc::new(self.finish())
}
/// Builds the array without resetting the builder.
fn finish_cloned(&self) -> ArrayRef {
Arc::new(self.finish_cloned())
}
/// Returns the builder as a non-mutable `Any` reference.
///
/// This is most useful when one wants to call non-mutable APIs on a specific builder
/// type. In this case, one can first cast this into a `Any`, and then use
/// `downcast_ref` to get a reference on the specific builder.
fn as_any(&self) -> &dyn Any {
self
}
/// Returns the builder as a mutable `Any` reference.
///
/// This is most useful when one wants to call mutable APIs on a specific builder
/// type. In this case, one can first cast this into a `Any`, and then use
/// `downcast_mut` to get a reference on the specific builder.
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
/// Returns the boxed builder as a box of `Any`.
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
}
impl StructBuilder {
/// Creates a new `StructBuilder`
pub fn new(fields: impl Into<Fields>, field_builders: Vec<Box<dyn ArrayBuilder>>) -> Self {
Self {
field_builders,
fields: fields.into(),
null_buffer_builder: NullBufferBuilder::new(0),
}
}
/// Creates a new `StructBuilder` from [`Fields`] and `capacity`
pub fn from_fields(fields: impl Into<Fields>, capacity: usize) -> Self {
let fields = fields.into();
let mut builders = Vec::with_capacity(fields.len());
for field in &fields {
builders.push(make_builder(field.data_type(), capacity));
}
Self::new(fields, builders)
}
/// Returns a mutable reference to the child field builder at index `i`.
/// Result will be `None` if the input type `T` provided doesn't match the actual
/// field builder's type.
pub fn field_builder<T: ArrayBuilder>(&mut self, i: usize) -> Option<&mut T> {
self.field_builders[i].as_any_mut().downcast_mut::<T>()
}
/// Returns a reference to field builders
pub fn field_builders(&self) -> &[Box<dyn ArrayBuilder>] {
&self.field_builders
}
/// Returns a mutable reference to field builders
pub fn field_builders_mut(&mut self) -> &mut [Box<dyn ArrayBuilder>] {
&mut self.field_builders
}
/// Returns the number of fields for the struct this builder is building.
pub fn num_fields(&self) -> usize {
self.field_builders.len()
}
/// Returns the fields for the struct this builder is building.
pub fn fields(&self) -> &Fields {
&self.fields
}
/// Appends an element (either null or non-null) to the struct. The actual elements
/// should be appended for each child sub-array in a consistent way.
#[inline]
pub fn append(&mut self, is_valid: bool) {
self.null_buffer_builder.append(is_valid);
}
/// Appends a null element to the struct.
#[inline]
pub fn append_null(&mut self) {
self.append(false)
}
/// Appends `n` `null`s into the builder.
#[inline]
pub fn append_nulls(&mut self, n: usize) {
self.null_buffer_builder.append_slice(&vec![false; n]);
}
/// Builds the `StructArray` and reset this builder.
pub fn finish(&mut self) -> StructArray {
self.validate_content();
if self.fields.is_empty() {
return StructArray::new_empty_fields(self.len(), self.null_buffer_builder.finish());
}
let arrays = self.field_builders.iter_mut().map(|f| f.finish()).collect();
let nulls = self.null_buffer_builder.finish();
StructArray::new(self.fields.clone(), arrays, nulls)
}
/// Builds the `StructArray` without resetting the builder.
pub fn finish_cloned(&self) -> StructArray {
self.validate_content();
if self.fields.is_empty() {
return StructArray::new_empty_fields(
self.len(),
self.null_buffer_builder.finish_cloned(),
);
}
let arrays = self
.field_builders
.iter()
.map(|f| f.finish_cloned())
.collect();
let nulls = self.null_buffer_builder.finish_cloned();
StructArray::new(self.fields.clone(), arrays, nulls)
}
/// Constructs and validates contents in the builder to ensure that
/// - fields and field_builders are of equal length
/// - the number of items in individual field_builders are equal to self.len()
fn validate_content(&self) {
if self.fields.len() != self.field_builders.len() {
panic!("Number of fields is not equal to the number of field_builders.");
}
self.field_builders.iter().enumerate().for_each(|(idx, x)| {
if x.len() != self.len() {
let builder = SchemaBuilder::from(&self.fields);
let schema = builder.finish();
panic!("{}", format!(
"StructBuilder ({}) and field_builder with index {} ({}) are of unequal lengths: ({} != {}).",
schema,
idx,
self.fields[idx].data_type(),
self.len(),
x.len()
));
}
});
}
/// Returns the current null buffer as a slice
pub fn validity_slice(&self) -> Option<&[u8]> {
self.null_buffer_builder.as_slice()
}
}
#[cfg(test)]
mod tests {
use std::any::type_name;
use super::*;
use arrow_buffer::Buffer;
use arrow_data::ArrayData;
use arrow_schema::Field;
use crate::{array::Array, types::ArrowDictionaryKeyType};
#[test]
fn test_struct_array_builder() {
let string_builder = StringBuilder::new();
let int_builder = Int32Builder::new();
let fields = vec![
Field::new("f1", DataType::Utf8, true),
Field::new("f2", DataType::Int32, true),
];
let field_builders = vec![
Box::new(string_builder) as Box<dyn ArrayBuilder>,
Box::new(int_builder) as Box<dyn ArrayBuilder>,
];
let mut builder = StructBuilder::new(fields, field_builders);
assert_eq!(2, builder.num_fields());
let string_builder = builder
.field_builder::<StringBuilder>(0)
.expect("builder at field 0 should be string builder");
string_builder.append_value("joe");
string_builder.append_null();
string_builder.append_null();
string_builder.append_value("mark");
string_builder.append_nulls(2);
string_builder.append_value("terry");
let int_builder = builder
.field_builder::<Int32Builder>(1)
.expect("builder at field 1 should be int builder");
int_builder.append_value(1);
int_builder.append_value(2);
int_builder.append_null();
int_builder.append_value(4);
int_builder.append_nulls(2);
int_builder.append_value(3);
builder.append(true);
builder.append(true);
builder.append_null();
builder.append(true);
builder.append_nulls(2);
builder.append(true);
let struct_data = builder.finish().into_data();
assert_eq!(7, struct_data.len());
assert_eq!(3, struct_data.null_count());
assert_eq!(&[75_u8], struct_data.nulls().unwrap().validity());
let expected_string_data = ArrayData::builder(DataType::Utf8)
.len(7)
.null_bit_buffer(Some(Buffer::from(&[73_u8])))
.add_buffer(Buffer::from_slice_ref([0, 3, 3, 3, 7, 7, 7, 12]))
.add_buffer(Buffer::from_slice_ref(b"joemarkterry"))
.build()
.unwrap();
let expected_int_data = ArrayData::builder(DataType::Int32)
.len(7)
.null_bit_buffer(Some(Buffer::from_slice_ref([75_u8])))
.add_buffer(Buffer::from_slice_ref([1, 2, 0, 4, 4, 4, 3]))
.build()
.unwrap();
assert_eq!(expected_string_data, struct_data.child_data()[0]);
assert_eq!(expected_int_data, struct_data.child_data()[1]);
assert!(struct_data.is_null(4));
assert!(struct_data.is_null(5));
}
#[test]
fn test_struct_array_builder_finish() {
let int_builder = Int32Builder::new();
let bool_builder = BooleanBuilder::new();
let fields = vec![
Field::new("f1", DataType::Int32, false),
Field::new("f2", DataType::Boolean, false),
];
let field_builders = vec![
Box::new(int_builder) as Box<dyn ArrayBuilder>,
Box::new(bool_builder) as Box<dyn ArrayBuilder>,
];
let mut builder = StructBuilder::new(fields, field_builders);
builder
.field_builder::<Int32Builder>(0)
.unwrap()
.append_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
builder
.field_builder::<BooleanBuilder>(1)
.unwrap()
.append_slice(&[
false, true, false, true, false, true, false, true, false, true,
]);
// Append slot values - all are valid.
for _ in 0..10 {
builder.append(true);
}
assert_eq!(10, builder.len());
let arr = builder.finish();
assert_eq!(10, arr.len());
assert_eq!(0, builder.len());
builder
.field_builder::<Int32Builder>(0)
.unwrap()
.append_slice(&[1, 3, 5, 7, 9]);
builder
.field_builder::<BooleanBuilder>(1)
.unwrap()
.append_slice(&[false, true, false, true, false]);
// Append slot values - all are valid.
for _ in 0..5 {
builder.append(true);
}
assert_eq!(5, builder.len());
let arr = builder.finish();
assert_eq!(5, arr.len());
assert_eq!(0, builder.len());
}
#[test]
fn test_build_fixed_size_list() {
const LIST_LENGTH: i32 = 4;
let fixed_size_list_dtype =
DataType::new_fixed_size_list(DataType::Int32, LIST_LENGTH, false);
let mut builder = make_builder(&fixed_size_list_dtype, 10);
let builder = builder
.as_any_mut()
.downcast_mut::<FixedSizeListBuilder<Box<dyn ArrayBuilder>>>();
match builder {
Some(builder) => {
assert_eq!(builder.value_length(), LIST_LENGTH);
assert!(
builder
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.is_some()
);
}
None => panic!("expected FixedSizeListBuilder, got a different builder type"),
}
}
#[test]
fn test_struct_array_builder_finish_cloned() {
let int_builder = Int32Builder::new();
let bool_builder = BooleanBuilder::new();
let fields = vec![
Field::new("f1", DataType::Int32, false),
Field::new("f2", DataType::Boolean, false),
];
let field_builders = vec![
Box::new(int_builder) as Box<dyn ArrayBuilder>,
Box::new(bool_builder) as Box<dyn ArrayBuilder>,
];
let mut builder = StructBuilder::new(fields, field_builders);
builder
.field_builder::<Int32Builder>(0)
.unwrap()
.append_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
builder
.field_builder::<BooleanBuilder>(1)
.unwrap()
.append_slice(&[
false, true, false, true, false, true, false, true, false, true,
]);
// Append slot values - all are valid.
for _ in 0..10 {
builder.append(true);
}
assert_eq!(10, builder.len());
let mut arr = builder.finish_cloned();
assert_eq!(10, arr.len());
assert_eq!(10, builder.len());
builder
.field_builder::<Int32Builder>(0)
.unwrap()
.append_slice(&[1, 3, 5, 7, 9]);
builder
.field_builder::<BooleanBuilder>(1)
.unwrap()
.append_slice(&[false, true, false, true, false]);
// Append slot values - all are valid.
for _ in 0..5 {
builder.append(true);
}
assert_eq!(15, builder.len());
arr = builder.finish();
assert_eq!(15, arr.len());
assert_eq!(0, builder.len());
}
#[test]
fn test_struct_array_builder_from_schema() {
let mut fields = vec![
Field::new("f1", DataType::Float32, false),
Field::new("f2", DataType::Utf8, false),
];
let sub_fields = vec![
Field::new("g1", DataType::Int32, false),
Field::new("g2", DataType::Boolean, false),
];
let struct_type = DataType::Struct(sub_fields.into());
fields.push(Field::new("f3", struct_type, false));
let mut builder = StructBuilder::from_fields(fields, 5);
assert_eq!(3, builder.num_fields());
assert!(builder.field_builder::<Float32Builder>(0).is_some());
assert!(builder.field_builder::<StringBuilder>(1).is_some());
assert!(builder.field_builder::<StructBuilder>(2).is_some());
}
#[test]
fn test_datatype_properties() {
let fields = Fields::from(vec![
Field::new("f1", DataType::Decimal128(1, 2), false),
Field::new(
"f2",
DataType::Timestamp(TimeUnit::Millisecond, Some("+00:00".into())),
false,
),
]);
let mut builder = StructBuilder::from_fields(fields.clone(), 1);
builder
.field_builder::<Decimal128Builder>(0)
.unwrap()
.append_value(1);
builder
.field_builder::<TimestampMillisecondBuilder>(1)
.unwrap()
.append_value(1);
builder.append(true);
let array = builder.finish();
assert_eq!(array.data_type(), &DataType::Struct(fields.clone()));
assert_eq!(array.column(0).data_type(), fields[0].data_type());
assert_eq!(array.column(1).data_type(), fields[1].data_type());
}
#[test]
fn test_struct_array_builder_from_dictionary_type_int8_key() {
test_struct_array_builder_from_dictionary_type_inner::<Int8Type>(DataType::Int8);
}
#[test]
fn test_struct_array_builder_from_dictionary_type_int16_key() {
test_struct_array_builder_from_dictionary_type_inner::<Int16Type>(DataType::Int16);
}
#[test]
fn test_struct_array_builder_from_dictionary_type_int32_key() {
test_struct_array_builder_from_dictionary_type_inner::<Int32Type>(DataType::Int32);
}
#[test]
fn test_struct_array_builder_from_dictionary_type_int64_key() {
test_struct_array_builder_from_dictionary_type_inner::<Int64Type>(DataType::Int64);
}
fn test_struct_array_builder_from_dictionary_type_inner<K: ArrowDictionaryKeyType>(
key_type: DataType,
) {
let dict_field = Field::new(
"f1",
DataType::Dictionary(Box::new(key_type), Box::new(DataType::Utf8)),
false,
);
let fields = vec![dict_field.clone()];
let expected_dtype = DataType::Struct(fields.into());
let cloned_dict_field = dict_field.clone();
let expected_child_dtype = dict_field.data_type();
let mut struct_builder = StructBuilder::from_fields(vec![cloned_dict_field], 5);
let Some(dict_builder) = struct_builder.field_builder::<StringDictionaryBuilder<K>>(0)
else {
panic!(
"Builder should be StringDictionaryBuilder<{}>",
type_name::<K>()
)
};
dict_builder.append_value("dict string");
struct_builder.append(true);
let array = struct_builder.finish();
assert_eq!(array.data_type(), &expected_dtype);
assert_eq!(array.column(0).data_type(), expected_child_dtype);
assert_eq!(array.column(0).len(), 1);
}
#[test]
#[should_panic(
expected = "Data type Dictionary(UInt64, Utf8) with key type UInt64 is not currently supported"
)]
fn test_struct_array_builder_from_schema_unsupported_type() {
let fields = vec![
Field::new("f1", DataType::UInt64, false),
Field::new(
"f2",
DataType::Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8)),
false,
),
];
let _ = StructBuilder::from_fields(fields, 5);
}
#[test]
#[should_panic(expected = "Dictionary value type Int32 is not currently supported")]
fn test_struct_array_builder_from_dict_with_unsupported_value_type() {
let fields = vec![Field::new(
"f1",
DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32)),
false,
)];
let _ = StructBuilder::from_fields(fields, 5);
}
#[test]
fn test_struct_array_builder_field_builder_type_mismatch() {
let int_builder = Int32Builder::with_capacity(10);
let fields = vec![Field::new("f1", DataType::Int32, false)];
let field_builders = vec![Box::new(int_builder) as Box<dyn ArrayBuilder>];
let mut builder = StructBuilder::new(fields, field_builders);
assert!(builder.field_builder::<BinaryBuilder>(0).is_none());
}
#[test]
#[should_panic(
expected = "StructBuilder (Field { \"f1\": Int32 }, Field { \"f2\": Boolean }) and field_builder with index 1 (Boolean) are of unequal lengths: (2 != 1)."
)]
fn test_struct_array_builder_unequal_field_builders_lengths() {
let mut int_builder = Int32Builder::with_capacity(10);
let mut bool_builder = BooleanBuilder::new();
int_builder.append_value(1);
int_builder.append_value(2);
bool_builder.append_value(true);
let fields = vec![
Field::new("f1", DataType::Int32, false),
Field::new("f2", DataType::Boolean, false),
];
let field_builders = vec![
Box::new(int_builder) as Box<dyn ArrayBuilder>,
Box::new(bool_builder) as Box<dyn ArrayBuilder>,
];
let mut builder = StructBuilder::new(fields, field_builders);
builder.append(true);
builder.append(true);
builder.finish();
}
#[test]
#[should_panic(expected = "Number of fields is not equal to the number of field_builders.")]
fn test_struct_array_builder_unequal_field_field_builders() {
let int_builder = Int32Builder::with_capacity(10);
let fields = vec![
Field::new("f1", DataType::Int32, false),
Field::new("f2", DataType::Boolean, false),
];
let field_builders = vec![Box::new(int_builder) as Box<dyn ArrayBuilder>];
let mut builder = StructBuilder::new(fields, field_builders);
builder.finish();
}
#[test]
#[should_panic(
expected = "Incorrect datatype for StructArray field \\\"timestamp\\\", expected Timestamp(ns, \\\"UTC\\\") got Timestamp(ns)"
)]
fn test_struct_array_mismatch_builder() {
let fields = vec![Field::new(
"timestamp",
DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".to_owned().into())),
false,
)];
let field_builders: Vec<Box<dyn ArrayBuilder>> =
vec![Box::new(TimestampNanosecondBuilder::new())];
let mut sa = StructBuilder::new(fields, field_builders);
sa.finish();
}
#[test]
fn test_empty() {
let mut builder = StructBuilder::new(Fields::empty(), vec![]);
builder.append(true);
builder.append(false);
let a1 = builder.finish_cloned();
let a2 = builder.finish();
assert_eq!(a1, a2);
assert_eq!(a1.len(), 2);
assert_eq!(a1.null_count(), 1);
assert!(a1.is_valid(0));
assert!(a1.is_null(1));
}
}
@@ -0,0 +1,487 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::builder::buffer_builder::{Int8BufferBuilder, Int32BufferBuilder};
use crate::builder::{ArrayBuilder, BufferBuilder};
use crate::{ArrayRef, ArrowPrimitiveType, UnionArray, make_array};
use arrow_buffer::NullBufferBuilder;
use arrow_buffer::{ArrowNativeType, Buffer, ScalarBuffer};
use arrow_data::ArrayDataBuilder;
use arrow_schema::{ArrowError, DataType, Field};
use std::any::Any;
use std::collections::BTreeMap;
use std::sync::Arc;
/// `FieldData` is a helper struct to track the state of the fields in the `UnionBuilder`.
#[derive(Debug)]
struct FieldData {
/// The type id for this field
type_id: i8,
/// The Arrow data type represented in the `values_buffer`, which is untyped
data_type: DataType,
/// A buffer containing the values for this field in raw bytes
values_buffer: Box<dyn FieldDataValues>,
/// The number of array slots represented by the buffer
slots: usize,
/// A builder for the null bitmap
null_buffer_builder: NullBufferBuilder,
}
/// A type-erased [`BufferBuilder`] used by [`FieldData`]
trait FieldDataValues: std::fmt::Debug + Send + Sync {
fn as_mut_any(&mut self) -> &mut dyn Any;
fn append_null(&mut self);
fn finish(&mut self) -> Buffer;
fn finish_cloned(&self) -> Buffer;
}
impl<T: ArrowNativeType> FieldDataValues for BufferBuilder<T> {
fn as_mut_any(&mut self) -> &mut dyn Any {
self
}
fn append_null(&mut self) {
self.advance(1)
}
fn finish(&mut self) -> Buffer {
self.finish()
}
fn finish_cloned(&self) -> Buffer {
Buffer::from_slice_ref(self.as_slice())
}
}
impl FieldData {
/// Creates a new `FieldData`.
fn new<T: ArrowPrimitiveType>(type_id: i8, data_type: DataType, capacity: usize) -> Self {
Self {
type_id,
data_type,
slots: 0,
values_buffer: Box::new(BufferBuilder::<T::Native>::new(capacity)),
null_buffer_builder: NullBufferBuilder::new(capacity),
}
}
/// Appends a single value to this `FieldData`'s `values_buffer`.
fn append_value<T: ArrowPrimitiveType>(&mut self, v: T::Native) {
self.values_buffer
.as_mut_any()
.downcast_mut::<BufferBuilder<T::Native>>()
.expect("Tried to append unexpected type")
.append(v);
self.null_buffer_builder.append(true);
self.slots += 1;
}
/// Appends a null to this `FieldData`.
fn append_null(&mut self) {
self.values_buffer.append_null();
self.null_buffer_builder.append(false);
self.slots += 1;
}
}
/// Builder for [`UnionArray`]
///
/// Example: **Dense Memory Layout**
///
/// ```
/// # use arrow_array::builder::UnionBuilder;
/// # use arrow_array::types::{Float64Type, Int32Type};
///
/// let mut builder = UnionBuilder::new_dense();
/// builder.append::<Int32Type>("a", 1).unwrap();
/// builder.append::<Float64Type>("b", 3.0).unwrap();
/// builder.append::<Int32Type>("a", 4).unwrap();
/// let union = builder.build().unwrap();
///
/// assert_eq!(union.type_id(0), 0);
/// assert_eq!(union.type_id(1), 1);
/// assert_eq!(union.type_id(2), 0);
///
/// assert_eq!(union.value_offset(0), 0);
/// assert_eq!(union.value_offset(1), 0);
/// assert_eq!(union.value_offset(2), 1);
/// ```
///
/// Example: **Sparse Memory Layout**
/// ```
/// # use arrow_array::builder::UnionBuilder;
/// # use arrow_array::types::{Float64Type, Int32Type};
///
/// let mut builder = UnionBuilder::new_sparse();
/// builder.append::<Int32Type>("a", 1).unwrap();
/// builder.append::<Float64Type>("b", 3.0).unwrap();
/// builder.append::<Int32Type>("a", 4).unwrap();
/// let union = builder.build().unwrap();
///
/// assert_eq!(union.type_id(0), 0);
/// assert_eq!(union.type_id(1), 1);
/// assert_eq!(union.type_id(2), 0);
///
/// assert_eq!(union.value_offset(0), 0);
/// assert_eq!(union.value_offset(1), 1);
/// assert_eq!(union.value_offset(2), 2);
/// ```
#[derive(Debug, Default)]
pub struct UnionBuilder {
/// The current number of slots in the array
len: usize,
/// Maps field names to `FieldData` instances which track the builders for that field
fields: BTreeMap<String, FieldData>,
/// Builder to keep track of type ids
type_id_builder: Int8BufferBuilder,
/// Builder to keep track of offsets (`None` for sparse unions)
value_offset_builder: Option<Int32BufferBuilder>,
initial_capacity: usize,
}
impl UnionBuilder {
/// Creates a new dense array builder.
pub fn new_dense() -> Self {
Self::with_capacity_dense(1024)
}
/// Creates a new sparse array builder.
pub fn new_sparse() -> Self {
Self::with_capacity_sparse(1024)
}
/// Creates a new dense array builder with capacity.
pub fn with_capacity_dense(capacity: usize) -> Self {
Self {
len: 0,
fields: Default::default(),
type_id_builder: Int8BufferBuilder::new(capacity),
value_offset_builder: Some(Int32BufferBuilder::new(capacity)),
initial_capacity: capacity,
}
}
/// Creates a new sparse array builder with capacity.
pub fn with_capacity_sparse(capacity: usize) -> Self {
Self {
len: 0,
fields: Default::default(),
type_id_builder: Int8BufferBuilder::new(capacity),
value_offset_builder: None,
initial_capacity: capacity,
}
}
/// Appends a null to this builder, encoding the null in the array
/// of the `type_name` child / field.
///
/// Since `UnionArray` encodes nulls as an entry in its children
/// (it doesn't have a validity bitmap itself), and where the null
/// is part of the final array, appending a NULL requires
/// specifying which field (child) to use.
#[inline]
pub fn append_null<T: ArrowPrimitiveType>(
&mut self,
type_name: &str,
) -> Result<(), ArrowError> {
self.append_option::<T>(type_name, None)
}
/// Appends a value to this builder.
#[inline]
pub fn append<T: ArrowPrimitiveType>(
&mut self,
type_name: &str,
v: T::Native,
) -> Result<(), ArrowError> {
self.append_option::<T>(type_name, Some(v))
}
fn append_option<T: ArrowPrimitiveType>(
&mut self,
type_name: &str,
v: Option<T::Native>,
) -> Result<(), ArrowError> {
let type_name = type_name.to_string();
let mut field_data = match self.fields.remove(&type_name) {
Some(data) => {
if data.data_type != T::DATA_TYPE {
return Err(ArrowError::InvalidArgumentError(format!(
"Attempt to write col \"{}\" with type {} doesn't match existing type {}",
type_name,
T::DATA_TYPE,
data.data_type
)));
}
data
}
None => match self.value_offset_builder {
Some(_) => FieldData::new::<T>(
self.fields.len() as i8,
T::DATA_TYPE,
self.initial_capacity,
),
// In the case of a sparse union, we should pass the maximum of the currently length and the capacity.
None => {
let mut fd = FieldData::new::<T>(
self.fields.len() as i8,
T::DATA_TYPE,
self.len.max(self.initial_capacity),
);
for _ in 0..self.len {
fd.append_null();
}
fd
}
},
};
self.type_id_builder.append(field_data.type_id);
match &mut self.value_offset_builder {
// Dense Union
Some(offset_builder) => {
offset_builder.append(field_data.slots as i32);
}
// Sparse Union
None => {
for (_, fd) in self.fields.iter_mut() {
// Append to all bar the FieldData currently being appended to
fd.append_null();
}
}
}
match v {
Some(v) => field_data.append_value::<T>(v),
None => field_data.append_null(),
}
self.fields.insert(type_name, field_data);
self.len += 1;
Ok(())
}
/// Builds this builder creating a new `UnionArray`.
pub fn build(self) -> Result<UnionArray, ArrowError> {
let mut children = Vec::with_capacity(self.fields.len());
let union_fields = self
.fields
.into_iter()
.map(
|(
name,
FieldData {
type_id,
data_type,
mut values_buffer,
slots,
mut null_buffer_builder,
},
)| {
let array_ref = make_array(unsafe {
ArrayDataBuilder::new(data_type.clone())
.add_buffer(values_buffer.finish())
.len(slots)
.nulls(null_buffer_builder.finish())
.build_unchecked()
});
children.push(array_ref);
(type_id, Arc::new(Field::new(name, data_type, false)))
},
)
.collect();
UnionArray::try_new(
union_fields,
self.type_id_builder.into(),
self.value_offset_builder.map(Into::into),
children,
)
}
/// Builds this builder creating a new `UnionArray` without consuming the builder.
///
/// This is used for the `finish_cloned` implementation in `ArrayBuilder`.
fn build_cloned(&self) -> Result<UnionArray, ArrowError> {
let mut children = Vec::with_capacity(self.fields.len());
let union_fields: Vec<_> = self
.fields
.iter()
.map(|(name, field_data)| {
let FieldData {
type_id,
data_type,
values_buffer,
slots,
null_buffer_builder,
} = field_data;
let array_ref = make_array(unsafe {
ArrayDataBuilder::new(data_type.clone())
.add_buffer(values_buffer.finish_cloned())
.len(*slots)
.nulls(null_buffer_builder.finish_cloned())
.build_unchecked()
});
children.push(array_ref);
(
*type_id,
Arc::new(Field::new(name.clone(), data_type.clone(), false)),
)
})
.collect();
UnionArray::try_new(
union_fields.into_iter().collect(),
ScalarBuffer::from(self.type_id_builder.as_slice().to_vec()),
self.value_offset_builder
.as_ref()
.map(|builder| ScalarBuffer::from(builder.as_slice().to_vec())),
children,
)
}
}
impl ArrayBuilder for UnionBuilder {
/// Returns the number of array slots in the builder
fn len(&self) -> usize {
self.len
}
/// Builds the array
fn finish(&mut self) -> ArrayRef {
// Even simpler - just move the builder using mem::take and replace with default
let builder = std::mem::take(self);
// Since UnionBuilder controls all invariants, this should never fail
Arc::new(builder.build().unwrap())
}
/// Builds the array without resetting the underlying builder
fn finish_cloned(&self) -> ArrayRef {
// We construct the UnionArray carefully to ensure try_new cannot fail.
// Since UnionBuilder controls all the invariants, this should never panic.
Arc::new(self.build_cloned().unwrap_or_else(|err| {
panic!("UnionBuilder::build_cloned failed unexpectedly: {}", err)
}))
}
/// Returns the builder as a non-mutable `Any` reference
fn as_any(&self) -> &dyn Any {
self
}
/// Returns the builder as a mutable `Any` reference
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
/// Returns the boxed builder as a box of `Any`
fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::array::Array;
use crate::cast::AsArray;
use crate::types::{Float64Type, Int32Type};
#[test]
fn test_union_builder_array_builder_trait() {
// Test that UnionBuilder implements ArrayBuilder trait
let mut builder = UnionBuilder::new_dense();
// Add some data
builder.append::<Int32Type>("a", 1).unwrap();
builder.append::<Float64Type>("b", 3.0).unwrap();
builder.append::<Int32Type>("a", 4).unwrap();
assert_eq!(builder.len(), 3);
// Test finish_cloned (non-destructive)
let array1 = builder.finish_cloned();
assert_eq!(array1.len(), 3);
// Verify values in cloned array
let union1 = array1.as_any().downcast_ref::<UnionArray>().unwrap();
assert_eq!(union1.type_ids(), &[0, 1, 0]);
assert_eq!(union1.offsets().unwrap().as_ref(), &[0, 0, 1]);
let int_array1 = union1.child(0).as_primitive::<Int32Type>();
let float_array1 = union1.child(1).as_primitive::<Float64Type>();
assert_eq!(int_array1.value(0), 1);
assert_eq!(int_array1.value(1), 4);
assert_eq!(float_array1.value(0), 3.0);
// Builder should still be usable after finish_cloned
builder.append::<Float64Type>("b", 5.0).unwrap();
assert_eq!(builder.len(), 4);
// Test finish (destructive)
let array2 = builder.finish();
assert_eq!(array2.len(), 4);
// Verify values in final array
let union2 = array2.as_any().downcast_ref::<UnionArray>().unwrap();
assert_eq!(union2.type_ids(), &[0, 1, 0, 1]);
assert_eq!(union2.offsets().unwrap().as_ref(), &[0, 0, 1, 1]);
let int_array2 = union2.child(0).as_primitive::<Int32Type>();
let float_array2 = union2.child(1).as_primitive::<Float64Type>();
assert_eq!(int_array2.value(0), 1);
assert_eq!(int_array2.value(1), 4);
assert_eq!(float_array2.value(0), 3.0);
assert_eq!(float_array2.value(1), 5.0);
}
#[test]
fn test_union_builder_type_erased() {
// Test type-erased usage with Box<dyn ArrayBuilder>
let mut builders: Vec<Box<dyn ArrayBuilder>> = vec![Box::new(UnionBuilder::new_sparse())];
// Downcast and use
let union_builder = builders[0]
.as_any_mut()
.downcast_mut::<UnionBuilder>()
.unwrap();
union_builder.append::<Int32Type>("x", 10).unwrap();
union_builder.append::<Float64Type>("y", 20.0).unwrap();
assert_eq!(builders[0].len(), 2);
let result = builders
.into_iter()
.map(|mut b| b.finish())
.collect::<Vec<_>>();
assert_eq!(result[0].len(), 2);
// Verify sparse union values
let union = result[0].as_any().downcast_ref::<UnionArray>().unwrap();
assert_eq!(union.type_ids(), &[0, 1]);
assert!(union.offsets().is_none()); // Sparse union has no offsets
let int_array = union.child(0).as_primitive::<Int32Type>();
let float_array = union.child(1).as_primitive::<Float64Type>();
assert_eq!(int_array.value(0), 10);
assert!(int_array.is_null(1)); // Null in sparse layout
assert!(float_array.is_null(0)); // Null in sparse layout
assert_eq!(float_array.value(1), 20.0);
}
}
File diff suppressed because it is too large Load Diff
+285
View File
@@ -0,0 +1,285 @@
// MIT License
//
// Copyright (c) 2020-2022 Oliver Margetts
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
// Copied from chronoutil crate
//! Contains utility functions for shifting Date objects.
use chrono::{DateTime, Datelike, Days, Months, TimeZone};
use std::cmp::Ordering;
/// Shift a date by the given number of months.
pub(crate) fn shift_months<D>(date: D, months: i32) -> D
where
D: Datelike + std::ops::Add<Months, Output = D> + std::ops::Sub<Months, Output = D>,
{
match months.cmp(&0) {
Ordering::Equal => date,
Ordering::Greater => date + Months::new(months as u32),
Ordering::Less => date - Months::new(months.unsigned_abs()),
}
}
/// Add the given number of months to the given datetime.
///
/// Returns `None` when it will result in overflow.
pub(crate) fn add_months_datetime<Tz: TimeZone>(
dt: DateTime<Tz>,
months: i32,
) -> Option<DateTime<Tz>> {
match months.cmp(&0) {
Ordering::Equal => Some(dt),
Ordering::Greater => dt.checked_add_months(Months::new(months as u32)),
Ordering::Less => dt.checked_sub_months(Months::new(months.unsigned_abs())),
}
}
/// Add the given number of days to the given datetime.
///
/// Returns `None` when it will result in overflow.
pub(crate) fn add_days_datetime<Tz: TimeZone>(dt: DateTime<Tz>, days: i32) -> Option<DateTime<Tz>> {
match days.cmp(&0) {
Ordering::Equal => Some(dt),
Ordering::Greater => dt.checked_add_days(Days::new(days as u64)),
Ordering::Less => dt.checked_sub_days(Days::new(days.unsigned_abs() as u64)),
}
}
/// Substract the given number of months to the given datetime.
///
/// Returns `None` when it will result in overflow.
pub(crate) fn sub_months_datetime<Tz: TimeZone>(
dt: DateTime<Tz>,
months: i32,
) -> Option<DateTime<Tz>> {
match months.cmp(&0) {
Ordering::Equal => Some(dt),
Ordering::Greater => dt.checked_sub_months(Months::new(months as u32)),
Ordering::Less => dt.checked_add_months(Months::new(months.unsigned_abs())),
}
}
/// Substract the given number of days to the given datetime.
///
/// Returns `None` when it will result in overflow.
pub(crate) fn sub_days_datetime<Tz: TimeZone>(dt: DateTime<Tz>, days: i32) -> Option<DateTime<Tz>> {
match days.cmp(&0) {
Ordering::Equal => Some(dt),
Ordering::Greater => dt.checked_sub_days(Days::new(days as u64)),
Ordering::Less => dt.checked_add_days(Days::new(days.unsigned_abs() as u64)),
}
}
#[cfg(test)]
mod tests {
use chrono::naive::{NaiveDate, NaiveDateTime, NaiveTime};
use super::*;
#[test]
fn test_shift_months() {
let base = NaiveDate::from_ymd_opt(2020, 1, 31).unwrap();
assert_eq!(
shift_months(base, 0),
NaiveDate::from_ymd_opt(2020, 1, 31).unwrap()
);
assert_eq!(
shift_months(base, 1),
NaiveDate::from_ymd_opt(2020, 2, 29).unwrap()
);
assert_eq!(
shift_months(base, 2),
NaiveDate::from_ymd_opt(2020, 3, 31).unwrap()
);
assert_eq!(
shift_months(base, 3),
NaiveDate::from_ymd_opt(2020, 4, 30).unwrap()
);
assert_eq!(
shift_months(base, 4),
NaiveDate::from_ymd_opt(2020, 5, 31).unwrap()
);
assert_eq!(
shift_months(base, 5),
NaiveDate::from_ymd_opt(2020, 6, 30).unwrap()
);
assert_eq!(
shift_months(base, 6),
NaiveDate::from_ymd_opt(2020, 7, 31).unwrap()
);
assert_eq!(
shift_months(base, 7),
NaiveDate::from_ymd_opt(2020, 8, 31).unwrap()
);
assert_eq!(
shift_months(base, 8),
NaiveDate::from_ymd_opt(2020, 9, 30).unwrap()
);
assert_eq!(
shift_months(base, 9),
NaiveDate::from_ymd_opt(2020, 10, 31).unwrap()
);
assert_eq!(
shift_months(base, 10),
NaiveDate::from_ymd_opt(2020, 11, 30).unwrap()
);
assert_eq!(
shift_months(base, 11),
NaiveDate::from_ymd_opt(2020, 12, 31).unwrap()
);
assert_eq!(
shift_months(base, 12),
NaiveDate::from_ymd_opt(2021, 1, 31).unwrap()
);
assert_eq!(
shift_months(base, 13),
NaiveDate::from_ymd_opt(2021, 2, 28).unwrap()
);
assert_eq!(
shift_months(base, -1),
NaiveDate::from_ymd_opt(2019, 12, 31).unwrap()
);
assert_eq!(
shift_months(base, -2),
NaiveDate::from_ymd_opt(2019, 11, 30).unwrap()
);
assert_eq!(
shift_months(base, -3),
NaiveDate::from_ymd_opt(2019, 10, 31).unwrap()
);
assert_eq!(
shift_months(base, -4),
NaiveDate::from_ymd_opt(2019, 9, 30).unwrap()
);
assert_eq!(
shift_months(base, -5),
NaiveDate::from_ymd_opt(2019, 8, 31).unwrap()
);
assert_eq!(
shift_months(base, -6),
NaiveDate::from_ymd_opt(2019, 7, 31).unwrap()
);
assert_eq!(
shift_months(base, -7),
NaiveDate::from_ymd_opt(2019, 6, 30).unwrap()
);
assert_eq!(
shift_months(base, -8),
NaiveDate::from_ymd_opt(2019, 5, 31).unwrap()
);
assert_eq!(
shift_months(base, -9),
NaiveDate::from_ymd_opt(2019, 4, 30).unwrap()
);
assert_eq!(
shift_months(base, -10),
NaiveDate::from_ymd_opt(2019, 3, 31).unwrap()
);
assert_eq!(
shift_months(base, -11),
NaiveDate::from_ymd_opt(2019, 2, 28).unwrap()
);
assert_eq!(
shift_months(base, -12),
NaiveDate::from_ymd_opt(2019, 1, 31).unwrap()
);
assert_eq!(
shift_months(base, -13),
NaiveDate::from_ymd_opt(2018, 12, 31).unwrap()
);
assert_eq!(
shift_months(base, 1265),
NaiveDate::from_ymd_opt(2125, 6, 30).unwrap()
);
}
#[test]
fn test_shift_months_with_overflow() {
let base = NaiveDate::from_ymd_opt(2020, 12, 31).unwrap();
assert_eq!(shift_months(base, 0), base);
assert_eq!(
shift_months(base, 1),
NaiveDate::from_ymd_opt(2021, 1, 31).unwrap()
);
assert_eq!(
shift_months(base, 2),
NaiveDate::from_ymd_opt(2021, 2, 28).unwrap()
);
assert_eq!(
shift_months(base, 12),
NaiveDate::from_ymd_opt(2021, 12, 31).unwrap()
);
assert_eq!(
shift_months(base, 18),
NaiveDate::from_ymd_opt(2022, 6, 30).unwrap()
);
assert_eq!(
shift_months(base, -1),
NaiveDate::from_ymd_opt(2020, 11, 30).unwrap()
);
assert_eq!(
shift_months(base, -2),
NaiveDate::from_ymd_opt(2020, 10, 31).unwrap()
);
assert_eq!(
shift_months(base, -10),
NaiveDate::from_ymd_opt(2020, 2, 29).unwrap()
);
assert_eq!(
shift_months(base, -12),
NaiveDate::from_ymd_opt(2019, 12, 31).unwrap()
);
assert_eq!(
shift_months(base, -18),
NaiveDate::from_ymd_opt(2019, 6, 30).unwrap()
);
}
#[test]
fn test_shift_months_datetime() {
let date = NaiveDate::from_ymd_opt(2020, 1, 31).unwrap();
let o_clock = NaiveTime::from_hms_opt(1, 2, 3).unwrap();
let base = NaiveDateTime::new(date, o_clock);
assert_eq!(
shift_months(base, 0).date(),
NaiveDate::from_ymd_opt(2020, 1, 31).unwrap()
);
assert_eq!(
shift_months(base, 1).date(),
NaiveDate::from_ymd_opt(2020, 2, 29).unwrap()
);
assert_eq!(
shift_months(base, 2).date(),
NaiveDate::from_ymd_opt(2020, 3, 31).unwrap()
);
assert_eq!(shift_months(base, 0).time(), o_clock);
assert_eq!(shift_months(base, 1).time(), o_clock);
assert_eq!(shift_months(base, 2).time(), o_clock);
}
}
File diff suppressed because it is too large Load Diff
+555
View File
@@ -0,0 +1,555 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Contains declarations to bind to the [C Stream Interface](https://arrow.apache.org/docs/format/CStreamInterface.html).
//!
//! This module has two main interfaces:
//! One interface maps C ABI to native Rust types, i.e. convert c-pointers, c_char, to native rust.
//! This is handled by [FFI_ArrowArrayStream].
//!
//! The second interface is used to import `FFI_ArrowArrayStream` as Rust implementation `RecordBatch` reader.
//! This is handled by `ArrowArrayStreamReader`.
//!
//! ```ignore
//! # use std::fs::File;
//! # use std::sync::Arc;
//! # use arrow::error::Result;
//! # use arrow::ffi_stream::{export_reader_into_raw, ArrowArrayStreamReader, FFI_ArrowArrayStream};
//! # use arrow::ipc::reader::FileReader;
//! # use arrow::record_batch::RecordBatchReader;
//! # fn main() -> Result<()> {
//! // create an record batch reader natively
//! let file = File::open("arrow_file").unwrap();
//! let reader = Box::new(FileReader::try_new(file).unwrap());
//!
//! // export it
//! let mut stream = FFI_ArrowArrayStream::empty();
//! unsafe { export_reader_into_raw(reader, &mut stream) };
//!
//! // consumed and used by something else...
//!
//! // import it
//! let stream_reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream).unwrap() };
//! let imported_schema = stream_reader.schema();
//!
//! let mut produced_batches = vec![];
//! for batch in stream_reader {
//! produced_batches.push(batch.unwrap());
//! }
//! Ok(())
//! }
//! ```
use arrow_schema::DataType;
use std::ffi::CStr;
use std::ptr::addr_of;
use std::{
ffi::CString,
os::raw::{c_char, c_int, c_void},
sync::Arc,
};
use arrow_data::ffi::FFI_ArrowArray;
use arrow_schema::{ArrowError, Schema, SchemaRef, ffi::FFI_ArrowSchema};
use crate::array::Array;
use crate::array::StructArray;
use crate::ffi::from_ffi_and_data_type;
use crate::record_batch::{RecordBatch, RecordBatchReader};
type Result<T> = std::result::Result<T, ArrowError>;
const ENOMEM: i32 = 12;
const EIO: i32 = 5;
const EINVAL: i32 = 22;
const ENOSYS: i32 = 78;
/// ABI-compatible struct for `ArrayStream` from C Stream Interface
/// See <https://arrow.apache.org/docs/format/CStreamInterface.html#structure-definitions>
/// This was created by bindgen
#[repr(C)]
#[derive(Debug)]
#[allow(non_camel_case_types)]
pub struct FFI_ArrowArrayStream {
/// C function to get schema from the stream
pub get_schema:
Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowSchema) -> c_int>,
/// C function to get next array from the stream
pub get_next: Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowArray) -> c_int>,
/// C function to get the error from last operation on the stream
pub get_last_error: Option<unsafe extern "C" fn(arg1: *mut Self) -> *const c_char>,
/// C function to release the stream
pub release: Option<unsafe extern "C" fn(arg1: *mut Self)>,
/// Private data used by the stream
pub private_data: *mut c_void,
}
unsafe impl Send for FFI_ArrowArrayStream {}
// callback used to drop [FFI_ArrowArrayStream] when it is exported.
unsafe extern "C" fn release_stream(stream: *mut FFI_ArrowArrayStream) {
if stream.is_null() {
return;
}
let stream = unsafe { &mut *stream };
stream.get_schema = None;
stream.get_next = None;
stream.get_last_error = None;
let private_data = unsafe { Box::from_raw(stream.private_data as *mut StreamPrivateData) };
drop(private_data);
stream.release = None;
}
struct StreamPrivateData {
batch_reader: Box<dyn RecordBatchReader + Send>,
last_error: Option<CString>,
}
// The callback used to get array schema
unsafe extern "C" fn get_schema(
stream: *mut FFI_ArrowArrayStream,
schema: *mut FFI_ArrowSchema,
) -> c_int {
ExportedArrayStream { stream }.get_schema(schema)
}
// The callback used to get next array
unsafe extern "C" fn get_next(
stream: *mut FFI_ArrowArrayStream,
array: *mut FFI_ArrowArray,
) -> c_int {
ExportedArrayStream { stream }.get_next(array)
}
// The callback used to get the error from last operation on the `FFI_ArrowArrayStream`
unsafe extern "C" fn get_last_error(stream: *mut FFI_ArrowArrayStream) -> *const c_char {
let mut ffi_stream = ExportedArrayStream { stream };
// The consumer should not take ownership of this string, we should return
// a const pointer to it.
match ffi_stream.get_last_error() {
Some(err_string) => err_string.as_ptr(),
None => std::ptr::null(),
}
}
impl Drop for FFI_ArrowArrayStream {
fn drop(&mut self) {
match self.release {
None => (),
Some(release) => unsafe { release(self) },
};
}
}
impl FFI_ArrowArrayStream {
/// Creates a new [`FFI_ArrowArrayStream`].
pub fn new(batch_reader: Box<dyn RecordBatchReader + Send>) -> Self {
let private_data = Box::new(StreamPrivateData {
batch_reader,
last_error: None,
});
Self {
get_schema: Some(get_schema),
get_next: Some(get_next),
get_last_error: Some(get_last_error),
release: Some(release_stream),
private_data: Box::into_raw(private_data) as *mut c_void,
}
}
/// Takes ownership of the pointed to [`FFI_ArrowArrayStream`]
///
/// This acts to [move] the data out of `raw_stream`, setting the release callback to NULL
///
/// # Safety
///
/// * `raw_stream` must be [valid] for reads and writes
/// * `raw_stream` must be properly aligned
/// * `raw_stream` must point to a properly initialized value of [`FFI_ArrowArrayStream`]
///
/// [move]: https://arrow.apache.org/docs/format/CDataInterface.html#moving-an-array
/// [valid]: https://doc.rust-lang.org/std/ptr/index.html#safety
pub unsafe fn from_raw(raw_stream: *mut FFI_ArrowArrayStream) -> Self {
unsafe { std::ptr::replace(raw_stream, Self::empty()) }
}
/// Creates a new empty [FFI_ArrowArrayStream]. Used to import from the C Stream Interface.
pub fn empty() -> Self {
Self {
get_schema: None,
get_next: None,
get_last_error: None,
release: None,
private_data: std::ptr::null_mut(),
}
}
}
struct ExportedArrayStream {
stream: *mut FFI_ArrowArrayStream,
}
impl ExportedArrayStream {
fn get_private_data(&mut self) -> &mut StreamPrivateData {
unsafe { &mut *((*self.stream).private_data as *mut StreamPrivateData) }
}
pub fn get_schema(&mut self, out: *mut FFI_ArrowSchema) -> i32 {
let private_data = self.get_private_data();
let reader = &private_data.batch_reader;
let schema = FFI_ArrowSchema::try_from(reader.schema().as_ref());
match schema {
Ok(schema) => {
unsafe { std::ptr::copy(addr_of!(schema), out, 1) };
std::mem::forget(schema);
0
}
Err(ref err) => {
private_data.last_error = Some(
CString::new(err.to_string()).expect("Error string has a null byte in it."),
);
get_error_code(err)
}
}
}
pub fn get_next(&mut self, out: *mut FFI_ArrowArray) -> i32 {
let private_data = self.get_private_data();
let reader = &mut private_data.batch_reader;
match reader.next() {
None => {
// Marks ArrowArray released to indicate reaching the end of stream.
unsafe { std::ptr::write(out, FFI_ArrowArray::empty()) }
0
}
Some(next_batch) => {
if let Ok(batch) = next_batch {
let struct_array = StructArray::from(batch);
let array = FFI_ArrowArray::new(&struct_array.to_data());
unsafe { std::ptr::write_unaligned(out, array) };
0
} else {
let err = &next_batch.unwrap_err();
private_data.last_error = Some(
CString::new(err.to_string()).expect("Error string has a null byte in it."),
);
get_error_code(err)
}
}
}
}
pub fn get_last_error(&mut self) -> Option<&CString> {
self.get_private_data().last_error.as_ref()
}
}
fn get_error_code(err: &ArrowError) -> i32 {
match err {
ArrowError::NotYetImplemented(_) => ENOSYS,
ArrowError::MemoryError(_) => ENOMEM,
ArrowError::IoError(_, _) => EIO,
_ => EINVAL,
}
}
/// A `RecordBatchReader` which imports Arrays from `FFI_ArrowArrayStream`.
///
/// Struct used to fetch `RecordBatch` from the C Stream Interface.
/// Its main responsibility is to expose `RecordBatchReader` functionality
/// that requires [FFI_ArrowArrayStream].
#[derive(Debug)]
pub struct ArrowArrayStreamReader {
stream: FFI_ArrowArrayStream,
schema: SchemaRef,
}
/// Gets schema from a raw pointer of `FFI_ArrowArrayStream`. This is used when constructing
/// `ArrowArrayStreamReader` to cache schema.
fn get_stream_schema(stream_ptr: *mut FFI_ArrowArrayStream) -> Result<SchemaRef> {
let mut schema = FFI_ArrowSchema::empty();
let ret_code = unsafe { (*stream_ptr).get_schema.unwrap()(stream_ptr, &mut schema) };
if ret_code == 0 {
let schema = Schema::try_from(&schema)?;
Ok(Arc::new(schema))
} else {
Err(ArrowError::CDataInterface(format!(
"Cannot get schema from input stream. Error code: {ret_code:?}"
)))
}
}
impl ArrowArrayStreamReader {
/// Creates a new `ArrowArrayStreamReader` from a `FFI_ArrowArrayStream`.
/// This is used to import from the C Stream Interface.
#[allow(dead_code)]
pub fn try_new(mut stream: FFI_ArrowArrayStream) -> Result<Self> {
if stream.release.is_none() {
return Err(ArrowError::CDataInterface(
"input stream is already released".to_string(),
));
}
let schema = get_stream_schema(&mut stream)?;
Ok(Self { stream, schema })
}
/// Creates a new `ArrowArrayStreamReader` from a raw pointer of `FFI_ArrowArrayStream`.
///
/// Assumes that the pointer represents valid C Stream Interfaces.
/// This function copies the content from the raw pointer and cleans up it to prevent
/// double-dropping. The caller is responsible for freeing up the memory allocated for
/// the pointer.
///
/// # Safety
///
/// See [`FFI_ArrowArrayStream::from_raw`]
pub unsafe fn from_raw(raw_stream: *mut FFI_ArrowArrayStream) -> Result<Self> {
Self::try_new(unsafe { FFI_ArrowArrayStream::from_raw(raw_stream) })
}
/// Get the last error from `ArrowArrayStreamReader`
fn get_stream_last_error(&mut self) -> Option<String> {
let get_last_error = self.stream.get_last_error?;
let error_str = unsafe { get_last_error(&mut self.stream) };
if error_str.is_null() {
return None;
}
let error_str = unsafe { CStr::from_ptr(error_str) };
Some(error_str.to_string_lossy().to_string())
}
}
impl Iterator for ArrowArrayStreamReader {
type Item = Result<RecordBatch>;
fn next(&mut self) -> Option<Self::Item> {
let mut array = FFI_ArrowArray::empty();
let ret_code = unsafe { self.stream.get_next.unwrap()(&mut self.stream, &mut array) };
if ret_code == 0 {
// The end of stream has been reached
if array.is_released() {
return None;
}
let result = unsafe {
from_ffi_and_data_type(array, DataType::Struct(self.schema().fields().clone()))
};
Some(result.and_then(|data| {
RecordBatch::try_new(self.schema.clone(), StructArray::from(data).into_parts().1)
}))
} else {
let last_error = self.get_stream_last_error();
let err = ArrowError::CDataInterface(last_error.unwrap());
Some(Err(err))
}
}
}
impl RecordBatchReader for ArrowArrayStreamReader {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
use arrow_schema::Field;
use crate::array::Int32Array;
use crate::ffi::from_ffi;
struct TestRecordBatchReader {
schema: SchemaRef,
iter: Box<dyn Iterator<Item = Result<RecordBatch>> + Send>,
}
impl TestRecordBatchReader {
pub fn new(
schema: SchemaRef,
iter: Box<dyn Iterator<Item = Result<RecordBatch>> + Send>,
) -> Box<TestRecordBatchReader> {
Box::new(TestRecordBatchReader { schema, iter })
}
}
impl Iterator for TestRecordBatchReader {
type Item = Result<RecordBatch>;
fn next(&mut self) -> Option<Self::Item> {
self.iter.next()
}
}
impl RecordBatchReader for TestRecordBatchReader {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
}
fn _test_round_trip_export(arrays: Vec<Arc<dyn Array>>) -> Result<()> {
let metadata = HashMap::from([("foo".to_owned(), "bar".to_owned())]);
let schema = Arc::new(Schema::new_with_metadata(
vec![
Field::new("a", arrays[0].data_type().clone(), true)
.with_metadata(metadata.clone()),
Field::new("b", arrays[1].data_type().clone(), true)
.with_metadata(metadata.clone()),
Field::new("c", arrays[2].data_type().clone(), true)
.with_metadata(metadata.clone()),
],
metadata,
));
let batch = RecordBatch::try_new(schema.clone(), arrays).unwrap();
let iter = Box::new(vec![batch.clone(), batch.clone()].into_iter().map(Ok)) as _;
let reader = TestRecordBatchReader::new(schema.clone(), iter);
// Export a `RecordBatchReader` through `FFI_ArrowArrayStream`
let mut ffi_stream = FFI_ArrowArrayStream::new(reader);
// Get schema from `FFI_ArrowArrayStream`
let mut ffi_schema = FFI_ArrowSchema::empty();
let ret_code = unsafe { get_schema(&mut ffi_stream, &mut ffi_schema) };
assert_eq!(ret_code, 0);
let exported_schema = Schema::try_from(&ffi_schema).unwrap();
assert_eq!(&exported_schema, schema.as_ref());
// Get array from `FFI_ArrowArrayStream`
let mut produced_batches = vec![];
loop {
let mut ffi_array = FFI_ArrowArray::empty();
let ret_code = unsafe { get_next(&mut ffi_stream, &mut ffi_array) };
assert_eq!(ret_code, 0);
// The end of stream has been reached
if ffi_array.is_released() {
break;
}
let array = unsafe { from_ffi(ffi_array, &ffi_schema) }.unwrap();
let record_batch = RecordBatch::try_new(
SchemaRef::from(exported_schema.clone()),
StructArray::from(array).into_parts().1,
)
.unwrap();
produced_batches.push(record_batch);
}
assert_eq!(produced_batches, vec![batch.clone(), batch]);
Ok(())
}
fn _test_round_trip_import(arrays: Vec<Arc<dyn Array>>) -> Result<()> {
let metadata = HashMap::from([("foo".to_owned(), "bar".to_owned())]);
let schema = Arc::new(Schema::new_with_metadata(
vec![
Field::new("a", arrays[0].data_type().clone(), true)
.with_metadata(metadata.clone()),
Field::new("b", arrays[1].data_type().clone(), true)
.with_metadata(metadata.clone()),
Field::new("c", arrays[2].data_type().clone(), true)
.with_metadata(metadata.clone()),
],
metadata,
));
let batch = RecordBatch::try_new(schema.clone(), arrays).unwrap();
let iter = Box::new(vec![batch.clone(), batch.clone()].into_iter().map(Ok)) as _;
let reader = TestRecordBatchReader::new(schema.clone(), iter);
// Import through `FFI_ArrowArrayStream` as `ArrowArrayStreamReader`
let stream = FFI_ArrowArrayStream::new(reader);
let stream_reader = ArrowArrayStreamReader::try_new(stream).unwrap();
let imported_schema = stream_reader.schema();
assert_eq!(imported_schema, schema);
let mut produced_batches = vec![];
for batch in stream_reader {
produced_batches.push(batch.unwrap());
}
assert_eq!(produced_batches, vec![batch.clone(), batch]);
Ok(())
}
#[test]
fn test_stream_round_trip_export() -> Result<()> {
let array = Int32Array::from(vec![Some(2), None, Some(1), None]);
let array: Arc<dyn Array> = Arc::new(array);
_test_round_trip_export(vec![array.clone(), array.clone(), array])
}
#[test]
fn test_stream_round_trip_import() -> Result<()> {
let array = Int32Array::from(vec![Some(2), None, Some(1), None]);
let array: Arc<dyn Array> = Arc::new(array);
_test_round_trip_import(vec![array.clone(), array.clone(), array])
}
#[test]
fn test_error_import() -> Result<()> {
let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
let iter = Box::new(vec![Err(ArrowError::MemoryError("".to_string()))].into_iter());
let reader = TestRecordBatchReader::new(schema.clone(), iter);
// Import through `FFI_ArrowArrayStream` as `ArrowArrayStreamReader`
let stream = FFI_ArrowArrayStream::new(reader);
let stream_reader = ArrowArrayStreamReader::try_new(stream).unwrap();
let imported_schema = stream_reader.schema();
assert_eq!(imported_schema, schema);
let mut produced_batches = vec![];
for batch in stream_reader {
produced_batches.push(batch);
}
// The results should outlive the lifetime of the stream itself.
assert_eq!(produced_batches.len(), 1);
assert!(produced_batches[0].is_err());
Ok(())
}
}
File diff suppressed because it is too large Load Diff
+295
View File
@@ -0,0 +1,295 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! The central type in Apache Arrow are arrays, which are a known-length sequence of values
//! all having the same type. This crate provides concrete implementations of each type, as
//! well as an [`Array`] trait that can be used for type-erasure.
//!
//! # Building an Array
//!
//! Most [`Array`] implementations can be constructed directly from iterators or [`Vec`]
//!
//! ```
//! # use arrow_array::{Int32Array, ListArray, StringArray};
//! # use arrow_array::types::Int32Type;
//! #
//! Int32Array::from(vec![1, 2]);
//! Int32Array::from(vec![Some(1), None]);
//! Int32Array::from_iter([1, 2, 3, 4]);
//! Int32Array::from_iter([Some(1), Some(2), None, Some(4)]);
//!
//! StringArray::from(vec!["foo", "bar"]);
//! StringArray::from(vec![Some("foo"), None]);
//! StringArray::from_iter([Some("foo"), None]);
//! StringArray::from_iter_values(["foo", "bar"]);
//!
//! ListArray::from_iter_primitive::<Int32Type, _, _>([
//! Some(vec![Some(1), None, Some(3)]),
//! None,
//! Some(vec![])
//! ]);
//! ```
//!
//! Additionally [`ArrayBuilder`](builder::ArrayBuilder) implementations can be
//! used to construct arrays with a push-based interface
//!
//! ```
//! # use arrow_array::Int16Array;
//! #
//! // Create a new builder with a capacity of 100
//! let mut builder = Int16Array::builder(100);
//!
//! // Append a single primitive value
//! builder.append_value(1);
//! // Append a null value
//! builder.append_null();
//! // Append a slice of primitive values
//! builder.append_slice(&[2, 3, 4]);
//!
//! // Build the array
//! let array = builder.finish();
//!
//! assert_eq!(5, array.len());
//! assert_eq!(2, array.value(2));
//! assert_eq!(&array.values()[3..5], &[3, 4])
//! ```
//!
//! # Low-level API
//!
//! Internally, arrays consist of one or more shared memory regions backed by a [`Buffer`],
//! the number and meaning of which depend on the arrays data type, as documented in
//! the [Arrow specification].
//!
//! For example, the type [`Int16Array`] represents an array of 16-bit integers and consists of:
//!
//! * An optional [`NullBuffer`] identifying any null values
//! * A contiguous [`ScalarBuffer<i16>`] of values
//!
//! Similarly, the type [`StringArray`] represents an array of UTF-8 strings and consists of:
//!
//! * An optional [`NullBuffer`] identifying any null values
//! * An offsets [`OffsetBuffer<i32>`] identifying valid UTF-8 sequences within the values buffer
//! * A values [`Buffer`] of UTF-8 encoded string data
//!
//! Array constructors such as [`PrimitiveArray::try_new`] provide the ability to cheaply
//! construct an array from these parts, with functions such as [`PrimitiveArray::into_parts`]
//! providing the reverse operation.
//!
//! ```
//! # use arrow_array::{Array, Int32Array, StringArray};
//! # use arrow_buffer::OffsetBuffer;
//! #
//! // Create a Int32Array from Vec without copying
//! let array = Int32Array::new(vec![1, 2, 3].into(), None);
//! assert_eq!(array.values(), &[1, 2, 3]);
//! assert_eq!(array.null_count(), 0);
//!
//! // Create a StringArray from parts
//! let offsets = OffsetBuffer::new(vec![0, 5, 10].into());
//! let array = StringArray::new(offsets, b"helloworld".into(), None);
//! let values: Vec<_> = array.iter().map(|x| x.unwrap()).collect();
//! assert_eq!(values, &["hello", "world"]);
//! ```
//!
//! As [`Buffer`], and its derivatives, can be created from [`Vec`] without copying, this provides
//! an efficient way to not only interoperate with other Rust code, but also implement kernels
//! optimised for the arrow data layout - e.g. by handling buffers instead of values.
//!
//! # Zero-Copy Slicing
//!
//! Given an [`Array`] of arbitrary length, it is possible to create an owned slice of this
//! data. Internally this just increments some ref-counts, and so is incredibly cheap
//!
//! ```rust
//! # use arrow_array::Int32Array;
//! let array = Int32Array::from_iter([1, 2, 3]);
//!
//! // Slice with offset 1 and length 2
//! let sliced = array.slice(1, 2);
//! assert_eq!(sliced.values(), &[2, 3]);
//! ```
//!
//! # Downcasting an Array
//!
//! Arrays are often passed around as a dynamically typed [`&dyn Array`] or [`ArrayRef`].
//! For example, [`RecordBatch`](`crate::RecordBatch`) stores columns as [`ArrayRef`].
//!
//! Whilst these arrays can be passed directly to the [`compute`], [`csv`], [`json`], etc... APIs,
//! it is often the case that you wish to interact with the concrete arrays directly.
//!
//! This requires downcasting to the concrete type of the array:
//!
//! ```
//! # use arrow_array::{Array, Float32Array, Int32Array};
//!
//! // Safely downcast an `Array` to an `Int32Array` and compute the sum
//! // using native i32 values
//! fn sum_int32(array: &dyn Array) -> i32 {
//! let integers: &Int32Array = array.as_any().downcast_ref().unwrap();
//! integers.iter().map(|val| val.unwrap_or_default()).sum()
//! }
//!
//! // Safely downcasts the array to a `Float32Array` and returns a &[f32] view of the data
//! // Note: the values for positions corresponding to nulls will be arbitrary (but still valid f32)
//! fn as_f32_slice(array: &dyn Array) -> &[f32] {
//! array.as_any().downcast_ref::<Float32Array>().unwrap().values()
//! }
//! ```
//!
//! The [`cast::AsArray`] extension trait can make this more ergonomic
//!
//! ```
//! # use arrow_array::Array;
//! # use arrow_array::cast::{AsArray, as_primitive_array};
//! # use arrow_array::types::Float32Type;
//!
//! fn as_f32_slice(array: &dyn Array) -> &[f32] {
//! array.as_primitive::<Float32Type>().values()
//! }
//! ```
//! # Alternatives to ChunkedArray Support
//!
//! The Rust implementation does not provide the ChunkedArray abstraction implemented by the Python
//! and C++ Arrow implementations. The recommended alternative is to use one of the following:
//! - `Vec<ArrayRef>` a simple, eager version of a `ChunkedArray`
//! - `impl Iterator<Item=ArrayRef>` a lazy version of a `ChunkedArray`
//! - `impl Stream<Item=ArrayRef>` a lazy async version of a `ChunkedArray`
//!
//! Similar patterns can be applied at the `RecordBatch` level. For example, [DataFusion] makes
//! extensive use of [RecordBatchStream].
//!
//! This approach integrates well into the Rust ecosystem, simplifies the implementation and
//! encourages the use of performant lazy and async patterns.
//! ```rust
//! use std::sync::Arc;
//! use arrow_array::{ArrayRef, Float32Array, RecordBatch, StringArray};
//! use arrow_array::cast::AsArray;
//! use arrow_array::types::Float32Type;
//! use arrow_schema::DataType;
//!
//! let batches = [
//! RecordBatch::try_from_iter(vec![
//! ("label", Arc::new(StringArray::from(vec!["A", "B", "C"])) as ArrayRef),
//! ("value", Arc::new(Float32Array::from(vec![0.1, 0.2, 0.3])) as ArrayRef),
//! ]).unwrap(),
//! RecordBatch::try_from_iter(vec![
//! ("label", Arc::new(StringArray::from(vec!["D", "E"])) as ArrayRef),
//! ("value", Arc::new(Float32Array::from(vec![0.4, 0.5])) as ArrayRef),
//! ]).unwrap(),
//! ];
//!
//! let labels: Vec<&str> = batches
//! .iter()
//! .flat_map(|batch| batch.column(0).as_string::<i32>())
//! .map(Option::unwrap)
//! .collect();
//!
//! let values: Vec<f32> = batches
//! .iter()
//! .flat_map(|batch| batch.column(1).as_primitive::<Float32Type>().values())
//! .copied()
//! .collect();
//!
//! assert_eq!(labels, ["A", "B", "C", "D", "E"]);
//! assert_eq!(values, [0.1, 0.2, 0.3, 0.4, 0.5]);
//!```
//! [`ScalarBuffer<T>`]: arrow_buffer::ScalarBuffer
//! [`ScalarBuffer<i16>`]: arrow_buffer::ScalarBuffer
//! [`OffsetBuffer<i32>`]: arrow_buffer::OffsetBuffer
//! [`NullBuffer`]: arrow_buffer::NullBuffer
//! [Arrow specification]: https://arrow.apache.org/docs/format/Columnar.html
//! [`&dyn Array`]: Array
//! [`NullBuffer`]: arrow_buffer::NullBuffer
//! [`Buffer`]: arrow_buffer::Buffer
//! [`compute`]: https://docs.rs/arrow/latest/arrow/compute/index.html
//! [`json`]: https://docs.rs/arrow/latest/arrow/json/index.html
//! [`csv`]: https://docs.rs/arrow/latest/arrow/csv/index.html
//! [DataFusion]: https://github.com/apache/arrow-datafusion
//! [RecordBatchStream]: https://docs.rs/datafusion/latest/datafusion/execution/trait.RecordBatchStream.html
#![doc(
html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg",
html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg"
)]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![deny(rustdoc::broken_intra_doc_links)]
#![warn(missing_docs)]
pub mod array;
pub use array::*;
mod record_batch;
pub use record_batch::{
RecordBatch, RecordBatchIterator, RecordBatchOptions, RecordBatchReader, RecordBatchWriter,
};
mod arithmetic;
pub use arithmetic::ArrowNativeTypeOp;
mod numeric;
pub use numeric::*;
mod scalar;
pub use scalar::*;
pub mod builder;
pub mod cast;
mod delta;
#[cfg(feature = "ffi")]
pub mod ffi;
#[cfg(feature = "ffi")]
pub mod ffi_stream;
pub mod iterator;
pub mod run_iterator;
pub mod temporal_conversions;
pub mod timezone;
mod trusted_len;
pub mod types;
#[cfg(test)]
mod tests {
use crate::builder::*;
#[test]
fn test_buffer_builder_availability() {
let _builder = Int8BufferBuilder::new(10);
let _builder = Int16BufferBuilder::new(10);
let _builder = Int32BufferBuilder::new(10);
let _builder = Int64BufferBuilder::new(10);
let _builder = UInt16BufferBuilder::new(10);
let _builder = UInt32BufferBuilder::new(10);
let _builder = Float32BufferBuilder::new(10);
let _builder = Float64BufferBuilder::new(10);
let _builder = TimestampSecondBufferBuilder::new(10);
let _builder = TimestampMillisecondBufferBuilder::new(10);
let _builder = TimestampMicrosecondBufferBuilder::new(10);
let _builder = TimestampNanosecondBufferBuilder::new(10);
let _builder = Date32BufferBuilder::new(10);
let _builder = Date64BufferBuilder::new(10);
let _builder = Time32SecondBufferBuilder::new(10);
let _builder = Time32MillisecondBufferBuilder::new(10);
let _builder = Time64MicrosecondBufferBuilder::new(10);
let _builder = Time64NanosecondBufferBuilder::new(10);
let _builder = IntervalYearMonthBufferBuilder::new(10);
let _builder = IntervalDayTimeBufferBuilder::new(10);
let _builder = IntervalMonthDayNanoBufferBuilder::new(10);
let _builder = DurationSecondBufferBuilder::new(10);
let _builder = DurationMillisecondBufferBuilder::new(10);
let _builder = DurationMicrosecondBufferBuilder::new(10);
let _builder = DurationNanosecondBufferBuilder::new(10);
}
}
+23
View File
@@ -0,0 +1,23 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::ArrowPrimitiveType;
/// A subtype of primitive type that represents numeric values.
pub trait ArrowNumericType: ArrowPrimitiveType {}
impl<T: ArrowPrimitiveType> ArrowNumericType for T {}
File diff suppressed because it is too large Load Diff
+384
View File
@@ -0,0 +1,384 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Idiomatic iterator for [`RunArray`](crate::RunArray)
use crate::{Array, TypedRunArray, array::ArrayAccessor, types::RunEndIndexType};
use arrow_buffer::ArrowNativeType;
/// The [`RunArrayIter`] provides an idiomatic way to iterate over the run array.
/// It returns Some(T) if there is a value or None if the value is null.
///
/// The iterator comes with a cost as it has to iterate over three arrays to determine
/// the value to be returned. The run_ends array is used to determine the index of the value.
/// The nulls array is used to determine if the value is null and the values array is used to
/// get the value.
///
/// Unlike other iterators in this crate, [`RunArrayIter`] does not use [`ArrayAccessor`]
/// because the run array accessor does binary search to access each value which is too slow.
/// The run array iterator can determine the next value in constant time.
///
#[derive(Debug)]
pub struct RunArrayIter<'a, R, V>
where
R: RunEndIndexType,
V: Sync + Send,
&'a V: ArrayAccessor,
<&'a V as ArrayAccessor>::Item: Default,
{
array: TypedRunArray<'a, R, V>,
current_front_logical: usize,
current_front_physical: usize,
current_back_logical: usize,
current_back_physical: usize,
}
impl<'a, R, V> RunArrayIter<'a, R, V>
where
R: RunEndIndexType,
V: Sync + Send,
&'a V: ArrayAccessor,
<&'a V as ArrayAccessor>::Item: Default,
{
/// create a new iterator
pub fn new(array: TypedRunArray<'a, R, V>) -> Self {
let current_front_physical = array.run_array().get_start_physical_index();
let current_back_physical = array.run_array().get_end_physical_index() + 1;
RunArrayIter {
array,
current_front_logical: array.offset(),
current_front_physical,
current_back_logical: array.offset() + array.len(),
current_back_physical,
}
}
}
impl<'a, R, V> Iterator for RunArrayIter<'a, R, V>
where
R: RunEndIndexType,
V: Sync + Send,
&'a V: ArrayAccessor,
<&'a V as ArrayAccessor>::Item: Default,
{
type Item = Option<<&'a V as ArrayAccessor>::Item>;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
if self.current_front_logical == self.current_back_logical {
return None;
}
// If current logical index is greater than current run end index then increment
// the physical index.
let run_ends = self.array.run_ends().values();
if self.current_front_logical >= run_ends[self.current_front_physical].as_usize() {
// As the run_ends is expected to be strictly increasing, there
// should be at least one logical entry in one physical entry. Because of this
// reason the next value can be accessed by incrementing physical index once.
self.current_front_physical += 1;
}
if self.array.values().is_null(self.current_front_physical) {
self.current_front_logical += 1;
Some(None)
} else {
self.current_front_logical += 1;
// Safety:
// The self.current_physical is kept within bounds of self.current_logical.
// The self.current_logical will not go out of bounds because of the check
// `self.current_logical = self.current_end_logical` above.
unsafe {
Some(Some(
self.array
.values()
.value_unchecked(self.current_front_physical),
))
}
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
(
self.current_back_logical - self.current_front_logical,
Some(self.current_back_logical - self.current_front_logical),
)
}
}
impl<'a, R, V> DoubleEndedIterator for RunArrayIter<'a, R, V>
where
R: RunEndIndexType,
V: Sync + Send,
&'a V: ArrayAccessor,
<&'a V as ArrayAccessor>::Item: Default,
{
fn next_back(&mut self) -> Option<Self::Item> {
if self.current_back_logical == self.current_front_logical {
return None;
}
self.current_back_logical -= 1;
let run_ends = self.array.run_ends().values();
if self.current_back_physical > 0
&& self.current_back_logical < run_ends[self.current_back_physical - 1].as_usize()
{
// As the run_ends is expected to be strictly increasing, there
// should be at least one logical entry in one physical entry. Because of this
// reason the next value can be accessed by decrementing physical index once.
self.current_back_physical -= 1;
}
Some(if self.array.values().is_null(self.current_back_physical) {
None
} else {
// Safety:
// The check `self.current_end_physical > 0` ensures the value will not underflow.
// Also self.current_end_physical starts with array.len() and
// decrements based on the bounds of self.current_end_logical.
unsafe {
Some(
self.array
.values()
.value_unchecked(self.current_back_physical),
)
}
})
}
}
/// all arrays have known size.
impl<'a, R, V> ExactSizeIterator for RunArrayIter<'a, R, V>
where
R: RunEndIndexType,
V: Sync + Send,
&'a V: ArrayAccessor,
<&'a V as ArrayAccessor>::Item: Default,
{
}
#[cfg(test)]
mod tests {
use rand::{Rng, rng, seq::SliceRandom};
use crate::{
Array, Int64RunArray, PrimitiveArray, RunArray,
array::{Int32Array, StringArray},
builder::PrimitiveRunBuilder,
types::{Int16Type, Int32Type},
};
fn build_input_array(size: usize) -> Vec<Option<i32>> {
// The input array is created by shuffling and repeating
// the seed values random number of times.
let mut seed: Vec<Option<i32>> = vec![
None,
None,
None,
Some(1),
Some(2),
Some(3),
Some(4),
Some(5),
Some(6),
Some(7),
Some(8),
Some(9),
];
let mut result: Vec<Option<i32>> = Vec::with_capacity(size);
let mut ix = 0;
let mut rng = rng();
// run length can go up to 8. Cap the max run length for smaller arrays to size / 2.
let max_run_length = 8_usize.min(1_usize.max(size / 2));
while result.len() < size {
// shuffle the seed array if all the values are iterated.
if ix == 0 {
seed.shuffle(&mut rng);
}
// repeat the items between 1 and 8 times. Cap the length for smaller sized arrays
let num = max_run_length.min(rng.random_range(1..=max_run_length));
for _ in 0..num {
result.push(seed[ix]);
}
ix += 1;
if ix == seed.len() {
ix = 0
}
}
result.resize(size, None);
result
}
#[test]
fn test_primitive_array_iter_round_trip() {
let mut input_vec = vec![
Some(32),
Some(32),
None,
Some(64),
Some(64),
Some(64),
Some(72),
];
let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
builder.extend(input_vec.iter().copied());
let ree_array = builder.finish();
let ree_array = ree_array.downcast::<Int32Array>().unwrap();
let output_vec: Vec<Option<i32>> = ree_array.into_iter().collect();
assert_eq!(input_vec, output_vec);
let rev_output_vec: Vec<Option<i32>> = ree_array.into_iter().rev().collect();
input_vec.reverse();
assert_eq!(input_vec, rev_output_vec);
}
#[test]
fn test_double_ended() {
let input_vec = vec![
Some(32),
Some(32),
None,
Some(64),
Some(64),
Some(64),
Some(72),
];
let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
builder.extend(input_vec);
let ree_array = builder.finish();
let ree_array = ree_array.downcast::<Int32Array>().unwrap();
let mut iter = ree_array.into_iter();
assert_eq!(Some(Some(32)), iter.next());
assert_eq!(Some(Some(72)), iter.next_back());
assert_eq!(Some(Some(32)), iter.next());
assert_eq!(Some(Some(64)), iter.next_back());
assert_eq!(Some(None), iter.next());
assert_eq!(Some(Some(64)), iter.next_back());
assert_eq!(Some(Some(64)), iter.next());
assert_eq!(None, iter.next_back());
assert_eq!(None, iter.next());
}
#[test]
fn test_run_iterator_comprehensive() {
// Test forward and backward iterator for different array lengths.
let logical_lengths = vec![1_usize, 2, 3, 4, 15, 16, 17, 63, 64, 65];
for logical_len in logical_lengths {
let input_array = build_input_array(logical_len);
let mut run_array_builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
run_array_builder.extend(input_array.iter().copied());
let run_array = run_array_builder.finish();
let typed_array = run_array.downcast::<Int32Array>().unwrap();
// test forward iterator
let mut input_iter = input_array.iter().copied();
let mut run_array_iter = typed_array.into_iter();
for _ in 0..logical_len {
assert_eq!(input_iter.next(), run_array_iter.next());
}
assert_eq!(None, run_array_iter.next());
// test reverse iterator
let mut input_iter = input_array.iter().rev().copied();
let mut run_array_iter = typed_array.into_iter().rev();
for _ in 0..logical_len {
assert_eq!(input_iter.next(), run_array_iter.next());
}
assert_eq!(None, run_array_iter.next());
}
}
#[test]
fn test_string_array_iter_round_trip() {
let input_vec = vec!["ab", "ab", "ba", "cc", "cc"];
let input_ree_array: Int64RunArray = input_vec.into_iter().collect();
let string_ree_array = input_ree_array.downcast::<StringArray>().unwrap();
// to and from iter, with a +1
let result: Vec<Option<String>> = string_ree_array
.into_iter()
.map(|e| {
e.map(|e| {
let mut a = e.to_string();
a.push('b');
a
})
})
.collect();
let result_asref: Vec<Option<&str>> = result.iter().map(|f| f.as_deref()).collect();
let expected_vec = vec![
Some("abb"),
Some("abb"),
Some("bab"),
Some("ccb"),
Some("ccb"),
];
assert_eq!(expected_vec, result_asref);
}
#[test]
#[cfg_attr(miri, ignore)] // Takes too long
fn test_sliced_run_array_iterator() {
let total_len = 80;
let input_array = build_input_array(total_len);
// Encode the input_array to run array
let mut builder =
PrimitiveRunBuilder::<Int16Type, Int32Type>::with_capacity(input_array.len());
builder.extend(input_array.iter().copied());
let run_array = builder.finish();
// test for all slice lengths.
for slice_len in 1..=total_len {
// test for offset = 0, slice length = slice_len
let sliced_run_array: RunArray<Int16Type> =
run_array.slice(0, slice_len).into_data().into();
let sliced_typed_run_array = sliced_run_array
.downcast::<PrimitiveArray<Int32Type>>()
.unwrap();
// Iterate on sliced typed run array
let actual: Vec<Option<i32>> = sliced_typed_run_array.into_iter().collect();
let expected: Vec<Option<i32>> = input_array.iter().take(slice_len).copied().collect();
assert_eq!(expected, actual);
// test for offset = total_len - slice_len, length = slice_len
let sliced_run_array: RunArray<Int16Type> = run_array
.slice(total_len - slice_len, slice_len)
.into_data()
.into();
let sliced_typed_run_array = sliced_run_array
.downcast::<PrimitiveArray<Int32Type>>()
.unwrap();
// Iterate on sliced typed run array
let actual: Vec<Option<i32>> = sliced_typed_run_array.into_iter().collect();
let expected: Vec<Option<i32>> = input_array
.iter()
.skip(total_len - slice_len)
.copied()
.collect();
assert_eq!(expected, actual);
}
}
}
+152
View File
@@ -0,0 +1,152 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::Array;
/// A possibly [`Scalar`] [`Array`]
///
/// This allows optimised binary kernels where one or more arguments are constant
///
/// ```
/// # use arrow_array::*;
/// # use arrow_buffer::{BooleanBuffer, MutableBuffer, NullBuffer};
/// # use arrow_schema::ArrowError;
/// #
/// fn eq_impl<T: ArrowPrimitiveType>(
/// a: &PrimitiveArray<T>,
/// a_scalar: bool,
/// b: &PrimitiveArray<T>,
/// b_scalar: bool,
/// ) -> BooleanArray {
/// let (array, scalar) = match (a_scalar, b_scalar) {
/// (true, true) | (false, false) => {
/// let len = a.len().min(b.len());
/// let nulls = NullBuffer::union(a.nulls(), b.nulls());
/// let buffer = BooleanBuffer::collect_bool(len, |idx| a.value(idx) == b.value(idx));
/// return BooleanArray::new(buffer, nulls);
/// }
/// (true, false) => (b, (a.null_count() == 0).then(|| a.value(0))),
/// (false, true) => (a, (b.null_count() == 0).then(|| b.value(0))),
/// };
/// match scalar {
/// Some(v) => {
/// let len = array.len();
/// let nulls = array.nulls().cloned();
/// let buffer = BooleanBuffer::collect_bool(len, |idx| array.value(idx) == v);
/// BooleanArray::new(buffer, nulls)
/// }
/// None => BooleanArray::new_null(array.len()),
/// }
/// }
///
/// pub fn eq(l: &dyn Datum, r: &dyn Datum) -> Result<BooleanArray, ArrowError> {
/// let (l_array, l_scalar) = l.get();
/// let (r_array, r_scalar) = r.get();
/// downcast_primitive_array!(
/// (l_array, r_array) => Ok(eq_impl(l_array, l_scalar, r_array, r_scalar)),
/// (a, b) => Err(ArrowError::NotYetImplemented(format!("{a} == {b}"))),
/// )
/// }
///
/// // Comparison of two arrays
/// let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
/// let b = Int32Array::from(vec![1, 2, 4, 7, 3]);
/// let r = eq(&a, &b).unwrap();
/// let values: Vec<_> = r.values().iter().collect();
/// assert_eq!(values, &[true, true, false, false, false]);
///
/// // Comparison of an array and a scalar
/// let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
/// let b = Int32Array::new_scalar(1);
/// let r = eq(&a, &b).unwrap();
/// let values: Vec<_> = r.values().iter().collect();
/// assert_eq!(values, &[true, false, false, false, false]);
pub trait Datum {
/// Returns the value for this [`Datum`] and a boolean indicating if the value is scalar
fn get(&self) -> (&dyn Array, bool);
}
impl<T: Array> Datum for T {
fn get(&self) -> (&dyn Array, bool) {
(self, false)
}
}
impl Datum for dyn Array {
fn get(&self) -> (&dyn Array, bool) {
(self, false)
}
}
impl Datum for &dyn Array {
fn get(&self) -> (&dyn Array, bool) {
(*self, false)
}
}
/// A wrapper around a single value [`Array`] that implements
/// [`Datum`] and indicates [compute] kernels should treat this array
/// as a scalar value (a single value).
///
/// Using a [`Scalar`] is often much more efficient than creating an
/// [`Array`] with the same (repeated) value.
///
/// See [`Datum`] for more information.
///
/// # Example
///
/// ```rust
/// # use arrow_array::{Scalar, Int32Array, ArrayRef};
/// # fn get_array() -> ArrayRef { std::sync::Arc::new(Int32Array::from(vec![42])) }
/// // Create a (typed) scalar for Int32Array for the value 42
/// let scalar = Scalar::new(Int32Array::from(vec![42]));
///
/// // Create a scalar using PrimtiveArray::scalar
/// let scalar = Int32Array::new_scalar(42);
///
/// // create a scalar from an ArrayRef (for dynamic typed Arrays)
/// let array: ArrayRef = get_array();
/// let scalar = Scalar::new(array);
/// ```
///
/// [compute]: https://docs.rs/arrow/latest/arrow/compute/index.html
#[derive(Debug, Copy, Clone)]
pub struct Scalar<T: Array>(T);
impl<T: Array> Scalar<T> {
/// Create a new [`Scalar`] from an [`Array`]
///
/// # Panics
///
/// Panics if `array.len() != 1`
pub fn new(array: T) -> Self {
assert_eq!(array.len(), 1);
Self(array)
}
/// Returns the inner array
#[inline]
pub fn into_inner(self) -> T {
self.0
}
}
impl<T: Array> Datum for Scalar<T> {
fn get(&self) -> (&dyn Array, bool) {
(&self.0, true)
}
}
@@ -0,0 +1,418 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Conversion methods for dates and times.
use crate::ArrowPrimitiveType;
use crate::timezone::Tz;
use arrow_schema::{DataType, TimeUnit};
use chrono::{DateTime, Duration, NaiveDate, NaiveDateTime, NaiveTime, TimeZone, Timelike, Utc};
/// Number of seconds in a day
pub const SECONDS_IN_DAY: i64 = 86_400;
/// Number of milliseconds in a second
pub const MILLISECONDS: i64 = 1_000;
/// Number of microseconds in a second
pub const MICROSECONDS: i64 = 1_000_000;
/// Number of nanoseconds in a second
pub const NANOSECONDS: i64 = 1_000_000_000;
/// Number of milliseconds in a day
pub const MILLISECONDS_IN_DAY: i64 = SECONDS_IN_DAY * MILLISECONDS;
/// Number of microseconds in a day
pub const MICROSECONDS_IN_DAY: i64 = SECONDS_IN_DAY * MICROSECONDS;
/// Number of nanoseconds in a day
pub const NANOSECONDS_IN_DAY: i64 = SECONDS_IN_DAY * NANOSECONDS;
/// Constant from chrono crate
///
/// Number of days between Januari 1, 1970 and December 31, 1 BCE which we define to be day 0.
/// 4 full leap year cycles until December 31, 1600 4 * 146097 = 584388
/// 1 day until January 1, 1601 1
/// 369 years until Januari 1, 1970 369 * 365 = 134685
/// of which floor(369 / 4) are leap years floor(369 / 4) = 92
/// except for 1700, 1800 and 1900 -3 +
/// --------
/// 719163
pub const UNIX_EPOCH_DAY: i64 = 719_163;
/// converts a `i32` representing a `date32` to [`NaiveDateTime`]
#[inline]
pub fn date32_to_datetime(v: i32) -> Option<NaiveDateTime> {
Some(DateTime::from_timestamp(v as i64 * SECONDS_IN_DAY, 0)?.naive_utc())
}
/// converts a `i64` representing a `date64` to [`NaiveDateTime`]
#[inline]
pub fn date64_to_datetime(v: i64) -> Option<NaiveDateTime> {
let (sec, milli_sec) = split_second(v, MILLISECONDS);
let datetime = DateTime::from_timestamp(
// extract seconds from milliseconds
sec,
// discard extracted seconds and convert milliseconds to nanoseconds
milli_sec * MICROSECONDS as u32,
)?;
Some(datetime.naive_utc())
}
/// converts a `i32` representing a `time32(s)` to [`NaiveDateTime`]
#[inline]
pub fn time32s_to_time(v: i32) -> Option<NaiveTime> {
NaiveTime::from_num_seconds_from_midnight_opt(v as u32, 0)
}
/// converts a `i32` representing a `time32(ms)` to [`NaiveDateTime`]
#[inline]
pub fn time32ms_to_time(v: i32) -> Option<NaiveTime> {
let v = v as i64;
NaiveTime::from_num_seconds_from_midnight_opt(
// extract seconds from milliseconds
(v / MILLISECONDS) as u32,
// discard extracted seconds and convert milliseconds to
// nanoseconds
(v % MILLISECONDS * MICROSECONDS) as u32,
)
}
/// converts a `i64` representing a `time64(us)` to [`NaiveDateTime`]
#[inline]
pub fn time64us_to_time(v: i64) -> Option<NaiveTime> {
NaiveTime::from_num_seconds_from_midnight_opt(
// extract seconds from microseconds
(v / MICROSECONDS) as u32,
// discard extracted seconds and convert microseconds to
// nanoseconds
(v % MICROSECONDS * MILLISECONDS) as u32,
)
}
/// converts a `i64` representing a `time64(ns)` to [`NaiveDateTime`]
#[inline]
pub fn time64ns_to_time(v: i64) -> Option<NaiveTime> {
NaiveTime::from_num_seconds_from_midnight_opt(
// extract seconds from nanoseconds
(v / NANOSECONDS) as u32,
// discard extracted seconds
(v % NANOSECONDS) as u32,
)
}
/// converts [`NaiveTime`] to a `i32` representing a `time32(s)`
#[inline]
pub fn time_to_time32s(v: NaiveTime) -> i32 {
v.num_seconds_from_midnight() as i32
}
/// converts [`NaiveTime`] to a `i32` representing a `time32(ms)`
#[inline]
pub fn time_to_time32ms(v: NaiveTime) -> i32 {
(v.num_seconds_from_midnight() as i64 * MILLISECONDS
+ v.nanosecond() as i64 * MILLISECONDS / NANOSECONDS) as i32
}
/// converts [`NaiveTime`] to a `i64` representing a `time64(us)`
#[inline]
pub fn time_to_time64us(v: NaiveTime) -> i64 {
v.num_seconds_from_midnight() as i64 * MICROSECONDS
+ v.nanosecond() as i64 * MICROSECONDS / NANOSECONDS
}
/// converts [`NaiveTime`] to a `i64` representing a `time64(ns)`
#[inline]
pub fn time_to_time64ns(v: NaiveTime) -> i64 {
v.num_seconds_from_midnight() as i64 * NANOSECONDS + v.nanosecond() as i64
}
/// converts a `i64` representing a `timestamp(s)` to [`NaiveDateTime`]
#[inline]
pub fn timestamp_s_to_datetime(v: i64) -> Option<NaiveDateTime> {
Some(DateTime::from_timestamp(v, 0)?.naive_utc())
}
/// Similar to timestamp_s_to_datetime but only compute `date`
#[inline]
pub fn timestamp_s_to_date(secs: i64) -> Option<NaiveDateTime> {
let days = secs.div_euclid(86_400) + UNIX_EPOCH_DAY;
if days < i32::MIN as i64 || days > i32::MAX as i64 {
return None;
}
let date = NaiveDate::from_num_days_from_ce_opt(days as i32)?;
Some(date.and_time(NaiveTime::default()).and_utc().naive_utc())
}
/// Similar to timestamp_s_to_datetime but only compute `time`
#[inline]
pub fn timestamp_s_to_time(secs: i64) -> Option<NaiveDateTime> {
let secs = secs.rem_euclid(86_400);
let time = NaiveTime::from_num_seconds_from_midnight_opt(secs as u32, 0)?;
Some(
DateTime::<Utc>::from_naive_utc_and_offset(
NaiveDateTime::new(NaiveDate::default(), time),
Utc,
)
.naive_utc(),
)
}
/// converts a `i64` representing a `timestamp(ms)` to [`NaiveDateTime`]
#[inline]
pub fn timestamp_ms_to_datetime(v: i64) -> Option<NaiveDateTime> {
let (sec, milli_sec) = split_second(v, MILLISECONDS);
let datetime = DateTime::from_timestamp(
// extract seconds from milliseconds
sec,
// discard extracted seconds and convert milliseconds to nanoseconds
milli_sec * MICROSECONDS as u32,
)?;
Some(datetime.naive_utc())
}
/// converts a `i64` representing a `timestamp(us)` to [`NaiveDateTime`]
#[inline]
pub fn timestamp_us_to_datetime(v: i64) -> Option<NaiveDateTime> {
let (sec, micro_sec) = split_second(v, MICROSECONDS);
let datetime = DateTime::from_timestamp(
// extract seconds from microseconds
sec,
// discard extracted seconds and convert microseconds to nanoseconds
micro_sec * MILLISECONDS as u32,
)?;
Some(datetime.naive_utc())
}
/// converts a `i64` representing a `timestamp(ns)` to [`NaiveDateTime`]
#[inline]
pub fn timestamp_ns_to_datetime(v: i64) -> Option<NaiveDateTime> {
let (sec, nano_sec) = split_second(v, NANOSECONDS);
let datetime = DateTime::from_timestamp(
// extract seconds from nanoseconds
sec, // discard extracted seconds
nano_sec,
)?;
Some(datetime.naive_utc())
}
#[inline]
pub(crate) fn split_second(v: i64, base: i64) -> (i64, u32) {
(v.div_euclid(base), v.rem_euclid(base) as u32)
}
/// converts a `i64` representing a `duration(s)` to [`Duration`]
#[inline]
#[deprecated(since = "55.2.0", note = "Use `try_duration_s_to_duration` instead")]
pub fn duration_s_to_duration(v: i64) -> Duration {
Duration::try_seconds(v).unwrap()
}
/// converts a `i64` representing a `duration(s)` to [`Option<Duration>`]
#[inline]
pub fn try_duration_s_to_duration(v: i64) -> Option<Duration> {
Duration::try_seconds(v)
}
/// converts a `i64` representing a `duration(ms)` to [`Duration`]
#[inline]
#[deprecated(since = "55.2.0", note = "Use `try_duration_ms_to_duration` instead")]
pub fn duration_ms_to_duration(v: i64) -> Duration {
Duration::try_seconds(v).unwrap()
}
/// converts a `i64` representing a `duration(ms)` to [`Option<Duration>`]
#[inline]
pub fn try_duration_ms_to_duration(v: i64) -> Option<Duration> {
Duration::try_milliseconds(v)
}
/// converts a `i64` representing a `duration(us)` to [`Duration`]
#[inline]
pub fn duration_us_to_duration(v: i64) -> Duration {
Duration::microseconds(v)
}
/// converts a `i64` representing a `duration(ns)` to [`Duration`]
#[inline]
pub fn duration_ns_to_duration(v: i64) -> Duration {
Duration::nanoseconds(v)
}
/// Converts an [`ArrowPrimitiveType`] to [`NaiveDateTime`]
pub fn as_datetime<T: ArrowPrimitiveType>(v: i64) -> Option<NaiveDateTime> {
match T::DATA_TYPE {
DataType::Date32 => date32_to_datetime(v as i32),
DataType::Date64 => date64_to_datetime(v),
DataType::Time32(_) | DataType::Time64(_) => None,
DataType::Timestamp(unit, _) => match unit {
TimeUnit::Second => timestamp_s_to_datetime(v),
TimeUnit::Millisecond => timestamp_ms_to_datetime(v),
TimeUnit::Microsecond => timestamp_us_to_datetime(v),
TimeUnit::Nanosecond => timestamp_ns_to_datetime(v),
},
// interval is not yet fully documented [ARROW-3097]
DataType::Interval(_) => None,
_ => None,
}
}
/// Converts an [`ArrowPrimitiveType`] to [`DateTime<Tz>`]
pub fn as_datetime_with_timezone<T: ArrowPrimitiveType>(v: i64, tz: Tz) -> Option<DateTime<Tz>> {
let naive = as_datetime::<T>(v)?;
Some(Utc.from_utc_datetime(&naive).with_timezone(&tz))
}
/// Converts an [`ArrowPrimitiveType`] to [`NaiveDate`]
pub fn as_date<T: ArrowPrimitiveType>(v: i64) -> Option<NaiveDate> {
as_datetime::<T>(v).map(|datetime| datetime.date())
}
/// Converts an [`ArrowPrimitiveType`] to [`NaiveTime`]
pub fn as_time<T: ArrowPrimitiveType>(v: i64) -> Option<NaiveTime> {
match T::DATA_TYPE {
DataType::Time32(unit) => {
// safe to immediately cast to u32 as `self.value(i)` is positive i32
let v = v as u32;
match unit {
TimeUnit::Second => time32s_to_time(v as i32),
TimeUnit::Millisecond => time32ms_to_time(v as i32),
_ => None,
}
}
DataType::Time64(unit) => match unit {
TimeUnit::Microsecond => time64us_to_time(v),
TimeUnit::Nanosecond => time64ns_to_time(v),
_ => None,
},
DataType::Timestamp(_, _) => as_datetime::<T>(v).map(|datetime| datetime.time()),
DataType::Date32 | DataType::Date64 => NaiveTime::from_hms_opt(0, 0, 0),
DataType::Interval(_) => None,
_ => None,
}
}
/// Converts an [`ArrowPrimitiveType`] to [`Duration`]
pub fn as_duration<T: ArrowPrimitiveType>(v: i64) -> Option<Duration> {
match T::DATA_TYPE {
DataType::Duration(unit) => match unit {
TimeUnit::Second => try_duration_s_to_duration(v),
TimeUnit::Millisecond => try_duration_ms_to_duration(v),
TimeUnit::Microsecond => Some(duration_us_to_duration(v)),
TimeUnit::Nanosecond => Some(duration_ns_to_duration(v)),
},
_ => None,
}
}
#[cfg(test)]
mod tests {
use crate::temporal_conversions::{
NANOSECONDS, date64_to_datetime, split_second, timestamp_ms_to_datetime,
timestamp_ns_to_datetime, timestamp_s_to_date, timestamp_s_to_datetime,
timestamp_s_to_time, timestamp_us_to_datetime,
};
use chrono::DateTime;
#[test]
fn test_timestamp_func() {
let timestamp = 1234;
let datetime = timestamp_s_to_datetime(timestamp).unwrap();
let expected_date = datetime.date();
let expected_time = datetime.time();
assert_eq!(
timestamp_s_to_date(timestamp).unwrap().date(),
expected_date
);
assert_eq!(
timestamp_s_to_time(timestamp).unwrap().time(),
expected_time
);
}
#[test]
fn negative_input_timestamp_ns_to_datetime() {
assert_eq!(
timestamp_ns_to_datetime(-1),
DateTime::from_timestamp(-1, 999_999_999).map(|x| x.naive_utc())
);
assert_eq!(
timestamp_ns_to_datetime(-1_000_000_001),
DateTime::from_timestamp(-2, 999_999_999).map(|x| x.naive_utc())
);
}
#[test]
fn negative_input_timestamp_us_to_datetime() {
assert_eq!(
timestamp_us_to_datetime(-1),
DateTime::from_timestamp(-1, 999_999_000).map(|x| x.naive_utc())
);
assert_eq!(
timestamp_us_to_datetime(-1_000_001),
DateTime::from_timestamp(-2, 999_999_000).map(|x| x.naive_utc())
);
}
#[test]
fn negative_input_timestamp_ms_to_datetime() {
assert_eq!(
timestamp_ms_to_datetime(-1),
DateTime::from_timestamp(-1, 999_000_000).map(|x| x.naive_utc())
);
assert_eq!(
timestamp_ms_to_datetime(-1_001),
DateTime::from_timestamp(-2, 999_000_000).map(|x| x.naive_utc())
);
}
#[test]
fn negative_input_date64_to_datetime() {
assert_eq!(
date64_to_datetime(-1),
DateTime::from_timestamp(-1, 999_000_000).map(|x| x.naive_utc())
);
assert_eq!(
date64_to_datetime(-1_001),
DateTime::from_timestamp(-2, 999_000_000).map(|x| x.naive_utc())
);
}
#[test]
fn test_split_seconds() {
let (sec, nano_sec) = split_second(100, NANOSECONDS);
assert_eq!(sec, 0);
assert_eq!(nano_sec, 100);
let (sec, nano_sec) = split_second(123_000_000_456, NANOSECONDS);
assert_eq!(sec, 123);
assert_eq!(nano_sec, 456);
let (sec, nano_sec) = split_second(-1, NANOSECONDS);
assert_eq!(sec, -1);
assert_eq!(nano_sec, 999_999_999);
let (sec, nano_sec) = split_second(-123_000_000_001, NANOSECONDS);
assert_eq!(sec, -124);
assert_eq!(nano_sec, 999_999_999);
}
}
+358
View File
@@ -0,0 +1,358 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Timezone for timestamp arrays
use arrow_schema::ArrowError;
use chrono::FixedOffset;
pub use private::{Tz, TzOffset};
/// Parses a fixed offset of the form "+09:00", "-09" or "+0930"
fn parse_fixed_offset(tz: &str) -> Option<FixedOffset> {
let bytes = tz.as_bytes();
let mut values = match bytes.len() {
// [+-]XX:XX
6 if bytes[3] == b':' => [bytes[1], bytes[2], bytes[4], bytes[5]],
// [+-]XXXX
5 => [bytes[1], bytes[2], bytes[3], bytes[4]],
// [+-]XX
3 => [bytes[1], bytes[2], b'0', b'0'],
_ => return None,
};
values.iter_mut().for_each(|x| *x = x.wrapping_sub(b'0'));
if values.iter().any(|x| *x > 9) {
return None;
}
let secs =
(values[0] * 10 + values[1]) as i32 * 60 * 60 + (values[2] * 10 + values[3]) as i32 * 60;
match bytes[0] {
b'+' => FixedOffset::east_opt(secs),
b'-' => FixedOffset::west_opt(secs),
_ => None,
}
}
#[cfg(feature = "chrono-tz")]
mod private {
use super::*;
use chrono::offset::TimeZone;
use chrono::{LocalResult, NaiveDate, NaiveDateTime, Offset};
use std::fmt::Display;
use std::str::FromStr;
/// An [`Offset`] for [`Tz`]
#[derive(Debug, Copy, Clone)]
pub struct TzOffset {
tz: Tz,
offset: FixedOffset,
}
impl std::fmt::Display for TzOffset {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
self.offset.fmt(f)
}
}
impl Offset for TzOffset {
fn fix(&self) -> FixedOffset {
self.offset
}
}
/// An Arrow [`TimeZone`]
#[derive(Debug, Copy, Clone)]
pub struct Tz(TzInner);
#[derive(Debug, Copy, Clone)]
enum TzInner {
Timezone(chrono_tz::Tz),
Offset(FixedOffset),
}
impl FromStr for Tz {
type Err = ArrowError;
fn from_str(tz: &str) -> Result<Self, Self::Err> {
match parse_fixed_offset(tz) {
Some(offset) => Ok(Self(TzInner::Offset(offset))),
None => Ok(Self(TzInner::Timezone(tz.parse().map_err(|e| {
ArrowError::ParseError(format!("Invalid timezone \"{tz}\": {e}"))
})?))),
}
}
}
impl Display for Tz {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self.0 {
TzInner::Timezone(tz) => tz.fmt(f),
TzInner::Offset(offset) => offset.fmt(f),
}
}
}
macro_rules! tz {
($s:ident, $tz:ident, $b:block) => {
match $s.0 {
TzInner::Timezone($tz) => $b,
TzInner::Offset($tz) => $b,
}
};
}
impl TimeZone for Tz {
type Offset = TzOffset;
fn from_offset(offset: &Self::Offset) -> Self {
offset.tz
}
fn offset_from_local_date(&self, local: &NaiveDate) -> LocalResult<Self::Offset> {
tz!(self, tz, {
tz.offset_from_local_date(local).map(|x| TzOffset {
tz: *self,
offset: x.fix(),
})
})
}
fn offset_from_local_datetime(&self, local: &NaiveDateTime) -> LocalResult<Self::Offset> {
tz!(self, tz, {
tz.offset_from_local_datetime(local).map(|x| TzOffset {
tz: *self,
offset: x.fix(),
})
})
}
fn offset_from_utc_date(&self, utc: &NaiveDate) -> Self::Offset {
tz!(self, tz, {
TzOffset {
tz: *self,
offset: tz.offset_from_utc_date(utc).fix(),
}
})
}
fn offset_from_utc_datetime(&self, utc: &NaiveDateTime) -> Self::Offset {
tz!(self, tz, {
TzOffset {
tz: *self,
offset: tz.offset_from_utc_datetime(utc).fix(),
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::{Timelike, Utc};
#[test]
fn test_with_timezone() {
let vals = [
Utc.timestamp_millis_opt(37800000).unwrap(),
Utc.timestamp_millis_opt(86339000).unwrap(),
];
assert_eq!(10, vals[0].hour());
assert_eq!(23, vals[1].hour());
let tz: Tz = "America/Los_Angeles".parse().unwrap();
assert_eq!(2, vals[0].with_timezone(&tz).hour());
assert_eq!(15, vals[1].with_timezone(&tz).hour());
}
#[test]
fn test_using_chrono_tz_and_utc_naive_date_time() {
let sydney_tz = "Australia/Sydney".to_string();
let tz: Tz = sydney_tz.parse().unwrap();
let sydney_offset_without_dst = FixedOffset::east_opt(10 * 60 * 60).unwrap();
let sydney_offset_with_dst = FixedOffset::east_opt(11 * 60 * 60).unwrap();
// Daylight savings ends
// When local daylight time was about to reach
// Sunday, 4 April 2021, 3:00:00 am clocks were turned backward 1 hour to
// Sunday, 4 April 2021, 2:00:00 am local standard time instead.
// Daylight savings starts
// When local standard time was about to reach
// Sunday, 3 October 2021, 2:00:00 am clocks were turned forward 1 hour to
// Sunday, 3 October 2021, 3:00:00 am local daylight time instead.
// Sydney 2021-04-04T02:30:00+11:00 is 2021-04-03T15:30:00Z
let utc_just_before_sydney_dst_ends = NaiveDate::from_ymd_opt(2021, 4, 3)
.unwrap()
.and_hms_nano_opt(15, 30, 0, 0)
.unwrap();
assert_eq!(
tz.offset_from_utc_datetime(&utc_just_before_sydney_dst_ends)
.fix(),
sydney_offset_with_dst
);
// Sydney 2021-04-04T02:30:00+10:00 is 2021-04-03T16:30:00Z
let utc_just_after_sydney_dst_ends = NaiveDate::from_ymd_opt(2021, 4, 3)
.unwrap()
.and_hms_nano_opt(16, 30, 0, 0)
.unwrap();
assert_eq!(
tz.offset_from_utc_datetime(&utc_just_after_sydney_dst_ends)
.fix(),
sydney_offset_without_dst
);
// Sydney 2021-10-03T01:30:00+10:00 is 2021-10-02T15:30:00Z
let utc_just_before_sydney_dst_starts = NaiveDate::from_ymd_opt(2021, 10, 2)
.unwrap()
.and_hms_nano_opt(15, 30, 0, 0)
.unwrap();
assert_eq!(
tz.offset_from_utc_datetime(&utc_just_before_sydney_dst_starts)
.fix(),
sydney_offset_without_dst
);
// Sydney 2021-04-04T03:30:00+11:00 is 2021-10-02T16:30:00Z
let utc_just_after_sydney_dst_starts = NaiveDate::from_ymd_opt(2022, 10, 2)
.unwrap()
.and_hms_nano_opt(16, 30, 0, 0)
.unwrap();
assert_eq!(
tz.offset_from_utc_datetime(&utc_just_after_sydney_dst_starts)
.fix(),
sydney_offset_with_dst
);
}
#[test]
fn test_timezone_display() {
let test_cases = ["UTC", "America/Los_Angeles", "-08:00", "+05:30"];
for &case in &test_cases {
let tz: Tz = case.parse().unwrap();
assert_eq!(tz.to_string(), case);
}
}
}
}
#[cfg(not(feature = "chrono-tz"))]
mod private {
use super::*;
use chrono::offset::TimeZone;
use chrono::{LocalResult, NaiveDate, NaiveDateTime, Offset};
use std::str::FromStr;
/// An [`Offset`] for [`Tz`]
#[derive(Debug, Copy, Clone)]
pub struct TzOffset(FixedOffset);
impl std::fmt::Display for TzOffset {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
self.0.fmt(f)
}
}
impl Offset for TzOffset {
fn fix(&self) -> FixedOffset {
self.0
}
}
/// An Arrow [`TimeZone`]
#[derive(Debug, Copy, Clone)]
pub struct Tz(FixedOffset);
impl FromStr for Tz {
type Err = ArrowError;
fn from_str(tz: &str) -> Result<Self, Self::Err> {
let offset = parse_fixed_offset(tz).ok_or_else(|| {
ArrowError::ParseError(format!(
"Invalid timezone \"{tz}\": only offset based timezones supported without chrono-tz feature"
))
})?;
Ok(Self(offset))
}
}
impl TimeZone for Tz {
type Offset = TzOffset;
fn from_offset(offset: &Self::Offset) -> Self {
Self(offset.0)
}
fn offset_from_local_date(&self, local: &NaiveDate) -> LocalResult<Self::Offset> {
self.0.offset_from_local_date(local).map(TzOffset)
}
fn offset_from_local_datetime(&self, local: &NaiveDateTime) -> LocalResult<Self::Offset> {
self.0.offset_from_local_datetime(local).map(TzOffset)
}
fn offset_from_utc_date(&self, utc: &NaiveDate) -> Self::Offset {
TzOffset(self.0.offset_from_utc_date(utc).fix())
}
fn offset_from_utc_datetime(&self, utc: &NaiveDateTime) -> Self::Offset {
TzOffset(self.0.offset_from_utc_datetime(utc).fix())
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::{NaiveDate, Offset, TimeZone};
#[test]
fn test_with_offset() {
let t = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap();
let tz: Tz = "-00:00".parse().unwrap();
assert_eq!(tz.offset_from_utc_date(&t).fix().local_minus_utc(), 0);
let tz: Tz = "+00:00".parse().unwrap();
assert_eq!(tz.offset_from_utc_date(&t).fix().local_minus_utc(), 0);
let tz: Tz = "-10:00".parse().unwrap();
assert_eq!(
tz.offset_from_utc_date(&t).fix().local_minus_utc(),
-10 * 60 * 60
);
let tz: Tz = "+09:00".parse().unwrap();
assert_eq!(
tz.offset_from_utc_date(&t).fix().local_minus_utc(),
9 * 60 * 60
);
let tz = "+09".parse::<Tz>().unwrap();
assert_eq!(
tz.offset_from_utc_date(&t).fix().local_minus_utc(),
9 * 60 * 60
);
let tz = "+0900".parse::<Tz>().unwrap();
assert_eq!(
tz.offset_from_utc_date(&t).fix().local_minus_utc(),
9 * 60 * 60
);
let err = "+9:00".parse::<Tz>().unwrap_err().to_string();
assert!(err.contains("Invalid timezone"), "{}", err);
}
}
+78
View File
@@ -0,0 +1,78 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer, bit_util};
/// Creates two [`Buffer`]s from an iterator of `Option`.
/// The first buffer corresponds to a bitmap buffer, the second one
/// corresponds to a values buffer.
/// # Safety
/// The caller must ensure that `iterator` is `TrustedLen`.
#[inline]
pub(crate) unsafe fn trusted_len_unzip<I, P, T>(iterator: I) -> (Buffer, Buffer)
where
T: ArrowNativeType,
P: std::borrow::Borrow<Option<T>>,
I: Iterator<Item = P>,
{
let (_, upper) = iterator.size_hint();
let upper = upper.expect("trusted_len_unzip requires an upper limit");
let len = upper * std::mem::size_of::<T>();
let mut null = MutableBuffer::from_len_zeroed(upper.saturating_add(7) / 8);
let mut buffer = MutableBuffer::new(len);
let dst_null = null.as_mut_ptr();
let mut dst = buffer.as_mut_ptr() as *mut T;
for (i, item) in iterator.enumerate() {
let item = item.borrow();
if let Some(item) = item {
unsafe { std::ptr::write(dst, *item) };
unsafe { bit_util::set_bit_raw(dst_null, i) };
} else {
unsafe { std::ptr::write(dst, T::default()) };
}
dst = unsafe { dst.add(1) };
}
assert_eq!(
unsafe { dst.offset_from(buffer.as_ptr() as *mut T) as usize },
upper,
"Trusted iterator length was not accurately reported"
);
unsafe { buffer.set_len(len) };
(null.into(), buffer.into())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn trusted_len_unzip_good() {
let vec = [Some(1u32), None];
let (null, buffer) = unsafe { trusted_len_unzip(vec.iter()) };
assert_eq!(null.as_slice(), &[0b00000001]);
assert_eq!(buffer.as_slice(), &[1u8, 0, 0, 0, 0, 0, 0, 0]);
}
#[test]
#[should_panic(expected = "trusted_len_unzip requires an upper limit")]
fn trusted_len_unzip_panic() {
let iter = std::iter::repeat(Some(4i32));
unsafe { trusted_len_unzip(iter) };
}
}
File diff suppressed because it is too large Load Diff