[
{
"path": ".gitignore",
"content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\npip-wheel-metadata/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n.python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n*.parquet\n\n.serverless/\nnode_modules/\n\n.DS_Store\n.node-version\n.python-version"
},
{
"path": "LICENSE",
"content": "MIT License\n\nCopyright (c) 2023 Bauplan Labs\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
},
{
"path": "README.md",
"content": "# Quack-reduce\nA playground for running duckdb as a stateless query engine over a data lake. \nThe idea is to have a zero-maintenance, [very fast](https://www.loom.com/share/96f1fd938c814d0a825facb215546f03) and almost free data engine for small analytics apps. \nThis repo is the companion code for this [blog post](https://towardsdatascience.com/a-serverless-query-engine-from-spare-parts-bd6320f10353).\n\nPlease refer to the blog post for more background information and details on the use case.\n\n## Quick Start ..ε=(｡ﾉ･ω･)ﾉ \n\nIf you read the [blog post](https://towardsdatascience.com/a-serverless-query-engine-from-spare-parts-bd6320f10353) and know already what we are up to, follow the quick setup steps below to run everything in no time.\n\n### Setup your account\n\nMake sure you have:\n\n- A working AWS account and an access key with [sufficient priviledges to deploy a lambda instance](https://www.serverless.com/framework/docs/providers/aws/guide/credentials) -- this could be the `AdministratorAccess` policy in AWS IAM, or something more fine grained;\n- [Docker](https://docs.docker.com/get-docker/) installed and running on your machine;\n- Python 3.9+ and Node.js properly installed on your machine;\n- A `profiles.yaml` file on your local machine to run the dbt project.\n\nIn the `src` folder, you should copy `local.env` to `.env` (do *not* commit it) and fill it with proper values:\n\n| value | type | description | example |\n|-----------------------|------|------------------------------------------------------|--------------------------:|\n| AWS_ACCESS_KEY_ID | str | User key for AWS access | AKIAIO... |\n| AWS_SECRET_ACCESS_KEY | str | Secret key for AWS access | wJalr/... |\n| AWS_DEFAULT_REGION | str | AWS region for the deployment | us-east-1 |\n| S3_BUCKET_NAME | str | Bucket to host the data (must be unique) | my-duck-bucket-130dcqda0u |\n\nThese variables will be used by the setup script and the runner to communicate with AWS services. Make sure the user has the permissions to:\n\n- create a bucket and upload files to it;\n- invoke the lambda below.\n\n### Run the project\nFrom the `src` folder:\n\n>**1. Create the DuckDB Lambda:** run `make nodejs-init` and then `make serverless-deploy`. Note that `src/serverless.yml` is configured to use `arm64`. This does a local Docker build, so if you're on an `x86_64` machine, it will fail. Replace `arm64` with `x86_64` as needed. After deployment, you can test the lambda is working from the [console](https://www.loom.com/share/97785a387af84924b830b9e0f35d8a1e).\n\n>**2. Build the Python env:** run `make python-init`.\n\n>**3. Download the data and upload it to S3:** run `make run_me_first` (check your S3 bucket and make sure you find a `partitioned` folder with [this structure](images/s3.png)).\n\n>**4. Test the serverless query engine:** run `make test`.\n\n>**5. Set up your dbt profile:** to run dbt locally, set up a dbt [profile](https://docs.getdbt.com/docs/core/connection-profiles) named `duckdb-taxi` (see [here](https://github.com/jwills/dbt-duckdb) for examples):\n```yaml\n# ~/.dbt/profiles.yml\nduckdb-taxi:\n outputs:\n dev:\n type: duckdb\n path: ':memory:'\n extensions:\n - httpfs\n - parquet\n target: dev\n```\n>**6. Run the dbt project:** run `make dbt-run`.\n\n>**7. Run the Analytics app:** run `make dashboard`.\n\n![](\"images/demo.gif\")
\n\nNote that every time the input field in the dashboard changes, we run a full round-trip on our engine in the back: it can be *this* fast!\n\n***\n\n## Project Overview\n\nIf you want to give it a try, follow the instructions in the `Quick Start` section above to get the system up and running. The rest of the README explores in more details the various components:\n\n* the lambda function running the queries;\n* interactions from a local script;\n* a serverless BI application;\n* running massive parallel workloads through the engine.\n\n![](\"images/flow.png\")
\n\n> NOTE: this project (including this README!) is written for pedagogical purposes and it is not production-ready (or even well tested!): our main goal is to provide a reference implementation of few key concepts as a starting point for future projects - so, sorry for being a bit verbose and perhaps pedantic at times.\n\n### Duckdb lambda\n\nThe `src/serverless` folder is a standard [serverless](https://www.serverless.com/framework/) project, to build an AWS lambda function with Duckdb on it. It has three main components:\n\n- a Dockerfile, which starts from the public AWS lambda image for Python (`public.ecr.aws/lambda/python:3.9`) and add the few dependencies we need;\n- an `app.py` file, containing the actual code our lambda will execute;\n- a `../serverless.yml` file, which ties all these things together in the infra-as-code fashion, and allows us to deploy and manage the function from the CLI.\n\nThe cloud setup is done for you when you run `make nodejs-init` and `make serverless-deploy` (Step 1 in the setup list above). The first time, deployment will take a while as it needs to create the image, ship it to AWS and [create the stack](images/serverless.png) - note that this is _a \"one-off\" thing_.\n\n> NOTE: you may get a `403 Forbidden` error when building the image: in our experience, this usually goes away with `aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws`.\n\n### Interacting with the engine\n\nWe can use a simple Python script to interact with our engine. First, we can test the system with a hard-coded query. Make sure you run Step 2 and Step 3 in the quick start list (to setup Python and the dataset): now, we can test everything is working with `make test`.\n\nIf all looks good, you can now run arbitrary queries (replacing `MY_BUCKET_NAME` with your value) just by using the provided `python quack.py` script; make sure to manually activate your venv with `source ./.venv/bin/activate`, e.g. you can run\n\n`python quack.py -q \"SELECT pickup_location_id AS location_id, COUNT(*) AS counts FROM read_parquet(['s3://MY_BUCKET_NAME/dataset/taxi_2019_04.parquet']) WHERE pickup_at >= '2019-04-01' AND pickup_at < '2019-04-03' GROUP BY 1 ORDER BY 2 DESC\"`\n\nto get the most popular pickup location (IDs) for the first few days of April. \n\nSince the amount of data that can be returned by a lambda is limited, the lambda will automatically limit your rows if you don't specific a limit in the script. You can get more data back with:\n\n`python quack.py -q ... -limit 100`\n\nbut be mindful of the infrastructure constraints!\n\n### Serverless BI architecture (Optional)\n\nIf you want to see how this architecture can bridge the gap between offline pipelines preparing artifacts, and real-time querying for BI (or other use cases), you can simulate how a dbt project may prepare a view that is querable in a dashboard, through our engine (check our blog post for some more context on this use case). \n\nThe quickest setup is running dbt locally, so you need to set up a dbt [profile](https://docs.getdbt.com/docs/core/connection-profiles) named `duckdb-taxi` (see [here](https://github.com/jwills/dbt-duckdb) for examples):\n\n```yaml\n# ~/.dbt/profiles.yml\nduckdb-taxi:\n outputs:\n dev:\n type: duckdb\n path: ':memory:'\n extensions:\n - httpfs\n - parquet\n target: dev\n```\n\n> NOTE: since we run dbt through `make`, there is no need to add credentials to the `extensions`. If you prefer to run it manually, your dbt profile should look more like this:\n\n```yaml\n# ~/.dbt/profiles.yml\nduckdb-taxi:\n outputs:\n dev:\n type: duckdb\n path: ':memory:'\n extensions:\n - httpfs\n - parquet\n settings:\n s3_region: us-east-1\n s3_access_key_id: YOUR_S3_USER\n s3_secret_access_key: YOUR_S3_KEY\n target: dev\n```\n\nAfter the dbt setup is completed, you can use again the `make` file to run a \"batch pipeline\" that produces an artifact in S3 from raw data: just type `make dbt-run` to materialize our view as a parquet file:\n\n![](\"images/dashboard.png\")
\n\n> NOTE: different warehouses would need different configurations to export the node to the same location, e.g. [Snowflake](https://docs.snowflake.com/en/user-guide/script-data-load-transform-parquet). \n\nTo run the front-end (a dashboard built with streamlit querying the view we materialized) run `make dashboard`. A page should open in the browser, displaying a chart:\n\n![](\"images/streamlit.png\")
\n\nYou can use the form to interact in real time with the dataset (video [here](https://www.loom.com/share/9d5de3ba822a445d9d117225c1b0307f)), through the serverless infrastructure we built.\n\n### From quack to quack-reduce (Optional)\n\nThe staless execution of SQL over an object storage (and therefore, using duckdb not really as a db, but basically as \"just\" a query engine) coupled with the parallel nature of AWS lambdas opens up interesting optimization possibilities.\n\nIn particular, we could rephrase (some) SQL queries through a map-reduce programming pattern *with other SQL queries*, and execute them all at the same time. To consider a trivial example, a query such as:\n\n`SELECT COUNT(*) FROM myTable WHERE DATE BETWEEN 04/01/2022 AND 04/05/2022`\n\ncan be rewritten as the SUM of the results of these smaller queries:\n\n`SELECT COUNT(*) FROM myTable WHERE DATE BETWEEN 04/01/2022 AND 04/02/2022` +\n`SELECT COUNT(*) FROM myTable WHERE DATE BETWEEN 04/02/2022 AND 04/03/2022` +\n...\n\nAs the number of files increases (as in a typical hive-partitioned data lake), scanning the object storage (in duckdb syntax `parquet_scan('folder/', HIVE_PARTITIONING=1)`) may take much longer than reading single _k_ files directly through ideally _k_ parallel functions, drastically improving query performances.\n\nTo test out this hypothesis, we built a script that compares the same engine across different deployment patterns - local, remote etc. You can run the bechmarks with default values with `make benchmark`. The script is minimal, but should be enough to give you a feeling of how the different setups perform compared to each other, and the trade-offs involved (check the code for how it's built, but don't expect much!).\n\n[A typical run](https://www.loom.com/share/18a060b89a6a4f6d814e06ffa2674b13) will result in something like this [table](images/benchmarks.png) (numbers will vary).\n\nPlease refer to the blogpost for more musings on this opportunity (and the non-trivial associated challenges).\n\n> NOTE: if you have never raised your concurrency limits on AWS lambda, you may need to request through the console for an increase in parallel execution, otherwise AWS will not allowed the scaling out of the function.\n\n## What's next?\n\nIf you like what you've seen so far, you may wonder what you could do next! There's a million ways to improve this design, some of which more obvious than others - as a non-exhaustive list (\"left as an excercise to the reader\"), this is where we would start:\n\n* if you always query the same table (say, a view for your dashboard), you may want to leverage the `cold` / `warm` pattern in the lambda code to store the table in memory when cold, and read from there (instead of parquet) when warm;\n\n* when you move from one file to multiple files, scanning parquet folders is a huge overhead: wouldn't it be nice to know where to look? While HIVE partitioning is great, modern table formats (e.g. Iceberg) are even better, so you could think of combine their table scan properties with our serverless engine. Performance aside, if you have queried `quack.py`, you know how tedious it is to fully remember the full file name every time: leveraging catalogs like Iceberg, Glue, Nessie etc. would make the experience more \"database-like\";\n\n* try out other use cases! For example, consider this recent [event collection](https://github.com/fal-ai/fal-events) platform. If you modify it to a dump-to-s3-then-query pattern (leveraging the engine we built with this repo), you end up with a lambda-only version of the [Snowflake architecture](https://github.com/jacopotagliabue/paas-data-ingestion) we open sourced some time ago - an end-to-end analytics platform running without servers;\n\n* while we now run the query in memory and return a subset of row from the lambda, this pattern is certainly not perfect: on the one hand, sometime we may wish to write back the result of a query (dbt-style, so to speak); on the other, even if analytics queries are often aggregates, result tables may still grow big (row-wise): writing them to s3 and have the client stream back rows from there may be a nice feature to add!\n\n## License\n\nAll the code is released without warranty, \"as is\" under a MIT License.\n\nThis started as a fun week-end project and should be treated with the appropriate sense of humour.\n"
},
{
"path": "src/Makefile",
"content": "include .env\n\n# need bash because we use the \"source\" command (otherwise it fails when make defaults to /bin/sh)\nSHELL=bash\n\nnodejs-init:\n\tnpm install\n.PHONY: nodejs-init\n\nserverless-deploy:\n\tnpx serverless deploy\n.PHONY: serverless-deploy\n\npython-init:\n\tpython3 -m venv ./.venv && source ./.venv/bin/activate && pip install -r requirements.txt\n.PHONY: python-init\n\nrun_me_first:\n\tsource ./.venv/bin/activate && python3 run_me_first.py\n.PHONY: run_me_first\n\ntest:\n\tsource ./.venv/bin/activate && python3 quack.py\n.PHONY: test\n\ntest-distinct:\n\tsource ./.venv/bin/activate && python3 quack.py -q \"SELECT pickup_location_id AS location_id, COUNT(*) AS counts FROM read_parquet(['s3://${S3_BUCKET_NAME}/dataset/taxi_2019_04.parquet']) WHERE pickup_at >= '2019-04-01' AND pickup_at < '2019-04-03' GROUP BY 1 ORDER BY 2 DESC\"\n.PHONY: test-distinct\n\nbenchmark:\n\tsource ./.venv/bin/activate && python3 benchmark.py\n.PHONY: benchmark\n\ndbt-run:\n\tsource ./.venv/bin/activate && cd dashboard/dbt && S3_BUCKET_NAME=${S3_BUCKET_NAME} dbt run\n.PHONY: dbt-run\n\ndbt-docs:\n\tsource ./.venv/bin/activate && cd dashboard/dbt && S3_BUCKET_NAME=${S3_BUCKET_NAME} dbt docs generate && dbt docs serve\n.PHONY: dbt-docs\n\ndashboard:\n\tsource ./.venv/bin/activate && cd dashboard && streamlit run dashboard.py\n.PHONY: dashboard\n"
},
{
"path": "src/benchmark.py",
"content": "\"\"\"\n\nThis Python script benchmarks the performance of running queries using duckdb as engine and \nan object storage as data source. \n\nThe script is minimal and not very configurable, but should be enough to give you a \nfeeling of how the different setups perform and the trade-offs involved.\n\nPlease note we basically start with 2019-04-01 and based on the -d flag we add days to the date.\nSo by increasing -d you will increase the amount of data to be processed. The map reduce version\nmanually unpacks the queries into queries-by-date and then runs them in parallel.\n\nDue to cold start, the first run of the serverless version may be slower than the others, so you should\nre-run the same script multiple times to get a better idea of the performance.\n\n\"\"\"\n\nimport os\nimport duckdb\nimport json\nimport statistics\nimport time\nfrom fastcore.parallel import parallel\nfrom quack import invoke_lambda, display_table\nfrom dotenv import load_dotenv\nfrom collections import defaultdict\nfrom rich.console import Console\n\n\n# get the environment variables from the .env file\nload_dotenv()\n\n\ndef run_benchmarks(\n bucket: str,\n repetitions: int,\n threads: int,\n days: int,\n is_debug: bool = False\n):\n test_location_id = 237\n execution_times = []\n # NOTE: as usual we re-use the same naming convention as in the setup script\n # and all the others\n partitioned_dataset_scan = f\"s3://{bucket}/partitioned/*/*.parquet\"\n # run the map reduce version\n repetition_times = []\n print(\"\\n====> Running map reduce version\")\n for i in range(repetitions):\n start_time = time.time()\n map_reduce_results = run_map_reduce(\n bucket=bucket,\n days=days,\n threads=threads,\n is_debug=is_debug\n )\n repetition_times.append(time.time() - start_time)\n time.sleep(3)\n \n execution_times.append({\n 'type': 'map_reduce',\n 'mean': round(sum(repetition_times) / len(repetition_times), 3),\n 'std': round(statistics.stdev(repetition_times), 3),\n 'test location': map_reduce_results[test_location_id]\n })\n # run the standard serverless version\n repetition_times = []\n print(\"\\n====> Running serverless duckdb\")\n for i in range(repetitions):\n start_time = time.time()\n results = run_serverless_lambda(\n partitioned_dataset_scan=partitioned_dataset_scan,\n days=days,\n is_debug=is_debug\n )\n repetition_times.append(time.time() - start_time)\n time.sleep(3)\n\n execution_times.append({\n 'type': 'serverless',\n 'mean': round(sum(repetition_times) / len(repetition_times), 3),\n 'std': round(statistics.stdev(repetition_times), 3),\n 'test location': results[test_location_id]\n })\n # run a local db querying the data lake\n repetition_times = []\n print(\"\\n====> Running local duckdb\")\n for i in range(repetitions):\n start_time = time.time()\n # just re-use the code inside the lambda without thinking too much ;-)\n con = duckdb.connect(database=':memory:')\n con.execute(f\"\"\"\n INSTALL httpfs;\n LOAD httpfs;\n SET s3_region='{os.environ.get('AWS_DEFAULT_REGION', 'us-east-1')}';\n SET s3_access_key_id='{os.environ['AWS_ACCESS_KEY_ID']}';\n SET s3_secret_access_key='{os.environ['AWS_SECRET_ACCESS_KEY']}';\n \"\"\")\n local_results = run_local_db(\n con=con,\n partitioned_dataset_scan=partitioned_dataset_scan,\n days=days,\n is_debug=is_debug\n )\n del con\n repetition_times.append(time.time() - start_time)\n time.sleep(3)\n \n execution_times.append({\n 'type': 'local',\n 'mean': round(sum(repetition_times) / len(repetition_times), 3),\n 'std': round(statistics.stdev(repetition_times), 3),\n 'test location': local_results[test_location_id]\n })\n # make sure the results are the same\n assert results[test_location_id] == map_reduce_results[test_location_id] == local_results[test_location_id], \"The results are not the same!\"\n\n # display results in a table\n console = Console()\n display_table(console, execution_times, title=\"Benchmarks\", color=\"cyan\")\n\n # all done, say goodbye\n print(\"All done! See you, duck cowboy!\")\n return \n\n\ndef run_local_db(\n con,\n partitioned_dataset_scan: str,\n days: int,\n is_debug: bool\n):\n single_query = \"\"\"\n SELECT \n pickup_location_id AS location_id, \n COUNT(*) AS counts \n FROM \n parquet_scan('{}', HIVE_PARTITIONING=1)\n WHERE \n DATE >= '2019-04-01' AND DATE < '2019-04-{}'\n GROUP BY 1\n \"\"\".format(\n partitioned_dataset_scan,\n \"{:02d}\".format(1 + days)\n )\n if is_debug:\n print(single_query)\n # just re-use the code inside the lambda with no particular changes\n _df = con.execute(single_query).df()\n _df = _df.head(1000)\n records = _df.to_dict('records')\n\n return { row['location_id']: row['counts'] for row in records }\n\n\ndef run_serverless_lambda(\n partitioned_dataset_scan: str,\n days: int,\n is_debug: bool\n):\n single_query = \"\"\"\n SELECT \n pickup_location_id AS location_id, \n COUNT(*) AS counts \n FROM \n parquet_scan('{}', HIVE_PARTITIONING=1)\n WHERE \n DATE >= '2019-04-01' AND DATE < '2019-04-{}'\n GROUP BY 1\n \"\"\".format(\n partitioned_dataset_scan,\n \"{:02d}\".format(1 + days)\n )\n if is_debug:\n print(single_query)\n response = invoke_lambda(json.dumps({ 'q': single_query, 'limit': 1000}))\n if 'errorMessage' in response:\n print(response['errorMessage'])\n raise Exception(\"There was an error in the serverless invocation\")\n records = response['data']['records']\n\n return { row['location_id']: row['counts'] for row in records }\n\n\ndef run_map_reduce(\n bucket: str,\n days: int,\n threads: int,\n is_debug: bool\n):\n query = \"\"\"\n SELECT \n pickup_location_id AS location_id, \n COUNT(*) AS counts \n FROM \n read_parquet('{}', HIVE_PARTITIONING=1)\n WHERE \n DATE >= '2019-04-{}' AND DATE < '2019-04-{}'\n GROUP BY 1\n \"\"\".strip()\n # prepare the queries for the map step\n queries = prepare_map_queries(query, bucket, days)\n if is_debug:\n print(queries[:3])\n assert len(queries) == days, \"The number of queries is not correct\"\n # run the queries in parallel\n payloads = [json.dumps({'q': q, 'limit': 1000 }) for q in queries]\n _results = parallel(\n invoke_lambda, \n payloads,\n n_workers=threads)\n # check for errors in ANY response\n if any(['errorMessage' in response for response in _results]):\n print(next(response['errorMessage'] for response in _results if 'errorMessage' in response))\n raise Exception(\"There was an error in the parallel invocation\")\n # do the \"reduce\" step in code\n results = defaultdict(lambda: 0)\n # loop over the results\n for response in _results:\n records = response['data']['records']\n for row in records:\n results[row['location_id']] += row['counts']\n \n return results\n\n\ndef prepare_map_queries(\n query: str,\n bucket: str,\n days: int\n ):\n # template for parquet scan\n queries = []\n for i in range(1, days + 1):\n start_day_as_str = \"{:02d}\".format(i)\n end_day_as_str = \"{:02d}\".format(i + 1)\n parquet_scan = f\"s3://{bucket}/partitioned/date=2019-04-{start_day_as_str}/*.parquet\"\n queries.append(query.format(parquet_scan, start_day_as_str, end_day_as_str))\n \n return queries\n\nif __name__ == \"__main__\":\n # make sure the envs are set\n assert 'S3_BUCKET_NAME' in os.environ, \"Please set the S3_BUCKET_NAME environment variable\"\n assert 'AWS_ACCESS_KEY_ID' in os.environ, \"Please set the AWS_ACCESS_KEY_ID environment variable\"\n assert 'AWS_SECRET_ACCESS_KEY' in os.environ, \"Please set the AWS_SECRET_ACCESS_KEY environment variable\"\n # get args from command line\n import argparse\n parser = argparse.ArgumentParser()\n parser.add_argument(\n \"-n\",\n type=int,\n help=\"number of repetitions\", \n default=3)\n # note: without reserved concurrency, too much concurrency will cause errors\n parser.add_argument(\n \"-t\",\n type=int,\n help=\"concurrent queries for map reduce\", \n default=20)\n parser.add_argument(\n \"-d\",\n type=int,\n help=\"number of days in April to query\", \n default=10)\n parser.add_argument(\n \"--debug\", \n action=\"store_true\",\n help=\"increase output verbosity\",\n default=False)\n args = parser.parse_args()\n # run the main function\n run_benchmarks(\n bucket=os.environ['S3_BUCKET_NAME'],\n repetitions=args.n,\n threads=args.t,\n days=args.d,\n is_debug=args.debug\n )"
},
{
"path": "src/dashboard/dashboard.py",
"content": "\"\"\"\n\nSimple dashboard for the taxi data based on Streamlit.\n\nIt re-uses through hugly imports the code from the quack.py script, and use seaborn to plot the data.\n\n\"\"\"\n\nimport sys\nimport os\nimport matplotlib.pyplot as plt\nimport pandas as pd\nimport seaborn as sns\nimport streamlit as st\nfrom dotenv import load_dotenv\n\n\n# get the environment variables from the .env file\nload_dotenv('../.env')\nassert 'S3_BUCKET_NAME' in os.environ, \"Please set the S3_BUCKET_NAME environment variable\"\nS3_BUCKET_NAME = os.environ['S3_BUCKET_NAME']\n# note: this is the same file we exported in the top_pickup_locations.sql query\n# as part of our data transformation pipeline\nPARQUET_FILE = f\"s3://{S3_BUCKET_NAME}/dashboard/my_view.parquet\"\n\n# import querying functoin from the runner\nsys.path.insert(0,'..')\nfrom quack import fetch_all\n# build up the dashboard\nst.markdown(\"# Trip Dashboard\")\nst.write(\"This dashboard shows KPIs for our taxi business.\")\nst.header(\"Top pickup locations (map id) by number of trips\")\n# hardcode the columns\nCOLS = ['PICKUP_LOCATION_ID', 'TRIPS']\n\n# get the total row count\nquery = f\"SELECT COUNT(*) AS C FROM read_parquet(['{PARQUET_FILE}'])\"\ndf, metadata = fetch_all(query, limit=1, display=False, is_debug=False)\nst.write(f\"Total row count: {df['C'][0]}\")\n\n# get the interactive chart\nbase_query = f\"\"\"\n SELECT \n location_id AS {COLS[0]}, \n counts AS {COLS[1]} \n FROM \n read_parquet(['{PARQUET_FILE}'])\n \"\"\".strip()\ntop_k = st.text_input('# of pickup locations', '5')\n# add a limit to the query based on the user input\nfinal_query = \"{} LIMIT {};\".format(base_query, top_k).format(top_k)\ndf, metadata = fetch_all(final_query, limit=int(top_k), display=False, is_debug=False)\n\n# if no error is returned, we plot the data\nif df is not None:\n fig = plt.figure(figsize=(10,5))\n sns.barplot(\n x = COLS[0],\n y = COLS[1],\n data = df,\n order=df.sort_values(COLS[1],ascending = False)[COLS[0]])\n plt.xticks(rotation=70)\n plt.tight_layout()\n st.pyplot(fig)\nelse:\n st.write(\"Sorry, something went wrong :-(\")\n\n# display metadata\nst.write(f\"Roundtrip ms: {metadata['roundtrip_time']}\")\nst.write(f\"Query exec. time ms: {metadata['timeMs']}\")\nst.write(f\"Lambda is warm: {metadata['warm']}\")\n \n"
},
{
"path": "src/dashboard/dbt/analysis/.gitkeep",
"content": ""
},
{
"path": "src/dashboard/dbt/dbt_project.yml",
"content": "name: 'taxi_dashboard'\nversion: '1.0.0'\n\nconfig-version: 2\nprofile: 'duckdb-taxi'\nsource-paths: [\"models\"]\nanalysis-paths: [\"analysis\"]\ntest-paths: [\"tests\"]\ndata-paths: [\"data\"]\nmacro-paths: [\"macros\"]\nsnapshot-paths: [\"snapshots\"]\n\ntarget-path: \"target\" \nclean-targets: \n - \"target\"\n - \"dbt_modules\"\n\nmodels:\n taxi:\n foundation:\n +materialized: view\n"
},
{
"path": "src/dashboard/dbt/macros/.gitkeep",
"content": ""
},
{
"path": "src/dashboard/dbt/models/taxi/top_pickup_locations.sql",
"content": "{{ config(materialized='external', location=\"s3://{{ env_var('S3_BUCKET_NAME') }}/dashboard/my_view.parquet\") }}\n\nSELECT \n location_id,\n counts\nFROM\n {{ ref('trips_by_pickup_location') }}\nORDER BY 2 DESC\nLIMIT 200"
},
{
"path": "src/dashboard/dbt/models/taxi/trips_by_pickup_location.sql",
"content": "SELECT \n pickup_location_id AS location_id, \n COUNT(*) AS counts \nFROM \n read_parquet(['s3://{{ env_var('S3_BUCKET_NAME') }}/dataset/taxi_2019_04.parquet']) \nGROUP BY 1"
},
{
"path": "src/dashboard/dbt/snapshots/.gitkeep",
"content": ""
},
{
"path": "src/dashboard/dbt/tests/.gitkeep",
"content": ""
},
{
"path": "src/data/.gitkeep",
"content": ""
},
{
"path": "src/local.env",
"content": "AWS_ACCESS_KEY_ID=\nAWS_SECRET_ACCESS_KEY=\nAWS_DEFAULT_REGION=us-east-1\nS3_BUCKET_NAME="
},
{
"path": "src/package.json",
"content": "{\n \"dependencies\": {\n \"serverless\": \"^3.30.1\",\n \"serverless-iam-roles-per-function\": \"^3.2.0\"\n }\n}\n"
},
{
"path": "src/quack.py",
"content": "\"\"\"\n\nPython script to interact with the serverless architecture. Query and limit parameters\ncan be passed through the command line, or the script can be run without parameters\nto check the status of the lambda (it will return the results from a pre-defined query).\n\n\nCheck the README.md for more details.\n\n\"\"\"\n\n\nimport os\nimport time\nimport boto3\nimport pandas as pd\nimport json\nfrom rich.console import Console\nfrom rich.table import Table\nfrom dotenv import load_dotenv\n\n\n# get the environment variables from the .env file\nload_dotenv()\n# we don't allow to display more than 10 rows in the terminal\nMAX_ROWS_IN_TERMINAL = 10\n# instantiate the boto3 client to communicate with the lambda\nlambda_client = boto3.client('lambda')\n\n\ndef invoke_lambda(json_payload_as_str: str):\n \"\"\"\n Invoke our duckdb lambda function. Note that the payload is a string,\n so the method should be called with json.dumps(payload)\n \"\"\"\n response = lambda_client.invoke(\n # the name of the lambda function should match what you have in your console\n # if you did not change the serverless.yml file, it should be this one:\n FunctionName='quack-reduce-lambda-dev-duckdb',\n InvocationType='RequestResponse',\n LogType='Tail',\n Payload=json_payload_as_str\n )\n\n # return response as dict\n return json.loads(response['Payload'].read().decode(\"utf-8\"))\n\n\ndef fetch_all(\n query: str,\n limit: int,\n display: bool=False,\n is_debug = False\n)-> pd.DataFrame:\n \"\"\"\n Get results from lambda and display them\n \"\"\"\n if is_debug:\n print(f\"Running query: {query}, with limit: {limit}\")\n # run the query\n start_time = time.time()\n response = invoke_lambda(json.dumps({'q': query, 'limit': limit}))\n roundtrip_time = int((time.time() - start_time) * 1000.0)\n # check for errors first\n if 'errorMessage' in response:\n print(f\"Error: {response['errorMessage']}\")\n # just raise an exception now as we don't have a proper error handling\n raise Exception(response['errorMessage'])\n # no error returned, display the results\n if is_debug:\n print(f\"Debug reponse: {response}\")\n\n rows = response['data']['records']\n # add the roundtrip time to the metadata\n response['metadata']['roundtrip_time'] = roundtrip_time\n # display in the console if required\n if display:\n console = Console()\n display_query_metadata(console, response['metadata'])\n display_table(console, rows)\n \n # return the results as a pandas dataframe and metadata\n return pd.DataFrame(rows), response['metadata']\n\n\ndef display_query_metadata(\n console: Console, \n metadata: dict\n ):\n \"\"\"\n Display the metadata returned by the lambda - we receive a dictionary with\n few properties (total time, echo of the query, is warm, etc.)\n \"\"\"\n # NOTE: we cut to 25 max the field values, to avoid the table to be too wide\n values = [{ 'Field': k, 'Value': str(v)[:50] } for k, v in metadata.items()]\n display_table(console, values, title=\"Metadata\", color=\"cyan\")\n\n return\n\n\ndef display_table(\n console: Console,\n rows: list, \n title: str=\"My query\", \n color: str=\"green\"\n ):\n \"\"\"\n We receive a list of rows, each row is a dict with the column names as keys.\n We use rich (https://rich.readthedocs.io/en/stable/tables.html) to display a nice table in the terminal.\n \"\"\"\n # build the table\n table = Table(title=title)\n # buld the header\n cols = list(rows[0].keys())\n for col in cols:\n table.add_column(col, justify=\"left\", style=color, no_wrap=True)\n # add the rows\n for row in rows[:MAX_ROWS_IN_TERMINAL]:\n # NOTE: we need to render str\n table.add_row(*[str(row[col]) for col in cols])\n # diplay the table\n console.print(table)\n\n return\n\n\ndef runner(\n bucket: str,\n query: str=None,\n limit: int=10,\n is_debug: bool = False\n):\n \"\"\"\n Run queries against our serverless (and stateless) database.\n\n We basically use duckdb not as a database much, but as an engine, and use\n object storage to store artifacts, like tables.\n\n If query and limits are not specified, we overwrite them with sensible choices.\n \"\"\"\n # if no query is specified, we run a simple count to verify that the lambda is working\n if query is None:\n # NOTE: the file path, after the bucket, should be the same as the one we have\n # in the run_me_first.py script. If you changed it there, you should change it here\n target_file = f\"s3://{bucket}/dataset/taxi_2019_04.parquet\"\n query = f\"SELECT COUNT(*) AS COUNTS FROM read_parquet(['{target_file}'])\"\n # since this is a test query, we force debug to be True\n rows, metadata = fetch_all(query, limit, display=True, is_debug=True)\n else:\n # run the query as it is\n rows, metadata = fetch_all(query, limit, display=True, is_debug=is_debug)\n\n return\n\n\nif __name__ == \"__main__\":\n assert 'S3_BUCKET_NAME' in os.environ, \"Please set the S3_BUCKET_NAME environment variable\"\n # get args from command line\n import argparse\n # declare basic arguments\n parser = argparse.ArgumentParser()\n parser.add_argument(\n \"-q\",\n type=str,\n help=\"query\", \n default=None)\n parser.add_argument(\n \"-limit\",\n type=int,\n help=\"max rows to return from the lambda\",\n default=10)\n parser.add_argument(\n \"--debug\", \n action=\"store_true\",\n help=\"increase output verbosity\",\n default=False)\n args = parser.parse_args()\n # run the main function\n runner(\n bucket=os.environ['S3_BUCKET_NAME'],\n query=args.q,\n limit=args.limit,\n is_debug=args.debug\n )"
},
{
"path": "src/requirements.txt",
"content": "requests==2.28.2\nstreamlit==1.20.0\npython-dotenv==1.0.0\nfsspec==2023.4.0\ns3fs==2023.4.0\ndbt-duckdb==1.4.1\nfastcore==1.5.27\nboto3==1.26.3\nmatplotlib==3.6.3\nseaborn==0.12.0"
},
{
"path": "src/run_me_first.py",
"content": "\"\"\"\n\nPython script to run a one-time setup for testing the serverless duckdb architecture.\n\nCheck the README.md for more details and for the prerequisites.\n\n\"\"\"\n\n\nimport os\n\nimport boto3\nimport requests\nimport pandas as pd\nfrom dotenv import load_dotenv\n\n\n# get the environment variables from the .env file\nload_dotenv()\n\n\ndef donwload_data(url: str, target_file: str):\n \"\"\"\n Download a file from a url and save it to a target file.\n \"\"\"\n r = requests.get(url)\n open(target_file, 'wb').write(r.content)\n\n return True\n\n\ndef download_taxi_data():\n \"\"\"\n Download the taxi data from the duckdb repo - if the file disappears, \n you can of course replace it with any other version of the same dataset.\n \"\"\"\n print('Downloading the taxi dataset')\n \n url = 'https://github.com/cwida/duckdb-data/releases/download/v1.0/taxi_2019_04.parquet'\n file_name = 'data/taxi_2019_04.parquet'\n donwload_data(url, file_name)\n\n return file_name\n\n\ndef upload_file_to_bucket(s3_client, file_name, bucket, object_name=None):\n \"\"\"\n Upload a file to an S3 bucket.\n \"\"\"\n from botocore.exceptions import ClientError\n \n try:\n print(f\"Uploading {object_name}\")\n response = s3_client.upload_file(file_name, bucket, object_name)\n except ClientError as e:\n print(f\"Error uploading file {file_name} to bucket {bucket} with error {e}\")\n return False\n \n return True\n\n\ndef upload_datasets(s3_client, bucket: str, taxi_dataset_path: str):\n \"\"\"\n Upload the datasets to the bucket, first as one parquet file, then as\n a directory of parquet files with hive partitioning.\n \"\"\"\n file_name = os.path.basename(taxi_dataset_path)\n # upload file as is, a single parquet file in the data/ folder of the target bucket\n is_uploaded = upload_file_to_bucket(\n s3_client, \n taxi_dataset_path,\n bucket,\n object_name=f\"dataset/{file_name}\"\n )\n is_uploaded = upload_partioned_dataset(\n bucket, \n taxi_dataset_path\n )\n\n return\n\n\ndef upload_partioned_dataset(\n bucket: str,\n taxi_dataset_path: str,\n partition_col: str = 'date'\n ):\n \"\"\"\n Use pandas to read the parquet file, then save it again as a directory\n on our s3 bucket. The final directory will have a subdirectory for each\n value of the partition column, and each subdirectory will contain parquet files.\n \"\"\"\n\n df = pd.read_parquet(taxi_dataset_path)\n df[partition_col] = pd.to_datetime(df['pickup_at']).dt.date\n target_folder = os.path.join('s3://', bucket, 'partitioned')\n print(f\"Saving data with hive partitioning ({partition_col}) in {target_folder}\")\n df.to_parquet(target_folder, partition_cols=[partition_col])\n \n return True\n\n\ndef setup_project():\n # check vars are ok\n assert 'S3_BUCKET_NAME' in os.environ, \"Please set the S3_BUCKET_NAME environment variable\"\n AWS_ACCESS_KEY_ID = os.environ.get('AWS_ACCESS_KEY_ID')\n AWS_SECRET_ACCESS_KEY = os.environ.get('AWS_SECRET_ACCESS_KEY')\n AWS_PROFILE = os.environ.get('AWS_PROFILE')\n assert (AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY) or AWS_PROFILE, \"Please set the AWS_ACCESS_KEY_ID & AWS_SECRET_ACCESS_KEY (or the AWS_PROFILE) environment variables\"\n \n # first download the data\n taxi_dataset_path = download_taxi_data()\n # upload the data to the bucket\n s3_client = boto3.client('s3')\n upload_datasets(\n s3_client,\n os.environ['S3_BUCKET_NAME'],\n taxi_dataset_path\n )\n # all done\n print(\"All done! See you, duck cowboy!\")\n return \n\n\nif __name__ == \"__main__\":\n setup_project()"
},
{
"path": "src/serverless/Dockerfile",
"content": "FROM public.ecr.aws/lambda/python:3.9\n\n# Install pip and other dependencies\nRUN pip3 install --upgrade pip \\\n && yum install gcc gcc-c++ -y \\\n && pip3 install pandas==1.5.3 duckdb==0.7.1 --target \"${LAMBDA_TASK_ROOT}\"\n\nENV HOME=/home/aws\n\nRUN mkdir /home/aws && python3 -c \"import duckdb; duckdb.query('INSTALL httpfs;');\"\n\nCOPY app.py ${LAMBDA_TASK_ROOT}\n\n# Set the CMD to the lambda handler\nCMD [ \"app.handler\" ]"
},
{
"path": "src/serverless/app.py",
"content": "import uuid\nimport os\nimport time\nimport duckdb\nimport pandas as pd\n\n\ncon = None # global conn object - we re-use this across calls\nDEFAULT_LIMIT = 20 # if we don't specify a limit, we will return at most 20 results\n\n\ndef return_duckdb_connection():\n \"\"\"\n Return a duckdb connection object\n \"\"\"\n duckdb_connection = duckdb.connect(database=':memory:')\n duckdb_connection.execute(f\"\"\"\n LOAD httpfs;\n SET s3_region='{os.environ['AWS_REGION']}';\n SET s3_session_token='{os.environ['AWS_SESSION_TOKEN']}';\n \"\"\"\n )\n\n return duckdb_connection\n\n\ndef handler(event, context):\n \"\"\"\n Run a SQL query in a memory db as a serverless function\n \"\"\"\n\n is_warm = False\n # run a timer for info\n start = time.time()\n global con\n if not con:\n # create a new connection\n con = return_duckdb_connection()\n else:\n # return to the caller the status of the lambda\n is_warm = True\n\n # get the query to be executed from the payload\n event_query = event.get('q', None)\n limit = int(event.get('limit', DEFAULT_LIMIT))\n results = []\n if not event_query:\n print(\"No query provided, will return empty results\")\n else:\n # execute the query and return a pandas dataframe\n _df = con.execute(event_query).df()\n # take rows up the limit, to avoid crashing the lambda\n # by returning too many results\n _df = _df.head(limit)\n results = convert_records_to_json(_df)\n \n # return response to the client with metadata\n return wrap_response(start, event_query, results, is_warm)\n\n\ndef convert_records_to_json(_df):\n if len(_df) > 0:\n # convert timestamp to string to avoid serialization issues\n cols = [col for col in _df.columns if _df[col].dtype == 'datetime64[ns]']\n _df = _df.astype({_: str for _ in cols})\n\n return _df.to_dict('records')\n\n\ndef wrap_response(start, event_query, results, is_warm):\n \"\"\"\n Wrap the response in a format that can be used by the client\n \"\"\"\n return {\n \"metadata\": {\n \"timeMs\": int((time.time() - start) * 1000.0),\n \"epochMs\": int(time.time() * 1000),\n \"eventId\": str(uuid.uuid4()),\n \"query\": event_query,\n \"warm\": is_warm\n },\n \"data\": {\n \"records\": results\n }\n }\n"
},
{
"path": "src/serverless.yml",
"content": "service: quack-reduce-lambda\nuseDotenv: true\n\nprovider:\n name: aws\n region: ${env:AWS_DEFAULT_REGION, 'us-east-1'}\n architecture: arm64\n memorySize: 3008\n timeout: 600\n ecr:\n images:\n quackimageblog:\n path: ./serverless\n platform: linux/arm64\n\nfunctions:\n duckdb:\n ephemeralStorageSize: 3008\n image:\n name: quackimageblog\n iamRoleStatements:\n - Effect: Allow\n Action:\n - s3:GetObject\n - s3:PutObject\n Resource:\n - Fn::Join:\n - ''\n - - Fn::GetAtt:\n - S3Bucket\n - Arn\n - '/*'\n - Effect: Allow\n Action:\n - s3:ListBucket\n Resource:\n - Fn::GetAtt:\n - S3Bucket\n - Arn\n\nresources:\n Resources:\n S3Bucket:\n Type: AWS::S3::Bucket\n DeletionPolicy: Retain\n Properties:\n BucketName: ${env:S3_BUCKET_NAME}\n\nplugins:\n - serverless-iam-roles-per-function\n"
}
]