From d20460b1996d900dd12f76ec0eee42573bb5ed81 Mon Sep 17 00:00:00 2001 From: Gleb Krivosheev - RMS <52839573+GKrivosheev-rms@users.noreply.github.com> Date: Sat, 5 Feb 2022 00:03:26 -0800 Subject: [PATCH 1/2] [Issue #4] Make writng of DataFrames more flexible - Added overload extension with ability to write an IEnumerable to a single file - Add ability to write DataFrame re-entrrantly to the same file - Make Nullability of columns depending on underlying type (stable despite changes to row data) --- ParquetSharp.DataFrame/DataFrameExtensions.cs | 98 +++++++++++++------ 1 file changed, 70 insertions(+), 28 deletions(-) diff --git a/ParquetSharp.DataFrame/DataFrameExtensions.cs b/ParquetSharp.DataFrame/DataFrameExtensions.cs index 745ce05..4de0376 100644 --- a/ParquetSharp.DataFrame/DataFrameExtensions.cs +++ b/ParquetSharp.DataFrame/DataFrameExtensions.cs @@ -8,9 +8,35 @@ namespace ParquetSharp public static class DataFrameExtensions { /// - /// Write a DataFrame in Parquet format + /// Writes a Dataframe in Parquet format using existing writer. /// - /// The DataFrame to write + /// DataFrame to write. + /// Writer to use. + /// Maximum number of rows per row group + public static void Write(DataFrame dataFrame, ParquetFileWriter fileWriter, int rowGroupSize = 1024 * 1024) + { + long numRows = dataFrame.Rows.Count; + long offset = 0L; + + while (offset < numRows) + { + int batchSize = (int)Math.Min(numRows - offset, rowGroupSize); + using RowGroupWriter rowGroupWriter = fileWriter.AppendRowGroup(); + foreach (DataFrameColumn dataFrameColumn in dataFrame.Columns) + { + using ColumnWriter columnWriter = rowGroupWriter.NextColumn(); + using LogicalColumnWriter logicalWriter = columnWriter.LogicalWriter(); + logicalWriter.Apply(new DataFrameWriter(dataFrameColumn, offset, batchSize)); + } + + offset += batchSize; + } + } + + /// + /// Write DataFrames in Parquet format + /// + /// DataFrames to write /// Path to write to /// Optional writer properties that override the default properties /// Mapping from column names to Parquet logical types, @@ -18,48 +44,64 @@ public static class DataFrameExtensions /// to specify the precision and scale to use. /// Maximum number of rows per row group public static void ToParquet( - this DataFrame dataFrame, string path, WriterProperties? writerProperties = null, + this IEnumerable dataFrames, string path, WriterProperties? writerProperties = null, IReadOnlyDictionary? logicalTypeOverrides = null, int rowGroupSize = 1024 * 1024) { - var schemaColumns = dataFrame.Columns.Select(col => GetSchemaColumn( - col, logicalTypeOverrides != null && logicalTypeOverrides.TryGetValue(col.Name, out var logicalType) ? logicalType : null)).ToArray(); - using var fileWriter = writerProperties == null - ? new ParquetFileWriter(path, schemaColumns) - : new ParquetFileWriter(path, schemaColumns, writerProperties); + if (dataFrames is null) + throw new ArgumentNullException(nameof(dataFrames)); - var numRows = dataFrame.Rows.Count; - var offset = 0L; + DataFrame firstDataFrame = dataFrames.First(); + using ParquetFileWriter fileWriter = GetParquetFileWriter(path, writerProperties, logicalTypeOverrides, firstDataFrame); - while (offset < numRows) + foreach (DataFrame dataFrame in dataFrames) { - var batchSize = (int)Math.Min(numRows - offset, rowGroupSize); - using var rowGroupWriter = fileWriter.AppendRowGroup(); - foreach (var dataFrameColumn in dataFrame.Columns) - { - using var columnWriter = rowGroupWriter.NextColumn(); - using var logicalWriter = columnWriter.LogicalWriter(); - logicalWriter.Apply(new DataFrameWriter(dataFrameColumn, offset, batchSize)); - } - - offset += batchSize; + Write(dataFrame, fileWriter, rowGroupSize); } fileWriter.Close(); } + /// + /// Write a DataFrame in Parquet format + /// + /// The DataFrame to write + /// Path to write to + /// Optional writer properties that override the default properties + /// Mapping from column names to Parquet logical types, + /// overriding the default logical types. When writing decimal columns, a logical type must be provided + /// to specify the precision and scale to use. + /// Maximum number of rows per row group + public static void ToParquet(this DataFrame dataFrame, string path, WriterProperties? writerProperties = null, + IReadOnlyDictionary? logicalTypeOverrides = null, int rowGroupSize = 1024 * 1024) + { + using ParquetFileWriter fileWriter = GetParquetFileWriter(path, writerProperties, logicalTypeOverrides, dataFrame); + Write(dataFrame, fileWriter, rowGroupSize); + + fileWriter.Close(); + } + + private static ParquetFileWriter GetParquetFileWriter( + string path, WriterProperties writerProperties, + IReadOnlyDictionary logicalTypeOverrides, DataFrame firstDataFrame) + { + Column[] schemaColumns = firstDataFrame.Columns.Select(col => GetSchemaColumn( + col, + (logicalTypeOverrides != null && logicalTypeOverrides.TryGetValue(col.Name, out LogicalType logicalType)) + ? logicalType : null)).ToArray(); + + return writerProperties == null + ? new ParquetFileWriter(path, schemaColumns) + : new ParquetFileWriter(path, schemaColumns, writerProperties); + } + private static Column GetSchemaColumn(DataFrameColumn column, LogicalType? logicalTypeOverride) { - var dataType = column.DataType; - var nullable = column.NullCount > 0; + Type dataType = column.DataType; if (dataType == typeof(decimal) && logicalTypeOverride == null) - { throw new ArgumentException($"Logical type override must be specified for decimal column '{column.Name}'"); - } - if (nullable && dataType.IsValueType) - { + if (dataType.IsValueType && Nullable.GetUnderlyingType(dataType) != null) dataType = typeof(Nullable<>).MakeGenericType(dataType); - } return new Column(dataType, column.Name, logicalTypeOverride); } From e18577670819981f9a69977f177f3152b03bcd2f Mon Sep 17 00:00:00 2001 From: Gleb Krivosheev - RMS <52839573+GKrivosheev-rms@users.noreply.github.com> Date: Sat, 5 Feb 2022 00:10:19 -0800 Subject: [PATCH 2/2] Turn the write into an extension method. --- ParquetSharp.DataFrame/DataFrameExtensions.cs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ParquetSharp.DataFrame/DataFrameExtensions.cs b/ParquetSharp.DataFrame/DataFrameExtensions.cs index 4de0376..74df692 100644 --- a/ParquetSharp.DataFrame/DataFrameExtensions.cs +++ b/ParquetSharp.DataFrame/DataFrameExtensions.cs @@ -13,7 +13,7 @@ public static class DataFrameExtensions /// DataFrame to write. /// Writer to use. /// Maximum number of rows per row group - public static void Write(DataFrame dataFrame, ParquetFileWriter fileWriter, int rowGroupSize = 1024 * 1024) + public static void ToParquet(this DataFrame dataFrame, ParquetFileWriter fileWriter, int rowGroupSize = 1024 * 1024) { long numRows = dataFrame.Rows.Count; long offset = 0L; @@ -55,7 +55,7 @@ public static void ToParquet( foreach (DataFrame dataFrame in dataFrames) { - Write(dataFrame, fileWriter, rowGroupSize); + dataFrame.ToParquet(fileWriter, rowGroupSize); } fileWriter.Close(); @@ -75,7 +75,7 @@ public static void ToParquet(this DataFrame dataFrame, string path, WriterProper IReadOnlyDictionary? logicalTypeOverrides = null, int rowGroupSize = 1024 * 1024) { using ParquetFileWriter fileWriter = GetParquetFileWriter(path, writerProperties, logicalTypeOverrides, dataFrame); - Write(dataFrame, fileWriter, rowGroupSize); + dataFrame.ToParquet(fileWriter, rowGroupSize); fileWriter.Close(); }