In AI data processing scenarios, the analysis and manipulation of multimodal data (such as text, images, point clouds, audio, and more) have become core requirements. Daft, as a lightweight and efficient distributed multimodal data processing framework, relies on user-defined functions (UDFs) as a key capability to support these complex operations. Within the Daft ecosystem, using UDFs to implement customized processing of multimodal data (for example, image feature extraction, text format cleansing, cross-modal data association, and so on) is not only a common industry practice, but also highly accessible due to the framework's low entry barrier. Developers do not need to deeply understand its underlying distributed logic; by simply encapsulating business logic in concise code, they can quickly integrate custom data processing operations into Daft's data pipeline, significantly lowering the technical threshold for multimodal data processing.
To meet different development habits and requirements across various scenarios, Daft provides diverse usage options, including Batch UDF, Row-wise UDF, Async UDF, and more, while supporting seamless switching between local debugging and distributed deployment.
Caution: Efficient execution of Daft UDFs is not "out-of-the-box"; their performance is highly dependent on parameter configuration tailored to the business—especially the two core parameters, batch_size (batch size) and concurrency (number of concurrent tasks). The configuration of these two parameters must be based on an understanding of Daft's operating mechanism: batch_size determines the maximum amount of data processed per UDF invocation. If set too high, it may cause memory overflow; if too low, it increases function call overhead. concurrency controls the number of UDF tasks executed in parallel. If set too high, it may lead to resource contention (such as CPU, GPU, or network bandwidth bottlenecks); if too low, distributed resources cannot be fully utilized. If parameter settings do not match the business data volume (such as the size of a single data item, total data volume) or hardware resources, it may result in decreased task concurrency efficiency in mild cases, or task blocking or crashes in severe cases. Therefore, when developers use Daft UDFs, they need to adjust parameters appropriately based on the specific business scenario (such as the type of multimodal data and the complexity of processing logic) and resource configuration in order to maximize Daft's distributed advantages.
Daft supports running UDFs in either single-node mode or distributed mode. There are some subtle differences in the scheduling mechanism. Running UDFs in distributed mode relies on Ray's distributed scheduling capabilities. Next, we will use an example to illustrate the aspects to consider when configuring UDF parameters for running a Daft Batch UDF in distributed mode with Ray.
tip
The current Daft version is v0.6.4.
import logging import time import os def configure_logging(): logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", datefmt='%Y-%m-%d %H:%M:%S.%s'.format()) logging.getLogger("tracing.span").setLevel(logging.WARNING) logging.getLogger("daft_io.stats").setLevel(logging.WARNING) logging.getLogger("DaftStatisticsManager").setLevel(logging.WARNING) logging.getLogger("DaftFlotillaScheduler").setLevel(logging.WARNING) logging.getLogger("DaftFlotillaDispatcher").setLevel(logging.WARNING) logging.getLogger("MyBatchUdf").setLevel(logging.WARNING) configure_logging() import daft from daft import col, Series @daft.udf(return_dtype=daft.DataType.string()) class MyBatchUdf: def __init__(self, sleep: int = 0): self._id = id(self) self.counter = 0 self.sleep = sleep self.logger = logging.getLogger(f"MyBatchUdf-{self._id}") def __call__(self, a: Series) -> Series: self.counter += len(a) self.logger.info(f"Actor got {len(a)} records, total: {self.counter}") if self.sleep > 0: time.sleep(self.sleep) return a if __name__ == "__main__": if os.getenv("DAFT_RUNNER", "native") == "ray": import ray ray.init(runtime_env={"worker_process_setup_hook": configure_logging}) partition = 2 udf_concurrency = 4 udf_batch_size = 200 sleep_seconds = 2 paths = [f"file-{i}" for i in range(1000)] df = daft.from_pydict({"path": paths}) # Repartition not supported on the NativeRunner df = df.repartition(partition, "path") udf = (MyBatchUdf .override_options(batch_size=udf_batch_size) .with_concurrency(udf_concurrency) .with_init_args(sleep=sleep_seconds)) df = df.with_column("batch", udf(col("path"))) df.collect()
In the example provided above, we implemented a MyBatchUdf that includes a counter to record the number of data entries processed by the current UDF. Additionally, this UDF accepts a column of type String and returns the column itself. The main data processing workflow is as follows:
path.from_pydict contains only one partition, so we repartition the DataFrame based on the path column. After repartitioning, the DataFrame contains two partitions.Let us first examine the issues that arise when parameter configuration is not optimal.
When we run this task in Ray single-node mode, we obtain the following log output. The log shows that a total of four UDFActor instances were started, but the number of data entries processed by each was [110, 200, 290, 400], which does not match our expectation of 250 entries per instance. Additionally, we can see that two UDF instances were invoked twice, and the time between invocations matches the configured sleep time, indicating that the same UDF instance serially processes multiple batches of data it receives. If the execution time for each data entry is consistent, then the UDF instance processing 400 entries becomes the long tail, and only one concurrent instance will be processing data in the later stages of the task.
(UDFActor pid=80511) 2025-09-28 16:40:48.1759048848 - MyBatchUdf-5155349264 - INFO - Actor got 200 records, total: 200 (UDFActor pid=80510) 2025-09-28 16:40:48.1759048848 - MyBatchUdf-4467594448 - INFO - Actor got 200 records, total: 200 (UDFActor pid=80509) 2025-09-28 16:40:48.1759048848 - MyBatchUdf-5098474896 - INFO - Actor got 110 records, total: 110 (UDFActor pid=80512) 2025-09-28 16:40:48.1759048848 - MyBatchUdf-4548224592 - INFO - Actor got 200 records, total: 200 (UDFActor pid=80510) 2025-09-28 16:40:50.1759048850 - MyBatchUdf-4467594448 - INFO - Actor got 90 records, total: 290 (UDFActor pid=80512) 2025-09-28 16:40:50.1759048850 - MyBatchUdf-4548224592 - INFO - Actor got 200 records, total: 400
The data skew observed here is related to the execution of repartition (the specific reason will be explained later). Let us see what happens if the repartition operation is not performed. From the following logs, we can see that the four concurrent processes each handled [200, 200, 200, 400] records. Although there are five batches, so one of the concurrent processes handles two batches, this situation appears to be better than the previous one.
(UDFActor pid=96995) 2025-09-28 16:55:06.1759049706 - MyBatchUdf-5049017488 - INFO - Actor got 200 records, total: 200 (UDFActor pid=96996) 2025-09-28 16:55:06.1759049706 - MyBatchUdf-5237938000 - INFO - Actor got 200 records, total: 200 (UDFActor pid=96993) 2025-09-28 16:55:06.1759049706 - MyBatchUdf-5112045840 - INFO - Actor got 200 records, total: 200 (UDFActor pid=96998) 2025-09-28 16:55:06.1759049706 - MyBatchUdf-5887815120 - INFO - Actor got 200 records, total: 200 (UDFActor pid=96998) 2025-09-28 16:55:08.1759049708 - MyBatchUdf-5887815120 - INFO - Actor got 200 records, total: 400
But what if the task is running on a Ray cluster with multiple workers? When this task is run on a Ray cluster with two workers, as shown in the logs below, only two concurrent processes actually handled the data, with the number of records processed being [400, 600]. The concurrency for data processing did not reach the expected level.
(UDFActor pid=13825) 2025-09-28 17:08:52.1759050532 - MyBatchUdf-6049781264 - INFO - Actor got 200 records, total: 200 (UDFActor pid=13827) 2025-09-28 17:08:52.1759050532 - MyBatchUdf-5088798160 - INFO - Actor got 200 records, total: 200 (UDFActor pid=13825) 2025-09-28 17:08:54.1759050534 - MyBatchUdf-6049781264 - INFO - Actor got 200 records, total: 400 (UDFActor pid=13827) 2025-09-28 17:08:54.1759050534 - MyBatchUdf-5088798160 - INFO - Actor got 200 records, total: 400 (UDFActor pid=13825) 2025-09-28 17:08:56.1759050536 - MyBatchUdf-6049781264 - INFO - Actor got 200 records, total: 600
In addition to mismatches between the number of partitions and the number of workers causing lower-than-expected concurrency, an inappropriate batch_size setting can also result in concurrency not being achieved. Suppose we set batch_size to 500; from the logs below, we can see that even when running on a single node, only two concurrent processes handled the data, with the number of records processed being [500, 500]. The concurrency for data processing did not reach the expected level.
(UDFActor pid=37087) 2025-09-28 17:29:33.1759051773 - MyBatchUdf-5215119504 - INFO - Actor got 500 records, total: 500 (UDFActor pid=37086) 2025-09-28 17:29:33.1759051773 - MyBatchUdf-4822819216 - INFO - Actor got 500 records, total: 500
From the above examples, we can see that an unreasonable configuration of the number of DataFrame partitions, the batch size of the UDF, and the concurrency setting of the UDF can lead to data skew and the issue of not achieving the desired concurrency. Below, we will explain the operating mechanism of Daft and the specific reasons from a theoretical perspective.
When running a Daft job in distributed mode, the entire job code is first constructed as a logical execution plan (LogicalPlan). When df.collect() is executed, a materialization operation occurs. The entire logical plan is optimized by the built-in Daft optimizer, then converted into a distributed physical execution plan (DistributedPhysicalPlan), and finally handed over to DistributedPhysicalPlanRunner to execute the entire plan.
DistributedPhysicalPlanRunner converts the physical plan into an executable pipeline, where each node in the plan corresponds to a PipelineNode. Ultimately, this pipeline is split into one or more SwordfishTask. Under normal circumstances, the number of SwordfishTask instances generated corresponds to the number of partitions in the DataFrame. These tasks are first distributed to RaySwordfishActor according to the relevant scheduling strategy (each Worker launches a RawSwordfishActor). When the task scheduling strategy is set to Spread, SwordfishTask instances are evenly assigned to multiple RaySwordfishActor based on the idle status of the Workers.
RaySwordfishActor is responsible for executing SwordfishTask, using NativeExecutor to run the LocalPhysicalPlan contained in the task. (At this point, the execution logic for Native Runner and Ray Runner is essentially the same. The difference is that in a distributed scenario, each SwordfishTask contains only the data for a single partition (the size of psets is 1), whereas Native Runner includes data from all partitions.) Each Swordfish Task contains the entire stage plan. When executing a UDF operator, the data in the current SwordfishTask is split according to batch_size into individual morsels/MicroPartitions, which are then scheduled to UDFActors on the current worker (executing UDFActor.eval_input).
Therefore, there are four key parameters to configure here:
RaySwordfishActor, and a single RaySwordfishActor can run multiple SwordfishTask instances "simultaneously" (asynchronously executing LocalPhysicalPlan). When the plan runner schedules tasks to each RaySwordfishActor, it determines whether to dispatch a task based on the current load of RaySwordfishActor. Tasks that are not dispatched remain in a waiting state and can only be sent after previous tasks have finished. Therefore, the number of workers and the total CPU count of the workers determine the maximum concurrency for task execution.SwordfishTask, so the number of SwordfishTask instances determines the maximum parallelism of the entire pipeline. If the number of SwordfishTask instances is less than the number of RaySwordfishActor, some RaySwordfishActor will remain idle because they do not receive any tasks. When the number of DataFrame partitions is relatively small, you can use repartition or into_partition to increase the number of DataFrame partitions, thereby increasing the number of tasks. You can also use into_batches to split a partition into multiple tasks, with each task containing part of the original partition's data, thus increasing the maximum concurrency for task execution.UDFActor instances are launched. When the workers in the Ray cluster are homogeneous, UDFActor instances are generally distributed evenly across each worker. If a worker has a UDFActor for the given UDF (for example, when UDF concurrency is less than the number of workers, some workers will not launch a UDFActor for that UDF), the corresponding RaySwordfishActor will only schedule morsels to the UDFActor on the current worker. Therefore, when the number of morsels split by RaySwordfishActor is less than the number of UDFActor instances on the current worker, some UDFActor instances will remain idle because they do not receive any data, resulting in UDFActor idling.RaySwordfishActor, which cannot fully utilize all UDFActor instances on the worker. When batch_size is set to a smaller value, it increases the overhead of function calls. Therefore, it is necessary to adjust this value based on the actual business scenario and optimization objectives.In Scenario 1, we encountered a case of uneven data processing. When we used repartition to repartition a dataset containing 1,000 records into two partitions, the resulting partitions were not perfectly balanced. This is because we used the hash value of the path as the repartition key, so the degree of skew depends on the distribution of the hash values. In the above example, after repartitioning, we obtained two partitions containing 510 and 490 records, respectively. Each partition was then split by the Daft engine with batch size = 200 into [200, 200, 110, 200, 200, 90] six batches, which were then scheduled across four concurrent tasks. As can be seen, the splitting of each partition is independent, and data from different partitions cannot be merged. Therefore, the last batch of each partition contains far fewer records than the other batches.
In Scenario 2, we encountered a situation where actual concurrency was lower than expected. This is because the original DataFrame had only one partition, resulting in only one task, which was scheduled to a single worker. Meanwhile, the concurrent UDF instances were already "evenly" distributed across all workers during the plan phase. When RaySwordfishActor schedules batches/morsels, it prioritizes scheduling to the local worker, causing only the UDFActor instance on one worker to process data. Therefore, only when the number of tasks is greater than or equal to the number of workers can we maximize concurrency. So, is there a way to both increase the number of tasks/partitions and balance the data volume across different tasks/partitions? When there are similar requirements, you can consider using the into_batches API. This API splits the data in a partition into different tasks according to the execution size, and then distributes them to different RaySwordfishActor instances. For example, when executing df.into_batches(200).with_column("batch", udf(col("path"))) in a single-worker Ray cluster, you will get the following result:
(UDFActor pid=29906) 2025-09-28 19:15:04.1759058104 - MyBatchUdf-4816300688 - INFO - Actor got 200 records, total: 200 (UDFActor pid=29907) 2025-09-28 19:15:04.1759058104 - MyBatchUdf-4952482832 - INFO - Actor got 200 records, total: 200 (UDFActor pid=29909) 2025-09-28 19:15:04.1759058104 - MyBatchUdf-4548368464 - INFO - Actor got 200 records, total: 200 (UDFActor pid=29906) 2025-09-28 19:15:06.1759058106 - MyBatchUdf-4816300688 - INFO - Actor got 200 records, total: 400 (UDFActor pid=29906) 2025-09-28 19:15:08.1759058108 - MyBatchUdf-4816300688 - INFO - Actor got 200 records, total: 600
From the above log, we can see that the data is evenly split into five batches. However, this does not match the expected even scheduling across four concurrent instances; only three concurrent instances received data, and one instance processed three batches (the result is random, so it is possible that two instances each processed two batches), resulting in an imbalance in data processing among concurrent instances. The reason for this phenomenon is consistent with Scenario 3. Although the data for the five tasks is balanced, after each task is split by RaySwordfishActor, only one morsel is generated, which is then randomly assigned to a UDFActor. Because data processing between tasks is independent, in extreme cases, all five morsels generated by these tasks could be scheduled to the same UDFActor. Therefore, when we expect maximum concurrency, it is important to ensure that after splitting the data of a task, it can be evenly distributed to each UDFActor on the current worker. For example, when we set udf_batch_size to 50, each task's data will be split into four morsels, which are then evenly distributed to each UDFActor using a round-robin strategy.
(UDFActor pid=43813) 2025-09-28 19:27:04.1759058824 - MyBatchUdf-5061156240 - INFO - Actor got 50 records, total: 50 (UDFActor pid=43815) 2025-09-28 19:27:04.1759058824 - MyBatchUdf-5090240720 - INFO - Actor got 50 records, total: 50 (UDFActor pid=43814) 2025-09-28 19:27:04.1759058824 - MyBatchUdf-5906907600 - INFO - Actor got 50 records, total: 50 (UDFActor pid=43816) 2025-09-28 19:27:04.1759058824 - MyBatchUdf-4546699600 - INFO - Actor got 50 records, total: 50 (UDFActor pid=43813) 2025-09-28 19:27:06.1759058826 - MyBatchUdf-5061156240 - INFO - Actor got 50 records, total: 100 (UDFActor pid=43815) 2025-09-28 19:27:06.1759058826 - MyBatchUdf-5090240720 - INFO - Actor got 50 records, total: 100 (UDFActor pid=43814) 2025-09-28 19:27:06.1759058826 - MyBatchUdf-5906907600 - INFO - Actor got 50 records, total: 100 (UDFActor pid=43816) 2025-09-28 19:27:06.1759058826 - MyBatchUdf-4546699600 - INFO - Actor got 50 records, total: 100 (UDFActor pid=43813) 2025-09-28 19:27:08.1759058828 - MyBatchUdf-5061156240 - INFO - Actor got 50 records, total: 150 (UDFActor pid=43815) 2025-09-28 19:27:08.1759058828 - MyBatchUdf-5090240720 - INFO - Actor got 50 records, total: 150 (UDFActor pid=43814) 2025-09-28 19:27:08.1759058828 - MyBatchUdf-5906907600 - INFO - Actor got 50 records, total: 150 (UDFActor pid=43816) 2025-09-28 19:27:08.1759058828 - MyBatchUdf-4546699600 - INFO - Actor got 50 records, total: 150 (UDFActor pid=43813) 2025-09-28 19:27:10.1759058830 - MyBatchUdf-5061156240 - INFO - Actor got 50 records, total: 200 (UDFActor pid=43815) 2025-09-28 19:27:10.1759058830 - MyBatchUdf-5090240720 - INFO - Actor got 50 records, total: 200 (UDFActor pid=43814) 2025-09-28 19:27:10.1759058830 - MyBatchUdf-5906907600 - INFO - Actor got 50 records, total: 200 (UDFActor pid=43816) 2025-09-28 19:27:10.1759058830 - MyBatchUdf-4546699600 - INFO - Actor got 50 records, total: 200 (UDFActor pid=43813) 2025-09-28 19:27:12.1759058832 - MyBatchUdf-5061156240 - INFO - Actor got 50 records, total: 250 (UDFActor pid=43815) 2025-09-28 19:27:12.1759058832 - MyBatchUdf-5090240720 - INFO - Actor got 50 records, total: 250 (UDFActor pid=43814) 2025-09-28 19:27:12.1759058832 - MyBatchUdf-5906907600 - INFO - Actor got 50 records, total: 250 (UDFActor pid=43816) 2025-09-28 19:27:12.1759058832 - MyBatchUdf-4546699600 - INFO - Actor got 50 records, total: 250
It appears that the combination of into_batches and UDF batch_size not only ensures balanced data distribution for each task, but also fully utilizes the concurrency of each UDFActor. However, into_batches is not a universal solution. For example, in scenarios with a single-machine, single-partition DataFrame, it is not suitable to use into_batches. Since there is only one partition and one worker, there is neither data imbalance nor a need to distribute data across multiple workers. Therefore, there is no need to increase the number of tasks using into_batches, which would only add overhead from function calls.
Additionally, when executing into_batches, even if the DataFrame or partition has sufficient data, the tasks generated may not all meet the corresponding batch size. For example, if a DataFrame has two partitions, with the first partition containing 180 records and the second partition containing 220 records, when splitting with batch_size=200, three tasks will be generated, containing 180, 200, and 20 records respectively.into_batches not only splits a large partition into multiple tasks, but also has the capability to merge several small partitions into larger tasks. For example, suppose a DataFrame has 16 partitions, each containing only 20 records. When df.into_batches(200) is executed, two tasks will be generated, each containing 160 records, rather than one task with 200 records and another with 120 records. The reason behind this is that the distributed plan runner splits the logic before and after into_batches into two stages. It first materializes the result of the first stage, then batches the results. When the accumulated number of records reaches or exceeds 80% of the batch size, a new task is created. Therefore, in the example above, two tasks are generated, each containing 160 records.
In different application scenarios, the desired outcome can be achieved by adjusting the concurrency of the UDF and the batch size parameter.
For example, to maximize UDF concurrency, suppose there are M Workers,
repartition or into_partition, or the number of generated tasks can be increased using into_batches.repartition/into_partition or into_batches to balance the data across tasks. The former is suitable for scenarios with a large number of partitions (if the number of partitions is too small, there can be significant differences in data volume between partitions. The general logic of repartition is: each upstream partition splits its data into N chunks, and each downstream partition then merges 1/N of the chunks from upstream partitions). When N is small, each chunk contains a large amount of data, so even if partitions differ by only one chunk, there will still be chunk-level differences in data volume. The latter is suitable for scenarios with small batches.The above discussion only covers how to maximize concurrent computation by adjusting the number of DataFrame partitions, UDF concurrency, and batch size. In practice, to run UDFs efficiently, you also need to pay attention to the load during UDF execution, such as CPU and memory usage. If the UDF consumes a large amount of memory while processing data, you should reduce the UDF batch size to avoid out-of-memory (OOM) issues. Reducing the batch size may increase the number of morsels, and when the UDFActor processes these morsels, backpressure can occur. Therefore, to maximize concurrency, you should increase UDF concurrency. When UDF processing involves a large amount of IO operations, consider using asynchronous operations and coroutines within the UDF to improve CPU utilization, and increase the batch size to reduce the number of calls to the UDF. Increasing the batch size may reduce the number of morsels. If this number falls below the number of UDF instances on a worker, some instances may not receive data and will remain idle.
For scenarios such as merging small files or fragmented upstream partitions, you can use into_batches to merge partitions, thereby reducing the number of calls to UDF instances.