![\"A](\"media/foursquare-os-places-density-2024.png\")
In the weeks since my previous post on Working with Arrow and DuckDB in Rust,\nI've found a few gripes that I'd like to address.
\nquery_arrow and stream_arrowIn the previous post, I used the query_arrow API.\nIt's pretty straightforward and gives you iterator-compatible access to the query results.\nHowever, there's one small problem: its memory consumption scales roughly linearly with your result set.
This isn't a problem for many uses of DuckDB, but if your datasets are in the tens or hundreds of gigabytes\nand you're wanting to process a large number of rows, the RAM requirements can be excessive.\nThe memory profile of query_arrow seems to be "create all of the RecordBatches upfront\nand keep them around for as long as you hold the Arrow handle.
Disclaimer
\nI have not done extensive allocation-level memory profiling as of this writing.\nIt's quite possible that I've missed something, but this seems to be what's happening\nfrom watching Activity Monitor.\nPlease let me know if I've misrepresented anything!
\nFortunately, DuckDB also has another API: stream_arrow.\nThis appears to allocate RecordBatches on demand rather than all at once.\nThere is also some overhead, which I'll revisit later that varies with result size.\nBut overall, profiling indicates that stream_arrow requires significantly less RAM over the life of a large Arrow iterator.
Unfortunately, none of the above information about memory consumption appears to be documented,\nand there are no (serious) code samples demonstrating the use of stream_arrow!
\n\n[!question] Down the rabbit hole...\nDigging into the code in duckdb-rs raises even more questions,\nsince several underlying C functions, like
\nduckdb_execute_prepared_streaming\nare marked as deprecated.\nPresumably, alternatives are being developed or the methods are just not stable yet.
SchemaRefThe signature of stream_arrow is a bit different from that of query_arrow.\nHere's what it looks like as of crate version 1.1.1:
pub fn stream_arrow<P: Params>(\n &mut self,\n params: P,\n schema: SchemaRef,\n) -> Result<ArrowStream<'_>>\nThis looks pretty familiar at first if you've used query_arrow,\nbut there's a new third parameter: schema.\nSchemaRef is just a type alias for Arc<Schema>.\nArrow objects have a schema associated with them,\nso this is a reasonable detail for a low-level API.\nBut DuckDB is fine at inferring this when needed!\nSurely there is a way of getting it from a query, right?\n(After all, query_arrow has to do something similar, but doesn't burden the caller.)
My first attempt at getting a Schema object was to call the schema() method on Statement.\nThe Statement type in duckdb-rs is actually a high-level wrapper around RawStatement,\nand at the time of this writing, the schema getter hides an unwrap.\nThe docs do tell you this (using a somewhat nonstandard heading?),\nbut basically you can't get a schema without executing a query.\nI wish they used the Typestate pattern\nor at least made the result an Option, but alas...
This leaves developers with three options.
\nStatement that is the same SQL, but with a LIMIT 0 clause at the end.Manually constructing the schema is a non-starter for me.\nA program which has a hand-written code dependency on a SQL string is a terrible idea\non several levels.\nBesides, DuckDB clearly can infer the schema in query_arrow, so why not here?
The second idea is, amusingly, what ChatGPT o1 suggested (after half a dozen prompts;\nit seems like it will just confidently refuse to fetch documentation now,\nand hallucinates new APIs based off its outdated training data).\nThe basic idea is to add LIMIT 0 to the end of the original query\nso it's able to get the schema, but doesn't actually return any results.
fn fetch_schema_for_query(db: &Connection, sql: &str) -> duckdb::Result<SchemaRef> {\n // Append "LIMIT 0" to the original query, so we don't actually fetch anything\n // NB: This does NOT handle cases such as the original query ending in a semicolon!\n let schema_sql = format!("{} LIMIT 0", sql);\n\n let mut statement = db.prepare(&schema_sql)?;\n let arrow_result = statement.query_arrow([])?;\n\n Ok(arrow_result.get_schema())\n}\nThere is nothing fundamentally unsound about this approach.\nBut it requires string manipulation, which is less than ideal.\nThere is also at least one obvious edge case.
\nThe third option is not as straightforward as I expected it to be.\nAt first, I tried the row_count method,\nbut internally this just calls a single FFI function.\nThis doesn't actually update the internal schema field.\nYou really do need to run through a more "normal" execution path.
A solution that seems reasonably clean is to do what the docs say and call stmt.execute().\nIt's a bit strange to do this on a SELECT query to be honest,\nbut the API does indeed internally mutate the Schema property,\nand returns a row count.\nSo it seems semantically equivalent to a SELECT COUNT(*) FROM (...)\n(and in my case, getting the row count was helpful too).
In my testing, it appears that this may actually allocate a non-trivial amount of memory,\nwhich may be mildly surprising.\nHowever, the max amount of memory we require during execution is definitely less overall.\nAny ideas why this is?
\nstream_arrowLet's bring what we've learned into a "real" example.
\n// let sql = "SELECT * FROM table;";\nlet mut stmt = conn.prepare(sql)?;\n// Execute the query (so we have a usable schema)\nlet size = stmt.execute([])?;\n// Now we run the "real" query using `stream_arrow`.\n// This returned in a few hundred milliseconds for my dataset.\nlet mut arrow = stmt.stream_arrow([], stmt.schema())?;\n// Iterate over arrow...\nWhen you structure your code like this rather than using the easier query_arrow,\nyou can significantly reduce your memory footprint for large datasets.\nIn my testing, there was no appreciable impact on performance.
The above leaves me with a few open questions.\nFirst, with my use case (a dataset of around 12GB of Parquet files), execute took several seconds.\nThe "real" stream_arrow query took a few hundred milliseconds.\nWhat's going on here?\nPerhaps it's doing a scan and/or caching some data initially the way to make subsequent queries faster?
Additionally, the memory profile does have a "spike" which makes me wonder what exactly each step loads into RAM,\nand thus, the memory requirements for working with extremely large datasets.\nIn my testing, adding a WHERE clause that significantly reduces the result set\nDOES reduce the memory footprint.\nThat's somewhat worrying to me, since it implies there is still measurable overhead\nproportional to the dataset size.\nWhat practical limits does this impose on dataset size?
Note
\nAn astute reader may be asking whether the memory profile of the LIMIT 0 and execute approaches are equivalent.\nThe answer appears to be yes.
I've opened issue #418\nasking for clarification.\nIf any readers have any insights, post them in the issue thread!
\n", "summary": "", "date_published": "2024-12-31T00:00:00-00:00", "image": "", "authors": [ { "name": "Ian Wagner", "url": "https://fosstodon.org/@ianthetechie", "avatar": "media/avi.jpeg" } ], "tags": [ "rust", "apache arrow", "parquet", "duckdb", "big data", "data engineering" ], "language": "en" }, { "id": "https://ianwwagner.com//how-and-why-to-work-with-arrow-and-duckdb-in-rust.html", "url": "https://ianwwagner.com//how-and-why-to-work-with-arrow-and-duckdb-in-rust.html", "title": "How (and why) to work with Arrow and DuckDB in Rust", "content_html": "My day job involves wrangling a lot of data very fast.\nI've heard a lot of people raving about several technologies like DuckDB,\n(Geo)Parquet, and Apache Arrow recently.\nBut despite being an "early adopter,"\nit took me quite a while to figure out how and why to leverage these practiclaly.
\nLast week, a few things "clicked" for me, so I'd like to share what I learned in case it helps you.
\n(Geo)Parquet is quite possibly the best understood tech in the mix.\nIt is not exactly new.\nParquet has been around for quite a while in the big data ecosystem.\nIf you need a refresher, the Cloud-optimized Geospatial Formats Guide\ngives a great high-level overview.
\nHere are the stand-out features:
\nPractically speaking, parquet lets you can distribute large datasets in one or more files\nwhich will be significantly smaller and faster to query than other familiar formats.
\nThe value proposition is clear for big data processing.\nIf you're trying to get a record of all traffic accidents in California,\nor find the hottest restaurants in Paris based on a multi-terabyte dataset,\nparquet provides clear advantages.\nYou can skip row groups within the parquet file or even whole files\nto narrow your search!\nAnd since datasets can be split across files,\nyou can keep adding to the dataset over time, parallelize queries,\nand other nice things.
\nBut what if you're not doing these high-level analytical things?\nWhy not just use a more straightforward format like CSV\nthat avoids the need to "rotate" back into rows\nfor non-aggregation use cases?\nHere are a few reasons to like Parquet:
\nThe upshot of all this is that it generally gets both easier and faster\nto work with your data...\nprovided that you have the right tools to leverage it.
\nDuckDB describes itself as an in-process, portable, feature-rich, and fast database\nfor analytical workloads.\nDuckDB was that tool that triggered my "lightbulb moment" last week.\nFoursquare, an app which I've used for a decade or more,\nrecently released an open data set,\nwhich was pretty cool!\nIt was also in Parquet format (just like Overture's data sets).
\nYou can't just open up a Parquet file in a text editor or spreadsheet software like you can a CSV.\nMy friend Oliver released a web-based demo\na few weeks ago which lets you inspect the data on a map at the point level.\nBut to do more than spot checking, you'll probably want a database that can work with Parquet.\nAnd that's where DuckDB comes in.
\nI understood the in-process part of DuckDB's value proposition right away.\nIt's similar to SQLite, where you don't have to go through a server\nor over an HTTP connection.\nThis is both simpler to reason about and is usually quite a bit faster\nthan having to call out to a separate service!
\nDuckDB is pretty quick to compile from source.\nYou probably don't need to muck around with this if you're just using the CLI,\nbut I wanted to eventually use it embedded in some Rust code.\nCompiling from source turned out to be the easiest way to get their crate working.\nIt looks for a shared library by default, but I couldn't get this working after a brew install.\nThis was mildly annoying, but on the other hand,\nvendoring the library does make consistent Docker builds easier 🤷🏻♂️
DuckDB includes a mind boggling number of features.\nNot in a confusing way; more in a Python stdlib way where just about everything you'd want is already there.\nYou can query a whole directory (or bucket) of CSV files,\na Postgres database, SQLite, or even an OpenStreetMap PBF file 🤯\nYou can even write a SQL query against a glob expression of Parquet files in S3\nas your "table."\nThat's really cool!\n(If you've been around the space, you may recognize this concept from\nAWS Athena and others.)
\nWriting a query against a local directory of files is actually really fast!\nIt does a bit of munging upfront, and yes,\nit's not quite as fast as if you'd prepped the data into a clean table,\nbut you actually can run quite efficient queries this way locally!
\nWhen running a query against local data,\nDuckDB will make liberal use of your system memory\n(the default is 80% of system RAM)\nand as many CPUs as you can throw at it.\nBut it will reward you with excellent response times,\ncourtesy of the "vectorized" query engine.\nWhat I've heard of the design reminds me of how array-oriented programming languages like APL\n(or less esoteric libraries like numpy) are often implemented.
\nI was able to do some spatial aggregation operations\n(bucketing a filtered list of locations by H3 index)\nin about 10 seconds on a dataset of more than 40 million rows!\n(The full dataset is over 100 million rows, so I also got to see the selective reading in action.)\nThat piqued my interest, to say the least.\n(Here's the result of that query, visualized).
\n
And now for the final buzzword in DuckDB's marketing: analytical.\nDuckDB frequently describes itself as optimized for OLAP (OnLine Analytical Processing) workloads.\nThis is contrasted with OLTP (OnLine Transaction Processing).\nWikipedia will tell you some differences\nin a lot of sweepingly broad terms, like being used for "business reporting" and read operations\nrather than "transactions."
\nWhen reaching for a definition, many sources focus on things like aggregation queries\nas a differentiator.\nThis didn't help, since most of my use cases involve slurping most or all of the data set.\nThe DuckDB marketing and docs didn't help clarify things either.
\nLet me know on Mastodon if you have a better explanatation of what an "analytical" database is 🤣
\nI think a better explanation is probably 1) you do mostly read queries,\nand 2) it can execute highly parallel queries.\nSo far, DuckDB has been excellent for both the "aggregate" and the "iterative" use case.\nI assume it's just not the best choice per se if your workload is a lot of single-record writes?
\nEmbedding DuckDB in a Rust project allowed me to deliver something with a better end-user experience,\nis easier to maintain,\nand saved writing hundreds of lines of code in the process.
\nMost general-purpose languages like Python and Rust\ndon't have primitives for expressing things like joins across datasets.\nDuckDB, like most database systems, does!\nYes, I could write some code using the parquet crate\nthat would filter across a nested directory tree of 5,000 files.\nBut DuckDB does that out of the box!
It feels like this is a "regex moment" for data processing.\nJust like you don't (usually) need to hand-roll string processing,\nthere's now little reason to hand-roll data aggregation.
\nFor the above visualization, I used the Rust DuckDB crate for the data processing,\nconverted the results to JSON,\nand served it up from an Axum web server.\nAll in a single binary!\nThat's lot nicer than a bash script that executes SQL,\ndumps to a file, and then starts up a Python or Node web server!\nAnd breaks when you don't have Python or Node installed,\nyour OS changes its default shell,\nyou forget that some awk flag doesn't work on the GNU version,\nand so on.
\nThe final thing I want to touch on is Apache Arrow.\nThis is yet another incredibly useful technology which I've been following for a while,\nbut never quite figured out how to properly use until last week.
\nArrow is a language-independent memory format\nthat's optimized for efficient analytic operations on modern CPUs and GPUs.\nThe core idea is that, rather than having to convert data from one format to another (this implies copying!),\nArrow defines as shared memory format which many systems understand.\nIn practice, this ends up being a bunch of standards which define common representations for different types,\nand libraries for working with them.\nFor example, the GeoArrow spec\nbuilds on the Arrow ecosystem to enable operations on spatial data in a common memory format.\nPretty cool!
\nIt turns out that copying and format shifting data can really eat into your processing times.\nArrow helps you sidestep that by reducing the amount of both you'll need to do,\nand by working on data in groups.
\nArrow is mostly hidden from view beneath other libraries.\nSo most of the time, especially if you're writing in a very high level language like Python,\nyou won't even see it.
\nBut if you're writing something at a slightly lower level,\nit's something you may have to touch for critical sections.\nThe DuckDB crate\nincludes an Arrow API\nwhich will give you an iterator over RecordBatches.\nThis is pretty convenient, since you can use DuckDB to gather all your data\nand just consume the stream of batches!
So, how do we work with RecordBatches?\nThe Arrow ecosystem, like Parquet, takes a lot of work to understand,\nand using the low-level libraries directly is difficult.\nEven as a seasoned Rustacean, I found the docs rather obtuse.
After some searching, I finally found serde_arrow.\nIt builds on the serde ecosystem with easy-to-use methods that operate on RecordBatches.\nFinally; something I can use!
I was initilaly worried about how performant the shift from columns to rows + any (minimal) serde overhead would be,\nbut this turned out to not be an issue.
Here's how the code looks:
\nserde_arrow::from_record_batch::<Vec<FoursquarePlaceRecord>>(&batch)\nA few combinators later and you've got a proper data pipeline!
\nWhat this ultimately enabled for me was being able to get a lot closer to "scripting"\na pipeline in Rust.\nMost people turn to Python or JavaScript for tasks like this,\nbut Rust has something to add: strong typing and all the related guarantees which can only come with some level of formalism.\nBut that doesn't necessarily have to get in the way of productivity!
\nHopefully this sparks some ideas for making your next data pipeline both fast and correct.
\n", "summary": "", "date_published": "2024-12-08T00:00:00-00:00", "image": "media/foursquare-os-places-density-2024.png", "authors": [ { "name": "Ian Wagner", "url": "https://fosstodon.org/@ianthetechie", "avatar": "media/avi.jpeg" } ], "tags": [ "rust", "apache arrow", "parquet", "duckdb", "big data", "data engineering", "gis" ], "language": "en" } ] }