So what exactly is a Data Sketch and how does it help?

Just to recap the original article, imagine you are asked to find the unique number of users who visited a site today. You would ideally do a count distinct user query on the sessions table and it would workout fine, unless the number of events in the order of millions or billions. If thats the case your query would have to visit every record in the table to ensure it has the exact count of all unique users. This could either break the database or take forever depending on the table size and the compute being used. But we could use a data sketch and get an approximate count in under a second. How is that? For this we can use an analogy.

Imagine a large corporate like Walmart, who has millions of employees working each day. Now there’s an inspector who’s job is to find out how many employees checked in yesterday. Note that an employee could checkin in multiple times for different shifts, so his duplicate entry will not be counted. Now he can manually note down whenever an employee checks in for the day on his notebook. Whenever someone enters the office he could check his list if the person had already checked in for the day or not, and make an entry accordingly. This can work but it’ll be slow and tiresome. So he uses another method. One that gets an approximate count as he doesn’t need the exact number. He assigns each employee a unique tag when they check in. The tag could be any number between 1 to 4 billion (We need a large number space to avoid tag collisions). The rules are:

• The number tag will be unique and exclusive to that employee

• The number tag generator randomly picks a number from a UNIFORM distribution.

So when the employees start checking in, they are assigned different tags: 600, 11000, 565732, 573, 64968, 395345.. and so on from the space of 4,000,000,000.

What the inspector does, is that he keeps track of only the LOWEST K tags. Lets say its 1000. That is enough for him to get the approximate count of the employees who checked in for the day.

How does that happen?

Here’s the logic: Given that these tags are assigned in a uniform distribution, we can sort of assume that the gaps among all the tags are even. So what does that mean? It means that the gap between each number tag on a scale on 0 to 4 billion is around Total Space / Total tags. This also means that the gap between the first 1000 tags is the same as all the tags. i.e K’th tag / K. i.e V(K) / K.

Does that make sense? If the 1000th lowest number tag is 2,779,252 then we do 2,779,252 / 1000 = 2779. That is the average assumed gap between tags. Now to get the total count, we just need to divide this gap with the total space. i.e 4,000,000,000 / 2779 = 1,439,366. This is the approximate count of all the unique tags, which is the unique number of employees who checked in today.

It’s sort of like this diagram above. If we arrange the dots in an organized manner, we can figure out an approx population count immediately by multiplying the number of rows and columns. As long as we know the number of rows and columns, we can arrive at the total count quickly, no matter how big the table gets. While when left unorganized (left) we would have to count each dot individually.

Hopefully you have a good conceptual idea of what a data sketch is. We can verify the data sketch with an actual code snippet.

import heapq as hq
import mmh3
import numpy as np

# ── Constants ─────────────────────────────────────────────────────────────────
RNG_SEED        = 25
MIN_EMPLOYEE_ID = 1
MAX_EMPLOYEE_ID = 4_000_000   # Walmart employee IDs in [1, 4M]
NUM_CHECKINS    = 2_000_000   # total check-ins today (duplicates expected)
K               = 1_000       # keep K lowest tags; error ≈ 1/√K ≈ 3.2%
TAG_SPACE       = 1 << 32     # tags in [1, ~4 billion] — must be >> unique employees
                              # to avoid collisions (1M tag space < 1.57M uniques = broken)

# ── Step 1: Simulate check-ins ────────────────────────────────────────────────
# 2M check-ins from employees with IDs in [1, 4M] — duplicates expected
rng      = np.random.default_rng(RNG_SEED)
checkins = rng.integers(MIN_EMPLOYEE_ID, MAX_EMPLOYEE_ID + 1, size=NUM_CHECKINS)

# ── Step 2: Build KMV sketch ──────────────────────────────────────────────────
# Each employee gets a deterministic integer tag in [1, TAG_SPACE] via hash.
# Same employee → same tag always, so duplicate check-ins are naturally skipped.
# Inspector keeps only the K lowest tags — O(K) memory regardless of stream size.
heap     = []    # max-heap (negated) — heap[0] = -(current K-th lowest tag)
heap_set = set() # mirrors heap for O(1) lookup — size stays ≤ K

for emp_id in checkins:
    tag = (mmh3.hash64(str(int(emp_id)), signed=False)[0] % TAG_SPACE) + 1

    if tag in heap_set:
        continue  # same tag already tracked (same employee checking in again)

    if len(heap) < K:
        hq.heappush(heap, -tag)
        heap_set.add(tag)
    elif tag < -heap[0]:          # smaller than current K-th tag — swap in
        evicted = -hq.heapreplace(heap, -tag)
        heap_set.discard(evicted)
        heap_set.add(tag)

