Skip to content

Commit

Permalink
add tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@ariesdevil
ariesdevil committed Jan 31, 2024
1 parent 5681a1b commit 09a9450
Show file tree
Hide file tree
Showing 3 changed files with 114 additions and 0 deletions.
68 changes: 68 additions & 0 deletions src/query/expression/src/converts/arrow2/from.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -142,6 +143,7 @@ fn arrow_type_to_table_type(ty: &ArrowDataType, is_nullable: bool) -> Result<Tab
ARROW_EXT_TYPE_EMPTY_MAP => 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)?,
},
_ => {
Expand Down Expand Up @@ -708,6 +710,72 @@ impl Column {
unsafe { std::mem::transmute::<Buffer<i64>, Buffer<u64>>(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::<databend_common_arrow::arrow::array::BinaryArray<i32>>()
.expect(
"fail to read `Geometry` from arrow: array should be `BinaryArray<i32>`",
);
let offsets = arrow_col
.offsets()
.buffer()
.iter()
.map(|x| *x as u64)
.collect::<Vec<_>>();
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::<databend_common_arrow::arrow::array::BinaryArray<i64>>()
.expect(
"fail to read `Geometry` from arrow: array should be `BinaryArray<i64>`",
);
let offsets = arrow_col.offsets().clone().into_inner();
let offsets =
unsafe { std::mem::transmute::<Buffer<i64>, Buffer<u64>>(offsets) };
Column::Geometry(BinaryColumn::new(arrow_col.values().clone(), offsets))
}
(DataType::Geometry, ArrowDataType::Binary) => {
let arrow_col = arrow_col
.as_any()
.downcast_ref::<databend_common_arrow::arrow::array::BinaryArray<i32>>()
.expect(
"fail to read `Geometry` from arrow: array should be `BinaryArray<i32>`",
);
let offsets = arrow_col
.offsets()
.buffer()
.iter()
.map(|x| *x as u64)
.collect::<Vec<_>>();
Column::Geometry(BinaryColumn::new(
arrow_col.values().clone(),
offsets.into(),
))
}
(DataType::Geometry, ArrowDataType::LargeBinary) => {
let arrow_col = arrow_col
.as_any()
.downcast_ref::<databend_common_arrow::arrow::array::BinaryArray<i64>>()
.expect(
"fail to read `Geometry` from arrow: array should be `BinaryArray<i64>`",
);
let offsets = arrow_col.offsets().clone().into_inner();
let offsets =
unsafe { std::mem::transmute::<Buffer<i64>, Buffer<u64>>(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)?
}
Expand Down
23 changes: 23 additions & 0 deletions src/query/functions/tests/it/scalars/testdata/geometry.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<Float64, Float64>(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<Geometry>(st_makepoint<Float64, Float64>(to_float64<Decimal(2, 1)>(7.0_d128(2,1)), to_float64<Decimal(2, 1)>(minus<Decimal(2, 1)>(8.0_d128(2,1)))))
Expand Down
23 changes: 23 additions & 0 deletions tests/sqllogictests/suites/query/02_function/02_0060_function_geometry.test
View file
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 09a9450

Please sign in to comment.