Skip to content

Commit 3f1ea0a

Browse files
MauseMause
authored
feat: add interval read support (#291)
* feat(arrow): interval read support * add support for binding intervals * add chrono support for Duration/TimeDelta * further tests * break down units * add interval to test_all_types * switch to struct variant * fix: write out correct extraction logic * fix: clippy * fix: add chrono feature * build: remove duplicate dependency * chore: improve nanos error message * chore: overflow nanos into seconds * chore: add todo * chore: name magic number
1 parent f85893f commit 3f1ea0a

File tree

9 files changed

+172
-13
lines changed

9 files changed

+172
-13
lines changed

‎Cargo.toml‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ extensions-full = ["httpfs", "json", "parquet", "vtab-full"]
3636
buildtime_bindgen = ["libduckdb-sys/buildtime_bindgen"]
3737
modern-full = ["chrono", "serde_json", "url", "r2d2", "uuid", "polars"]
3838
polars = ["dep:polars"]
39+
chrono = ["dep:chrono", "num-integer"]
3940

4041
[dependencies]
4142
# time = { version = "0.3.2", features = ["formatting", "parsing"], optional = true }
@@ -52,14 +53,15 @@ memchr = "2.3"
5253
uuid = { version = "1.0", optional = true }
5354
smallvec = "1.6.1"
5455
cast = { version = "0.3", features = ["std"] }
55-
arrow = { version = "50", default-features = false, features = ["prettyprint", "ffi"] }
56+
arrow = { version = "51", default-features = false, features = ["prettyprint", "ffi"] }
5657
rust_decimal = "1.14"
5758
strum = { version = "0.25", features = ["derive"] }
5859
r2d2 = { version = "0.8.9", optional = true }
5960
calamine = { version = "0.22.0", optional = true }
6061
num = { version = "0.4", optional = true, default-features = false, features = ["std"] }
6162
duckdb-loadable-macros = { version = "0.1.1", path="./duckdb-loadable-macros", optional = true }
6263
polars = { version = "0.35.4", features = ["dtype-full"], optional = true}
64+
num-integer = {version = "0.1.46", optional = true}
6365

6466
[dev-dependencies]
6567
doc-comment = "0.3"

‎src/lib.rs‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -570,6 +570,8 @@ doc_comment::doctest!("../README.md");
570570

571571
#[cfg(test)]
572572
mod test {
573+
use crate::types::Value;
574+
573575
use super::*;
574576
use std::{error::Error as StdError, fmt};
575577

@@ -1297,6 +1299,26 @@ mod test {
12971299
Ok(())
12981300
}
12991301

1302+
#[test]
1303+
fn round_trip_interval() -> Result<()> {
1304+
let db = checked_memory_handle();
1305+
db.execute_batch("CREATE TABLE foo (t INTERVAL);")?;
1306+
1307+
let d = Value::Interval {
1308+
months: 1,
1309+
days: 2,
1310+
nanos: 3,
1311+
};
1312+
db.execute("INSERT INTO foo VALUES (?)", [d])?;
1313+
1314+
let mut stmt = db.prepare("SELECT t FROM foo")?;
1315+
let mut rows = stmt.query([])?;
1316+
let row = rows.next()?.unwrap();
1317+
let d: Value = row.get_unwrap(0);
1318+
assert_eq!(d, d);
1319+
Ok(())
1320+
}
1321+
13001322
#[test]
13011323
fn test_database_name_to_string() -> Result<()> {
13021324
assert_eq!(DatabaseName::Main.to_string(), "main");

‎src/row.rs‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -542,15 +542,29 @@ impl<'stmt> Row<'stmt> {
542542
}
543543
ValueRef::Time64(types::TimeUnit::Microsecond, array.value(row))
544544
}
545+
DataType::Interval(unit) => match unit {
546+
IntervalUnit::MonthDayNano => {
547+
let array = column
548+
.as_any()
549+
.downcast_ref::<array::IntervalMonthDayNanoArray>()
550+
.unwrap();
551+
552+
if array.is_null(row) {
553+
return ValueRef::Null;
554+
}
555+
556+
let value = array.value(row);
557+
558+
// TODO: remove this manual conversion once arrow-rs bug is fixed
559+
let months = (value) as i32;
560+
let days = (value >> 32) as i32;
561+
let nanos = (value >> 64) as i64;
562+
563+
ValueRef::Interval { months, days, nanos }
564+
}
565+
_ => unimplemented!("{:?}", unit),
566+
},
545567
// TODO: support more data types
546-
// DataType::Interval(unit) => match unit {
547-
// IntervalUnit::DayTime => {
548-
// make_string_interval_day_time!(column, row)
549-
// }
550-
// IntervalUnit::YearMonth => {
551-
// make_string_interval_year_month!(column, row)
552-
// }
553-
// },
554568
// DataType::List(_) => make_string_from_list!(column, row),
555569
// DataType::Dictionary(index_type, _value_type) => match **index_type {
556570
// DataType::Int8 => dict_array_value_to_string::<Int8Type>(column, row),

‎src/statement.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -497,6 +497,10 @@ impl Statement<'_> {
497497
};
498498
ffi::duckdb_bind_timestamp(ptr, col as u64, ffi::duckdb_timestamp { micros })
499499
},
500+
ValueRef::Interval { months, days, nanos } => unsafe {
501+
let micros = nanos / 1_000;
502+
ffi::duckdb_bind_interval(ptr, col as u64, ffi::duckdb_interval { months, days, micros })
503+
},
500504
_ => unreachable!("not supported: {}", value.data_type()),
501505
};
502506
result_from_duckdb_prepare(rc, ptr)