# ── Step 3: Estimate unique headcount ─────────────────────────────────────────
# The K-th lowest tag sits at position K in the sorted uniform distribution.
# Average gap between tags = kth_tag / K
# Total unique employees  ≈ TAG_SPACE / gap = TAG_SPACE * K / kth_tag

kth_tag  = -heap[0]
gap      = kth_tag / K
estimate = TAG_SPACE / gap          # = TAG_SPACE * K / kth_tag

# ── Step 4: Verify ────────────────────────────────────────────────────────────
exact      = np.unique(checkins).size
error_rate = abs(estimate - exact) / exact

print(f"Tag space         : 1 to {TAG_SPACE:,}") # 1 to 4,294,967,296
print(f"K-th tag (K={K}) : {kth_tag:,}") # 2,779,252
print(f"Avg gap           : {gap:,.0f}") # 2,779
print(f"Estimate          : {estimate:,.0f}") # 1,545,368
print(f"Exact             : {exact:,}") # 1,574,043
print(f"Error             : {error_rate:.2%}  (expected ~{1/K**0.5:.2%})") # 1.82%  (expected ~3.16%)

Its all about probability

This can be confusing at first. How do we evenly distribute these numbers and ensure they are unique? The secret that this algo works is due to the magic inside the hash function (MurmurHash3). That is the crux of this solution. The hash function needs to be well tested and vetted to generate a uniform distribution of numbers on a large scale. The MurmurHash3 hash function in this case does that job. It distributes the hash output numbers equally across the space in complete isolation. And that happens due to the law of large numbers. We will dive deeper into that in another article, but to be brief, this algorithm ensures that every tag (hash) it generates has an equal chance of being placed anywhere on the scale. And with the probability laws in effect, these odds for these tags to get evenly distributed are massive (given the data points are large too).

Using this method, we can quickly estimate the counts for a metric. It comes in handy when we just need approximations and not absolute numbers. Especially on dashboards or reports.

These functions are available in most of the data storage platforms. In Spark you can use it to aggregate counts across clusters without shuffling data. In Postgres (postgresql-hll extension), you can use it to aggregate counts across days, weeks, months etc without manually going through the rows. So the next, time you require quick aggregations, consider using Data Sketches instead of count(). It’ll definitely save you a lot of compute cycles.

Anyone who has worked with PySpark has probably come across User Defined Functions. UDF’s are functions in PySpark which can be used to do very specific custom tasks that might be difficult or not possible in native PySpark. Here’s a simple example:

This is a simple dataframe with username and date of birth:

We will create a Python UDF here to infer the users current age via using this date of birth.

from pyspark.sql.functions import udf
import datetime as dt
from dateutil import relativedelta as rd
from pyspark.sql.types import IntegerType

@udf(IntegerType())
def infer_age_udf(date_of_birth):
    try:
        return rd.relativedelta(dt.datetime.now(), date_of_birth).years
    except Exception:
        return None


users_df_udf = users_df.withColumn("age", infer_age_udf(col("date_of_birth")))

This is the end result:

This a common UDF implementation we often see in PySpark. UDF’ s however comes with a hidden underlying cost. Most new developers assume that UDF’s are just a quick way to implement Python logic and get something done quickly. PySpark usually gives this impression that it’s all within the Python ecosystem, so it runs one and the same. But that is an assumption far from the truth.

So what happens when you use a UDF in PySpark?

Contrary to our beliefs, the PySpark application does not really run Python. Sure it has a Python syntax, but internally, it is just a Python wrapper communicating with a Scala application running on the JVM. When we call a built-in PySpark function, we are ultimately calling a highly optimized Scala function.

Scala is a compiled language that runs on the JVM. Because Spark has full control over this native code, its Catalyst Optimizer can inspect and optimize it to a massive extent. This is what makes Spark so fast.

The exception to this rule is UDFs. When we create a custom Python UDF, it does not automatically translate into Scala. Spark has no Scala equivalent for your custom Python logic. To run this code, the Spark Executor JVM has to exclusively spin up a separate Python worker process on each node.

The major overhead here is data serialization. To pass data out of the JVM and into the Python worker, Spark has to serialize (pickle) the data. The Python worker executes your code, and then the results must be serialized again and sent back to the JVM.

Furthermore, because this execution happens inside a black-box Python process outside the JVM, Spark's Catalyst Optimizer cannot inspect or optimize your code. You are completely at the mercy of the Python worker.

