fixing lint & more testing

paradedb · Aug 7, 2024 · b9f8e91 · b9f8e91
1 parent 1d0fc36
commit b9f8e91
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git a/src/fdw/base.rs b/src/fdw/base.rs
@@ -149,6 +149,7 @@ pub trait BaseFdw {
             self.get_target_columns().clone().into_iter().enumerate()
         {
             let batch_column = current_batch.column(column_index);
+
             let cell = batch_column.get_cell(
                 current_batch_index,
                 target_column.type_oid,

diff --git a/src/schema/cell.rs b/src/schema/cell.rs
@@ -1089,6 +1089,22 @@ where
                         None => Ok(None),
                     }
                 }
+                DataType::Date32 => match self.get_primitive_value::<Date32Array>(index)? {
+                    Some(timestamp_in_days) => {
+                        Ok(arrow_date32_to_postgres_timestamps(timestamp_in_days)?
+                            .map(Timestamp::from)
+                            .map(Cell::Timestamp))
+                    }
+                    None => Ok(None),
+                },
+                DataType::Date64 => match self.get_primitive_value::<Date64Array>(index)? {
+                    Some(timestamp_in_milliseconds) => Ok(arrow_date64_to_postgres_timestamps(
+                        timestamp_in_milliseconds,
+                    )?
+                    .map(Timestamp::from)
+                    .map(Cell::Timestamp)),
+                    None => Ok(None),
+                },
                 unsupported => Err(DataTypeError::DataTypeMismatch(
                     name.to_string(),
                     unsupported.clone(),

diff --git a/tests/scan.rs b/tests/scan.rs
@@ -52,9 +52,12 @@ fn date_time_record_batch() -> Result<(RecordBatch, FieldSpec, Vec<String>)> {
         ("date64_col", DataType::Date64, false, "date"),
     ]);
     let dates = vec![
-        "2023-04-01 21:10:00 +0000".to_string(), "2023-04-01 22:08:00 +0000".to_string(),
-        "2023-04-02 04:55:00 +0000".to_string(), "2023-04-02 11:45:00 +0000".to_string(),
-        "2023-04-03 01:20:00 +0000".to_string(), "2023-04-03 12:30:00 +0000".to_string(),
+        "2023-04-01 21:10:00 +0000".to_string(),
+        "2023-04-01 22:08:00 +0000".to_string(),
+        "2023-04-02 04:55:00 +0000".to_string(),
+        "2023-04-02 11:45:00 +0000".to_string(),
+        "2023-04-03 01:20:00 +0000".to_string(),
+        "2023-04-03 12:30:00 +0000".to_string(),
     ];
     let (dates_i32, dates_i64): (Vec<_>, Vec<_>) = dates
         .iter()
@@ -361,7 +364,6 @@ async fn test_date_functions_support_with_local_file(
     let fetched_rows =
         "SELECT DATE_PART('day', date32_col), DATE_TRUNC('day', date64_col) FROM dates"
             .fetch_result::<(f64, chrono::DateTime<Utc>)>(&mut conn)?;
-
     assert_eq!(expected_rows.len(), fetched_rows.len());
     assert_eq!(expected_rows, fetched_rows);