From 09a9450f234067db1d170bb64d2f34ea478d18c1 Mon Sep 17 00:00:00 2001 From: Yijun Zhao Date: Wed, 31 Jan 2024 23:05:03 +0800 Subject: [PATCH] add tests --- .../expression/src/converts/arrow2/from.rs | 68 +++++++++++++++++++ .../tests/it/scalars/testdata/geometry.txt | 23 +++++++ .../02_0060_function_geometry.test | 23 +++++++ 3 files changed, 114 insertions(+) create mode 100644 tests/sqllogictests/suites/query/02_function/02_0060_function_geometry.test diff --git a/src/query/expression/src/converts/arrow2/from.rs b/src/query/expression/src/converts/arrow2/from.rs index 03f4e818a9108..497c4f571f5d2 100644 --- a/src/query/expression/src/converts/arrow2/from.rs +++ b/src/query/expression/src/converts/arrow2/from.rs @@ -24,6 +24,7 @@ use databend_common_exception::Result; use super::ARROW_EXT_TYPE_BITMAP; use super::ARROW_EXT_TYPE_EMPTY_ARRAY; use super::ARROW_EXT_TYPE_EMPTY_MAP; +use super::ARROW_EXT_TYPE_GEOMETRY; use super::ARROW_EXT_TYPE_VARIANT; use crate::types::array::ArrayColumn; use crate::types::binary::BinaryColumn; @@ -142,6 +143,7 @@ fn arrow_type_to_table_type(ty: &ArrowDataType, is_nullable: bool) -> Result TableDataType::EmptyMap, ARROW_EXT_TYPE_BITMAP => TableDataType::Bitmap, ARROW_EXT_TYPE_VARIANT => TableDataType::Variant, + ARROW_EXT_TYPE_GEOMETRY => TableDataType::Geometry, _ => arrow_type_to_table_type(data_type, is_nullable)?, }, _ => { @@ -708,6 +710,72 @@ impl Column { unsafe { std::mem::transmute::, Buffer>(offsets) }; Column::Bitmap(BinaryColumn::new(arrow_col.values().clone(), offsets)) } + ( + DataType::Geometry, + ArrowDataType::Extension(name, box ArrowDataType::Binary, None), + ) if name == ARROW_EXT_TYPE_GEOMETRY => { + let arrow_col = arrow_col + .as_any() + .downcast_ref::>() + .expect( + "fail to read `Geometry` from arrow: array should be `BinaryArray`", + ); + let offsets = arrow_col + .offsets() + .buffer() + .iter() + .map(|x| *x as u64) + .collect::>(); + Column::Geometry(BinaryColumn::new( + arrow_col.values().clone(), + offsets.into(), + )) + } + ( + DataType::Geometry, + ArrowDataType::Extension(name, box ArrowDataType::LargeBinary, None), + ) if name == ARROW_EXT_TYPE_GEOMETRY => { + let arrow_col = arrow_col + .as_any() + .downcast_ref::>() + .expect( + "fail to read `Geometry` from arrow: array should be `BinaryArray`", + ); + let offsets = arrow_col.offsets().clone().into_inner(); + let offsets = + unsafe { std::mem::transmute::, Buffer>(offsets) }; + Column::Geometry(BinaryColumn::new(arrow_col.values().clone(), offsets)) + } + (DataType::Geometry, ArrowDataType::Binary) => { + let arrow_col = arrow_col + .as_any() + .downcast_ref::>() + .expect( + "fail to read `Geometry` from arrow: array should be `BinaryArray`", + ); + let offsets = arrow_col + .offsets() + .buffer() + .iter() + .map(|x| *x as u64) + .collect::>(); + Column::Geometry(BinaryColumn::new( + arrow_col.values().clone(), + offsets.into(), + )) + } + (DataType::Geometry, ArrowDataType::LargeBinary) => { + let arrow_col = arrow_col + .as_any() + .downcast_ref::>() + .expect( + "fail to read `Geometry` from arrow: array should be `BinaryArray`", + ); + let offsets = arrow_col.offsets().clone().into_inner(); + let offsets = + unsafe { std::mem::transmute::, Buffer>(offsets) }; + Column::Geometry(BinaryColumn::new(arrow_col.values().clone(), offsets)) + } (data_type, ArrowDataType::Extension(_, arrow_type, _)) => { from_arrow_with_arrow_type(arrow_col, arrow_type, data_type)? } diff --git a/src/query/functions/tests/it/scalars/testdata/geometry.txt b/src/query/functions/tests/it/scalars/testdata/geometry.txt index 5bc779e83eb88..54f4ff3b50572 100644 --- a/src/query/functions/tests/it/scalars/testdata/geometry.txt +++ b/src/query/functions/tests/it/scalars/testdata/geometry.txt @@ -16,6 +16,29 @@ output domain : Undefined output : "POINT(7 -8)" +ast : st_makepoint(a, b) +raw expr : st_makepoint(a::Float64, b::Float64) +checked expr : st_makepoint(a, b) +evaluation: ++--------+---------+---------+--------------+ +| | a | b | Output | ++--------+---------+---------+--------------+ +| Type | Float64 | Float64 | Geometry | +| Domain | {1..=3} | {1..=3} | Undefined | +| Row 0 | 1 | 1 | "POINT(1 1)" | +| Row 1 | 2 | 2 | "POINT(2 2)" | +| Row 2 | 3 | 3 | "POINT(3 3)" | ++--------+---------+---------+--------------+ +evaluation (internal): ++--------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| Column | Data | ++--------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| a | Float64([1, 2, 3]) | +| b | Float64([1, 2, 3]) | +| Output | BinaryColumn { data: 0x0101000000000000000000f03f000000000000f03f010100000000000000000000400000000000000040010100000000000000000008400000000000000840, offsets: [0, 21, 42, 63] } | ++--------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + + ast : to_string(st_makepoint(7.0, -8.0)) raw expr : to_string(st_makepoint(7.0, minus(8.0))) checked expr : to_string(st_makepoint(to_float64(7.0_d128(2,1)), to_float64(minus(8.0_d128(2,1))))) diff --git a/tests/sqllogictests/suites/query/02_function/02_0060_function_geometry.test b/tests/sqllogictests/suites/query/02_function/02_0060_function_geometry.test new file mode 100644 index 0000000000000..5412f6e52a088 --- /dev/null +++ b/tests/sqllogictests/suites/query/02_function/02_0060_function_geometry.test @@ -0,0 +1,23 @@ +statement ok +DROP TABLE IF EXISTS t1 + +statement ok +CREATE TABLE t1 (lat Float64, lon Float64) ENGINE=Memory + +statement ok +INSERT INTO t1 VALUES(55.77922738, 37.63098076), (5.77922738, 7.63098076) + +query T +SELECT to_string(st_makepoint(lat, lon)) FROM t1 +---- +POINT(55.77922738 37.63098076) +POINT(5.77922738 7.63098076) + +query T +SELECT st_makepoint(lat, lon)::String FROM t1 +---- +POINT(55.77922738 37.63098076) +POINT(5.77922738 7.63098076) + +statement ok +DROP TABLE IF EXISTS t1