1use std::collections::HashMap;
19use std::sync::Arc;
20
21use crate::protobuf_common as protobuf;
22use crate::protobuf_common::{
23 EmptyMessage, arrow_type::ArrowTypeEnum, scalar_value::Value,
24};
25use arrow::array::{ArrayRef, RecordBatch};
26use arrow::csv::{QuoteStyle, WriterBuilder};
27use arrow::datatypes::{
28 DataType, Field, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalUnit, Schema,
29 SchemaRef, TimeUnit, UnionMode,
30};
31use arrow::ipc::writer::{
32 DictionaryTracker, IpcDataGenerator, IpcWriteContext, IpcWriteOptions,
33};
34use datafusion_common::parsers::CsvQuoteStyle;
35use datafusion_common::{
36 Column, ColumnStatistics, Constraint, Constraints, DFSchema, DFSchemaRef,
37 DataFusionError, JoinSide, ScalarValue, Statistics,
38 config::{
39 CsvOptions, JsonOptions, ParquetCdcOptions, ParquetColumnOptions, ParquetOptions,
40 TableParquetOptions,
41 },
42 file_options::{csv_writer::CsvWriterOptions, json_writer::JsonWriterOptions},
43 parsers::CompressionTypeVariant,
44 plan_datafusion_err,
45 stats::Precision,
46};
47
48#[derive(Debug)]
49pub enum Error {
50 General(String),
51
52 InvalidScalarValue(ScalarValue),
53
54 InvalidScalarType(DataType),
55
56 InvalidTimeUnit(TimeUnit),
57
58 NotImplemented(String),
59}
60
61impl std::error::Error for Error {}
62
63impl std::fmt::Display for Error {
64 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
65 match self {
66 Self::General(desc) => write!(f, "General error: {desc}"),
67 Self::InvalidScalarValue(value) => {
68 write!(f, "{value:?} is invalid as a DataFusion scalar value")
69 }
70 Self::InvalidScalarType(data_type) => {
71 write!(f, "{data_type} is invalid as a DataFusion scalar type")
72 }
73 Self::InvalidTimeUnit(time_unit) => {
74 write!(
75 f,
76 "Only TimeUnit::Microsecond and TimeUnit::Nanosecond are valid time units, found: {time_unit:?}"
77 )
78 }
79 Self::NotImplemented(s) => {
80 write!(f, "Not implemented: {s}")
81 }
82 }
83 }
84}
85
86impl From<Error> for DataFusionError {
87 fn from(e: Error) -> Self {
88 plan_datafusion_err!("{}", e)
89 }
90}
91
92impl TryFrom<&Field> for protobuf::Field {
93 type Error = Error;
94
95 fn try_from(field: &Field) -> Result<Self, Self::Error> {
96 let arrow_type = field.data_type().try_into()?;
97 Ok(Self {
98 name: field.name().to_owned(),
99 arrow_type: Some(Box::new(arrow_type)),
100 nullable: field.is_nullable(),
101 children: Vec::new(),
102 metadata: field.metadata().clone(),
103 })
104 }
105}
106
107impl TryFrom<&DataType> for protobuf::ArrowType {
108 type Error = Error;
109
110 fn try_from(val: &DataType) -> Result<Self, Self::Error> {
111 let arrow_type_enum: ArrowTypeEnum = val.try_into()?;
112 Ok(Self {
113 arrow_type_enum: Some(arrow_type_enum),
114 })
115 }
116}
117
118impl TryFrom<&DataType> for ArrowTypeEnum {
119 type Error = Error;
120
121 fn try_from(val: &DataType) -> Result<Self, Self::Error> {
122 let res = match val {
123 DataType::Null => Self::None(EmptyMessage {}),
124 DataType::Boolean => Self::Bool(EmptyMessage {}),
125 DataType::Int8 => Self::Int8(EmptyMessage {}),
126 DataType::Int16 => Self::Int16(EmptyMessage {}),
127 DataType::Int32 => Self::Int32(EmptyMessage {}),
128 DataType::Int64 => Self::Int64(EmptyMessage {}),
129 DataType::UInt8 => Self::Uint8(EmptyMessage {}),
130 DataType::UInt16 => Self::Uint16(EmptyMessage {}),
131 DataType::UInt32 => Self::Uint32(EmptyMessage {}),
132 DataType::UInt64 => Self::Uint64(EmptyMessage {}),
133 DataType::Float16 => Self::Float16(EmptyMessage {}),
134 DataType::Float32 => Self::Float32(EmptyMessage {}),
135 DataType::Float64 => Self::Float64(EmptyMessage {}),
136 DataType::Timestamp(time_unit, timezone) => {
137 Self::Timestamp(protobuf::Timestamp {
138 time_unit: protobuf::TimeUnit::from(time_unit) as i32,
139 timezone: timezone.as_deref().unwrap_or("").to_string(),
140 })
141 }
142 DataType::Date32 => Self::Date32(EmptyMessage {}),
143 DataType::Date64 => Self::Date64(EmptyMessage {}),
144 DataType::Time32(time_unit) => {
145 Self::Time32(protobuf::TimeUnit::from(time_unit) as i32)
146 }
147 DataType::Time64(time_unit) => {
148 Self::Time64(protobuf::TimeUnit::from(time_unit) as i32)
149 }
150 DataType::Duration(time_unit) => {
151 Self::Duration(protobuf::TimeUnit::from(time_unit) as i32)
152 }
153 DataType::Interval(interval_unit) => {
154 Self::Interval(protobuf::IntervalUnit::from(interval_unit) as i32)
155 }
156 DataType::Binary => Self::Binary(EmptyMessage {}),
157 DataType::BinaryView => Self::BinaryView(EmptyMessage {}),
158 DataType::FixedSizeBinary(size) => Self::FixedSizeBinary(*size),
159 DataType::LargeBinary => Self::LargeBinary(EmptyMessage {}),
160 DataType::Utf8 => Self::Utf8(EmptyMessage {}),
161 DataType::Utf8View => Self::Utf8View(EmptyMessage {}),
162 DataType::LargeUtf8 => Self::LargeUtf8(EmptyMessage {}),
163 DataType::List(item_type) => Self::List(Box::new(protobuf::List {
164 field_type: Some(Box::new(item_type.as_ref().try_into()?)),
165 })),
166 DataType::FixedSizeList(item_type, size) => {
167 Self::FixedSizeList(Box::new(protobuf::FixedSizeList {
168 field_type: Some(Box::new(item_type.as_ref().try_into()?)),
169 list_size: *size,
170 }))
171 }
172 DataType::LargeList(item_type) => Self::LargeList(Box::new(protobuf::List {
173 field_type: Some(Box::new(item_type.as_ref().try_into()?)),
174 })),
175 DataType::ListView(item_type) => Self::ListView(Box::new(protobuf::List {
176 field_type: Some(Box::new(item_type.as_ref().try_into()?)),
177 })),
178 DataType::LargeListView(item_type) => {
179 Self::LargeListView(Box::new(protobuf::List {
180 field_type: Some(Box::new(item_type.as_ref().try_into()?)),
181 }))
182 }
183 DataType::Struct(struct_fields) => Self::Struct(protobuf::Struct {
184 sub_field_types: convert_arc_fields_to_proto_fields(struct_fields)?,
185 }),
186 DataType::Union(fields, union_mode) => {
187 let union_mode = match union_mode {
188 UnionMode::Sparse => protobuf::UnionMode::Sparse,
189 UnionMode::Dense => protobuf::UnionMode::Dense,
190 };
191 Self::Union(protobuf::Union {
192 union_types: convert_arc_fields_to_proto_fields(
193 fields.iter().map(|(_, item)| item),
194 )?,
195 union_mode: union_mode.into(),
196 type_ids: fields.iter().map(|(x, _)| x as i32).collect(),
197 })
198 }
199 DataType::Dictionary(key_type, value_type) => {
200 Self::Dictionary(Box::new(protobuf::Dictionary {
201 key: Some(Box::new(key_type.as_ref().try_into()?)),
202 value: Some(Box::new(value_type.as_ref().try_into()?)),
203 }))
204 }
205 DataType::Decimal32(precision, scale) => {
206 Self::Decimal32(protobuf::Decimal32Type {
207 precision: *precision as u32,
208 scale: *scale as i32,
209 })
210 }
211 DataType::Decimal64(precision, scale) => {
212 Self::Decimal64(protobuf::Decimal64Type {
213 precision: *precision as u32,
214 scale: *scale as i32,
215 })
216 }
217 DataType::Decimal128(precision, scale) => {
218 Self::Decimal128(protobuf::Decimal128Type {
219 precision: *precision as u32,
220 scale: *scale as i32,
221 })
222 }
223 DataType::Decimal256(precision, scale) => {
224 Self::Decimal256(protobuf::Decimal256Type {
225 precision: *precision as u32,
226 scale: *scale as i32,
227 })
228 }
229 DataType::Map(field, sorted) => Self::Map(Box::new(protobuf::Map {
230 field_type: Some(Box::new(field.as_ref().try_into()?)),
231 keys_sorted: *sorted,
232 })),
233 DataType::RunEndEncoded(run_ends_field, values_field) => {
234 Self::RunEndEncoded(Box::new(protobuf::RunEndEncoded {
235 run_ends_field: Some(Box::new(run_ends_field.as_ref().try_into()?)),
236 values_field: Some(Box::new(values_field.as_ref().try_into()?)),
237 }))
238 }
239 };
240
241 Ok(res)
242 }
243}
244
245impl From<Column> for protobuf::Column {
246 fn from(c: Column) -> Self {
247 Self {
248 relation: c.relation.map(|relation| protobuf::ColumnRelation {
249 relation: relation.to_string(),
250 }),
251 name: c.name,
252 }
253 }
254}
255
256impl From<&Column> for protobuf::Column {
257 fn from(c: &Column) -> Self {
258 c.clone().into()
259 }
260}
261
262impl TryFrom<&Schema> for protobuf::Schema {
263 type Error = Error;
264
265 fn try_from(schema: &Schema) -> Result<Self, Self::Error> {
266 Ok(Self {
267 columns: convert_arc_fields_to_proto_fields(schema.fields())?,
268 metadata: schema.metadata.clone(),
269 })
270 }
271}
272
273impl TryFrom<SchemaRef> for protobuf::Schema {
274 type Error = Error;
275
276 fn try_from(schema: SchemaRef) -> Result<Self, Self::Error> {
277 Ok(Self {
278 columns: convert_arc_fields_to_proto_fields(schema.fields())?,
279 metadata: schema.metadata.clone(),
280 })
281 }
282}
283
284impl TryFrom<&DFSchema> for protobuf::DfSchema {
285 type Error = Error;
286
287 fn try_from(s: &DFSchema) -> Result<Self, Self::Error> {
288 let columns = s
289 .iter()
290 .map(|(qualifier, field)| {
291 Ok(protobuf::DfField {
292 field: Some(field.as_ref().try_into()?),
293 qualifier: qualifier.map(|r| protobuf::ColumnRelation {
294 relation: r.to_string(),
295 }),
296 })
297 })
298 .collect::<Result<Vec<_>, Error>>()?;
299 Ok(Self {
300 columns,
301 metadata: s.metadata().clone(),
302 })
303 }
304}
305
306impl TryFrom<&DFSchemaRef> for protobuf::DfSchema {
307 type Error = Error;
308
309 fn try_from(s: &DFSchemaRef) -> Result<Self, Self::Error> {
310 s.as_ref().try_into()
311 }
312}
313
314impl TryFrom<&ScalarValue> for protobuf::ScalarValue {
315 type Error = Error;
316
317 fn try_from(val: &ScalarValue) -> Result<Self, Self::Error> {
318 let data_type = val.data_type();
319 match val {
320 ScalarValue::Boolean(val) => {
321 create_proto_scalar(val.as_ref(), &data_type, |s| Value::BoolValue(*s))
322 }
323 ScalarValue::Float16(val) => {
324 create_proto_scalar(val.as_ref(), &data_type, |s| {
325 Value::Float32Value((*s).into())
326 })
327 }
328 ScalarValue::Float32(val) => {
329 create_proto_scalar(val.as_ref(), &data_type, |s| Value::Float32Value(*s))
330 }
331 ScalarValue::Float64(val) => {
332 create_proto_scalar(val.as_ref(), &data_type, |s| Value::Float64Value(*s))
333 }
334 ScalarValue::Int8(val) => {
335 create_proto_scalar(val.as_ref(), &data_type, |s| {
336 Value::Int8Value(*s as i32)
337 })
338 }
339 ScalarValue::Int16(val) => {
340 create_proto_scalar(val.as_ref(), &data_type, |s| {
341 Value::Int16Value(*s as i32)
342 })
343 }
344 ScalarValue::Int32(val) => {
345 create_proto_scalar(val.as_ref(), &data_type, |s| Value::Int32Value(*s))
346 }
347 ScalarValue::Int64(val) => {
348 create_proto_scalar(val.as_ref(), &data_type, |s| Value::Int64Value(*s))
349 }
350 ScalarValue::UInt8(val) => {
351 create_proto_scalar(val.as_ref(), &data_type, |s| {
352 Value::Uint8Value(*s as u32)
353 })
354 }
355 ScalarValue::UInt16(val) => {
356 create_proto_scalar(val.as_ref(), &data_type, |s| {
357 Value::Uint16Value(*s as u32)
358 })
359 }
360 ScalarValue::UInt32(val) => {
361 create_proto_scalar(val.as_ref(), &data_type, |s| Value::Uint32Value(*s))
362 }
363 ScalarValue::UInt64(val) => {
364 create_proto_scalar(val.as_ref(), &data_type, |s| Value::Uint64Value(*s))
365 }
366 ScalarValue::Utf8(val) => {
367 create_proto_scalar(val.as_ref(), &data_type, |s| {
368 Value::Utf8Value(s.to_owned())
369 })
370 }
371 ScalarValue::LargeUtf8(val) => {
372 create_proto_scalar(val.as_ref(), &data_type, |s| {
373 Value::LargeUtf8Value(s.to_owned())
374 })
375 }
376 ScalarValue::Utf8View(val) => {
377 create_proto_scalar(val.as_ref(), &data_type, |s| {
378 Value::Utf8ViewValue(s.to_owned())
379 })
380 }
381 ScalarValue::List(arr) => {
382 encode_scalar_nested_value(arr.to_owned() as ArrayRef, val)
383 }
384 ScalarValue::LargeList(arr) => {
385 encode_scalar_nested_value(arr.to_owned() as ArrayRef, val)
386 }
387 ScalarValue::FixedSizeList(arr) => {
388 encode_scalar_nested_value(arr.to_owned() as ArrayRef, val)
389 }
390 ScalarValue::ListView(arr) => {
391 encode_scalar_nested_value(arr.to_owned() as ArrayRef, val)
392 }
393 ScalarValue::LargeListView(arr) => {
394 encode_scalar_nested_value(arr.to_owned() as ArrayRef, val)
395 }
396 ScalarValue::Struct(arr) => {
397 encode_scalar_nested_value(arr.to_owned() as ArrayRef, val)
398 }
399 ScalarValue::Map(arr) => {
400 encode_scalar_nested_value(arr.to_owned() as ArrayRef, val)
401 }
402 ScalarValue::Date32(val) => {
403 create_proto_scalar(val.as_ref(), &data_type, |s| Value::Date32Value(*s))
404 }
405 ScalarValue::TimestampMicrosecond(val, tz) => {
406 create_proto_scalar(val.as_ref(), &data_type, |s| {
407 Value::TimestampValue(protobuf::ScalarTimestampValue {
408 timezone: tz.as_deref().unwrap_or("").to_string(),
409 value: Some(
410 protobuf::scalar_timestamp_value::Value::TimeMicrosecondValue(
411 *s,
412 ),
413 ),
414 })
415 })
416 }
417 ScalarValue::TimestampNanosecond(val, tz) => {
418 create_proto_scalar(val.as_ref(), &data_type, |s| {
419 Value::TimestampValue(protobuf::ScalarTimestampValue {
420 timezone: tz.as_deref().unwrap_or("").to_string(),
421 value: Some(
422 protobuf::scalar_timestamp_value::Value::TimeNanosecondValue(
423 *s,
424 ),
425 ),
426 })
427 })
428 }
429 ScalarValue::Decimal32(val, p, s) => match *val {
430 Some(v) => {
431 let array = v.to_be_bytes();
432 let vec_val: Vec<u8> = array.to_vec();
433 Ok(protobuf::ScalarValue {
434 value: Some(Value::Decimal32Value(protobuf::Decimal32 {
435 value: vec_val,
436 p: *p as i64,
437 s: *s as i64,
438 })),
439 })
440 }
441 None => Ok(protobuf::ScalarValue {
442 value: Some(Value::NullValue((&data_type).try_into()?)),
443 }),
444 },
445 ScalarValue::Decimal64(val, p, s) => match *val {
446 Some(v) => {
447 let array = v.to_be_bytes();
448 let vec_val: Vec<u8> = array.to_vec();
449 Ok(protobuf::ScalarValue {
450 value: Some(Value::Decimal64Value(protobuf::Decimal64 {
451 value: vec_val,
452 p: *p as i64,
453 s: *s as i64,
454 })),
455 })
456 }
457 None => Ok(protobuf::ScalarValue {
458 value: Some(Value::NullValue((&data_type).try_into()?)),
459 }),
460 },
461 ScalarValue::Decimal128(val, p, s) => match *val {
462 Some(v) => {
463 let array = v.to_be_bytes();
464 let vec_val: Vec<u8> = array.to_vec();
465 Ok(protobuf::ScalarValue {
466 value: Some(Value::Decimal128Value(protobuf::Decimal128 {
467 value: vec_val,
468 p: *p as i64,
469 s: *s as i64,
470 })),
471 })
472 }
473 None => Ok(protobuf::ScalarValue {
474 value: Some(Value::NullValue((&data_type).try_into()?)),
475 }),
476 },
477 ScalarValue::Decimal256(val, p, s) => match *val {
478 Some(v) => {
479 let array = v.to_be_bytes();
480 let vec_val: Vec<u8> = array.to_vec();
481 Ok(protobuf::ScalarValue {
482 value: Some(Value::Decimal256Value(protobuf::Decimal256 {
483 value: vec_val,
484 p: *p as i64,
485 s: *s as i64,
486 })),
487 })
488 }
489 None => Ok(protobuf::ScalarValue {
490 value: Some(Value::NullValue((&data_type).try_into()?)),
491 }),
492 },
493 ScalarValue::Date64(val) => {
494 create_proto_scalar(val.as_ref(), &data_type, |s| Value::Date64Value(*s))
495 }
496 ScalarValue::TimestampSecond(val, tz) => {
497 create_proto_scalar(val.as_ref(), &data_type, |s| {
498 Value::TimestampValue(protobuf::ScalarTimestampValue {
499 timezone: tz.as_deref().unwrap_or("").to_string(),
500 value: Some(
501 protobuf::scalar_timestamp_value::Value::TimeSecondValue(*s),
502 ),
503 })
504 })
505 }
506 ScalarValue::TimestampMillisecond(val, tz) => {
507 create_proto_scalar(val.as_ref(), &data_type, |s| {
508 Value::TimestampValue(protobuf::ScalarTimestampValue {
509 timezone: tz.as_deref().unwrap_or("").to_string(),
510 value: Some(
511 protobuf::scalar_timestamp_value::Value::TimeMillisecondValue(
512 *s,
513 ),
514 ),
515 })
516 })
517 }
518 ScalarValue::IntervalYearMonth(val) => {
519 create_proto_scalar(val.as_ref(), &data_type, |s| {
520 Value::IntervalYearmonthValue(*s)
521 })
522 }
523 ScalarValue::Null => Ok(protobuf::ScalarValue {
524 value: Some(Value::NullValue((&data_type).try_into()?)),
525 }),
526
527 ScalarValue::Binary(val) => {
528 create_proto_scalar(val.as_ref(), &data_type, |s| {
529 Value::BinaryValue(s.to_owned())
530 })
531 }
532 ScalarValue::BinaryView(val) => {
533 create_proto_scalar(val.as_ref(), &data_type, |s| {
534 Value::BinaryViewValue(s.to_owned())
535 })
536 }
537 ScalarValue::LargeBinary(val) => {
538 create_proto_scalar(val.as_ref(), &data_type, |s| {
539 Value::LargeBinaryValue(s.to_owned())
540 })
541 }
542 ScalarValue::FixedSizeBinary(length, val) => {
543 create_proto_scalar(val.as_ref(), &data_type, |s| {
544 Value::FixedSizeBinaryValue(protobuf::ScalarFixedSizeBinary {
545 values: s.to_owned(),
546 length: *length,
547 })
548 })
549 }
550
551 ScalarValue::Time32Second(v) => {
552 create_proto_scalar(v.as_ref(), &data_type, |v| {
553 Value::Time32Value(protobuf::ScalarTime32Value {
554 value: Some(
555 protobuf::scalar_time32_value::Value::Time32SecondValue(*v),
556 ),
557 })
558 })
559 }
560
561 ScalarValue::Time32Millisecond(v) => {
562 create_proto_scalar(v.as_ref(), &data_type, |v| {
563 Value::Time32Value(protobuf::ScalarTime32Value {
564 value: Some(
565 protobuf::scalar_time32_value::Value::Time32MillisecondValue(
566 *v,
567 ),
568 ),
569 })
570 })
571 }
572
573 ScalarValue::Time64Microsecond(v) => {
574 create_proto_scalar(v.as_ref(), &data_type, |v| {
575 Value::Time64Value(protobuf::ScalarTime64Value {
576 value: Some(
577 protobuf::scalar_time64_value::Value::Time64MicrosecondValue(
578 *v,
579 ),
580 ),
581 })
582 })
583 }
584
585 ScalarValue::Time64Nanosecond(v) => {
586 create_proto_scalar(v.as_ref(), &data_type, |v| {
587 Value::Time64Value(protobuf::ScalarTime64Value {
588 value: Some(
589 protobuf::scalar_time64_value::Value::Time64NanosecondValue(
590 *v,
591 ),
592 ),
593 })
594 })
595 }
596
597 ScalarValue::IntervalDayTime(val) => {
598 let value = if let Some(v) = val {
599 let (days, milliseconds) = IntervalDayTimeType::to_parts(*v);
600 Value::IntervalDaytimeValue(protobuf::IntervalDayTimeValue {
601 days,
602 milliseconds,
603 })
604 } else {
605 Value::NullValue((&data_type).try_into()?)
606 };
607
608 Ok(protobuf::ScalarValue { value: Some(value) })
609 }
610
611 ScalarValue::IntervalMonthDayNano(v) => {
612 let value = if let Some(v) = v {
613 let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(*v);
614 Value::IntervalMonthDayNano(protobuf::IntervalMonthDayNanoValue {
615 months,
616 days,
617 nanos,
618 })
619 } else {
620 Value::NullValue((&data_type).try_into()?)
621 };
622
623 Ok(protobuf::ScalarValue { value: Some(value) })
624 }
625
626 ScalarValue::DurationSecond(v) => {
627 let value = match v {
628 Some(v) => Value::DurationSecondValue(*v),
629 None => Value::NullValue((&data_type).try_into()?),
630 };
631 Ok(protobuf::ScalarValue { value: Some(value) })
632 }
633 ScalarValue::DurationMillisecond(v) => {
634 let value = match v {
635 Some(v) => Value::DurationMillisecondValue(*v),
636 None => Value::NullValue((&data_type).try_into()?),
637 };
638 Ok(protobuf::ScalarValue { value: Some(value) })
639 }
640 ScalarValue::DurationMicrosecond(v) => {
641 let value = match v {
642 Some(v) => Value::DurationMicrosecondValue(*v),
643 None => Value::NullValue((&data_type).try_into()?),
644 };
645 Ok(protobuf::ScalarValue { value: Some(value) })
646 }
647 ScalarValue::DurationNanosecond(v) => {
648 let value = match v {
649 Some(v) => Value::DurationNanosecondValue(*v),
650 None => Value::NullValue((&data_type).try_into()?),
651 };
652 Ok(protobuf::ScalarValue { value: Some(value) })
653 }
654
655 ScalarValue::Union(val, df_fields, mode) => {
656 let mut fields =
657 Vec::<protobuf::UnionField>::with_capacity(df_fields.len());
658 for (id, field) in df_fields.iter() {
659 let field_id = id as i32;
660 let field = Some(field.as_ref().try_into()?);
661 let field = protobuf::UnionField { field_id, field };
662 fields.push(field);
663 }
664 let mode = match mode {
665 UnionMode::Sparse => 0,
666 UnionMode::Dense => 1,
667 };
668 let value = match val {
669 None => None,
670 Some((_id, v)) => Some(Box::new(v.as_ref().try_into()?)),
671 };
672 let val = protobuf::UnionValue {
673 value_id: val.as_ref().map(|(id, _v)| *id as i32).unwrap_or(0),
674 value,
675 fields,
676 mode,
677 };
678 let val = Value::UnionValue(Box::new(val));
679 let val = protobuf::ScalarValue { value: Some(val) };
680 Ok(val)
681 }
682
683 ScalarValue::Dictionary(index_type, val) => {
684 let value: protobuf::ScalarValue = val.as_ref().try_into()?;
685 Ok(protobuf::ScalarValue {
686 value: Some(Value::DictionaryValue(Box::new(
687 protobuf::ScalarDictionaryValue {
688 index_type: Some(index_type.as_ref().try_into()?),
689 value: Some(Box::new(value)),
690 },
691 ))),
692 })
693 }
694
695 ScalarValue::RunEndEncoded(run_ends_field, values_field, val) => {
696 Ok(protobuf::ScalarValue {
697 value: Some(Value::RunEndEncodedValue(Box::new(
698 protobuf::ScalarRunEndEncodedValue {
699 run_ends_field: Some(run_ends_field.as_ref().try_into()?),
700 values_field: Some(values_field.as_ref().try_into()?),
701 value: Some(Box::new(val.as_ref().try_into()?)),
702 },
703 ))),
704 })
705 }
706 }
707 }
708}
709
710impl From<&TimeUnit> for protobuf::TimeUnit {
711 fn from(val: &TimeUnit) -> Self {
712 match val {
713 TimeUnit::Second => protobuf::TimeUnit::Second,
714 TimeUnit::Millisecond => protobuf::TimeUnit::Millisecond,
715 TimeUnit::Microsecond => protobuf::TimeUnit::Microsecond,
716 TimeUnit::Nanosecond => protobuf::TimeUnit::Nanosecond,
717 }
718 }
719}
720
721impl From<&IntervalUnit> for protobuf::IntervalUnit {
722 fn from(interval_unit: &IntervalUnit) -> Self {
723 match interval_unit {
724 IntervalUnit::YearMonth => protobuf::IntervalUnit::YearMonth,
725 IntervalUnit::DayTime => protobuf::IntervalUnit::DayTime,
726 IntervalUnit::MonthDayNano => protobuf::IntervalUnit::MonthDayNano,
727 }
728 }
729}
730
731impl From<Constraints> for protobuf::Constraints {
732 fn from(value: Constraints) -> Self {
733 let constraints = value.into_iter().map(|item| item.into()).collect();
734 protobuf::Constraints { constraints }
735 }
736}
737
738impl From<Constraint> for protobuf::Constraint {
739 fn from(value: Constraint) -> Self {
740 let res = match value {
741 Constraint::PrimaryKey(indices) => {
742 let indices = indices.into_iter().map(|item| item as u64).collect();
743 protobuf::constraint::ConstraintMode::PrimaryKey(
744 protobuf::PrimaryKeyConstraint { indices },
745 )
746 }
747 Constraint::Unique(indices) => {
748 let indices = indices.into_iter().map(|item| item as u64).collect();
749 protobuf::constraint::ConstraintMode::PrimaryKey(
750 protobuf::PrimaryKeyConstraint { indices },
751 )
752 }
753 };
754 protobuf::Constraint {
755 constraint_mode: Some(res),
756 }
757 }
758}
759
760impl From<&Precision<usize>> for protobuf::Precision {
761 fn from(s: &Precision<usize>) -> protobuf::Precision {
762 match s {
763 Precision::Exact(val) => protobuf::Precision {
764 precision_info: protobuf::PrecisionInfo::Exact.into(),
765 val: Some(crate::protobuf_common::ScalarValue {
766 value: Some(Value::Uint64Value(*val as u64)),
767 }),
768 },
769 Precision::Inexact(val) => protobuf::Precision {
770 precision_info: protobuf::PrecisionInfo::Inexact.into(),
771 val: Some(crate::protobuf_common::ScalarValue {
772 value: Some(Value::Uint64Value(*val as u64)),
773 }),
774 },
775 Precision::Absent => protobuf::Precision {
776 precision_info: protobuf::PrecisionInfo::Absent.into(),
777 val: Some(crate::protobuf_common::ScalarValue { value: None }),
778 },
779 }
780 }
781}
782
783impl From<&Precision<ScalarValue>> for protobuf::Precision {
784 fn from(s: &Precision<ScalarValue>) -> protobuf::Precision {
785 match s {
786 Precision::Exact(val) => protobuf::Precision {
787 precision_info: protobuf::PrecisionInfo::Exact.into(),
788 val: val.try_into().ok(),
789 },
790 Precision::Inexact(val) => protobuf::Precision {
791 precision_info: protobuf::PrecisionInfo::Inexact.into(),
792 val: val.try_into().ok(),
793 },
794 Precision::Absent => protobuf::Precision {
795 precision_info: protobuf::PrecisionInfo::Absent.into(),
796 val: Some(crate::protobuf_common::ScalarValue { value: None }),
797 },
798 }
799 }
800}
801
802impl From<&Statistics> for protobuf::Statistics {
803 fn from(s: &Statistics) -> protobuf::Statistics {
804 let column_stats = s.column_statistics.iter().map(|s| s.into()).collect();
805 protobuf::Statistics {
806 num_rows: Some(protobuf::Precision::from(&s.num_rows)),
807 total_byte_size: Some(protobuf::Precision::from(&s.total_byte_size)),
808 column_stats,
809 }
810 }
811}
812
813impl From<&ColumnStatistics> for protobuf::ColumnStats {
814 fn from(s: &ColumnStatistics) -> protobuf::ColumnStats {
815 protobuf::ColumnStats {
816 min_value: Some(protobuf::Precision::from(&s.min_value)),
817 max_value: Some(protobuf::Precision::from(&s.max_value)),
818 sum_value: Some(protobuf::Precision::from(&s.sum_value)),
819 null_count: Some(protobuf::Precision::from(&s.null_count)),
820 distinct_count: Some(protobuf::Precision::from(&s.distinct_count)),
821 byte_size: Some(protobuf::Precision::from(&s.byte_size)),
822 }
823 }
824}
825
826impl From<JoinSide> for protobuf::JoinSide {
827 fn from(t: JoinSide) -> Self {
828 match t {
829 JoinSide::Left => protobuf::JoinSide::LeftSide,
830 JoinSide::Right => protobuf::JoinSide::RightSide,
831 JoinSide::None => protobuf::JoinSide::None,
832 }
833 }
834}
835
836impl From<&CompressionTypeVariant> for protobuf::CompressionTypeVariant {
837 fn from(value: &CompressionTypeVariant) -> Self {
838 match value {
839 CompressionTypeVariant::GZIP => Self::Gzip,
840 CompressionTypeVariant::BZIP2 => Self::Bzip2,
841 CompressionTypeVariant::XZ => Self::Xz,
842 CompressionTypeVariant::ZSTD => Self::Zstd,
843 CompressionTypeVariant::UNCOMPRESSED => Self::Uncompressed,
844 }
845 }
846}
847
848impl From<CsvQuoteStyle> for protobuf::CsvQuoteStyle {
849 fn from(value: CsvQuoteStyle) -> Self {
850 match value {
851 CsvQuoteStyle::Necessary => Self::Necessary,
852 CsvQuoteStyle::Always => Self::Always,
853 CsvQuoteStyle::NonNumeric => Self::NonNumeric,
854 CsvQuoteStyle::Never => Self::Never,
855 }
856 }
857}
858
859impl From<QuoteStyle> for protobuf::CsvQuoteStyle {
860 fn from(value: QuoteStyle) -> Self {
861 match value {
862 QuoteStyle::Necessary => Self::Necessary,
863 QuoteStyle::Always => Self::Always,
864 QuoteStyle::NonNumeric => Self::NonNumeric,
865 QuoteStyle::Never => Self::Never,
866 _ => Self::Necessary,
867 }
868 }
869}
870
871impl TryFrom<&CsvWriterOptions> for protobuf::CsvWriterOptions {
872 type Error = DataFusionError;
873
874 fn try_from(opts: &CsvWriterOptions) -> datafusion_common::Result<Self, Self::Error> {
875 Ok(csv_writer_options_to_proto(
876 &opts.writer_options,
877 &opts.compression,
878 ))
879 }
880}
881
882impl TryFrom<&JsonWriterOptions> for protobuf::JsonWriterOptions {
883 type Error = DataFusionError;
884
885 fn try_from(
886 opts: &JsonWriterOptions,
887 ) -> datafusion_common::Result<Self, Self::Error> {
888 let compression: protobuf::CompressionTypeVariant = opts.compression.into();
889 Ok(protobuf::JsonWriterOptions {
890 compression: compression.into(),
891 })
892 }
893}
894
895impl TryFrom<&ParquetOptions> for protobuf::ParquetOptions {
896 type Error = DataFusionError;
897
898 fn try_from(value: &ParquetOptions) -> datafusion_common::Result<Self, Self::Error> {
899 Ok(protobuf::ParquetOptions {
900 enable_page_index: value.enable_page_index,
901 pruning: value.pruning,
902 skip_metadata: value.skip_metadata,
903 metadata_size_hint_opt: value.metadata_size_hint.map(|v| protobuf::parquet_options::MetadataSizeHintOpt::MetadataSizeHint(v as u64)),
904 pushdown_filters: value.pushdown_filters,
905 reorder_filters: value.reorder_filters,
906 force_filter_selections: value.force_filter_selections,
907 data_pagesize_limit: value.data_pagesize_limit as u64,
908 write_batch_size: value.write_batch_size as u64,
909 writer_version: value.writer_version.to_string(),
910 compression_opt: value.compression.clone().map(protobuf::parquet_options::CompressionOpt::Compression),
911 dictionary_enabled_opt: value.dictionary_enabled.map(protobuf::parquet_options::DictionaryEnabledOpt::DictionaryEnabled),
912 dictionary_page_size_limit: value.dictionary_page_size_limit as u64,
913 statistics_enabled_opt: value.statistics_enabled.clone().map(protobuf::parquet_options::StatisticsEnabledOpt::StatisticsEnabled),
914 max_row_group_size: value.max_row_group_size as u64,
915 max_in_list_size: value.max_in_list_size as u64,
916 created_by: value.created_by.clone(),
917 column_index_truncate_length_opt: value.column_index_truncate_length.map(|v| protobuf::parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(v as u64)),
918 statistics_truncate_length_opt: value.statistics_truncate_length.map(|v| protobuf::parquet_options::StatisticsTruncateLengthOpt::StatisticsTruncateLength(v as u64)),
919 data_page_row_count_limit: value.data_page_row_count_limit as u64,
920 encoding_opt: value.encoding.clone().map(protobuf::parquet_options::EncodingOpt::Encoding),
921 bloom_filter_on_read: value.bloom_filter_on_read,
922 bloom_filter_on_write: value.bloom_filter_on_write,
923 bloom_filter_fpp_opt: value.bloom_filter_fpp.map(protobuf::parquet_options::BloomFilterFppOpt::BloomFilterFpp),
924 bloom_filter_ndv_opt: value.bloom_filter_ndv.map(protobuf::parquet_options::BloomFilterNdvOpt::BloomFilterNdv),
925 allow_single_file_parallelism: value.allow_single_file_parallelism,
926 maximum_parallel_row_group_writers: value.maximum_parallel_row_group_writers as u64,
927 maximum_buffered_record_batches_per_stream: value.maximum_buffered_record_batches_per_stream as u64,
928 schema_force_view_types: value.schema_force_view_types,
929 binary_as_string: value.binary_as_string,
930 skip_arrow_metadata: value.skip_arrow_metadata,
931 coerce_int96_opt: value.coerce_int96.clone().map(protobuf::parquet_options::CoerceInt96Opt::CoerceInt96),
932 coerce_int96_tz_opt: value.coerce_int96_tz.clone().map(protobuf::parquet_options::CoerceInt96TzOpt::CoerceInt96Tz),
933 max_predicate_cache_size_opt: value.max_predicate_cache_size.map(|v| protobuf::parquet_options::MaxPredicateCacheSizeOpt::MaxPredicateCacheSize(v as u64)),
934 max_row_group_bytes_opt: value.max_row_group_bytes.map(|v| protobuf::parquet_options::MaxRowGroupBytesOpt::MaxRowGroupBytes(v.get() as u64)),
935 content_defined_chunking: Some((&value.content_defined_chunking).into()),
936 })
937 }
938}
939
940impl From<&ParquetCdcOptions> for protobuf::ParquetCdcOptions {
941 fn from(value: &ParquetCdcOptions) -> Self {
942 protobuf::ParquetCdcOptions {
943 enabled: value.enabled,
944 min_chunk_size: value.min_chunk_size as u64,
945 max_chunk_size: value.max_chunk_size as u64,
946 norm_level: value.norm_level,
947 }
948 }
949}
950
951impl TryFrom<&ParquetColumnOptions> for protobuf::ParquetColumnOptions {
952 type Error = DataFusionError;
953
954 fn try_from(
955 value: &ParquetColumnOptions,
956 ) -> datafusion_common::Result<Self, Self::Error> {
957 Ok(protobuf::ParquetColumnOptions {
958 compression_opt: value
959 .compression
960 .clone()
961 .map(protobuf::parquet_column_options::CompressionOpt::Compression),
962 dictionary_enabled_opt: value
963 .dictionary_enabled
964 .map(protobuf::parquet_column_options::DictionaryEnabledOpt::DictionaryEnabled),
965 statistics_enabled_opt: value
966 .statistics_enabled
967 .clone()
968 .map(protobuf::parquet_column_options::StatisticsEnabledOpt::StatisticsEnabled),
969 encoding_opt: value
970 .encoding
971 .clone()
972 .map(protobuf::parquet_column_options::EncodingOpt::Encoding),
973 bloom_filter_enabled_opt: value
974 .bloom_filter_enabled
975 .map(protobuf::parquet_column_options::BloomFilterEnabledOpt::BloomFilterEnabled),
976 bloom_filter_fpp_opt: value
977 .bloom_filter_fpp
978 .map(protobuf::parquet_column_options::BloomFilterFppOpt::BloomFilterFpp),
979 bloom_filter_ndv_opt: value
980 .bloom_filter_ndv
981 .map(protobuf::parquet_column_options::BloomFilterNdvOpt::BloomFilterNdv),
982 })
983 }
984}
985
986impl TryFrom<&TableParquetOptions> for protobuf::TableParquetOptions {
987 type Error = DataFusionError;
988 fn try_from(
989 value: &TableParquetOptions,
990 ) -> datafusion_common::Result<Self, Self::Error> {
991 let column_specific_options = value
992 .column_specific_options
993 .iter()
994 .map(|(k, v)| {
995 Ok(protobuf::ParquetColumnSpecificOptions {
996 column_name: k.into(),
997 options: Some(v.try_into()?),
998 })
999 })
1000 .collect::<datafusion_common::Result<Vec<_>>>()?;
1001 let key_value_metadata = value
1002 .key_value_metadata
1003 .iter()
1004 .filter_map(|(k, v)| v.as_ref().map(|v| (k.clone(), v.clone())))
1005 .collect::<HashMap<String, String>>();
1006
1007 let global: protobuf::ParquetOptions = (&value.global).try_into()?;
1008
1009 Ok(protobuf::TableParquetOptions {
1010 global: Some(global),
1011 column_specific_options,
1012 key_value_metadata,
1013 })
1014 }
1015}
1016
1017impl TryFrom<&CsvOptions> for protobuf::CsvOptions {
1018 type Error = DataFusionError; fn try_from(opts: &CsvOptions) -> datafusion_common::Result<Self, Self::Error> {
1021 let compression: protobuf::CompressionTypeVariant = opts.compression.into();
1022 let quote_style: protobuf::CsvQuoteStyle = opts.quote_style.into();
1023 Ok(protobuf::CsvOptions {
1024 has_header: opts.has_header.map_or_else(Vec::new, |h| vec![h as u8]),
1025 delimiter: vec![opts.delimiter],
1026 quote: vec![opts.quote],
1027 terminator: opts.terminator.map_or_else(Vec::new, |e| vec![e]),
1028 escape: opts.escape.map_or_else(Vec::new, |e| vec![e]),
1029 double_quote: opts.double_quote.map_or_else(Vec::new, |h| vec![h as u8]),
1030 newlines_in_values: opts
1031 .newlines_in_values
1032 .map_or_else(Vec::new, |h| vec![h as u8]),
1033 compression: compression.into(),
1034 schema_infer_max_rec: opts.schema_infer_max_rec.map(|h| h as u64),
1035 date_format: opts.date_format.clone().unwrap_or_default(),
1036 datetime_format: opts.datetime_format.clone().unwrap_or_default(),
1037 timestamp_format: opts.timestamp_format.clone().unwrap_or_default(),
1038 timestamp_tz_format: opts.timestamp_tz_format.clone().unwrap_or_default(),
1039 time_format: opts.time_format.clone().unwrap_or_default(),
1040 null_value: opts.null_value.clone().unwrap_or_default(),
1041 null_regex: opts.null_regex.clone().unwrap_or_default(),
1042 comment: opts.comment.map_or_else(Vec::new, |h| vec![h]),
1043 truncated_rows: opts.truncated_rows.map_or_else(Vec::new, |h| vec![h as u8]),
1044 compression_level: opts.compression_level,
1045 quote_style: quote_style.into(),
1046 ignore_leading_whitespace: opts
1047 .ignore_leading_whitespace
1048 .map_or_else(Vec::new, |h| vec![h as u8]),
1049 ignore_trailing_whitespace: opts
1050 .ignore_trailing_whitespace
1051 .map_or_else(Vec::new, |h| vec![h as u8]),
1052 })
1053 }
1054}
1055
1056impl TryFrom<&JsonOptions> for protobuf::JsonOptions {
1057 type Error = DataFusionError;
1058
1059 fn try_from(opts: &JsonOptions) -> datafusion_common::Result<Self, Self::Error> {
1060 let compression: protobuf::CompressionTypeVariant = opts.compression.into();
1061 Ok(protobuf::JsonOptions {
1062 compression: compression.into(),
1063 schema_infer_max_rec: opts.schema_infer_max_rec.map(|h| h as u64),
1064 compression_level: opts.compression_level,
1065 newline_delimited: Some(opts.newline_delimited),
1066 })
1067 }
1068}
1069
1070fn create_proto_scalar<I, T: FnOnce(&I) -> Value>(
1073 v: Option<&I>,
1074 null_arrow_type: &DataType,
1075 constructor: T,
1076) -> Result<protobuf::ScalarValue, Error> {
1077 let value = v
1078 .map(constructor)
1079 .unwrap_or(Value::NullValue(null_arrow_type.try_into()?));
1080
1081 Ok(protobuf::ScalarValue { value: Some(value) })
1082}
1083
1084fn encode_scalar_nested_value(
1087 arr: ArrayRef,
1088 val: &ScalarValue,
1089) -> Result<protobuf::ScalarValue, Error> {
1090 let batch = RecordBatch::try_from_iter(vec![("field_name", arr)]).map_err(|e| {
1091 Error::General(format!(
1092 "Error creating temporary batch while encoding nested ScalarValue: {e}"
1093 ))
1094 })?;
1095
1096 let ipc_gen = IpcDataGenerator {};
1097 let mut dict_tracker = DictionaryTracker::new(false);
1098 let write_options = IpcWriteOptions::default();
1099 ipc_gen.schema_to_bytes_with_dictionary_tracker(
1102 batch.schema().as_ref(),
1103 &mut dict_tracker,
1104 &write_options,
1105 );
1106 let mut compression_context = IpcWriteContext::default();
1107 let (encoded_dictionaries, encoded_message) = ipc_gen
1108 .encode(
1109 &batch,
1110 &mut dict_tracker,
1111 &write_options,
1112 &mut compression_context,
1113 )
1114 .map_err(|e| {
1115 Error::General(format!("Error encoding nested ScalarValue as IPC: {e}"))
1116 })?;
1117
1118 let schema: protobuf::Schema = batch.schema().try_into()?;
1119
1120 let scalar_list_value = protobuf::ScalarNestedValue {
1121 ipc_message: encoded_message.ipc_message,
1122 arrow_data: encoded_message.arrow_data,
1123 dictionaries: encoded_dictionaries
1124 .into_iter()
1125 .map(|data| protobuf::scalar_nested_value::Dictionary {
1126 ipc_message: data.ipc_message,
1127 arrow_data: data.arrow_data,
1128 })
1129 .collect(),
1130 schema: Some(schema),
1131 };
1132
1133 match val {
1134 ScalarValue::List(_) => Ok(protobuf::ScalarValue {
1135 value: Some(Value::ListValue(scalar_list_value)),
1136 }),
1137 ScalarValue::LargeList(_) => Ok(protobuf::ScalarValue {
1138 value: Some(Value::LargeListValue(scalar_list_value)),
1139 }),
1140 ScalarValue::FixedSizeList(_) => Ok(protobuf::ScalarValue {
1141 value: Some(Value::FixedSizeListValue(scalar_list_value)),
1142 }),
1143 ScalarValue::ListView(_) => Ok(protobuf::ScalarValue {
1144 value: Some(Value::ListViewValue(scalar_list_value)),
1145 }),
1146 ScalarValue::LargeListView(_) => Ok(protobuf::ScalarValue {
1147 value: Some(Value::LargeListViewValue(scalar_list_value)),
1148 }),
1149 ScalarValue::Struct(_) => Ok(protobuf::ScalarValue {
1150 value: Some(Value::StructValue(scalar_list_value)),
1151 }),
1152 ScalarValue::Map(_) => Ok(protobuf::ScalarValue {
1153 value: Some(Value::MapValue(scalar_list_value)),
1154 }),
1155 _ => unreachable!(),
1156 }
1157}
1158
1159fn convert_arc_fields_to_proto_fields<'a, I>(
1161 fields: I,
1162) -> Result<Vec<protobuf::Field>, Error>
1163where
1164 I: IntoIterator<Item = &'a Arc<Field>>,
1165{
1166 fields
1167 .into_iter()
1168 .map(|field| field.as_ref().try_into())
1169 .collect::<Result<Vec<_>, Error>>()
1170}
1171
1172pub(crate) fn csv_writer_options_to_proto(
1173 csv_options: &WriterBuilder,
1174 compression: &CompressionTypeVariant,
1175) -> protobuf::CsvWriterOptions {
1176 let compression: protobuf::CompressionTypeVariant = compression.into();
1177 let quote_style: protobuf::CsvQuoteStyle = csv_options.quote_style().into();
1178 protobuf::CsvWriterOptions {
1179 compression: compression.into(),
1180 delimiter: (csv_options.delimiter() as char).to_string(),
1181 has_header: csv_options.header(),
1182 date_format: csv_options.date_format().unwrap_or("").to_owned(),
1183 datetime_format: csv_options.datetime_format().unwrap_or("").to_owned(),
1184 timestamp_format: csv_options.timestamp_format().unwrap_or("").to_owned(),
1185 time_format: csv_options.time_format().unwrap_or("").to_owned(),
1186 null_value: csv_options.null().to_owned(),
1187 quote: (csv_options.quote() as char).to_string(),
1188 escape: (csv_options.escape() as char).to_string(),
1189 double_quote: csv_options.double_quote(),
1190 quote_style: quote_style.into(),
1191 ignore_leading_whitespace: csv_options.ignore_leading_whitespace(),
1192 ignore_trailing_whitespace: csv_options.ignore_trailing_whitespace(),
1193 }
1194}