Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add geometry data type #352

Merged
merged 1 commit into from
Feb 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions bindings/nodejs/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,7 @@ impl ToNapiValue for Value {
databend_driver::Value::Tuple(_) => String::to_napi_value(env, format!("{}", val.0)),
databend_driver::Value::Bitmap(s) => String::to_napi_value(env, s),
databend_driver::Value::Variant(s) => String::to_napi_value(env, s),
databend_driver::Value::Geometry(s) => String::to_napi_value(env, s),
}
}
}
Expand Down
1 change: 1 addition & 0 deletions bindings/python/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ impl IntoPy<PyObject> for Value {
}
databend_driver::Value::Bitmap(s) => s.into_py(py),
databend_driver::Value::Variant(s) => s.into_py(py),
databend_driver::Value::Geometry(s) => s.into_py(py),
}
}
}
Expand Down
4 changes: 4 additions & 0 deletions cli/src/ast/tokenizer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -286,6 +286,8 @@ pub enum TokenKind {
BETWEEN,
#[token("BIGINT", ignore(ascii_case))]
BIGINT,
#[token("BITMAP", ignore(ascii_case))]
BITMAP,
#[token("BOOL", ignore(ascii_case))]
BOOL,
#[token("BOOLEAN", ignore(ascii_case))]
Expand Down Expand Up @@ -478,6 +480,8 @@ pub enum TokenKind {
FUSE,
#[token("GET", ignore(ascii_case))]
GET,
#[token("GEOMETRY", ignore(ascii_case))]
GEOMETRY,
#[token("GLOBAL", ignore(ascii_case))]
GLOBAL,
#[token("GRAPH", ignore(ascii_case))]
Expand Down
10 changes: 2 additions & 8 deletions cli/src/display.rs
Original file line number Diff line number Diff line change
Expand Up @@ -567,13 +567,7 @@ fn create_table(
let mut header = Vec::with_capacity(column_count);
let mut aligns = Vec::with_capacity(column_count);

render_head(
schema,
&mut widths,
&mut column_map,
&mut header,
&mut aligns,
);
render_head(schema, &mut widths, &column_map, &mut header, &mut aligns);
table.set_header(header);

// render the top rows
Expand Down Expand Up @@ -681,7 +675,7 @@ fn create_table(
fn render_head(
schema: SchemaRef,
widths: &mut [usize],
column_map: &mut Vec<i32>,
column_map: &[i32],
header: &mut Vec<Cell>,
aligns: &mut Vec<CellAlignment>,
) {
Expand Down
1 change: 1 addition & 0 deletions sql/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ flight-sql = ["dep:arrow-array", "dep:arrow-schema", "dep:tonic"]
databend-client = { workspace = true }

chrono = { version = "0.4", default-features = false }
geozero = { version = "0.12.0", features = ["default", "with-wkb"] }
glob = "0.3"
itertools = "0.12"
jsonb = "0.3"
Expand Down
8 changes: 8 additions & 0 deletions sql/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use geozero::error::GeozeroError;

#[derive(Debug)]
pub struct ConvertError {
target: &'static str,
Expand Down Expand Up @@ -186,3 +188,9 @@ impl From<ConvertError> for Error {
Error::Convert(e)
}
}

impl From<GeozeroError> for Error {
fn from(e: GeozeroError) -> Self {
Error::Parsing(e.to_string())
}
}
6 changes: 6 additions & 0 deletions sql/src/schema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@ pub(crate) const ARROW_EXT_TYPE_EMPTY_MAP: &str = "EmptyMap";
pub(crate) const ARROW_EXT_TYPE_VARIANT: &str = "Variant";
#[cfg(feature = "flight-sql")]
pub(crate) const ARROW_EXT_TYPE_BITMAP: &str = "Bitmap";
#[cfg(feature = "flight-sql")]
pub(crate) const ARROW_EXT_TYPE_GEOMETRY: &str = "Geometry";

use databend_client::response::SchemaField as APISchemaField;

Expand Down Expand Up @@ -86,6 +88,7 @@ pub enum DataType {
Tuple(Vec<DataType>),
Variant,
Bitmap,
Geometry,
// Generic(usize),
}

Expand Down Expand Up @@ -139,6 +142,7 @@ impl std::fmt::Display for DataType {
}
DataType::Variant => write!(f, "Variant"),
DataType::Bitmap => write!(f, "Bitmap"),
DataType::Geometry => write!(f, "Geometry"),
}
}
}
Expand Down Expand Up @@ -246,6 +250,7 @@ impl TryFrom<&TypeDesc<'_>> for DataType {
}
"Variant" => DataType::Variant,
"Bitmap" => DataType::Bitmap,
"Geometry" => DataType::Geometry,
_ => return Err(Error::Parsing(format!("Unknown type: {:?}", desc))),
};
Ok(dt)
Expand Down Expand Up @@ -289,6 +294,7 @@ impl TryFrom<&Arc<ArrowField>> for Field {
ARROW_EXT_TYPE_EMPTY_MAP => DataType::EmptyMap,
ARROW_EXT_TYPE_VARIANT => DataType::Variant,
ARROW_EXT_TYPE_BITMAP => DataType::Bitmap,
ARROW_EXT_TYPE_GEOMETRY => DataType::Geometry,
_ => {
return Err(Error::Parsing(format!(
"Unsupported extension datatype for arrow field: {:?}",
Expand Down
28 changes: 26 additions & 2 deletions sql/src/value.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,9 @@ use crate::{
schema::{DecimalDataType, DecimalSize},
};

use geozero::wkb::FromWkb;
use geozero::wkb::WkbDialect;
use geozero::wkt::Ewkt;
use std::fmt::Write;

// Thu 1970-01-01 is R.D. 719163
Expand All @@ -30,7 +33,7 @@ const NULL_VALUE: &str = "NULL";
use {
crate::schema::{
ARROW_EXT_TYPE_BITMAP, ARROW_EXT_TYPE_EMPTY_ARRAY, ARROW_EXT_TYPE_EMPTY_MAP,
ARROW_EXT_TYPE_VARIANT, EXTENSION_KEY,
ARROW_EXT_TYPE_GEOMETRY, ARROW_EXT_TYPE_VARIANT, EXTENSION_KEY,
},
arrow_array::{
Array as ArrowArray, BinaryArray, BooleanArray, Date32Array, Decimal128Array,
Expand Down Expand Up @@ -77,6 +80,7 @@ pub enum Value {
Tuple(Vec<Value>),
Bitmap(String),
Variant(String),
Geometry(String),
}

impl Value {
Expand Down Expand Up @@ -126,6 +130,7 @@ impl Value {
}
Self::Bitmap(_) => DataType::Bitmap,
Self::Variant(_) => DataType::Variant,
Self::Geometry(_) => DataType::Geometry,
}
}
}
Expand Down Expand Up @@ -191,6 +196,7 @@ impl TryFrom<(&DataType, &str)> for Value {
)),
DataType::Bitmap => Ok(Self::Bitmap(v.to_string())),
DataType::Variant => Ok(Self::Variant(v.to_string())),
DataType::Geometry => Ok(Self::Geometry(v.to_string())),

DataType::Nullable(inner) => {
if v == NULL_VALUE {
Expand Down Expand Up @@ -244,6 +250,18 @@ impl TryFrom<(&ArrowField, &Arc<dyn ArrowArray>, usize)> for Value {
None => Err(ConvertError::new("bitmap", format!("{:?}", array)).into()),
};
}
ARROW_EXT_TYPE_GEOMETRY => {
if field.is_nullable() && array.is_null(seq) {
return Ok(Value::Null);
}
return match array.as_any().downcast_ref::<LargeBinaryArray>() {
Some(array) => {
let wkt = parse_geometry(array.value(seq))?;
Ok(Value::Geometry(wkt))
}
None => Err(ConvertError::new("geometry", format!("{:?}", array)).into()),
};
}
_ => {
return Err(ConvertError::new(
"extension",
Expand Down Expand Up @@ -605,7 +623,7 @@ fn encode_value(f: &mut std::fmt::Formatter<'_>, val: &Value, raw: bool) -> std:
Value::Boolean(b) => write!(f, "{}", b),
Value::Number(n) => write!(f, "{}", n),
Value::Binary(s) => write!(f, "{}", hex::encode_upper(s)),
Value::String(s) | Value::Bitmap(s) | Value::Variant(s) => {
Value::String(s) | Value::Bitmap(s) | Value::Variant(s) | Value::Geometry(s) => {
if raw {
write!(f, "{}", s)
} else {
Expand Down Expand Up @@ -800,3 +818,9 @@ pub fn parse_decimal(text: &str, size: DecimalSize) -> Result<NumberValue> {
}
}
}

pub fn parse_geometry(raw_data: &[u8]) -> Result<String> {
let mut data = std::io::Cursor::new(raw_data);
let wkt = Ewkt::from_wkb(&mut data, WkbDialect::Ewkb);
wkt.map(|g| g.0).map_err(|e| e.into())
}
Loading