Skip to content

Commit 4a6a7ea

Browse files
AdamGSAdamGS
committed
Less code
Signed-off-by: Adam Gutglick <adam@spiraldb.com>
1 parent bff3860 commit 4a6a7ea

2 files changed

Lines changed: 40 additions & 206 deletions

File tree

‎encodings/parquet-variant/src/fns/json_to_variant.rs‎

Lines changed: 9 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,10 @@ impl ScalarFnVTable for JsonToVariant {
134134
fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
135135
let input_dtype = &arg_dtypes[0];
136136
vortex_ensure!(
137-
is_json_string_dtype(input_dtype),
137+
input_dtype.is_utf8()
138+
|| input_dtype
139+
.as_extension_opt()
140+
.is_some_and(|ext_dtype| ext_dtype.is::<Json>()),
138141
"JsonToVariant input must be Utf8 or a Json extension, found {input_dtype}"
139142
);
140143

@@ -152,7 +155,11 @@ impl ScalarFnVTable for JsonToVariant {
152155

153156
let storage = if input.dtype().is_utf8() {
154157
input
155-
} else if is_json_string_dtype(input.dtype()) {
158+
} else if input
159+
.dtype()
160+
.as_extension_opt()
161+
.is_some_and(|ext_dtype| ext_dtype.is::<Json>())
162+
{
156163
input
157164
.execute::<ExtensionArray>(ctx)?
158165
.storage_array()
@@ -188,15 +195,6 @@ impl ScalarFnVTable for JsonToVariant {
188195
}
189196
}
190197

191-
/// Returns whether `dtype` is acceptable input for [`JsonToVariant`]: `Utf8` or the [`Json`]
192-
/// extension dtype.
193-
fn is_json_string_dtype(dtype: &DType) -> bool {
194-
dtype.is_utf8()
195-
|| dtype
196-
.as_extension_opt()
197-
.is_some_and(|ext_dtype| ext_dtype.is::<Json>())
198-
}
199-
200198
/// A list of `(path, dtype)` directives describing which Variant paths to shred and as what
201199
/// type.
202200
///
@@ -301,7 +299,6 @@ mod tests {
301299
use vortex_array::assert_nth_scalar_is_null;
302300
use vortex_array::dtype::Nullability;
303301
use vortex_array::dtype::PType;
304-
use vortex_array::dtype::extension::ExtDType;
305302
use vortex_array::expr::Expression;
306303
use vortex_array::expr::proto::ExprSerializeProtoExt;
307304
use vortex_array::expr::root;
@@ -357,33 +354,6 @@ mod tests {
357354
Ok(())
358355
}
359356

360-
#[test]
361-
fn options_roundtrip_serialization() -> VortexResult<()> {
362-
let specs = [
363-
ShreddingSpec::empty(),
364-
ShreddingSpec::try_new([
365-
(VariantPath::field("a"), i64_dtype()),
366-
(
367-
VariantPath::new([
368-
VariantPathElement::field("b"),
369-
VariantPathElement::field("c"),
370-
]),
371-
DType::Utf8(Nullability::NonNullable),
372-
),
373-
])?,
374-
];
375-
376-
for spec in specs {
377-
let options = JsonToVariantOptions::new(spec);
378-
let metadata = JsonToVariant
379-
.serialize(&options)?
380-
.ok_or_else(|| vortex_err!("expected metadata"))?;
381-
let actual = JsonToVariant.deserialize(&metadata, &VortexSession::empty())?;
382-
assert_eq!(actual, options);
383-
}
384-
Ok(())
385-
}
386-
387357
#[test]
388358
fn expression_roundtrip_serialization() -> VortexResult<()> {
389359
let expr: Expression = json_to_variant(root(), shred_field_as_i64("a")?);
@@ -394,64 +364,6 @@ mod tests {
394364
Ok(())
395365
}
396366

397-
#[test]
398-
fn options_display() -> VortexResult<()> {
399-
assert_eq!(JsonToVariantOptions::unshredded().to_string(), "");
400-
assert_eq!(
401-
JsonToVariantOptions::new(ShreddingSpec::try_new([
402-
(VariantPath::field("a"), i64_dtype()),
403-
(
404-
VariantPath::field("b"),
405-
DType::Utf8(Nullability::NonNullable)
406-
),
407-
])?)
408-
.to_string(),
409-
"$.a as i64?, $.b as utf8"
410-
);
411-
Ok(())
412-
}
413-
414-
#[test]
415-
fn return_dtype_propagates_nullability() {
416-
let expr = json_to_variant(root(), ShreddingSpec::empty());
417-
assert_eq!(
418-
expr.return_dtype(&DType::Utf8(Nullability::NonNullable))
419-
.unwrap(),
420-
DType::Variant(Nullability::NonNullable)
421-
);
422-
assert_eq!(
423-
expr.return_dtype(&DType::Utf8(Nullability::Nullable))
424-
.unwrap(),
425-
DType::Variant(Nullability::Nullable)
426-
);
427-
}
428-
429-
#[test]
430-
fn return_dtype_accepts_json_extension() -> VortexResult<()> {
431-
let json_dtype = DType::Extension(
432-
ExtDType::<Json>::try_new(EmptyMetadata, DType::Utf8(Nullability::Nullable))?.erased(),
433-
);
434-
let expr = json_to_variant(root(), ShreddingSpec::empty());
435-
assert_eq!(
436-
expr.return_dtype(&json_dtype)?,
437-
DType::Variant(Nullability::Nullable)
438-
);
439-
Ok(())
440-
}
441-
442-
#[test]
443-
fn return_dtype_rejects_non_string_input() {
444-
let expr = json_to_variant(root(), ShreddingSpec::empty());
445-
let err = expr
446-
.return_dtype(&DType::Bool(Nullability::NonNullable))
447-
.unwrap_err();
448-
assert!(
449-
err.to_string()
450-
.contains("JsonToVariant input must be Utf8 or a Json extension"),
451-
"unexpected error: {err}"
452-
);
453-
}
454-
455367
#[test]
456368
fn converts_utf8_json_rows() -> VortexResult<()> {
457369
let input =

‎encodings/parquet-variant/src/fns/variant_to_json.rs‎

Lines changed: 31 additions & 109 deletions
Original file line numberDiff line numberDiff line change
@@ -173,13 +173,10 @@ mod tests {
173173
use vortex_array::arrays::PrimitiveArray;
174174
use vortex_array::arrays::VarBinViewArray;
175175
use vortex_array::arrays::Variant;
176-
use vortex_array::arrays::VariantArray;
177176
use vortex_array::arrays::extension::ExtensionArrayExt;
178177
use vortex_array::arrays::variant::VariantArrayExt;
179178
use vortex_array::dtype::Nullability;
180179
use vortex_array::dtype::PType;
181-
use vortex_array::expr::Expression;
182-
use vortex_array::expr::proto::ExprSerializeProtoExt;
183180
use vortex_array::expr::root;
184181
use vortex_array::scalar_fn::fns::variant_get::VariantPath;
185182
use vortex_array::scalar_fn::fns::variant_get::VariantPathElement;
@@ -205,12 +202,6 @@ mod tests {
205202
))
206203
}
207204

208-
fn execute_variant_to_json(input: ArrayRef) -> VortexResult<ArrayRef> {
209-
input
210-
.apply(&variant_to_json(root()))?
211-
.execute::<ArrayRef>(&mut SESSION.create_execution_ctx())
212-
}
213-
214205
fn json_strings(array: &ArrayRef) -> VortexResult<Vec<Option<String>>> {
215206
let mut ctx = SESSION.create_execution_ctx();
216207
let ext = array.clone().execute::<ExtensionArray>(&mut ctx)?;
@@ -245,55 +236,6 @@ mod tests {
245236
.execute::<ArrayRef>(&mut SESSION.create_execution_ctx())
246237
}
247238

248-
#[test]
249-
fn expression_roundtrip_serialization() -> VortexResult<()> {
250-
let expr: Expression = variant_to_json(root());
251-
let proto = expr.serialize_proto()?;
252-
let actual = Expression::from_proto(&proto, &SESSION)?;
253-
254-
assert_eq!(actual, expr);
255-
Ok(())
256-
}
257-
258-
#[test]
259-
fn return_dtype_is_json_extension() -> VortexResult<()> {
260-
let expr = variant_to_json(root());
261-
assert_eq!(
262-
expr.return_dtype(&DType::Variant(Nullability::NonNullable))?,
263-
json_dtype(Nullability::NonNullable)?
264-
);
265-
assert_eq!(
266-
expr.return_dtype(&DType::Variant(Nullability::Nullable))?,
267-
json_dtype(Nullability::Nullable)?
268-
);
269-
Ok(())
270-
}
271-
272-
#[test]
273-
fn return_dtype_rejects_non_variant_input() {
274-
let expr = variant_to_json(root());
275-
let err = expr
276-
.return_dtype(&DType::Utf8(Nullability::NonNullable))
277-
.unwrap_err();
278-
assert!(
279-
err.to_string()
280-
.contains("VariantToJson input must be Variant"),
281-
"unexpected error: {err}"
282-
);
283-
}
284-
285-
#[test]
286-
fn formats_sql() {
287-
let expr = variant_to_json(root());
288-
assert_eq!(expr.to_string(), "variant_to_json($)");
289-
}
290-
291-
#[test]
292-
fn is_fallible() {
293-
let expr = variant_to_json(root());
294-
assert!(expr.signature().is_fallible());
295-
}
296-
297239
#[test]
298240
fn renders_unshredded_values() -> VortexResult<()> {
299241
let input = unshredded_variant([
@@ -303,7 +245,9 @@ mod tests {
303245
PqVariant::Null,
304246
])?;
305247

306-
let result = execute_variant_to_json(input)?;
248+
let result = input
249+
.apply(&variant_to_json(root()))?
250+
.execute::<ArrayRef>(&mut SESSION.create_execution_ctx())?;
307251

308252
assert_eq!(result.dtype(), &json_dtype(Nullability::NonNullable)?);
309253
assert_eq!(
@@ -323,7 +267,9 @@ mod tests {
323267
let input =
324268
json_rows_to_variant(vec![Some("1"), None, Some("null")], ShreddingSpec::empty())?;
325269

326-
let result = execute_variant_to_json(input)?;
270+
let result = input
271+
.apply(&variant_to_json(root()))?
272+
.execute::<ArrayRef>(&mut SESSION.create_execution_ctx())?;
327273

328274
assert_eq!(result.dtype(), &json_dtype(Nullability::Nullable)?);
329275
assert_eq!(
@@ -351,7 +297,12 @@ mod tests {
351297

352298
#[test]
353299
fn unshreds_typed_value_only_storage() -> VortexResult<()> {
354-
let result = execute_variant_to_json(typed_value_only_variant()?)?;
300+
let result = {
301+
let input = typed_value_only_variant()?;
302+
input
303+
.apply(&variant_to_json(root()))?
304+
.execute::<ArrayRef>(&mut SESSION.create_execution_ctx())
305+
}?;
355306

356307
assert_eq!(
357308
json_strings(&result)?,
@@ -383,7 +334,9 @@ mod tests {
383334
"fixture must be shredded"
384335
);
385336

386-
let result = execute_variant_to_json(input)?;
337+
let result = input
338+
.apply(&variant_to_json(root()))?
339+
.execute::<ArrayRef>(&mut SESSION.create_execution_ctx())?;
387340

388341
assert_eq!(
389342
json_strings(&result)?,
@@ -396,30 +349,6 @@ mod tests {
396349
Ok(())
397350
}
398351

399-
#[test]
400-
fn renders_canonical_variant_with_shredded_child() -> VortexResult<()> {
401-
let mut ctx = SESSION.create_execution_ctx();
402-
let canonical = typed_value_only_variant()?
403-
.execute::<VariantArray>(&mut ctx)?
404-
.into_array();
405-
assert!(
406-
canonical.as_::<Variant>().shredded().is_some(),
407-
"fixture must carry a canonical shredded child"
408-
);
409-
410-
let result = execute_variant_to_json(canonical)?;
411-
412-
assert_eq!(
413-
json_strings(&result)?,
414-
vec![
415-
Some("10".to_string()),
416-
Some("20".to_string()),
417-
Some("30".to_string()),
418-
]
419-
);
420-
Ok(())
421-
}
422-
423352
/// Shreds `rows` per `spec`, then canonicalizes so the typed values are lifted into a logical
424353
/// shredded child (as a file read-back would produce).
425354
fn canonical_shredded(rows: Vec<Option<&str>>, spec: ShreddingSpec) -> VortexResult<ArrayRef> {
@@ -439,14 +368,22 @@ mod tests {
439368
spec: ShreddingSpec,
440369
) -> VortexResult<()> {
441370
let unshredded = json_rows_to_variant(rows.clone(), ShreddingSpec::empty())?;
442-
let want = json_strings(&execute_variant_to_json(unshredded)?)?;
371+
let want = json_strings(
372+
&unshredded
373+
.apply(&variant_to_json(root()))?
374+
.execute::<ArrayRef>(&mut SESSION.create_execution_ctx())?,
375+
)?;
443376

444377
let canonical = canonical_shredded(rows, spec)?;
445378
assert!(
446379
canonical.as_::<Variant>().shredded().is_some(),
447380
"fixture must carry a canonical shredded child"
448381
);
449-
let got = json_strings(&execute_variant_to_json(canonical)?)?;
382+
let got = json_strings(
383+
&canonical
384+
.apply(&variant_to_json(root()))?
385+
.execute::<ArrayRef>(&mut SESSION.create_execution_ctx())?,
386+
)?;
450387

451388
assert_eq!(got, want);
452389
Ok(())
@@ -474,7 +411,9 @@ mod tests {
474411
"fixture must carry a canonical shredded child"
475412
);
476413

477-
let result = execute_variant_to_json(canonical)?;
414+
let result = canonical
415+
.apply(&variant_to_json(root()))?
416+
.execute::<ArrayRef>(&mut SESSION.create_execution_ctx())?;
478417

479418
assert_eq!(
480419
json_strings(&result)?,
@@ -525,32 +464,15 @@ mod tests {
525464
)
526465
}
527466

528-
#[test]
529-
fn json_round_trip_normalizes_whitespace_and_key_order() -> VortexResult<()> {
530-
let input = json_rows_to_variant(
531-
vec![Some(r#"{ "b" : 1 , "a" : 2 }"#), Some("[ 1 , 2 , 3 ]")],
532-
ShreddingSpec::empty(),
533-
)?;
534-
535-
let result = execute_variant_to_json(input)?;
536-
537-
assert_eq!(
538-
json_strings(&result)?,
539-
vec![
540-
Some(r#"{"a":2,"b":1}"#.to_string()),
541-
Some("[1,2,3]".to_string()),
542-
]
543-
);
544-
Ok(())
545-
}
546-
547467
#[test]
548468
fn variant_only_types_are_stringified_so_reparsing_loses_types() -> VortexResult<()> {
549469
let date =
550470
NaiveDate::from_ymd_opt(2026, 6, 11).ok_or_else(|| vortex_err!("invalid test date"))?;
551471
let input = unshredded_variant([PqVariant::from(date)])?;
552472

553-
let rendered = execute_variant_to_json(input)?;
473+
let rendered = input
474+
.apply(&variant_to_json(root()))?
475+
.execute::<ArrayRef>(&mut SESSION.create_execution_ctx())?;
554476
let json = json_strings(&rendered)?;
555477
assert_eq!(json, vec![Some(r#""2026-06-11""#.to_string())]);
556478

0 commit comments

Comments
 (0)