Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/physical-plan/src/common.rs
Line
Count
Source (jump to first uncovered line)
1
// Licensed to the Apache Software Foundation (ASF) under one
2
// or more contributor license agreements.  See the NOTICE file
3
// distributed with this work for additional information
4
// regarding copyright ownership.  The ASF licenses this file
5
// to you under the Apache License, Version 2.0 (the
6
// "License"); you may not use this file except in compliance
7
// with the License.  You may obtain a copy of the License at
8
//
9
//   http://www.apache.org/licenses/LICENSE-2.0
10
//
11
// Unless required by applicable law or agreed to in writing,
12
// software distributed under the License is distributed on an
13
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
// KIND, either express or implied.  See the License for the
15
// specific language governing permissions and limitations
16
// under the License.
17
18
//! Defines common code used in execution plans
19
20
use std::fs;
21
use std::fs::{metadata, File};
22
use std::path::{Path, PathBuf};
23
use std::sync::Arc;
24
25
use super::SendableRecordBatchStream;
26
use crate::stream::RecordBatchReceiverStream;
27
use crate::{ColumnStatistics, Statistics};
28
29
use arrow::datatypes::Schema;
30
use arrow::ipc::writer::{FileWriter, IpcWriteOptions};
31
use arrow::record_batch::RecordBatch;
32
use arrow_array::Array;
33
use datafusion_common::stats::Precision;
34
use datafusion_common::{plan_err, DataFusionError, Result};
35
use datafusion_execution::memory_pool::MemoryReservation;
36
37
use futures::{StreamExt, TryStreamExt};
38
use parking_lot::Mutex;
39
40
/// [`MemoryReservation`] used across query execution streams
41
pub(crate) type SharedMemoryReservation = Arc<Mutex<MemoryReservation>>;
42
43
/// Create a vector of record batches from a stream
44
3.32k
pub async fn collect(stream: SendableRecordBatchStream) -> Result<Vec<RecordBatch>> {
45
3.53k
    stream.try_collect::<Vec<_>>().await
46
3.31k
}
47
48
/// Recursively builds a list of files in a directory with a given extension
49
0
pub fn build_checked_file_list(dir: &str, ext: &str) -> Result<Vec<String>> {
50
0
    let mut filenames: Vec<String> = Vec::new();
51
0
    build_file_list_recurse(dir, &mut filenames, ext)?;
52
0
    if filenames.is_empty() {
53
0
        return plan_err!("No files found at {dir} with file extension {ext}");
54
0
    }
55
0
    Ok(filenames)
56
0
}
57
58
/// Recursively builds a list of files in a directory with a given extension
59
0
pub fn build_file_list(dir: &str, ext: &str) -> Result<Vec<String>> {
60
0
    let mut filenames: Vec<String> = Vec::new();
61
0
    build_file_list_recurse(dir, &mut filenames, ext)?;
62
0
    Ok(filenames)
63
0
}
64
65
/// Recursively build a list of files in a directory with a given extension with an accumulator list
66
0
fn build_file_list_recurse(
67
0
    dir: &str,
68
0
    filenames: &mut Vec<String>,
69
0
    ext: &str,
70
0
) -> Result<()> {
71
0
    let metadata = metadata(dir)?;
72
0
    if metadata.is_file() {
73
0
        if dir.ends_with(ext) {
74
0
            filenames.push(dir.to_string());
75
0
        }
76
    } else {
77
0
        for entry in fs::read_dir(dir)? {
78
0
            let entry = entry?;
79
0
            let path = entry.path();
80
0
            if let Some(path_name) = path.to_str() {
81
0
                if path.is_dir() {
82
0
                    build_file_list_recurse(path_name, filenames, ext)?;
83
0
                } else if path_name.ends_with(ext) {
84
0
                    filenames.push(path_name.to_string());
85
0
                }
86
            } else {
87
0
                return plan_err!("Invalid path");
88
            }
89
        }
90
    }
91
0
    Ok(())
92
0
}
93
94
/// If running in a tokio context spawns the execution of `stream` to a separate task
95
/// allowing it to execute in parallel with an intermediate buffer of size `buffer`
96
40
pub(crate) fn spawn_buffered(
97
40
    mut input: SendableRecordBatchStream,
98
40
    buffer: usize,
99
40
) -> SendableRecordBatchStream {
100
    // Use tokio only if running from a multi-thread tokio context
101
40
    match tokio::runtime::Handle::try_current() {
102
0
        Ok(handle)
103
40
            if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread =>
104
0
        {
105
0
            let mut builder = RecordBatchReceiverStream::builder(input.schema(), buffer);
106
0
107
0
            let sender = builder.tx();
108
0
109
0
            builder.spawn(async move {
110
0
                while let Some(item) = input.next().await {
111
0
                    if sender.send(item).await.is_err() {
112
                        // receiver dropped when query is shutdown early (e.g., limit) or error,
113
                        // no need to return propagate the send error.
114
0
                        return Ok(());
115
0
                    }
116
                }
117
118
0
                Ok(())
119
0
            });
120
0
121
0
            builder.build()
122
        }
123
40
        _ => input,
124
    }
125
40
}
126
127
/// Computes the statistics for an in-memory RecordBatch
128
///
129
/// Only computes statistics that are in arrows metadata (num rows, byte size and nulls)
130
/// and does not apply any kernel on the actual data.
131
30
pub fn compute_record_batch_statistics(
132
30
    batches: &[Vec<RecordBatch>],
133
30
    schema: &Schema,
134
30
    projection: Option<Vec<usize>>,
135
30
) -> Statistics {
136
30
    let nb_rows = batches.iter().flatten().map(RecordBatch::num_rows).sum();
137
138
30
    let projection = match projection {
139
2
        Some(p) => p,
140
28
        None => (0..schema.fields().len()).collect(),
141
    };
142
143
30
    let total_byte_size = batches
144
30
        .iter()
145
30
        .flatten()
146
103
        .map(|b| {
147
103
            projection
148
103
                .iter()
149
104
                .map(|index| b.column(*index).get_array_memory_size())
150
103
                .sum::<usize>()
151
103
        })
152
30
        .sum();
153
30
154
30
    let mut null_counts = vec![0; projection.len()];
155
156
105
    for partition in 
batches.iter()30
{
157
208
        for 
batch103
in partition {
158
104
            for (stat_index, col_index) in 
projection.iter().enumerate()103
{
159
104
                null_counts[stat_index] += batch.column(*col_index).null_count();
160
104
            }
161
        }
162
    }
163
30
    let column_statistics = null_counts
164
30
        .into_iter()
165
34
        .map(|null_count| {
166
34
            let mut s = ColumnStatistics::new_unknown();
167
34
            s.null_count = Precision::Exact(null_count);
168
34
            s
169
34
        })
170
30
        .collect();
171
30
172
30
    Statistics {
173
30
        num_rows: Precision::Exact(nb_rows),
174
30
        total_byte_size: Precision::Exact(total_byte_size),
175
30
        column_statistics,
176
30
    }
177
30
}
178
179
/// Write in Arrow IPC format.
180
pub struct IPCWriter {
181
    /// path
182
    pub path: PathBuf,
183
    /// inner writer
184
    pub writer: FileWriter<File>,
185
    /// batches written
186
    pub num_batches: usize,
187
    /// rows written
188
    pub num_rows: usize,
189
    /// bytes written
190
    pub num_bytes: usize,
191
}
192
193
impl IPCWriter {
194
    /// Create new writer
195
47
    pub fn new(path: &Path, schema: &Schema) -> Result<Self> {
196
47
        let file = File::create(path).map_err(|e| {
197
0
            DataFusionError::Execution(format!(
198
0
                "Failed to create partition file at {path:?}: {e:?}"
199
0
            ))
200
47
        })
?0
;
201
        Ok(Self {
202
            num_batches: 0,
203
            num_rows: 0,
204
            num_bytes: 0,
205
47
            path: path.into(),
206
47
            writer: FileWriter::try_new(file, schema)
?0
,
207
        })
208
47
    }
209
210
    /// Create new writer with IPC write options
211
0
    pub fn new_with_options(
212
0
        path: &Path,
213
0
        schema: &Schema,
214
0
        write_options: IpcWriteOptions,
215
0
    ) -> Result<Self> {
216
0
        let file = File::create(path).map_err(|e| {
217
0
            DataFusionError::Execution(format!(
218
0
                "Failed to create partition file at {path:?}: {e:?}"
219
0
            ))
220
0
        })?;
221
        Ok(Self {
222
            num_batches: 0,
223
            num_rows: 0,
224
            num_bytes: 0,
225
0
            path: path.into(),
226
0
            writer: FileWriter::try_new_with_options(file, schema, write_options)?,
227
        })
228
0
    }
229
    /// Write one single batch
230
51
    pub fn write(&mut self, batch: &RecordBatch) -> Result<()> {
231
51
        self.writer.write(batch)
?0
;
232
51
        self.num_batches += 1;
233
51
        self.num_rows += batch.num_rows();
234
51
        let num_bytes: usize = batch.get_array_memory_size();
235
51
        self.num_bytes += num_bytes;
236
51
        Ok(())
237
51
    }
238
239
    /// Finish the writer
240
47
    pub fn finish(&mut self) -> Result<()> {
241
47
        self.writer.finish().map_err(Into::into)
242
47
    }
243
244
    /// Path write to
245
0
    pub fn path(&self) -> &Path {
246
0
        &self.path
247
0
    }
248
}
249
250
/// Checks if the given projection is valid for the given schema.
251
683
pub fn can_project(
252
683
    schema: &arrow_schema::SchemaRef,
253
683
    projection: Option<&Vec<usize>>,
254
683
) -> Result<()> {
255
683
    match projection {
256
0
        Some(columns) => {
257
0
            if columns
258
0
                .iter()
259
0
                .max()
260
0
                .map_or(false, |&i| i >= schema.fields().len())
261
            {
262
0
                Err(arrow_schema::ArrowError::SchemaError(format!(
263
0
                    "project index {} out of bounds, max field {}",
264
0
                    columns.iter().max().unwrap(),
265
0
                    schema.fields().len()
266
0
                ))
267
0
                .into())
268
            } else {
269
0
                Ok(())
270
            }
271
        }
272
683
        None => Ok(()),
273
    }
274
683
}
275
276
#[cfg(test)]
277
mod tests {
278
    use super::*;
279
280
    use arrow::{
281
        array::{Float32Array, Float64Array, UInt64Array},
282
        datatypes::{DataType, Field},
283
    };
284
285
    #[test]
286
1
    fn test_compute_record_batch_statistics_empty() -> Result<()> {
287
1
        let schema = Arc::new(Schema::new(vec![
288
1
            Field::new("f32", DataType::Float32, false),
289
1
            Field::new("f64", DataType::Float64, false),
290
1
        ]));
291
1
        let stats = compute_record_batch_statistics(&[], &schema, Some(vec![0, 1]));
292
1
293
1
        assert_eq!(stats.num_rows, Precision::Exact(0));
294
1
        assert_eq!(stats.total_byte_size, Precision::Exact(0));
295
1
        Ok(())
296
1
    }
297
298
    #[test]
299
1
    fn test_compute_record_batch_statistics() -> Result<()> {
300
1
        let schema = Arc::new(Schema::new(vec![
301
1
            Field::new("f32", DataType::Float32, false),
302
1
            Field::new("f64", DataType::Float64, false),
303
1
            Field::new("u64", DataType::UInt64, false),
304
1
        ]));
305
1
        let batch = RecordBatch::try_new(
306
1
            Arc::clone(&schema),
307
1
            vec![
308
1
                Arc::new(Float32Array::from(vec![1., 2., 3.])),
309
1
                Arc::new(Float64Array::from(vec![9., 8., 7.])),
310
1
                Arc::new(UInt64Array::from(vec![4, 5, 6])),
311
1
            ],
312
1
        )
?0
;
313
314
        // just select f32,f64
315
1
        let select_projection = Some(vec![0, 1]);
316
1
        let byte_size = batch
317
1
            .project(&select_projection.clone().unwrap())
318
1
            .unwrap()
319
1
            .get_array_memory_size();
320
1
321
1
        let actual =
322
1
            compute_record_batch_statistics(&[vec![batch]], &schema, select_projection);
323
1
324
1
        let expected = Statistics {
325
1
            num_rows: Precision::Exact(3),
326
1
            total_byte_size: Precision::Exact(byte_size),
327
1
            column_statistics: vec![
328
1
                ColumnStatistics {
329
1
                    distinct_count: Precision::Absent,
330
1
                    max_value: Precision::Absent,
331
1
                    min_value: Precision::Absent,
332
1
                    null_count: Precision::Exact(0),
333
1
                },
334
1
                ColumnStatistics {
335
1
                    distinct_count: Precision::Absent,
336
1
                    max_value: Precision::Absent,
337
1
                    min_value: Precision::Absent,
338
1
                    null_count: Precision::Exact(0),
339
1
                },
340
1
            ],
341
1
        };
342
1
343
1
        assert_eq!(actual, expected);
344
1
        Ok(())
345
1
    }
346
347
    #[test]
348
1
    fn test_compute_record_batch_statistics_null() -> Result<()> {
349
1
        let schema =
350
1
            Arc::new(Schema::new(vec![Field::new("u64", DataType::UInt64, true)]));
351
1
        let batch1 = RecordBatch::try_new(
352
1
            Arc::clone(&schema),
353
1
            vec![Arc::new(UInt64Array::from(vec![Some(1), None, None]))],
354
1
        )
?0
;
355
1
        let batch2 = RecordBatch::try_new(
356
1
            Arc::clone(&schema),
357
1
            vec![Arc::new(UInt64Array::from(vec![Some(1), Some(2), None]))],
358
1
        )
?0
;
359
1
        let byte_size = batch1.get_array_memory_size() + batch2.get_array_memory_size();
360
1
        let actual =
361
1
            compute_record_batch_statistics(&[vec![batch1], vec![batch2]], &schema, None);
362
1
363
1
        let expected = Statistics {
364
1
            num_rows: Precision::Exact(6),
365
1
            total_byte_size: Precision::Exact(byte_size),
366
1
            column_statistics: vec![ColumnStatistics {
367
1
                distinct_count: Precision::Absent,
368
1
                max_value: Precision::Absent,
369
1
                min_value: Precision::Absent,
370
1
                null_count: Precision::Exact(3),
371
1
            }],
372
1
        };
373
1
374
1
        assert_eq!(actual, expected);
375
1
        Ok(())
376
1
    }
377
}