So what can we do about it?

Native Spark (The best approach)

In most instances, UDF’s can be completely avoided by using Spark native functions. String manipulations, mathematical calculations, dictionary accessors can all be implemented in PySpark itself. The following code will give us the exact same results without touching a UDF, utilizing native PySpark code but giving us a 13x speedup:

from pyspark.sql.functions import col, current_date, months_between, floor

users_df_native = users_df.withColumn(
    "age",
    floor(months_between(current_date(), col("date_of_birth")) / 12)
)

Pandas UDF (Vectorized operations)

If you must use Python, you can speed up processing by using Pandas UDFs. Pandas UDFs utilize Apache Arrow which does not require data serialization for copying data and leverages vectorized operations which speed up processing.

from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import IntegerType
import json
import pandas as pd
import datetime as dt

@pandas_udf(IntegerType())
def infer_age_pdf(dob: pd.Series) -> pd.Series:
    today = dt.datetime.now()
    dob = pd.to_datetime(dob)
    
    age = (today.year - dob.dt.year)
    birthday_past = (today.month < dob.dt.month) | ((today.month == dob.dt.month) & (today.day < dob.dt.day))

    # Calculate age and subtract 1 if the birthday hasn't occurred yet
    return age - birthday_past.astype(int)

users_df_pdf = users_df.withColumn("age", infer_age_pdf(col("date_of_birth")))

Scala UDF

Another option is to write your UDF in Scala. These are extremely performant and fit natively into the Spark ecosystem. The JVM executes them directly, meaning there is zero Python serialization overhead.

%scala

import org.apache.spark.sql.api.java.UDF1
import java.time.LocalDate
import java.time.Period
import java.sql.Date

class InferAge extends UDF1[java.sql.Date, Integer] {
  def call(dob: java.sql.Date): Integer = {
    if (dob != null) {
      val birthDate = dob.toLocalDate
      val today = LocalDate.now()
      val age = Period.between(birthDate, today).getYears
      age
    } else null
  }
}
spark.udf.register("infer_age_scala", new InferAge(), org.apache.spark.sql.types.IntegerType)

from pyspark.sql.functions import expr

users_df_scala = users_df.withColumn("age", expr("infer_age_scala(date_of_birth)"))

Benchmarks

Here are some benchmarks for this Dataframe of 1,000,000 rows. You can see the massive difference in performance among these methods. The native Scala options are extremely fast, followed by the vectorized pandas UDF implementation. The Python UDF is around 13x slower due to the Python serialization and execution overhead.

These Python UDF alternatives will work most of the time. There are a few exceptions like using Python machine learning libraries which don’t have Scala/Java equivalents. Those are acceptable cases where we would need Python UDF’s to get the job done.

When writing code, ensure you keep these things in mind as and when dealing with your data. These tiny optimization make a huge difference during heavy compute operations.

If you’d like to have a tool to ensure such issues are highlighted in your code, you can use CatalystOps within your VS code.

This article is just to journal and summarize my learnings from this chapter on partitioning.

What is Sharding and Partitioning?

Sharding is a process of dividing a database into smaller batches and distributing them across multiple machines. Partitioning is how we divide that data into smaller batches to improve query efficiency. Although the terms are used interchangeably, sharding often refers to horizontal partitioning (split by rows instead of columns) across different machines.

But why do we do that? Well, the size of your database without sharding is limited by the size of your physical hard drive. When your database reaches that physical limit (i.e., 1TB, 2TB, etc.), you will encounter a problem. So instead of storing all this data on a single machine, we can distribute it across 10 machines, each handling it independently.

Sharding allows a larger number of users to access the database simultaneously. Instead of being limited by the read/write speed of one hard drive, we will be limited by the combined read/write speed of multiple hard drives.

It also distributes the risk of losing data across the drives. If the data were sitting on one drive, then we could have to rely only on the health of that drive. If it fails, we lose all the data. With a sharded database, if a drive fails, we will lose only the data in that one shard.

So how do we go about distributing this data? There are various partitioning techniques. Let’s have a look at those.

Partitioning Strategies

Random / Equal Partitioning

Given that we want to distribute data across the nodes, we could perform an equal split of the data and push it to different nodes. While this will solve our distribution problem, we will have a tough time finding a record during queries, as we would have no idea which partition it is present in.

Round Robin Partitioning

In round robin partitioning, we equally balance the shards by storing the next incoming record of data in the next available partition in the cycle.

For example, if we have 3 nodes/shards and 4 records, then:

1. Record 1 goes to partition 1

