questdb/egress/arrow/reader.rs
1/*******************************************************************************
2 * ___ _ ____ ____
3 * / _ \ _ _ ___ ___| |_| _ \| __ )
4 * | | | | | | |/ _ \/ __| __| | | | _ \
5 * | |_| | |_| | __/\__ \ |_| |_| | |_) |
6 * \__\_\\__,_|\___||___/\__|____/|____/
7 *
8 * Copyright (c) 2014-2019 Appsicle
9 * Copyright (c) 2019-2025 QuestDB
10 *
11 * Licensed under the Apache License, Version 2.0 (the "License");
12 * you may not use this file except in compliance with the License.
13 * You may obtain a copy of the License at
14 *
15 * http://www.apache.org/licenses/LICENSE-2.0
16 *
17 * Unless required by applicable law or agreed to in writing, software
18 * distributed under the License is distributed on an "AS IS" BASIS,
19 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20 * See the License for the specific language governing permissions and
21 * limitations under the License.
22 *
23 ******************************************************************************/
24
25//! Streaming `RecordBatchReader` adapter over a [`Cursor`].
26
27use arrow::array::{RecordBatch, RecordBatchReader};
28use arrow::datatypes::SchemaRef;
29use arrow::error::ArrowError;
30
31use crate::egress::Cursor;
32use crate::egress::arrow::convert::external_arrow_error;
33use crate::error::{Error, ErrorCode};
34
35/// Adapter implementing [`arrow::array::RecordBatchReader`] over a
36/// [`Cursor`]. Snapshots the first batch's Arrow schema at construction
37/// and poisons on mid-stream schema drift. Failover semantics inherit
38/// from [`Cursor::next_batch`](crate::egress::Cursor::next_batch).
39pub struct CursorRecordBatchReader<'r, 'c> {
40 cursor: &'c mut Cursor<'r>,
41 schema: SchemaRef,
42 pending: Option<RecordBatch>,
43 poisoned: bool,
44 /// `Cursor::failover_resets()` at the point `schema` was pinned. A
45 /// transparently-replayed query (re-read from `batch_seq 0` on a new
46 /// endpoint) bumps the cursor counter; the replayed batch 0 is accepted
47 /// without an internal drift check but must still match the pinned
48 /// schema, since a [`RecordBatchReader`]'s schema is fixed for its
49 /// lifetime. `fetch_all_arrow` reads the same counter to discard its
50 /// partial accumulation.
51 resets_at_pin: u32,
52}
53
54impl<'r, 'c> CursorRecordBatchReader<'r, 'c> {
55 pub(crate) fn new(cursor: &'c mut Cursor<'r>) -> Result<Self, Error> {
56 let first = cursor.next_arrow_batch_inner(None, false)?.ok_or_else(|| {
57 Error::new(
58 ErrorCode::NoSchema,
59 "no batch produced; nothing to snapshot",
60 )
61 })?;
62 let schema = first.schema();
63 let resets_at_pin = cursor.failover_resets();
64 Ok(Self {
65 cursor,
66 schema,
67 pending: Some(first),
68 poisoned: false,
69 resets_at_pin,
70 })
71 }
72
73 /// Snapshotted schema. Same as the [`RecordBatchReader::schema`]
74 /// trait method, exposed for callers without the trait imported.
75 pub fn schema(&self) -> SchemaRef {
76 self.schema.clone()
77 }
78
79 /// Reconnect count observed by the underlying cursor. `fetch_all_arrow`
80 /// polls this between batches: an increase means the query was replayed
81 /// from scratch, so anything accumulated so far must be dropped.
82 pub(crate) fn failover_resets(&self) -> u32 {
83 self.cursor.failover_resets()
84 }
85}
86
87impl Iterator for CursorRecordBatchReader<'_, '_> {
88 type Item = Result<RecordBatch, ArrowError>;
89
90 fn next(&mut self) -> Option<Self::Item> {
91 if self.poisoned {
92 return None;
93 }
94 if let Some(rb) = self.pending.take() {
95 return Some(Ok(rb));
96 }
97 // A transparent mid-query failover re-reads the result from
98 // `batch_seq 0` on a new endpoint. Pass `None` (no drift check) for
99 // that first replayed frame so the new node's batch 0 isn't rejected,
100 // then require it to match the pinned schema: a RecordBatchReader's
101 // schema must be stable for its lifetime, so a genuinely different
102 // post-failover schema is surfaced as drift, not silently swapped in.
103 let drift_check = if self.cursor.failover_resets() == self.resets_at_pin {
104 Some(&self.schema)
105 } else {
106 None
107 };
108 match self.cursor.next_arrow_batch_inner(drift_check, false) {
109 Ok(Some(rb)) => {
110 if self.cursor.failover_resets() != self.resets_at_pin {
111 if rb.schema() != self.schema {
112 self.poisoned = true;
113 return Some(Err(external_arrow_error(Error::new(
114 ErrorCode::SchemaDrift,
115 "post-failover replay returned a different schema; \
116 a RecordBatchReader schema must be stable for the \
117 reader's lifetime; use Cursor::next_arrow_batch to \
118 handle drift explicitly",
119 ))));
120 }
121 self.resets_at_pin = self.cursor.failover_resets();
122 } else if has_tentative_array(&self.schema) && rb.schema() != self.schema {
123 self.poisoned = true;
124 return Some(Err(external_arrow_error(Error::new(
125 ErrorCode::SchemaDrift,
126 "tentative→firm ndim upgrade is not representable in \
127 RecordBatchReader (schema must be stable for the \
128 reader's lifetime); use Cursor::next_arrow_batch \
129 to handle drift explicitly",
130 ))));
131 }
132 Some(Ok(rb))
133 }
134 Ok(None) => {
135 self.poisoned = true;
136 None
137 }
138 Err(e) => {
139 self.poisoned = true;
140 Some(Err(external_arrow_error(e)))
141 }
142 }
143 }
144}
145
146/// True if any field carries
147/// [`arrow_metadata::ARRAY_DIM_TENTATIVE`](crate::arrow_metadata::ARRAY_DIM_TENTATIVE).
148/// Gates the tentative→firm ndim mid-stream upgrade.
149pub fn has_tentative_array(schema: &SchemaRef) -> bool {
150 schema.fields().iter().any(|f| {
151 f.metadata()
152 .get(crate::arrow_metadata::ARRAY_DIM_TENTATIVE)
153 .is_some_and(|v| v == "true")
154 })
155}
156
157impl RecordBatchReader for CursorRecordBatchReader<'_, '_> {
158 fn schema(&self) -> SchemaRef {
159 self.schema.clone()
160 }
161}
162
163/// Downcast an [`ArrowError`] produced by this adapter to the
164/// underlying [`Error`]. Returns `None` for foreign Arrow errors.
165pub fn try_downcast_questdb(err: &ArrowError) -> Option<&Error> {
166 match err {
167 ArrowError::ExternalError(boxed) => boxed.downcast_ref::<Error>(),
168 _ => None,
169 }
170}