1use super::{
3 ClaimRelationKind, ClaimRelationRecord, MAX_RESEARCH_ROWS, ResearchAuthority, ResearchCategory,
4 ResearchClaimRecord, ResearchDecisionKind, ResearchDecisionRecord, ResearchSubjectKind,
5 ResearchSuppressionRecord,
6};
7use crate::{
8 KnowledgeError, fixed_32_at, fixed_column, invalid, optional_fixed_16, require_schema,
9 required_i64, required_text, required_u32, string_column, timestamp_column, u32_column,
10 uuid_at, uuid_field,
11};
12use arrow::{
13 array::{
14 Array, ArrayRef, FixedSizeBinaryBuilder, StringArray, TimestampMicrosecondArray,
15 UInt32Array, UInt64Array,
16 },
17 datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit},
18 record_batch::RecordBatch,
19};
20use std::sync::{Arc, LazyLock};
21use uuid::Uuid;
22
23pub const RESEARCH_RECORD_VERSION: u32 = 1;
25
26pub static RESEARCH_CLAIM_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
28 Arc::new(Schema::new(vec![
29 uuid_field("assertion_uuid", false),
30 uuid_field("conceptual_uuid", false),
31 Field::new("category", DataType::Utf8, false),
32 uuid_field("creator_uuid", false),
33 uuid_field("run_uuid", true),
34 uuid_field("origin_branch_uuid", true),
35 uuid_field("origin_version_uuid", true),
36 uuid_field("provenance_uuid", false),
37 Field::new(
38 "recorded_at",
39 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
40 false,
41 ),
42 Field::new("contract_version", DataType::UInt32, false),
43 ]))
44});
45
46pub(super) fn claim_batch(rows: &[ResearchClaimRecord]) -> Result<RecordBatch, KnowledgeError> {
47 RecordBatch::try_new(
48 Arc::clone(&RESEARCH_CLAIM_SCHEMA),
49 vec![
50 uuid_array(rows.iter().map(|row| Some(row.assertion_uuid)))?,
51 uuid_array(rows.iter().map(|row| Some(row.conceptual_uuid)))?,
52 Arc::new(StringArray::from_iter_values(
53 rows.iter().map(|row| row.category.as_str()),
54 )),
55 uuid_array(rows.iter().map(|row| Some(row.creator_uuid)))?,
56 uuid_array(rows.iter().map(|row| row.run_uuid))?,
57 uuid_array(rows.iter().map(|row| row.origin_branch_uuid))?,
58 uuid_array(rows.iter().map(|row| row.origin_version_uuid))?,
59 uuid_array(rows.iter().map(|row| Some(row.provenance_uuid)))?,
60 Arc::new(
61 TimestampMicrosecondArray::from_iter_values(rows.iter().map(|row| row.recorded_at))
62 .with_timezone("UTC"),
63 ),
64 Arc::new(UInt32Array::from(vec![RESEARCH_RECORD_VERSION; rows.len()])),
65 ],
66 )
67 .map_err(Into::into)
68}
69
70pub(super) fn claim_rows(
71 batches: &[RecordBatch],
72) -> Result<Vec<ResearchClaimRecord>, KnowledgeError> {
73 let mut rows = Vec::new();
74 for batch in batches {
75 require_schema(batch, &RESEARCH_CLAIM_SCHEMA, "research.claim.schema")?;
76 bound_batch(rows.len(), batch.num_rows())?;
77 let assertion_uuid = fixed_column(batch, "assertion_uuid")?;
78 let conceptual_uuid = fixed_column(batch, "conceptual_uuid")?;
79 let category = string_column(batch, "category")?;
80 let creator_uuid = fixed_column(batch, "creator_uuid")?;
81 let run_uuid = fixed_column(batch, "run_uuid")?;
82 let origin_branch_uuid = fixed_column(batch, "origin_branch_uuid")?;
83 let origin_version_uuid = fixed_column(batch, "origin_version_uuid")?;
84 let provenance_uuid = fixed_column(batch, "provenance_uuid")?;
85 let recorded_at = timestamp_column(batch, "recorded_at")?;
86 let contract_version = u32_column(batch, "contract_version")?;
87 for row in 0..batch.num_rows() {
88 if required_u32(contract_version, row, "contract_version")? != RESEARCH_RECORD_VERSION {
89 return Err(invalid("research.contract_version", "unsupported version"));
90 }
91 rows.push(ResearchClaimRecord {
92 assertion_uuid: uuid_at(assertion_uuid, row, "assertion_uuid")?,
93 conceptual_uuid: uuid_at(conceptual_uuid, row, "conceptual_uuid")?,
94 category: ResearchCategory::parse(required_text(category, row, "category")?)?,
95 creator_uuid: uuid_at(creator_uuid, row, "creator_uuid")?,
96 run_uuid: optional_fixed_16(run_uuid, row, "run_uuid")?,
97 origin_branch_uuid: optional_fixed_16(
98 origin_branch_uuid,
99 row,
100 "origin_branch_uuid",
101 )?,
102 origin_version_uuid: optional_fixed_16(
103 origin_version_uuid,
104 row,
105 "origin_version_uuid",
106 )?,
107 provenance_uuid: uuid_at(provenance_uuid, row, "provenance_uuid")?,
108 recorded_at: required_i64(recorded_at, row, "recorded_at")?,
109 });
110 }
111 }
112 Ok(rows)
113}
114
115pub static CLAIM_RELATION_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
117 Arc::new(Schema::new(vec![
118 uuid_field("relation_uuid", false),
119 uuid_field("source_assertion_uuid", false),
120 uuid_field("target_assertion_uuid", false),
121 Field::new("kind", DataType::Utf8, false),
122 uuid_field("creator_uuid", false),
123 uuid_field("provenance_uuid", false),
124 Field::new(
125 "recorded_at",
126 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
127 false,
128 ),
129 Field::new("contract_version", DataType::UInt32, false),
130 ]))
131});
132
133pub(super) fn relation_batch(rows: &[ClaimRelationRecord]) -> Result<RecordBatch, KnowledgeError> {
134 RecordBatch::try_new(
135 Arc::clone(&CLAIM_RELATION_SCHEMA),
136 vec![
137 uuid_array(rows.iter().map(|row| Some(row.relation_uuid)))?,
138 uuid_array(rows.iter().map(|row| Some(row.source_assertion_uuid)))?,
139 uuid_array(rows.iter().map(|row| Some(row.target_assertion_uuid)))?,
140 Arc::new(StringArray::from_iter_values(
141 rows.iter().map(|row| row.kind.as_str()),
142 )),
143 uuid_array(rows.iter().map(|row| Some(row.creator_uuid)))?,
144 uuid_array(rows.iter().map(|row| Some(row.provenance_uuid)))?,
145 Arc::new(
146 TimestampMicrosecondArray::from_iter_values(rows.iter().map(|row| row.recorded_at))
147 .with_timezone("UTC"),
148 ),
149 Arc::new(UInt32Array::from(vec![RESEARCH_RECORD_VERSION; rows.len()])),
150 ],
151 )
152 .map_err(Into::into)
153}
154
155pub(super) fn relation_rows(
156 batches: &[RecordBatch],
157) -> Result<Vec<ClaimRelationRecord>, KnowledgeError> {
158 let mut rows = Vec::new();
159 for batch in batches {
160 require_schema(batch, &CLAIM_RELATION_SCHEMA, "research.relation.schema")?;
161 bound_batch(rows.len(), batch.num_rows())?;
162 let relation_uuid = fixed_column(batch, "relation_uuid")?;
163 let source_assertion_uuid = fixed_column(batch, "source_assertion_uuid")?;
164 let target_assertion_uuid = fixed_column(batch, "target_assertion_uuid")?;
165 let kind = string_column(batch, "kind")?;
166 let creator_uuid = fixed_column(batch, "creator_uuid")?;
167 let provenance_uuid = fixed_column(batch, "provenance_uuid")?;
168 let recorded_at = timestamp_column(batch, "recorded_at")?;
169 let contract_version = u32_column(batch, "contract_version")?;
170 for row in 0..batch.num_rows() {
171 if required_u32(contract_version, row, "contract_version")? != RESEARCH_RECORD_VERSION {
172 return Err(invalid("research.contract_version", "unsupported version"));
173 }
174 rows.push(ClaimRelationRecord {
175 relation_uuid: uuid_at(relation_uuid, row, "relation_uuid")?,
176 source_assertion_uuid: uuid_at(
177 source_assertion_uuid,
178 row,
179 "source_assertion_uuid",
180 )?,
181 target_assertion_uuid: uuid_at(
182 target_assertion_uuid,
183 row,
184 "target_assertion_uuid",
185 )?,
186 kind: ClaimRelationKind::parse(required_text(kind, row, "kind")?)?,
187 creator_uuid: uuid_at(creator_uuid, row, "creator_uuid")?,
188 provenance_uuid: uuid_at(provenance_uuid, row, "provenance_uuid")?,
189 recorded_at: required_i64(recorded_at, row, "recorded_at")?,
190 });
191 }
192 }
193 Ok(rows)
194}
195
196pub static RESEARCH_DECISION_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
198 Arc::new(Schema::new(vec![
199 Field::new("sequence", DataType::UInt64, false),
200 uuid_field("decision_uuid", false),
201 uuid_field("operation_uuid", false),
202 Field::new("request_sha256", DataType::FixedSizeBinary(32), false),
203 uuid_field("project_uuid", false),
204 uuid_field("community_uuid", true),
205 uuid_field("context_uuid", false),
206 Field::new("subject_kind", DataType::Utf8, false),
207 uuid_field("subject_uuid", false),
208 Field::new("kind", DataType::Utf8, false),
209 uuid_field("creator_uuid", false),
210 uuid_field("source_version_uuid", true),
211 Field::new(
212 "recorded_at",
213 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
214 false,
215 ),
216 Field::new("contract_version", DataType::UInt32, false),
217 ]))
218});
219
220pub(super) fn decision_batch(
221 rows: &[ResearchDecisionRecord],
222) -> Result<RecordBatch, KnowledgeError> {
223 RecordBatch::try_new(
224 Arc::clone(&RESEARCH_DECISION_SCHEMA),
225 vec![
226 Arc::new(UInt64Array::from_iter_values(
227 rows.iter().map(|row| row.sequence),
228 )),
229 uuid_array(rows.iter().map(|row| Some(row.decision_uuid)))?,
230 uuid_array(rows.iter().map(|row| Some(row.operation_uuid)))?,
231 hash_array(rows.iter().map(|row| row.request_sha256))?,
232 uuid_array(rows.iter().map(|row| Some(row.authority.project_uuid)))?,
233 uuid_array(rows.iter().map(|row| row.authority.community_uuid))?,
234 uuid_array(rows.iter().map(|row| Some(row.authority.context_uuid)))?,
235 Arc::new(StringArray::from_iter_values(
236 rows.iter().map(|row| row.subject_kind.as_str()),
237 )),
238 uuid_array(rows.iter().map(|row| Some(row.subject_uuid)))?,
239 Arc::new(StringArray::from_iter_values(
240 rows.iter().map(|row| row.kind.as_str()),
241 )),
242 uuid_array(rows.iter().map(|row| Some(row.creator_uuid)))?,
243 uuid_array(rows.iter().map(|row| row.source_version_uuid))?,
244 Arc::new(
245 TimestampMicrosecondArray::from_iter_values(rows.iter().map(|row| row.recorded_at))
246 .with_timezone("UTC"),
247 ),
248 Arc::new(UInt32Array::from(vec![RESEARCH_RECORD_VERSION; rows.len()])),
249 ],
250 )
251 .map_err(Into::into)
252}
253
254pub(super) fn decision_rows(
255 batches: &[RecordBatch],
256) -> Result<Vec<ResearchDecisionRecord>, KnowledgeError> {
257 let mut rows = Vec::new();
258 for batch in batches {
259 require_schema(batch, &RESEARCH_DECISION_SCHEMA, "research.decision.schema")?;
260 bound_batch(rows.len(), batch.num_rows())?;
261 let sequence = u64_column(batch, "sequence")?;
262 let decision_uuid = fixed_column(batch, "decision_uuid")?;
263 let operation_uuid = fixed_column(batch, "operation_uuid")?;
264 let request_sha256 = fixed_column(batch, "request_sha256")?;
265 let project_uuid = fixed_column(batch, "project_uuid")?;
266 let community_uuid = fixed_column(batch, "community_uuid")?;
267 let context_uuid = fixed_column(batch, "context_uuid")?;
268 let subject_kind = string_column(batch, "subject_kind")?;
269 let subject_uuid = fixed_column(batch, "subject_uuid")?;
270 let kind = string_column(batch, "kind")?;
271 let creator_uuid = fixed_column(batch, "creator_uuid")?;
272 let source_version_uuid = fixed_column(batch, "source_version_uuid")?;
273 let recorded_at = timestamp_column(batch, "recorded_at")?;
274 let contract_version = u32_column(batch, "contract_version")?;
275 for row in 0..batch.num_rows() {
276 if required_u32(contract_version, row, "contract_version")? != RESEARCH_RECORD_VERSION {
277 return Err(invalid("research.contract_version", "unsupported version"));
278 }
279 rows.push(ResearchDecisionRecord {
280 authority: ResearchAuthority {
281 project_uuid: uuid_at(project_uuid, row, "project_uuid")?,
282 community_uuid: optional_fixed_16(community_uuid, row, "community_uuid")?,
283 context_uuid: uuid_at(context_uuid, row, "context_uuid")?,
284 },
285 sequence: required_u64(sequence, row, "sequence")?,
286 decision_uuid: uuid_at(decision_uuid, row, "decision_uuid")?,
287 operation_uuid: uuid_at(operation_uuid, row, "operation_uuid")?,
288 request_sha256: fixed_32_at(request_sha256, row, "request_sha256")?,
289 subject_kind: ResearchSubjectKind::parse(required_text(
290 subject_kind,
291 row,
292 "subject_kind",
293 )?)?,
294 subject_uuid: uuid_at(subject_uuid, row, "subject_uuid")?,
295 kind: ResearchDecisionKind::parse(required_text(kind, row, "kind")?)?,
296 creator_uuid: uuid_at(creator_uuid, row, "creator_uuid")?,
297 source_version_uuid: optional_fixed_16(
298 source_version_uuid,
299 row,
300 "source_version_uuid",
301 )?,
302 recorded_at: required_i64(recorded_at, row, "recorded_at")?,
303 });
304 }
305 }
306 Ok(rows)
307}
308
309pub static CLAIM_SUPPRESSION_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
311 Arc::new(Schema::new(vec![
312 uuid_field("suppression_uuid", false),
313 uuid_field("assertion_uuid", false),
314 uuid_field("context_uuid", false),
315 uuid_field("creator_uuid", false),
316 uuid_field("provenance_uuid", false),
317 Field::new(
318 "recorded_at",
319 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
320 false,
321 ),
322 Field::new("contract_version", DataType::UInt32, false),
323 ]))
324});
325
326pub(super) fn suppression_batch(
327 rows: &[ResearchSuppressionRecord],
328) -> Result<RecordBatch, KnowledgeError> {
329 RecordBatch::try_new(
330 Arc::clone(&CLAIM_SUPPRESSION_SCHEMA),
331 vec![
332 uuid_array(rows.iter().map(|row| Some(row.suppression_uuid)))?,
333 uuid_array(rows.iter().map(|row| Some(row.assertion_uuid)))?,
334 uuid_array(rows.iter().map(|row| Some(row.context_uuid)))?,
335 uuid_array(rows.iter().map(|row| Some(row.creator_uuid)))?,
336 uuid_array(rows.iter().map(|row| Some(row.provenance_uuid)))?,
337 Arc::new(
338 TimestampMicrosecondArray::from_iter_values(rows.iter().map(|row| row.recorded_at))
339 .with_timezone("UTC"),
340 ),
341 Arc::new(UInt32Array::from(vec![RESEARCH_RECORD_VERSION; rows.len()])),
342 ],
343 )
344 .map_err(Into::into)
345}
346
347pub(super) fn suppression_rows(
348 batches: &[RecordBatch],
349) -> Result<Vec<ResearchSuppressionRecord>, KnowledgeError> {
350 let mut rows = Vec::new();
351 for batch in batches {
352 require_schema(
353 batch,
354 &CLAIM_SUPPRESSION_SCHEMA,
355 "research.suppression.schema",
356 )?;
357 bound_batch(rows.len(), batch.num_rows())?;
358 let suppression_uuid = fixed_column(batch, "suppression_uuid")?;
359 let assertion_uuid = fixed_column(batch, "assertion_uuid")?;
360 let context_uuid = fixed_column(batch, "context_uuid")?;
361 let creator_uuid = fixed_column(batch, "creator_uuid")?;
362 let provenance_uuid = fixed_column(batch, "provenance_uuid")?;
363 let recorded_at = timestamp_column(batch, "recorded_at")?;
364 let contract_version = u32_column(batch, "contract_version")?;
365 for row in 0..batch.num_rows() {
366 if required_u32(contract_version, row, "contract_version")? != RESEARCH_RECORD_VERSION {
367 return Err(invalid("research.contract_version", "unsupported version"));
368 }
369 rows.push(ResearchSuppressionRecord {
370 suppression_uuid: uuid_at(suppression_uuid, row, "suppression_uuid")?,
371 assertion_uuid: uuid_at(assertion_uuid, row, "assertion_uuid")?,
372 context_uuid: uuid_at(context_uuid, row, "context_uuid")?,
373 creator_uuid: uuid_at(creator_uuid, row, "creator_uuid")?,
374 provenance_uuid: uuid_at(provenance_uuid, row, "provenance_uuid")?,
375 recorded_at: required_i64(recorded_at, row, "recorded_at")?,
376 });
377 }
378 }
379 Ok(rows)
380}
381
382fn uuid_array(rows: impl Iterator<Item = Option<Uuid>>) -> Result<ArrayRef, KnowledgeError> {
383 let mut builder = FixedSizeBinaryBuilder::new(16);
384 for row in rows {
385 match row {
386 Some(id) => builder.append_value(id.as_bytes())?,
387 None => builder.append_null(),
388 }
389 }
390 Ok(Arc::new(builder.finish()))
391}
392fn hash_array(rows: impl Iterator<Item = [u8; 32]>) -> Result<ArrayRef, KnowledgeError> {
393 let mut builder = FixedSizeBinaryBuilder::new(32);
394 for row in rows {
395 builder.append_value(row)?;
396 }
397 Ok(Arc::new(builder.finish()))
398}
399fn bound_batch(existing: usize, incoming: usize) -> Result<(), KnowledgeError> {
400 let observed = existing.saturating_add(incoming);
401 if observed > MAX_RESEARCH_ROWS {
402 return Err(KnowledgeError::Limit {
403 participant: "research",
404 observed,
405 limit: MAX_RESEARCH_ROWS,
406 });
407 }
408 Ok(())
409}
410fn u64_column<'a>(
411 batch: &'a RecordBatch,
412 name: &'static str,
413) -> Result<&'a UInt64Array, KnowledgeError> {
414 batch
415 .column_by_name(name)
416 .and_then(|column| column.as_any().downcast_ref())
417 .ok_or_else(|| invalid(name, "missing or wrong Arrow type"))
418}
419fn required_u64(
420 array: &UInt64Array,
421 row: usize,
422 field: &'static str,
423) -> Result<u64, KnowledgeError> {
424 if array.is_null(row) {
425 return Err(invalid(field, "must not be null"));
426 }
427 Ok(array.value(row))
428}