2. Record 2 goes to partition 2

3. Record 3 goes to partition 3

4. Record 4 goes to partition 1

and so on. This way, all the shards have an equal amount of data. The issue with the approach is that while querying, the database has no idea on which shard the record it seeks lives on. It will have to make a scatter-gather query where all the nodes receive the request and start searching for the record.

List Partitioning

A better approach to partitioning is based on a particular key. For example, say we have a list of people with their names and country of origin. We could partition the list by country of origin. All people from India could be in one partition, all the people from Korea in another, etc. This approach will help us to easily find the node that has a particular person by matching them with the country of origin. The issue with this approach, though, is that it’s prone to skewness and hotspots. Which means that if 80% of people in the data are from India, they will all end up in one node, and the rest of 20% will be across other nodes. So the majority of the system load might be taken up by just one node, which is not desirable. Also, a condition for this approach to succeed is that the value of this attribute should have high cardinality. i.e If there are people only from 4 countries and there are 10 nodes, the data will be stored on only 4 nodes. The other 6 nodes will not have any records assigned to them.

Hash Partitioning

In the above example, we partitioned the database by country. In such cases, there might be a large number of partitions created and only a few nodes. i.e 195 partitions (nodes), 3 nodes. Now, we could store multiple partitions on one node, but how can we find the data when there is a query? The solution to this is a hash-based partitioning scheme. In hash-based partitioning, we insert the key into a hash function, which spits out the node number that the record should go to. For example, if we partition by country, and we have 3 nodes, then India might be assigned to node 1, Korea to node 2, the UK to 3, Denmark to 2, and Germany to 1, and so on. When there is a query, we can use the same hash function to find which node the data lives on.

Using a decent hashing function, we could distribute the records pretty uniformly across the nodes. The only drawback of hash partitioning is that it destroys ordering. Given that the records are mapped to nodes in a pseudo-random order, if we make a query like fetch the last 100 records of a user, then the database engine will have to fetch the last 100 records from each node, sort them out in the main node, and then pick the latest 100, which is very inefficient.

Range Partitioning

In the range-based partitioning, the data is divided by a key range. For example, if we have a list of users and we range partition by country name, then all the names starting from A-C might be put in node 1, D-F in another, G-K in another. This way, we would know which nodes contain which ranges of the key, and we can fetch and write to these nodes accordingly. An advantage of this method is that if one key dominates the dataset, then it can be split into a smaller group and pushed to another node. For examples if the I names are a large set, you can split it as I - Iq, Ir-U into another node. We could also use this to repartition a node in case the data grows too large.

The one issue with range partitioning, however, might arise during writes. For example, i.e in user event activities, if we partition the data by timestamp (i.e., one day in one partition), all the writes for the one day will flow to one node, overloading it, while the others remain idle. Hence, it is very important to choose the right partition keys here.

Composite Key Partitioning

Composite keys comprise of two parts. Partition key and Sort key. One functions at the inter-node level (routing data to the right server) and the other at the intra-node level (organizing data on the disk). Composite Key partitioning helps us solve the issues faced in Hash and Range partitioning by grouping and sorting related data together on one node. For example, for user event activities flowing in, data may first be partitioned by the user and then stored in the nodes. Within those nodes, the data for these users will be grouped and sorted by timestamp.

For example, say we have user data that contains name, country, and timestamp. We can have a primary partition key that will partition the data by User. Within these partitions, the data will be grouped by user and sorted using the timestamp field, i.e., Carol’s records will be stored together and sorted by timestamp. During a range query, for example: “All records for Carol from 31st December 2025 between 11.00 AM and 12.00 PM”, the database will find the node by the user name and then fetch all her records between the timestamps of 11.00 AM and 12.00 PM as they are already sorted and stored together.

Measurements to do while partitioning

While choosing partitioning and indexing schemes, it is important to take a step back and question what exactly we are trying to achieve. We need to have a look at the data, the frequency, and the type of read and write queries, what the hotspots are, etc. Once we know this, we can implement a partitioning scheme.

Some factors to keep in mind are:

• For a range query, the best performance can be achieved if the data to be fetched for that query sits in one node in one partition adjacent to each other.

• For multiple read queries, it’s best that they can be distributed across multiple nodes, and each node can serve the requested data without requiring data from other nodes.

• For write queries, it’s best to have each write request hit different nodes so that the load is distributed.

Of course, we cannot have a perfect setup for this, and there will always be a compromise, but it’s best to keep this in mind when designing a system. The type of queries will give you a good idea of how to design the scheme.