datafusion_common/
test_util.rs

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//! Utility functions to make testing DataFusion based crates easier
19
20use crate::arrow::util::pretty::pretty_format_batches_with_options;
21use arrow::array::{ArrayRef, RecordBatch};
22use arrow::error::ArrowError;
23use std::fmt::Display;
24use std::{error::Error, path::PathBuf};
25
26/// Converts a vector or array into an ArrayRef.
27pub trait IntoArrayRef {
28    fn into_array_ref(self) -> ArrayRef;
29}
30
31pub fn format_batches(results: &[RecordBatch]) -> Result<impl Display, ArrowError> {
32    let datafusion_format_options = crate::config::FormatOptions::default();
33
34    let arrow_format_options: arrow::util::display::FormatOptions =
35        (&datafusion_format_options).try_into().unwrap();
36
37    pretty_format_batches_with_options(results, &arrow_format_options)
38}
39
40/// Compares formatted output of a record batch with an expected
41/// vector of strings, with the result of pretty formatting record
42/// batches. This is a macro so errors appear on the correct line
43///
44/// Designed so that failure output can be directly copy/pasted
45/// into the test code as expected results.
46///
47/// Expects to be called about like this:
48///
49/// `assert_batches_eq!(expected_lines: &[&str], batches: &[RecordBatch])`
50///
51/// # Example
52/// ```
53/// # use std::sync::Arc;
54/// # use arrow::record_batch::RecordBatch;
55/// # use arrow::array::{ArrayRef, Int32Array};
56/// # use datafusion_common::assert_batches_eq;
57/// let col: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
58/// let batch = RecordBatch::try_from_iter([("column", col)]).unwrap();
59/// // Expected output is a vec of strings
60/// let expected = vec![
61///     "+--------+",
62///     "| column |",
63///     "+--------+",
64///     "| 1      |",
65///     "| 2      |",
66///     "+--------+",
67/// ];
68/// // compare the formatted output of the record batch with the expected output
69/// assert_batches_eq!(expected, &[batch]);
70/// ```
71#[macro_export]
72macro_rules! assert_batches_eq {
73    ($EXPECTED_LINES: expr, $CHUNKS: expr) => {
74        let expected_lines: Vec<String> =
75            $EXPECTED_LINES.iter().map(|&s| s.into()).collect();
76
77        let formatted = $crate::test_util::format_batches($CHUNKS)
78            .unwrap()
79            .to_string();
80
81        let actual_lines: Vec<&str> = formatted.trim().lines().collect();
82
83        assert_eq!(
84            expected_lines, actual_lines,
85            "\n\nexpected:\n\n{:#?}\nactual:\n\n{:#?}\n\n",
86            expected_lines, actual_lines
87        );
88    };
89}
90
91pub fn batches_to_string(batches: &[RecordBatch]) -> String {
92    let actual = format_batches(batches).unwrap().to_string();
93
94    actual.trim().to_string()
95}
96
97pub fn batches_to_sort_string(batches: &[RecordBatch]) -> String {
98    let actual_lines = format_batches(batches).unwrap().to_string();
99
100    let mut actual_lines: Vec<&str> = actual_lines.trim().lines().collect();
101
102    // sort except for header + footer
103    let num_lines = actual_lines.len();
104    if num_lines > 3 {
105        actual_lines.as_mut_slice()[2..num_lines - 1].sort_unstable()
106    }
107
108    actual_lines.join("\n")
109}
110
111/// Compares formatted output of a record batch with an expected
112/// vector of strings in a way that order does not matter.
113/// This is a macro so errors appear on the correct line
114///
115/// See [`assert_batches_eq`] for more details and example.
116///
117/// Expects to be called about like this:
118///
119/// `assert_batch_sorted_eq!(expected_lines: &[&str], batches: &[RecordBatch])`
120#[macro_export]
121macro_rules! assert_batches_sorted_eq {
122    ($EXPECTED_LINES: expr, $CHUNKS: expr) => {
123        let mut expected_lines: Vec<String> =
124            $EXPECTED_LINES.iter().map(|&s| s.into()).collect();
125
126        // sort except for header + footer
127        let num_lines = expected_lines.len();
128        if num_lines > 3 {
129            expected_lines.as_mut_slice()[2..num_lines - 1].sort_unstable()
130        }
131
132        let formatted = $crate::test_util::format_batches($CHUNKS)
133            .unwrap()
134            .to_string();
135        // fix for windows: \r\n -->
136
137        let mut actual_lines: Vec<&str> = formatted.trim().lines().collect();
138
139        // sort except for header + footer
140        let num_lines = actual_lines.len();
141        if num_lines > 3 {
142            actual_lines.as_mut_slice()[2..num_lines - 1].sort_unstable()
143        }
144
145        assert_eq!(
146            expected_lines, actual_lines,
147            "\n\nexpected:\n\n{:#?}\nactual:\n\n{:#?}\n\n",
148            expected_lines, actual_lines
149        );
150    };
151}
152
153/// A macro to assert that one string is contained within another with
154/// a nice error message if they are not.
155///
156/// Usage: `assert_contains!(actual, expected)`
157///
158/// Is a macro so test error
159/// messages are on the same line as the failure;
160///
161/// Both arguments must be convertible into Strings ([`Into`]<[`String`]>)
162#[macro_export]
163macro_rules! assert_contains {
164    ($ACTUAL: expr, $EXPECTED: expr) => {
165        let actual_value: String = $ACTUAL.into();
166        let expected_value: String = $EXPECTED.into();
167        assert!(
168            actual_value.contains(&expected_value),
169            "Can not find expected in actual.\n\nExpected:\n{}\n\nActual:\n{}",
170            expected_value,
171            actual_value
172        );
173    };
174}
175
176/// A macro to assert that one string is NOT contained within another with
177/// a nice error message if they are.
178///
179/// Usage: `assert_not_contains!(actual, unexpected)`
180///
181/// Is a macro so test error
182/// messages are on the same line as the failure;
183///
184/// Both arguments must be convertible into Strings ([`Into`]<[`String`]>)
185#[macro_export]
186macro_rules! assert_not_contains {
187    ($ACTUAL: expr, $UNEXPECTED: expr) => {
188        let actual_value: String = $ACTUAL.into();
189        let unexpected_value: String = $UNEXPECTED.into();
190        assert!(
191            !actual_value.contains(&unexpected_value),
192            "Found unexpected in actual.\n\nUnexpected:\n{}\n\nActual:\n{}",
193            unexpected_value,
194            actual_value
195        );
196    };
197}
198
199/// Returns the datafusion test data directory, which is by default rooted at `datafusion/core/tests/data`.
200///
201/// The default can be overridden by the optional environment
202/// variable `DATAFUSION_TEST_DATA`
203///
204/// panics when the directory can not be found.
205///
206/// Example:
207/// ```
208/// let testdata = datafusion_common::test_util::datafusion_test_data();
209/// let csvdata = format!("{}/window_1.csv", testdata);
210/// assert!(std::path::PathBuf::from(csvdata).exists());
211/// ```
212pub fn datafusion_test_data() -> String {
213    match get_data_dir("DATAFUSION_TEST_DATA", "../../datafusion/core/tests/data") {
214        Ok(pb) => pb.display().to_string(),
215        Err(err) => panic!("failed to get arrow data dir: {err}"),
216    }
217}
218
219/// Returns the arrow test data directory, which is by default stored
220/// in a git submodule rooted at `testing/data`.
221///
222/// The default can be overridden by the optional environment
223/// variable `ARROW_TEST_DATA`
224///
225/// panics when the directory can not be found.
226///
227/// Example:
228/// ```
229/// let testdata = datafusion_common::test_util::arrow_test_data();
230/// let csvdata = format!("{}/csv/aggregate_test_100.csv", testdata);
231/// assert!(std::path::PathBuf::from(csvdata).exists());
232/// ```
233pub fn arrow_test_data() -> String {
234    match get_data_dir("ARROW_TEST_DATA", "../../testing/data") {
235        Ok(pb) => pb.display().to_string(),
236        Err(err) => panic!("failed to get arrow data dir: {err}"),
237    }
238}
239
240/// Returns the parquet test data directory, which is by default
241/// stored in a git submodule rooted at
242/// `parquet-testing/data`.
243///
244/// The default can be overridden by the optional environment variable
245/// `PARQUET_TEST_DATA`
246///
247/// panics when the directory can not be found.
248///
249/// Example:
250/// ```
251/// let testdata = datafusion_common::test_util::parquet_test_data();
252/// let filename = format!("{}/binary.parquet", testdata);
253/// assert!(std::path::PathBuf::from(filename).exists());
254/// ```
255#[cfg(feature = "parquet")]
256pub fn parquet_test_data() -> String {
257    match get_data_dir("PARQUET_TEST_DATA", "../../parquet-testing/data") {
258        Ok(pb) => {
259            let mut path = pb.display().to_string();
260            if cfg!(target_os = "windows") {
261                // Replace backslashes (Windows paths; avoids some test issues).
262                path = path.replace("\\", "/");
263            }
264            path
265        }
266        Err(err) => panic!("failed to get parquet data dir: {err}"),
267    }
268}
269
270/// Returns a directory path for finding test data.
271///
272/// udf_env: name of an environment variable
273///
274/// submodule_dir: fallback path (relative to CARGO_MANIFEST_DIR)
275///
276///  Returns either:
277/// The path referred to in `udf_env` if that variable is set and refers to a directory
278/// The submodule_data directory relative to CARGO_MANIFEST_PATH
279pub fn get_data_dir(
280    udf_env: &str,
281    submodule_data: &str,
282) -> Result<PathBuf, Box<dyn Error>> {
283    // Try user defined env.
284    if let Ok(dir) = std::env::var(udf_env) {
285        let trimmed = dir.trim().to_string();
286        if !trimmed.is_empty() {
287            let pb = PathBuf::from(trimmed);
288            if pb.is_dir() {
289                return Ok(pb);
290            } else {
291                return Err(format!(
292                    "the data dir `{}` defined by env {} not found",
293                    pb.display(),
294                    udf_env
295                )
296                .into());
297            }
298        }
299    }
300
301    // The env is undefined or its value is trimmed to empty, let's try default dir.
302
303    // env "CARGO_MANIFEST_DIR" is "the directory containing the manifest of your package",
304    // set by `cargo run` or `cargo test`, see:
305    // https://doc.rust-lang.org/cargo/reference/environment-variables.html
306    let dir = env!("CARGO_MANIFEST_DIR");
307
308    let pb = PathBuf::from(dir).join(submodule_data);
309    if pb.is_dir() {
310        Ok(pb)
311    } else {
312        Err(format!(
313            "env `{}` is undefined or has empty value, and the pre-defined data dir `{}` not found\n\
314             HINT: try running `git submodule update --init`",
315            udf_env,
316            pb.display(),
317        ).into())
318    }
319}
320
321#[macro_export]
322macro_rules! create_array {
323    (Boolean, $values: expr) => {
324        std::sync::Arc::new($crate::arrow::array::BooleanArray::from($values))
325    };
326    (Int8, $values: expr) => {
327        std::sync::Arc::new($crate::arrow::array::Int8Array::from($values))
328    };
329    (Int16, $values: expr) => {
330        std::sync::Arc::new($crate::arrow::array::Int16Array::from($values))
331    };
332    (Int32, $values: expr) => {
333        std::sync::Arc::new($crate::arrow::array::Int32Array::from($values))
334    };
335    (Int64, $values: expr) => {
336        std::sync::Arc::new($crate::arrow::array::Int64Array::from($values))
337    };
338    (UInt8, $values: expr) => {
339        std::sync::Arc::new($crate::arrow::array::UInt8Array::from($values))
340    };
341    (UInt16, $values: expr) => {
342        std::sync::Arc::new($crate::arrow::array::UInt16Array::from($values))
343    };
344    (UInt32, $values: expr) => {
345        std::sync::Arc::new($crate::arrow::array::UInt32Array::from($values))
346    };
347    (UInt64, $values: expr) => {
348        std::sync::Arc::new($crate::arrow::array::UInt64Array::from($values))
349    };
350    (Float16, $values: expr) => {
351        std::sync::Arc::new($crate::arrow::array::Float16Array::from($values))
352    };
353    (Float32, $values: expr) => {
354        std::sync::Arc::new($crate::arrow::array::Float32Array::from($values))
355    };
356    (Float64, $values: expr) => {
357        std::sync::Arc::new($crate::arrow::array::Float64Array::from($values))
358    };
359    (Utf8, $values: expr) => {
360        std::sync::Arc::new($crate::arrow::array::StringArray::from($values))
361    };
362}
363
364/// Creates a record batch from literal slice of values, suitable for rapid
365/// testing and development.
366///
367/// **Deprecated**: prefer the upstream macro from `arrow`,
368/// [`arrow::array::record_batch`], which now supports both the literal slice
369/// form shown below and a variable/expression form.
370///
371/// Example:
372/// ```
373/// use arrow::array::record_batch;
374/// let batch = record_batch!(
375///     ("a", Int32, vec![1, 2, 3]),
376///     ("b", Float64, vec![Some(4.0), None, Some(5.0)]),
377///     ("c", Utf8, vec!["alpha", "beta", "gamma"])
378/// );
379/// ```
380#[deprecated(since = "55.0.0", note = "Use `arrow::array::record_batch` instead")]
381#[macro_export]
382macro_rules! record_batch {
383    ($(($name: expr, $type: ident, $values: expr)),*) => {
384        {
385            let schema = std::sync::Arc::new($crate::arrow::datatypes::Schema::new(vec![
386                $(
387                    $crate::arrow::datatypes::Field::new($name, $crate::arrow::datatypes::DataType::$type, true),
388                )*
389            ]));
390
391            let batch = $crate::arrow::array::RecordBatch::try_new(
392                schema,
393                vec![$(
394                    $crate::create_array!($type, $values),
395                )*]
396            );
397
398            batch
399        }
400    }
401}
402
403pub mod array_conversion {
404    use arrow::array::ArrayRef;
405
406    use super::IntoArrayRef;
407
408    impl IntoArrayRef for Vec<bool> {
409        fn into_array_ref(self) -> ArrayRef {
410            create_array!(Boolean, self)
411        }
412    }
413
414    impl IntoArrayRef for Vec<Option<bool>> {
415        fn into_array_ref(self) -> ArrayRef {
416            create_array!(Boolean, self)
417        }
418    }
419
420    impl IntoArrayRef for &[bool] {
421        fn into_array_ref(self) -> ArrayRef {
422            create_array!(Boolean, self.to_vec())
423        }
424    }
425
426    impl IntoArrayRef for &[Option<bool>] {
427        fn into_array_ref(self) -> ArrayRef {
428            create_array!(Boolean, self.to_vec())
429        }
430    }
431
432    impl IntoArrayRef for Vec<i8> {
433        fn into_array_ref(self) -> ArrayRef {
434            create_array!(Int8, self)
435        }
436    }
437
438    impl IntoArrayRef for Vec<Option<i8>> {
439        fn into_array_ref(self) -> ArrayRef {
440            create_array!(Int8, self)
441        }
442    }
443
444    impl IntoArrayRef for &[i8] {
445        fn into_array_ref(self) -> ArrayRef {
446            create_array!(Int8, self.to_vec())
447        }
448    }
449
450    impl IntoArrayRef for &[Option<i8>] {
451        fn into_array_ref(self) -> ArrayRef {
452            create_array!(Int8, self.to_vec())
453        }
454    }
455
456    impl IntoArrayRef for Vec<i16> {
457        fn into_array_ref(self) -> ArrayRef {
458            create_array!(Int16, self)
459        }
460    }
461
462    impl IntoArrayRef for Vec<Option<i16>> {
463        fn into_array_ref(self) -> ArrayRef {
464            create_array!(Int16, self)
465        }
466    }
467
468    impl IntoArrayRef for &[i16] {
469        fn into_array_ref(self) -> ArrayRef {
470            create_array!(Int16, self.to_vec())
471        }
472    }
473
474    impl IntoArrayRef for &[Option<i16>] {
475        fn into_array_ref(self) -> ArrayRef {
476            create_array!(Int16, self.to_vec())
477        }
478    }
479
480    impl IntoArrayRef for Vec<i32> {
481        fn into_array_ref(self) -> ArrayRef {
482            create_array!(Int32, self)
483        }
484    }
485
486    impl IntoArrayRef for Vec<Option<i32>> {
487        fn into_array_ref(self) -> ArrayRef {
488            create_array!(Int32, self)
489        }
490    }
491
492    impl IntoArrayRef for &[i32] {
493        fn into_array_ref(self) -> ArrayRef {
494            create_array!(Int32, self.to_vec())
495        }
496    }
497
498    impl IntoArrayRef for &[Option<i32>] {
499        fn into_array_ref(self) -> ArrayRef {
500            create_array!(Int32, self.to_vec())
501        }
502    }
503
504    impl IntoArrayRef for Vec<i64> {
505        fn into_array_ref(self) -> ArrayRef {
506            create_array!(Int64, self)
507        }
508    }
509
510    impl IntoArrayRef for Vec<Option<i64>> {
511        fn into_array_ref(self) -> ArrayRef {
512            create_array!(Int64, self)
513        }
514    }
515
516    impl IntoArrayRef for &[i64] {
517        fn into_array_ref(self) -> ArrayRef {
518            create_array!(Int64, self.to_vec())
519        }
520    }
521
522    impl IntoArrayRef for &[Option<i64>] {
523        fn into_array_ref(self) -> ArrayRef {
524            create_array!(Int64, self.to_vec())
525        }
526    }
527
528    impl IntoArrayRef for Vec<u8> {
529        fn into_array_ref(self) -> ArrayRef {
530            create_array!(UInt8, self)
531        }
532    }
533
534    impl IntoArrayRef for Vec<Option<u8>> {
535        fn into_array_ref(self) -> ArrayRef {
536            create_array!(UInt8, self)
537        }
538    }
539
540    impl IntoArrayRef for &[u8] {
541        fn into_array_ref(self) -> ArrayRef {
542            create_array!(UInt8, self.to_vec())
543        }
544    }
545
546    impl IntoArrayRef for &[Option<u8>] {
547        fn into_array_ref(self) -> ArrayRef {
548            create_array!(UInt8, self.to_vec())
549        }
550    }
551
552    impl IntoArrayRef for Vec<u16> {
553        fn into_array_ref(self) -> ArrayRef {
554            create_array!(UInt16, self)
555        }
556    }
557
558    impl IntoArrayRef for Vec<Option<u16>> {
559        fn into_array_ref(self) -> ArrayRef {
560            create_array!(UInt16, self)
561        }
562    }
563
564    impl IntoArrayRef for &[u16] {
565        fn into_array_ref(self) -> ArrayRef {
566            create_array!(UInt16, self.to_vec())
567        }
568    }
569
570    impl IntoArrayRef for &[Option<u16>] {
571        fn into_array_ref(self) -> ArrayRef {
572            create_array!(UInt16, self.to_vec())
573        }
574    }
575
576    impl IntoArrayRef for Vec<u32> {
577        fn into_array_ref(self) -> ArrayRef {
578            create_array!(UInt32, self)
579        }
580    }
581
582    impl IntoArrayRef for Vec<Option<u32>> {
583        fn into_array_ref(self) -> ArrayRef {
584            create_array!(UInt32, self)
585        }
586    }
587
588    impl IntoArrayRef for &[u32] {
589        fn into_array_ref(self) -> ArrayRef {
590            create_array!(UInt32, self.to_vec())
591        }
592    }
593
594    impl IntoArrayRef for &[Option<u32>] {
595        fn into_array_ref(self) -> ArrayRef {
596            create_array!(UInt32, self.to_vec())
597        }
598    }
599
600    impl IntoArrayRef for Vec<u64> {
601        fn into_array_ref(self) -> ArrayRef {
602            create_array!(UInt64, self)
603        }
604    }
605
606    impl IntoArrayRef for Vec<Option<u64>> {
607        fn into_array_ref(self) -> ArrayRef {
608            create_array!(UInt64, self)
609        }
610    }
611
612    impl IntoArrayRef for &[u64] {
613        fn into_array_ref(self) -> ArrayRef {
614            create_array!(UInt64, self.to_vec())
615        }
616    }
617
618    impl IntoArrayRef for &[Option<u64>] {
619        fn into_array_ref(self) -> ArrayRef {
620            create_array!(UInt64, self.to_vec())
621        }
622    }
623
624    impl IntoArrayRef for Vec<half::f16> {
625        fn into_array_ref(self) -> ArrayRef {
626            create_array!(Float16, self)
627        }
628    }
629
630    impl IntoArrayRef for Vec<Option<half::f16>> {
631        fn into_array_ref(self) -> ArrayRef {
632            create_array!(Float16, self)
633        }
634    }
635
636    impl IntoArrayRef for &[half::f16] {
637        fn into_array_ref(self) -> ArrayRef {
638            create_array!(Float16, self.to_vec())
639        }
640    }
641
642    impl IntoArrayRef for &[Option<half::f16>] {
643        fn into_array_ref(self) -> ArrayRef {
644            create_array!(Float16, self.to_vec())
645        }
646    }
647
648    impl IntoArrayRef for Vec<f32> {
649        fn into_array_ref(self) -> ArrayRef {
650            create_array!(Float32, self)
651        }
652    }
653
654    impl IntoArrayRef for Vec<Option<f32>> {
655        fn into_array_ref(self) -> ArrayRef {
656            create_array!(Float32, self)
657        }
658    }
659
660    impl IntoArrayRef for &[f32] {
661        fn into_array_ref(self) -> ArrayRef {
662            create_array!(Float32, self.to_vec())
663        }
664    }
665
666    impl IntoArrayRef for &[Option<f32>] {
667        fn into_array_ref(self) -> ArrayRef {
668            create_array!(Float32, self.to_vec())
669        }
670    }
671
672    impl IntoArrayRef for Vec<f64> {
673        fn into_array_ref(self) -> ArrayRef {
674            create_array!(Float64, self)
675        }
676    }
677
678    impl IntoArrayRef for Vec<Option<f64>> {
679        fn into_array_ref(self) -> ArrayRef {
680            create_array!(Float64, self)
681        }
682    }
683
684    impl IntoArrayRef for &[f64] {
685        fn into_array_ref(self) -> ArrayRef {
686            create_array!(Float64, self.to_vec())
687        }
688    }
689
690    impl IntoArrayRef for &[Option<f64>] {
691        fn into_array_ref(self) -> ArrayRef {
692            create_array!(Float64, self.to_vec())
693        }
694    }
695
696    impl IntoArrayRef for Vec<&str> {
697        fn into_array_ref(self) -> ArrayRef {
698            create_array!(Utf8, self)
699        }
700    }
701
702    impl IntoArrayRef for Vec<Option<&str>> {
703        fn into_array_ref(self) -> ArrayRef {
704            create_array!(Utf8, self)
705        }
706    }
707
708    impl IntoArrayRef for &[&str] {
709        fn into_array_ref(self) -> ArrayRef {
710            create_array!(Utf8, self.to_vec())
711        }
712    }
713
714    impl IntoArrayRef for &[Option<&str>] {
715        fn into_array_ref(self) -> ArrayRef {
716            create_array!(Utf8, self.to_vec())
717        }
718    }
719
720    impl IntoArrayRef for Vec<String> {
721        fn into_array_ref(self) -> ArrayRef {
722            create_array!(Utf8, self)
723        }
724    }
725
726    impl IntoArrayRef for Vec<Option<String>> {
727        fn into_array_ref(self) -> ArrayRef {
728            create_array!(Utf8, self)
729        }
730    }
731
732    impl IntoArrayRef for &[String] {
733        fn into_array_ref(self) -> ArrayRef {
734            create_array!(Utf8, self.to_vec())
735        }
736    }
737
738    impl IntoArrayRef for &[Option<String>] {
739        fn into_array_ref(self) -> ArrayRef {
740            create_array!(Utf8, self.to_vec())
741        }
742    }
743}
744
745#[cfg(test)]
746mod tests {
747    use crate::cast::{as_float64_array, as_int32_array, as_string_array};
748    use crate::error::Result;
749
750    use super::*;
751    use std::env;
752
753    #[test]
754    fn test_data_dir() {
755        let udf_env = "get_data_dir";
756        let cwd = env::current_dir().unwrap();
757
758        let existing_pb = cwd.join("..");
759        let existing = existing_pb.display().to_string();
760        let existing_str = existing.as_str();
761
762        let non_existing = cwd.join("non-existing-dir").display().to_string();
763        let non_existing_str = non_existing.as_str();
764
765        unsafe {
766            env::set_var(udf_env, non_existing_str);
767            let res = get_data_dir(udf_env, existing_str);
768            assert!(res.is_err());
769
770            env::set_var(udf_env, "");
771            let res = get_data_dir(udf_env, existing_str);
772            assert!(res.is_ok());
773            assert_eq!(res.unwrap(), existing_pb);
774
775            env::set_var(udf_env, " ");
776            let res = get_data_dir(udf_env, existing_str);
777            assert!(res.is_ok());
778            assert_eq!(res.unwrap(), existing_pb);
779
780            env::set_var(udf_env, existing_str);
781            let res = get_data_dir(udf_env, existing_str);
782            assert!(res.is_ok());
783            assert_eq!(res.unwrap(), existing_pb);
784
785            env::remove_var(udf_env);
786            let res = get_data_dir(udf_env, non_existing_str);
787            assert!(res.is_err());
788
789            let res = get_data_dir(udf_env, existing_str);
790            assert!(res.is_ok());
791            assert_eq!(res.unwrap(), existing_pb);
792        }
793    }
794
795    #[test]
796    #[cfg(feature = "parquet")]
797    fn test_happy() {
798        let res = arrow_test_data();
799        assert!(PathBuf::from(res).is_dir());
800
801        let res = parquet_test_data();
802        assert!(PathBuf::from(res).is_dir());
803    }
804
805    #[test]
806    #[expect(
807        deprecated,
808        reason = "testing the deprecated record_batch! macro itself"
809    )]
810    fn test_create_record_batch() -> Result<()> {
811        use arrow::array::Array;
812
813        let batch = record_batch!(
814            ("a", Int32, vec![1, 2, 3, 4]),
815            ("b", Float64, vec![Some(4.0), None, Some(5.0), None]),
816            ("c", Utf8, vec!["alpha", "beta", "gamma", "delta"])
817        )?;
818
819        assert_eq!(3, batch.num_columns());
820        assert_eq!(4, batch.num_rows());
821
822        let values: Vec<_> = as_int32_array(batch.column(0))?
823            .values()
824            .iter()
825            .map(|v| v.to_owned())
826            .collect();
827        assert_eq!(values, vec![1, 2, 3, 4]);
828
829        let values: Vec<_> = as_float64_array(batch.column(1))?
830            .values()
831            .iter()
832            .map(|v| v.to_owned())
833            .collect();
834        assert_eq!(values, vec![4.0, 0.0, 5.0, 0.0]);
835
836        let nulls: Vec<_> = as_float64_array(batch.column(1))?
837            .nulls()
838            .unwrap()
839            .iter()
840            .collect();
841        assert_eq!(nulls, vec![true, false, true, false]);
842
843        let values: Vec<_> = as_string_array(batch.column(2))?.iter().flatten().collect();
844        assert_eq!(values, vec!["alpha", "beta", "gamma", "delta"]);
845
846        Ok(())
847    }
848}
datafusion_common/test_util.rs

datafusion_common/
test_util.rs
