pg2parquet 0.2.2

Command line tool for exporting PostgreSQL tables or queries into Parquet files
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

use parquet::data_type::{Int32Type, Int64Type, ByteArrayType};
use parquet::file::{writer::SerializedFileWriter, properties::WriterProperties};
use parquet::schema::types::Type;
use parquet::basic::{Encoding, Compression, self, Repetition, ConvertedType};
use parquet::format::{self, StringType};
use std::mem;
use std::{fs, sync::Arc};

#[cfg(not(debug_assertions))]
pub fn create_something(_file_path: &std::path::PathBuf) {
	println!("Disabled in release build")
}

#[cfg(debug_assertions)]
pub fn create_something(file_path: &std::path::PathBuf) {
	let file = fs::File::create(&file_path).unwrap();
	let props =
		WriterProperties::builder()
			.set_compression(Compression::SNAPPY)
		.build().arc();

	let schema =
		Type::group_type_builder("root")
			.with_fields(vec![
				Type::primitive_type_builder("id", basic::Type::INT32).build().unwrap().arc(),

				Type::group_type_builder("nested")
					.with_fields(vec![
						Type::primitive_type_builder("x", basic::Type::INT32)
							.with_repetition(Repetition::REPEATED).build().unwrap().arc(),
						Type::primitive_type_builder("y", basic::Type::INT64)
							.with_converted_type(ConvertedType::TIMESTAMP_MICROS)
							.build().unwrap().arc(),
					])
					.with_repetition(Repetition::OPTIONAL)
					.build().unwrap().arc(),

			])
			.build().unwrap().arc();

	let mut writer =
		SerializedFileWriter::new(file, schema, props).unwrap();

	let mut rg = writer.next_row_group().unwrap();

	let mut col_id = rg.next_column().unwrap().unwrap();
	col_id.typed::<Int32Type>().write_batch(
		&vec![1, 0, 2],
		Some(&vec![1, 0, 1]),
		None
	).unwrap();
	col_id.close().unwrap();

	let mut col_nested_x = rg.next_column().unwrap().unwrap();
	col_nested_x.typed::<Int32Type>().write_batch(
		&vec![     1, 2, 3, 4,       10, 0,     0],
		Some(&vec![2, 2, 2, 2,        2, 1,     0]),
		Some(&vec![0, 1, 1, 1,        0, 1,     0])
	).unwrap();
	col_nested_x.close().unwrap();

	let now_ms = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_millis() as i64;
	let mut y = rg.next_column().unwrap().unwrap();
	y.typed::<Int64Type>().write_batch(
		&vec![     now_ms,   0,     0],
		Some(&vec![2,        1,     0]),
		None
	).unwrap();
	y.close().unwrap();
	rg.close().unwrap();
	writer.close().unwrap();

	// todo!()
}


trait IntoArc<T> {
	fn arc(self) -> Arc<T>;
}
impl <T> IntoArc<T> for T {
	fn arc(self) -> Arc<T> {
		Arc::new(self)
	}
}