‎src/test_all_types.rs‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@ fn test_all_types() -> crate::Result<()> {
1818
// union is currently blocked by https://github.com/duckdb/duckdb/pull/11326
1919
"union",
2020
// these remaining types are not yet supported by duckdb-rs
21-
"interval",
2221
"small_enum",
2322
"medium_enum",
2423
"large_enum",
@@ -219,6 +218,25 @@ fn test_single(idx: &mut i32, column: String, value: ValueRef) {
219218
1 => assert_eq!(value, ValueRef::Blob(&[3, 245])),
220219
_ => assert_eq!(value, ValueRef::Null),
221220
},
221+
"interval" => match idx {
222+
0 => assert_eq!(
223+
value,
224+
ValueRef::Interval {
225+
months: 0,
226+
days: 0,
227+
nanos: 0
228+
}
229+
),
230+
1 => assert_eq!(
231+
value,
232+
ValueRef::Interval {
233+
months: 999,
234+
days: 999,
235+
nanos: 999999999000
236+
}
237+
),
238+
_ => assert_eq!(value, ValueRef::Null),
239+
},
222240
_ => todo!("{column:?}"),
223241
}
224242
}

‎src/types/chrono.rs‎

Lines changed: 77 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,15 @@
11
//! Convert most of the [Time Strings](http://sqlite.org/lang_datefunc.html) to chrono types.
22
3-
use chrono::{DateTime, Local, NaiveDate, NaiveDateTime, NaiveTime, TimeZone, Utc};
3+
use chrono::{DateTime, Duration, Local, NaiveDate, NaiveDateTime, NaiveTime, TimeZone, Utc};
4+
use num_integer::Integer;
45

56
use crate::{
67
types::{FromSql, FromSqlError, FromSqlResult, TimeUnit, ToSql, ToSqlOutput, ValueRef},
78
Result,
89
};
910

11+
use super::Value;
12+
1013
/// ISO 8601 calendar date without timezone => "YYYY-MM-DD"
1114
impl ToSql for NaiveDate {
1215
#[inline]
@@ -126,13 +129,55 @@ impl FromSql for DateTime<Local> {
126129
}
127130
}
128131

