Fix

Author: trueleo

datafusion-postgres/src/encoder/struct_encoder.rs

@@ -2,7 +2,11 @@ use std::{error::Error, sync::Arc};
 
 use bytes::{BufMut, BytesMut};
 use datafusion::arrow::array::{Array, StructArray};
-use pgwire::{api::results::FieldFormat, error::PgWireResult, types::ToSqlText};
+use pgwire::{
+    api::results::FieldFormat,
+    error::PgWireResult,
+    types::{ToSqlText, QUOTE_CHECK, QUOTE_ESCAPE},
+};
 use postgres_types::{Field, IsNull, ToSql, Type};
 
 use super::{encode_value, EncodedValue};
@@ -58,7 +62,21 @@ impl super::Encoder for StructEncoder {
             if self.curr_col == 0 {
                 self.row_buffer.put_slice(b"(");
             }
-            value.to_sql_text(data_type, &mut self.row_buffer)?;
+            // encode value in an intermediate buf
+            let mut buf = BytesMut::new();
+            value.to_sql_text(data_type, &mut buf)?;
+            let encoded_value_as_str = String::from_utf8_lossy(&buf);
+            if QUOTE_CHECK.is_match(&encoded_value_as_str) {
+                self.row_buffer.put_u8(b'"');
+                self.row_buffer.put_slice(
+                    QUOTE_ESCAPE
+                        .replace_all(&encoded_value_as_str, r#"\$1"#)
+                        .as_bytes(),
+                );
+                self.row_buffer.put_u8(b'"');
+            } else {
+                self.row_buffer.put_slice(&buf);
+            }
             if self.curr_col == self.num_cols - 1 {
                 self.row_buffer.put_slice(b")");
             } else {

tests-integration/all_types.parquet

@@ -18,6 +18,7 @@
             pa.field("nested_str", pa.string()),
             pa.field("nested_date32", pa.date32()),
             pa.field("nested_timestamp", pa.timestamp("ms")),
+            pa.field("nested_array", pa.list_(pa.int64())),
         ]
     )
 )
@@ -81,6 +82,7 @@
                     "nested_str": "x",
                     "nested_date32": date(2012, 1, 1),
                     "nested_timestamp": datetime(2012, 1, 1),
+                    "nested_array": [1, 2],
                 }
             ],
             [
@@ -89,6 +91,7 @@
                     "nested_str": "y",
                     "nested_date32": date(2012, 1, 2),
                     "nested_timestamp": datetime(2012, 1, 2),
+                    "nested_array": [3, 4],
                 }
             ],
             [
@@ -97,6 +100,7 @@
                     "nested_str": "z",
                     "nested_date32": date(2012, 1, 3),
                     "nested_timestamp": datetime(2012, 1, 3),
+                    "nested_array": [5, 6],
                 }
             ],
         ],

tests-integration/create_arrow_testfile.py

@@ -23,9 +23,9 @@ def assert_select_all(results: list[psycopg.rows.Row], format: str):
                 else ("10", "x", "2012-01-01", "2012-01-01 00:00:00.000000")
             ),
             (
-                [(10, "x", date(2012, 1, 1), datetime(2012, 1, 1))]
+                [(10, "x", date(2012, 1, 1), datetime(2012, 1, 1), [1, 2])]
                 if format == "binary"
-                else [("10", "x", "2012-01-01", "2012-01-01 00:00:00.000000")]
+                else [("10", "x", "2012-01-01", "2012-01-01 00:00:00.000000", "{1,2}")]
             ),
         ),
         (
@@ -41,13 +41,11 @@ def assert_select_all(results: list[psycopg.rows.Row], format: str):
                 if format == "binary"
                 else ("20", "y", "2012-01-02", "2012-01-02 00:00:00.000000")
             ),
-            [
-                (
-                    (20, "y", date(2012, 1, 2), datetime(2012, 1, 2, 0, 0))
-                    if format == "binary"
-                    else ("20", "y", "2012-01-02", "2012-01-02 00:00:00.000000")
-                ),
-            ],
+            (
+                [(20, "y", date(2012, 1, 2), datetime(2012, 1, 2, 0, 0), [3, 4])]
+                if format == "binary"
+                else [("20", "y", "2012-01-02", "2012-01-02 00:00:00.000000", "{3,4}")]
+            ),
         ),
         (
             3,
@@ -62,13 +60,11 @@ def assert_select_all(results: list[psycopg.rows.Row], format: str):
                 if format == "binary"
                 else ("30", "z", "2012-01-03", "2012-01-03 00:00:00.000000")
             ),
-            [
-                (
-                    (30, "z", date(2012, 1, 3), datetime(2012, 1, 3, 0, 0))
-                    if format == "binary"
-                    else ("30", "z", "2012-01-03", "2012-01-03 00:00:00.000000")
-                ),
-            ],
+            (
+                [(30, "z", date(2012, 1, 3), datetime(2012, 1, 3, 0, 0), [5, 6])]
+                if format == "binary"
+                else [("30", "z", "2012-01-03", "2012-01-03 00:00:00.000000", "{5,6}")]
+            ),
         ),
     ]