8000 refactor: enhance datetime and date conversion logic in python_value_… · kosiew/datafusion-python@6fbafcd · GitHub
[go: up one dir, main page]

Skip to content

Commit 6fbafcd

Browse files
committed
refactor: enhance datetime and date conversion logic in python_value_to_scalar_value function
1 parent dc86e77 commit 6fbafcd

File tree

1 file changed

+122
-16
lines changed

1 file changed

+122
-16
lines changed

src/dataframe.rs

Lines changed: 122 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -934,26 +934,45 @@ fn python_value_to_scalar_value(value: &PyObject, py: Python) -> PyDataFusionRes
934934

935935
if let Ok(datetime_cls) = datetime_result {
936936
if let Ok(true) = value.is_instance(datetime_cls) {
937-
if value.is_instance_of::<pyo3::types::PyDateTime>(py) {
938-
if let Ok(naive_dt) = value.extract::<chrono::NaiveDateTime>(py) {
939-
return Ok(ScalarValue::TimestampNanosecond(
940-
Some(naive_dt.timestamp_nanos()),
941-
None,
942-
));
937+
if let Ok(dt) = value.cast_as::<pyo3::types::PyDateTime>(py) {
938+
// Convert Python datetime to timestamp in nanoseconds
939+
let year = dt.get_year() as i32;
940+
let month = dt.get_month() as u8;
941+
let day = dt.get_day() as u8;
942+
let hour = dt.get_hour() as u8;
943+
let minute = dt.get_minute() as u8;
944+
let second = dt.get_second() as u8;
945+
let micro = dt.get_microsecond() as u32;
946+
947+
// Use DataFusion's timestamp conversion logic
948+
if let Ok(ts) =
949+
date_to_timestamp(year, month, day, hour, minute, second, micro * 1000)
950+
{
951+
return Ok(ScalarValue::TimestampNanosecond(Some(ts), None));
943952
}
944953
}
945-
// Check for date (not datetime)
946-
let date_result = py.import("datetime").and_then(|m| m.getattr("date"< 10000 /span>));
947-
if let Ok(date_cls) = date_result {
948-
if let Ok(true) = value.is_instance(date_cls) {
949-
if let Ok(naive_date) = value.extract::<chrono::NaiveDate>(py) {
950-
return Ok(ScalarValue::Date32(Some(
951-
naive_date.num_days_from_ce() - 719163, // Convert from CE to Unix epoch
952-
)));
953-
}
954+
955+
let msg = "Failed to convert Python datetime";
956+
return Err(PyDataFusionError::Common(msg.to_string()));
957+
}
958+
}
959+
960+
// Check for date (not datetime)
961+
let date_result = py.import("datetime").and_then(|m| m.getattr("date"));
962+
if let Ok(date_cls) = date_result {
963+
if let Ok(true) = value.is_instance(date_cls) {
964+
if let Ok(date) = value.cast_as::<pyo3::types::PyDate>(py) {
965+
let year = date.get_year() as i32;
966+
let month = date.get_month() as u8;
967+
let day = date.get_day() as u8;
968+
969+
// Calculate days since Unix epoch (1970-01-01)
970+
if let Ok(days) = date_to_days_since_epoch(year, month, day) {
971+
return Ok(ScalarValue::Date32(Some(days)));
954972
}
955973
}
956-
let msg = "Unsupported datetime type format";
974+
975+
let msg = "Failed to convert Python date";
957976
return Err(PyDataFusionError::Common(msg.to_string()));
958977
}
959978
}
@@ -973,3 +992,90 @@ fn python_value_to_scalar_value(value: &PyObject, py: Python) -> PyDataFusionRes
973992
}
974993
}
975994
}
995+
996+
/// Helper function to convert date components to timestamp in nanoseconds
997+
fn date_to_timestamp(
998+
year: i32,
999+
month: u8,
1000+
day: u8,
1001+
hour: u8,
1002+
minute: u8,
1003+
second: u8,
1004+
nano: u32,
1005+
) -> Result<i64, String> {
1006+
// This is a simplified implementation
1007+
// For production code, consider using a more complete date/time library
1008+
1009+
// Number of days in each month (non-leap year)
1010+
const DAYS_IN_MONTH: [u8; 12] = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31];
1011+
1012+
// Validate inputs
1013+
if month < 1 || month > 12 {
1014+
return Err("Invalid month".to_string());
1015+
}
1016+
1017+
let max_days = if month == 2 && is_leap_year(year) {
1018+
29
1019+
} else {
1020+
DAYS_IN_MONTH[(month - 1) as usize]
1021+
};
1022+
1023+
if day < 1 || day > max_days {
1024+
return Err("Invalid day".to_string());
1025+
}
1026+
1027+
if hour > 23 || minute > 59 || second > 59 {
1028+
return Err("Invalid time".to_string());
1029+
}
1030+
1031+
// Calculate days since epoch
1032+
let days = date_to_days_since_epoch(year, month, day)?;
1033+
1034+
// Convert to seconds and add time components
1035+
let seconds =
1036+
341A days as i64 * 86400 + (hour as i64) * 3600 + (minute as i64) * 60 + (second as i64);
1037+
1038+
// Convert to nanoseconds
1039+
Ok(seconds * 1_000_000_000 + nano as i64)
1040+
}
1041+
1042+
/// Helper function to check if a year is a leap year
1043+
fn is_leap_year(year: i32) -> bool {
1044+
(year % 4 == 0 && year % 100 != 0) || (year % 400 == 0)
1045+
}
1046+
1047+
/// Helper function to convert date to days since Unix epoch (1970-01-01)
1048+
fn date_to_days_since_epoch(year: i32, month: u8, day: u8) -> Result<i32, String> {
1049+
// This is a simplified implementation to calculate days since epoch
1050+
if year < 1970 {
1051+
return Err("Dates before 1970 not supported in this implementation".to_string());
1052+
}
1053+
1054+
let mut days = 0;
1055+
1056+
// Add days for each year since 1970
1057+
for y in 1970..year {
1058+
days += if is_leap_year(y) { 366 } else { 365 };
1059+
}
1060+
1061+
// Add days for each month in the current year
1062+
for m in 1..month {
1063+
days += match m {
1064+
1 | 3 | 5 | 7 | 8 | 10 | 12 => 31,
1065+
4 | 6 | 9 | 11 => 30,
1066+
2 => {
1067+
if is_leap_year(year) {
1068+
29
1069+
} else {
1070+
28
1071+
}
1072+
}
1073+
_ => return Err("Invalid month".to_string()),
1074+
};
1075+
}
1076+
1077+
// Add days in current month
1078+
days += day as i32 - 1; // Subtract 1 because we're counting from the start of the month
1079+
1080+
Ok(days)
1081+
}

0 commit comments

Comments
 (0)
0