DataFrame.cache() does not work in distributed environments

Is your feature request related to a problem or challenge?

DataFrame.cache() does not work in distributed environments as it will materialise whole table in memory.

pub async fn cache(self) -> Result<DataFrame> {
    let context = SessionContext::new_with_state((*self.session_state).clone());
    // The schema is consistent with the output
    let plan = self.clone().create_physical_plan().await?;
    let schema = plan.schema();
    let task_ctx = Arc::new(self.task_ctx());
    let partitions = collect_partitioned(plan, task_ctx).await?;
    let mem_table = MemTable::try_new(schema, partitions)?;
    context.read_table(Arc::new(mem_table))
}

This does not work in distributed environments such as ballista.

Also note, cache will eagerly materialise, even there is no downstream consumers (which may not be a big problem) but it does not follow semantics of spark.cache(..)

Describe the solution you'd like

Ideally cache should be represented with logical plan node, which would be resolved first time cache is needed. This approach may be a bit complicated to implement, and datafusion may not really benefit from it.
As an easier to implement alternative, we could provide a cache factory (at session_state maybe which would provide same functionality as it is currently if not overridden or provide user specified logic, such as returning LogicalPlan::Extension or similar, and leave query planner/user to deal with cache materialisation decision.
EDIT: Another alternative would be to have a configuration option, like datafusion.execution.local_cache which would decide should it use current cache logic or return new logical plan element like LogicalPlan::Cache { id: string, lineage: LogicalPlan }

pub async fn cache(self) -> Result<DataFrame> {
   if `datafusion.execution.local_cache` == true {
     let context = SessionContext::new_with_state((*self.session_state).clone());
     // The schema is consistent with the output
     let plan = self.clone().create_physical_plan().await?;
     let schema = plan.schema();
     let task_ctx = Arc::new(self.task_ctx());
     let partitions = collect_partitioned(plan, task_ctx).await?;
     let mem_table = MemTable::try_new(schema, partitions)?;
     context.read_table(Arc::new(mem_table))
  } else {
   let lineage = self... // get current logical plan from df up to this point
   Ok(
     LogicalPlan::Cache(Cache { id : uuid(), lineage)  
   )
  }
}

Describe alternatives you've considered

Users can create DataFrameExt which would provide a new method like distributed_cache which would return LogicalPlan::Extension with DistributedCacheExtension. This alternative is the simplest to implement (not affecting datafusion) but user would need to change code when moving from datafusion to ballista

Additional context

No response

Hey. Can i take this one? Please assign

Feel free to take this issue by leaving a comment with just the word take

See reference in contributor guide: https://datafusion.apache.org/contributor-guide/index.html#open-contribution-and-assigning-tickets

Edited issue description with another option:

Another alternative would be to have a configuration option, like datafusion.execution.local_cache which would decide should it use current cache logic or return new logical plan element like LogicalPlan::Cache { id: string, lineage: LogicalPlan }

pub async fn cache(self) -> Result<DataFrame> {
   if `datafusion.execution.local_cache` == true {
     let context = SessionContext::new_with_state((*self.session_state).clone());
     // The schema is consistent with the output
     let plan = self.clone().create_physical_plan().await?;
     let schema = plan.schema();
     let task_ctx = Arc::new(self.task_ctx());
     let partitions = collect_partitioned(plan, task_ctx).await?;
     let mem_table = MemTable::try_new(schema, partitions)?;
     context.read_table(Arc::new(mem_table))
  } else {
   let lineage = self... // get current logical plan from df up to this point
   Ok(
     LogicalPlan::Cache(Cache { id : uuid(), lineage)  
   )
  }
}

take

feat: add configurable cache mode (local_cache) with LogicalPlan::Cac…

@alamb

@alamb if you have time it would be great to get your opinion about this issue as solving it may require some coding.

In terms of designing an API, I think it would help me to understand what behavior you want in Ballista when a user runs DataFrame::cache (aka how wll you use this API)?

For example do you want it to write to fies distributed across the nodes?

It seems to me at the moment DataFrame::cache is basically a shorthand for

"CREATE TABLE <temp> AS SELECT * FROM <dataframe>

The current implementation has a hard coded MemTable and the orchestration (creation of the MemTable / batches) is in the DataFrame::cache function itself.

If you want to move the logic to LogicalPlans you'll probably also have to add an equivalent physical node that handles the orchestration

Ballista will create cache partitions locally on some of the executors. Handling the physical part of execution is specific to the implementor, such as Ballista in this case.

Overall flow would be something like:

Ballista receives LogicalPlan::Cache
LogicalPlan::Cache is then converted to BallistaCacheReadExec , this part is handled with BallistaQueryPlanner (there is a bit more to it, a job can be started which would create cache with BallistaCacheWriteExec if not materialised already, task will be pending until cache job creation is in progress, cache partitions can be re-created in case of node failures ...)

As the physical execution is tied to an external system (cache implementer) I believe we do not need to bring CachePhysicalExec to DataFusion, we just need to provide a LogicalPlan::Cache.

With the proposed solution we would be able to keep the current behaviour, or we can delegate cache handling to the external system if we wish. So if the user disables local cache datafusion.execution.local_cache=false it would be up to them to provide a query planner which would know how to handle LogicalPlan::Cache

Probably it would make more sense to name datafusion.execution.local_cache=false, datafusion.execution.external_cache=true

I believe we do not need to bring CachePhysicalExec to DataFusion, we just need to provide a LogicalPlan::Cache.

That makes sense to me and the idea of being able to specially handle LogicalPlan nodes seems reasonable to me

I dug around and it seems like we already do this for various DDL here:
https://github.com/apache/datafusion/blob/cb69574c0ff326ff975c58234b124904add33cfc/datafusion/core/src/execution/context/mod.rs#L664-L663

What if we instead changed DataFrame::cache to use a LogicalPlan::DDL / CreateMemoryTable https://github.com/apache/datafusion/blob/a571b7d44436c5c2bf641444857e8358edd9d8a8/datafusion/expr/src/logical_plan/ddl.rs#L41-L40

Then ballista could use the same mechanism?

Create memory table may work, input property may be used as lineage, table reference can capture cache id and session I'd.

One thing I need to check if DDL statements will be propagated to custom query planner, if I recall correctly they might be intercepted

looking at spark, cache is represented as InMemoryRelation :

from pyspark.sql import SparkSession
spark : SparkSession = SparkSession.builder.getOrCreate()
# %%
df0 = spark.read.parquet("./tpch-data-sf10/lineitem")
df1 = df0.filter("l_partkey == 1")
df2 = df1.cache()
df2.explain()

will produce:

== Optimized Logical Plan ==
InMemoryRelation [l_orderkey#0L, l_partkey#1L, l_suppkey#2L, l_linenumber#3, l_quantity#4, l_extendedprice#5, l_discount#6, l_tax#7, l_returnflag#8, l_linestatus#9, l_shipdate#10, l_commitdate#11, l_receiptdate#12, l_shipinstruct#13, l_shipmode#14, l_comment#15], StorageLevel(disk, memory, deserialized, 1 replicas)
   +- *(1) Filter (isnotnull(l_partkey#1L) AND (l_partkey#1L = 1))
      +- *(1) ColumnarToRow
         +- FileScan parquet [l_orderkey#0L,l_partkey#1L,l_suppkey#2L,l_linenumber#3,l_quantity#4,l_extendedprice#5,l_discount#6,l_tax#7,l_returnflag#8,l_linestatus#9,l_shipdate#10,l_commitdate#11,l_receiptdate#12,l_shipinstruct#13,l_shipmode#14,l_comment#15] Batched: true, DataFilters: [isnotnull(l_partkey#1L), (l_partkey#1L = 1)], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/Users/marko/TMP/tpch-data-sf10/lineitem], PartitionFilters: [], PushedFilters: [IsNotNull(l_partkey), EqualTo(l_partkey,1)], ReadSchema: struct<l_orderkey:bigint,l_partkey:bigint,l_suppkey:bigint,l_linenumber:int,l_quantity:decimal(15...

Looks like create memory table will be handled by datafusion context

datafusion/datafusion/core/src/execution/context/mod.rs

Line 662 in fd7df66

    
               pub async fn execute_logical_plan(&self, plan: LogicalPlan) -> Result<DataFrame> {

datafusion/datafusion/core/src/execution/context/mod.rs

Lines 807 to 808 in fd7df66

    
           async fn create_memory_table(&self, cmd: CreateMemoryTable) -> Result<DataFrame> { 
        
               let CreateMemoryTable {

So it looks like create memory table will not be propagated to custom query planner.

Looks like create memory table will be handled by datafusion context

datafusion/datafusion/core/src/execution/context/mod.rs

Line 662 in fd7df66
 pub async fn execute_logical_plan(&self, plan: LogicalPlan) -> Result<DataFrame> { 
datafusion/datafusion/core/src/execution/context/mod.rs

Lines 807 to 808 in fd7df66

async fn create_memory_table(&self, cmd: CreateMemoryTable) -> Result {
let CreateMemoryTable {
So it looks like create memory table will not be propagated to custom query planner.

What we do in InfluxDB is basically override sql and specially handle all the DDL commands -- maybe ballista could do the same 🤔

datafusion/datafusion/core/src/execution/context/mod.rs
Lines 807 to 808 in fd7df66
async fn create_memory_table(&self, cmd: CreateMemoryTable) -> Result {
let CreateMemoryTable {
So it looks like create memory table will not be propagated to custom query planner.

What we do in InfluxDB is basically override sql and specially handle all the DDL commands -- maybe ballista could do the same 🤔

I'm trying to keep Ballista’s SessionContext the same as DataFusion’s SessionContext. I believe there are a lot of benefits to sharing the same context.

if we use LogicalPlan::TableScan( ) where we implement custom InMemoryTableSource (implementing TableSource) may make sense. TableSource exposes TableSource::get_logical_plan which could capture lineage. This would be quite aligned with spark, and would not introduce new logical plan node 🤔

@milenkovicm

Hi @milenkovicm , I am interested in this task and would like to understand the proposal better. Is the idea that

DataFrame::cache will create a new logical plan with a TableScan node as the root and current logical plan of the dataframe (via self.plan) as its child?
To capture the lineage, TableScan would be constructed with a custom InMemoryTableSource that overrides the get_logical_plan method from the trait TableSource?
Would DataFrame::cache returns ctx.execute_logical_plan(new_plan)?

Thanks a lot!

@jizezhang

Hi @jizezhang, thanks for sharing interest,

Hi @milenkovicm , I am interested in this task and would like to understand the proposal better. Is the idea that

DataFrame::cache will create a new logical plan with a TableScan node as the root and current logical plan of the dataframe (via self.plan) as its child?

yes, new logical plan node to be created, now its question which one to be creted

To capture the lineage, TableScan would be constructed with a custom InMemoryTableSource that overrides the get_logical_plan method from the trait TableSource?

I have tried that, it did not work as expected, I don't remember reason why

Would DataFrame::cache returns ctx.execute_logical_plan(new_plan)?

Not necessarily to execute logical plan right then, it should just return logical plan node and let user decide when to run the plan .

I believe the simplest approach would be to create a factory method which user can configure

Something like:

pub async fn cache(self) -> Result<DataFrame> {
   if let Some(cache_producer) = self.state().cache_producer() {
    cache_producer(&self)
  } else {
    // this is current behaviour
     let context = SessionContext::new_with_state((*self.session_state).clone());
     // The schema is consistent with the output
     let plan = self.clone().create_physical_plan().await?;
     let schema = plan.schema();
     let task_ctx = Arc::new(self.task_ctx());
     let partitions = collect_partitioned(plan, task_ctx).await?;
     let mem_table = MemTable::try_new(schema, partitions)?;
     context.read_table(Arc::new(mem_table))
  }
}

cache_producer would be a method accepting DataFrame and returning DataFrame, probably a logical plan extension.

Does it make sense? Let me know what you think

@milenkovicm

Hi @milenkovicm thanks for replying. I do have some questions:

Is the approach with cache_producer related to the discussion above on creating a new logical node/plan at all? or in this case it would be up to the user to encapsulate their custom logic in cache_producer?
How would a user pass cache_producer, e.g. via SessionState?

Thanks!

Hi @milenkovicm thanks for replying. I do have some questions:

Is the approach with cache_producer related to the discussion above on creating a new logical node/plan at all? or in this case it would be up to the user to encapsulate their custom logic in cache_producer?

yes, just a logical plan node, no custom logic

How would a user pass cache_producer, e.g. via SessionState?

maybe we could copy approach function factory follows ?

Thanks!

Thank you. Actually I would like to ask another question:

cache_producer would be a method accepting DataFrame and returning DataFrame, probably a logical plan extension.

By "logical plan extension", are you referring to

datafusion/datafusion/expr/src/logical_plan/plan.rs

Lines 3186 to 3189 in 552dbe4

    
           pub struct Extension { 
        
               /// The runtime extension operator 
        
               pub node: Arc<dyn UserDefinedLogicalNode>, 
        
           }

Is the idea that we create some UserDefinedLogicalNode that takes current dataframe's logical plan (self.plan) as inputs, and then return a dataframe with a new LogicalPlan::Extension that wraps the node? or how does the extension relate to cache_producer being a method that accepts and returns a dataframe? Could you elaborate a bit more? Thanks!

yes i believe your understanding is right. user can create their own cache implementation, for doing that they would need their own UserDefinedLogicalNode which will represent their cache implementation. regarding input and outputs ... its up to you to experiment and see which one makes most sense, having input as logical plan and state and dataframe as output may be the best approach

@milenkovicm

@milenkovicm Please review the PR when you get a chance. Thanks.

I have updated the PR to address the comments. Please review when you get a chance. Thank you.

feat: allow custom caching via logical node (#18688)

Navigation Menu

Repository navigation

`DataFrame.cache()` does not work in distributed environments #17297

Is your feature request related to a problem or challenge?

Describe the solution you'd like

Describe alternatives you've considered

Additional context

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	async fn create_memory_table(&self, cmd: CreateMemoryTable) -> Result<DataFrame> {
	let CreateMemoryTable {

	pub struct Extension {
	/// The runtime extension operator
	pub node: Arc<dyn UserDefinedLogicalNode>,
	}

DataFrame.cache() does not work in distributed environments #17297

Description

Is your feature request related to a problem or challenge?

Describe the solution you'd like

Describe alternatives you've considered

Additional context

Activity

Jefffrey commented on Aug 24, 2025

milenkovicm commented on Aug 24, 2025

MrGranday commented on Aug 24, 2025

milenkovicm commented on Aug 26, 2025

alamb commented on Aug 27, 2025

milenkovicm commented on Aug 27, 2025

alamb commented on Sep 4, 2025

milenkovicm commented on Sep 4, 2025

milenkovicm commented on Sep 5, 2025

milenkovicm commented on Sep 5, 2025

alamb commented on Sep 8, 2025

milenkovicm commented on Sep 8, 2025

milenkovicm commented on Sep 17, 2025

jizezhang commented on Nov 9, 2025

milenkovicm commented on Nov 9, 2025

jizezhang commented on Nov 10, 2025

milenkovicm commented on Nov 10, 2025

jizezhang commented on Nov 11, 2025

milenkovicm commented on Nov 11, 2025

jizezhang commented on Nov 14, 2025

jizezhang commented on Nov 16, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

`DataFrame.cache()` does not work in distributed environments #17297