132+
impl FromSql for Duration {
133+
fn column_result(value: ValueRef<'_>) -> FromSqlResult<Self> {
134+
match value {
135+
ValueRef::Interval { months, days, nanos } => {
136+
let days = days + (months * 30);
137+
let (additional_seconds, nanos) = nanos.div_mod_floor(&NANOS_PER_SECOND);
138+
let seconds = additional_seconds + (i64::from(days) * 24 * 3600);
139+
140+
match nanos.try_into() {
141+
Ok(nanos) => {
142+
if let Some(duration) = Duration::new(seconds, nanos) {
143+
Ok(duration)
144+
} else {
145+
Err(FromSqlError::Other("Invalid duration".into()))
146+
}
147+
}
148+
Err(err) => Err(FromSqlError::Other(format!("Invalid duration: {err}").into())),
149+
}
150+
}
151+
_ => Err(FromSqlError::InvalidType),
152+
}
153+
}
154+
}
155+
156+
const DAYS_PER_MONTH: i64 = 30;
157+
const SECONDS_PER_DAY: i64 = 24 * 3600;
158+
const NANOS_PER_SECOND: i64 = 1_000_000_000;
159+
const NANOS_PER_DAY: i64 = SECONDS_PER_DAY * NANOS_PER_SECOND;
160+
161+
impl ToSql for Duration {
162+
fn to_sql(&self) -> Result<ToSqlOutput<'_>> {
163+
let nanos = self.num_nanoseconds().unwrap();
164+
let (days, nanos) = nanos.div_mod_floor(&NANOS_PER_DAY);
165+
let (months, days) = days.div_mod_floor(&DAYS_PER_MONTH);
166+
Ok(ToSqlOutput::Owned(Value::Interval {
167+
months: months.try_into().unwrap(),
168+
days: days.try_into().unwrap(),
169+
nanos,
170+
}))
171+
}
172+
}
173+
129174
#[cfg(test)]
130175
mod test {
131176
use crate::{
132-
types::{FromSql, ValueRef},
177+
types::{FromSql, ToSql, ToSqlOutput, ValueRef},
133178
Connection, Result,
134179
};
135-
use chrono::{DateTime, Duration, Local, NaiveDate, NaiveDateTime, NaiveTime, TimeZone, Utc};
180+
use chrono::{DateTime, Duration, Local, NaiveDate, NaiveDateTime, NaiveTime, TimeDelta, TimeZone, Utc};
136181

137182
fn checked_memory_handle() -> Result<Connection> {
138183
let db = Connection::open_in_memory()?;
@@ -216,6 +261,35 @@ mod test {
216261
Ok(())
217262
}
218263

264+
#[test]
265+
fn test_time_delta_roundtrip() {
266+
roundtrip_type(TimeDelta::new(3600, 0).unwrap());
267+
roundtrip_type(TimeDelta::new(3600, 1000).unwrap());
268+
}
269+
270+
#[test]
271+
fn test_time_delta() -> Result<()> {
272+
let db = checked_memory_handle()?;
273+
let td = TimeDelta::new(3600, 0).unwrap();
274+
275+
let row: Result<TimeDelta> = db.query_row("SELECT ?", [td], |row| Ok(row.get(0)))?;
276+
277+
assert_eq!(row.unwrap(), td);
278+
279+
Ok(())
280+
}
281+
282+
fn roundtrip_type<T: FromSql + ToSql + Eq + std::fmt::Debug>(td: T) {
283+
let sqled = td.to_sql().unwrap();
284+
let value = match sqled {
285+
ToSqlOutput::Borrowed(v) => v,
286+
ToSqlOutput::Owned(ref v) => ValueRef::from(v),
287+
};
288+
let reversed = FromSql::column_result(value).unwrap();
289+
290+
assert_eq!(td, reversed);
291+
}
292+
219293
#[test]
220294
fn test_date_time_local() -> Result<()> {
221295
let db = checked_memory_handle()?;

‎src/types/mod.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,8 @@ pub enum Type {
144144
Date32,
145145
/// TIME64
146146
Time64,
147+
/// INTERVAL
148+
Interval,
147149
/// Any
148150
Any,
149151
}
@@ -170,6 +172,7 @@ impl fmt::Display for Type {
170172
Type::Blob => f.pad("Blob"),
171173
Type::Date32 => f.pad("Date32"),
172174
Type::Time64 => f.pad("Time64"),
175+
Type::Interval => f.pad("Interval"),
173176
Type::Any => f.pad("Any"),
174177
}
175178
}

‎src/types/value.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,15 @@ pub enum Value {
4646
Date32(i32),
4747
/// The value is a time64
4848
Time64(TimeUnit, i64),
49+
/// The value is an interval (month, day, nano)
50+
Interval {
51+
/// months
52+
months: i32,
53+
/// days
54+
days: i32,
55+
/// nanos
56+
nanos: i64,
57+
},
4958
}
5059

5160
impl From<Null> for Value {
@@ -212,6 +221,7 @@ impl Value {
212221
Value::Blob(_) => Type::Blob,
213222
Value::Date32(_) => Type::Date32,
214223
Value::Time64(..) => Type::Time64,
224+
Value::Interval { .. } => Type::Interval,
215225
}
216226
}
217227
}

‎src/types/value_ref.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,15 @@ pub enum ValueRef<'a> {
6161
Date32(i32),
6262
/// The value is a time64
6363
Time64(TimeUnit, i64),
64+
/// The value is an interval (month, day, nano)
65+
Interval {
66+
/// months
67+
months: i32,
68+
/// days
69+
days: i32,
70+
/// nanos
71+
nanos: i64,
72+
},
6473
}
6574

6675
impl ValueRef<'_> {
@@ -87,6 +96,7 @@ impl ValueRef<'_> {
8796
ValueRef::Blob(_) => Type::Blob,
8897
ValueRef::Date32(_) => Type::Date32,
8998
ValueRef::Time64(..) => Type::Time64,
99+
ValueRef::Interval { .. } => Type::Interval,
90100
}
91101
}
92102
}
@@ -140,6 +150,7 @@ impl From<ValueRef<'_>> for Value {
140150
ValueRef::Blob(b) => Value::Blob(b.to_vec()),
141151
ValueRef::Date32(d) => Value::Date32(d),
142152
ValueRef::Time64(t, d) => Value::Time64(t, d),
153+
ValueRef::Interval { months, days, nanos } => Value::Interval { months, days, nanos },
143154
}
144155
}
145156
}
@@ -181,6 +192,7 @@ impl<'a> From<&'a Value> for ValueRef<'a> {
181192
Value::Blob(ref b) => ValueRef::Blob(b),
182193
Value::Date32(d) => ValueRef::Date32(d),
183194
Value::Time64(t, d) => ValueRef::Time64(t, d),
195+
Value::Interval { months, days, nanos } => ValueRef::Interval { months, days, nanos },
184196
}
185197
}
186198
}

0 commit comments

Comments
 (0)