test

Author: forsaken628

src/query/service/src/servers/http/v1/http_query_handlers.rs

@@ -1005,7 +1005,7 @@ impl Header for BodyFormat {
     {
         if let Some(v) = values.next() {
             match v.to_str() {
-                Ok("application/vnd.apache.arrow.file") => return Ok(BodyFormat::Arrow),
+                Ok(s) if s == BodyFormat::Arrow.content_type() => return Ok(BodyFormat::Arrow),
                 Err(_) => return Err(headers::Error::invalid()),
                 _ => {}
             };
@@ -1019,7 +1019,7 @@ impl Header for BodyFormat {
 }
 
 impl BodyFormat {
-    pub fn content_type(&self) -> &'static str {
+    pub const fn content_type(&self) -> &'static str {
         match self {
             BodyFormat::Json => "application/json",
             BodyFormat::Arrow => "application/vnd.apache.arrow.stream",

src/query/service/src/servers/http/v1/query/blocks_serializer.rs

@@ -15,7 +15,10 @@
 use std::cell::RefCell;
 use std::ops::DerefMut;
 
+use arrow_ipc::writer::IpcWriteOptions;
 use arrow_ipc::writer::StreamWriter;
+use arrow_ipc::CompressionType;
+use arrow_ipc::MetadataVersion;
 use databend_common_exception::Result;
 use databend_common_expression::types::date::date_to_string;
 use databend_common_expression::types::interval::interval_to_string;
@@ -105,7 +108,9 @@ impl BlocksSerializer {
 
     pub fn to_arrow_ipc(&self, schema: arrow_schema::Schema) -> Result<Vec<u8>> {
         let mut buf = Vec::new();
-        let mut writer = StreamWriter::try_new(&mut buf, &schema)?;
+        let opts = IpcWriteOptions::try_new(8, false, MetadataVersion::V5)?
+            .try_with_compression(Some(CompressionType::LZ4_FRAME))?;
+        let mut writer = StreamWriter::try_new_with_options(&mut buf, &schema, opts)?;
 
         let mut data_schema = None;
         for (block, _) in &self.columns {

tests/nox/noxfile.py

@@ -61,7 +61,7 @@ def run_jdbc_test(session, driver_version, main_version):
 
 @nox.session
 def test_suites(session):
-    session.install("pytest", "requests", "pytest-asyncio")
+    session.install("pytest", "requests", "pytest-asyncio", "pyarrow")
     # Usage: nox -s test_suites -- suites/1_stateful/09_http_handler/test_09_0007_session.py::test_session
     session.run("pytest", *session.posargs)
 

tests/nox/suites/1_stateful/09_http_handler/test_09_0015_arrow_ipc.py

@@ -0,0 +1,59 @@
+import json
+
+import requests
+import pyarrow.ipc as ipc
+
+auth = ("root", "")
+
+
+def do_query(query, session, pagination):
+    url = f"http://localhost:8000/v1/query"
+    payload = {
+        "sql": query,
+    }
+    if session:
+        payload["session"] = session
+    if pagination:
+        payload["pagination"] = pagination
+    headers = {
+        "Accept": "application/vnd.apache.arrow.stream",
+        "Content-Type": "application/json",
+    }
+
+    return requests.post(url, headers=headers, json=payload, auth=auth)
+
+
+def test_arrow_ipc():
+    pagination = {
+        "max_rows_per_page": 20,
+    }
+    resp = do_query("select * from numbers(97)", session=None, pagination=pagination)
+
+    # print("content", len(resp.content))
+    # IpcWriteOptions(alignment 64 compression None) content: 1672
+    # IpcWriteOptions(alignment 8 compression lz4) content: 1448
+
+    rows = 0
+    with ipc.open_stream(resp.content) as reader:
+        header = json.loads(reader.schema.metadata[b"response_header"])
+        assert header["error"] == None
+        for batch in reader:
+            rows += batch.num_rows
+
+    for _ in range(30):
+        if header.get("next_uri") == None:
+            break
+
+        uri = f"http://localhost:8000/{header['next_uri']}"
+        resp = requests.get(
+            uri, auth=auth, headers={"Accept": "application/vnd.apache.arrow.stream"}
+        )
+        with ipc.open_stream(resp.content) as reader:
+            header = json.loads(reader.schema.metadata[b"response_header"])
+            assert header["error"] == None
+            for batch in reader:
+                rows += batch.num_rows
+                if rows < 96:
+                    assert batch.num_rows == 20
+
+    assert rows == 97