Skip to main content

polars_stream/physical_plan/
mod.rs

1use std::num::NonZeroUsize;
2use std::sync::Arc;
3
4use polars_core::frame::DataFrame;
5#[cfg(any(
6    feature = "dtype-date",
7    feature = "dtype-datetime",
8    feature = "dtype-time"
9))]
10use polars_core::prelude::DataType;
11use polars_core::prelude::{IdxSize, InitHashMaps, PlHashMap, PlIndexMap, SortMultipleOptions};
12use polars_core::schema::{Schema, SchemaRef};
13use polars_error::PolarsResult;
14use polars_io::RowIndex;
15use polars_io::cloud::CloudOptions;
16use polars_ops::frame::JoinArgs;
17#[cfg(any(
18    feature = "dtype-date",
19    feature = "dtype-datetime",
20    feature = "dtype-time"
21))]
22use polars_plan::dsl::StrptimeOptions;
23use polars_plan::dsl::deletion::DeletionFilesList;
24use polars_plan::dsl::{
25    CastColumnsPolicy, ColumnsUdf, FileSinkOptions, JoinTypeOptionsIR, MissingColumnsPolicy,
26    PartitionedSinkOptionsIR, PredicateFileSkip, ScanSources, TableStatistics,
27};
28use polars_plan::plans::expr_ir::ExprIR;
29use polars_plan::plans::hive::HivePartitionsDf;
30use polars_plan::plans::{AExpr, DataFrameUdf, DynamicPred, FunctionArgMap, IR};
31
32mod fmt;
33mod io;
34mod lower_expr;
35mod lower_group_by;
36mod lower_ir;
37mod to_graph;
38
39pub use fmt::{NodeStyle, visualize_plan};
40use polars_plan::prelude::PlanCallback;
41#[cfg(feature = "dynamic_group_by")]
42use polars_time::DynamicGroupOptions;
43use polars_time::{ClosedWindow, Duration};
44use polars_utils::arena::{Arena, Node};
45use polars_utils::pl_str::PlSmallStr;
46use polars_utils::slice_enum::Slice;
47use polars_utils::{UnitVec, unitvec};
48use slotmap::{SecondaryMap, SlotMap};
49pub use to_graph::physical_plan_to_graph;
50
51pub use self::lower_ir::StreamingLowerIRContext;
52use crate::nodes::io_sources::multi_scan::components::forbid_extra_columns::ForbidExtraColumns;
53use crate::nodes::io_sources::multi_scan::components::projection::builder::ProjectionBuilder;
54use crate::nodes::io_sources::multi_scan::reader_interface::builder::FileReaderBuilder;
55use crate::physical_plan::lower_expr::ExprCache;
56
57slotmap::new_key_type! {
58    /// Key used for physical nodes.
59    pub struct PhysNodeKey;
60}
61
62impl PhysNodeKey {
63    pub fn as_ffi(&self) -> u64 {
64        self.0.as_ffi()
65    }
66}
67
68/// A node in the physical plan.
69///
70/// A physical plan is created when the `IR` is translated to a directed
71/// acyclic graph of operations that can run on the streaming engine.
72#[derive(Clone, Debug)]
73pub struct PhysNode {
74    output_schemas: UnitVec<Arc<Schema>>,
75    kind: PhysNodeKind,
76}
77
78impl PhysNode {
79    pub fn new(output_schema: Arc<Schema>, kind: PhysNodeKind) -> Self {
80        Self {
81            output_schemas: unitvec![output_schema],
82            kind,
83        }
84    }
85
86    pub fn new_multi_output(output_schemas: UnitVec<Arc<Schema>>, kind: PhysNodeKind) -> Self {
87        Self {
88            output_schemas,
89            kind,
90        }
91    }
92
93    pub fn output_schema(&self, port_idx: usize) -> &Arc<Schema> {
94        &self.output_schemas[port_idx]
95    }
96
97    pub fn output_schema_mut(&mut self, port_idx: usize) -> &mut Arc<Schema> {
98        &mut self.output_schemas[port_idx]
99    }
100
101    pub fn kind(&self) -> &PhysNodeKind {
102        &self.kind
103    }
104}
105
106/// A handle representing a physical stream of data with a fixed schema in the
107/// physical plan. It consists of a reference to a physical node as well as the
108/// output port on that node to connect to receive the stream.
109#[derive(Clone, Debug, Copy, PartialEq, Eq, Hash)]
110pub struct PhysStream {
111    pub node: PhysNodeKey,
112    pub port: usize,
113}
114
115impl PhysStream {
116    #[allow(unused)]
117    pub fn new(node: PhysNodeKey, port: usize) -> Self {
118        Self { node, port }
119    }
120
121    // Convenience method to refer to the first output port of a physical node.
122    pub fn first(node: PhysNodeKey) -> Self {
123        Self { node, port: 0 }
124    }
125
126    pub fn output_schema<'sm>(&self, sm: &'sm SlotMap<PhysNodeKey, PhysNode>) -> &'sm Arc<Schema> {
127        sm[self.node].output_schema(self.port)
128    }
129
130    pub fn output_schema_mut<'sm>(
131        &self,
132        sm: &'sm mut SlotMap<PhysNodeKey, PhysNode>,
133    ) -> &'sm mut Arc<Schema> {
134        sm[self.node].output_schema_mut(self.port)
135    }
136}
137
138/// Behaviour when handling multiple DataFrames with different heights.
139
140#[derive(Clone, Debug, Copy)]
141#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
142#[cfg_attr(
143    feature = "physical_plan_visualization_schema",
144    derive(schemars::JsonSchema)
145)]
146pub enum ZipBehavior {
147    /// Fill the shorter DataFrames with nulls to the height of the longest DataFrame.
148    NullExtend,
149    /// All inputs must be the same height, or have length 1 in which case they are broadcast.
150    Broadcast,
151    /// Raise an error if the DataFrames have different heights.
152    Strict,
153}
154
155#[derive(Clone, Debug)]
156pub enum PhysNodeKind {
157    InMemorySource {
158        df: Arc<DataFrame>,
159        disable_morsel_split: bool,
160    },
161
162    Select {
163        input: PhysStream,
164        selectors: Vec<ExprIR>,
165        extend_original: bool,
166    },
167
168    InputIndependentSelect {
169        selectors: Vec<ExprIR>,
170    },
171
172    WithRowIndex {
173        input: PhysStream,
174        name: PlSmallStr,
175        offset: Option<IdxSize>,
176    },
177
178    Reduce {
179        input: PhysStream,
180        exprs: Vec<ExprIR>,
181    },
182
183    StreamingSlice {
184        input: PhysStream,
185        offset: usize,
186        length: usize,
187    },
188
189    NegativeSlice {
190        input: PhysStream,
191        offset: i64,
192        length: usize,
193    },
194
195    DynamicSlice {
196        input: PhysStream,
197        offset: PhysStream,
198        length: PhysStream,
199    },
200
201    Shift {
202        input: PhysStream,
203        offset: PhysStream,
204        fill: Option<PhysStream>,
205    },
206
207    Filter {
208        input: PhysStream,
209        predicate: ExprIR,
210    },
211
212    SimpleProjection {
213        input: PhysStream,
214        // Key: output column, value: input column.
215        columns: PlIndexMap<PlSmallStr, PlSmallStr>,
216    },
217
218    InMemorySink {
219        input: PhysStream,
220    },
221
222    CallbackSink {
223        input: PhysStream,
224        function: PlanCallback<DataFrame, bool>,
225        maintain_order: bool,
226        chunk_size: Option<NonZeroUsize>,
227    },
228
229    FileSink {
230        input: PhysStream,
231        options: FileSinkOptions,
232    },
233
234    PartitionedSink {
235        input: PhysStream,
236        options: PartitionedSinkOptionsIR,
237    },
238
239    SinkMultiple {
240        sinks: Vec<PhysNodeKey>,
241    },
242
243    /// Generic fallback for (as-of-yet) unsupported streaming mappings.
244    /// Fully sinks all data to an in-memory data frame and uses the in-memory
245    /// engine to perform the map.
246    InMemoryMap {
247        input: PhysStream,
248        map: Arc<dyn DataFrameUdf>,
249        /// A formatted string of what the in-memory map is. This usually calls format on the IR.
250        format_str: Option<String>,
251    },
252
253    Map {
254        input: PhysStream,
255        map: Arc<dyn DataFrameUdf>,
256        /// A formatted string of what the map is. This usually calls format on the IR.
257        format_str: Option<String>,
258    },
259
260    ColumnarFunction {
261        inputs: Vec<PhysStream>,
262        func: Arc<dyn ColumnsUdf>,
263        arg_map: Option<FunctionArgMap>,
264        output_name: PlSmallStr,
265        format_str: Option<String>,
266    },
267
268    /// Streaming strptime without an explicit format.
269    #[cfg(any(
270        feature = "dtype-date",
271        feature = "dtype-datetime",
272        feature = "dtype-time"
273    ))]
274    StrptimeInfer {
275        input: PhysStream,
276        dtype: DataType,
277        options: StrptimeOptions,
278
279        /// Ambiguous can be `raise`, `earliest`, `latest` and `null`.
280        ///
281        /// If it is broadcast and it is `raise` or `null`, we can actually execute it in this
282        /// node. So
283        /// - `false` -> "null"
284        /// - `true`  -> "raise"
285        ambiguous_is_raise: bool,
286    },
287
288    SortedGroupBy {
289        input: PhysStream,
290        key: PlSmallStr,
291        aggs: Vec<ExprIR>,
292        slice: Option<(IdxSize, IdxSize)>,
293    },
294
295    Sort {
296        input: PhysStream,
297        by_column: Vec<ExprIR>,
298        slice: Option<(i64, usize)>,
299        sort_options: SortMultipleOptions,
300    },
301
302    TopK {
303        input: PhysStream,
304        k: PhysStream,
305        by_column: Vec<ExprIR>,
306        reverse: Vec<bool>,
307        nulls_last: Vec<bool>,
308        dyn_pred: Option<DynamicPred>,
309    },
310
311    Repeat {
312        value: PhysStream,
313        repeats: PhysStream,
314    },
315
316    #[cfg(feature = "cum_agg")]
317    CumAgg {
318        input: PhysStream,
319        kind: crate::nodes::cum_agg::CumAggKind,
320    },
321
322    // Parameter is the input stream
323    GatherEvery {
324        input: PhysStream,
325        n: usize,
326        offset: usize,
327    },
328    ForwardFill {
329        input: PhysStream,
330        limit: Option<IdxSize>,
331    },
332    BackwardFill {
333        input: PhysStream,
334        limit: Option<IdxSize>,
335    },
336    #[cfg(feature = "interpolate")]
337    Interpolate {
338        input: PhysStream,
339        method: polars_ops::series::InterpolationMethod,
340    },
341    Rle(PhysStream),
342    RleId(PhysStream),
343    SortedUnique {
344        input: PhysStream,
345        keys: Vec<PlSmallStr>,
346    },
347    PeakMinMax {
348        input: PhysStream,
349        is_peak_max: bool,
350    },
351    IsSorted {
352        input: PhysStream,
353        descending: Option<bool>,
354        nulls_last: Option<bool>,
355        output_name: PlSmallStr,
356    },
357
358    OrderedUnion {
359        inputs: Vec<PhysStream>,
360    },
361
362    UnorderedUnion {
363        inputs: Vec<PhysStream>,
364    },
365
366    Zip {
367        inputs: Vec<PhysStream>,
368        zip_behavior: ZipBehavior,
369    },
370
371    #[allow(unused)]
372    Multiplexer {
373        input: PhysStream,
374    },
375
376    MultiScan {
377        scan_sources: ScanSources,
378
379        file_reader_builder: Arc<dyn FileReaderBuilder>,
380        cloud_options: Option<Arc<CloudOptions>>,
381
382        /// Columns to project from the file.
383        file_projection_builder: ProjectionBuilder,
384        /// Final output schema of morsels being sent out of MultiScan.
385        output_schema: SchemaRef,
386
387        row_index: Option<RowIndex>,
388        pre_slice: Option<Slice>,
389        predicate: Option<ExprIR>,
390        predicate_file_skip_applied: Option<PredicateFileSkip>,
391
392        hive_parts: Option<HivePartitionsDf>,
393        include_file_paths: Option<PlSmallStr>,
394        cast_columns_policy: CastColumnsPolicy,
395        missing_columns_policy: MissingColumnsPolicy,
396        forbid_extra_columns: Option<ForbidExtraColumns>,
397
398        deletion_files: Option<DeletionFilesList>,
399        table_statistics: Option<TableStatistics>,
400
401        /// Schema of columns contained in the file. Does not contain external columns (e.g. hive / row_index).
402        file_schema: SchemaRef,
403        disable_morsel_split: bool,
404    },
405
406    #[cfg(feature = "python")]
407    PythonScan {
408        options: polars_plan::plans::python::PythonOptions,
409    },
410
411    GroupBy {
412        inputs: Vec<PhysStream>,
413        // Must have the same schema when applied for each input.
414        key_per_input: Vec<Vec<ExprIR>>,
415        // Must be a 'simple' expression, a singular column feeding into a single aggregate, or Len.
416        aggs_per_input: Vec<Vec<ExprIR>>,
417    },
418
419    #[cfg(feature = "dynamic_group_by")]
420    DynamicGroupBy {
421        input: PhysStream,
422        options: DynamicGroupOptions,
423        aggs: Vec<ExprIR>,
424        slice: Option<(IdxSize, IdxSize)>,
425    },
426
427    #[cfg(feature = "dynamic_group_by")]
428    RollingGroupBy {
429        input: PhysStream,
430        index_column: PlSmallStr,
431        period: Duration,
432        offset: Duration,
433        closed: ClosedWindow,
434        slice: Option<(IdxSize, IdxSize)>,
435        aggs: Vec<ExprIR>,
436    },
437
438    #[cfg(feature = "is_first_distinct")]
439    IsFirstDistinct {
440        input: PhysStream,
441        out_name: PlSmallStr,
442        columns: Vec<PlSmallStr>,
443    },
444
445    EquiJoin {
446        input_left: PhysStream,
447        input_right: PhysStream,
448        left_on: Vec<ExprIR>,
449        right_on: Vec<ExprIR>,
450        args: JoinArgs,
451    },
452
453    MergeJoin {
454        input_left: PhysStream,
455        input_right: PhysStream,
456        left_on: Vec<PlSmallStr>,
457        right_on: Vec<PlSmallStr>,
458        tmp_left_key_col: Option<PlSmallStr>,
459        tmp_right_key_col: Option<PlSmallStr>,
460        descending: bool,
461        nulls_last: bool,
462        keys_row_encoded: bool,
463        args: JoinArgs,
464    },
465
466    SemiAntiJoin {
467        input_left: PhysStream,
468        input_right: PhysStream,
469        left_on: Vec<ExprIR>,
470        right_on: Vec<ExprIR>,
471        args: JoinArgs,
472        output_bool: bool,
473    },
474
475    CrossJoin {
476        input_left: PhysStream,
477        input_right: PhysStream,
478        args: JoinArgs,
479    },
480
481    AsOfJoin {
482        input_left: PhysStream,
483        input_right: PhysStream,
484        left_on: PlSmallStr,
485        right_on: PlSmallStr,
486        tmp_left_key_col: Option<PlSmallStr>,
487        tmp_right_key_col: Option<PlSmallStr>,
488        by_descending: Option<Vec<bool>>,
489        by_nulls_last: Option<Vec<bool>>,
490        args: JoinArgs,
491    },
492
493    #[cfg(feature = "iejoin")]
494    RangeJoin {
495        input_left: PhysStream,
496        input_right: PhysStream,
497        left_on: Vec<PlSmallStr>,
498        right_on: Vec<PlSmallStr>,
499        tmp_left_key_cols: Vec<Option<PlSmallStr>>,
500        tmp_right_key_cols: Vec<Option<PlSmallStr>>,
501        descending: bool,
502        args: JoinArgs,
503        options: polars_ops::frame::IEJoinOptions,
504    },
505
506    /// Generic fallback for (as-of-yet) unsupported streaming joins.
507    /// Fully sinks all data to in-memory data frames and uses the in-memory
508    /// engine to perform the join.
509    InMemoryJoin {
510        input_left: PhysStream,
511        input_right: PhysStream,
512        left_on: Vec<ExprIR>,
513        right_on: Vec<ExprIR>,
514        args: JoinArgs,
515        options: Option<JoinTypeOptionsIR>,
516    },
517
518    #[cfg(feature = "merge_sorted")]
519    MergeSorted {
520        input_left: PhysStream,
521        input_right: PhysStream,
522        maintain_order: bool,
523    },
524
525    Gather {
526        input: PhysStream,
527        idxs: PhysStream,
528        null_on_oob: bool,
529    },
530
531    #[cfg(feature = "ewma")]
532    EwmMean {
533        input: PhysStream,
534        options: polars_ops::series::EWMOptions,
535    },
536
537    #[cfg(feature = "ewma")]
538    EwmSum {
539        input: PhysStream,
540        options: polars_ops::series::EWMOptions,
541    },
542
543    #[cfg(feature = "ewma")]
544    EwmVar {
545        input: PhysStream,
546        options: polars_ops::series::EWMOptions,
547    },
548
549    #[cfg(feature = "ewma")]
550    EwmStd {
551        input: PhysStream,
552        options: polars_ops::series::EWMOptions,
553    },
554}
555
556fn visit_node_inputs_mut(
557    roots: Vec<PhysNodeKey>,
558    phys_sm: &mut SlotMap<PhysNodeKey, PhysNode>,
559    mut visit: impl FnMut(&mut PhysStream),
560) {
561    let mut to_visit = roots;
562    let mut seen: SecondaryMap<PhysNodeKey, ()> =
563        to_visit.iter().copied().map(|n| (n, ())).collect();
564    macro_rules! rec {
565        ($n:expr) => {
566            let n = $n;
567            if seen.insert(n, ()).is_none() {
568                to_visit.push(n)
569            }
570        };
571    }
572    while let Some(node) = to_visit.pop() {
573        match &mut phys_sm[node].kind {
574            PhysNodeKind::InMemorySource { .. }
575            | PhysNodeKind::MultiScan { .. }
576            | PhysNodeKind::InputIndependentSelect { .. } => {},
577            #[cfg(feature = "python")]
578            PhysNodeKind::PythonScan { .. } => {},
579            PhysNodeKind::Select { input, .. }
580            | PhysNodeKind::WithRowIndex { input, .. }
581            | PhysNodeKind::Reduce { input, .. }
582            | PhysNodeKind::StreamingSlice { input, .. }
583            | PhysNodeKind::NegativeSlice { input, .. }
584            | PhysNodeKind::Filter { input, .. }
585            | PhysNodeKind::SimpleProjection { input, .. }
586            | PhysNodeKind::InMemorySink { input }
587            | PhysNodeKind::CallbackSink { input, .. }
588            | PhysNodeKind::FileSink { input, .. }
589            | PhysNodeKind::PartitionedSink { input, .. }
590            | PhysNodeKind::InMemoryMap { input, .. }
591            | PhysNodeKind::SortedGroupBy { input, .. }
592            | PhysNodeKind::Map { input, .. }
593            | PhysNodeKind::Sort { input, .. }
594            | PhysNodeKind::Multiplexer { input }
595            | PhysNodeKind::GatherEvery { input, .. }
596            | PhysNodeKind::ForwardFill { input, .. }
597            | PhysNodeKind::BackwardFill { input, .. }
598            | PhysNodeKind::Rle(input)
599            | PhysNodeKind::RleId(input)
600            | PhysNodeKind::SortedUnique { input, .. }
601            | PhysNodeKind::PeakMinMax { input, .. }
602            | PhysNodeKind::IsSorted { input, .. } => {
603                rec!(input.node);
604                visit(input);
605            },
606
607            #[cfg(feature = "interpolate")]
608            PhysNodeKind::Interpolate { input, .. } => {
609                rec!(input.node);
610                visit(input);
611            },
612
613            #[cfg(feature = "is_first_distinct")]
614            PhysNodeKind::IsFirstDistinct { input, .. } => {
615                rec!(input.node);
616                visit(input);
617            },
618
619            #[cfg(any(
620                feature = "dtype-date",
621                feature = "dtype-datetime",
622                feature = "dtype-time"
623            ))]
624            PhysNodeKind::StrptimeInfer { input, .. } => {
625                rec!(input.node);
626                visit(input);
627            },
628
629            #[cfg(feature = "dynamic_group_by")]
630            PhysNodeKind::DynamicGroupBy { input, .. } => {
631                rec!(input.node);
632                visit(input);
633            },
634            #[cfg(feature = "dynamic_group_by")]
635            PhysNodeKind::RollingGroupBy { input, .. } => {
636                rec!(input.node);
637                visit(input);
638            },
639
640            #[cfg(feature = "cum_agg")]
641            PhysNodeKind::CumAgg { input, .. } => {
642                rec!(input.node);
643                visit(input);
644            },
645
646            PhysNodeKind::InMemoryJoin {
647                input_left,
648                input_right,
649                ..
650            }
651            | PhysNodeKind::EquiJoin {
652                input_left,
653                input_right,
654                ..
655            }
656            | PhysNodeKind::MergeJoin {
657                input_left,
658                input_right,
659                ..
660            }
661            | PhysNodeKind::SemiAntiJoin {
662                input_left,
663                input_right,
664                ..
665            }
666            | PhysNodeKind::CrossJoin {
667                input_left,
668                input_right,
669                ..
670            }
671            | PhysNodeKind::AsOfJoin {
672                input_left,
673                input_right,
674                ..
675            } => {
676                rec!(input_left.node);
677                rec!(input_right.node);
678                visit(input_left);
679                visit(input_right);
680            },
681
682            #[cfg(feature = "iejoin")]
683            PhysNodeKind::RangeJoin {
684                input_left,
685                input_right,
686                ..
687            } => {
688                rec!(input_left.node);
689                rec!(input_right.node);
690                visit(input_left);
691                visit(input_right);
692            },
693
694            #[cfg(feature = "merge_sorted")]
695            PhysNodeKind::MergeSorted {
696                input_left,
697                input_right,
698                ..
699            } => {
700                rec!(input_left.node);
701                rec!(input_right.node);
702                visit(input_left);
703                visit(input_right);
704            },
705
706            PhysNodeKind::Gather { input, idxs, .. } => {
707                rec!(input.node);
708                rec!(idxs.node);
709                visit(input);
710                visit(idxs);
711            },
712
713            PhysNodeKind::TopK { input, k, .. } => {
714                rec!(input.node);
715                rec!(k.node);
716                visit(input);
717                visit(k);
718            },
719
720            PhysNodeKind::DynamicSlice {
721                input,
722                offset,
723                length,
724            } => {
725                rec!(input.node);
726                rec!(offset.node);
727                rec!(length.node);
728                visit(input);
729                visit(offset);
730                visit(length);
731            },
732
733            PhysNodeKind::Shift {
734                input,
735                offset,
736                fill,
737            } => {
738                rec!(input.node);
739                rec!(offset.node);
740                if let Some(fill) = fill {
741                    rec!(fill.node);
742                }
743                visit(input);
744                visit(offset);
745                if let Some(fill) = fill {
746                    visit(fill);
747                }
748            },
749
750            PhysNodeKind::Repeat { value, repeats } => {
751                rec!(value.node);
752                rec!(repeats.node);
753                visit(value);
754                visit(repeats);
755            },
756
757            PhysNodeKind::GroupBy { inputs, .. }
758            | PhysNodeKind::OrderedUnion { inputs }
759            | PhysNodeKind::UnorderedUnion { inputs }
760            | PhysNodeKind::Zip { inputs, .. }
761            | PhysNodeKind::ColumnarFunction { inputs, .. } => {
762                for input in inputs {
763                    rec!(input.node);
764                    visit(input);
765                }
766            },
767
768            PhysNodeKind::SinkMultiple { sinks } => {
769                for sink in sinks {
770                    rec!(*sink);
771                    visit(&mut PhysStream::first(*sink));
772                }
773            },
774
775            #[cfg(feature = "ewma")]
776            PhysNodeKind::EwmMean { input, options: _ }
777            | PhysNodeKind::EwmSum { input, options: _ }
778            | PhysNodeKind::EwmVar { input, options: _ }
779            | PhysNodeKind::EwmStd { input, options: _ } => {
780                rec!(input.node);
781                visit(input)
782            },
783        }
784    }
785}
786
787fn insert_multiplexers(roots: Vec<PhysNodeKey>, phys_sm: &mut SlotMap<PhysNodeKey, PhysNode>) {
788    let mut refcount: PlIndexMap<_, usize> = PlIndexMap::new();
789    visit_node_inputs_mut(roots.clone(), phys_sm, |i| {
790        *refcount.entry(*i).or_insert(0) += 1;
791    });
792
793    let mut multiplexer_map: PlHashMap<PhysStream, PhysStream> = refcount
794        .into_iter()
795        .filter(|(_stream, refcount)| *refcount > 1)
796        .map(|(stream, refcount)| {
797            let input_schema = Arc::clone(stream.output_schema(phys_sm));
798            let multiplexer_node = phys_sm.insert(PhysNode::new_multi_output(
799                (0..refcount).map(|_| Arc::clone(&input_schema)).collect(),
800                PhysNodeKind::Multiplexer { input: stream },
801            ));
802            (stream, PhysStream::first(multiplexer_node))
803        })
804        .collect();
805
806    visit_node_inputs_mut(roots, phys_sm, |i| {
807        if let Some(m) = multiplexer_map.get_mut(i) {
808            *i = *m;
809            m.port += 1;
810        }
811    });
812}
813
814fn split_multiplexers(roots: Vec<PhysNodeKey>, phys_sm: &mut SlotMap<PhysNodeKey, PhysNode>) {
815    let mut refcount: SecondaryMap<PhysNodeKey, usize> = SecondaryMap::new();
816    visit_node_inputs_mut(roots.clone(), phys_sm, |i| {
817        *refcount.entry(i.node).unwrap().or_insert(0) += 1;
818    });
819
820    let mut split_map: SecondaryMap<PhysNodeKey, PhysNode> = SecondaryMap::new();
821    for (k, n) in phys_sm.iter() {
822        if let PhysNodeKind::Multiplexer { input } = n.kind {
823            if let PhysNodeKind::InMemorySource { .. } = phys_sm[input.node].kind {
824                split_map.insert(k, phys_sm[input.node].clone());
825            }
826        }
827    }
828
829    let mut replacements: SecondaryMap<PhysNodeKey, Vec<PhysStream>> = split_map
830        .into_iter()
831        .map(|(k, n)| {
832            let repls = (0..refcount[k]).map(|_| PhysStream::first(phys_sm.insert(n.clone())));
833            (k, repls.collect())
834        })
835        .collect();
836
837    visit_node_inputs_mut(roots, phys_sm, |i| {
838        if let Some(r) = replacements.get_mut(i.node) {
839            *i = r.pop().unwrap();
840        }
841    });
842}
843
844pub fn build_physical_plan(
845    root: Node,
846    ir_arena: &mut Arena<IR>,
847    expr_arena: &mut Arena<AExpr>,
848    phys_sm: &mut SlotMap<PhysNodeKey, PhysNode>,
849    ctx: StreamingLowerIRContext<'_>,
850) -> PolarsResult<PhysNodeKey> {
851    let mut schema_cache = PlHashMap::with_capacity(ir_arena.len());
852    let mut expr_cache = ExprCache::with_capacity(expr_arena.len());
853    let mut cache_nodes = PlHashMap::new();
854    let phys_root = lower_ir::lower_ir(
855        root,
856        ir_arena,
857        expr_arena,
858        phys_sm,
859        &mut schema_cache,
860        &mut expr_cache,
861        &mut cache_nodes,
862        ctx,
863        None,
864    )?;
865    insert_multiplexers(vec![phys_root.node], phys_sm);
866    split_multiplexers(vec![phys_root.node], phys_sm);
867    Ok(phys_root.node)
868}