1use std::collections::HashMap;
11
12use serde::{Deserialize, Serialize};
13use serde_json::{Map, Value};
14
15use crate::error::Error;
16use crate::spec::{Column, Data, FieldType, Spec};
17
18pub type Row = Map<String, Value>;
19
20#[derive(Debug, Deserialize)]
23#[serde(untagged)]
24pub enum DataDoc {
25 Envelope {
26 columns: Vec<EnvColumn>,
27 rows: Vec<Vec<Value>>,
28 #[serde(default)]
29 truncated: Option<bool>,
30 #[serde(default)]
31 total_rows: Option<u64>,
32 },
33 Rows(Vec<Row>),
34}
35
36#[derive(Debug, Deserialize)]
38pub struct EnvColumn {
39 pub name: String,
40 #[serde(default, rename = "type")]
41 pub ty: Option<String>,
42}
43
44impl From<&Column> for EnvColumn {
45 fn from(c: &Column) -> Self {
46 EnvColumn {
47 name: c.name.clone(),
48 ty: c.ty.clone(),
49 }
50 }
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
54#[serde(rename_all = "snake_case")]
55pub enum DataSource {
56 InlineValues,
57 InlineColumns,
58 StdinValues,
59 StdinColumns,
60}
61
62#[derive(Debug, Serialize)]
63pub struct DataProvenance {
64 pub source: DataSource,
65 pub truncated: Option<bool>,
66 pub total_rows: Option<u64>,
67}
68
69#[derive(Debug)]
72pub struct Table {
73 pub rows: Vec<Row>,
74 pub declared: HashMap<String, FieldType>,
75 pub provenance: DataProvenance,
76}
77
78pub fn parse_data_doc(s: &str) -> Result<DataDoc, Error> {
81 serde_json::from_str(s).map_err(|e| {
82 Error::Data(format!(
83 "cannot parse stdin as a data document; expected \
84 {{\"columns\":[{{\"name\",\"type\"?}}...],\"rows\":[[...]...]}} or a JSON array \
85 of row objects: {e}"
86 ))
87 })
88}
89
90pub fn declared_field_type(t: &str) -> FieldType {
106 match t.to_ascii_uppercase().as_str() {
107 "INT64" | "INTEGER" | "INT" | "SMALLINT" | "BIGINT" | "FLOAT64" | "FLOAT" | "DOUBLE"
108 | "NUMERIC" | "BIGNUMERIC" | "DECIMAL" | "REAL" | "NUMBER" => FieldType::Quantitative,
109 "DATE" | "DATETIME" | "TIMESTAMP" | "TIME" => FieldType::Temporal,
110 "BOOLEAN" | "BOOL" => FieldType::Nominal,
111 _ => FieldType::Nominal,
112 }
113}
114
115pub fn resolve(spec: &Spec, stdin: Option<DataDoc>) -> Result<Table, Error> {
118 match (&spec.data, stdin) {
119 (Some(_), Some(_)) => Err(Error::Spec(
120 "data provided twice: the spec has inline `data` and a data document \
121 arrived on stdin; remove one"
122 .into(),
123 )),
124 (Some(data), None) => resolve_inline(data),
125 (None, Some(doc)) => resolve_stdin(doc),
126 (None, None) => Err(Error::Spec(
127 "no data: the spec has no `data` and nothing arrived on stdin; add \
128 data.values or data.columns+rows to the spec, or pipe a data document"
129 .into(),
130 )),
131 }
132}
133
134fn resolve_inline(data: &Data) -> Result<Table, Error> {
139 match (&data.values, &data.columns, &data.rows) {
140 (Some(values), None, None) => finish(
142 values.clone(),
143 HashMap::new(),
144 DataSource::InlineValues,
145 None,
146 None,
147 ),
148 (None, Some(columns), Some(rows)) => {
149 let env: Vec<EnvColumn> = columns.iter().map(EnvColumn::from).collect();
152 let (rows, declared) = columnar_to_rows(&env, rows)?;
153 finish(rows, declared, DataSource::InlineColumns, None, None)
154 }
155 _ => Err(Error::Spec(
156 "data must contain either `values`, or `columns` and `rows`".into(),
157 )),
158 }
159}
160
161fn resolve_stdin(doc: DataDoc) -> Result<Table, Error> {
163 match doc {
164 DataDoc::Envelope {
165 columns,
166 rows,
167 truncated,
168 total_rows,
169 } => {
170 let (rows, declared) = columnar_to_rows(&columns, &rows)?;
171 finish(
172 rows,
173 declared,
174 DataSource::StdinColumns,
175 truncated,
176 total_rows,
177 )
178 }
179 DataDoc::Rows(rows) => finish(rows, HashMap::new(), DataSource::StdinValues, None, None),
180 }
181}
182
183fn columnar_to_rows(
186 columns: &[EnvColumn],
187 rows: &[Vec<Value>],
188) -> Result<(Vec<Row>, HashMap<String, FieldType>), Error> {
189 let mut declared = HashMap::new();
190 let mut seen = std::collections::HashSet::new();
191 for col in columns {
192 if !seen.insert(col.name.as_str()) {
193 return Err(Error::Data(format!(
194 "duplicate column name \"{}\"",
195 col.name
196 )));
197 }
198 if let Some(ty) = &col.ty {
199 declared.insert(col.name.clone(), declared_field_type(ty));
200 }
201 }
202
203 let want = columns.len();
204 let mut out = Vec::with_capacity(rows.len());
205 for (i, row) in rows.iter().enumerate() {
206 if row.len() != want {
207 return Err(Error::Data(format!(
208 "row {i} has {} values but {want} columns are declared",
209 row.len()
210 )));
211 }
212 let mut obj = Row::new();
213 for (col, val) in columns.iter().zip(row) {
214 obj.insert(col.name.clone(), val.clone());
215 }
216 out.push(obj);
217 }
218 Ok((out, declared))
219}
220
221fn finish(
223 rows: Vec<Row>,
224 declared: HashMap<String, FieldType>,
225 source: DataSource,
226 truncated: Option<bool>,
227 total_rows: Option<u64>,
228) -> Result<Table, Error> {
229 if rows.is_empty() {
230 return Err(Error::Data(
231 "data has no rows; provide at least one row".into(),
232 ));
233 }
234 Ok(Table {
235 rows,
236 declared,
237 provenance: DataProvenance {
238 source,
239 truncated,
240 total_rows,
241 },
242 })
243}
244
245#[cfg(test)]
246mod tests {
247 use super::*;
248 use serde_json::json;
249
250 fn spec_with_values() -> Spec {
251 serde_json::from_str(
252 r#"{"data":{"values":[{"x":1}]},"mark":"bar",
253 "encoding":{"x":{"field":"x"},"y":{"field":"y"}}}"#,
254 )
255 .expect("fixture spec parses")
256 }
257
258 fn spec_with_empty_values() -> Spec {
259 serde_json::from_str(
260 r#"{"data":{"values":[]},"mark":"bar",
261 "encoding":{"x":{"field":"x"},"y":{"field":"y"}}}"#,
262 )
263 .expect("fixture spec parses")
264 }
265
266 #[test]
267 fn parses_bare_row_array() {
268 let doc = parse_data_doc(r#"[{"a": 1}, {"a": 2}]"#).expect("bare array parses");
269 match doc {
270 DataDoc::Rows(rows) => assert_eq!(rows.len(), 2),
271 other => panic!("expected Rows, got {other:?}"),
272 }
273 }
274
275 #[test]
276 fn envelope_tolerates_unknown_keys() {
277 let doc = parse_data_doc(
279 r#"{"columns":[{"name":"day","type":"STRING"},{"name":"n","type":"INT64"}],
280 "rows":[["mon",32]],
281 "query":{"job_id":"abc","note":"ignored"}}"#,
282 )
283 .expect("envelope with extra keys parses");
284 match doc {
285 DataDoc::Envelope { columns, rows, .. } => {
286 assert_eq!(columns.len(), 2);
287 assert_eq!(rows.len(), 1);
288 }
289 other => panic!("expected Envelope, got {other:?}"),
290 }
291 }
292
293 #[test]
294 fn envelope_truncated_and_total_rows_reach_provenance() {
295 let doc = parse_data_doc(
296 r#"{"columns":[{"name":"n"}],"rows":[[1]],"truncated":true,"total_rows":123}"#,
297 )
298 .expect("envelope parses");
299 let table = resolve_stdin(doc).expect("resolves");
300 assert_eq!(table.provenance.source, DataSource::StdinColumns);
301 assert_eq!(table.provenance.truncated, Some(true));
302 assert_eq!(table.provenance.total_rows, Some(123));
303 }
304
305 #[test]
306 fn bare_rows_provenance_has_no_envelope_fields() {
307 let doc = parse_data_doc(r#"[{"a":1}]"#).expect("parses");
308 let table = resolve_stdin(doc).expect("resolves");
309 assert_eq!(table.provenance.source, DataSource::StdinValues);
310 assert_eq!(table.provenance.truncated, None);
311 assert_eq!(table.provenance.total_rows, None);
312 assert!(table.declared.is_empty());
313 }
314
315 #[test]
316 fn columnar_zips_rows_in_column_order() {
317 let doc = parse_data_doc(
318 r#"{"columns":[{"name":"day","type":"STRING"},{"name":"n","type":"INT64"}],
319 "rows":[["mon",32],["tue",78]]}"#,
320 )
321 .expect("parses");
322 let table = resolve_stdin(doc).expect("resolves");
323 assert_eq!(table.rows.len(), 2);
324 assert_eq!(table.rows[0].get("day"), Some(&json!("mon")));
325 assert_eq!(table.rows[0].get("n"), Some(&json!(32)));
326 assert_eq!(table.rows[1].get("day"), Some(&json!("tue")));
327 assert_eq!(table.rows[1].get("n"), Some(&json!(78)));
328 let keys: Vec<&String> = table.rows[0].keys().collect();
330 assert_eq!(keys, vec!["day", "n"]);
331 assert_eq!(table.declared.get("day"), Some(&FieldType::Nominal));
332 assert_eq!(table.declared.get("n"), Some(&FieldType::Quantitative));
333 }
334
335 #[test]
336 fn duplicate_column_name_errors() {
337 let doc = parse_data_doc(r#"{"columns":[{"name":"a"},{"name":"a"}],"rows":[[1,2]]}"#)
338 .expect("parses");
339 let err = resolve_stdin(doc).expect_err("duplicate must error");
340 insta::assert_snapshot!(err.to_string(), @r###"duplicate column name "a""###);
341 }
342
343 #[test]
344 fn row_length_mismatch_errors() {
345 let doc =
346 parse_data_doc(r#"{"columns":[{"name":"a"},{"name":"b"}],"rows":[[1,2],[1,2,3]]}"#)
347 .expect("parses");
348 let err = resolve_stdin(doc).expect_err("length mismatch must error");
349 insta::assert_snapshot!(
350 err.to_string(),
351 @"row 1 has 3 values but 2 columns are declared"
352 );
353 }
354
355 #[test]
356 fn empty_rows_errors_columnar() {
357 let doc = parse_data_doc(r#"{"columns":[{"name":"a"}],"rows":[]}"#).expect("parses");
358 let err = resolve_stdin(doc).expect_err("empty rows must error");
359 insta::assert_snapshot!(err.to_string(), @"data has no rows; provide at least one row");
360 }
361
362 #[test]
363 fn empty_rows_errors_bare_array() {
364 let doc = parse_data_doc(r#"[]"#).expect("parses");
365 let err = resolve_stdin(doc).expect_err("empty rows must error");
366 insta::assert_snapshot!(err.to_string(), @"data has no rows; provide at least one row");
367 }
368
369 #[test]
370 fn empty_rows_errors_inline_values() {
371 let err = resolve(&spec_with_empty_values(), None).expect_err("empty inline must error");
372 insta::assert_snapshot!(err.to_string(), @"data has no rows; provide at least one row");
373 }
374
375 #[test]
376 fn inline_values_resolve_to_inline_provenance() {
377 let table = resolve(&spec_with_values(), None).expect("resolves");
378 assert_eq!(table.provenance.source, DataSource::InlineValues);
379 assert_eq!(table.provenance.truncated, None);
380 assert_eq!(table.provenance.total_rows, None);
381 assert!(table.declared.is_empty());
382 assert_eq!(table.rows.len(), 1);
383 }
384
385 #[test]
386 fn data_provided_twice_errors() {
387 let doc = parse_data_doc(r#"[{"a":1}]"#).expect("parses");
388 let err = resolve(&spec_with_values(), Some(doc)).expect_err("data twice must error");
389 insta::assert_snapshot!(
390 err.to_string(),
391 @"data provided twice: the spec has inline `data` and a data document arrived on stdin; remove one"
392 );
393 }
394
395 #[test]
396 fn inline_columnar_resolves_with_declared_types() {
397 let spec: Spec = serde_json::from_str(
398 r#"{"data":{"columns":[{"name":"day","type":"STRING"},{"name":"n","type":"INT64"}],
399 "rows":[["mon",32],["tue",78]]},
400 "mark":"bar","encoding":{"x":{"field":"day"},"y":{"field":"n"}}}"#,
401 )
402 .expect("inline columnar spec parses");
403 let table = resolve(&spec, None).expect("resolves");
404 assert_eq!(table.provenance.source, DataSource::InlineColumns);
405 assert_eq!(table.rows.len(), 2);
406 assert_eq!(table.rows[0].get("day"), Some(&json!("mon")));
407 assert_eq!(table.rows[0].get("n"), Some(&json!(32)));
408 assert_eq!(table.declared.get("day"), Some(&FieldType::Nominal));
409 assert_eq!(table.declared.get("n"), Some(&FieldType::Quantitative));
410 }
411
412 #[test]
413 fn inline_data_both_forms_errors() {
414 let spec: Spec = serde_json::from_str(
415 r#"{"data":{"values":[{"a":1}],"columns":[{"name":"a"}],"rows":[[1]]},
416 "mark":"bar","encoding":{"x":{"field":"a"},"y":{"field":"a"}}}"#,
417 )
418 .expect("spec parses");
419 let err = resolve(&spec, None).expect_err("both forms must error");
420 insta::assert_snapshot!(
421 err.to_string(),
422 @"data must contain either `values`, or `columns` and `rows`"
423 );
424 }
425
426 #[test]
427 fn no_data_anywhere_errors() {
428 let spec: Spec = serde_json::from_str(
429 r#"{"mark":"bar","encoding":{"x":{"field":"a"},"y":{"field":"b"}}}"#,
430 )
431 .expect("spec without data parses");
432 let err = resolve(&spec, None).expect_err("no data must error");
433 insta::assert_snapshot!(
434 err.to_string(),
435 @"no data: the spec has no `data` and nothing arrived on stdin; add data.values or data.columns+rows to the spec, or pipe a data document"
436 );
437 }
438
439 #[test]
440 fn declared_field_type_mapping() {
441 for t in [
443 "INT64",
444 "INTEGER",
445 "INT",
446 "SMALLINT",
447 "BIGINT",
448 "FLOAT64",
449 "FLOAT",
450 "DOUBLE",
451 "NUMERIC",
452 "BIGNUMERIC",
453 "DECIMAL",
454 "REAL",
455 "NUMBER",
458 ] {
459 assert_eq!(declared_field_type(t), FieldType::Quantitative, "{t}");
460 }
461 for t in ["DATE", "DATETIME", "TIMESTAMP", "TIME"] {
463 assert_eq!(declared_field_type(t), FieldType::Temporal, "{t}");
464 }
465 assert_eq!(declared_field_type("STRING"), FieldType::Nominal);
468 assert_eq!(declared_field_type("BOOL"), FieldType::Nominal);
469 assert_eq!(declared_field_type("BOOLEAN"), FieldType::Nominal);
470 assert_eq!(declared_field_type("boolean"), FieldType::Nominal);
471 assert_eq!(declared_field_type("number"), FieldType::Quantitative);
472 assert_eq!(declared_field_type("whatever"), FieldType::Nominal);
473 assert_eq!(declared_field_type("int64"), FieldType::Quantitative);
475 assert_eq!(declared_field_type("Date"), FieldType::Temporal);
476 assert_eq!(declared_field_type("String"), FieldType::Nominal);
477 }
478}