From fbf90a724cc8077f4bf6ad5f51399c23a078b895 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alja=C5=BE=20Mur=20Er=C5=BEen?= Date: Wed, 21 Feb 2024 14:12:26 +0100 Subject: [PATCH] feat: sqlite decimal --- Cargo.lock | 8 +++--- connector_arrow/Cargo.toml | 2 +- connector_arrow/src/sqlite/append.rs | 26 +++++++++++++++++-- connector_arrow/src/sqlite/mod.rs | 2 ++ connector_arrow/src/util/decimal.rs | 33 +++++++++++++++++++++++++ connector_arrow/src/util/mod.rs | 1 + connector_arrow/tests/it/test_sqlite.rs | 1 + flake.nix | 17 ++++++++----- 8 files changed, 77 insertions(+), 13 deletions(-) create mode 100644 connector_arrow/src/util/decimal.rs diff --git a/Cargo.lock b/Cargo.lock index 49baa5a..1b3d3a6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -624,9 +624,9 @@ dependencies = [ [[package]] name = "duckdb" -version = "0.10.0" +version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e0288a770d9d49f8d262573c00c7c78aa3f1d3df8935a23fef881dc8da7b277" +checksum = "326e8f84acb4d57c4025637f77e89dc3eee0e25b3a79c21cfd8b72db5ecd3c97" dependencies = [ "arrow", "cast", @@ -1031,9 +1031,9 @@ checksum = "13e3bf6590cbc649f4d1a3eefc9d5d6eb746f5200ffb04e5e142700b8faa56e7" [[package]] name = "libduckdb-sys" -version = "0.10.0" +version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d5c3ce660d33747e6c88abacc209abfe906d94ed6962795985abe55758ce076c" +checksum = "666dbbb7ac31fc49fc2535df5d6efe9ce09815783a0ad98eea3564d2f0223974" dependencies = [ "autocfg", "flate2", diff --git a/connector_arrow/Cargo.toml b/connector_arrow/Cargo.toml index 9f13d08..5c52f09 100644 --- a/connector_arrow/Cargo.toml +++ b/connector_arrow/Cargo.toml @@ -47,7 +47,7 @@ default-features = false optional = true [dependencies.duckdb] -version = "0.10" +version = "0.9.2" default-features = false optional = true diff --git a/connector_arrow/src/sqlite/append.rs b/connector_arrow/src/sqlite/append.rs index 42f44cc..126cef2 100644 --- a/connector_arrow/src/sqlite/append.rs +++ b/connector_arrow/src/sqlite/append.rs @@ -126,6 +126,30 @@ impl ConsumeTy for Vec { } } +impl ConsumeTy for Vec { + fn consume(&mut self, ty: &DataType, value: i128) { + self.push(Value::Text(crate::util::decimal::decimal128_to_string( + ty, value, + ))); + } + + fn consume_null(&mut self) { + self.push(Value::Null); + } +} + +impl ConsumeTy for Vec { + fn consume(&mut self, ty: &DataType, value: i256) { + self.push(Value::Text(crate::util::decimal::decimal256_to_string( + ty, value, + ))); + } + + fn consume_null(&mut self) { + self.push(Value::Null); + } +} + impl_consume_ty!(BooleanType, Value::Integer, i64::from); impl_consume_ty!(Int8Type, Value::Integer, i64::from); impl_consume_ty!(Int16Type, Value::Integer, i64::from); @@ -164,8 +188,6 @@ impl_consume_unsupported!( DurationMillisecondType, DurationMicrosecondType, DurationNanosecondType, - Decimal128Type, - Decimal256Type, ) ); diff --git a/connector_arrow/src/sqlite/mod.rs b/connector_arrow/src/sqlite/mod.rs index 1969001..0a9335c 100644 --- a/connector_arrow/src/sqlite/mod.rs +++ b/connector_arrow/src/sqlite/mod.rs @@ -59,6 +59,8 @@ impl Connection for SQLiteConnection { DataType::LargeBinary => Some(DataType::LargeBinary), DataType::Utf8 => Some(DataType::LargeUtf8), DataType::LargeUtf8 => Some(DataType::LargeUtf8), + DataType::Decimal128(_, _) => Some(DataType::LargeUtf8), + DataType::Decimal256(_, _) => Some(DataType::LargeUtf8), _ => None, } } diff --git a/connector_arrow/src/util/decimal.rs b/connector_arrow/src/util/decimal.rs new file mode 100644 index 0000000..338b868 --- /dev/null +++ b/connector_arrow/src/util/decimal.rs @@ -0,0 +1,33 @@ +use arrow::datatypes::{i256, DataType}; + +pub fn decimal128_to_string(ty: &DataType, value: i128) -> String { + let DataType::Decimal128(_, scale) = ty else { + unreachable!(); + }; + insert_dot(scale, value.to_string()) +} + +pub fn decimal256_to_string(ty: &DataType, value: i256) -> String { + let DataType::Decimal256(_, scale) = ty else { + unreachable!(); + }; + insert_dot(scale, value.to_string()) +} + +fn insert_dot(scale: &i8, value: String) -> String { + if *scale > 0 { + let minus_len: usize = if value.starts_with('-') { 1 } else { 0 }; + + let fracs = *scale as usize; + if fracs >= value.len() - minus_len { + let (minus, value) = value.split_at(minus_len); + format!("{minus}0.{}{value}", "0".repeat(fracs - value.len())) + } else { + let mut value = value; + value.insert(value.len() - fracs, '.'); + value + } + } else { + value + } +} diff --git a/connector_arrow/src/util/mod.rs b/connector_arrow/src/util/mod.rs index 8c69b1c..a9eef5b 100644 --- a/connector_arrow/src/util/mod.rs +++ b/connector_arrow/src/util/mod.rs @@ -3,6 +3,7 @@ mod arrow_reader; pub mod coerce; +pub mod decimal; mod row_collect; mod row_reader; mod row_writer; diff --git a/connector_arrow/tests/it/test_sqlite.rs b/connector_arrow/tests/it/test_sqlite.rs index af7c8d9..1410f0f 100644 --- a/connector_arrow/tests/it/test_sqlite.rs +++ b/connector_arrow/tests/it/test_sqlite.rs @@ -20,6 +20,7 @@ fn query_01() { #[case::int("roundtrip::int", spec::int())] #[case::uint("roundtrip::uint", spec::uint())] #[case::float("roundtrip::float", spec::float())] +#[case::decimal("roundtrip::decimal", spec::decimal())] fn roundtrip(#[case] table_name: &str, #[case] spec: spec::ArrowGenSpec) { let mut conn = init(); super::tests::roundtrip(&mut conn, table_name, spec); diff --git a/flake.nix b/flake.nix index 92c3372..d4f031d 100644 --- a/flake.nix +++ b/flake.nix @@ -19,12 +19,14 @@ dontCheckPython = drv: drv.overridePythonAttrs (old: { doCheck = false; }); essentials = with pkgs; [ - fenix_pkgs.stable.cargo - fenix_pkgs.stable.clippy - fenix_pkgs.stable.rust-src - fenix_pkgs.stable.rustc - fenix_pkgs.stable.rustfmt - fenix_pkgs.stable.rust-analyzer + (fenix_pkgs.complete.withComponents [ + "cargo" + "clippy" + "rust-src" + "rustc" + "rustfmt" + "rust-analyzer" + ]) clang # tools @@ -50,6 +52,9 @@ { devShells.default = pkgs.mkShell { buildInputs = essentials ++ dbs; + + # linking duckdb needs this + LD_LIBRARY_PATH = "${pkgs.stdenv.cc.cc.lib}/lib"; }; }); }